"""
Cohort-Anchored Robust Event Study (Liu 2025, arXiv 2509.01829). [@cohort_anchored2025]

Standard event-study TWFE with leads/lags is fragile under staggered
adoption because untreated leads/lags from later cohorts contaminate
earlier cohorts' coefficients. The cohort-anchored estimator computes
event-time effects *separately within each treatment cohort* (the
"anchor"), then averages with cohort weights, producing event-study
coefficients that are robust to staggered timing without requiring
the user to choose between CS / SA / BJS.
"""

from __future__ import annotations

from typing import List, Optional

import numpy as np
from ..core._bootstrap import bootstrap_se as _bootstrap_se
import pandas as pd
from scipy import stats

from ..core.results import CausalResult


def cohort_anchored_event_study(
    data: pd.DataFrame,
    y: str,
    treat: str,
    time: str,
    id: str,
    leads: int = 4,
    lags: int = 4,
    cluster: Optional[str] = None,
    alpha: float = 0.05,
) -> CausalResult:
    """
    Cohort-anchored event-study estimator (Rambachan-Roth successor).

    Parameters
    ----------
    data : pd.DataFrame
        Long-format panel.
    y : str
    treat : str
        First-treatment-period column (0 = never-treated).
    time : str
    id : str
    leads, lags : int
        Number of pre/post event-time periods to estimate.
    cluster : str, optional
        Cluster column for SE; defaults to ``id``.
    alpha : float

    Returns
    -------
    CausalResult
        ``estimate`` is the average post-treatment effect across event
        times 0..lags. Per-event-time coefficients in
        ``model_info['event_study']`` (DataFrame with columns
        ``rel_time``, ``att``, ``se``, ``ci_low``, ``ci_high``).

    References
    ----------
    arXiv 2509.01829, *Cohort-Anchored Robust Inference for
    Event-Study with Staggered Adoption* (2025).
    """
    df = data[[y, treat, time, id] + ([cluster] if cluster else [])] \
        .dropna().reset_index(drop=True)
    cluster_col = cluster or id

    treat_arr = df[treat].to_numpy()
    treated_mask = treat_arr > 0
    cohorts = sorted(df.loc[treated_mask, treat].unique())
    if not cohorts:
        raise ValueError("No treated cohorts found (all treat == 0).")

    # Per-cohort, per-event-time ATT
    rel_times = list(range(-leads, lags + 1))
    rows = []
    cohort_weights = []
    for c in cohorts:
        cohort_units = df.loc[df[treat] == c, id].unique()
        # Control = never-treated units only (clean comparison)
        control_units = df.loc[df[treat] == 0, id].unique()
        if len(control_units) == 0:
            continue
        cohort_df = df[df[id].isin(np.concatenate([cohort_units, control_units]))]
        for k in rel_times:
            t_target = c + k
            sub = cohort_df[cohort_df[time] == t_target]
            if len(sub) < 4:
                continue
            sub = sub.copy()
            sub['_treated'] = sub[id].isin(cohort_units).astype(int)
            # Reference period: t = c - 1
            ref = cohort_df[cohort_df[time] == c - 1]
            if len(ref) < 4:
                continue
            ref = ref.copy()
            ref['_treated'] = ref[id].isin(cohort_units).astype(int)
            # ATT(c, k) = (Y_c,k - Y_0,k) - (Y_c,c-1 - Y_0,c-1)
            try:
                m = sub.groupby('_treated')[y].mean()
                m_ref = ref.groupby('_treated')[y].mean()
                att = float(
                    (m.get(1, np.nan) - m.get(0, np.nan))
                    - (m_ref.get(1, np.nan) - m_ref.get(0, np.nan))
                )
                rows.append({
                    'cohort': c, 'rel_time': k, 'att': att,
                    'n': int(len(sub)),
                })
            except Exception:
                continue
        cohort_weights.append((c, int((df[treat] == c).sum())))

    if not rows:
        raise ValueError("No cohort-time cells could be estimated.")
    cohort_atts = pd.DataFrame(rows)

    # Aggregate per event-time across cohorts (weighted by cohort size)
    cw = pd.DataFrame(cohort_weights, columns=['cohort', 'cw'])
    cohort_atts = cohort_atts.merge(cw, on='cohort', how='left')
    es_rows = []
    for k in rel_times:
        sub = cohort_atts[cohort_atts['rel_time'] == k]
        if sub.empty or not np.isfinite(sub['att']).any():
            continue
        finite = sub['att'].dropna()
        weights = sub.loc[finite.index, 'cw']
        att_k = float(np.average(finite, weights=weights))
        # Cluster-bootstrap SE: resample clusters (user-specified
        # ``cluster`` column, or the unit id by default) and recompute
        # the ATT(c, k) from scratch for each draw. This honors a
        # user-supplied cluster level (e.g. state, region, firm)
        # instead of silently collapsing to cohort-level resampling.
        rng = np.random.default_rng(0)
        cluster_ids = df[cluster_col].unique()
        boot = np.full(200, np.nan)
        for b in range(200):
            sampled = rng.choice(cluster_ids, size=len(cluster_ids), replace=True)
            # Resample full observations belonging to sampled clusters
            pieces = [df[df[cluster_col] == cid] for cid in sampled]
            if not pieces:
                continue
            df_b = pd.concat(pieces, ignore_index=True)
            try:
                att_vals_b = []
                w_vals_b = []
                for c_b in cohorts:
                    cohort_units_b = df_b.loc[df_b[treat] == c_b, id].unique()
                    control_units_b = df_b.loc[df_b[treat] == 0, id].unique()
                    if len(cohort_units_b) == 0 or len(control_units_b) == 0:
                        continue
                    sub_b = df_b[(df_b[time] == c_b + k)
                                  & df_b[id].isin(np.concatenate([cohort_units_b, control_units_b]))]
                    ref_b = df_b[(df_b[time] == c_b - 1)
                                  & df_b[id].isin(np.concatenate([cohort_units_b, control_units_b]))]
                    if len(sub_b) < 2 or len(ref_b) < 2:
                        continue
                    sub_b = sub_b.assign(_tr=sub_b[id].isin(cohort_units_b).astype(int))
                    ref_b = ref_b.assign(_tr=ref_b[id].isin(cohort_units_b).astype(int))
                    m_b = sub_b.groupby('_tr')[y].mean()
                    mr_b = ref_b.groupby('_tr')[y].mean()
                    att_b = (
                        (m_b.get(1, np.nan) - m_b.get(0, np.nan))
                        - (mr_b.get(1, np.nan) - mr_b.get(0, np.nan))
                    )
                    if np.isfinite(att_b):
                        att_vals_b.append(att_b)
                        w_vals_b.append(int((df_b[treat] == c_b).sum()))
                if att_vals_b:
                    boot[b] = float(np.average(att_vals_b, weights=w_vals_b))
            except Exception:
                pass
        se_k = _bootstrap_se(boot, label="did.cohort_anchored")
        z_crit = float(stats.norm.ppf(1 - alpha / 2))
        es_rows.append({
            'rel_time': k, 'att': att_k, 'se': se_k,
            'ci_low': att_k - z_crit * se_k,
            'ci_high': att_k + z_crit * se_k,
        })
    event_study_df = pd.DataFrame(es_rows)

    # Headline: simple average across post periods (rel_time >= 0)
    post = event_study_df[event_study_df['rel_time'] >= 0]
    if post.empty:
        att_avg = float(event_study_df['att'].mean())
        se_avg = float(event_study_df['se'].mean()) or 1e-6
    else:
        att_avg = float(post['att'].mean())
        # Conservative SE under independence approximation
        se_avg = float(np.sqrt((post['se'] ** 2).sum()) / len(post)) or 1e-6

    z_crit = float(stats.norm.ppf(1 - alpha / 2))
    ci = (att_avg - z_crit * se_avg, att_avg + z_crit * se_avg)
    z = att_avg / se_avg if se_avg > 0 else 0.0
    pvalue = float(2 * (1 - stats.norm.cdf(abs(z))))

    _result = CausalResult(
        method="Cohort-Anchored Event Study (staggered-robust)",
        estimand="ATT (avg post)",
        estimate=att_avg,
        se=se_avg,
        pvalue=pvalue,
        ci=ci,
        alpha=alpha,
        n_obs=len(df),
        model_info={
            'estimator': 'cohort_anchored_event_study',
            'event_study': event_study_df,
            'n_cohorts': len(cohorts),
            'reference': 'arXiv 2509.01829 (2025)',
        },
        _citation_key='cohort_anchored',
    )
    try:
        from ..output._lineage import attach_provenance as _attach_prov
        _attach_prov(
            _result,
            function="sp.did.cohort_anchored_event_study",
            params={
                "y": y, "treat": treat, "time": time, "id": id,
                "leads": leads, "lags": lags,
                "cluster": cluster, "alpha": alpha,
            },
            data=data,
            overwrite=False,
        )
    except Exception:  # pragma: no cover
        pass
    return _result


CausalResult._CITATIONS['cohort_anchored'] = (
    "@article{cohort_anchored2025,\n"
    "  title={Cohort-Anchored Robust Inference for Event-Study with "
    "Staggered Adoption},\n"
    "  author={Anonymous},\n"
    "  journal={arXiv preprint arXiv:2509.01829},\n"
    "  year={2025}\n"
    "}"
)
