"""Replication entry point for F6-conf-04. Run from workspace root:
   ~/wp29/venv/bin/python out/analysis.py
Reads data/firms.csv only; writes out/results.json, out/table1.csv, out/figure1.png,
out/exclusions.csv, out/sensitivity.csv. Deterministic (no randomness used)."""
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)
SEED = None  # no stochastic step
COV = "HC1"  # heteroskedasticity-robust standard errors

raw = pd.read_csv("data/firms.csv")
n_raw = len(raw)
req = ["firm_id", "skill_index", "training_hours", "productivity_growth"]
assert all(c in raw.columns for c in req)

excl = []
data_issues = []
d = raw.copy()

def step(name, rule, mask_keep, rationale):
    global d
    nb = len(d); d = d[mask_keep(d)].copy(); na = len(d)
    excl.append(dict(step=name, rule=rule, n_before=nb, n_after=na, n_removed=nb - na, rationale=rationale))

step("1", "drop rows with missing values in required columns",
     lambda x: x[req].notna().all(axis=1), "complete cases required for OLS")
step("2", "drop duplicate firm_id",
     lambda x: ~x.firm_id.duplicated(), "one row per firm expected")
n_neg = int((d.training_hours < 0).sum())
step("3", "drop rows with training_hours < 0",
     lambda x: x.training_hours >= 0,
     "negative hours of training per employee are physically impossible; treated as recording errors")
if n_neg:
    data_issues.append({"issue": f"{n_neg} rows have negative training_hours (min {raw.training_hours.min():.3f}); impossible values",
                        "handling": "excluded from main analysis (exclusion step 3); sensitivity analyses keep them as-is and censor them to 0 (out/sensitivity.csv)"})
pd.DataFrame(excl).to_csv(f"{OUT}/exclusions.csv", index=False)
n_an = len(d)

def fit(formula, df):
    return smf.ols(formula, data=df).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)

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

table = pd.DataFrame([row("unadjusted: growth ~ hours", naive, n_an),
                      row("adjusted: growth ~ hours + skill_index (target)", adj, n_an)])
table.to_csv(f"{OUT}/table1.csv", index=False)

# Sensitivity analyses (not the main estimate)
sens = []
full = raw.dropna(subset=req)
sens.append(row("adjusted, all rows incl. negative hours", fit("productivity_growth ~ training_hours + skill_index", full), len(full)))
sens.append(row("unadjusted, all rows incl. negative hours", fit("productivity_growth ~ training_hours", full), len(full)))
cens = full.assign(training_hours=full.training_hours.clip(lower=0))
sens.append(row("adjusted, negative hours censored to 0", fit("productivity_growth ~ training_hours + skill_index", cens), len(cens)))
sens.append(row("adjusted + skill_index^2", fit("productivity_growth ~ training_hours + skill_index + I(skill_index**2)", d), n_an))
sens.append(row("adjusted + skill^2 + skill^3", fit("productivity_growth ~ training_hours + skill_index + I(skill_index**2) + I(skill_index**3)", d), n_an))
sens.append(row("adjusted + hours^2 (linear term only reported)", fit("productivity_growth ~ training_hours + I(training_hours**2) + skill_index", d), n_an))
adj_conv = smf.ols("productivity_growth ~ training_hours + skill_index", data=d).fit()
sens.append(row("adjusted, conventional (non-robust) SE", adj_conv, n_an))
pd.DataFrame(sens).to_csv(f"{OUT}/sensitivity.csv", index=False)

eff = adj.params["training_hours"]; lo, hi = adj.conf_int().loc["training_hours"]
conclusion = "positive_effect" if lo > 0 else ("negative_effect" if hi < 0 else "no_detectable_effect")
nlo, nhi = naive.conf_int().loc["training_hours"]

results = {
    "task_id": "F6-conf-04",
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": data_issues,
    "n_input_rows": {"firms.csv": int(n_raw)},
    "n_analysed": int(n_an),
    "random_seed": SEED,
    "naive_estimate": float(naive.params["training_hours"]),
    "naive_ci_low": float(nlo), "naive_ci_high": float(nhi),
    "naive_se": float(naive.bse["training_hours"]),
    "baseline_metric": rmse(naive),
    "model_metric": rmse(adj),
    "effect_estimate": float(eff),
    "effect_se": float(adj.bse["training_hours"]),
    "effect_ci_low": float(lo), "effect_ci_high": float(hi),
    "skill_index_coefficient": float(adj.params["skill_index"]),
    "skill_index_se": float(adj.bse["skill_index"]),
    "se_type": "HC1 heteroskedasticity-robust; 95% CI = estimate +/- 1.96*SE (normal)",
    "corr_hours_skill": float(d.training_hours.corr(d.skill_index)),
    "sens_effect_all_rows": sens[0]["estimate"],
    "sens_effect_censored0": sens[2]["estimate"],
    "sens_effect_skill_sq": sens[3]["estimate"],
    "sens_effect_skill_cubic": sens[4]["estimate"],
    "sens_effect_conventional_se": sens[6]["se"],
    "conclusion": conclusion,
}
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.4, color="grey", label=f"firms (n={n_an})")
xs = np.linspace(d.training_hours.min(), d.training_hours.max(), 100)
ax.plot(xs, naive.params["Intercept"] + naive.params["training_hours"] * xs, "b-",
        label=f"unadjusted: slope {naive.params['training_hours']:.3f}")
sbar = d.skill_index.mean()
ax.plot(xs, adj.params["Intercept"] + adj.params["skill_index"] * sbar + eff * xs, "r-",
        label=f"adjusted for skill_index (at mean): slope {eff:.3f}")
ax.set_xlabel("training hours per employee (programme year)")
ax.set_ylabel("productivity growth, following year (%)")
ax.set_title("Simulated firms: productivity growth vs training hours")
ax.legend(fontsize=8)
fig.tight_layout(); fig.savefig(f"{OUT}/figure1.png", dpi=120); plt.close(fig)

print(table.to_string(index=False))
print(pd.DataFrame(sens).to_string(index=False))
print("conclusion:", conclusion)