# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
import numpy as np
from sklearn.base import BaseEstimator
from chainladder.tails import TailConstant
from chainladder.development import Development
from chainladder.core.io import EstimatorIO
from chainladder.core.common import Common


class MethodBase(BaseEstimator, EstimatorIO, Common):
    def validate_X(self, X):
        obj = X.copy()
        if "ldf_" not in obj:
            obj = Development().fit_transform(obj)
        if len(obj.ddims) - len(obj.ldf_.ddims) == 1:
            obj = TailConstant().fit_transform(obj)
        return obj

    def _align_cdf(self, ultimate, sample_weight=None):
        """ Vertically align CDF to ultimate vector to origin period latest
        diagonal.
        """
        xp = ultimate.get_array_module()
        from chainladder.utils.utility_functions import num_to_nan

        if self.cdf_.key_labels != ultimate.key_labels and len(self.ldf_.index) > 1:
            level = list(set(self.cdf_.key_labels).intersection(ultimate.key_labels))
            idx = (
                ultimate.index[level]
                .merge(self.cdf_.index[level].reset_index(), how="left", on=level)[
                    "index"
                ]
                .values
            )
            cdf = self.cdf_.values[list(idx.astype(int)), ..., : ultimate.shape[-1]]
        else:
            cdf = self.cdf_.values[..., : ultimate.shape[-1]]
        a = ultimate.iloc[0, 0] * 0 + ultimate.nan_triangle
        if a.array_backend == "sparse":
            a = a - a[a.valuation < a.valuation_date]
        a = a.set_backend(ultimate.array_backend)
        if sample_weight:
            ultimate.values = xp.nan_to_num(ultimate.values * a.values) + xp.nan_to_num(
                sample_weight.values * a.values
            )
        else:
            ultimate.values = xp.nan_to_num(ultimate.values * a.values)
        ultimate.values = num_to_nan(ultimate.values)
        ultimate = ultimate / ultimate
        cdf = ultimate * cdf
        cdf = cdf.latest_diagonal.values
        return cdf

    def _set_ult_attr(self, ultimate):
        """ Ultimate scaffolding """
        xp = ultimate.get_array_module()
        if ultimate.array_backend != "sparse":
            ultimate.values[~xp.isfinite(ultimate.values)] = xp.nan
        ultimate.ddims = np.array([9999])
        ultimate._set_slicers()
        ultimate.valuation_date = ultimate.valuation.max()
        return ultimate

    @property
    def ldf_(self):
        return self.X_.ldf_

    @property
    def latest_diagonal(self):
        return self.X_.latest_diagonal

    def fit(self, X, y=None, sample_weight=None):
        """Applies the chainladder technique to triangle **X**

        Parameters
        ----------
        X : Triangle
            The data used to compute the mean and standard deviation
            used for later scaling along the features axis.
        y : Ignored
        sample_weight : Ignored

        Returns
        -------
        self : object
            Returns the instance itself.
        """
        self.X_ = self.validate_X(X)
        if sample_weight:
            self.sample_weight_ = sample_weight.set_backend(self.X_.array_backend)
        else:
            self.sample_weight_ = sample_weight
        return self

    def predict(self, X, sample_weight=None):
        """Predicts the chainladder ultimate on a new triangle **X**

        Parameters
        ----------
        X : Triangle
            The data used to compute the mean and standard deviation
            used for later scaling along the features axis.
        sample_weight : Triangle
            For exposure-based methods, the exposure to be used for predictions

        Returns
        -------
        X_new: Triangle

        """
        obj = X.copy()
        xp = obj.get_array_module()
        obj.ldf_ = self.ldf_
        if sample_weight:
            sample_weight = sample_weight.set_backend(obj.array_backend)
        obj.ultimate_ = self._get_ultimate(obj, sample_weight)
        return obj

    def fit_predict(self, X, y=None, sample_weight=None):
        self.fit(X, y, sample_weight)
        return self.predict(X, sample_weight)

    def _include_process_variance(self):
        if hasattr(self.X_, "_get_process_variance"):
            full = self.full_triangle_
            obj = self.X_._get_process_variance(full)
            self.ultimate_.values = obj.values[..., -1:]
            process_var = obj - full
        else:
            process_var = None
        return process_var
