"""
Data-loading helpers for the HaN dose-prediction models.

Trimmed from the original research toolkit: only the functions the released
data loader and inference script actually call are kept, which drops the
cv2 / matplotlib / SimpleITK plotting dependencies.

Derived from the data pipeline of:
    Riqiang Gao, Bin Lou, Zhoubing Xu, Dorin Comaniciu, Ali Kamen.
    "Flexible-CM GAN: Towards Precise 3D Dose Prediction in Radiotherapy." CVPR 2023.
"""

import os

import numpy as np
import torch
import nibabel as nib
from monai.transforms import (
    Compose,
    Resized,
    RandFlipd,
    RandRotated,
    SpatialPadd,
    SpatialCropd,
    RandSpatialCropd,
)


def combine_oar(tmp_dict, need_list,norm_oar = True, OAR_DICT = None):

    '''
    this function is used to support the data loader. 

    tmp_dict: the dictionary of the data loaded from the npz file
    need_list: the list of the OARs needed to be combined
    norm_oar: if True, the OARs will be normalized to the same scale
    OAR_DICT: the dictionary of the OARs, the key is the name of the OAR, the value is the index of the OAR in the combined
    '''
    
    comb_oar = torch.zeros(tmp_dict['img'].shape)  
    cat_oar = torch.zeros([len(OAR_DICT)] + list(tmp_dict['img'].shape)[1:])
    for key in OAR_DICT.keys():
        
        if key not in need_list:
            continue

        if key in tmp_dict.keys():
            single_oar = tmp_dict[key]
        else:
            single_oar = torch.zeros(tmp_dict['img'].shape)

        cat_oar[OAR_DICT[key]-1: OAR_DICT[key]]  = single_oar

        if norm_oar:
            comb_oar = torch.maximum(comb_oar, single_oar.round() * (1.0 + 4.0 * OAR_DICT[key] / 30))
        else:
            comb_oar = torch.maximum(comb_oar, single_oar.round() * OAR_DICT[key])  # changed in this version

    return comb_oar, cat_oar


def combine_ptv(tmp_dict,  scaled_dose_dict):

    prescribed_dose = []
    
    cat_ptv = torch.zeros([3] + list(tmp_dict['img'].shape)[1:])
    prescribed_dose = [0] * 3

    comb_ptv = torch.zeros(tmp_dict['img'].shape)

    cnt = 0
    for key in scaled_dose_dict.keys():
       
       tmp_ptv = tmp_dict[key] * scaled_dose_dict[key]  
       
       prescribed_dose[cnt] =  scaled_dose_dict[key] 
       
       cat_ptv[cnt] = tmp_ptv 
       comb_ptv = torch.maximum(comb_ptv, tmp_ptv)

       cnt += 1

    # sort the cat_ptv according to the prescribed dose
    paired = [(cat_ptv[i], prescribed_dose[i]) for i in range(len(prescribed_dose))]
    paired_sorted = sorted(paired, key=lambda x: x[1], reverse=True)
    cat_ptv = torch.stack([x[0] for x in paired_sorted])
    prescribed_dose = [x[1] for x in paired_sorted]
    
    return comb_ptv, prescribed_dose, cat_ptv 


def tr_augmentation(KEYS, in_size, out_size, crop_center):
    return Compose([
        SpatialCropd(keys = KEYS, roi_center = crop_center, roi_size = [int(in_size[0] * 1.2), int(in_size[1] * 1.2), int(in_size[2] * 1.2)], allow_missing_keys = True),
        SpatialPadd(keys = KEYS, spatial_size = [int(in_size[0] * 1.2), int(in_size[1] * 1.2), int(in_size[2] * 1.2)], mode = 'constant', allow_missing_keys = True),
        RandSpatialCropd(keys = KEYS, roi_size = [int(in_size[0] * 0.85), int(in_size[1] * 0.85), int(in_size[2] * 0.85)], max_roi_size = [int(in_size[0] * 1.2), int(in_size[1] * 1.2), int(in_size[2] * 1.2)], random_center = True, random_size = True, allow_missing_keys = True), 
        RandRotated(keys = KEYS, prob=0.8, range_x= 1, range_y = 0.2, range_z = 0.2, allow_missing_keys = True),
        RandFlipd(keys = KEYS, prob = 0.4, spatial_axis = 0, allow_missing_keys = True), 
        RandFlipd(keys = KEYS, prob = 0.4, spatial_axis = 1, allow_missing_keys = True), 
        RandFlipd(keys = KEYS, prob = 0.4, spatial_axis = 2, allow_missing_keys = True), 
        Resized(keys = KEYS, spatial_size = out_size, allow_missing_keys = True), 
    ])

def tt_augmentation(KEYS, in_size, out_size, crop_center):
    return Compose([
        SpatialCropd(keys = KEYS, roi_center = crop_center, roi_size = in_size, allow_missing_keys = True),
        SpatialPadd(keys = KEYS, spatial_size = in_size, mode = 'constant', allow_missing_keys = True),
        Resized(keys = KEYS, spatial_size = out_size, allow_missing_keys = True), 
    ])

def load_nifti(path):
    """
    Load a .nii or .nii.gz file and return a numpy array.
    torch shape: [B, C, Z, Y, X] where B=1, C=1, Z=depth, Y=height, X=width
    """
    nii = nib.load(path)
    return np.transpose(nii.get_fdata(), (2, 1, 0)).astype(np.float32)