"""F1-conf-03: information-respecting LPM of next-year default on leverage.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/firm_panel.csv; writes out/results.json, out/analytic.csv, out/exclusions.csv."""
import json, os
import numpy as np
import pandas as pd
import statsmodels.api as sm

IN = "data/firm_panel.csv"
OUT = "out"
os.makedirs(OUT, exist_ok=True)

raw = pd.read_csv(IN, dtype={"firm_id": str})
n_raw = len(raw)
ledger, issues = [], []

def step(name, rule, before, after, why):
    ledger.append(dict(step=name, rule=rule, n_before=before, n_after=after,
                       n_removed=before - after, rationale=why))

# Step 1: column exclusion (information rule) -- no rows removed
excluded = [{"variable": "rating_2025",
             "reason": "Final-vintage field assigned in the 2025 review using realised default history "
                       "through 2022 and average sample leverage; known only in 2025, not at end of origin "
                       "year t. Including it would leak the outcome (look-ahead)."}]
df = raw.drop(columns=["rating_2025"])
step("1_drop_rating_2025", "drop column rating_2025 (known 2025 > end of year t)", n_raw, len(df),
     "Information-availability rule; column removal, no rows removed")
issues.append({"issue": "Column rating_2025 is a final-vintage (2025) field built from realised defaults "
                        "and average leverage: outcome leakage / look-ahead",
               "handling": "Excluded from analytic dataset and model; listed in variables_excluded"})

# Step 2: exact duplicate rows (assembly error)
n_dup_exact = int(df.duplicated().sum())
key_dup_before = int(raw.duplicated(["firm_id", "year"]).sum())
b = len(df)
df = df.drop_duplicates().reset_index(drop=True)
step("2_drop_exact_duplicates", "drop exact duplicate rows (all columns identical), keep first", b, len(df),
     "Assembly error: same firm-year repeated with identical values; one intended row per firm-year")
issues.append({"issue": f"{key_dup_before} duplicated (firm_id, year) rows in the extract; all are exact "
                        f"duplicates of another row ({n_dup_exact} exact duplicates)",
               "handling": "Kept one copy of each firm-year (drop_duplicates)"})
assert not df.duplicated(["firm_id", "year"]).any(), "conflicting firm-year duplicates remain"

# Step 3: unobserved outcome
b = len(df)
n_na_by_year = df[df.default_next.isna()].year.value_counts().to_dict()
df = df[df.default_next.notna()].copy()
step("3_drop_missing_outcome", "drop rows with default_next missing (not imputed)", b, len(df),
     "Outcome year not observed at data cut (end 2022); missing by year: "
     + ", ".join(f"{k}:{v}" for k, v in sorted(n_na_by_year.items())))
issues.append({"issue": "default_next missing for "
                        + ", ".join(f"{k}: {v} rows" for k, v in sorted(n_na_by_year.items()))
                        + " (outcome year 2023 beyond data cut)",
               "handling": "Excluded from estimation sample; outcomes not imputed"})

# Step 4: missing predictors (check)
preds = ["leverage", "cash_ratio", "log_sales"]
b = len(df)
df = df.dropna(subset=preds)
step("4_drop_missing_predictors", "drop rows with any missing predictor", b, len(df),
     "Complete-case requirement for OLS (none expected)")
df["default_next"] = df["default_next"].astype(int)
df = df.sort_values(["firm_id", "year"]).reset_index(drop=True)

# Estimation: OLS LPM, firm-clustered SE
X = sm.add_constant(df[preds])
fit = sm.OLS(df["default_next"], X).fit(cov_type="cluster", cov_kwds={"groups": df["firm_id"]})

df[["firm_id", "year"] + preds + ["default_next"]].to_csv(f"{OUT}/analytic.csv", index=False)
pd.DataFrame(ledger).to_csv(f"{OUT}/exclusions.csv", index=False)

res = {
    "task_id": "F1-conf-03",
    "n_analysed": int(len(df)),
    "n_firms": int(df.firm_id.nunique()),
    "predictors_used": preds,
    "variables_excluded": excluded,
    "beta_leverage": float(fit.params["leverage"]),
    "se_leverage_cluster": float(fit.bse["leverage"]),
    "r2": float(fit.rsquared),
    "default_rate": float(df.default_next.mean()),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"firm_panel.csv": int(n_raw)},
    "random_seed": None,
    # supplementary (not required by TASK.md)
    "supp_ci95_leverage_cluster": [float(x) for x in fit.conf_int().loc["leverage"]],
    "supp_se_leverage_hc1": float(sm.OLS(df["default_next"], X).fit(cov_type="HC1").bse["leverage"]),
    "supp_beta_cash_ratio": float(fit.params["cash_ratio"]),
    "supp_beta_log_sales": float(fit.params["log_sales"]),
    "supp_intercept": float(fit.params["const"]),
    "supp_year_range": [int(df.year.min()), int(df.year.max())],
}
with open(f"{OUT}/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps({k: res[k] for k in ["n_analysed", "n_firms", "beta_leverage",
                                       "se_leverage_cluster", "r2", "default_rate"]}, indent=2))