#!/usr/bin/env python
# -*- coding: utf-8 -*-

import pdb

import matplotlib
import pydicom
import pymedphys
import numpy as np
from pymedphys import gamma
import matplotlib.pyplot as plt
import pymedphys._dicom.structure as struct
from pymedphys._dicom.coords import xyz_axes_from_dataset
from pymedphys._dicom.dose import get_dose_grid_structure_mask, dose_from_dataset


def get_dose_grid_structure_mask(
        structure_name: str,
        structure_dataset: "pydicom.Dataset",
        dose_dataset: "pydicom.Dataset",
):
    """Determines the 3D boolean mask defining whether or not a grid
    point is inside or outside of a defined structure.

    In its current implementation the dose grid and the planes upon
    which the structures are defined need to be aligned. This is due to
    the implementation only stepping through each structure plane and
    undergoing a 2D mask on the respective dose grid. In order to
    undergo a mask when the contours and dose grids do not align
    inter-slice contour interpolation would be required.

    For now, having two contours for the same structure name on a single
    slice is also not supported.

    Parameters
    ----------
    structure_name
        The name of the structure for which the mask is to be created
    structure_dataset : pydicom.Dataset
        An RT Structure DICOM object containing the respective
        structures.
    dose_dataset : pydicom.Dataset
        An RT Dose DICOM object from which the grid mask coordinates are
        determined.

    Raises
    ------
    ValueError
        If an unsupported contour is provided or the dose grid does not
        align with the structure planes.

    """
    x_dose, y_dose, z_dose = xyz_axes_from_dataset(dose_dataset)

    xx, yy = np.meshgrid(x_dose, y_dose)
    points = np.swapaxes(np.vstack([xx.ravel(), yy.ravel()]), 0, 1)

    x_structure, y_structure, z_structure = struct.pull_structure(
        structure_name, structure_dataset
    )

    structure_z_values = []
    for item in z_structure:
        item = np.unique(item)
        if len(item) != 1:
            raise ValueError("Only one z value per contour supported")
        structure_z_values.append(item[0])

    structure_z_values = np.sort(structure_z_values)
    unique_structure_z_values = np.unique(structure_z_values)

    # if np.any(structure_z_values != unique_structure_z_values):
    #     raise ValueError("Only one contour per slice is currently supported")

    sorted_dose_z = np.sort(z_dose)

    first_dose_index = np.where(sorted_dose_z == structure_z_values[0])[0][0]
    for i, z_val in enumerate(structure_z_values):
        dose_index = first_dose_index + i
        # if structure_z_values[i] != sorted_dose_z[dose_index]:
        #     raise ValueError(
        #         "Only contours where both, there are no gaps in the "
        #         "z-axis of the contours, and the contour axis and dose "
        #         "axis, are aligned are supported."
        #     )

    mask_yxz = np.zeros((len(y_dose), len(x_dose), len(z_dose)), dtype=bool)

    for structure_index, z_val in enumerate(structure_z_values):
        dose_index = int(np.where(z_dose == z_val)[0])

        if z_structure[structure_index][0] != z_dose[dose_index]:
            raise ValueError("Structure and dose indices do not align")

        structure_polygon = matplotlib.path.Path(
            [
                (x_structure[structure_index][i], y_structure[structure_index][i])
                for i in range(len(x_structure[structure_index]))
            ]
        )

        # This logical "or" here is actually in place for the case where
        # there may be multiple contours on the one slice. That's not
        # going to be used at the moment however, as that case is not
        # yet supported in the logic above.
        mask_yxz[:, :, dose_index] = mask_yxz[:, :, dose_index] | (
            structure_polygon.contains_points(points).reshape(len(y_dose), len(x_dose))
        )

    mask_xyz = np.swapaxes(mask_yxz, 0, 1)
    mask_zyx = np.swapaxes(mask_xyz, 0, 2)

    return mask_zyx


