# 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/.

from chainladder.methods.chainladder import Chainladder
from sklearn.base import BaseEstimator, TransformerMixin
import numpy as np
from chainladder.utils.cupy import cp
import copy
import warnings
from chainladder.core import EstimatorIO

class BerquistSherman(BaseEstimator, TransformerMixin, EstimatorIO):
    """
    Class to alter the inner diagonals of a Triangle using the methods described
    by Berquist and Sherman.

    Parameters
    ----------
    paid_amount : str
        The triangle column associated with paid amounts
    incurred_amount : str
        The triangle column associated with incurred amounts
    reported_count : str
        The triangle column associated with reported claim count
    closed_count : str
        The triangle column associated with closed claim counts
    trend : float (default=0.0)
        Trend rate underlying average open case reserves.
    reported_count_estimator : Estimator
        An IBNR estimator for reported_count used to calculate closed_count
        disposal rates.  Estimator can be a Pipeline.  If None selected then
        basic Chainladder model is used.

    Attributes
    ----------
    adjusted_triangle_ : Triangle
        A set of triangles represented by each simulation
    disposal_rate_ : Triangle
        The disposal rates of closed claims based on the reported_count_estimator
    a_ : Triangle
        Two-period Exponential intercept parameters
    b_ : Triangle
        Two-period Exponential slope parameters
    """
    def __init__(self, paid_amount=None, incurred_amount=None, reported_count=None,
                 closed_count=None, trend=0.0, reported_count_estimator=None):
        self.paid_amount = paid_amount
        self.incurred_amount = incurred_amount
        self.reported_count = reported_count
        self.closed_count = closed_count
        self.trend = trend
        self.reported_count_estimator = reported_count_estimator


    def fit(self, X, y=None, sample_weight=None):
        xp = cp.get_array_module(X.values)
        if not (self.paid_amount in X.columns and self.incurred_amount in X.columns
                and self.reported_count in X.columns and self.closed_count in X.columns):
            raise ValueError('Must enter values valid columns for paid_amount, incurred_amount, reported_count and closed_count')
        obj = copy.copy(X)
        paid_amount = self.paid_amount
        incurred_amount = self.incurred_amount
        reported_count = self.reported_count
        closed_count = self.closed_count
        reported_count_estimator = self.reported_count_estimator
        # Case reserve adequacy adjustment
        open_count = obj[reported_count] - obj[closed_count]
        avg_case = (obj[incurred_amount] - obj[paid_amount]) / open_count
        adj_avg_case = (
            avg_case.trend(1/(1+self.trend)-1, axis='valuation')/ avg_case *
            avg_case[avg_case.valuation==avg_case.valuation_date].sum('origin'))
        adj_incurred_amount = adj_avg_case * open_count + obj[paid_amount]

        # Paid and closed claim adjustments
        if reported_count_estimator is None:
            reported_count_estimator = Chainladder()
        reported_count_estimator.fit(obj[reported_count])
        if reported_count_estimator.__class__.__name__ == 'Pipeline':
            rep_cnt_ult = reported_count_estimator.named_steps[
                reported_count_estimator.steps[-1][0]].ultimate_
        else:
            rep_cnt_ult = reported_count_estimator.ultimate_
        disposal_rate = (obj[closed_count] / rep_cnt_ult)
        adj_closed_clm = (
            (disposal_rate*0+1)*disposal_rate[
                disposal_rate.valuation==X.valuation_date].mean('origin') *
            rep_cnt_ult)
        adj_closed_clm.valuation_date = disposal_rate.valuation_date = X.valuation_date

        # Two-period exponential regression
        y = obj[paid_amount]
        x = obj[closed_count]
        x0 = x.values[...,:-1]
        x1 = x.values[..., 1:]
        y0 = np.log(y.values[...,:-1])
        y1 = np.log(y.values[..., 1:])
        b = ((x0*y0+x1*y1)/2-(x0+x1)/2*(y0+y1)/2) / \
            ((x0**2+x1**2)/2-((x0+x1)/2)**2)
        a = np.exp((y0+y1)/2-b*(x0+x1)/2)

        # Need to Ignore warnings for NaN cells
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            lookup = np.maximum(
                np.concatenate(
                    [(adj_closed_clm.values[...,i:i+1]>obj[closed_count].values).sum(
                        axis=-1, keepdims=True)
                     for i in range(adj_closed_clm.shape[-1])], axis=-1),1)-1
        a = np.concatenate(
            [a[..., i, lookup[0, 0, i:i+1, :]]
             for i in range(lookup.shape[-2])], -2) * \
            adj_closed_clm._nan_triangle()[None,None,...]
        b = np.concatenate(
            [b[..., i, lookup[0, 0, i:i+1, :]]
             for i in range(lookup.shape[-2])], -2) * \
            adj_closed_clm._nan_triangle()[None,None,...]
        # Adjust paids
        adj_paid_claims = adj_closed_clm*0+np.exp(adj_closed_clm.values*b)*a
        adj_paid_claims = (
            y[y.valuation==y.valuation_date] +
            adj_paid_claims[adj_paid_claims.valuation<adj_paid_claims.valuation_date])

        adjusted_triangle_ = copy.deepcopy(X)
        adjusted_triangle_[paid_amount] = adj_paid_claims
        adjusted_triangle_[incurred_amount] = adj_incurred_amount
        adjusted_triangle_[closed_count] = adj_closed_clm
        self.adjusted_triangle_ = adjusted_triangle_
        self.disposal_rate_ = disposal_rate
        self.a_ = a
        self.b_ = b
        return self


    def transform(self, X):
        """ If X and self are of different shapes, align self to X, else
        return self.

        Parameters
        ----------
        X : Triangle
            The triangle to be transformed

        Returns
        -------
            X_new : New triangle with transformed attributes.
        """
        X_new = copy.deepcopy(X)
        X_new[self.paid_amount] = self.adjusted_triangle_[self.paid_amount]
        X_new[self.incurred_amount] = self.adjusted_triangle_[self.incurred_amount]
        X_new[self.reported_count] = self.adjusted_triangle_[self.reported_count]
        X_new[self.closed_count] = self.adjusted_triangle_[self.closed_count]
        X_new.a_ = self.a_
        X_new.b_ = self.b_
        return X_new

    def set_params(self, **params):
        from chainladder.utils.utility_functions import read_json
        if type(params['reported_count_estimator']) is str:
            params['reported_count_estimator'] = read_json(params['reported_count_estimator'])
        return super().set_params(**params)
