import os
import nibabel as nib
import numpy as np
import json

'''
DCM,HCM,MINF,NOR,RV
'''

def load_nii(nii_path):
    data = nib.load(nii_path)
    img = data.get_data()
    affine = data.affine
    header = data.header
    return img ,affine,header


def read_Infocfg(cfg_path):
    patient_info = {}




    with open (cfg_path) as f_in:
        for line in f_in:
            l = line.rstrip().split(":")
            #l is the list of the patient_info
            print(len(l))
            patient_info[l[0]] = l[1]
        # print(patient_info)
        '''
        ['ED', ' 1']
        ['ES', ' 12']
        ['Group', ' DCM']
        ['Height', ' 184.0']
        ['NbFrame', ' 30']
        ['Weight', ' 95.0']
        {'ED': ' 1', 'ES': ' 12', 'Group': ' DCM', 'Height': ' 184.0', 'NbFrame': ' 30', 'Weight': ' 95.0'}
        '''
    return patient_info

def read_json(fpath):
    with open(fpath,'r') as f:
        obj = json.load(f)
    return obj

def write_json(obj, fpath):
    with open(fpath,'w') as f:
        json.dump(obj,f,indent=4)

def write_json_append(obj, fpath):
    with open(fpath,'a') as f:
        json.dump(obj,f,indent=4)

def gen_alldatalist(path):
    filelist = []
    for dir in os.listdir(path):

        for file in os.listdir(os.path.join(root_path,dir)):
            filelist.append(os.path.join(root_path,dir,file))
    #the length of the filelist is 1902
    # print(len(filelist))
    filelist.sort()
    out_dir = os.path.dirname(os.path.abspath(__file__))
    #/home/ffbian/chencheng/XieheCardiac/2DUNet/UNet/libs/datasets
    write_json(filelist, os.path.join(out_dir, "./acdcjson/ACDCDataList.json"))

def gen_every_kind_datalist(kind_path):
    filelist = []

    for file in os.listdir(kind_path):
        filelist.append(os.path.join(kind_path,file))

    filelist.sort()
    out_dir = os.path.dirname(os.path.abspath(__file__))
    #/home/ffbian/chencheng/XieheCardiac/2DUNet/UNet/libs/datasets
    write_json(filelist, os.path.join(out_dir, "./acdcjson/{}DataList.json".format(kind_path[-2:])))

def generate_train_test_list(json_file_path):
    # json_file = "/home/fcheng/Cardia/DataList.json"
    # json_file = "/home/ffbian/chencheng/XieheCardiac/2DUNet/UNet/libs/datasets/ACDCDataList.json"
    fileslist = read_json(json_file_path)

    nums = len(fileslist)
    train_ind = set(np.random.choice(nums, size=int(np.ceil(0.8 * nums)), replace=False))
    test_ind = set(np.arange(nums)) - train_ind

    test_ind = list(test_ind)
    test_ind.sort()

    train_list = [fileslist[fl] for fl in train_ind]
    test_list = [fileslist[fl] for fl in test_ind]

    out_dir = os.path.dirname(os.path.abspath(__file__))
    print(out_dir)
    write_json(train_list, os.path.join(out_dir, "./acdcjson/RVtrain.json"))
    write_json(test_list, os.path.join(out_dir, "./acdcjson/RVtest.json"))
    write_json_append(train_list, os.path.join(out_dir, "./acdcjson/train.json"))
    write_json_append(test_list, os.path.join(out_dir, "./acdcjson/test.json"))
    # write_json(test_list, "/home/ffbian/chencheng/XieheCardiac/2DUNet/UNet/libs/datasets/differentkind/kuodatest.json")


if __name__ == "__main__":
    # read_Infocfg(os.path.join(root_path,"./DCM/patient001/Info.cfg"))
    '''
    a,b,c = load_nii(os.path.join(root_path,"./patient001/patient001_frame01.nii.gz"))
    print(np.unique(a))

    print(c)
    '''
    root_path = "/home/ffbian/chencheng/MICCAIACDC2017/processed_acdc_dataset/hdf5_files/Bykind"
    kind_path = "/home/ffbian/chencheng/MICCAIACDC2017/processed_acdc_dataset/hdf5_files/Bykind/RV"

    # gen_alldatalist(root_path)
    json_file_path = "/home/ffbian/chencheng/MICCAIACDC2017/mycode/libs/dataset/acdcjson/RVDataList.json"

    # generate_train_test_list(json_file_path)
    # gen_every_kind_datalist(kind_path)
    result = read_json(json_file_path)
    print(result)