def dose_within_struct(struct_name, dcm_struct, dcm_dose):
    all_dose = dose_from_dataset(dcm_dose)  # (z,y,x)
    # print(all_dose.shape)
    #  print(f'max dose={all_dose.max()}; min dose={all_dose.min()}')
    # mask in (y,x,z)
    mask = get_dose_grid_structure_mask(structure_name=struct_name, structure_dataset=dcm_struct, dose_dataset=dcm_dose)
    # mask = np.rollaxis(mask, -1, 0)  # (y,x,z)->(z,y,x)
    # organ_dose = all_dose[mask] # 1-D vector
    # organ_dose = np.where(mask, all_dose, 0) #(z,y,x)
    # print(organ_dose.shape)
    # print(mask.sum())
    #  print((all_dose[mask]).max())
    return all_dose[mask]


def bin_cumsum(dose):
    hist = np.histogram(dose, 100)
    freq = hist[0]
    bin_edge = hist[1]
    bin_mid = (bin_edge[1::] + bin_edge[:-1:]) / 2

    cumulative = np.cumsum(freq[::-1])
    cumulative = cumulative[::-1]
    bin_mid = np.append([0], bin_mid)

    cumulative = np.append(cumulative[0], cumulative)
    percent_cumulative = cumulative / cumulative[0] * 100
    return [bin_mid, percent_cumulative]


def V_constraint_OAR(dcm_dose, dose_threshold=30):
    '''
    dose: dose of OAR
    Return: relative volume > threshold_dose
    Note: gradient can not be backwarded '''
    # dose = to_np(dose)
    dcm_dose.sort()
    limitv = np.interp(dose_threshold, dcm_dose, range(len(dcm_dose)))
    volume = (len(dcm_dose) - limitv) * 100. / len(dcm_dose)
    return volume


def calculate_pass_rate(gamma_array):
    valid_gamma = gamma_array[np.invert(np.isnan(gamma_array))]
    percent_pass = 100 * np.sum(valid_gamma < 1) / len(valid_gamma)
    return {'percent_pass': percent_pass,
            'valid_points_num': valid_gamma.size,
            'passed_points_num': np.sum(valid_gamma < 1),
            'failed_points_num': np.sum(valid_gamma >= 1)}


class StructGamma1():
    def __init__(self, rtstruct_fn, dose_ref_fn, dose_tps_fn):
        gamma_options = {
            'dose_percent_threshold': 3,  # 3%
            'distance_mm_threshold': 3,  # 3mm
            'lower_percent_dose_cutoff': 5,
            'interp_fraction': 10,  # Should be 10 or more for more accurate results
            'max_gamma': 2,
            # 'random_subset': True,  # True or None (default)
            'local_gamma': True,
            # 'ram_available': 2**29,  # 1/2 GB
            # 'quiet': True,
        }

        #
        self.points_num = {}

        # get all structures
        self.dcm_struct = pydicom.read_file(str(rtstruct_fn), force=True)
        self.structs = struct.list_structures(self.dcm_struct)

        # get geometry and dose from ref and eval doses
        self.dcm_dose_ref = pydicom.read_file(str(dose_ref_fn), force=True)
        self.dcm_dose_tps = pydicom.read_file(str(dose_tps_fn), force=True)
        # self.dcm_dose_evl2 = pydicom.read_file(str(dose_evl2_fn), force=True)

        axes_ref, dose_ref = pymedphys.dicom.zyx_and_dose_from_dataset(self.dcm_dose_ref)
        axes_tps, dose_tps = pymedphys.dicom.zyx_and_dose_from_dataset(self.dcm_dose_tps)
        # axes_evl2, dose_evl2 = pymedphys.dicom.zyx_and_dose_from_dataset(self.dcm_dose_evl2)

        # gamma for all dose
        gamma1 = pymedphys.gamma(axes_ref, dose_ref, axes_tps, dose_tps, **gamma_options)
        self.gamma_inf = calculate_pass_rate(gamma1)
        # gamma2 = pymedphys.gamma(axes_ref, dose_ref, axes_evl2, dose_evl2, **gamma_options)
        # self.gamma_inf = calculate_pass_rate(gamma2)

        # gamma for each struct
        self.struct_gamma_inf = self._get_struct_gamma(gamma1)

        # print
        self._print()

    def _get_struct_gamma(self, gamma1):
        struct_gamma_inf = {}
        for struct_name in self.structs:
            print(struct_name)

            # get sturcture mask in (y,x,z) order
            mask = get_dose_grid_structure_mask(struct_name, self.dcm_struct, self.dcm_dose_ref)
            # mask = np.rollaxis(mask, -1, 0)  # (y,x,z)->(z,y,x)
            self.points_num[struct_name] = np.sum(mask)  # recode sturct points number

            # 3D struct gamma in (z,y,x) order
            struct_gamma1 = np.where(mask, gamma1, np.nan)  # (z,y,x)
            struct_gamma_inf[struct_name] = calculate_pass_rate(struct_gamma1)
        return struct_gamma_inf

    def _print(self):
        print('------------------------------------TPS plan gamma-----------------------------------')
        for k, v in self.gamma_inf.items():
            print(k, v)
        print()

        print('------------------------------------TPS struct gamma-----------------------------------')
        for struct_name, gamma_inf in self.struct_gamma_inf.items():
            print(f'struct_name={struct_name}, points_num={self.points_num[struct_name]}')
            for k, v in gamma_inf.items():
                print(k, v)
            print()
        print('------------------------------------TPS gamma end-----------------------------------')


