from piscat.InputOutput.reading_videos import DirectoryType, video_reader
from piscat.Preproccessing.filtering import Filters
from piscat.Localization.particle_localization import PSFsExtraction
from piscat.InputOutput.cpu_configurations import CPUConfigurations

from joblib import Parallel, delayed
from skimage.morphology import area_closing, local_minima
from skimage import feature
from tqdm.autonotebook import tqdm

import pandas as pd
import math
import numpy as np



class BinaryToiSCATLocalization:
    def __init__(self, video_binary, video_iSCAT, df_postion=None, sigma=1.7, area_threshold=4, internal_parallel_flag=True):
        """
        This class is used for localized PSFs based on a binary mask generated by the anomaly.

        Parameters
        ----------
        video_binary: NDArray
            The binary video is the anomaly's output.

        video_iSCAT: NDArray
            The video that the intensity of center PSFs should read from that.

        df_postion: pandas dataframe
            The initial data frame contains the locations of PSFs ( x, y, frame, sigma).

        sigma: float
            The size of the Gaussian kernel that was convolved to a binary mask to find the center of mass.

        area_threshold: int
            The minimum number of connected anomalous pixels that can survive after the morphological procedure is applied.

        internal_parallel_flag: bool
            Internal flag for activating parallel computation. Default is True!
        """

        self.cpu = CPUConfigurations()
        self.video_binary = video_binary
        self.video_iSCAT = video_iSCAT
        self.flag_parallel = internal_parallel_flag

        if video_binary is not None:
            self.sigma = sigma
            if self.cpu.parallel_active and internal_parallel_flag:
                print("\n---start area closing with Parallel---")
                video_closing = Parallel(n_jobs=self.cpu.n_jobs, backend=self.cpu.backend)(delayed(self._kernel_area_closing)(f_, area_threshold=area_threshold) for f_ in tqdm(range(video_binary.shape[0])))
            else:
                print("\n---start area closing without Parallel---")
                video_closing = [area_closing(video_binary[f_], area_threshold=area_threshold) for f_ in tqdm(range(video_binary.shape[0]))]

            self.vide_closing = np.asarray(video_closing)
            self.blure_video = Filters(video=self.vide_closing, inter_flag_parallel_active=internal_parallel_flag).gaussian(sigma=sigma)

            if self.cpu.parallel_active and internal_parallel_flag:
                print("\n---start area local_minima with Parallel---")
                local_minima_ = Parallel(n_jobs=self.cpu.n_jobs, backend=self.cpu.backend)(delayed(self._kernel_local_minima)(f_) for f_ in tqdm(range(self.blure_video.shape[0])))
            else:
                print("\n---start area local_minima without Parallel---")
                local_minima_ = [local_minima(self.blure_video[f_], allow_borders=False, indices=True) for f_ in tqdm(range(self.blure_video.shape[0]))]
            df_pos = self.data_handeling(local_minima_)
            self.df_pos = self.remove_over_lap_localization(df_PSFs=df_pos)
        elif df_postion is not None:
            self.df_pos = df_postion

    def _kernel_area_closing(self, f_, area_threshold=4):
        return area_closing(self.video_binary[f_], area_threshold=area_threshold)

    def _kernel_local_minima(self, f_):
        return local_minima(self.blure_video[f_], indices=True, allow_borders=False)

    def data_handeling(self, local_minima_):
        local_minima_array = np.asarray(local_minima_)
        list_x = []
        list_y = []
        list_f = []
        list_center_intensity = []
        for r_ in range(local_minima_array.shape[0]):
            tmp_ = local_minima_array[r_, :]
            X = tmp_[1].tolist()
            Y = tmp_[0].tolist()
            for x_, y_ in zip(X, Y):
                list_center_intensity = self.blure_video[r_, y_, x_]
                list_x.append(x_)
                list_y.append(y_)
                list_f.append(r_)
        dic_pos = {'y': list_y, 'x': list_x, 'frame': list_f, 'center_intensity': list_center_intensity, 'sigma': self.sigma}

        df_pos = pd.DataFrame.from_dict(dic_pos)
        return df_pos

    def list_frames(self, df_PSFs):
        list_frames = df_PSFs['frame'].tolist()
        list_frames = np.sort(np.unique(list_frames))
        list_frames = list_frames.tolist()
        return list_frames

    def remove_over_lap_localization(self, df_PSFs, threshold=0):
        """
        This filter removes false detections on side lobes of PSFs caused by the localization algorithm by comparing center intensity contrast.

        Parameters
        ----------
        df_PSFs: pandas dataframe
            The data frame contains PSFs locations( x, y, frame, sigma, center_intensity)

        threshold: float
            It specifies the portion of the overlay that two PSFs must have to remove from the list.

        Returns
        -------
        filter_df_PSFs: pandas dataframe
            The filter data frame contains PSFs locations( x, y, frame, sigma, center_intensity)
        """
        if df_PSFs.shape[0] == 0 or df_PSFs is None:
            print('---data frames is empty!---')
            return df_PSFs

        if type(df_PSFs) is pd.core.frame.DataFrame:
            df_PSFs = df_PSFs
        else:
            raise ValueError('Input does not have correct type! This function needs panda data frames.')

        list_frames = self.list_frames(df_PSFs)

        self.particles_after_closeFilter = df_PSFs
        self.index_particles_filter = self.particles_after_closeFilter.index

        num_particles = self.particles_after_closeFilter.shape[0]

        print('\n---Cleaning the df_PSFs that have side lobs without parallel loop---')
        self.point_1 = np.zeros((1, 2), dtype=np.float64)
        self.point_2 = np.zeros((1, 2), dtype=np.float64)
        self.remove_list_close = []

        for frame_num in tqdm(list_frames):
            particle = self.particles_after_closeFilter.loc[self.particles_after_closeFilter['frame'] == frame_num]
            index_list = [index for index in particle.index]
            particle_X = particle['x'].tolist()
            particle_Y = particle['y'].tolist()
            particle_sigma = particle['sigma'].tolist()
            particle_center_intensity = particle['center_intensity'].tolist()

            if len(index_list) != 1:
                for i_ in range(len(particle_X)):
                    self.point_1[0, 0] = particle_X[i_]
                    self.point_1[0, 1] = particle_Y[i_]
                    sigma_1 = particle_sigma[i_]

                    count_ = i_ + 1
                    while count_ <= (len(particle_X) - 1):

                        self.point_2[0, 0] = particle_X[count_]
                        self.point_2[0, 1] = particle_Y[count_]
                        sigma_2 = particle_sigma[count_]

                        distance = math.sqrt(((self.point_1[0, 0] - self.point_2[0, 0]) ** 2) + (
                                (self.point_1[0, 1] - self.point_2[0, 1]) ** 2))
                        tmp = (math.sqrt(2) * (sigma_1 + sigma_2))
                        if distance <= ((math.sqrt(2) * (sigma_1 + sigma_2)) - (threshold * tmp)):

                            intesity_1 = particle_center_intensity[i_]
                            intesity_2 = particle_center_intensity[count_]

                            if np.abs(intesity_1) == np.abs(intesity_2):
                                self.remove_list_close.append(index_list[i_])
                            if np.abs(intesity_1) > np.abs(intesity_2):
                                self.remove_list_close.append(index_list[count_])
                            if intesity_1 < intesity_2:
                                self.remove_list_close.append(index_list[i_])
                        count_ = count_ + 1

        remove_list = list(set(self.remove_list_close))
        self.particles_after_closeFilter = self.particles_after_closeFilter.drop(remove_list, axis=0, errors='ignore')
        print("\nNumber of PSFs before filters = {}".format(num_particles))
        print("\nNumber of PSFs after filters = {}".format(self.particles_after_closeFilter.shape[0]))

        return self.particles_after_closeFilter.reset_index(drop=True)

    def gaussian2D_fit_iSCAT(self, scale=12, internal_parallel_flag=True):
        """
        Using 2D-Gaussian fit localisation on the ROI.

        Parameters
        ----------
        scale: int
            The ROI around PSFs is defined using this scale, which is based on their sigmas.

        internal_parallel_flag: bool
            Internal flag for activating parallel computation. Default is True!

        """
        loc_ = PSFsExtraction(video=self.video_iSCAT)
        df_fit_pos = loc_.fit_Gaussian2D_wrapper(PSF_List=self.df_pos, scale=scale, internal_parallel_flag=internal_parallel_flag)
        df_fit_pos['center_intensity'] = None

        x_fit = df_fit_pos['Fit_X-Center'].tolist()
        y_fit = df_fit_pos['Fit_Y-Center'].tolist()
        y_fit_error = df_fit_pos['Fit_errors_Y-Center'].tolist()
        x_fit_error = df_fit_pos['Fit_errors_X-Center'].tolist()
        frame_fit = df_fit_pos['frame'].tolist()
        center_intensity = []

        for i_ in range(len(frame_fit)):
            if x_fit[i_] != np.nan or y_fit[i_] != np.nan or frame_fit[i_] != np.nan:
                if y_fit_error[i_] <= 1 and x_fit_error[i_] <= 1:
                    center_intensity.append(
                        self.video_iSCAT[int(frame_fit[i_]), int(x_fit[i_]), int(y_fit[i_])])
                else:
                    center_intensity.append(None)
            else:
                center_intensity.append(None)

        df_fit_pos['center_intensity'] = center_intensity

        return df_fit_pos

    def local_extremum_in_window(self, scale=2):
        """
        Reading the ROI's extreme intensity as the PSF's center intensity.

        Parameters
        ----------
        scale: int
            The ROI around PSFs is defined using this scale, which is based on their sigmas.
        """
        self.df_pos['center_intensity'] = None
        self.df2numpy = self.df_pos.to_numpy()

        self.sigma_0 = self.df2numpy[:, 4].tolist()
        self.frame_num = self.df2numpy[:, 2].tolist()
        self.p_x = self.df2numpy[:, 1].tolist()
        self.p_y = self.df2numpy[:, 0].tolist()

        if False:
            print("\n---start local extremum in window with Parallel---")
            center_intensity = Parallel(n_jobs=self.cpu.n_jobs, backend=self.cpu.backend)(
                delayed(self._kernekl_local_extremum_in_window)(f_, scale) for f_ in tqdm(range(self.df2numpy.shape[0])))
        else:
            print("\n---start local extremum in window without Parallel---")
            center_intensity = [self._kernekl_local_extremum_in_window(i_, scale) for i_ in tqdm(range(self.df2numpy.shape[0]))]
        self.df_pos['center_intensity'] = center_intensity

    def local_dog_in_window(self, min_sigma=1.5, max_sigma=1.8, sigma_ratio=1.1, threshold=2.5e-5, overlap=0, scale=8):
        """
        Based on the DOG localization, fine localization, and reading the center intensity of PSFs intensity..

        Parameters
        ----------
        min_sigma: float, list of floats
            The is the minimum standard deviation for the kernel. The lower the value, the smaller blobs can be detected.
            The standard deviations of the filter are given for each axis in sequence or with a single number which is considered for both axis.

        max_sigma: float, list of floats
            The is the maximum standard deviation for the kernel. The higher the value, the bigger blobs can be detected.
            The standard deviations of the filter are given for each axis in sequence or with a single number
            which is considered for both axis.

        sigma_ratio: float
            The ratio between the standard deviation of Kernels which is used for computing the DoG.

        threshold: float
            The absolute lower bound for scale space maxima. Local maxima smaller than thresh are ignored. Reduce this
            to detect blobs with less intensities.

        overlap: float
            A value between 0 and 1. If the area of two blobs are overlapping by a fraction greater than threshold, smaller blobs are eliminated.


        scale: int
            The ROI around PSFs is defined using this scale, which is based on their sigmas.
        """
        self.min_sigma = min_sigma
        self.max_sigma = max_sigma
        self.sigma_ratio = sigma_ratio
        self.threshold = threshold
        self.overlap = overlap
        
        self.df_pos['center_intensity'] = None
        self.df2numpy = self.df_pos.to_numpy()

        self.sigma_0 = self.df2numpy[:, 4].tolist()
        self.frame_num = self.df2numpy[:, 2].tolist()
        self.p_x = self.df2numpy[:, 0].tolist()
        self.p_y = self.df2numpy[:, 1].tolist()

        if self.cpu.parallel_active and self.flag_parallel:
            print("\n---start fine dog localization in window with Parallel---")
            center_intensity = Parallel(n_jobs=self.cpu.n_jobs, backend=self.cpu.backend)(
                delayed(self._kernekl_local_dog_in_window)(f_, scale) for f_ in tqdm(range(self.df2numpy.shape[0])))
        else:
            print("\n---start dog in window without Parallel---")
            center_intensity = [self._kernekl_local_dog_in_window(i_, scale) for i_ in tqdm(range(self.df2numpy.shape[0]))]
        self.df_pos['center_intensity'] = center_intensity
        return self.df_pos

    def _kernekl_local_extremum_in_window(self, i_, scale=2):
        sigma_0 = self.sigma_0[i_]
        frame_num = self.frame_num[i_]
        p_x = self.p_x[i_]
        p_y = self.p_y[i_]

        window_size = round(scale * np.sqrt(2) * sigma_0)

        img_size_x = self.video_iSCAT.shape[1]
        img_size_y = self.video_iSCAT.shape[2]

        start_y = np.max([0, p_y - window_size])
        start_x = np.max([0, p_x - window_size])

        end_y = np.min([img_size_y-1, p_y + window_size])
        end_x = np.min([img_size_x-1, p_x + window_size])

        window_frame = self.video_iSCAT[int(frame_num), int(start_y):int(end_y),
                       int(start_x):int(end_x)]

        if window_frame.shape[0] > 0 and window_frame.shape[1] > 0:
            if np.mean(window_frame) >= 0:
                center_intensity = window_frame.max()
            else:
                center_intensity = window_frame.min()
        else:
            center_intensity = None
        return center_intensity

    def _kernekl_local_dog_in_window(self, i_, scale=12):
        sigma_0 = self.sigma_0[i_]
        frame_num = self.frame_num[i_]
        p_x = self.p_x[i_]
        p_y = self.p_y[i_]

        window_size = round(scale * np.sqrt(2) * sigma_0)

        img_size_x = self.video_iSCAT.shape[1]
        img_size_y = self.video_iSCAT.shape[2]

        start_y = np.max([0, p_y - window_size])
        start_x = np.max([0, p_x - window_size])

        end_y = np.min([img_size_y-1, p_y + window_size])
        end_x = np.min([img_size_x-1, p_x + window_size])

        window_frame = self.video_iSCAT[int(frame_num), int(start_y):int(end_y),
                       int(start_x):int(end_x)]

        PSF_l = PSFsExtraction(video=np.expand_dims(window_frame, axis=0))
        df_PSFs = PSF_l.psf_detection(function='dog',  # function='log', 'doh', 'dog'
                                      min_sigma=self.min_sigma, max_sigma=self.max_sigma,
                                      sigma_ratio=self.sigma_ratio,
                                      threshold=self.threshold,
                                      overlap=self.overlap, mode='BOTH')

        if df_PSFs is not None:
            if df_PSFs.shape[0] == 1:

                x = df_PSFs['x'].to_list()
                y = df_PSFs['y'].to_list()
                center_intensity = window_frame[int(y[0]), int(x[0])]
                return center_intensity

            else:
                # x = df_PSFs['x'].to_list()
                # y = df_PSFs['y'].to_list()
                # x_ = np.mean(x)
                # y_ = np.mean(y)
                # center_intensity = window_frame[int(x_), int(y_)]
                return None

        else:
            return None

