import re

# import lmfit
import numpy as np
import pandas as pd
from scipy.optimize import least_squares


class Calibrate:
    def __init__(self, model):
        """Initialize Calibration class.

        Parameters
        ----------
        model : ttim.Model
            model to calibrate
        """

        self.model = model
        self.parameters = pd.DataFrame(
            columns=["optimal", "std", "perc_std", "pmin", "pmax", "initial", "parray"]
        )
        self.seriesdict = {}
        self.seriesinwelldict = {}

    def set_parameter(self, name=None, initial=0, pmin=-np.inf, pmax=np.inf):
        """Set parameter to be optimized.

        Parameters
        ----------
        name : str
            parameter name, can include layer information.
            name can be 'kaq', 'Saq' or 'c'. A number after the parameter
            name denotes the layer number, i.e. 'kaq0' refers to the hydraulic
            conductivity of layer 0.
            name also supports layer ranges, entered by adding a '_' and a
            layer number, i.e. 'kaq0_3' denotes conductivity for layers 0 up to
            and including 3.
        initial : float, optional
            initial value for the parameter (the default is 0)
        pmin : float, optional
            lower bound for parameter value (the default is -np.inf)
        pmax : float, optional
            upper bound for paramater value (the default is np.inf)
        """

        assert isinstance(name, str), "Error: name must be string"
        # find numbers in name str for support layer ranges
        layers_from_name = re.findall(r"\d+", name)
        p = None
        if "_" in name:
            fromlay, tolay = [int(i) for i in layers_from_name]
            if name[:3] == "kaq":
                p = self.model.aq.kaq[fromlay : tolay + 1]
            elif name[:3] == "Saq":
                p = self.model.aq.Saq[fromlay : tolay + 1]
            elif name[0] == "c":
                p = self.model.aq.c[fromlay : tolay + 1]
            elif name[:3] == "Sll":
                p = self.model.aq.Sll[fromlay : tolay + 1]
            elif name[0:8] == "kzoverkh":
                p = self.model.aq.kzoverkh[fromlay : tolay + 1]
        else:
            layer = int(layers_from_name[0])
            # Set, kaq, Saq, c
            if name[:3] == "kaq":
                p = self.model.aq.kaq[layer : layer + 1]
            elif name[:3] == "Saq":
                p = self.model.aq.Saq[layer : layer + 1]
            elif name[0] == "c":
                p = self.model.aq.c[layer : layer + 1]
            elif name[:3] == "Sll":
                p = self.model.aq.Sll[layer : layer + 1]
            elif name[0:8] == "kzoverkh":
                p = self.model.aq.kzoverkh[layer : layer + 1]
        if p is None:  # no parameter set
            print("parameter name not recognized or no parameter ref supplied")
            return
        self.parameters.loc[name] = {
            "optimal": initial,
            "std": None,
            "perc_std": None,
            "pmin": pmin,
            "pmax": pmax,
            "initial": initial,
            "parray": p[:],
        }

    def set_parameter_by_reference(
        self, name=None, parameter=None, initial=0, pmin=-np.inf, pmax=np.inf
    ):
        """Set parameter to be optimized.

        Parameters
        ----------
        name : str
            parameter name
        parameter : np.array
            array reference containing the parameter to be optimized. must be
            specified as reference, i.e. w.rc[0:]
        initial : float, optional
            initial value for the parameter (the default is 0)
        pmin : float, optional
            lower bound for parameter value (the default is -np.inf)
        pmax : float, optional
            upper bound for paramater value (the default is np.inf)
        """
        assert isinstance(name, str), "Error: name must be string"
        if parameter is not None:
            assert isinstance(
                parameter, np.ndarray
            ), "Error: parameter needs to be numpy array"
            p = parameter
        self.parameters.loc[name] = {
            "optimal": initial,
            "std": None,
            "perc_std": None,
            "pmin": pmin,
            "pmax": pmax,
            "initial": initial,
            "parray": p[:],
        }

    def series(self, name, x, y, layer, t, h, weights=None):
        """Method to add observations to Calibration object.

        Parameters
        ----------
        name : str
            name of series
        x : float
            x-coordinate
        y : float
            y-coordinate
        layer : int
            layer number, 0-indexed
        t : np.array
            array containing timestamps of timeseries
        h : np.array
            array containing timeseries values, i.e. head observations
        """

        s = Series(x, y, layer, t, h, weights=weights)
        self.seriesdict[name] = s

    def seriesinwell(self, name, element, t, h):
        """Method to add observations to Calibration object.

        Parameters
        ----------
        name : str
            name of series
        element: element object with headinside function
        t : np.array
            array containing timestamps of timeseries
        h : np.array
            array containing timeseries values, i.e. head observations
        """

        e = SeriesInWell(element, t, h)
        self.seriesinwelldict[name] = e

    def residuals(self, p, printdot=False, weighted=True, layers=None, series=None):
        """Method to calculate residuals given certain parameters.

        Parameters
        ----------
        p : np.array
            array containing parameter values
        printdot : bool, optional
            print dot for each function call

        Returns
        -------
        np.array
            array containing all residuals
        """

        if printdot:
            print(".", end="")
        # set the values of the variables

        if printdot == 7:
            print(p)

        if layers is None:
            layers = range(self.model.aq.naq)

        for i, k in enumerate(self.parameters.index):
            # [:] needed to do set value in array
            self.parameters.loc[k, "parray"][:] = p[i]

        self.model.solve(silent=True)

        rv = np.empty(0)
        cal_series = self.seriesdict.keys() if series is None else series
        for key in cal_series:
            s = self.seriesdict[key]
            if s.layer not in layers:
                continue
            h = self.model.head(s.x, s.y, s.t, layers=s.layer)
            w = s.weights if ((s.weights is not None) and weighted) else np.ones_like(h)
            rv = np.append(rv, (s.h - h) * w)
        for key in self.seriesinwelldict:
            s = self.seriesinwelldict[key]
            h = s.element.headinside(s.t)[0]
            rv = np.append(rv, s.h - h)
        return rv

    def residuals_lmfit(self, lmfitparams, printdot=False):
        vals = lmfitparams.valuesdict()
        p = np.array([vals[k] for k in self.parameters.index])
        # p = np.array([vals[k] for k in vals])
        return self.residuals(p, printdot)

    def fit_least_squares(self, report=True, diff_step=1e-4, xtol=1e-8, method="lm"):
        self.fitresult = least_squares(
            self.residuals,
            self.parameters.initial.values,
            args=(True,),
            bounds=(self.parameters.pmin.values, self.parameters.pmax.values),
            method=method,
            diff_step=diff_step,
            xtol=xtol,
            x_scale="jac",
        )
        print("", flush=True)
        # Call residuals to specify optimal values for model
        res = self.residuals(self.fitresult.x)
        for ipar in self.parameters.index:
            self.parameters.loc[ipar, "optimal"] = self.parameters.loc[ipar, "parray"][
                0
            ]
        nparam = len(self.fitresult.x)
        H = self.fitresult.jac.T @ self.fitresult.jac
        sigsq = np.var(res, ddof=nparam)
        self.covmat = np.linalg.inv(H) * sigsq
        self.sig = np.sqrt(np.diag(self.covmat))
        D = np.diag(1 / self.sig)
        self.cormat = D @ self.covmat @ D
        self.parameters["std"] = self.sig
        self.parameters["perc_std"] = self.sig / self.parameters["optimal"] * 100
        if report:
            print(self.parameters)
            print(self.sig)
            print(self.covmat)
            print(self.cormat)

    def fit_lmfit(self, report=True, printdot=True):
        import lmfit

        self.lmfitparams = lmfit.Parameters()
        for name in self.parameters.index:
            p = self.parameters.loc[name]
            self.lmfitparams.add(name, value=p["initial"], min=p["pmin"], max=p["pmax"])
        fit_kws = {"epsfcn": 1e-4}
        self.fitresult = lmfit.minimize(
            self.residuals_lmfit,
            self.lmfitparams,
            method="leastsq",
            kws={"printdot": printdot},
            **fit_kws,
        )
        print("", flush=True)
        print(self.fitresult.message)
        if self.fitresult.success:
            for name in self.parameters.index:
                self.parameters.loc[
                    name, "optimal"
                ] = self.fitresult.params.valuesdict()[name]
            if hasattr(self.fitresult, "covar"):
                self.parameters["std"] = np.sqrt(np.diag(self.fitresult.covar))
                self.parameters["perc_std"] = (
                    100 * self.parameters["std"] / np.abs(self.parameters["optimal"])
                )
            else:
                self.parameters["std"] = np.nan
                self.parameters["perc_std"] = np.nan
        if report:
            print(lmfit.fit_report(self.fitresult))

    def fit(self, report=True, printdot=True):
        # current default fitting routine
        return self.fit_lmfit(report, printdot)

    def rmse(self, weighted=True, layers=None):
        """Calculate root-mean-squared-error.

        Returns
        -------
        float
            return rmse value
        """

        r = self.residuals(
            self.parameters["optimal"].values, weighted=weighted, layers=layers
        )
        return np.sqrt(np.mean(r**2))


class Series:
    def __init__(self, x, y, layer, t, h, weights=None):
        self.x = x
        self.y = y
        self.layer = layer
        self.t = t
        self.h = h
        self.weights = weights


class SeriesInWell:
    def __init__(self, element, t, h):
        self.element = element
        self.t = t
        self.h = h