class StructGamma2():
    def __init__(self, rtstruct_fn, dose_ref_fn, dose_pred_fn):
        gamma_options = {
            'dose_percent_threshold': 3,  # 3%
            'distance_mm_threshold': 3,  # 3mm
            'lower_percent_dose_cutoff': 5,
            'interp_fraction': 10,  # Should be 10 or more for more accurate results
            'max_gamma': 2,
            # 'random_subset': True,  # True or None (default)
            'local_gamma': True,
            # 'ram_available': 2**29,  # 1/2 GB
            # 'quiet': True,
        }

        #
        self.points_num = {}

        # get all structures
        self.dcm_struct = pydicom.read_file(str(rtstruct_fn), force=True)
        self.structs = struct.list_structures(self.dcm_struct)

        # get geometry and dose from ref and eval doses
        self.dcm_dose_ref = pydicom.read_file(str(dose_ref_fn), force=True)
        # self.dcm_dose_evl1 = pydicom.read_file(str(dose_evl1_fn), force=True)
        self.dcm_dose_pred = pydicom.read_file(str(dose_pred_fn), force=True)

        axes_ref, dose_ref = pymedphys.dicom.zyx_and_dose_from_dataset(self.dcm_dose_ref)
        # axes_evl1, dose_evl1 = pymedphys.dicom.zyx_and_dose_from_dataset(self.dcm_dose_evl1)
        axes_pred, dose_pred = pymedphys.dicom.zyx_and_dose_from_dataset(self.dcm_dose_pred)

        # gamma for all dose
        # gamma1 = pymedphys.gamma(axes_ref, dose_ref, axes_evl1, dose_evl1, **gamma_options)
        # self.gamma_inf = calculate_pass_rate(gamma1)
        gamma2 = pymedphys.gamma(axes_ref, dose_ref, axes_pred, dose_pred, **gamma_options)
        self.gamma_inf = calculate_pass_rate(gamma2)

        # gamma for each struct
        self.struct_gamma_inf = self._get_struct_gamma(gamma2)

        # print
        self._print()

    def _get_struct_gamma(self, gamma2):
        struct_gamma_inf = {}
        for struct_name in self.structs:
            print(struct_name)

            # get sturcture mask in (y,x,z) order
            mask = get_dose_grid_structure_mask(struct_name, self.dcm_struct, self.dcm_dose_ref)
            mask = np.rollaxis(mask, -1, 0)  # (y,x,z)->(z,y,x)
            self.points_num[struct_name] = np.sum(mask)  # recode sturct points number

            # 3D struct gamma in (z,y,x) order
            struct_gamma2 = np.where(mask, gamma2, np.nan)  # (z,y,x)
            struct_gamma_inf[struct_name] = calculate_pass_rate(struct_gamma2)
        return struct_gamma_inf

    def _print(self):
        print('------------------------------------Pred plan gamma-----------------------------------')
        for k, v in self.gamma_inf.items():
            print(k, v)
        print()

        print('------------------------------------Pred struct gamma-----------------------------------')
        for struct_name, gamma_inf in self.struct_gamma_inf.items():
            print(f'struct_name={struct_name}, points_num={self.points_num[struct_name]}')
            for k, v in gamma_inf.items():
                print(k, v)
            print()
        print('------------gamma2 end-----------')


