import os
import os.path
from os import environ
import sys
import json
import subprocess
import time
import nibabel as nib

# +++++++++++++ Conversion imports +++++++++++++++++++++++++
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.path.abspath(".."))
# +++++++++++++ Conversion imports +++++++++++++++++++++++++

from utils import *
from dicom_to_nii import convert_ct_dicom_to_nii, convert_transform_mr_to_nii, PatientList, save_images
from nii_to_dicom import convert_nii_to_dicom, integer_to_onehot
from predict_nnunet import predictNNUNet

def predict(tempPath, patient_id, regSeriesInstanceUID, runInterpreter):
    
    # Important: Check the input parameters #################
    if not patient_id or patient_id == "":
        sys.exit("No Patient dataset loaded: Load the patient dataset in Study Management.")
        
    if not regSeriesInstanceUID or regSeriesInstanceUID == "":
        sys.exit("No series instance UID for Modality 'REG' file. Check for REG file in your study")
        
    dir_base = os.path.join(tempPath, patient_id)
    createdir(dir_base)
    
    dir_ct_dicom = os.path.join(dir_base, 'ct_dicom')
    createdir(dir_ct_dicom)

    dir_mr_dicom = os.path.join(dir_base, 'mr_dicom')
    createdir(dir_mr_dicom)
    
    dir_reg_dicom = os.path.join(dir_base, 'reg_dicom')
    createdir(dir_reg_dicom)
        
    nnUNet_raw = os.path.join(os.getcwd(), 'nnUNet_raw')    
    nnUNet_preprocessed = os.path.join(os.getcwd(), 'nnUNet_preprocessed')
    RESULTS_FOLDER = os.path.join(os.getcwd(), 'nnUNet_trained_models')
    dataset = "Dataset103_EPTN_T1_CT_all_structures" 
    # IMPORTANT: data set modality: MR or CT ######################
    predictType='MR' 

    # IMPORTANT DOT Remove ########################################
    os.environ['nnUNet_raw'] = nnUNet_raw
    os.environ['nnUNet_preprocessed'] = nnUNet_preprocessed
    os.environ['nnUNet_results'] = RESULTS_FOLDER
    
    # Important ++++++++++++++++++++++++++++++++++++++++++++++++
    # Import the lib after setting environ parameters
    # import nnunet.inference.predict_simple as nnunetpredict
    
    print('** The python enviornment path: ', os.environ["PATH"])
    
    # For nnunet version 2
    import nnunetv2.inference.predict_from_raw_data as nnunetpredict
    # ###########################################################

    # predicted files
    predictedNiiFile = os.path.join(tempPath, patient_id, 'predict_nii')
    createdir(predictedNiiFile)
    
    predictedDicom = os.path.join(tempPath, patient_id, 'predicted_dicom')
    createdir(predictedDicom)
    
    predictedDicomFile = os.path.join(predictedDicom, 'predicted_rtstruct.dcm')
        
    print('** Use python interpreter: ', runInterpreter)
    print('** Patient name: ', patient_id)
    print('** REG series instance UID: ', regSeriesInstanceUID)
    
    # Convert CT image to NII #############
    startTime = time.time()
    
    if predictType == 'CT':
        
        dir_dicom_to_nii = os.path.join(nnUNet_raw, 'nnUNet_raw_data', 'Dataset098_HAN_nodes')
        createdir(dir_dicom_to_nii)
    
        downloadSeriesInstanceByModality(instanceID, dir_ct_dicom, "CT")
        print("Loading CT from Orthanc done: ", time.time()-startTime)
        
        # Convert CT image to NII #############
        refCT= convert_ct_dicom_to_nii(dir_dicom=dir_ct_dicom, dir_nii=dir_dicom_to_nii, outputname='1a_001_0000.nii.gz', newvoxelsize = None)
        print("Convert CT image to NII Done: ",  time.time()-startTime)
        
        # new version 2:  
        cmd = [modelPath, '-i', dir_dicom_to_nii, '-o', predictedNiiFile, '-d', dataset, '-tr', 'nnUNetTrainer_650epochs', '-c', '3d_fullres', '-f', '0']
        
        out = subprocess.check_output(cmd)
        # Important ########################
        sys.argv = cmd
        
        # #### nnunet version 2 #############
        nnunetpredict.predict_entry_point() 
        print("Prediction CT done",  time.time()-startTime)
        
        niiFile = os.path.join(predictedNiiFile, '1a_001.nii.gz')
        
        # POSTPROCESSING TO CONVERT FROM INTEGERS TO 2**i, ADD CONTOURS EXISTS, AND SMOOTH        
        integer_to_onehot(niiFile)
        print("POST processing convert from integers done: ", time.time()-startTime)
        
        startTime = time.time()
        convert_nii_to_dicom(dicomctdir=dir_ct_dicom, predictedNiiFile=niiFile, predictedDicomFile=predictedDicomFile,
                             predicted_structures=predicted_structures, rtstruct_colors=rtstruct_colors, refCT=refCT)
        
        print("Convert CT predicted NII to DICOM done: ", time.time()-startTime)
                
    elif predictType == 'MR':
        
        dir_dicom_to_nii = os.path.join(nnUNet_raw, 'nnUNet_raw_data',dataset)
        createdir(dir_dicom_to_nii)
        
        # Download the REG dicom ##############
        downloadSeriesInstanceByModality(regSeriesInstanceUID, dir_reg_dicom, "REG")
        print("Loading REG from Orthanc done: ", time.time()-startTime)
        
        # Download the MR dicom ###############
        # Read the mr study instance UID from the download REG dicom   
        mrSeriesInstanceUID = getSeriesInstanceUIDFromRegDicom(dir_reg_dicom, regSeriesInstanceUID)
        
        downloadSeriesInstanceByModality(mrSeriesInstanceUID, dir_mr_dicom, "MR")
        print("Loading MR from Orthanc done: ", time.time()-startTime)
        
        # Execute REG tranformation ###########
        ctSeriesInstanceUIDFromRegDicom = getCTSeriesInstanceUIDFromRegDicom(dir_reg_dicom, regSeriesInstanceUID)
        print("CT Series Instance UID referenced by Reg dicom: ", ctSeriesInstanceUIDFromRegDicom)
        
        downloadSeriesInstanceByModality(ctSeriesInstanceUIDFromRegDicom, dir_ct_dicom, "CT")
        
        Patients = PatientList() 
        Patients.list_dicom_files(dir_ct_dicom, 1)
        patient = Patients.list[0]
        patient_name = patient.PatientInfo.PatientName
        patient.import_patient_data(newvoxelsize=None)
        CT = patient.CTimages[0]

        startTime = time.time()
        mr_reg = regMatrixTransformation(dir_mr_dicom, reg_file_path=dir_reg_dicom, regSeriesInstanceUID=regSeriesInstanceUID, CT=CT)
        print("Transforming MR data done (OpenTPS.Core)")
        
        # Convert transform MR image to NII ##################
        refMR = convert_transform_mr_to_nii(dir_mr_dicom=dir_mr_dicom, tranform_mr = mr_reg, dir_nii=dir_dicom_to_nii, outputname='1a_001_0000.nii.gz', CT=CT)
        refCT= convert_ct_dicom_to_nii(dir_dicom=dir_ct_dicom, dir_nii=dir_dicom_to_nii, outputname='1a_001_0001.nii.gz', newvoxelsize = None)
        print("Convert CT image to NII Done: ",  time.time()-startTime)
        print("Convert transform MR image to NII Done: ",  time.time()-startTime)
        
        
        print("## start MR running prediction ###############")
        startTime = time.time()
        # modelPath = '..\\..\\python_environments\\prediction-3.10.9\\Scripts\\nnUNetv2_predict.exe' 
        # cmd = [modelPath, '-i', dir_dicom_to_nii, '-o', predictedNiiFile, '-d', '99', '-c', '3d_fullres' , '--disable_tta', '-tr', 'nnUNetTrainer_650epochs', '-f', '1, 4']
        
        predictNNUNet(os.path.join(RESULTS_FOLDER,dataset, 'nnUNetTrainer_650epochs__nnUNetPlans__3d_fullres'),
                      dir_dicom_to_nii,
                      predictedNiiFile,
                      [1]) 
        
        print("Prediction MR done",  time.time()-startTime)
        
        startTime = time.time()

        predicted_structures = ["background", "BRAIN", "AMYGDALAE", "BRAINSTEM", "CAUDATENUCLEI", "CEREBELLUM", "CHIASM", "COCHLEAS", "CORNEAS", "CORPUSCALLOSUM", "FORNICES", "GLANDPINEAL", "HIPPOCAMPI", "HYPOTHALAMI", "LACRIMALGLANDS", "LENSES", "OPTICNERVES", "ORBITOFRONTALS", "PITUITARY", "RETINAS", "THALAMI", "VSCCs"]
        rtstruct_colors = [[255,0,0]]*len(predicted_structures)
        
        niiFile = os.path.join(predictedNiiFile, '1a_001.nii.gz')
        
        # POSTPROCESSING TO CONVERT FROM INTEGERS TO 2**i, ADD CONTOURS EXISTS, AND SMOOTH        
        integer_to_onehot(niiFile)
        print("POST processing convert from integers done: ", time.time()-startTime)
        
        # Convert CT image to NII #############
        
        
        
        convert_nii_to_dicom(dicomctdir=dir_ct_dicom, predictedNiiFile=niiFile, predictedDicomFile=predictedDicomFile, 
                            predicted_structures=predicted_structures, rtstruct_colors=rtstruct_colors, refCT=refCT)
    else:
        print("Not supported yet")
        
    
    startTime = time.time()
    uploadDicomToOrthanc(predictedDicomFile)
    print("Upload predicted result to Orthanc done: ",  time.time()-startTime)

