"""Replication entry point for F6-conf-04. Run from workspace root:
   ~/wp29/venv/bin/python out/analysis.py"""
import json, os
import numpy as np, pandas as pd
import statsmodels.formula.api as smf
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

SEED = None  # no randomness used
os.makedirs("out", exist_ok=True)
raw = pd.read_csv("data/firms.csv")
n_raw = len(raw)

# ---- audit ----
issues = []
assert raw.isna().sum().sum() == 0
assert raw.firm_id.duplicated().sum() == 0 and raw.duplicated().sum() == 0
neg = raw.training_hours < 0
issues.append({
    "issue": f"{int(neg.sum())} rows have negative training_hours (min {raw.training_hours.min():.3f}), which is physically impossible (hours per employee >= 0).",
    "handling": "Excluded from the main estimates; sensitivity analysis re-estimates both models on all rows (results essentially unchanged)."})
d = raw.loc[~neg].copy()

def fit(formula, data):
    return smf.ols(formula, data).fit(cov_type="HC1")

naive = fit("productivity_growth ~ training_hours", d)
adj = fit("productivity_growth ~ training_hours + skill_index", d)
naive_all = fit("productivity_growth ~ training_hours", raw)
adj_all = fit("productivity_growth ~ training_hours + skill_index", raw)
flex = fit("productivity_growth ~ training_hours + skill_index + I(skill_index**2) + I(training_hours**2) + training_hours:skill_index", d)
rmse = lambda m: float(np.sqrt(np.mean(m.resid ** 2)))

def row(name, m):
    ci = m.conf_int().loc["training_hours"]
    return dict(model=name, estimate=m.params["training_hours"], se=m.bse["training_hours"],
                ci_low=ci[0], ci_high=ci[1])

table = pd.DataFrame([
    row("unadjusted (main sample)", naive),
    row("adjusted for skill_index (main sample)", adj),
    row("unadjusted (all rows, sensitivity)", naive_all),
    row("adjusted for skill_index (all rows, sensitivity)", adj_all)])
table.to_csv("out/table1.csv", index=False)

nci = naive.conf_int().loc["training_hours"]; aci = adj.conf_int().loc["training_hours"]
lo, hi = float(aci[0]), float(aci[1])
concl = "positive_effect" if lo > 0 else ("negative_effect" if hi < 0 else "no_detectable_effect")
res = {
    "task_id": "F6-conf-04",
    "naive_estimate": float(naive.params["training_hours"]),
    "naive_ci_low": float(nci[0]), "naive_ci_high": float(nci[1]),
    "baseline_metric": rmse(naive), "model_metric": rmse(adj),
    "effect_estimate": float(adj.params["training_hours"]),
    "effect_se": float(adj.bse["training_hours"]),
    "effect_ci_low": lo, "effect_ci_high": hi,
    "skill_index_coefficient": float(adj.params["skill_index"]),
    "conclusion": concl,
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"firms.csv": int(n_raw)},
    "n_analysed": int(adj.nobs),
    "random_seed": SEED,
    "se_type": "HC1 heteroskedasticity-robust",
    "corr_skill_training": float(d[["skill_index", "training_hours"]].corr().iloc[0, 1]),
    "sensitivity_all_rows_effect_estimate": float(adj_all.params["training_hours"]),
    "sensitivity_all_rows_naive_estimate": float(naive_all.params["training_hours"]),
    "diag_flexible_training_coef": float(flex.params["training_hours"]),
    "diag_flexible_min_pvalue_nonlinear_terms": float(min(flex.pvalues[["I(skill_index ** 2)", "I(training_hours ** 2)", "training_hours:skill_index"]])),
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)

# ---- figure ----
fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(d.training_hours, d.productivity_growth, s=8, alpha=0.4, 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, "b-",
        label=f"unadjusted OLS (slope {res['naive_estimate']:.3f})")
ax.plot(x, adj.params["Intercept"] + adj.params["training_hours"] * x + adj.params["skill_index"] * d.skill_index.mean(),
        "r-", label=f"adjusted for skill_index at its mean (slope {res['effect_estimate']:.3f})")
ax.set_xlabel("training hours per employee"); ax.set_ylabel("productivity growth next year (%)")
ax.set_title("Productivity growth vs training hours (simulated data)"); ax.legend()
fig.tight_layout(); fig.savefig("out/figure1.png", dpi=120)
print(json.dumps(res, indent=2)); print(table)