# hello fellow human, this script is used to save kspace data to disk
# You may ask why? Well, as it turns out having h5py read the entire .h5 file
# and then just accessing the kspace data as numpy array takes around 50 seconds for a single file
# and that's just too slow for me. So I'm going to save the kspace data to disk as numpy arrays

import h5py
import huggingface_hub as hfh
import numpy as np
# datasets
# osbm/fastmri-prostate 
# osbm/fastmri-brain
# osbm/fastmri-knee

# files in the dataset
# prostate
# - training_T2_1/file_prostate_AXT2_0002.h5
# - training_T2_1/file_prostate_AXT2_0015.h5
# brain
# - multicoil_train/file_brain_AXFLAIR_200_6002442.h5
# - multicoil_train/file_brain_AXFLAIR_200_6002487.h5
# knee
# - singlecoil_train/file1000015.h5
# - multicoil_train/file1000015.h5

# Download files

file_paths = {
    "prostate1": hfh.hf_hub_download(
        repo_id="osbm/fastmri-prostate",
        filename="training_T2_1/file_prostate_AXT2_0002.h5",
        repo_type="dataset",
        cache_dir="./data"
    ),
    "prostate2": hfh.hf_hub_download(
        repo_id="osbm/fastmri-prostate",
        filename="training_T2_1/file_prostate_AXT2_0015.h5",
        repo_type="dataset",
        cache_dir="./data"
    ),
    "brain1": hfh.hf_hub_download(
        repo_id="osbm/fastmri-brain",
        filename="multicoil_train/file_brain_AXFLAIR_200_6002442.h5",
        repo_type="dataset",
        cache_dir="./data"
    ),
    "brain2": hfh.hf_hub_download(
        repo_id="osbm/fastmri-brain",
        filename="multicoil_train/file_brain_AXFLAIR_200_6002487.h5",
        repo_type="dataset",
        cache_dir="./data"
    ),
    "knee1": hfh.hf_hub_download(
        repo_id="osbm/fastmri-knee",
        filename="singlecoil_train/file1000015.h5",
        repo_type="dataset",
        cache_dir="./data"
    ),
    "knee2": hfh.hf_hub_download(
        repo_id="osbm/fastmri-knee",
        filename="multicoil_train/file1000015.h5",
        repo_type="dataset",
        cache_dir="./data"
    )
}

for key, file_path in file_paths.items():
    print(f"{key}: {file_path}")

    file = h5py.File(file_path, "r")
    kspace = file["kspace"][()]

    print(kspace.shape)
    if key.startswith("prostate"):
        kspace = kspace[0, :, :, :] + kspace[1, :, :, :]

    print(kspace.shape)

    np.save(f"./data/{key}_kspace.npy", kspace)