from locale import normalize
import numpy as np
import pickle
from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from ..nonconformist import IcpRegressor
from ..nonconformist import RegressorNc
from ..nonconformist import RegressorNormalizer, AbsErrorErrFunc
from ..utils import Progbar


class PredictionInterval(BaseEstimator, RegressorMixin):
    """Class PredictionInterval: Obtain prediction intervals.

    Attributes:

        obj: an object;
            fitted object containing methods `fit` and `predict`

        method: a string;
            method for constructing the prediction intervals.
            Currently "splitconformal" (default) and "localconformal"

        level: a float;
            Confidence level for prediction intervals. Default is 0.95,
            equivalent to a miscoverage error of 0.05

        seed: an integer;
            Reproducibility of fit (there's a random split between fitting and calibration data)
    """

    def __init__(self, obj, method="splitconformal", level=0.95, seed=123):

        self.obj = obj
        self.method = method
        self.level = level
        self.seed = seed
        self.quantile_ = None
        self.icp_ = None

    def fit(self, X, y):
        """Fit the `method` to training data (X, y).

        Args:

            X: array-like, shape = [n_samples, n_features];
                Training set vectors, where n_samples is the number
                of samples and n_features is the number of features.

            y: array-like, shape = [n_samples, ]; Target values.

        """

        X_train, X_calibration, y_train, y_calibration = train_test_split(
            X, y, test_size=0.5, random_state=self.seed
        )

        if self.method == "splitconformal":

            n_samples_calibration = X_calibration.shape[0]
            q = self.level * (1 + 1 / n_samples_calibration)
            self.obj.fit(X_train, y_train)
            preds_calibration = self.obj.predict(X_calibration)
            absolute_residuals = np.abs(y_calibration - preds_calibration)
            try:
                # numpy version >= 1.22
                self.quantile_ = np.quantile(
                    a=absolute_residuals, q=q, method="higher"
                )
            except:
                # numpy version < 1.22
                self.quantile_ = np.quantile(
                    a=absolute_residuals, q=q, interpolation="higher"
                )

        if self.method == "localconformal":

            mad_estimator = RandomForestRegressor()
            normalizer = RegressorNormalizer(
                self.obj, mad_estimator, AbsErrorErrFunc()
            )
            nc = RegressorNc(self.obj, AbsErrorErrFunc(), normalizer)

            self.icp_ = IcpRegressor(nc)
            self.icp_.fit(X_train, y_train)
            self.icp_.calibrate(X_calibration, y_calibration)

        return self

    def predict(self, X, return_pi=False):
        """Obtain predictions and prediction intervals

        Args:

            X: array-like, shape = [n_samples, n_features];
                Testing set vectors, where n_samples is the number
                of samples and n_features is the number of features.

            return_pi: boolean
                Whether the prediction interval is returned or not.
                Default is False, for compatibility with other _estimators_.
                If True, a tuple containing the predictions + lower and upper
                bounds is returned.

        """

        pred = self.obj.predict(X)

        if self.method == "splitconformal":

            if return_pi:

                return pred, (pred - self.quantile_), (pred + self.quantile_)

            else:

                return pred

        if self.method == "localconformal":

            if return_pi:

                predictions_bounds = self.icp_.predict(
                    X, significance=1 - self.level
                )
                return pred, predictions_bounds[:, 0], predictions_bounds[:, 1]

            else:

                return pred
