import numpy as np
import xarray as xr


class ImageAnalyser():

    def __init__(self) -> None:
        self._image_name = {
            'atoms': 'atoms',
            'background': 'background',
            'dark': 'dark',
            'OD':'OD',
        }
        self._center = None
        self._span = None
        self._fraction = None
    
    @property
    def image_name(self):
        return self._image_name
    
    @image_name.setter
    def image_name(self, value):
        self._image_name.update(value)

    @property
    def center(self):
        return self._center
    
    @center.setter
    def center(self, value):
        self._center = value

    @property
    def span(self):
        return self._span

    @span.setter
    def span(self, value):
        self._span = value

    @property
    def fraction(self):
        return self._fraction

    @fraction.setter
    def fraction(self, value):
        self._fraction = value
    
    def get_offset_from_corner(self, dataArray, x_fraction=None, y_fraction=None, fraction=None, xAxisName='x', yAxisName='y'):
        
        if fraction is None:
            if x_fraction is None:
                x_fraction = self._fraction[0]
        
            if y_fraction is None:
                y_fraction = self._fraction[1]
        else:
            x_fraction = fraction[0]
            y_fraction = fraction[1]

        x_number = dataArray[xAxisName].shape[0]
        y_number = dataArray[yAxisName].shape[0]

        mean = dataArray.isel(x=slice(0, int(x_number * x_fraction)), y=slice(0 , int(y_number * y_fraction))).mean(dim=[xAxisName, yAxisName])
        mean += dataArray.isel(x=slice(0, int(x_number * x_fraction)), y=slice(int(y_number - y_number * y_fraction) , int(y_number))).mean(dim=[xAxisName, yAxisName])
        mean += dataArray.isel(x=slice(int(x_number - x_number * x_fraction) , int(x_number)), y=slice(0 , int(y_number * y_fraction))).mean(dim=[xAxisName, yAxisName])
        mean += dataArray.isel(x=slice(int(x_number - x_number * x_fraction) , int(x_number)), y=slice(int(y_number - y_number * y_fraction) , int(y_number))).mean(dim=[xAxisName, yAxisName])
    
        return mean / 4
    
    def substract_offset(self, dataArray, **kwargs):
        return dataArray - self.get_offset_from_corner(dataArray, **kwargs)
    
    def crop_image(self, dataSet, center=None, span=None):

        if center is None:
            center = self._center
        if span is None:
            span = self._span

        x_start = int(center[0] - span[0] / 2)
        x_end = int(center[0] + span[0] / 2)
        y_end = int(center[1] + span[1] / 2)
        y_start = int(center[1] - span[1] / 2)

        return dataSet.isel(x=slice(x_start, x_end), y=slice(y_start, y_end))
    
    def get_OD(self, imageAtom, imageBackground, imageDrak):

        numerator = np.atleast_1d(imageBackground - imageDrak)
        denominator = np.atleast_1d(imageAtom - imageDrak)
        
        numerator[numerator == 0] = 1
        denominator[denominator == 0] = 1
        imageOD = np.abs(np.divide(denominator, numerator)) 
        imageOD= -np.log(imageOD)

        if len(imageOD) == 1:
            return imageOD[0]
        else:
            return imageOD

    def get_Ncount(self, dataSet, dim=['x', 'y'], **kwargs):
        return dataSet.sum(dim=['x', 'y'], **kwargs)

    def get_absorption_images(self, dataSet, dask='allowed', keep_attrs=True, **kwargs):

        kwargs.update(
            {
                'dask': dask,
                'keep_attrs': keep_attrs,
            }
        )

        dataSet = dataSet.assign(
            {
                self._image_name['OD']: xr.apply_ufunc(self.get_OD, dataSet[self._image_name['atoms']], dataSet[self._image_name['background']], dataSet[self._image_name['dark']], **kwargs)
            }
        )

        return dataSet

    def remove_background(self, dataSet, dask='allowed', keep_attrs=True, **kwargs):

        kwargs.update(
            {
                'dask': dask,
                'keep_attrs': keep_attrs,
            }
        )

        xr.apply_ufunc(self.get_OD, dataSet[self._image_name['atoms']], dataSet[self._image_name['background']], dataSet[self._image_name['dark']], **kwargs)