"""
Analysis Script 02: Difference-in-Differences Regression

Estimates DD models using pyfixest:
1. Main DD (fatal crashes)
2. With controls
3. South only
4. Non-South
5. Alternative outcome (serious crashes)

Output: analysis/output/tables/dd_results.tex
"""

import pandas as pd
from pathlib import Path

try:
    import pyfixest as pf
except ImportError:
    print("ERROR: pyfixest not installed. Install with: pip install pyfixest")
    raise

# Set project root
ROOT = Path(__file__).parent.parent.parent

# Define paths
input_file = ROOT / "build/output/analysis_panel.csv"
output_file = ROOT / "analysis/output/tables/dd_results.tex"

# Read analysis panel
panel = pd.read_csv(input_file)

# Col 1: Main result
m1 = pf.feols("fatal_crashes ~ post_treated | state_fips + year",
              vcov={"CRV1": "state_fips"}, data=panel)

# Col 2: With controls
m2 = pf.feols("fatal_crashes ~ post_treated + log_pop + median_income"
              " | state_fips + year",
              vcov={"CRV1": "state_fips"}, data=panel)

# Col 3: South only
m3 = pf.feols("fatal_crashes ~ post_treated + log_pop | state_fips + year",
              vcov={"CRV1": "state_fips"},
              data=panel[panel['region'] == 'South'])

# Col 4: Non-South
m4 = pf.feols("fatal_crashes ~ post_treated + log_pop | state_fips + year",
              vcov={"CRV1": "state_fips"},
              data=panel[panel['region'] != 'South'])

# Col 5: Serious crashes as outcome
m5 = pf.feols("serious_crashes ~ post_treated + log_pop | state_fips + year",
              vcov={"CRV1": "state_fips"}, data=panel)

# Export table
output_file.parent.mkdir(parents=True, exist_ok=True)
pf.etable(
    [m1, m2, m3, m4, m5],
    type="tex",
    file_name=str(output_file),
    labels={
        "fatal_crashes": "Fatal Crashes",
        "serious_crashes": "Serious Crashes",
        "post_treated": "Treat x Post",
        "log_pop": "Log(Population)",
        "median_income": "Median Income",
    },
)

print(f"Saved dd_results.tex")
