import numpy as np
from shapely.geometry import LinearRing, Polygon, Point
from scipy import sparse
from copy import deepcopy
from scipy.sparse import csr_matrix
import itertools
from patchify import patchify
from typing import List, Union
from .ct import CT
from .beam import Beams
from .structures import Structures


class InfluenceMatrix:
    """
    class of influence matrix

    - **Attributes** ::

        :param beamlet_width_mm: The width of each beamlet in millimeters. Default is 2.5. Must be a multiple of 2.5.
        :param beamlet_height_mm: The height of each beamlet in millimeters. Default is 2.5. Must be a multiple of 2.5.
        :param opt_vox_xyz_res_mm: A list of integers representing the resolution for optimization voxels in the x, y, and z axes. When set to None, the original optimization voxel resolution will be used.
        :param target_structure: The target struct_name for creating the beamlets in the influence matrix. Default is 'PTV'.
        :param beamlets_dict: A dictionary containing the information about beamlets used in the plan.
        :param opt_beamlets_PTV_margin_mm: A float value representing the margin_mm around the PTV that was used in creating beamlets
        :param A: influence matrix that is generated by the get_influence_matrix() method.
        :param opt_voxels_dict: A dictionary containing the information about optimization voxels

    - **Methods** ::
        :dose_1d_to_3d(sol)
            Convert dose_1d from 1d to 3d and return dose_1d in 3d
        :dose_3d_to_1d(dose_3d)
            Convert dose_1d from 3d to 1d voxels and return dose_1d in 1d
        :fluence_2d_to_1d(fluence_2d)
            Create vector of intensities from 2d fluence maps
        :fluence_1d_to_2d(fluence_1d)
            From vector of intensities create 2d fluence maps

    """

    def __init__(self, ct: CT, structs: Structures, beams: Beams,
                 beamlet_width_mm: float = None, beamlet_height_mm: float = None, opt_vox_xyz_res_mm: List[float] = None,
                 is_full: bool = False, target_structure: str = 'PTV', opt_beamlets_PTV_margin_mm: float = 3) -> None:
        """
        Create a influence matrix object for Influence Matrix class based upon beamlet resolution and opt_vox_xyz_res_mm

        :param ct: object of class CT
        :param structs: object of class Structures
        :param beams: Object of class Beams
        :param beamlet_width_mm: beamlet width in mm. It should be multiple of 2.5, defaults to 2.5
        :param beamlet_height_mm: beamlet height in mm. It should be multiple of 2.5, defaults to 2.5
        :param target_structure: target struct_name for creating BEV beamlets, defaults to 'PTV'
        :param opt_vox_xyz_res_mm: It down-samples optimization voxels as factor of ct resolution
                e.g. opt_vox_xyz_res = [5*ct.res.x,5*ct.res.y,1*ct.res.z]. It will down-sample optimization voxels with 5 * ct res. in x direction, 5 * ct res. in y direction and 1*ct res. in z direction.
                defaults to None. When None it will use the original optimization voxel resolution.
        :param is_full: Load full or sparse matrix. defaults to False. If True, will load full matrix

        """
        if beamlet_width_mm is None and beamlet_height_mm is None:
            self.beamlet_width_mm = beams.get_beamlet_width()
            self.beamlet_height_mm = beams.get_beamlet_height()
        else:
            if beamlet_width_mm % 2.5 == 0 and beamlet_height_mm % 2.5 == 0:
                self.beamlet_width_mm = beamlet_width_mm
                self.beamlet_height_mm = beamlet_height_mm
            else:
                raise ValueError('beamlet_width_mm and beamlet_height_mm should be multiple of 2.5')
        self._beams = beams
        self._structs = structs
        self._ct = ct

        down_sample_xyz = None  # Temporary variable to check if we want to down sample or not
        if opt_vox_xyz_res_mm is not None:
            down_sample_xyz = [round(i / j) for i, j in zip(opt_vox_xyz_res_mm, ct.get_ct_res_xyz_mm())]
            if np.all(np.array(down_sample_xyz) == 1):  # if all are 1 then no down sample
                down_sample_xyz = None

        self._down_sample_xyz = down_sample_xyz
        self.is_full = is_full

        # create deepcopy of the object or else it will modify the structs object
        if hasattr(structs, 'opt_voxels_dict'):
            self.opt_voxels_dict = deepcopy(structs.opt_voxels_dict)
            # del structs.opt_voxels_dict  # remove opt_voxels_dict from structures
        else:
            self.opt_voxels_dict = deepcopy(self.opt_voxels_dict)

        # creating deepcopy so that it doesnt modify beams object
        if 'beamlets' in beams.beams_dict:
            self.beamlets_dict = deepcopy(beams.beams_dict['beamlets'])
        else:
            self.beamlets_dict = deepcopy(self.beamlets_dict['beamlets'])

        for i in range(len(beams.beams_dict['ID'])):
            self.beamlets_dict[i]['beam_id'] = beams.beams_dict['ID'][i]  # save beam_id in beamlet_dict

        self.opt_beamlets_PTV_margin_mm = opt_beamlets_PTV_margin_mm
        self.opt_voxels_dict['ct_origin_xyz_mm'] = ct.ct_dict['origin_xyz_mm']  # store ct origin from plan
        print('Creating BEV..')
        self.preprocess_beams(structure=target_structure)
        if self._down_sample_xyz is not None:
            self.pre_process_voxels()  # create new optimization voxel indices based on down-sample resolution

        if not self.is_full:
            print('Loading sparse influence matrix...')
            self.A = self.get_influence_matrix()  # create sparse influence matrix
            self.sparse_tol = float(beams.beams_dict['influenceMatrixSparse_tol'][0])
        else:
            print('Loading full influence matrix..')
            self.A = self.get_influence_matrix(self.is_full)  # create full matrix
        self._vox_map = None
        self._vox_weights = None
        print('Done')

    def get_voxel_info(self, row_number):
        row_dict = {}
        if row_number <= self.A.shape[0]:
            for i in range(len(self.opt_voxels_dict['name'])):
                if row_number in list(self.opt_voxels_dict['voxel_idx'][i]):
                    row_dict.setdefault('structures', []).append(self.opt_voxels_dict['name'][i])

            patch = np.where(
                self.opt_voxels_dict['ct_to_dose_voxel_map'][0] == row_number)  # get dose voxel index patch

            # get center of patch. It should be equivalent dose voxel center
            z_ind = (min(patch[0]) + max(patch[0])) / 2
            y_ind = (min(patch[1]) + max(patch[1])) / 2
            x_ind = (min(patch[2]) + max(patch[2])) / 2
            ct_res = self.opt_voxels_dict['ct_voxel_resolution_xyz_mm']
            ct_orig = self.opt_voxels_dict['ct_origin_xyz_mm']
            xyz_mm = [ct_orig[0] + x_ind * ct_res[0], ct_orig[1] + y_ind * ct_res[1], ct_orig[2] + z_ind * ct_res[2]]
            row_dict['position_xyz_mm'] = xyz_mm
        else:
            raise ValueError('invalid row number {}'.format(row_number))
        return row_dict

    def get_beamlet_info(self, col_number: int) -> dict:
        """

        :param col_number: col number/beamlet number in influence matrix
        :return: dictionary containing information about the beamlet
        """
        col_dict = {}
        for ind in range(len(self.beamlets_dict)):
            if col_number in range(self.beamlets_dict[ind]['start_beamlet_idx'],
                                   self.beamlets_dict[ind]['end_beamlet_idx'] + 1):
                if ind > 0:
                    prev_beam_beamlet = self.beamlets_dict[ind - 1]['end_beamlet_idx'] + 1
                else:
                    prev_beam_beamlet = 0
                col_dict['beam_id'] = self.beamlets_dict[ind]['beam_id']
                col_dict['position_x_mm'] = self.beamlets_dict[ind]['position_x_mm'][0][col_number - prev_beam_beamlet]
                col_dict['position_y_mm'] = self.beamlets_dict[ind]['position_y_mm'][0][col_number - prev_beam_beamlet]
                col_dict['width_mm'] = self.beamlets_dict[ind]['width_mm'][0][col_number - prev_beam_beamlet]
                col_dict['height_mm'] = self.beamlets_dict[ind]['height_mm'][0][col_number - prev_beam_beamlet]
        return col_dict

    def dose_1d_to_3d(self, sol: dict = None, dose_1d: np.array = None) -> np.ndarray:
        """
        Create 3d dose_1d from dose_1d in 1d voxels. 3d array have same resolution as CT

        :param sol: solution dictionary from optimization
        :param dose_1d: dose_1d in 1d. Optional
        :return: dose_1d in 3d


        """
        dose_vox_map = self.opt_voxels_dict['ct_to_dose_voxel_map'][0]
        # dose_1d = my_plan.opt_voxels_dict['dose_1d']
        if dose_1d is None:
            if 'dose_1d' not in sol:
                dose_1d = sol['inf_matrix'].A * sol['optimal_intensity']  # multiply it with num fractions
            else:
                dose_1d = sol['dose_1d']
        dose_3d = np.zeros_like(dose_vox_map, dtype='float32')
        inds = np.unique(dose_vox_map[dose_vox_map >= 0])
        a = np.where(np.isin(dose_vox_map, inds))
        dose_3d[a] = dose_1d[dose_vox_map[a]]
        return dose_3d

    def dose_3d_to_1d(self, dose_3d: np.ndarray) -> np.array:
        """
        Get dose_1d in 1d voxels for the given influence matrix from 3d dose_1d

        :param dose_3d: 3d dose_1d
        :return: dose_1d in 1d voxel indices
        """
        dose_vox_map = self.opt_voxels_dict['ct_to_dose_voxel_map'][0]
        inds = np.unique(dose_vox_map[dose_vox_map >= 0])
        dose_1d = np.zeros_like(inds, dtype=float)
        a = np.where(np.isin(dose_vox_map, inds))
        dose_1d[dose_vox_map[a]] = dose_3d[a]
        return dose_1d

    def fluence_2d_to_1d(self, fluence_2d: List[np.ndarray]) -> np.array:
        """
        Create vector of intensities from 2d fluence maps

        :param fluence_2d: 2d fluence as list of nd array with same length as number of beams_dict
        :return: fluence in 1d for beamlet indices
        """
        fluence_1d = np.zeros((self.A.shape[1]))
        for ind in range(len(self.beamlets_dict)):
            maps = self.beamlets_dict[ind]['beamlet_idx_2d_finest_grid']
            numRows = np.size(maps, 0)
            numCols = np.size(maps, 1)
            # wMaps = np.zeros((numRows, numCols))
            for r in range(numRows):
                for c in range(numCols):
                    if maps[r, c] >= 0:
                        curr = maps[r, c]
                        fluence_1d[curr] = fluence_2d[ind][r, c]
        return fluence_1d

    def fluence_1d_to_2d(self, fluence_1d: np.array = None, sol: dict = None) -> List[np.ndarray]:
        """
        Create 2d fluence maps from vector of intensities.

        :param fluence_1d: 2d fluence as vector
        :param sol: solution from optimization
        :return: 2d fluence as a list for each beam

        """
        if fluence_1d is None:
            fluence_1d = sol['optimal_intensity']

        wMaps = []
        for ind in range(len(self.beamlets_dict)):
            maps = self.beamlets_dict[ind]['beamlet_idx_2d_finest_grid']
            numRows = np.size(maps, 0)
            numCols = np.size(maps, 1)
            wMaps.append(np.zeros((numRows, numCols)))
            # wMaps = np.zeros((numRows, numCols))
            for r in range(numRows):
                for c in range(numCols):
                    if maps[r, c] >= 0:
                        curr = maps[r, c]
                        wMaps[ind][r, c] = fluence_1d[curr]
                        # wMaps[r, c] = optimal_intensity[curr]

        return wMaps

    @staticmethod
    def sol_change_inf_matrix(sol: dict, inf_matrix) -> dict:
        """
        Create a new solution by changing the basis of current solution.
        It will create a solution with same number of beamlets and voxels as inf_matrix


        :param sol: solution for which influence matrix is changed
        :param inf_matrix: object of class Influence matrix
        :return: new solution dictionary having same number of beamlets and voxels as inf_matrix
        """
        new_sol = dict()
        if sol['optimal_intensity'].shape[0] < inf_matrix.A.shape[1]:
            optimal_intensity = sol['inf_matrix'].fluence_1d_to_2d(fluence_1d=sol['optimal_intensity'])
            new_sol['optimal_intensity'] = inf_matrix.fluence_2d_to_1d(optimal_intensity)
        elif sol['optimal_intensity'].shape[0] == inf_matrix.A.shape[1]:
            new_sol['optimal_intensity'] = sol['optimal_intensity']
        else:
            raise ValueError("Beamlet resolution should be greater than or equal to beamlets for inf_matrix")

        # new_sol['dose_1d'] = inf_matrix.A * new_sol['optimal_intensity']

        # dose_3d = sol['inf_matrix'].dose_1d_to_3d(dose_1d=sol['dose_1d'])
        # new_sol['dose_1d'] = inf_matrix.dose_3d_to_1d(dose_3d=dose_3d)

        new_sol['inf_matrix'] = inf_matrix
        return new_sol

    def create_down_sample(self, beamlet_width_mm: float = None, beamlet_height_mm: float = None,
                           opt_vox_xyz_res_mm: List[float] = None,
                           overwrite: bool = False, remove_corner_beamlets: bool = False):
        if overwrite:
            new_inf_matrix = self
        else:
            new_inf_matrix = deepcopy(self)
        if beamlet_width_mm is None:
            beamlet_width_mm = self.beamlet_width_mm
            beamlet_height_mm = self.beamlet_height_mm
        new_inf_matrix.beamlet_width_mm = beamlet_width_mm
        new_inf_matrix.beamlet_height_mm = beamlet_height_mm
        new_inf_matrix.beamlets_dict = deepcopy(self._beams.beams_dict['beamlets'])
        for i in range(len(self._beams.beams_dict['ID'])):
            new_inf_matrix.beamlets_dict[i]['beam_id'] = deepcopy(self._beams.beams_dict['ID'][i])  # save beam_id in beamlet_dict

        for i in range(len(new_inf_matrix.beamlets_dict)):
            new_inf_matrix.beamlets_dict[i]['beamlet_idx_2d_finest_grid'] = deepcopy(self.beamlets_dict[i]['beamlet_idx_2d_finest_grid'])
        down_sample_xyz = None  # Temporary variable to check if we want to down sample or not
        if opt_vox_xyz_res_mm is not None:
            down_sample_xyz = [round(i / j) for i, j in zip(opt_vox_xyz_res_mm, self.opt_voxels_dict['ct_voxel_resolution_xyz_mm'])]
            if np.all(np.array(down_sample_xyz) == 1):  # if all are 1 then no down sample
                down_sample_xyz = None
        new_inf_matrix._down_sample_xyz = down_sample_xyz
        new_inf_matrix.preprocess_beams(remove_corner_beamlets=remove_corner_beamlets)
        new_inf_matrix.pre_process_voxels()
        A = new_inf_matrix.get_influence_matrix(is_full=new_inf_matrix.is_full)
        new_inf_matrix.A = A
        return new_inf_matrix

    def get_influence_matrix(self, is_full=False):
        """

        Load influence matrix based on the beamlets and voxels.

        :param beams: object of class Beams
        :param is_full: get full or sparse matrix. Default to True.
        :return: full or sparse matrix
        """
        data_beamlet_width = self._beams.get_beamlet_width()
        data_beamlet_height = self._beams.get_beamlet_height()

        if not is_full:
            if self.beamlet_width_mm > data_beamlet_width or \
                    self.beamlet_height_mm > data_beamlet_height or self._down_sample_xyz is not None:
                inf_matrix_sparse = self.A
            else:
                # deepcopy so that it doesnt modify
                inf_matrix_sparse = deepcopy(self._beams.beams_dict['influenceMatrixSparse'])
                # if 'influenceMatrixSparse' in self._beams.beams_dict:
                del self._beams.beams_dict['influenceMatrixSparse']

            for ind in range(len(self.beamlets_dict)):
                opt_beamlets = self.beamlets_dict[ind]['opt_beamlets_ids']

                if ind == 0:
                    if self.beamlet_width_mm > data_beamlet_width or self.beamlet_height_mm > data_beamlet_height:
                        print('parsing influence matrix for beam {}'.format(ind))
                        inf_matrix = sparse.hstack(
                            [csr_matrix(inf_matrix_sparse[:, np.unique(opt_beamlets[i])].sum(axis=1)) for i in
                             range(len(opt_beamlets))], format='csr')
                    else:
                        if self._down_sample_xyz is None:
                            inf_matrix = inf_matrix_sparse[ind][:, opt_beamlets]
                else:
                    if self.beamlet_width_mm > data_beamlet_width or self.beamlet_height_mm > data_beamlet_height:
                        print('parsing influence matrix for beam {}'.format(ind))
                        inf_matrix_2 = sparse.hstack(
                            [csr_matrix(inf_matrix_sparse[:, np.unique(opt_beamlets[i])].sum(axis=1)) for i in
                             range(len(opt_beamlets))], format='csr')
                        inf_matrix = sparse.hstack(
                            [inf_matrix, inf_matrix_2], format='csr')
                    else:
                        if self._down_sample_xyz is None:
                            inf_matrix = sparse.hstack(
                                [inf_matrix, inf_matrix_sparse[ind][:, opt_beamlets]], format='csr')
            # if del_org_matrix:
            if self._down_sample_xyz is not None:
                if self.beamlet_width_mm <= data_beamlet_width or self.beamlet_height_mm <= data_beamlet_height:
                    inf_matrix = inf_matrix_sparse
                print('creating influence matrix for down sample voxels..')
                inf_matrix = sparse.vstack(
                    [csr_matrix((sparse.diags(self._vox_weights[i]).dot(inf_matrix[self._vox_map[i], :])).sum(axis=0))
                     for i
                     in
                     range(len(self._vox_map))], format='csr')
        else:
            if self.beamlet_width_mm > data_beamlet_width or self.beamlet_height_mm > data_beamlet_height or self._down_sample_xyz is not None:
                inf_matrix_full = self.A
            else:
                inf_matrix_full = self._beams.beams_dict['influenceMatrixFull']
                del self._beams.beams_dict['influenceMatrixFull']

            for ind in range(len(self.beamlets_dict)):
                opt_beamlets = self.beamlets_dict[ind]['opt_beamlets_ids']

                if ind == 0:
                    if self.beamlet_width_mm > data_beamlet_width or self.beamlet_height_mm > data_beamlet_height:
                        print('parsing full influence matrix for beam {}'.format(ind))
                        inf_matrix = np.hstack(
                            [inf_matrix_full[:, np.unique(opt_beamlets[i])].sum(axis=1) for i in
                             range(len(opt_beamlets))])
                    else:
                        if self._down_sample_xyz is None:
                            inf_matrix = inf_matrix_full[ind][:, opt_beamlets]
                else:
                    if self.beamlet_width_mm > data_beamlet_width or self.beamlet_height_mm > data_beamlet_height:
                        print('parsing full influence matrix for beam {}'.format(ind))
                        inf_matrix_2 = np.hstack(
                            [inf_matrix_full[:, np.unique(opt_beamlets[i])].sum(axis=1) for i in
                             range(len(opt_beamlets))])
                        inf_matrix = np.hstack([inf_matrix, inf_matrix_2])
                    else:
                        if self._down_sample_xyz is None:
                            inf_matrix = np.hstack([inf_matrix, inf_matrix_full[ind][:, opt_beamlets]])

                # down sampling voxels
            if self._down_sample_xyz is not None:
                if self.beamlet_width_mm <= data_beamlet_width or self.beamlet_height_mm <= data_beamlet_height:
                    inf_matrix = inf_matrix_full
                print('creating influence matrix for down sample voxels..')
                inf_matrix = np.vstack(
                    [(np.diags(self._vox_weights[i]).dot(inf_matrix[self._vox_map[i], :])).sum(axis=0)
                     for i in range(len(self._vox_map))])

        return inf_matrix

    def create_BEV_mask_from_contours(self, ind: int, structure: str = 'PTV',
                                      margin_mm: float = None) -> np.ndarray:
        """

        Since beams object contain projection of struct_name contours on BEV, this function helps to create mask from those contours on BEV

        :param beams: object of class Beams
        :param ind: indices of the beam in beamlet dictionary
        :param structure: struct_name for which contours to be converted to mask
        :param margin_mm: expand the contours in mm. defaults to None. If it is not none, it will expand the contours
        :return:
        """
        if margin_mm is None and structure == 'PTV':
            margin_mm = self.opt_beamlets_PTV_margin_mm
        elif margin_mm is None:
            margin_mm = 0

        # get PTV contour from data
        contours = self._beams.beams_dict['BEV_structure_contour_points'][ind][structure]
        # ind = my_plan.beamlets_dict['ID'].index(beam_id)
        # contours = my_plan._beams_contours[ind][struct_name]

        # for each contour create polygon and get beamlets inside the polygon and create mask for it
        for count_num in range(len(contours)):
            polygon = []
            for j in contours[count_num]:
                polygon.append((j[0], j[1]))
            r = LinearRing(polygon)
            s = Polygon(r)
            # add margin_mm around polygon
            if margin_mm == 0:
                shapely_poly = s
            else:
                shapely_poly = Polygon(s.buffer(margin_mm), [r])
            #             poly_coordinates = np.array(list(t.exterior.coords))

            # create beam map from beamlet coordinates and mask the beamlets inside shapely polygon
            beamlets = self.beamlets_dict[ind]
            x_positions = beamlets['position_x_mm'][0]
            y_positions = beamlets['position_y_mm'][0]
            x_min_max_sort = np.sort(np.unique(x_positions))
            y_max_min_sort = np.sort(np.unique(y_positions))[::-1]
            points = []
            for i in range(len(beamlets['id'][0])):
                x_coord = beamlets['position_x_mm'][0][i]
                y_coord = beamlets['position_y_mm'][0][i]
                points.append(Point(x_coord, y_coord))
            valid_points = []
            valid_points.extend(filter(shapely_poly.contains, points))
            x_and_y = [(a.x, a.y) for a in valid_points]
            XX, YY = np.meshgrid(x_min_max_sort, y_max_min_sort)
            w_all = np.column_stack((x_positions, y_positions))
            if count_num == 0:
                mask = np.zeros_like(XX, dtype=bool)
            for row in range(XX.shape[0]):
                for col in range(XX.shape[1]):
                    ind = np.where((w_all[:, 0] == XX[row, col]) & (w_all[:, 1] == YY[row, col]))[0][0]

                    # if (beamlets[ind]['position_x_mm'], beamlets[ind]['position_y_mm']) in x_and_y:
                    if (beamlets['position_x_mm'][0][ind], beamlets['position_y_mm'][0][ind]) in x_and_y:
                        # beam_map[row, col] = beamlets[ind]['id']
                        mask[row, col] = True
        return mask

    def preprocess_beams(self, structure='PTV', remove_corner_beamlets=False):
        """
        Preprocess beams to create beamlets dictionary based on beamlet resolution.

        :param beams: object of class Beams
        :param structure: projection of the struct_name on BEV for which beamlets to be considered for creating influence matrix. defaults to PTV
        :param remove_corner_beamlets: If remove corner beamlet is true, it will remove the corner beamlets during down sampling.
        :return: beamlets for processing influence matrix
        """

        for ind in range(len(self.beamlets_dict)):
            # ind = my_plan.beamlets_dict['ID'].index(beam_id)
            mask = self.create_BEV_mask_from_contours(ind=ind, structure=structure,
                                                      margin_mm=self.opt_beamlets_PTV_margin_mm)
            beam_2d_grid = self.create_beamlet_idx_2d_finest_grid(ind=ind)
            beamlets = self.beamlets_dict[ind]

            # creating beam_map of original resolution so that mask can be multiplied with it
            beam_map = self.get_orig_res_2d_grid(ind)
            beam_map = np.multiply((beam_map + int(1)), mask)  # add and subtract one to maintain 0th beamlet
            beam_map = beam_map - int(1)  # subtract one again to get original beamlets

            # get mask and beam_map in 2.5mm resolution
            beamlet_ind = np.unique(beam_map.flatten())
            beamlet_ind = beamlet_ind[beamlet_ind >= 0]
            a = np.where(np.isin(beam_2d_grid, beamlet_ind))
            mask_2d_grid = np.zeros_like(beam_2d_grid, dtype=bool)
            mask_2d_grid[a] = True
            beam_2d_grid = (beam_2d_grid + int(1)) * mask_2d_grid  # add and subtract 1 to retain ids
            beam_2d_grid = beam_2d_grid - int(1)
            beam_map = beam_2d_grid

            # get opt beamlets for down_sample grid as list of original inf_matrix beamlet
            if self.beamlet_width_mm > self._beams.get_beamlet_width() or \
                    self.beamlet_height_mm > self._beams.get_beamlet_height():
                down_sample_2d_grid = self.down_sample_2d_grid(ind=ind, beamlet_width_mm=self.beamlet_width_mm,
                                                               beamlet_height_mm=self.beamlet_height_mm)
                down_sample_2d_grid = (down_sample_2d_grid + int(1)) * mask_2d_grid  # add and subtract 1 to retain ids
                down_sample_2d_grid = down_sample_2d_grid - int(1)
                down_sample_beamlets, counts = np.unique(np.sort(down_sample_2d_grid[down_sample_2d_grid >= 0]),
                                                         return_counts=True)
                # remove corner beamlets in coarse resolution and updating down sample map
                if remove_corner_beamlets:
                    count_ind = np.where(counts >= (self.beamlet_width_mm / 2.5) * (self.beamlet_height_mm / 2.5))
                    down_sample_beamlets = down_sample_beamlets[count_ind]
                    # updating down sample grid after removing corner beamlets
                    inds_down_sample = down_sample_2d_grid == down_sample_beamlets[:, None,
                                                              None]  # keep only down sample beamlets and remove others
                    down_sample_2d_grid[~np.any(inds_down_sample, axis=0)] = -1

                a = np.where(np.isin(down_sample_2d_grid, down_sample_beamlets))
                mask_2d_grid = np.zeros_like(beam_2d_grid, dtype=bool)
                mask_2d_grid[a] = True
                # actual_beamlets = beam_2d_grid[a]
                actual_beamlets = self.beamlets_dict[ind]['beamlet_idx_2d_finest_grid'][a]
                sampled_beamlets = down_sample_2d_grid[a]
                b = [np.where(sampled_beamlets == down_sample_beamlets[i]) for i in range(len(down_sample_beamlets))]
                opt_beamlets = [actual_beamlets[i] for i in b]

                beam_map = down_sample_2d_grid
            else:
                opt_beamlets = np.unique(np.sort(beam_map[beam_map >= 0]))

            # make beamlets continuous
            std_map = self.sort_beamlets(beam_map)
            if ind == 0:
                beam_map = std_map
            else:
                beam_map = std_map + int(
                    np.amax(self.beamlets_dict[ind - 1]['beamlet_idx_2d_finest_grid']) + 1) * mask_2d_grid
            standInd = np.unique(np.sort(beam_map.flatten()))
            self.beamlets_dict[ind]['beamlet_idx_2d_finest_grid'] = beam_map
            # my_plan.beamlets_dict.setdefault('structure_mask_2dgrid', []).append(mask_2d_grid)
            self.beamlets_dict[ind]['start_beamlet_idx'] = standInd[1]
            self.beamlets_dict[ind]['end_beamlet_idx'] = np.amax(self.beamlets_dict[ind]['beamlet_idx_2d_finest_grid'])
            self.beamlets_dict[ind]['opt_beamlets_ids'] = opt_beamlets

            # update beamlet dict
            if self.beamlet_width_mm > self._beams.get_beamlet_width() or \
                    self.beamlet_height_mm > self._beams.get_beamlet_height():
                pass
            else:
                self.beamlets_dict[ind]['position_x_mm'][0] = beamlets['position_x_mm'][0][opt_beamlets]
                self.beamlets_dict[ind]['position_y_mm'][0] = beamlets['position_y_mm'][0][opt_beamlets]
                self.beamlets_dict[ind]['width_mm'][0] = beamlets['width_mm'][0][opt_beamlets]
                self.beamlets_dict[ind]['height_mm'][0] = beamlets['height_mm'][0][opt_beamlets]
                self.beamlets_dict[ind]['MLC_leaf_idx'][0] = beamlets['MLC_leaf_idx'][0][opt_beamlets]
            del self.beamlets_dict[ind]['id']
            # my_plan.beams_dict.setdefault('opt_beamlets_ids', []).append(standInd[1:])

    def create_beamlet_idx_2d_finest_grid(self, ind: int) -> np.ndarray:
        """
        Create beamlet idx in 2d grid of 2.5mm resolution. i.e. each element in matrix is 2.5mm*2.5mm

        :param ind: indices for the beam in beamlets dictionary
        :return:
        """
        # ind = my_plan.beamlets_dict['ID'].index(beam_id)
        beamlets = self.beamlets_dict[ind]
        x_positions = beamlets['position_x_mm'][0] - beamlets['width_mm'][0] / 2
        y_positions = beamlets['position_y_mm'][0] + beamlets['height_mm'][0] / 2
        right_ind = np.argmax(x_positions)
        bottom_ind = np.argmin(y_positions)
        w_all = np.column_stack((x_positions, y_positions))  # top left corners of all beamlets
        x_coord = np.arange(np.min(x_positions), np.max(x_positions) + beamlets['width_mm'][0][right_ind], 2.5)
        y_coord = np.arange(np.max(y_positions), np.min(y_positions) - beamlets['height_mm'][0][bottom_ind], -2.5)
        XX, YY = np.meshgrid(x_coord, y_coord)
        beamlet_idx_2d_grid = np.ones_like(XX, dtype=int)
        beamlet_idx_2d_grid = beamlet_idx_2d_grid * int(-1)
        for row in range(XX.shape[0]):
            for col in range(XX.shape[1]):
                ind = np.where((w_all[:, 0] == XX[row, col]) & (w_all[:, 1] == YY[row, col]))
                if np.size(ind) > 0:
                    ind = ind[0][0]
                    num_width = int(beamlets['width_mm'][0][ind] / 2.5)
                    num_height = int(beamlets['height_mm'][0][ind] / 2.5)
                    beamlet_idx_2d_grid[row:row + num_height, col:col + num_width] = ind

        return beamlet_idx_2d_grid

    def down_sample_2d_grid(self, ind: int, beamlet_width_mm: float = 5.0,
                            beamlet_height_mm: float = 5.0) -> np.ndarray:
        """
        create down sample beam map based on beamlet width and height

        :param ind: indices for the beam in beamlets dictionary
        :param beamlet_width_mm:
        :param beamlet_height_mm:
        :return: down sampled 2d grid

        """
        # ind = my_plan.beamlets_dict['ID'].index(beam_id)
        beamlets = self.beamlets_dict[ind]
        x_positions = beamlets['position_x_mm'][0] - beamlets['width_mm'][0] / 2
        y_positions = beamlets['position_y_mm'][0] + beamlets['height_mm'][0] / 2
        right_ind = np.argmax(x_positions)
        bottom_ind = np.argmin(y_positions)
        # w_all = np.column_stack((x_positions, y_positions))  # top left corners of all beamlets
        x_coord = np.arange(np.min(x_positions), np.max(x_positions) + beamlets['width_mm'][0][right_ind], 2.5)
        y_coord = np.arange(np.max(y_positions), np.min(y_positions) - beamlets['height_mm'][0][bottom_ind], -2.5)
        XX, YY = np.meshgrid(x_coord, y_coord)
        beamlet_resample_2d_grid = np.ones_like(XX, dtype=int)
        beamlet_resample_2d_grid = beamlet_resample_2d_grid * int(-1)
        row_col_covered = []
        ind = 0
        for row in range(XX.shape[0]):
            for col in range(XX.shape[1]):
                # ind = np.where((w_all[:, 0] == XX[row, col]) & (w_all[:, 1] == YY[row, col]))
                # if np.size(ind) > 0:
                #     ind = ind[0][0]
                if [row, col] not in row_col_covered:
                    num_width = int(beamlet_width_mm / 2.5)
                    num_height = int(beamlet_height_mm / 2.5)
                    beamlet_resample_2d_grid[row:row + num_height, col:col + num_width] = ind
                    ind = ind + 1
                    rows = np.arange(row, row + num_height)
                    cols = np.arange(col, col + num_width)
                    for r in itertools.product(rows, cols):
                        row_col_covered.append([r[0], r[1]])
        return beamlet_resample_2d_grid

    def get_orig_res_2d_grid(self, ind: int) -> np.ndarray:
        """
        Create beam map 2d grid in the original resolution where each element in matrix is the size of respective beamlet.

        :param ind: indices for the beam in beamlets dictionary
        :return: beam map 2d grid in original resolution

        """
        beamlets = self.beamlets_dict[ind]
        x_positions = beamlets['position_x_mm'][0]
        y_positions = beamlets['position_y_mm'][0]
        x_min_max_sort = np.sort(np.unique(x_positions))
        y_max_min_sort = np.sort(np.unique(y_positions))[::-1]
        XX, YY = np.meshgrid(x_min_max_sort, y_max_min_sort)
        w_all = np.column_stack((x_positions, y_positions))
        beam_map = np.zeros_like(XX, dtype=int)
        for row in range(XX.shape[0]):
            for col in range(XX.shape[1]):
                b_ind = np.where((w_all[:, 0] == XX[row, col]) & (w_all[:, 1] == YY[row, col]))[0][0]
                # beam_map[row, col] = beamlets[ind]['id']+int(1)# adding one so that 0th beamlet is retained
                beam_map[row, col] = b_ind  # + int(1)  # adding one so that 0th beamlet is retained
        # beam_map = np.multiply(beam_map, mask)
        # beam_map = beam_map - int(1)  # subtract one again to get original beamlets
        return beam_map

    def get_bev_2d_grid(self, beam_id: Union[int, List[int]] = None, ind: int = None,
                        finest_grid: bool = False) -> Union[np.ndarray, List[np.ndarray]]:

        """
        get BEV in 2d grid in original resolution where beamlet index is the column number in influence matrix

        :param ind: idx of the beam in beamlets_dict.
        It can be int or List[int]. If int, ndarray is return, If List[int], list[ndarray] is returned
        :param beam_id: beam_id of the beam.
        It can be int or List[int]. If int, ndarray is return, If List[int], list[ndarray] is returned
        :param finest_grid:Default to False. If set to true, it will return 2d grid in finest resolution
        :return: ndarray/List[ndarray]
        """

        if beam_id is not None:
            ind = []
            if isinstance(beam_id, int):
                ind = [i for i in range(len(self.beamlets_dict)) if
                       self.beamlets_dict[i]['beam_id'] == beam_id]
                if not ind:
                    raise ValueError("invalid beam id {}".format(beam_id))
            elif isinstance(beam_id, list):
                for idx in beam_id:
                    try:
                        ind_1 = [i for i in range(len(self.beamlets_dict)) if
                                 self.beamlets_dict[i]['beam_id'] == idx][0]
                        ind.append(ind_1)
                    except:
                        raise ValueError("invalid beam id {}".format(idx))
        # Bug fix. in case ind is int make it as list
        if isinstance(ind, int):
            ind = [ind]
        beam_orig = []
        for b in ind:
            beam_map = self.beamlets_dict[b]['beamlet_idx_2d_finest_grid']
            if not finest_grid:
                rowsNoRepeat = [0]
                for i in range(1, np.size(beam_map, 0)):
                    if (beam_map[i, :] != beam_map[rowsNoRepeat[-1], :]).any():
                        rowsNoRepeat.append(i)
                colsNoRepeat = [0]
                for j in range(1, np.size(beam_map, 1)):
                    if (beam_map[:, j] != beam_map[:, colsNoRepeat[-1]]).any():
                        colsNoRepeat.append(j)
                beam_map = beam_map[np.ix_(np.asarray(rowsNoRepeat), np.asarray(colsNoRepeat))]

                # add this part in case remove corner beamlets create issue for getting original resolution
                remove_row = []
                remove_col = []
                for row in range(beam_map.shape[0]):
                    if row < beam_map.shape[0] - 1:
                        prev_elem = beam_map[row, :][beam_map[row, :] > -1]
                        next_elem = beam_map[row + 1, :][beam_map[row + 1, :] > -1]
                        matching_elements = np.intersect1d(prev_elem, next_elem)
                        if len(matching_elements) >= 1:
                            if len(prev_elem) <= len(next_elem):
                                remove_row.append(row)
                            else:
                                remove_row.append(row + 1)
                for col in range(beam_map.shape[1]):
                    if col < beam_map.shape[1] - 1:
                        prev_elem = beam_map[:, col][beam_map[:, col] > -1]
                        next_elem = beam_map[:, col + 1][beam_map[:, col + 1] > -1]
                        matching_elements = np.intersect1d(prev_elem, next_elem)
                        if len(matching_elements) >= 1:
                            if len(prev_elem) <= len(next_elem):
                                remove_col.append(col)
                            else:
                                remove_col.append(col + 1)
                mask_rows = np.ones(beam_map.shape[0], dtype=bool)
                mask_columns = np.ones(beam_map.shape[1], dtype=bool)
                if len(remove_row) > 0:
                    mask_rows[remove_row] = False
                if len(remove_col) > 0:
                    mask_columns[remove_col] = False
                beam_map = beam_map[mask_rows][:, mask_columns]

            beam_orig.append(beam_map)
        if len(beam_orig) == 1:
            beam_orig = beam_orig[0]  # return ndarray in case if it is not list

        return beam_orig

    def down_sample_voxels(self, down_sample_xyz: List[int] = None):
        """
        This method creates new ct to dose_1d voxel map array based on down sampled voxels.
        It also outputs map between down sampled voxel indices to original voxel indices for down sampling influence matrix
        :param down_sample_xyz: It down-samples optimization voxels as factor of ct resolution
                e.g. _down_sample_xyz = [5,5,1]. It will down-sample optimization voxels with 5 * ct res. in x direction, 5 * ct res. in y direction and 1*ct res. in z direction.
                defaults to None. When None it will use the original optimization voxel resolution.
        :return: voxel map(list containing map between down-sampled indices to original indices), voxel weight (weighted average of voxels in down sample
         and ct to dose_1d voxel map
        """

        vox_3d = self.opt_voxels_dict['ct_to_dose_voxel_map'][0]
        # dose_to_ct_int = np.round(np.array(self.opt_voxels_dict['dose_voxel_resolution_xyz_mm']) /
        #                           np.array(self.opt_voxels_dict['ct_voxel_resolution_xyz_mm']))
        # dose_to_ct_int = dose_to_ct_int.astype(int)

        # using patchify
        use_patchify = True
        if use_patchify:
            print('reindexing voxels...')
            vox_map = []
            vox_weights = []
            slices, height, width = down_sample_xyz[2], down_sample_xyz[1], down_sample_xyz[0]
            patches = patchify(vox_3d, (slices, height, width), step=(slices, height, width))
            # pat = np.vstack(patches)
            # pat = np.vstack(pat)
            pat = patches.reshape(-1, slices, height, width)
            count = 0
            for i in range(np.shape(pat)[0]):
                if np.any(pat[i] > -1):
                    vox_inds, weights = np.unique(np.sort(pat[i][pat[i] >= 0]), return_counts=True)
                    weight = weights / np.sum(weights)  # calculate weight for each voxel
                    # vox_map[vox_inds] = np.column_stack((vox_inds, count * np.ones_like(vox_inds), weight))
                    vox_map.append(vox_inds)
                    vox_weights.append(weight)
                    # pat[i] = np.ones((1, 5, 5), dtype=int)*int(count)
                    pat[i][pat[i] > -1] = int(count)
                    count = count + 1
            unfold_shape = patches.shape
            output_c = unfold_shape[0] * unfold_shape[3]
            output_h = unfold_shape[1] * unfold_shape[4]
            output_w = unfold_shape[2] * unfold_shape[5]
            pat_reshape = pat.reshape(patches.shape)
            down_sample_3d = pat_reshape.transpose([0, 3, 1, 4, 2, 5]).reshape(output_c, output_h, output_w)
            if down_sample_3d.shape != vox_3d.shape:
                down_sample_3d = np.pad(
                    down_sample_3d,
                    [(0, vox_3d.shape[i] - down_sample_3d.shape[i]) for i in range(len(down_sample_3d.shape))],
                    "constant", constant_values=-1)

        # using pytorch
        use_torch = False
        if use_torch:
            import torch
            vox_map = []
            vox_weights = []
            vox_3d_torch = torch.from_numpy(vox_3d)
            kc, kh, kw = down_sample_xyz[2], down_sample_xyz[1], down_sample_xyz[0]  # xyz kernel size
            dc, dh, dw = down_sample_xyz[2], down_sample_xyz[1], down_sample_xyz[0]  # xyz stride
            patches = vox_3d_torch.unfold(0, kc, dc).unfold(1, kh, dh).unfold(2, kw, dw)
            unfold_shape = patches.size()
            patches = patches.contiguous().view(-1, kc, kh, kw)
            count = 0
            for pat in patches:
                if torch.any(pat > -1):
                    elem, ind = torch.sort(pat[pat >= 0])
                    vox_inds, weights = torch.unique(elem, return_counts=True)
                    weight = weights / torch.sum(weights)  # calculate weight for each voxel
                    # vox_map[vox_inds] = np.column_stack((vox_inds, count * np.ones_like(vox_inds), weight))
                    vox_map.append(vox_inds.numpy())
                    vox_weights.append(weight.numpy())
                    # pat[i] = np.ones((1, 5, 5), dtype=int)*int(count)
                    pat[pat > -1] = count
                    count = count + 1
            # Reshape back
            patches_orig = patches.view(unfold_shape)
            output_c = unfold_shape[0] * unfold_shape[3]
            output_h = unfold_shape[1] * unfold_shape[4]
            output_w = unfold_shape[2] * unfold_shape[5]
            patches_orig = patches_orig.permute(0, 3, 1, 4, 2, 5).contiguous()
            down_sample_3d_torch = patches_orig.view(output_c, output_h, output_w)
            down_sample_3d = down_sample_3d_torch.numpy()
            if down_sample_3d.shape != vox_3d.shape:
                down_sample_3d = np.pad(
                    down_sample_3d,
                    [(0, vox_3d.shape[i] - down_sample_3d.shape[i]) for i in range(len(down_sample_3d.shape))],
                    "constant", constant_values=-1)

        return vox_map, vox_weights, down_sample_3d

    def pre_process_voxels(self):
        """
        Updates opt_voxels_dict based upon ct to dose_1d voxel map

        :param ct: object if class CT
        :param structs: object of class CT
        :return:
        """
        if self._down_sample_xyz is not None:
            vox_map, vox_weights, down_sample_3d = self.down_sample_voxels(down_sample_xyz=self._down_sample_xyz)
            self.opt_voxels_dict['ct_to_dose_voxel_map'][0] = down_sample_3d
            self._vox_map = vox_map
            self._vox_weights = vox_weights
            for structure_name in self._structs.structures_dict['name']:
                ind = self._structs.structures_dict['name'].index(structure_name)
                # deep copying mask so that it doesnt modify in structures
                vox_3d = deepcopy(self._structs.structures_dict['structure_mask_3d'][ind]) * self.opt_voxels_dict['ct_to_dose_voxel_map'][0]
                # my_plan.structures_dict['voxel_idx'][i] = np.unique(vox_3d[vox_3d > 0])
                vox, counts = np.unique(vox_3d[vox_3d > 0], return_counts=True)
                self.opt_voxels_dict['voxel_idx'][ind] = vox
                self.opt_voxels_dict['voxel_volume_cc'][ind] = counts * np.prod(
                    self._ct.get_ct_res_xyz_mm())/1000  # calculate weight for each voxel
                # self.opt_voxels_dict['voxel_volume_cc'][ind] = counts / np.max(counts)  # calculate weight for each voxel

    @staticmethod
    def sort_beamlets(b_map):
        c = b_map[b_map >= 0]
        c = np.unique(c)
        ind = np.arange(0, len(c))
        c_sort = np.sort(c)
        matrix_ind = [np.where(b_map == c_i) for c_i in c_sort]
        map_copy = b_map.copy()
        for i in range(len(ind)):
            map_copy[matrix_ind[i]] = ind[i]
        return map_copy

    # for voxels idx methods
    def set_opt_voxel_idx(self, plan_obj, structure_name: str) -> None:
        """
        Set opt_voxel_idx for the struct_name based upon ct_dose_voxel_map and struct_name mask

        :param plan_obj: object of class Plan
        :param structure_name: struct_name name
        :return: set the voxel idx in opt_voxels_dict for the struct_name
        """
        ind = plan_obj.structures.structures_dict['name'].index(structure_name)
        vox_3d = plan_obj.structures.structures_dict['structure_mask_3d'][ind] * \
                 self.opt_voxels_dict['ct_to_dose_voxel_map'][0]
        # my_plan.structures_dict['voxel_idx'][i] = np.unique(vox_3d[vox_3d > 0])
        vox, counts = np.unique(vox_3d[vox_3d > 0], return_counts=True)
        self.opt_voxels_dict['voxel_idx'].append(vox)
        # self.opt_voxels_dict['voxel_volume_cc'].append(counts / np.max(counts))  # calculate weight for each voxel
        self.opt_voxels_dict['voxel_volume_cc'].append(
            counts * np.prod(plan_obj.get_ct_res_xyz_mm())/1000)  # calculate weight for each voxel
        self.opt_voxels_dict['name'].append(structure_name)

    def get_fraction_of_vol_in_calc_box(self, structure_name: str):
        return self._structs.get_fraction_of_vol_in_calc_box(structure_name=structure_name)

    def get_opt_voxels_idx(self, structure_name: str) -> np.ndarray:
        """
        Get voxel index for struct_name
        :param structure_name: name of the struct_name in plan
        :return: voxel indexes for the struct_name
        """
        ind = self.opt_voxels_dict['name'].index(structure_name)
        vox_ind = self.opt_voxels_dict['voxel_idx'][ind]
        # vox_ind = np.where(my_plan.opt_voxels_dict['voxel_structure_map'][0][:, ind] == 1)[0]
        return vox_ind

    def get_opt_voxels_volume_cc(self, structure_name: str):
        """
         :param structure_name: name of the struct_name in plan
         :return: voxel size for the struct_name
         """
        ind = self.opt_voxels_dict['name'].index(structure_name)
        vox_weights = self.opt_voxels_dict['voxel_volume_cc'][ind]
        # vox_ind = np.where(my_plan.opt_voxels_dict['voxel_structure_map'][0][:, ind] == 1)[0]
        return vox_weights

    def get_all_beam_ids(self) -> List[int]:
        ids = [self.beamlets_dict[i]['beam_id'] for i in range(len(self.beamlets_dict))]
        return ids
