from __future__ import annotations

from typing import Optional

import numpy as np
from scipy.optimize import fsolve
from scipy.signal import savgol_filter

try:
    from scipy.integrate import cumulative_trapezoid
except ImportError:
    from scipy.integrate import cumtrapz as cumulative_trapezoid
from scipy.stats import linregress

from .points import Points

from ..dsc.tnmh_model import TNMH


class DSC(Points):
    def __init__(self, x, y, **kwargs):

        y = np.asarray(y).reshape(np.asarray(x).shape)

        x, unique = np.unique(x, return_index=True)
        if kwargs.get('y_err', None) is not None:
            _yerr = np.asarray(kwargs['y_err']).reshape(np.asarray(x).shape)
            kwargs['y_err'] = _yerr[unique]

        self.tg = {'onset': np.nan, 'mid': np.nan, 'end': np.nan, 'inflection': np.nan, 'fictive': np.nan}

        super().__init__(x, y[unique], **kwargs)

    def get_fictive_cp(self, glass: tuple[float, float], liquid: tuple[float, float]) -> ('DSC', float):
        min_glass, max_glass = min(glass), max(glass)
        min_liquid, max_liquid = min(liquid), max(liquid)

        region = self.copy()
        region.cut(min_glass, max_liquid)

        glass_regime = (min_glass < region.x) & (region.x < max_glass)
        regress = linregress(region.x[glass_regime], region.y[glass_regime])
        glass_extrapolation = regress.slope * region.x + regress.intercept

        liquid_regime = (min_liquid < region.x) & (region.x < max_liquid)
        regress2 = linregress(region.x[liquid_regime], region.y[liquid_regime])

        region.y -= glass_extrapolation
        real_area = cumulative_trapezoid(region.y, region.x, initial=0)
        real_area -= real_area[-1]

        t = regress2.intercept - regress.intercept
        m = regress2.slope - regress.slope
        c0 = 0.5 * m * region.x.max() ** 2 + t * region.x.max()

        def equiv(_x, _i):
            return (0.5 * m * _x ** 2 + t * _x - c0) - real_area[_i]

        def equiv_prime(_x, _i):
            return m * _x + t

        fictive_temperature = np.array(
            [fsolve(equiv, region.x[i], fprime=equiv_prime, args=(i,))[0] for i in range(len(region.x))])

        t_g_fictive = fictive_temperature[:20].mean()
        region.y = np.gradient(fictive_temperature, region.x)

        return region, t_g_fictive

    def calculate_tnmh(self, p0: list, glass: tuple[float, float], liquid: tuple[float, float],
                       tg: float = None, num_points: int = 200, return_fictive: bool = True) \
            -> ('FitResult', Optional[float], Optional[DSC]):

        dtf_dt, fictive_tg = self.get_fictive_cp(glass, liquid)
        if tg is None:
            tg = fictive_tg

        temp_equidist = np.linspace(dtf_dt.x[0], dtf_dt.x[-1], num_points)
        dtf_dt_equidist = np.interp(temp_equidist, dtf_dt.x, dtf_dt.y)

        from ..fit.minimizer import FitRoutine
        fitter = FitRoutine()
        fitter.set_model(TNMH)
        data = fitter.add_data(temp_equidist, dtf_dt_equidist)
        data.set_parameter(p0 + [tg, self.value], var=[True] * 4 + [False] * 2, default_bounds=True)

        res = fitter.run()[0]

        if return_fictive:
            return res, tg, dtf_dt
        else:
            return res

    def glass_transition(self, glass, liquid):
        low_idx = tuple(np.argmin(np.abs(self.x - g)) for g in glass)
        high_idx = tuple(np.argmin(np.abs(self.x - l)) for l in liquid)

        x = self.x[low_idx[0]:high_idx[1]]
        y = self.y[low_idx[0]:high_idx[1]]

        win_len = min(len(x) // 20, 51)
        if win_len % 2 == 0:
            win_len += 1
        yy = savgol_filter(y, window_length=win_len, polyorder=1, deriv=1) / np.mean(np.diff(x))

        high_idx = (high_idx[0] - low_idx[0], high_idx[1] - low_idx[0])
        low_idx = (0, low_idx[1] - low_idx[0])

        inflection = np.argmax(yy)

        p1 = linregress(x[low_idx[0]:low_idx[1]], y[low_idx[0]:low_idx[1]])
        glass_baseline = p1.slope * x + p1.intercept

        p2 = linregress(x[high_idx[0]:high_idx[1]], y[high_idx[0]:high_idx[1]])
        liquid_baseline = p2.slope * x + p2.intercept

        tangent_line = yy[inflection] * (x - x[inflection]) + y[inflection]

        onset = np.argmin(np.abs(tangent_line - glass_baseline))
        end = np.argmin(np.abs(tangent_line - liquid_baseline))

        midpoint = np.argmin(np.abs(y - 0.5 * (liquid_baseline[end] - glass_baseline[onset])))
        cut_tangent = np.where((tangent_line > y.min() - 1) & (tangent_line < y.max() + 1))

        glass = Points(x, glass_baseline, name=f'Glass baseline ({self.name})', value=self.value)
        tangent = Points(x[cut_tangent], tangent_line[cut_tangent], name=f'Tangent ({self.name})', value=self.value)
        liquid = Points(x, liquid_baseline, name=f'Liquid baseline ({self.name})', value=self.value)

        ret_dic = {
            'onset': (x[onset], glass_baseline[onset]),
            'midpoint': (x[midpoint], y[midpoint]),
            'end': (x[end], liquid_baseline[end]),
            'inflection': (x[inflection], y[inflection]),
        }

        self.tg.update(ret_dic)

        return ret_dic, glass, liquid, tangent
