"""
This module implements corfunc
"""


from typing import Optional, Tuple, Callable
from dataclasses import dataclass
from enum import Enum

import numpy as np
import scipy.optimize
from scipy.interpolate import interp1d
from scipy.signal import argrelextrema
from scipy.fftpack import dct
from scipy.integrate import trapezoid, cumulative_trapezoid

from sas.sascalc.corfunc.calculation_data import (TransformedData,
                                                  LamellarParameters,
                                                  SupplementaryParameters,
                                                  TangentMethod,
                                                  LongPeriodMethod,
                                                  ExtrapolationParameters,
                                                  SettableExtrapolationParameters,
                                                  Fittable,
                                                  PorodData,
                                                  GuinierData)


from sasdata.dataloader.data_info import Data1D
from sas.sascalc.corfunc.transform_thread import FourierThread
from sas.sascalc.corfunc.transform_thread import HilbertThread
from sas.sascalc.corfunc.smoothing import SmoothJoin


class CalculationError(Exception):
    """ Error doing calculation"""
    def __init__(self, msg: str):
        self.msg = msg
        super().__init__(msg)

class CorfuncCalculator:

    def __init__(self,
                 data: Optional[Data1D] = None,
                 extrapolation_parameters: Optional[SettableExtrapolationParameters] = None,
                 long_period_method: Optional[LongPeriodMethod] = None,
                 tangent_method: Optional[TangentMethod] = None):

        """
        Back-end for corfunc calculations

        :param data: Input data (Data1D)
        :param extrapolation_parameters: SettableExtrapolationParameters object containing the q values use to extrapolate
        """

        # Input data
        self._data = data

        # Input parameters
        self._extrapolation_parameters: Optional[SettableExtrapolationParameters] = extrapolation_parameters
        self.tangent_method: Optional[TangentMethod] = tangent_method
        self.long_period_method: Optional[LongPeriodMethod] = long_period_method

        # Fittable parameters
        self._background: Fittable[float] = Fittable()
        self._porod: Fittable[PorodData] = Fittable()
        self._guinier: Fittable[GuinierData] = Fittable()

        # Derived quantities
        self._background_subtracted: Optional[np.ndarray] = None
        self._extrapolation_function: Optional[SmoothJoin] = None
        self._extrapolation_data: Optional[Data1D] = None
        self._transformed_data: Optional[TransformedData] = None
        self._lamellar_parameters: Optional[LamellarParameters] = None
        self._supplementary_parameters: Optional[SupplementaryParameters] = None

    def reset_calculated_values(self):

        """ Resets the calculated values, but does not clear the data or reset the user specified parameters """

        # Fitted parameters
        self._background.clear()
        self._porod.clear()
        self._guinier.clear()

        # Derived quantities
        self._background_subtracted: Optional[np.ndarray] = None
        self._extrapolation_function: Optional[SmoothJoin] = None
        self._extrapolation_data: Optional[Data1D] = None
        self._transformed_data: Optional[TransformedData] = None
        self._lamellar_parameters: Optional[LamellarParameters] = None
        self._supplementary_parameters: Optional[SupplementaryParameters] = None



    #
    # Getters and setters
    #

    @property
    def extrapolation_parameters(self) -> Optional[ExtrapolationParameters]:
        if self._data is None or self._extrapolation_parameters is None:
            return None
        else:
            return ExtrapolationParameters(
                min(self._data.x),
                self._extrapolation_parameters.point_1,
                self._extrapolation_parameters.point_2,
                self._extrapolation_parameters.point_3,
                max(self._data.x))

    @extrapolation_parameters.setter
    def extrapolation_parameters(self, extrap: ExtrapolationParameters):
        self._extrapolation_parameters = SettableExtrapolationParameters(
            extrap.point_1,
            extrap.point_2,
            extrap.point_3)


    @property
    def data(self) -> Optional[Data1D]:
        return self._data

    @data.setter
    def data(self, data: Optional[Data1D]):
        if data is None:
            return

        # Only process data of the class Data1D
        if not issubclass(data.__class__, Data1D):
            raise ValueError("Correlation function cannot be computed with 2D data.")

        self._data = data

    @property
    def q_range(self) -> Tuple[float, float]:
        return self.data.x[0], self.data.x[-1]
    @property
    def background(self):
        if self._background is None:
            return None

        return self._background.data

    @background.setter
    def background(self, value: Optional[float]):
        self._background.data = value

    @property
    def guinier(self):
        if self._guinier is None:
            return None

        return self._guinier.data

    @guinier.setter
    def guinier(self, value: Optional[GuinierData]):
        self._guinier.data = value

    @property
    def porod(self):
        if self._porod is None:
            return None

        return self._porod.data

    @porod.setter
    def porod(self, value: Optional[PorodData]):
        self._porod.data = value

    @property
    def transformed(self):
        return self._transformed_data

    @property
    def lamellar_parameters(self) -> LamellarParameters:
        return self._lamellar_parameters

    @property
    def supplementary_parameters(self) -> Optional[SupplementaryParameters]:
        return self._supplementary_parameters

    @property
    def extrapolation_function(self) -> Optional[Callable]:
        return self._extrapolation_function

    @property
    def min_extrapolated(self) -> Optional[float]:
        if self._extrapolation_data is None:
            return None
        else:
            return np.min(self._extrapolation_data.y)

    @property
    def fit_background(self):
        return self._background.allow_fit

    @fit_background.setter
    def fit_background(self, value: bool):
        self._background.allow_fit = value

    @property
    def fit_guinier(self):
        return self._guinier.allow_fit

    @fit_guinier.setter
    def fit_guinier(self, value: bool):
        self._guinier.allow_fit = value

    @property
    def fit_porod(self):
        return self._porod.allow_fit

    @fit_porod.setter
    def fit_porod(self, value: bool):
        self._porod.allow_fit = value

    @property
    def extrapolated(self):
        return self._extrapolation_data

    #
    # Calculation Steps
    #


    def run(self):
        """ Execute the calculation"""
        self._calculate_background()
        self._calculate_background_subtracted()
        self._calculate_porod_parameters()
        self._calculate_guinier_parameters()
        self._calculate_extrapolation_function()
        self._calculate_extrapolation_data()
        self._calculate_transforms()
        self._calculate_parameters()



    # Calculation components

    def _calculate_background(self):
        """
        Compute the background level from the Porod region of the data,
        only do this if background fitting is allowed
        """
        if self._data is None:
            raise ValueError("Data not set")

        if self._extrapolation_parameters is None:
            raise ValueError("Extrapolation settings not specified")

        if self._background.allow_fit:

            q = self._data.x

            # Fit the last section only
            point_2 = self._extrapolation_parameters.point_2
            point_3 = self._extrapolation_parameters.point_3
            mask = np.logical_and(q > point_2, q < point_3)

            _, _, background = CorfuncCalculator.calculate_porod_parameters(q[mask], self._data.y[mask])

            self._background.data = background

    def _calculate_background_subtracted(self):
        """ Calculate the data with the background removed"""

        if self._data is None:
            raise ValueError("Data not set")

        if self._background.data is None:
            raise ValueError("Background value not set")

        self._background_subtracted = self._data.y - self._background.data


    def _calculate_porod_parameters(self):
        """
        Calculate the Porod parameters
        """

        if self._data is None:
            raise ValueError("Data not set")

        if self._extrapolation_parameters is None:
            raise ValueError("Extrapolation settings not specified")

        if self._porod.allow_fit:
            # Fit the last section only

            q = self.data.x

            point_2 = self._extrapolation_parameters.point_2
            point_3 = self._extrapolation_parameters.point_3
            mask = np.logical_and(q > point_2, q < point_3)

            # Returns an array where the 1st and 2nd elements are the values of k
            # and sigma for the best-fit Porod function
            K, sigma, _ = CorfuncCalculator.calculate_porod_parameters(q[mask], self.data.y[mask])

            self._porod.data = PorodData(K=K, sigma=sigma)



    def _calculate_guinier_parameters(self):
        # Smooths between the best-fit porod function and the data to produce a
        # better fitting curve
        # Returns parameters for the best-fit Guinier function


        if self._data is None:
            raise ValueError("Data not set")

        if self._extrapolation_parameters is None:
            raise ValueError("Extrapolation settings not specified")

        if self._background_subtracted is None:
            raise ValueError("Backgroundless data not set")

        if self._guinier.allow_fit:

            q = self.data.x
            mask = np.logical_and(q < self._extrapolation_parameters.point_1, 0 < q)

            g = CorfuncCalculator.calculate_guinier_parameters(q[mask], self._background_subtracted[mask])

            self._guinier.data = GuinierData(A=g[0], B=g[1])

    def _calculate_extrapolation_function(self):

        if self.data is None:
            raise ValueError("Data not set")

        if self._background.data is None:
            raise ValueError("Background is not set")

        if self._porod.data is None:
            raise ValueError("Porod parameters not set")

        if self._guinier.data is None:
            raise ValueError("Guinier parameters not set")

        data_function = interp1d(self.data.x, self.data.y) # Note that this still has background values

        # Smooth between data and Porod
        s1 = SmoothJoin(
                data_function,
                CorfuncCalculator.porod_function(
                    self._porod.data.K,
                    self._porod.data.sigma,
                    self._background.data),
                self._extrapolation_parameters.point_2,
                self._extrapolation_parameters.point_3)

        # Smooth between the previous function and Guinier
        s2 = SmoothJoin(
                CorfuncCalculator.guinier_function(
                    self._guinier.data.A,
                    self._guinier.data.B,
                    self._background.data),
                s1,
                self.data.x[0],
                self._extrapolation_parameters.point_1)

        self._extrapolation_function = s2

        # params = {'A': g[1], 'B': g[0], 'K': k, 'sigma': sigma}



    def _calculate_extrapolation_data(self):
        """ Numerically evaluate the extrapolation curve """

        if self.data is None:
            raise ValueError("Data not set")

        if self._extrapolation_function is None:
            raise ValueError("Extrapolation function not set")

        q = self.data.x

        extrapolated_q = np.arange(0, q[-1]*100, (q[1]-q[0]))
        extrapolated_I = self._extrapolation_function(extrapolated_q)

        self._extrapolation_data = Data1D(extrapolated_q, extrapolated_I)

    def _calculate_transforms(self):

        """ Calculate the transforms """

        if self.data is None:
            raise ValueError("Data not set")

        if self._background.data is None:
            raise ValueError("Background not set")

        if self._extrapolation_data is None:
            raise ValueError("Extrapolation data not set")


        qs = self._extrapolation_data.x
        iqs = self._extrapolation_data.y
        q = self.data.x
        background = self._background.data

        xs = np.pi * np.arange(len(qs), dtype=np.float32) / (q[1] - q[0]) / len(qs)

        # 1D Correlation Function
        gamma1 = dct((iqs - background) * (qs ** 2))
        Q = np.max(gamma1)
        gamma1 /= Q

        # 3D Correlation Function
        # gamma3(R) = 1/R int_{0}^{R} gamma1(x) dx
        # numerical approximation for increasing R using the trapezium rule
        # Note: SasView 4.x series limited the range to xs <= 1000.0

        gamma3 = cumulative_trapezoid(gamma1, xs) / xs[1:]
        gamma3 = np.hstack((1.0, gamma3))  # gamma3(0) is defined as 1

        # Interface Distribution function
        idf = dct(-qs ** 4 * (iqs - background))

        # Manually calculate IDF(0.0), since scipy DCT tends to give us a
        # very large negative value.

        #    IDF(x) = int_0^inf q^4 * I(q) * cos(q*x) * dq
        # => IDF(0) = int_0^inf q^4 * I(q) * dq

        idf[0] = trapezoid(-qs ** 4 * (iqs - background), qs)
        idf /= Q  # Normalise using scattering invariant

        transform1d = Data1D(xs, gamma1)
        transform3d = Data1D(xs, gamma3)
        idf = Data1D(xs, idf)

        self._transformed_data = TransformedData(transform1d, transform3d, idf)


    def _calculate_parameters(self):
        """
        Extract the interesting measurements from a correlation function
        """

        gamma_1 = self._transformed_data.gamma_1  # 1D transform
        idf = self._transformed_data.idf

        # Calculate indexes of maxima and minima
        z = gamma_1.x
        gamma = gamma_1.y
        gamma_fun = interp1d(z, gamma)

        maxs = argrelextrema(gamma, np.greater)[0]
        mins = argrelextrema(gamma, np.less)[0]

        # If there are no maxima, return None
        if len(maxs) == 0:
            raise CalculationError("No maxima found in data")

        max_values = gamma[maxs]
        largest_max = np.argmax(max_values)

        gamma_min = gamma[mins[0]]  # The value at the first minimum
        z_at_min = z[mins[0]]

        dgamma_dz = (gamma[2:]-gamma[:-2])/(z[2:]-z[:-2])  # 1st derivative of y


        # Find where the second derivative goes to zero
        #  * the IDF is the second derivative of gamma_1
        #  * ... but has a large DC component that needs to be ignored

        above_zero = idf.y[1:] > 0

        zero_crossings = \
            np.argwhere(
                np.logical_xor(
                    above_zero[1:],
                    above_zero[:-1]))[:, 0]

        inflection_point_index = zero_crossings[0] + 1 # +1 for ignoring DC, left side of crossing, not right

        #
        # Work out the tangent index based on the method specified
        #

        tangent_method = self.tangent_method

        if tangent_method is None:
            if inflection_point_index < mins[0]:
                tangent_method = TangentMethod.INFLECTION
            else:
                tangent_method = TangentMethod.HALF_MIN

        if tangent_method == TangentMethod.INFLECTION:
            tangent_index = inflection_point_index
        elif tangent_method == TangentMethod.HALF_MIN:
            tangent_index = mins[0] // 2
        else:
            raise ValueError(f"Unknown tangent calculation method: '{tangent_method}', options are {TangentMethod.options}")

        #
        # Work out the long period index based on the method specified
        #

        long_period_method = self.long_period_method

        if long_period_method is None:
            if len(maxs) > 0:
                long_period_method = LongPeriodMethod.MAX
            else:
                long_period_method = LongPeriodMethod.DOUBLE_MIN

        if long_period_method == LongPeriodMethod.MAX:
            long_period = z[maxs[largest_max]]
        elif long_period_method == LongPeriodMethod.DOUBLE_MIN:
            long_period = z_at_min * 2
        else:
            raise ValueError(f"Unknown long period calculation method: '{long_period_method}', options are {LongPeriodMethod.options}")

        # Try to calculate slope around linear_point using 80 data points
        tangent_region_lower = tangent_index - 40
        tangent_region_upper = tangent_index + 40

        # If too few data points to the left, use linear_point*2 data points
        if tangent_region_lower < 0:
            tangent_region_lower = 0
            tangent_region_upper = inflection_point_index * 2

        # If too few to right, use 2*(dy.size - linear_point) data points
        elif tangent_region_upper > len(dgamma_dz):
            tangent_region_upper = len(dgamma_dz)
            tangent_region_lower = 2*inflection_point_index - dgamma_dz.size

        # Slope at inflection point calculated by mean over inflection region
        tangent_slope = np.mean(dgamma_dz[tangent_region_lower:tangent_region_upper])  # Linear slope
        tangent_intercept = gamma[1:-1][tangent_index]-tangent_slope*z[1:-1][tangent_index]  # Linear intercept

        hard_block_thickness = (gamma_min - tangent_intercept) / tangent_slope  # Hard block thickness
        soft_block_thickness = long_period - hard_block_thickness

        # Find the data points where the graph is linear to within 1%
        mask = np.where(np.abs((gamma-(tangent_slope*z+tangent_intercept))/gamma) < 0.01)[0]

        if len(mask) == 0:  # Return garbage for bad fits
            raise CalculationError("No tangent values found")

        interface_thickness = z[mask[0]]  # Beginning of Linear Section
        core_thickness = z[mask[-1]]  # End of Linear Section

        local_crystallinity = hard_block_thickness / long_period

        gamma_max = gamma[mask[-1]]

        polydispersity_ryan = np.abs(gamma_min / gamma_max)  # Normalized depth of minimum
        polydispersity_stribeck = np.abs(local_crystallinity / ((local_crystallinity - 1) * gamma_max))  # Normalized depth of minimum

        self._supplementary_parameters = SupplementaryParameters(
            tangent_point_z=z[tangent_index],
            tangent_point_gamma=gamma[tangent_index],
            tangent_gradient=float(tangent_slope),
            first_minimum_z=z_at_min,
            first_minimum_gamma=gamma_min,
            first_maximum_z=long_period,
            first_maximum_gamma=gamma_fun(long_period),
            hard_block_z=hard_block_thickness,
            hard_block_gamma=gamma_min,
            interface_z=interface_thickness,
            core_z=core_thickness,
            z_range=(1 / self.data.x[-1], 1 / self.data.x[0]),
            gamma_range=(np.min(gamma), np.max(gamma)))

        self._lamellar_parameters = LamellarParameters(
            long_period=long_period,
            interface_thickness=interface_thickness,
            hard_block_thickness=hard_block_thickness,
            soft_block_thickness=soft_block_thickness,
            core_thickness=core_thickness,
            polydispersity_ryan=polydispersity_ryan,
            polydispersity_stribeck=polydispersity_stribeck,
            local_crystallinity=local_crystallinity)


    #
    # Utility functions / Static methods
    #

    @staticmethod
    def porod_function(K, sigma, background):
        def fun(q):
            return CorfuncCalculator.porod_fitting_function(q, K, sigma, background)

        return fun

    @staticmethod
    def guinier_function(A, B, background):
        def fun(q):
            return np.exp(A + B*q*q) + background

        return fun

    @staticmethod
    def porod_fitting_function(q, K, sigma, background):
        """Equation for the Porod region of the data"""
        return background + (K * q ** (-4)) * np.exp(-q ** 2 * sigma ** 2)

    @staticmethod
    def porod_fitting_function_expected(q, K, sigma, background):
        """ Expected value function used for fitting Porod region (has a q^2 weighting)"""
        return CorfuncCalculator.porod_fitting_function(q, K, sigma, background) * (q * q)

    @staticmethod
    def porod_fitting_function_observed(q, I):
        """ Observed value function used for fitting Porod region"""
        return I * q * q

    @staticmethod
    def calculate_guinier_parameters(q, I):
        """Fit the Guinier region of the curve using linear least squares"""
        A = np.vstack([np.ones(q.shape), q**2]).T

        return np.linalg.lstsq(A, np.log(I))[0]

    @staticmethod
    def calculate_porod_parameters(q, I):
        """Fit the Porod region of the curve"""
        fitp = scipy.optimize.curve_fit(
            CorfuncCalculator.porod_fitting_function_expected, q,
            CorfuncCalculator.porod_fitting_function_observed(q, I),
            bounds=([-np.inf, 0, -np.inf], [np.inf, np.inf, np.inf]))[0]

        k, sigma, bg = fitp
        return k, sigma, bg


