"""Replication entry point for F6-conf-04. Run from workspace root:
   ~/wp29/venv/bin/python out/analysis.py
Reads only data/firms.csv; writes out/results.json, out/table1.csv, out/figure1.png, out/exclusions.csv.
No randomness is used (random_seed = null)."""
import json, os
import numpy as np
import pandas as pd
import statsmodels.formula.api as smf
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

OUT = "out"
os.makedirs(OUT, exist_ok=True)
raw = pd.read_csv("data/firms.csv")
n_raw = len(raw)

# --- exclusions ---
excl = []
d = raw.dropna()
excl.append(("1", "drop rows with missing values", n_raw, len(d), "none expected; complete-case"))
n0 = len(d)
d = d[d["training_hours"] >= 0].copy()
excl.append(("2", "drop training_hours < 0", n0, len(d), "negative hours of training are physically impossible"))
pd.DataFrame([dict(step=s, rule=r, n_before=a, n_after=b, n_removed=a - b, rationale=why)
              for s, r, a, b, why in excl]).to_csv(f"{OUT}/exclusions.csv", index=False)
n_neg = int((raw["training_hours"] < 0).sum())

COV = "HC1"  # heteroskedasticity-robust SEs

def fit(formula, data):
    m = smf.ols(formula, data).fit(cov_type=COV)
    ci = m.conf_int(0.05)
    rmse = float(np.sqrt(np.mean(m.resid ** 2)))
    return m, ci, rmse

naive, naive_ci, naive_rmse = fit("productivity_growth ~ training_hours", d)
adj, adj_ci, adj_rmse = fit("productivity_growth ~ training_hours + skill_index", d)
# sensitivity analyses (not the main estimate)
naive_all, _, _ = fit("productivity_growth ~ training_hours", raw)
sens_all, sens_all_ci, _ = fit("productivity_growth ~ training_hours + skill_index", raw)
sens_int, sens_int_ci, _ = fit("productivity_growth ~ training_hours * skill_index", d)
sens_q, sens_q_ci, _ = fit("productivity_growth ~ training_hours + skill_index + I(skill_index**2)", d)
adj_conv = smf.ols("productivity_growth ~ training_hours + skill_index", d).fit()

b = float(adj.params["training_hours"]); lo, hi = map(float, adj_ci.loc["training_hours"])
conclusion = "positive_effect" if lo > 0 else ("negative_effect" if hi < 0 else "no_detectable_effect")

rows = []
def addrow(name, m, ci):
    rows.append(dict(model=name, estimate=float(m.params["training_hours"]), se=float(m.bse["training_hours"]),
                     ci_low=float(ci.loc["training_hours", 0]), ci_high=float(ci.loc["training_hours", 1])))
addrow("unadjusted (growth ~ hours)", naive, naive_ci)
addrow("adjusted (growth ~ hours + skill) [main]", adj, adj_ci)
addrow("sensitivity: adjusted, incl. negative-hours rows", sens_all, sens_all_ci)
addrow("sensitivity: adjusted + hours x skill (coef at skill=0)", sens_int, sens_int_ci)
addrow("sensitivity: adjusted + skill^2", sens_q, sens_q_ci)
pd.DataFrame(rows).to_csv(f"{OUT}/table1.csv", index=False)

results = {
    "task_id": "F6-conf-04",
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": [
        {"issue": f"{n_neg} rows have negative training_hours (physically impossible; all low skill_index)",
         "handling": "excluded from main estimate; sensitivity estimate including them reported in table1.csv"},
    ],
    "n_input_rows": {"firms.csv": int(n_raw)},
    "n_analysed": int(len(d)),
    "random_seed": None,
    "naive_estimate": float(naive.params["training_hours"]),
    "naive_ci_low": float(naive_ci.loc["training_hours", 0]),
    "naive_ci_high": float(naive_ci.loc["training_hours", 1]),
    "baseline_metric": naive_rmse,
    "model_metric": adj_rmse,
    "effect_estimate": b,
    "effect_se": float(adj.bse["training_hours"]),
    "effect_ci_low": lo,
    "effect_ci_high": hi,
    "skill_index_coefficient": float(adj.params["skill_index"]),
    "conclusion": conclusion,
    "se_type": "HC1 heteroskedasticity-robust",
    "effect_se_conventional": float(adj_conv.bse["training_hours"]),
    "naive_incl_negative_hours_estimate": float(naive_all.params["training_hours"]),
    "sensitivity_incl_negative_hours_estimate": float(sens_all.params["training_hours"]),
    "sensitivity_interaction_term": float(sens_int.params["training_hours:skill_index"]),
    "sensitivity_interaction_term_p": float(sens_int.pvalues["training_hours:skill_index"]),
    "sensitivity_skill_sq_estimate": float(sens_q.params["training_hours"]),
    "corr_skill_training": float(d["skill_index"].corr(d["training_hours"])),
    "adjusted_r2": float(adj.rsquared),
    "naive_r2": float(naive.rsquared),
}
with open(f"{OUT}/results.json", "w") as f:
    json.dump(results, f, indent=2)

# --- figure ---
fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(d["training_hours"], d["productivity_growth"], s=8, alpha=0.35, color="grey", label="firms")
x = np.linspace(d["training_hours"].min(), d["training_hours"].max(), 100)
ax.plot(x, naive.params["Intercept"] + naive.params["training_hours"] * x, color="tab:red",
        label=f"unadjusted: slope {results['naive_estimate']:.3f}")
sm = d["skill_index"].mean()
ax.plot(x, adj.params["Intercept"] + b * x + adj.params["skill_index"] * sm, color="tab:blue",
        label=f"adjusted for skill (at mean skill): slope {b:.3f}")
ax.set_xlabel("training hours per employee (programme year)")
ax.set_ylabel("productivity growth, following year (%)")
ax.set_title("Productivity growth vs training hours (simulated data)")
ax.legend()
fig.tight_layout()
fig.savefig(f"{OUT}/figure1.png", dpi=120)

print(json.dumps({k: results[k] for k in ["naive_estimate", "effect_estimate", "effect_se",
                                          "effect_ci_low", "effect_ci_high", "conclusion", "n_analysed"]}, indent=1))