#!/usr/bin/env python3
"""
Benchmark: MultiPeriodDiD event study (diff-diff MultiPeriodDiD).

Usage:
    python benchmark_multiperiod.py --data path/to/data.csv --output path/to/results.json \
        --n-pre 4 --n-post 4
"""

import argparse
import json
import os
import sys
from pathlib import Path

# IMPORTANT: Parse --backend and set environment variable BEFORE importing diff_diff
# This ensures the backend configuration is respected by all modules
def _get_backend_from_args():
    """Parse --backend argument without importing diff_diff."""
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--backend", default="auto", choices=["auto", "python", "rust"])
    args, _ = parser.parse_known_args()
    return args.backend

_requested_backend = _get_backend_from_args()
if _requested_backend in ("python", "rust"):
    os.environ["DIFF_DIFF_BACKEND"] = _requested_backend

# NOW import diff_diff and other dependencies (will see the env var)
import pandas as pd

# Add parent to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from diff_diff import MultiPeriodDiD, HAS_RUST_BACKEND
from benchmarks.python.utils import Timer


def parse_args():
    parser = argparse.ArgumentParser(description="Benchmark MultiPeriodDiD estimator")
    parser.add_argument("--data", required=True, help="Path to input CSV data")
    parser.add_argument("--output", required=True, help="Path to output JSON results")
    parser.add_argument(
        "--cluster", default="unit", help="Column to cluster standard errors on"
    )
    parser.add_argument(
        "--n-pre", type=int, required=True, help="Number of pre-treatment periods"
    )
    parser.add_argument(
        "--n-post", type=int, required=True, help="Number of post-treatment periods"
    )
    parser.add_argument(
        "--reference-period", type=int, default=None,
        help="Reference period (default: last pre-period = n_pre)"
    )
    parser.add_argument(
        "--backend", default="auto", choices=["auto", "python", "rust"],
        help="Backend to use: auto (default), python (pure Python), rust (Rust backend)"
    )
    return parser.parse_args()


def get_actual_backend() -> str:
    """Return the actual backend being used based on HAS_RUST_BACKEND."""
    return "rust" if HAS_RUST_BACKEND else "python"


def main():
    args = parse_args()

    # Get actual backend (already configured via env var before imports)
    actual_backend = get_actual_backend()
    print(f"Using backend: {actual_backend}")

    # Load data
    print(f"Loading data from: {args.data}")
    data = pd.read_csv(args.data)

    # Compute post_periods and reference_period from args
    all_periods = sorted(data["time"].unique())
    n_pre = args.n_pre
    post_periods = [p for p in all_periods if p > n_pre]
    ref_period = args.reference_period if args.reference_period is not None else n_pre

    print(f"All periods: {all_periods}")
    print(f"Post periods: {post_periods}")
    print(f"Reference period: {ref_period}")

    # Run benchmark
    print("Running MultiPeriodDiD estimation...")

    did = MultiPeriodDiD(robust=True, cluster=args.cluster)

    with Timer() as timer:
        results = did.fit(
            data,
            outcome="outcome",
            treatment="treated",
            time="time",
            post_periods=post_periods,
            reference_period=ref_period,
            absorb=["unit"],
        )

    total_time = timer.elapsed

    # Extract period effects (excluding reference period)
    period_effects = []
    for period, pe in sorted(results.period_effects.items()):
        event_time = period - ref_period
        period_effects.append({
            "period": int(period),
            "event_time": int(event_time),
            "att": float(pe.effect),
            "se": float(pe.se),
        })

    # Build output
    output = {
        "estimator": "diff_diff.MultiPeriodDiD",
        "backend": actual_backend,
        "cluster": args.cluster,
        # Average treatment effect (across post-periods)
        "att": float(results.avg_att),
        "se": float(results.avg_se),
        "pvalue": float(results.avg_p_value),
        "ci_lower": float(results.avg_conf_int[0]),
        "ci_upper": float(results.avg_conf_int[1]),
        # Reference period
        "reference_period": int(ref_period),
        # Period-level effects
        "period_effects": period_effects,
        # Timing
        "timing": {
            "estimation_seconds": total_time,
            "total_seconds": total_time,
        },
        # Metadata
        "metadata": {
            "n_units": int(data["unit"].nunique()),
            "n_periods": int(data["time"].nunique()),
            "n_obs": len(data),
            "n_pre": n_pre,
            "n_post": len(post_periods),
        },
    }

    # Write output
    print(f"Writing results to: {args.output}")
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w") as f:
        json.dump(output, f, indent=2)

    print(f"ATT: {results.avg_att:.6f}")
    print(f"SE:  {results.avg_se:.6f}")
    print(f"Completed in {total_time:.3f} seconds")
    return output


if __name__ == "__main__":
    main()
