"""Import the original deviations design and evaluate its joint F restriction."""
from pathlib import Path
import json, sys, warnings, re
from importlib.metadata import version
import numpy as np
import pandas as pd
import pyfixest as pf
from scipy import stats
from scipy.sparse import coo_matrix
from scipy.sparse.csgraph import connected_components


def run(config):
    sys.path.insert(0, str(Path(config['reference_repo']) / 'code'))
    from saturated import test_treatment_heterogeneity
    data = pd.read_csv(config['data'])
    unit, time = config['unit'], config['time']
    with warnings.catch_warnings():
        warnings.simplefilter('ignore')
        model = test_treatment_heterogeneity(
            data, outcome=config['outcome'], treatment=config['treatment'],
            unit_id=unit, time_id=time, retmod=True,
        )
        extra = config.get('covariates', [])
        weight = config.get('weight')
        if extra or weight:
            lhs, fe = model._fml.split('|')
            formula = lhs + (' + ' + ' + '.join(extra) if extra else '') + ' | ' + fe
            model = pf.feols(formula, model._data, vcov={'CRV1':config['cluster']}, weights=weight)
    names = pd.Series(model.coef().index)
    event = names.str.extract(r'\[T\.(-?\d+(?:\.\d+)?)\]', expand=False).astype(float)
    selected = names.str.contains('cohort_dummy_') & event.gt(0)
    horizon = config.get('test_horizon')
    if horizon is not None:
        selected &= event.le(horizon)
    keep = ~selected.to_numpy()
    q = int(selected.sum())
    if q == 0:
        raise ValueError('No post-treatment cohort deviations to test.')
    X, Y = np.asarray(model._X), np.asarray(model._Y).ravel()
    # Stable least squares on the imported, fixed-effect-residualized design.
    b, _, rank_u, _ = np.linalg.lstsq(X, Y, rcond=None)
    br, _, rank_r, _ = np.linalg.lstsq(X[:, keep], Y, rcond=None)
    u = Y - X @ b
    ur = Y - X[:, keep] @ br
    if rank_u != X.shape[1] or rank_u-rank_r != q:
        raise ValueError('The requested deviations are not independently identified.')
    units = pd.Categorical(data[unit]).codes
    times = pd.Categorical(data[time]).codes
    ni, nt = len(np.unique(units)), len(np.unique(times))
    graph = coo_matrix((np.ones(len(data)), (units, ni+times)), shape=(ni+nt, ni+nt))
    components = connected_components(graph, directed=False, return_labels=False)
    fe_rank = ni + nt - components
    df2 = int(len(data) - rank_u - fe_rank)
    rss_u, rss_r = float(u @ u), float(ur @ ur)
    mode = config.get('test_mode', 'classical')
    if mode == 'classical':
        F = ((rss_r-rss_u)/q) / (rss_u/df2)
        inference = 'Homoskedastic nested-model F'
        covariance_rank = q
    elif mode == 'clustered':
        # Sandwich scores from the stable least-squares solution.
        cluster_codes = pd.Categorical(data[config['cluster']]).codes
        G = len(np.unique(cluster_codes))
        scores = np.zeros((G, X.shape[1]))
        np.add.at(scores, cluster_codes, X*u[:,None])
        bread = np.linalg.inv(X.T @ X)
        influence = scores @ bread[:, selected.to_numpy()]
        vcov = influence.T @ influence * G/(G-1) * (len(data)-1)/df2
        covariance_rank = int(np.linalg.matrix_rank(vcov))
        if covariance_rank < q:
            raise ValueError(f'Clustered restriction covariance rank {covariance_rank} < {q}; choose a shorter test_horizon.')
        delta = b[selected.to_numpy()]
        F = float(delta @ np.linalg.solve(vcov, delta) / q)
        df2 = G-1
        inference = 'Cluster-robust F: ' + config['cluster']
    else:
        raise ValueError('test_mode must be classical or clustered')
    pvalue = float(stats.f.sf(F,q,df2))
    result = dict(F=float(F),df_num=q,df_den=df2,pvalue=pvalue,
                  log10_pvalue=float(stats.f.logsf(F,q,df2)/np.log(10)),
                  inference=inference,observations=len(data),design_columns=int(rank_u),
                  restricted_design_columns=int(rank_r),fixed_effect_rank=int(fe_rank),
                  covariance_rank=covariance_rank,rss_unrestricted=rss_u,rss_restricted=rss_r,
                  test_horizon=horizon,reference_cohort=int(data.loc[data[config['treatment']].eq(1)].groupby(unit)[time].min().min()),
                  pyfixest=version('pyfixest'),formulaic=version('formulaic'),
                  design_head=pd.DataFrame(X[:5,:6],columns=names.iloc[:6]).to_dict('split'),
                  formula=model._fml,restriction_names=names.loc[selected].tolist())
    Path(config['output']).write_text(json.dumps(result,indent=2)+'\n')
    print(json.dumps({k:v for k,v in result.items() if k not in ['design_head','formula','restriction_names']},indent=2))

if __name__ == '__main__':
    run(json.loads(sys.argv[1]))