import warnings

import numpy as np
import pandas as pd
import statsmodels.api as sm
from scipy.linalg import sqrtm
from scipy.stats import norm


class BLP:
    r"""Best linear predictor (BLP) with orthogonal signals.
    Mainly used for CATE and GATE estimation for IRM models.

    The Best Linear Predictor (BLP) targets the coefficient vector :math:`\beta_0` that minimizes the mean squared error
    between the true treatment effect function :math:`\tau(X)` and a linear combination of basis functions :math:`b(X)`:

    .. math::
        \beta_0 = \arg\min_{\beta \in \mathbb{R}^K} \mathbb{E}\Big[\big(\tau(X) - b(X)^\top \beta \big)^2\Big].

    This is characterized by the moment condition:

    .. math::
        \mathbb{E}[b(X)\psi] = \mathbb{E}[b(X)b(X)^\top]\beta_0,

    where :math:`\psi` is the orthogonal signal such that :math:`\mathbb{E}[\psi \mid X] = \tau(X)`.

    The estimator is obtained via OLS of the orthogonal signal on the basis:

    .. math::
        \hat{\beta} = (B^\top B)^{-1}B^\top\psi.

    **GATE (Group Average Treatment Effect)**

    When ``is_gate=True``, the basis consists of group indicators (dummy variables).
    In this case, the BLP coefficients correspond to the group means of the orthogonal signal,
    which approximate the GATEs:

    .. math::
        \hat{\beta}_k = \frac{1}{n_k}\sum_{i:G_i=k}\psi_i \approx \text{GATE}_k.

    **Confidence Intervals**

    Confidence intervals for any linear combination :math:`\hat{g} = A\hat{\beta}` are computed using the estimated covariance matrix :math:`\widehat{\Omega}`:

    .. math::
        \widehat{\operatorname{Var}}(\hat{g}) \approx A\widehat{\Omega}A^\top.

    Pointwise and joint confidence intervals (via Gaussian multiplier bootstrap) are supported.

    Parameters
    ----------
    orth_signal : :class:`numpy.array`
        The orthogonal signal to be predicted. Has to be of shape ``(n_obs,)``,
        where ``n_obs`` is the number of observations.

    basis : :class:`pandas.DataFrame`
        The basis for estimating the best linear predictor. Has to have the shape ``(n_obs, d)``,
        where ``n_obs`` is the number of observations and ``d`` is the number of predictors.

    is_gate : bool
        Indicates whether the basis is constructed for GATEs (dummy-basis).
        Default is ``False``.
    """

    def __init__(self, orth_signal, basis, is_gate=False):
        if not isinstance(orth_signal, np.ndarray):
            raise TypeError(f"The signal must be of np.ndarray type. Signal of type {str(type(orth_signal))} was passed.")

        if orth_signal.ndim != 1:
            raise ValueError(
                f"The signal must be of one dimensional. Signal of dimensions {str(orth_signal.ndim)} was passed."
            )

        if not isinstance(basis, pd.DataFrame):
            raise TypeError(f"The basis must be of DataFrame type. Basis of type {str(type(basis))} was passed.")

        if not basis.columns.is_unique:
            raise ValueError("Invalid pd.DataFrame: Contains duplicate column names.")

        self._orth_signal = orth_signal
        self._basis = basis
        self._is_gate = is_gate

        # initialize the score and the covariance
        self._blp_model = None
        self._blp_omega = None

    def __str__(self):
        class_name = self.__class__.__name__
        header = f"================== {class_name} Object ==================\n"
        fit_summary = str(self.summary)
        res = header + "\n------------------ Fit summary ------------------\n" + fit_summary
        return res

    @property
    def blp_model(self):
        """
        Best-Linear-Predictor model.
        """
        return self._blp_model

    @property
    def orth_signal(self):
        """
        Orthogonal signal.
        """
        return self._orth_signal

    @property
    def basis(self):
        """
        Basis.
        """
        return self._basis

    @property
    def blp_omega(self):
        """
        Covariance matrix.
        """
        return self._blp_omega

    @property
    def summary(self):
        """
        A summary for the best linear predictor effect after calling :meth:`fit`.
        """
        col_names = ["coef", "std err", "t", "P>|t|", "[0.025", "0.975]"]
        if self.blp_model is None:
            df_summary = pd.DataFrame(columns=col_names)
        else:
            summary_stats = {
                "coef": self.blp_model.params,
                "std err": self.blp_model.bse,
                "t": self.blp_model.tvalues,
                "P>|t|": self.blp_model.pvalues,
                "[0.025": self.blp_model.conf_int()[0],
                "0.975]": self.blp_model.conf_int()[1],
            }
            df_summary = pd.DataFrame(summary_stats, columns=col_names)
        return df_summary

    def fit(self, cov_type="HC0", diagnostic_data: bool = True, **kwargs):
        """
        Estimate BLP models.

        Parameters
        ----------
        cov_type : str
            The covariance type to be used in the estimation. Default is ``'HC0'``.
            See :meth:`statsmodels.regression.linear_model.OLS.fit` for more information.

        diagnostic_data : bool, default True
            Whether to include diagnostic data_contracts. (Currently not used for BLP).

        **kwargs: dict
            Additional keyword arguments to be passed to :meth:`statsmodels.regression.linear_model.OLS.fit`.

        Returns
        -------
        self : object
        """

        # fit the best-linear-predictor of the orthogonal signal with respect to the grid
        self._blp_model = sm.OLS(self._orth_signal, self._basis).fit(cov_type=cov_type, **kwargs)
        self._blp_omega = self._blp_model.cov_params().to_numpy()

        return self

    def confint(self, basis=None, joint=False, alpha=0.05, n_rep_boot=500):
        """
        Confidence intervals for the BLP model.

        Parameters
        ----------
        basis : :class:`pandas.DataFrame`
            The basis for constructing the confidence interval. Has to have the same form as the basis from
            the construction. If ``None`` is passed, if the basis is constructed for GATEs, the GATEs are returned.
            Else, the confidence intervals for the basis coefficients are returned (with pointwise cofidence intervals).
            Default is ``None``.

        joint : bool
            Indicates whether joint confidence intervals are computed.
            Default is ``False``.

        alpha : float
            The significance level.
            Default is ``0.05``.

        n_rep_boot : int
            The number of bootstrap repetitions (only relevant for joint confidence intervals).
            Default is ``500``.

        Returns
        -------
        df_ci : pd.DataFrame
            A data_contracts frame with the confidence interval(s).
        """
        if not isinstance(joint, bool):
            raise TypeError(f"joint must be True or False. Got {str(joint)}.")

        if not isinstance(alpha, float):
            raise TypeError(f"The significance level must be of float type. {str(alpha)} of type {str(type(alpha))} was passed.")
        if (alpha <= 0) | (alpha >= 1):
            raise ValueError(f"The significance level must be in (0,1). {str(alpha)} was passed.")

        if not isinstance(n_rep_boot, int):
            raise TypeError(
                "The number of bootstrap replications must be of int type. "
                f"{str(n_rep_boot)} of type {str(type(n_rep_boot))} was passed."
            )
        if n_rep_boot < 1:
            raise ValueError(f"The number of bootstrap replications must be positive. {str(n_rep_boot)} was passed.")

        if self._blp_model is None:
            raise ValueError("Apply fit() before confint().")

        gate_names = None
        # define basis if none is supplied
        if basis is None:
            if self._is_gate:
                # reduce to unique groups
                basis = pd.DataFrame(np.diag(v=np.full((self._basis.shape[1]), True)))
                gate_names = list(self._basis.columns.values)
            else:
                if joint:
                    warnings.warn("Returning pointwise confidence intervals for basis coefficients.", UserWarning)
                # return the confidence intervals for the basis coefficients
                ci = np.vstack(
                    (
                        self.blp_model.conf_int(alpha=alpha / 2)[0],
                        self.blp_model.params,
                        self.blp_model.conf_int(alpha=alpha / 2)[1],
                    )
                ).T
                df_ci = pd.DataFrame(
                    ci,
                    columns=["{:.1f} %".format(alpha / 2 * 100), "effect", "{:.1f} %".format((1 - alpha / 2) * 100)],
                    index=self._basis.columns,
                )
                return df_ci

        elif not (basis.shape[1] == self._basis.shape[1]):
            raise ValueError("Invalid basis: DataFrame has to have the exact same number and ordering of columns.")
        elif not list(basis.columns.values) == list(self._basis.columns.values):
            raise ValueError("Invalid basis: DataFrame has to have the exact same number and ordering of columns.")

        # blp of the orthogonal signal
        g_hat = self._blp_model.predict(basis)

        np_basis = basis.to_numpy()
        # calculate se for basis elements
        blp_se = np.sqrt((np.dot(np_basis, self._blp_omega) * np_basis).sum(axis=1))

        if joint:
            # calculate the maximum t-statistic with bootstrap
            normal_samples = np.random.normal(size=[basis.shape[1], n_rep_boot])
            bootstrap_samples = np.multiply(np.dot(np_basis, np.dot(sqrtm(self._blp_omega), normal_samples)).T, (1.0 / blp_se))

            max_t_stat = np.quantile(np.max(np.abs(bootstrap_samples), axis=0), q=1 - alpha)

            # Lower simultaneous CI
            g_hat_lower = g_hat - max_t_stat * blp_se
            # Upper simultaneous CI
            g_hat_upper = g_hat + max_t_stat * blp_se

        else:
            # Lower point-wise CI
            g_hat_lower = g_hat + norm.ppf(q=alpha / 2) * blp_se
            # Upper point-wise CI
            g_hat_upper = g_hat + norm.ppf(q=1 - alpha / 2) * blp_se

        ci = np.vstack((g_hat_lower, g_hat, g_hat_upper)).T
        df_ci = pd.DataFrame(
            ci,
            columns=["{:.1f} %".format(alpha / 2 * 100), "effect", "{:.1f} %".format((1 - alpha / 2) * 100)],
            index=basis.index,
        )

        if self._is_gate and gate_names is not None:
            df_ci.index = gate_names

        return df_ci
