"""Replication script for F6-conf-04. Run from workspace root:
    ~/wp29/venv/bin/python out/analysis.py
"""
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

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

issues = []
assert raw.isna().sum().sum() == 0
assert not raw.firm_id.duplicated().any()
neg = raw.training_hours < 0
n_neg = int(neg.sum())
if n_neg:
    issues.append({"issue": f"{n_neg} rows have negative training_hours (min {raw.training_hours.min():.3f}); "
                            "hours per employee cannot be negative, so these values are invalid",
                   "handling": "excluded from all main estimates; sensitivity analysis re-estimates with them included "
                               "and with them set to 0 (reported in report/table1)"})
d = raw[~neg].copy()

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

def rmse(m):
    return float(np.sqrt(np.mean(m.resid ** 2)))

naive = fit("productivity_growth ~ training_hours", d)
adj = fit("productivity_growth ~ training_hours + skill_index", d)
# sensitivity
adj_all = fit("productivity_growth ~ training_hours + skill_index", raw)
rc = raw.copy(); rc["training_hours"] = rc.training_hours.clip(lower=0)
adj_clip = fit("productivity_growth ~ training_hours + skill_index", rc)
adj_conv = d.pipe(lambda x: smf.ols("productivity_growth ~ training_hours + skill_index", x).fit())
adj_quad = fit("productivity_growth ~ training_hours + skill_index + I(skill_index**2)", d)

def row(name, m):
    ci = m.conf_int().loc["training_hours"]
    return {"model": name, "estimate": m.params["training_hours"], "se": m.bse["training_hours"],
            "ci_low": ci[0], "ci_high": ci[1], "n": int(m.nobs), "rmse": rmse(m)}

rows = [row("unadjusted (HC1)", naive), row("adjusted for skill_index (HC1) [main]", adj),
        row("adjusted, conventional SE", adj_conv),
        row("adjusted, incl. negative-hours rows (HC1)", adj_all),
        row("adjusted, negative hours set to 0 (HC1)", adj_clip),
        row("adjusted + skill_index^2 (HC1)", adj_quad)]
pd.DataFrame(rows).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_se_type": "HC1 heteroskedasticity-robust",
    "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": n_raw},
    "n_analysed": int(adj.nobs),
    "random_seed": SEED,
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)

fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(d.training_hours, d.productivity_growth, s=8, alpha=0.4, color="grey", label="firms (analysed)")
x = np.linspace(d.training_hours.min(), d.training_hours.max(), 100)
ax.plot(x, naive.params["Intercept"] + naive.params["training_hours"] * x, "r-",
        label=f"unadjusted OLS (slope {naive.params['training_hours']:.3f})")
sm = d.skill_index.mean()
ax.plot(x, adj.params["Intercept"] + adj.params["training_hours"] * x + adj.params["skill_index"] * sm, "b-",
        label=f"adjusted for skill_index at mean (slope {adj.params['training_hours']:.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(fontsize=8); fig.tight_layout(); fig.savefig("out/figure1.png", dpi=120)
print(json.dumps(res, indent=2)); print(pd.DataFrame(rows).to_string())