"""Replication script 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 matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import statsmodels.formula.api as smf

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

# ---- audit ----
assert raw.isna().sum().sum() == 0
assert not raw.firm_id.duplicated().any()
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 cannot be < 0); likely an untruncated simulation draw or recording error.",
    "handling": "Excluded from the main analysis; sensitivity analysis re-estimates with them included (results reported in table1.csv and report)."})
d = raw[~neg].copy()

# ---- models (HC1 robust SEs) ----
naive = smf.ols("productivity_growth ~ training_hours", d).fit(cov_type="HC1")
adj = smf.ols("productivity_growth ~ training_hours + skill_index", d).fit(cov_type="HC1")
adj_all = smf.ols("productivity_growth ~ training_hours + skill_index", raw).fit(cov_type="HC1")
naive_all = smf.ols("productivity_growth ~ training_hours", raw).fit(cov_type="HC1")
# diagnostic: nonlinearity / heterogeneity checks
quad = smf.ols("productivity_growth ~ training_hours + I(training_hours**2) + skill_index + I(skill_index**2)", d).fit(cov_type="HC1")
inter = smf.ols("productivity_growth ~ training_hours * skill_index", d).fit(cov_type="HC1")

rmse = lambda m: float(np.sqrt(np.mean(m.resid ** 2)))
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]}
table = pd.DataFrame([
    row("unadjusted (main sample)", naive),
    row("adjusted for skill_index (main sample) [targets estimand]", adj),
    row("unadjusted (all rows incl. negative hours)", naive_all),
    row("adjusted for skill_index (all rows incl. negative hours)", adj_all),
    row("adjusted + quadratics (diagnostic)", quad),
])
table.to_csv("out/table1.csv", index=False)

eci = adj.conf_int().loc["training_hours"]; nci = naive.conf_int().loc["training_hours"]
if eci[0] > 0: concl = "positive_effect"
elif eci[1] < 0: concl = "negative_effect"
else: concl = "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": float(eci[0]), "effect_ci_high": float(eci[1]),
    "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": None,
    "diagnostics": {
        "se_type": "HC1 heteroskedasticity-robust",
        "corr_training_skill": float(d[["training_hours", "skill_index"]].corr().iloc[0, 1]),
        "sensitivity_all_rows_effect": float(adj_all.params["training_hours"]),
        "sensitivity_all_rows_se": float(adj_all.bse["training_hours"]),
        "quad_training_sq_coef": float(quad.params["I(training_hours ** 2)"]),
        "quad_training_sq_p": float(quad.pvalues["I(training_hours ** 2)"]),
        "quad_skill_sq_p": float(quad.pvalues["I(skill_index ** 2)"]),
        "interaction_coef": float(inter.params["training_hours:skill_index"]),
        "interaction_p": float(inter.pvalues["training_hours:skill_index"]),
        "conventional_se_effect": float(smf.ols("productivity_growth ~ training_hours + skill_index", d).fit().bse["training_hours"]),
        "r2_naive": float(naive.rsquared), "r2_adjusted": float(adj.rsquared),
    },
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)

# ---- figure ----
fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(raw.training_hours[~neg], raw.productivity_growth[~neg], s=8, alpha=0.4, label="firms (analysed)")
ax.scatter(raw.training_hours[neg], raw.productivity_growth[neg], s=14, marker="x", c="k", label="excluded (negative hours)")
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: 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, "g-",
        label=f"adjusted for skill (at mean skill): slope {adj.params['training_hours']:.3f}")
ax.set_xlabel("training hours per employee"); ax.set_ylabel("productivity growth (%)")
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))