from glob import glob
from os.path import normpath, basename


def get_paths(main_dir):
    all_niftis = glob(main_dir, recursive=True)
    all_niftis = [a for a in all_niftis if "fd" not in a.lower()]
    all_niftis = [a for a in all_niftis if "seg" not in a.lower()]
    t2s = [i for i in all_niftis if "t2" in i.lower() and "tra" in i.lower()] 

    adcs = [i for i in all_niftis if "adc" in i.lower() and ("manual" not in i.lower())] 
    if not adcs:
        adcs = [i for i in all_niftis if "manual_adc" in i.lower()] 

    dwis = [i for i in all_niftis if ("diff" in i.lower() or "dwi" in i.lower())
        and ("b-2000" in i.lower() or "b-1400" in i.lower() or "b-1500" in i.lower())] 
    if not dwis:
        dwis = [i for i in all_niftis if "b-1400" in i.lower()] 

    return t2s, adcs, dwis

# check is pat is available in umcg lesions 2022 clinsig: pat_available
paths = glob('../../datasets/umcg_lesions_2022_clinsig/*.nii.gz')
pat_available = []
for path in paths:
    pat_id = basename(normpath(path))[:-7]
    pat_available.append(pat_id)

# read patients included in R script
pat_id_dir = '../pat_ids.txt'
pat_ids_cs = []
with open(pat_id_dir, 'r') as f:
        pat_ids_cs = [l.strip() for l in f.readlines()]

results = {}
paths_seg = []
paths_t2 = []
paths_dwi = []
paths_adc = []
for pat_id_cs in pat_ids_cs:
    t2s,adcs,dwis = get_paths(f"../../datasets/anonymized_mri/only_nii_directory/{pat_id_cs}/**/*.nii.gz")

    pat_id = pat_id_cs.replace("/","-") 
    if pat_id in pat_available:
        results[pat_id] = "p"
        paths_seg.append(f'/data/pca-rad/datasets/umcg_lesions_2022_clinsig/{pat_id}.nii.gz')
        paths_t2.append(t2s)
        paths_adc.append(adcs)
        paths_dwi.append(dwis)

    else:
        results[pat_id] = "missing"

with open('t2_test.txt','w') as output:
    for item in paths_t2:
        output.write("%s\n" % ''.join(item))

with open('dwi_test.txt','w') as output:
    for item in paths_dwi:
        output.write("%s\n" % ''.join(item))

with open('adc_test.txt','w') as output:
    for item in paths_adc:
        output.write("%s\n" % ''.join(item))