import pandas as pd
import numpy as np
from pathlib import Path

ROOT = Path(__file__).parent.parent.parent

crashes_file = ROOT / "build/output/crashes_state_year.csv"
demographics_file = ROOT / "build/output/demographics_state_year.csv"
policy_file = ROOT / "build/input/policy_adoptions.csv"
names_file = ROOT / "build/input/state_names.csv"
output_file = ROOT / "build/output/analysis_panel.csv"

# Load all datasets
crashes = pd.read_csv(crashes_file)
demographics = pd.read_csv(demographics_file)
policies = pd.read_csv(policy_file)
state_names = pd.read_csv(names_file)

# Merge 1: Crashes + Demographics (on state_fips + year)
panel = crashes.merge(demographics, on=['state_fips', 'year'], how='inner')
print(f"After merge 1: {len(panel)} rows")

# Merge 2: + Policy adoptions (on state_fips only — many:1)
panel = panel.merge(policies, on='state_fips', how='left')
print(f"After merge 2: {len(panel)} rows")

# Merge 3: + State names (on state_fips only — many:1)
panel = panel.merge(state_names, on='state_fips', how='left')
print(f"After merge 3: {len(panel)} rows")

# Create treatment indicator
panel['ever_treated'] = panel['adoption_year'].notna().astype(int)
panel['post_treated'] = (
    (panel['year'] >= panel['adoption_year']) &
    panel['adoption_year'].notna()
).astype(int)
panel['log_pop'] = np.log(panel['population'])

output_file.parent.mkdir(parents=True, exist_ok=True)
panel.to_csv(output_file, index=False)
print(f"Saved analysis_panel.csv: {len(panel)} rows")
