"""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.
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

OUT = "out"
os.makedirs(OUT, exist_ok=True)
df = pd.read_csv("data/firms.csv")
n_raw = len(df)

# ---- audit ----
issues = []
n_neg = int((df.training_hours < 0).sum())
if n_neg:
    issues.append({"issue": f"{n_neg} rows have negative training_hours (min {df.training_hours.min():.3f}), physically impossible",
                   "handling": "kept in main analysis (smooth continuous tail of simulated hours-skill relation; dropping would truncate on the treatment); sensitivity analysis excluding them reported in table1.csv"})
assert df.isna().sum().sum() == 0 and not df.firm_id.duplicated().any()

# exclusion ledger: no rows excluded in main analysis
pd.DataFrame([{"step": 1, "rule": "drop rows with missing values or duplicate firm_id", "n_before": n_raw,
               "n_after": n_raw, "n_removed": 0, "rationale": "none found; all rows retained"}]
             ).to_csv(f"{OUT}/exclusions.csv", index=False)

def fit(formula, data):
    return smf.ols(formula, data=data).fit(cov_type="HC1")

def rmse(m):
    return float(np.sqrt(np.mean(m.resid ** 2)))

naive = fit("productivity_growth ~ training_hours", df)
main = fit("productivity_growth ~ training_hours + skill_index", df)
sens_pos = fit("productivity_growth ~ training_hours + skill_index", df[df.training_hours >= 0])
sens_quad = fit("productivity_growth ~ training_hours + skill_index + I(skill_index**2)", df)
# conventional-SE version of main model for comparison
main_conv = smf.ols("productivity_growth ~ training_hours + skill_index", data=df).fit()

def row(name, m):
    ci = m.conf_int().loc["training_hours"]
    return {"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: y ~ hours (HC1)", naive),
        row("adjusted (main): y ~ hours + skill (HC1)", main),
        row("adjusted, conventional SE", main_conv),
        row("sensitivity: main excl. negative hours (HC1)", sens_pos),
        row("sensitivity: + skill^2 (HC1)", sens_quad)]
pd.DataFrame(rows).to_csv(f"{OUT}/table1.csv", index=False)

r_n, r_m = rows[0], rows[1]
lo, hi = r_m["ci_low"], r_m["ci_high"]
conclusion = "positive_effect" if lo > 0 else ("negative_effect" if hi < 0 else "no_detectable_effect")

results = {
    "task_id": "F6-conf-04",
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"firms.csv": n_raw},
    "n_analysed": int(main.nobs),
    "random_seed": None,
    "naive_estimate": r_n["estimate"],
    "naive_ci_low": r_n["ci_low"],
    "naive_ci_high": r_n["ci_high"],
    "baseline_metric": r_n["rmse"],
    "model_metric": r_m["rmse"],
    "effect_estimate": r_m["estimate"],
    "effect_se": r_m["se"],
    "effect_se_type": "HC1 heteroskedasticity-robust",
    "effect_ci_low": lo,
    "effect_ci_high": hi,
    "skill_index_coefficient": float(main.params["skill_index"]),
    "skill_index_se": float(main.bse["skill_index"]),
    "corr_skill_hours": float(df.skill_index.corr(df.training_hours)),
    "sens_excl_negative_hours_estimate": rows[3]["estimate"],
    "sens_excl_negative_hours_n": rows[3]["n"],
    "sens_skill_quadratic_estimate": rows[4]["estimate"],
    "sens_skill_quadratic_sq_coef_pvalue": float(sens_quad.pvalues["I(skill_index ** 2)"]),
    "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(df.training_hours, df.productivity_growth, s=6, alpha=0.4, color="grey", label="firms")
x = np.linspace(df.training_hours.min(), df.training_hours.max(), 100)
ax.plot(x, naive.params["Intercept"] + naive.params["training_hours"] * x, color="C3",
        label=f"unadjusted (slope {r_n['estimate']:.3f})")
ax.plot(x, main.params["Intercept"] + main.params["training_hours"] * x + main.params["skill_index"] * df.skill_index.mean(),
        color="C0", label=f"adjusted for skill, at mean skill (slope {r_m['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)")
fig.tight_layout()
fig.savefig(f"{OUT}/figure1.png", dpi=120)
print(json.dumps(results, indent=2))
print(pd.DataFrame(rows).to_string())