"""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.
No randomness is used (random_seed recorded as 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

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(ROOT, "out")
raw = pd.read_csv(os.path.join(ROOT, "data", "firms.csv"))
n_raw = len(raw)

# ---- exclusions ----
excl = []
d = raw.dropna(subset=["skill_index", "training_hours", "productivity_growth"])
excl.append(("1", "drop rows with missing values in analysis columns", n_raw, len(d),
             "estimation requires complete cases (none found)"))
n0 = len(d)
d = d[d.training_hours >= 0].copy()
excl.append(("2", "drop rows with training_hours < 0", n0, len(d),
             "negative training hours are physically impossible (invalid treatment value)"))
pd.DataFrame([{"step": s, "rule": r, "n_before": a, "n_after": b, "n_removed": a - b,
               "rationale": why} for s, r, a, b, why in excl]).to_csv(
    os.path.join(OUT, "exclusions.csv"), index=False)
n_neg = n0 - len(d)

COV = "HC1"  # heteroskedasticity-robust standard errors


def fit(formula, data):
    m = smf.ols(formula, data).fit(cov_type=COV)
    ci = m.conf_int(0.05).loc["training_hours"]
    rmse = float(np.sqrt(np.mean(m.resid ** 2)))
    return m, float(m.params["training_hours"]), float(m.bse["training_hours"]), float(ci[0]), float(ci[1]), rmse


# ---- baseline (P2): unadjusted OLS ----
mn, b_n, se_n, lo_n, hi_n, rmse_n = fit("productivity_growth ~ training_hours", d)
# ---- target model (P3): adjust for pre-programme skill_index ----
ma, b_a, se_a, lo_a, hi_a, rmse_a = fit("productivity_growth ~ training_hours + skill_index", d)
# ---- sensitivity: full sample incl. negative hours; flexible skill specification ----
_, b_nf, se_nf, lo_nf, hi_nf, _ = fit("productivity_growth ~ training_hours", raw)
_, b_af, se_af, lo_af, hi_af, _ = fit("productivity_growth ~ training_hours + skill_index", raw)
_, b_q, se_q, lo_q, hi_q, _ = fit(
    "productivity_growth ~ training_hours + skill_index + I(skill_index**2) + I(skill_index**3)", d)
# conventional (non-robust) SE for the target model, for comparison
se_a_conv = float(smf.ols("productivity_growth ~ training_hours + skill_index", d).fit().bse["training_hours"])

if lo_a > 0:
    concl = "positive_effect"
elif hi_a < 0:
    concl = "negative_effect"
else:
    concl = "no_detectable_effect"

table = pd.DataFrame([
    ("unadjusted (main sample, n=%d)" % len(d), b_n, se_n, lo_n, hi_n),
    ("adjusted for skill_index (main sample, n=%d) [target]" % len(d), b_a, se_a, lo_a, hi_a),
    ("adjusted, cubic in skill_index (main sample, n=%d)" % len(d), b_q, se_q, lo_q, hi_q),
    ("unadjusted (all rows incl. negative hours, n=%d)" % n_raw, b_nf, se_nf, lo_nf, hi_nf),
    ("adjusted for skill_index (all rows incl. negative hours, n=%d)" % n_raw, b_af, se_af, lo_af, hi_af),
], columns=["model", "estimate", "se", "ci_low", "ci_high"])
table.to_csv(os.path.join(OUT, "table1.csv"), index=False)

results = {
    "task_id": "F6-conf-04",
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": [
        {"issue": "%d rows have negative training_hours (min %.3f), which is physically impossible" % (n_neg, raw.training_hours.min()),
         "handling": "excluded from main analysis (exclusions.csv step 2); sensitivity with all %d rows reported (sensitivity_* keys, table1)" % n_raw},
        {"issue": "skill_index strongly correlated with both training_hours (r=%.2f) and productivity_growth (r=%.2f): confounding of the unadjusted slope" % (raw.skill_index.corr(raw.training_hours), raw.skill_index.corr(raw.productivity_growth)),
         "handling": "target model adjusts for skill_index (pre-programme covariate)"},
    ],
    "n_input_rows": {"firms.csv": int(n_raw)},
    "n_analysed": int(len(d)),
    "random_seed": None,
    "naive_estimate": b_n, "naive_se": se_n, "naive_ci_low": lo_n, "naive_ci_high": hi_n,
    "baseline_metric": rmse_n,
    "model_metric": rmse_a,
    "effect_estimate": b_a, "effect_se": se_a, "effect_ci_low": lo_a, "effect_ci_high": hi_a,
    "effect_se_type": "HC1 heteroskedasticity-robust",
    "effect_se_conventional": se_a_conv,
    "skill_index_coefficient": float(ma.params["skill_index"]),
    "skill_index_se": float(ma.bse["skill_index"]),
    "adjusted_r2": float(ma.rsquared), "naive_r2": float(mn.rsquared),
    "conclusion": concl,
    "sensitivity_cubic_skill_estimate": b_q, "sensitivity_cubic_skill_ci_low": lo_q, "sensitivity_cubic_skill_ci_high": hi_q,
    "sensitivity_fullsample_naive_estimate": b_nf,
    "sensitivity_fullsample_effect_estimate": b_af, "sensitivity_fullsample_effect_ci_low": lo_af, "sensitivity_fullsample_effect_ci_high": hi_af,
    "n_excluded_negative_hours": int(n_neg),
}
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(results, f, indent=2)

# ---- figure ----
fig, ax = plt.subplots(figsize=(7, 5))
sc = ax.scatter(d.training_hours, d.productivity_growth, c=d.skill_index, cmap="viridis", s=8, alpha=0.6)
plt.colorbar(sc, ax=ax, label="skill_index (pre-programme)")
xs = np.linspace(d.training_hours.min(), d.training_hours.max(), 100)
ax.plot(xs, mn.params["Intercept"] + b_n * xs, "r-", lw=2, label="unadjusted OLS: %.3f pp/h" % b_n)
sbar = d.skill_index.mean()
ax.plot(xs, ma.params["Intercept"] + b_a * xs + ma.params["skill_index"] * sbar, "k--", lw=2,
        label="adjusted for skill (at mean skill): %.3f pp/h" % b_a)
ax.set_xlabel("training hours per employee (programme year)")
ax.set_ylabel("productivity growth, following year (%)")
ax.set_title("Productivity growth vs training hours (n=%d)" % len(d))
ax.legend()
fig.tight_layout()
fig.savefig(os.path.join(OUT, "figure1.png"), dpi=120)

print("n_analysed=%d naive=%.4f [%.4f, %.4f] effect=%.4f (HC1 se %.4f) [%.4f, %.4f] conclusion=%s"
      % (len(d), b_n, lo_n, hi_n, b_a, se_a, lo_a, hi_a, concl))