import os
from monai.data import create_test_image_3d
import nibabel as nib
from monai.bundle.config_parser import ConfigParser
import numpy as np
from monai.apps.auto3dseg import DataAnalyzer

fake_datalist: dict[str, list[dict]] = {
    "testing": [{"image": "val_001.fake.nii.gz"}, {"image": "val_002.fake.nii.gz"}],
    "training": [
        {"fold": 0, "image": "tr_image_001.fake.nii.gz", "label": "tr_label_001.fake.nii.gz"},
        {"fold": 0, "image": "tr_image_002.fake.nii.gz", "label": "tr_label_002.fake.nii.gz"},
        {"fold": 0, "image": "tr_image_003.fake.nii.gz", "label": "tr_label_003.fake.nii.gz"},
        {"fold": 0, "image": "tr_image_004.fake.nii.gz", "label": "tr_label_004.fake.nii.gz"},
        {"fold": 1, "image": "tr_image_005.fake.nii.gz", "label": "tr_label_005.fake.nii.gz"},
        {"fold": 1, "image": "tr_image_006.fake.nii.gz", "label": "tr_label_006.fake.nii.gz"},
        {"fold": 1, "image": "tr_image_007.fake.nii.gz", "label": "tr_label_007.fake.nii.gz"},
        {"fold": 1, "image": "tr_image_008.fake.nii.gz", "label": "tr_label_008.fake.nii.gz"},
        {"fold": 2, "image": "tr_image_009.fake.nii.gz", "label": "tr_label_009.fake.nii.gz"},
        {"fold": 2, "image": "tr_image_010.fake.nii.gz", "label": "tr_label_010.fake.nii.gz"},
        {"fold": 2, "image": "tr_image_011.fake.nii.gz", "label": "tr_label_011.fake.nii.gz"},
        {"fold": 2, "image": "tr_image_012.fake.nii.gz", "label": "tr_label_012.fake.nii.gz"},
    ],
}
work_dir = "./build"
dataroot = "./build/data"

da_output_yaml = os.path.join(work_dir, "datastats.yaml")
data_src_cfg = os.path.join(work_dir, "data_src_cfg.yaml")

if not os.path.isdir(dataroot):
    os.makedirs(dataroot)

if not os.path.isdir(work_dir):
    os.makedirs(work_dir)

# Generate a fake dataset
for d in fake_datalist["testing"] + fake_datalist["training"]:
    im, seg = create_test_image_3d(24, 24, 24, rad_max=10, num_seg_classes=1)
    nib_image = nib.Nifti1Image(im, affine=np.eye(4))
    image_fpath = os.path.join(dataroot, d["image"])
    nib.save(nib_image, image_fpath)

    if "label" in d:
        nib_image = nib.Nifti1Image(seg, affine=np.eye(4))
        label_fpath = os.path.join(dataroot, d["label"])
        nib.save(nib_image, label_fpath)

# write to a json file
fake_json_datalist = os.path.join(dataroot, "fake_input.json")
ConfigParser.export_config_file(fake_datalist, fake_json_datalist)

da = DataAnalyzer(fake_json_datalist, dataroot, output_path=da_output_yaml)
da.get_all_case_stats()