#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Sep 17 22:28:02 2021

@author: dan
"""

import numpy as np
import pandas as pd
from py_tools import data as dt


def reghdfe(df, yvar, xvars, **kwargs):
    """
    High-dimensional fixed-effects regression (convenience wrapper).

    Builds a formula string from ``yvar`` and ``xvars`` and delegates to
    :func:`reghdfe_formula`.

    Parameters
    ----------
    df : pandas.DataFrame
        Input data.
    yvar : str
        Name of the dependent variable.
    xvars : list of str
        Names of independent variables.
    **kwargs
        Additional keyword arguments forwarded to :func:`reghdfe_formula`
        (e.g. ``fes``, ``weight_var``, ``se_type``, ``cluster``).

    Returns
    -------
    coeffs : ndarray of shape (n_regressors,)
        OLS coefficient estimates on the demeaned variables.
    se : ndarray of shape (n_regressors,)
        Robust standard errors.
    """
    formula = "{0} ~ {1}".format(yvar, " + ".join(xvars))
    return reghdfe_formula(df, formula, **kwargs)


def reghdfe_formula(
    df, formula, fes=None, weight_var=None, se_type="robust", tol=1e-12, cluster=None
):
    """
    High-dimensional fixed-effects regression via within-transformation.

    Absorbs fixed effects iteratively (Frisch-Waugh-Lovell), removes singleton
    observations, then runs weighted OLS on the demeaned outcome and
    regressors.

    Parameters
    ----------
    df : pandas.DataFrame
        Input data.
    formula : str
        Patsy-compatible formula string (e.g. ``'y ~ x1 + x2'``).
    fes : list, optional
        Fixed effects to absorb.  Each element may be a string (single FE
        variable) or a list of strings (interaction FE).  Defaults to ``[]``
        (no fixed effects).
    weight_var : str, optional
        Column name of observation weights.  If ``None`` (default), unit
        weights are used.
    se_type : {'robust'}, optional
        Type of standard errors to compute.  Currently only ``'robust'``
        (HC0 sandwich) is supported.  Default is ``'robust'``.
    tol : float, optional
        Convergence tolerance passed to the within-transformation absorber.
        Default is ``1e-12``.
    cluster : object, optional
        Cluster specification for cluster-robust standard errors. Currently not
        implemented; passing a non-``None`` value raises ``NotImplementedError``.

    Returns
    -------
    coeffs : ndarray of shape (n_regressors,)
        OLS coefficient estimates on the demeaned data (intercept absorbed).
    se : ndarray of shape (n_regressors,)
        HC0 robust standard errors.

    Notes
    -----
    Singleton observations (groups with only one observation for any FE level)
    are dropped before estimation.  The degrees-of-freedom adjustment accounts
    for both the number of regressors and the number of absorbed fixed-effect
    groups.
    """
    if fes is None:
        fes = []
    if cluster is not None:
        raise NotImplementedError(
            "cluster-robust standard errors are not implemented in reghdfe_formula yet."
        )

    unweighted = weight_var is None

    fe_vars = []
    for fe in fes:
        if isinstance(fe, list):
            fe_vars += fe
        else:
            fe_vars.append(fe)

    fe_vars = sorted(list(set(fe_vars)))

    try:
        import patsy
    except ImportError as e:
        raise ImportError(
            "patsy is required for reghdfe_formula. "
            "Install it with: pip install py_tools[ml]"
        ) from e

    y, X = patsy.dmatrices(formula, df, return_type="dataframe")

    ix = (
        pd.notnull(y.iloc[:, 0])
        & np.all(pd.notnull(X), axis=1)
        & np.all(pd.notnull(df[fe_vars]), axis=1)
    )

    y = y.loc[ix]
    X = X.loc[ix]

    if fes:
        F = df.loc[ix, fe_vars].copy()

        nonsingle = np.ones(len(F), dtype=bool)
        # Remove singletons
        for fe in fes:
            this_count = F.groupby(fe).transform(len)
            nonsingle &= this_count.values[:, 0] > 1

        ix[ix] = nonsingle
        y = y.loc[nonsingle]
        X = X.loc[nonsingle]
        F = F.loc[nonsingle]

        # Count groups
        Ng = []
        for fe in fes:
            Ng.append(F.groupby(fe).ngroups)

        yvar = y.columns[0]
        xvars = X.columns

        data = pd.concat([y, X, F], axis=1)

        if weight_var is None:
            weight_var = "_w"
            data[weight_var] = 1.0
        else:
            data[weight_var] = df[weight_var]

        y = dt.absorb(
            data,
            fes,
            yvar,
            weight_var=weight_var,
            restore_mean=False,
            display=False,
            tol=tol,
        )

        X_list = []
        for xvar in xvars:
            if xvar != "Intercept":
                X_list.append(
                    dt.absorb(
                        data,
                        fes,
                        xvar,
                        weight_var=weight_var,
                        restore_mean=False,
                        display=False,
                        tol=tol,
                    )
                )
        X = pd.concat(X_list, axis=1)

    y = y.values
    X = X.values

    Nt, Nx = X.shape
    Nt_adj = Nt - Nx
    if fes:
        Nt_adj -= sum(Ng)
        Nt_adj += len(fes)

    if unweighted:
        XW = X.T
    else:
        XW = X.T * df.loc[ix, weight_var].values[np.newaxis, :]

    XWX = XW @ X
    coeffs = np.linalg.solve(XWX, XW @ y)

    e = y - X @ coeffs

    if se_type == "robust":
        XWe = XW * e.T
        S_robust = XWe @ XWe.T
        XWX_inv = np.linalg.inv(XWX)
        V_ols_robust = Nt * (XWX_inv @ (S_robust @ XWX_inv))
        se = np.sqrt(np.diag(V_ols_robust) / Nt_adj)

    return coeffs, se
