# -*- coding: utf-8 -*-
"""
Created on Sun Sep 15 17:19:21 2024

@author: salikha4
"""

import requests
import zipfile
import os
import pickle

def load_DeepMIMO_data(zip_file_name='dataset.zip', p_file_name='deepmimo_data.p', extract_path='./data'):
    url = "https://huggingface.co/datasets/sadjadalikhani/lwm/resolve/main/dataset.zip"
    # Step 1: Download the ZIP file from Hugging Face
    print(f"Downloading ZIP file from {url}...")
    response = requests.get(url)
    with open(zip_file_name, 'wb') as f:
        f.write(response.content)
    print(f"Downloaded ZIP file to {zip_file_name}.")

    # Step 2: Unzip the file
    print(f"Extracting ZIP file to {extract_path}...")
    with zipfile.ZipFile(zip_file_name, 'r') as zip_ref:
        zip_ref.extractall(extract_path)
    print(f"Extracted ZIP file contents to {extract_path}.")

    # Step 3: Load the .p file
    p_file_path = os.path.join(extract_path, p_file_name)
    print(f"Loading .p file from {p_file_path}...")
    with open(p_file_path, 'rb') as f:
        data = pickle.load(f)
    
    print("Data successfully loaded from the .p file.")
    return data