"""Replication entry point: run from workspace root with ~/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

SEED = None  # no randomness used
OUT = "out"
d_raw = pd.read_csv("data/firms.csv")
n_raw = len(d_raw)

# Exclusions
excl = []
d = d_raw.dropna()
excl.append(("1", "drop rows with missing values", n_raw, len(d), n_raw - len(d), "required fields"))
n0 = len(d); d = d[d.training_hours >= 0].copy()
excl.append(("2", "drop training_hours < 0", n0, len(d), n0 - len(d), "negative hours per employee are physically impossible (recording error)"))
pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv(f"{OUT}/exclusions.csv", index=False)

COV = "HC1"  # heteroskedasticity-robust SEs
def fit(formula, data):
    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 models
adj_all = fit("productivity_growth ~ training_hours + skill_index", d_raw)
naive_all = fit("productivity_growth ~ training_hours", d_raw)
adj_q = fit("productivity_growth ~ training_hours + skill_index + I(skill_index**2)", d)

def row(name, m):
    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(m.nobs), rmse=rmse(m))
rows = [row("unadjusted (main sample)", naive), row("adjusted for skill_index (main sample) [estimand]", adj),
        row("adjusted + skill_index^2 (sensitivity)", adj_q), row("unadjusted (all 1200 rows, sensitivity)", naive_all),
        row("adjusted for skill_index (all 1200 rows, sensitivity)", adj_all)]
pd.DataFrame(rows).to_csv(f"{OUT}/table1.csv", index=False)

r_n, r_a = rows[0], rows[1]
lo, hi = r_a["ci_low"], r_a["ci_high"]
conclusion = "positive_effect" if lo > 0 else ("negative_effect" if hi < 0 else "no_detectable_effect")
n_neg = int((d_raw.training_hours < 0).sum())
results = {
    "task_id": "F6-conf-04",
    "naive_estimate": r_n["estimate"], "naive_ci_low": r_n["ci_low"], "naive_ci_high": r_n["ci_high"],
    "naive_se": r_n["se"],
    "baseline_metric": r_n["rmse"], "model_metric": r_a["rmse"],
    "effect_estimate": r_a["estimate"], "effect_se": r_a["se"], "effect_se_type": "HC1 heteroskedasticity-robust",
    "effect_ci_low": lo, "effect_ci_high": hi,
    "skill_index_coefficient": float(adj.params["skill_index"]),
    "skill_index_se": float(adj.bse["skill_index"]),
    "conclusion": conclusion,
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": [{"issue": f"{n_neg} rows have negative training_hours (impossible values)",
                     "handling": "excluded from main analysis; sensitivity estimate on all 1200 rows reported (sensitivity_all_rows_effect)"},
                    {"issue": "strong confounding: skill_index correlates 0.69 with training_hours and -0.79 with productivity_growth",
                     "handling": "adjusted for pre-programme skill_index by OLS"}],
    "n_input_rows": {"firms.csv": n_raw},
    "n_analysed": int(adj.nobs),
    "random_seed": SEED,
    "sensitivity_all_rows_effect": rows[4]["estimate"], "sensitivity_all_rows_se": rows[4]["se"],
    "sensitivity_quadratic_skill_effect": rows[2]["estimate"], "sensitivity_quadratic_skill_se": rows[2]["se"],
    "corr_skill_training": float(d.skill_index.corr(d.training_hours)),
}
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="firms (n=%d)" % len(d))
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 {r_n['estimate']:.3f})")
sm = d.skill_index.mean()
ax.plot(x, adj.params["Intercept"] + adj.params["training_hours"] * x + adj.params["skill_index"] * sm, "r-",
        label=f"adjusted for skill_index at mean (slope {r_a['estimate']:.3f})")
ax.set_xlabel("training hours per employee"); ax.set_ylabel("productivity growth next year (%)")
ax.legend(); ax.set_title("Productivity growth vs training hours (simulated data)")
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))