# tempPath = 'C:\Temp\parrot_prediction'
# regSeriesInstanceUID = '1.2.246.352.205.5029381855449574337.1508502639685232062'
# runInterpreter = 'py3109'
# patientName = 'P0461C0006I7638639'

'''
Prediction parameters provided by the server. Select the parameters to be used for prediction:
    [1] tempPath: The path where the predict.py is stored,
    [2] patientname: python version,
    [3] ctSeriesInstanceUID: Series instance UID for data set with modality = CT. To predict 'MR' modality data, retrieve the CT UID by the code (see Precision Code)
    [4] rtStructSeriesInstanceUID: Series instance UID for modality = RTSTURCT
    [5] regSeriesInstanceUID: Series instance UID for modality = REG,
    [6] runInterpreter: The python version for the python environment
    [7] oarList: only for dose predciton. For contour predicion oarList = []
    [8] tvList:  only for dose prediction. For contour prediction tvList = []
''' 
if __name__ == '__main__':
    predict(tempPath=sys.argv[1], patient_id=sys.argv[2], regSeriesInstanceUID=sys.argv[5], runInterpreter=sys.argv[6])
    # predict(tempPath=tempPath, patient_id=patientName, regSeriesInstanceUID=regSeriesInstanceUID, runInterpreter=runInterpreter)