"""Replication entry point for F6-conf-04. Run from workspace root:
    ~/wp29/venv/bin/python out/analysis.py
Reads only data/firms.csv; writes out/results.json, out/table1.csv, out/figure1.png, out/exclusions.csv.
Deterministic: no randomness is used (random_seed = null)."""
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

DATA = "data/firms.csv"
OUT = "out"
os.makedirs(OUT, exist_ok=True)

raw = pd.read_csv(DATA)
n_raw = len(raw)
cols = ["firm_id", "skill_index", "training_hours", "productivity_growth"]
assert list(raw.columns) == cols

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

# Step 1: missing values (none expected)
nb = len(d); d = d.dropna(subset=cols); excl.append(("1", "drop rows with missing values", nb, len(d), "required variables must be observed"))
# Step 2: duplicate firm_id (none expected)
nb = len(d); d = d.drop_duplicates(subset="firm_id"); excl.append(("2", "drop duplicate firm_id", nb, len(d), "one row per firm"))
# Step 3: negative training hours (physically impossible)
n_neg = int((d.training_hours < 0).sum())
nb = len(d); d = d[d.training_hours >= 0]; excl.append(("3", "drop training_hours < 0", nb, len(d),
        "hours of training per employee cannot be negative; values are invalid (likely data error). Sensitivity with them retained is reported."))
if n_neg:
    data_issues.append({"issue": f"{n_neg} rows have negative training_hours (min {raw.training_hours.min():.3f}), which is physically impossible",
                        "handling": "excluded from the main analysis; sensitivity analysis re-estimates both models on all 1200 rows (keys sens_all_rows_*)"})
data_issues.append({"issue": "training_hours is strongly correlated with pre-programme skill_index (r = %.3f in all 1200 rows), and skill_index is strongly negatively related to productivity_growth: the unadjusted slope is confounded" % raw[["skill_index","training_hours"]].corr().iloc[0,1],
                    "handling": "skill_index (measured before the programme) is included as a control in the model that targets the estimand"})

pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "rationale"]).assign(
    n_removed=lambda x: x.n_before - x.n_after)[["step", "rule", "n_before", "n_after", "n_removed", "rationale"]].to_csv(
    os.path.join(OUT, "exclusions.csv"), index=False)

COV = "HC1"  # heteroskedasticity-robust standard errors
def fit(formula, df):
    return smf.ols(formula, 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)
# robustness checks (not the main estimate)
quad = fit("productivity_growth ~ training_hours + skill_index + I(skill_index**2)", d)
inter = 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)

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

table = pd.DataFrame([
    row("M1 unadjusted: growth ~ hours", naive),
    row("M2 adjusted (main): growth ~ hours + skill_index", adj),
    row("R1 adds skill_index^2", quad),
    row("R2 adds hours x skill_index (effect at skill_index=0)", inter),
    row("S1 unadjusted, all 1200 rows incl. negative hours", naive_all),
    row("S2 adjusted, all 1200 rows incl. negative hours", adj_all),
])
table.to_csv(os.path.join(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])
conclusion = "positive_effect" if lo > 0 else ("negative_effect" if hi < 0 else "no_detectable_effect")

res = {
    "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(adj.nobs),
    "random_seed": None,
    "naive_estimate": float(naive.params["training_hours"]),
    "naive_se": float(naive.bse["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"]),
    "skill_index_se": float(adj.bse["skill_index"]),
    "se_type": "HC1 heteroskedasticity-robust; 95% CI = estimate +/- 1.96*SE (normal)",
    "conclusion": conclusion,
    "corr_skill_training": float(d[["skill_index", "training_hours"]].corr().iloc[0, 1]),
    "robust_quadratic_skill_effect": float(quad.params["training_hours"]),
    "robust_quadratic_skill_sq_pvalue": float(quad.pvalues["I(skill_index ** 2)"]),
    "robust_interaction_coef": float(inter.params["training_hours:skill_index"]),
    "robust_interaction_pvalue": float(inter.pvalues["training_hours:skill_index"]),
    "sens_all_rows_naive_estimate": float(naive_all.params["training_hours"]),
    "sens_all_rows_effect_estimate": float(adj_all.params["training_hours"]),
    "sens_all_rows_effect_se": float(adj_all.bse["training_hours"]),
    "n_excluded_negative_hours": n_neg,
}
with open(os.path.join(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.35, color="grey", label="firms (n=%d)" % len(d))
xs = np.linspace(d.training_hours.min(), d.training_hours.max(), 100)
ax.plot(xs, naive.params["Intercept"] + naive.params["training_hours"] * xs, color="tab:red",
        label="unadjusted OLS (slope %.3f)" % naive.params["training_hours"])
sbar = d.skill_index.mean()
ax.plot(xs, adj.params["Intercept"] + adj.params["skill_index"] * sbar + adj.params["training_hours"] * xs,
        color="tab:blue", label="adjusted for skill_index, at mean skill (slope %.3f)" % adj.params["training_hours"])
ax.set_xlabel("training hours per employee (programme year)")
ax.set_ylabel("productivity growth, following year (%)")
ax.set_title("Training hours vs productivity growth (simulated data)")
ax.legend(fontsize=8)
fig.tight_layout()
fig.savefig(os.path.join(OUT, "figure1.png"), dpi=120)

print(f"n_raw={n_raw} n_analysed={res['n_analysed']}")
print(f"naive slope={res['naive_estimate']:.4f} [{res['naive_ci_low']:.4f}, {res['naive_ci_high']:.4f}]")
print(f"adjusted slope={res['effect_estimate']:.4f} (HC1 se {res['effect_se']:.4f}) [{lo:.4f}, {hi:.4f}] -> {conclusion}")