from __future__ import generators

import logging
import glob, os, functools
import sys
sys.path.append('../')


import SimpleITK as sitk
import numpy as np
import scipy
import nibabel as nib
import skimage
import matplotlib.pyplot as plt
import scipy.misc
from scipy import ndimage
from skimage.transform import resize,rescale
import cv2
import itk
import subprocess
from tqdm import tqdm
import pandas as pd
import warnings
import statistics
import torch
import csv
import os
import yaml

from HD_BET.run import run_hd_bet # git clone HDBET repo
from dataset.preprocess_utils import enhance, enhance_noN4
from dataset.preprocess_datasets_T1_to_2d import create_quantile_from_brain

warnings.filterwarnings('ignore')
cuda_device = '1'
os.environ['CUDA_VISIBLE_DEVICES'] = cuda_device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#torch.cuda.set_device(1)  # Set CUDA device to 0 (first GPU)
#net = net.cuda()

def select_template_based_on_age(age):
    for golden_file_path, age_values in age_ranges.items():
        if age_values['min_age'] <= int(age) and int(age) <= age_values['max_age']: 
            print(golden_file_path)
            return golden_file_path
        
def register_to_template(input_image_path, output_path, fixed_image_path,rename_id,create_subfolder=True):
    fixed_image = itk.imread(fixed_image_path, itk.F)

    # Import Parameter Map
    parameter_object = itk.ParameterObject.New()
    parameter_object.AddParameterFile('/media/sdb/divyanshu/divyanshu/aidan_segmentation/pediatric-brain-age-main/dataset/golden_image/mni_templates/Parameters_Rigid.txt')

    if "nii" in input_image_path and "._" not in input_image_path:
        print(input_image_path)

        # Call registration function
        try:        
            moving_image = itk.imread(input_image_path, itk.F)
            result_image, result_transform_parameters = itk.elastix_registration_method(
                fixed_image, moving_image,
                parameter_object=parameter_object,
                log_to_console=False)
            image_id = input_image_path.split("/")[-1]
            
            itk.imwrite(result_image, output_path+"/"+rename_id+".nii.gz")
                
            print("Registered ", rename_id)
        except:
            print("Cannot transform", rename_id)
            
def outlier_voting(numbers):
    mean = statistics.mean(numbers)
    stdev = statistics.stdev(numbers)

    threshold = stdev # *2  #*3
    
    good_nums_avg =[]
    for n in numbers:
        if n > mean + threshold or n < mean - threshold:
            continue
        else:
            good_nums_avg.append(n)
    
    #if len(good_nums_avg)<=3:
    #    print(len(good_nums_avg))
    return np.average(good_nums_avg)


#Data https://openneuro.org/datasets/ds000228/versions/1.1.0
img_path = '/media/sdb/divyanshu/divyanshu/aidan_segmentation/dummy_t1_preprocess/OAS2_0001_MR1.nii.gz'
data_path = "/media/sdb/divyanshu/divyanshu/longitudinal_fm/datasets/abide/data"
gt_age = 86  # age of subject
gender = "M" # gender
path_to = "/media/sdb/divyanshu/divyanshu/longitudinal_fm/datasets/abide/preprocessed_data" # save to

# MNI templates http://nist.mni.mcgill.ca/pediatric-atlases-4-5-18-5y/
age_ranges = {"/media/data/BrainIAC/src/BrainIAC/golden_image/mni_templates/nihpd_asym_04.5-08.5_t1w.nii" : {"min_age":3, "max_age":7},
                "/media/data/BrainIAC/src/BrainIAC/golden_image/mni_templates/nihpd_asym_07.5-13.5_t1w.nii": {"min_age":8, "max_age":13},
                "/media/data/BrainIAC/src/BrainIAC/golden_image/mni_templates/nihpd_asym_13.0-18.5_t1w.nii": {"min_age":14, "max_age":100}}


for eachimage in tqdm(os.listdir(data_path), desc="Processing images", unit="image"):
    if 1:#"sub" in eachimage:
        ## load image
        img_path = os.path.join(data_path, eachimage)
        nii= nib.load(img_path)
        image, affine  = nii.get_fdata(), nii.affine
        #plt.imshow(image[:,:,100])
        #print(nib.aff2axcodes(affine))

        # path to store registered image in
        new_path_to = path_to#path_to+"/"+img_path.split("/")[-1].split(".")[0]
        if eachimage in os.listdir(new_path_to):
            print("yay")
        else:
            if not os.path.exists(path_to):
                os.mkdir(path_to)
            if not os.path.exists(new_path_to):
                os.mkdir(new_path_to)

            # register image to MNI template
            golden_file_path = select_template_based_on_age(gt_age)
            print("Registering to template:", golden_file_path)
            #fun fact: the registering to the template pipeline is not deterministic
            register_to_template(img_path, new_path_to, golden_file_path,eachimage.split(".")[0]+"_"+"registered.nii.gz", create_subfolder=False)


            # enchance and normalize image
            #if not os.path.exists(new_path_to+"/no_z"):
            #    os.mkdir(new_path_to+"/no_z")
                
            image_sitk =  sitk.ReadImage(os.path.join(new_path_to, eachimage.split(".")[0]+"_"+"registered.nii.gz"))
            image_array  = sitk.GetArrayFromImage(image_sitk)
            image_array = enhance(image_array) # or enhance_noN4(image_array) if no bias field correction is needed
            image3 = sitk.GetImageFromArray(image_array)
            sitk.WriteImage(image3,os.path.join(new_path_to, eachimage.split(".")[0]+"_"+"registered_no_z.nii.gz"))

            #skull strip   ## when running this with rest of the preprocessing, change the src path to include the registered image path!!!!
            new_path_to = path_to
            run_hd_bet(os.path.join(new_path_to, eachimage.split(".")[0]+"_"+"registered_no_z.nii.gz"),os.path.join(new_path_to, eachimage),                      
                        mode="accurate", 
                        config_file='/media/sdb/divyanshu/divyanshu/aidan_segmentation/pediatric-brain-age-main/HD_BET/config.py',
                        device=device,
                        postprocess=False,
                        do_tta=True,
                        keep_mask=True, 
                        overwrite=True)