class DVH():
    def __init__(self, rtstruct_fn, dose_ref_fn, dose_tps_fn, dose_pred_fn):
        # get all structures
        self.dcm_struct = pydicom.read_file(str(rtstruct_fn), force=True)
        self.struct_names = struct.list_structures(self.dcm_struct)

        # all doses
        self.dcm_dose_ref = pydicom.read_file(dose_ref_fn, force=True)
        self.dcm_dose_tps = pydicom.read_file(dose_tps_fn, force=True)
        self.dcm_dose_pred = pydicom.read_file(dose_pred_fn, force=True)

        # struct doses
        self.hists = {}
        self.struct_doses = {}
        for struct_name in self.struct_names:
            self.struct_doses[f'{struct_name}_qa'] = dose_within_struct(struct_name, self.dcm_struct,
                                                                        self.dcm_dose_ref)
            self.struct_doses[f'{struct_name}_tps'] = dose_within_struct(struct_name, self.dcm_struct,
                                                                         self.dcm_dose_tps)
            self.struct_doses[f'{struct_name}_pred'] = dose_within_struct(struct_name, self.dcm_struct,
                                                                          self.dcm_dose_pred)

            self.hists[f'{struct_name}_qa'] = bin_cumsum(self.struct_doses[f'{struct_name}_qa'])
            self.hists[f'{struct_name}_tps'] = bin_cumsum(self.struct_doses[f'{struct_name}_tps'])
            self.hists[f'{struct_name}_pred'] = bin_cumsum(self.struct_doses[f'{struct_name}_pred'])

        # struct volum
        self.Volume = {}
        self.struct_volum = {}
        for struct_name in self.struct_names:
            self.struct_volum[f'{struct_name}_qa'] = dose_within_struct(struct_name, self.dcm_struct,
                                                                        self.dcm_dose_ref)
            self.struct_volum[f'{struct_name}_tps'] = dose_within_struct(struct_name, self.dcm_struct,
                                                                         self.dcm_dose_tps)
            self.struct_volum[f'{struct_name}_pred'] = dose_within_struct(struct_name, self.dcm_struct,
                                                                          self.dcm_dose_pred)

            self.Volume[f'{struct_name}_qa'] = V_constraint_OAR(self.struct_volum[f'{struct_name}_qa'])
            self.Volume[f'{struct_name}_tps'] = V_constraint_OAR(self.struct_volum[f'{struct_name}_tps'])
            self.Volume[f'{struct_name}_pred'] = V_constraint_OAR(self.struct_volum[f'{struct_name}_pred'])

        self._print()

        # dvh
        self.plot_dvh()

    def _print(self):
        for (struct_name, struct_dose) in self.struct_doses.items():
            if 'ref' in struct_name:
                D100 = np.percentile(struct_dose, q=100)
                D98 = np.percentile(struct_dose, q=2)
                D95 = np.percentile(struct_dose, q=95)
                D2 = np.percentile(struct_dose, q=98)
                Dmean = np.mean(struct_dose)
                Dmax = np.max(struct_dose)
                print(f'{struct_name}: Dmax={Dmax},Dmean={Dmean},D98={D98},D2={D2}\n')
            if 'tps' in struct_name:
                D100 = np.percentile(struct_dose, q=100)
                D98 = np.percentile(struct_dose, q=2)
                D95 = np.percentile(struct_dose, q=95)
                D2 = np.percentile(struct_dose, q=98)
                Dmean = np.mean(struct_dose)
                Dmax = np.max(struct_dose)
                print(f'{struct_name}: Dmax={Dmax},Dmean={Dmean},D98={D98},D2={D2}\n')
            if 'pred' in struct_name:
                D100 = np.percentile(struct_dose, q=100)
                D98 = np.percentile(struct_dose, q=2)
                D95 = np.percentile(struct_dose, q=95)
                D2 = np.percentile(struct_dose, q=98)
                Dmean = np.mean(struct_dose)
                Dmax = np.max(struct_dose)
                print(f'{struct_name}: Dmax={Dmax},Dmean={Dmean},D98={D98},D2={D2}\n')
        for (struct_name, struct_dose) in self.struct_volum.items():
            if 'ref' in struct_name:
                V30 = V_constraint_OAR(struct_dose, dose_threshold=30)
                V40 = V_constraint_OAR(struct_dose, dose_threshold=40)
                V45 = V_constraint_OAR(struct_dose, dose_threshold=45)
                V95 = V_constraint_OAR(struct_dose, dose_threshold=42.75)

                print(f'{struct_name}:V30={V30},V40={V40},V45={V45},V42.75Gy={V95},V45Gy={V45}\n')
            if 'tps' in struct_name:
                V30 = V_constraint_OAR(struct_dose, dose_threshold=30)
                V40 = V_constraint_OAR(struct_dose, dose_threshold=40)
                V45 = V_constraint_OAR(struct_dose, dose_threshold=45)
                V95 = V_constraint_OAR(struct_dose, dose_threshold=42.75)

                print(f'{struct_name}:V30={V30},V40={V40},V45={V45},V42.75Gy={V95},V45Gy={V45}\n')
            if 'pred' in struct_name:
                V30 = V_constraint_OAR(struct_dose, dose_threshold=30)
                V40 = V_constraint_OAR(struct_dose, dose_threshold=40)
                V50 = V_constraint_OAR(struct_dose, dose_threshold=45)
                V95 = V_constraint_OAR(struct_dose, dose_threshold=42.75)

                print(f'{struct_name}:V30={V30},V40={V40},V45={V45},V42.75Gy={V95},V45Gy={V45}\n')

    def plot_dvh(self):
        # pdb.set_trace()
        plt.figure(figsize=(20, 15))
        plt.grid(True)
        plt.grid(color='k',
                 linestyle=':',
                 linewidth=1,
                 alpha=0.3)

        for struct_name, (bin_mid, percent_cumulative) in self.hists.items():
            plt.plot(bin_mid, percent_cumulative, label=struct_name, linestyle='solid')
        #     supported values are '-', '--', '-.', ':', 'None', ' ', '', 'solid', 'dashed', 'dashdot', 'dotted'
        plt.title("DVH", fontsize=30)
        plt.xlabel("Dose (Gy)", fontsize=30)
        plt.ylabel("Relative Volume (%)", fontsize=30)
        plt.legend()
        # plt.show()
        plt.savefig('./resource/dvh.png', dpi=200)
        print('DVH Figure has been saved to ./resource/dvh.png')


if __name__ == '__main__':
    dose_ref_fn = r'/home/allen/Temp/PlanQA/test/RT192013/plan/RT192013_QAdose.dcm'
    dose_tps_fn = r'/home/allen/Temp/PlanQA/test/RT192013/plan/93817696_arcplanfinal_Dose.dcm'
    rtstruct_fn = r'/home/allen/Temp/PlanQA/test/RT192013/plan/93817696_StrctrSets.dcm'
    dose_pred_fn = r'/home/allen/Temp/PlanQA/test/RT192013/dose_1.2.826.0.1.3680043.8.274.1.1.8323329.4074.1651029265.433468.dcm'

    # calculate dvh
    print('**********dvh*************')
    DVH(rtstruct_fn, dose_ref_fn, dose_tps_fn, dose_pred_fn)

    # calculate gamma
    # print('\n**********Gamma*************')
    # StructGamma1(rtstruct_fn, dose_ref_fn, dose_tps_fn)
    # StructGamma2(rtstruct_fn, dose_ref_fn, dose_pred_fn)