def extract_lamellar_parameters(
        data: Data1D,
        guinier_to_data_transition_right: float,
        data_to_porod_transition_left: float,
        data_to_porod_transition_right: float,
        long_period_method: Optional[LongPeriodMethod]=None,
        tangent_method: Optional[TangentMethod]=None):

    """
    Extract lamellar parameters from data using the corfunc calculator

    :param data: Data1D object containing QSpace data
    :param guinier_to_data_transition_right: Q value at which to end the extrapolation from Guiner to data
    :param data_to_porod_transition_left: Q value at which to begin the extrapolation from data to Porod
    :param data_to_porod_transition_right: Q value at which to end the extrapolation from data to Porod
    :param long_period_method: LongPeriodMethod enum value specifying how to calculate the long period (None autodetects)
    :param tangent_method: TangentMethod enum value specifying how to calculate the long period (None autodetects)

    :returns: LamellarParameters object with calculated lamellar parameters
    """

    parameters = SettableExtrapolationParameters(
        guinier_to_data_transition_right,
        data_to_porod_transition_left,
        data_to_porod_transition_right)

    calculator = CorfuncCalculator(
                    data, parameters,
                    long_period_method=long_period_method,
                    tangent_method=tangent_method)

    calculator.run()

    return calculator.lamellar_parameters