import numpy as np
import datetime
import copy
from lmfit.models import PolynomialModel, PseudoVoigtModel


class Section(object):
    def __init__(self):
        self.x = None
        self.y_bgsub = None
        self.y_bg = None  # this is the bg from peakpo
        self.timestamp = None
        self.baseline_in_queue = []  # list of dict, value, constraints
        self.parameters = None
        self.fit_result = None
        self.peaks_in_queue = []  # list of dic, value, constraints
        self.peakinfo = {}
        self._component_cache_token = None
        self._component_cache_bgsub = None
        self._component_cache_with_bg = None

    def get_xrange(self):
        return (self.x.min(), self.x.max())

    def get_yrange(self, bgsub=False):
        if bgsub:
            return (self.y_bgsub.min(), self.y_bgsub.max())
        else:
            return ((self.y_bgsub + self.y_bg).min(),
                    (self.y_bgsub + self.y_bg).max())

    def clear_queue(self):
        self.peaks_in_queue[:] = []
        self.baseline_in_queue[:] = []

    def invalidate_fit_result(self):
        """use with caution"""
        # self.fit_result = None

    def fitted(self):
        if (self.fit_result is None):
            return False
        else:
            return True

    def get_peak_positions(self):
        if self.get_number_of_peaks_in_queue() == 0:
            return []
        else:
            x_c_lst = []
            for peak in self.peaks_in_queue:
                x_c_lst.append(peak['center'])
            return x_c_lst

    def set(self, x, y_bgsub, y_bg):
        """
        it accepts sectioned x, y_bgsub, y_bg only
        """
        self.x = x
        self.y_bgsub = y_bgsub
        self.y_bg = y_bg
        self._invalidate_component_cache()

    def _invalidate_component_cache(self):
        self._component_cache_token = None
        self._component_cache_bgsub = None
        self._component_cache_with_bg = None

    def _get_component_cache_token(self):
        return (
            id(self.fit_result),
            id(self.x),
            id(self.y_bg),
            None if self.x is None else len(self.x),
            None if self.y_bg is None else len(self.y_bg),
        )

    @staticmethod
    def pseudo_voigt_height_to_amplitude(height, sigma, fraction=0.5):
        sigma = max(float(sigma), 1e-5)
        fraction = min(max(float(fraction), 0.0), 1.0)
        gaussian_sigma = sigma / np.sqrt(2.0 * np.log(2.0))
        gaussian_norm = 1.0 / (gaussian_sigma * np.sqrt(2.0 * np.pi))
        lorentz_norm = 1.0 / (np.pi * sigma)
        peak_scale = ((1.0 - fraction) * gaussian_norm) + \
            (fraction * lorentz_norm)
        if peak_scale <= 0.0:
            return 0.0
        return float(height) / peak_scale

    def set_single_peak(self, x_center, fwhm, hkl=[0, 0, 0],
                        phase_name='unknown', y_center=None):
        if self.x is None or len(self.x) == 0:
            return False
        x_min = float(np.min(self.x))
        x_max = float(np.max(self.x))
        # Allow small numerical mismatch at boundaries.
        if len(self.x) > 1:
            tol = 0.5 * float(np.min(np.abs(np.diff(self.x))))
        else:
            tol = 0.0
        if (x_center > x_max + tol) or (x_center < x_min - tol):
            return False
        # Clamp near-boundary clicks into the valid section range.
        x_center = min(max(float(x_center), x_min), x_max)
        if y_center is None:
            y_center = self.get_nearest_intensity(x_center)
        y_center = max(float(y_center), 0.0)
        peak = {}
        peak['center'] = x_center
        peak['amplitude'] = self.pseudo_voigt_height_to_amplitude(
            y_center, fwhm, 0.5)
        peak['sigma'] = fwhm
        peak['fraction'] = 0.5
        peak['center_vary'] = True
        peak['amplitude_vary'] = True
        peak['sigma_vary'] = True
        peak['fraction_vary'] = True
        peak['phasename'] = phase_name
        peak['h'] = hkl[0]
        peak['k'] = hkl[1]
        peak['l'] = hkl[2]
        self.peaks_in_queue.append(peak)
        return True

    def get_order_of_baseline_in_queue(self):
        return self.baseline_in_queue.__len__() - 1

    def set_baseline(self, poly_order):
        old_baseline = copy.deepcopy(self.baseline_in_queue)
        new_baseline = []
        for i in range(poly_order + 1):
            factor = {}
            factor['value'] = 0.
            factor['vary'] = True
            new_baseline.append(factor)
        if old_baseline.__len__() == -1:
            self.baseline_in_queue == new_baseline
            return
        if old_baseline.__len__() >= new_baseline.__len__():
            max_iter = new_baseline.__len__()
        else:
            max_iter = old_baseline.__len__()
        for i in range(max_iter):
            new_baseline[i] = old_baseline[i]
        self.baseline_in_queue = new_baseline

    def peaks_exist(self):
        if self.peaks_in_queue == []:
            return False
        else:
            return True

    def remove_single_peak_nearby(self, x):
        diffs = []
        x_c_lst = self.get_peak_positions()
        index = (np.abs(np.asarray(x_c_lst) - x)).argmin()
        self.peaks_in_queue.pop(index)

    def prepare_for_fitting(self, poly_order, maxwidth, centerrange):
        """
        :param x_center: numpy array of initial x values at picked centers
        :param y_center: numpy array of initial y values at picked centers
        :param fwhm: single float number for initial fwhm value
        """
        self.set_baseline(poly_order)
        baseline_mod = PolynomialModel(poly_order, prefix='b_')
        mod = baseline_mod
        pars = baseline_mod.make_params()
        peakinfo = {}
        for i in range(poly_order + 1):
            prefix = "b_c{0:d}".format(i)
            pars[prefix].set(
                value=self.baseline_in_queue[i]['value'],
                vary=self.baseline_in_queue[i]['vary'])
        i = 0
        for peak in self.peaks_in_queue:
            prefix = "p{0:d}_".format(i)
            peak_mod = PseudoVoigtModel(prefix=prefix, )
            pars.update(peak_mod.make_params())
            pars[prefix + 'center'].set(
                value=peak['center'], min=peak['center']-centerrange,
                max=peak['center']+centerrange,
                vary=peak['center_vary'])
            pars[prefix + 'sigma'].set(
                value=peak['sigma'], min=0.0, vary=peak['sigma_vary'],
                max=maxwidth)
            pars[prefix + 'amplitude'].set(
                value=peak['amplitude'], min=0, vary=peak['amplitude_vary'])
            pars[prefix + 'fraction'].set(
                value=peak['fraction'], min=0., max=1.,
                vary=peak['fraction_vary'])
            peakinfo[prefix + 'phasename'] = peak['phasename']
            peakinfo[prefix + 'h'] = peak['h']
            peakinfo[prefix + 'k'] = peak['k']
            peakinfo[prefix + 'l'] = peak['l']
            mod += peak_mod
            i += 1
        self.parameters = pars
        self.peakinfo = peakinfo
        self.fit_model = mod

    def conduct_fitting(self):
        out = self.fit_model.fit(
            self.y_bgsub, self.parameters, x=self.x)
        self.fit_result = copy.deepcopy(out)
        self._invalidate_component_cache()
        self.timestamp = str(datetime.datetime.now())[:-7]
        self.copy_fit_result_to_queue()
        if self.fit_result is None:
            return False
        else:
            return True

    def get_fit_result(self):
        return self.fit_result.params

    def get_timestamp(self):
        return self.timestamp

    def copy_fit_result_to_queue(self):
        n_peaks = self.get_number_of_peaks_in_queue()
        # self.clear_queue()
        i = 0
        for peak in self.peaks_in_queue:
            prefix = "p{0:d}_".format(i)
            peak['center'] = self.fit_result.params[prefix + 'center'].value
            peak['amplitude'] = self.fit_result.params[
                prefix + 'amplitude'].value
            peak['sigma'] = self.fit_result.params[
                prefix + 'sigma'].value
            peak['fraction'] = self.fit_result.params[
                prefix + 'fraction'].value
            peak['phasename'] = self.peakinfo[prefix + 'phasename']
            peak['h'] = self.peakinfo[prefix + 'h']
            peak['k'] = self.peakinfo[prefix + 'k']
            peak['l'] = self.peakinfo[prefix + 'l']
            i += 1
        i = 0
        for factor in self.baseline_in_queue:
            prefix = "b_c{0:d}".format(i)
            factor['value'] = self.fit_result.params[prefix].value
            i += 1

    def get_number_of_peaks_in_queue(self):
        return self.peaks_in_queue.__len__()

    def _queue_peak_profile(self, peak):
        x = np.asarray(self.x, dtype=float)
        if x.size == 0:
            return np.asarray([], dtype=float)
        center = float(peak['center'])
        sigma = max(float(peak['sigma']), 1e-5)
        fraction = min(max(float(peak.get('fraction', 0.5)), 0.0), 1.0)
        amplitude = max(float(peak['amplitude']), 0.0)
        dx = x - center
        gaussian_sigma = sigma / np.sqrt(2.0 * np.log(2.0))
        gaussian = np.exp(-0.5 * np.square(dx / gaussian_sigma))
        lorentz = 1.0 / (1.0 + np.square(dx / sigma))
        gaussian_norm = 1.0 / (gaussian_sigma * np.sqrt(2.0 * np.pi))
        lorentz_norm = 1.0 / (np.pi * sigma)
        return amplitude * (
            ((1.0 - fraction) * gaussian_norm * gaussian) +
            (fraction * lorentz_norm * lorentz)
        )

    def get_estimated_profiles(self, bgsub=False):
        profiles = {}
        if self.x is None or self.peaks_in_queue == []:
            return profiles
        for i, peak in enumerate(self.peaks_in_queue):
            profile = self._queue_peak_profile(peak)
            if not bgsub:
                profile = profile + self.y_bg
            profiles[f"p{i:d}_"] = profile
        return profiles

    def get_estimated_total_profile(self, bgsub=False):
        if self.x is None:
            return np.asarray([], dtype=float)
        total = np.zeros_like(np.asarray(self.x, dtype=float), dtype=float)
        for peak in self.peaks_in_queue:
            total = total + self._queue_peak_profile(peak)
        if not bgsub:
            total = total + self.y_bg
        return total

    def get_individual_profiles(self, bgsub=False):
        """
        return_value['p1_']
        return_value['b_']
        """
        token = self._get_component_cache_token()
        if token != self._component_cache_token:
            comps = self.fit_result.eval_components(x=self.x)
            self._component_cache_bgsub = comps
            self._component_cache_with_bg = None
            self._component_cache_token = token
        if bgsub:
            return self._component_cache_bgsub
        if self._component_cache_with_bg is None:
            bg_comps = {}
            for key, value in self._component_cache_bgsub.items():
                bg_comps[key] = value + self.y_bg
            self._component_cache_with_bg = bg_comps
        return self._component_cache_with_bg

    def get_fit_profile(self, bgsub=False):
        if bgsub:
            return self.fit_result.best_fit
        else:
            return self.fit_result.best_fit + self.y_bg

    def get_fit_residue(self, bgsub=False):
        return self.y_bgsub - self.fit_result.best_fit + \
            self.get_fit_residue_baseline(bgsub=bgsub)

    def get_fit_residue_baseline(self, bgsub=False):
        if bgsub:
            return 0
            # return self.y_bgsub.min()
        else:
            return self.y_bg.min()

    def get_nearest_intensity(self, x_pick):
        index = (np.abs(np.asarray(self.x) - x_pick)).argmin()
        return self.y_bgsub[index]

    def get_nearest_xy(self, x_pick):
        index = (np.abs(np.asarray(self.x) - x_pick)).argmin()
        return self.x_bgsub[index], self.y_bgsub[index]
