"""Replication entry point for F6-conf-04. Run from workspace root."""
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"
os.makedirs(OUT, exist_ok=True)
raw = pd.read_csv("data/firms.csv")
n_raw = len(raw)

# Exclusion ledger
rows = []
d = raw.dropna()
rows.append(("1", "drop rows with missing values", n_raw, len(d), "no imputation allowed"))
n0 = len(d)
d = d[d["training_hours"] >= 0].copy()
rows.append(("2", "drop training_hours < 0", n0, len(d), "negative hours are physically impossible; treated as data errors"))
excl = pd.DataFrame(rows, columns=["step", "rule", "n_before", "n_after", "rationale"])
excl["n_removed"] = excl["n_before"] - excl["n_after"]
excl[["step", "rule", "n_before", "n_after", "n_removed", "rationale"]].to_csv(f"{OUT}/exclusions.csv", index=False)
n_neg = int((raw["training_hours"] < 0).sum())

COV = "HC1"
naive = smf.ols("productivity_growth ~ training_hours", d).fit(cov_type=COV)
adj = smf.ols("productivity_growth ~ training_hours + skill_index", d).fit(cov_type=COV)
# sensitivity analyses
adj_full = smf.ols("productivity_growth ~ training_hours + skill_index", raw).fit(cov_type=COV)
adj_quad = smf.ols("productivity_growth ~ training_hours + skill_index + I(skill_index**2)", d).fit(cov_type=COV)
adj_conv = smf.ols("productivity_growth ~ training_hours + skill_index", d).fit()

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]))

table = pd.DataFrame([
    row("unadjusted (baseline): growth ~ hours", naive),
    row("adjusted (main): growth ~ hours + skill_index", adj),
    row("sensitivity: adjusted, full sample incl. negative hours", adj_full),
    row("sensitivity: adjusted + skill_index^2", adj_quad),
    row("sensitivity: adjusted, conventional SE", adj_conv),
])
table.to_csv(f"{OUT}/table1.csv", index=False)

rmse = lambda m: float(np.sqrt(np.mean(m.resid ** 2)))
lo, hi = float(adj.conf_int().loc["training_hours", 0]), float(adj.conf_int().loc["training_hours", 1])
concl = "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": [{"issue": f"{n_neg} rows have negative training_hours (physically impossible)",
                     "handling": "excluded from main analysis; full-sample estimate reported as sensitivity in table1.csv"}],
    "n_input_rows": {"firms.csv": int(n_raw)},
    "n_analysed": int(adj.nobs),
    "random_seed": SEED,
    "naive_estimate": float(naive.params["training_hours"]),
    "naive_ci_low": float(naive.conf_int().loc["training_hours", 0]),
    "naive_ci_high": float(naive.conf_int().loc["training_hours", 1]),
    "naive_se": float(naive.bse["training_hours"]),
    "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"]),
    "se_type": "heteroskedasticity-robust HC1",
    "conclusion": concl,
    "sensitivity_full_sample_estimate": float(adj_full.params["training_hours"]),
    "sensitivity_quadratic_skill_estimate": float(adj_quad.params["training_hours"]),
    "sensitivity_conventional_se": float(adj_conv.bse["training_hours"]),
    "corr_skill_training": float(d["skill_index"].corr(d["training_hours"])),
}
with open(f"{OUT}/results.json", "w") as f:
    json.dump(res, f, indent=2)

fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(d["training_hours"], d["productivity_growth"], s=8, alpha=0.4, color="grey", label="firms")
xs = np.linspace(d["training_hours"].min(), d["training_hours"].max(), 100)
ax.plot(xs, naive.params["Intercept"] + naive.params["training_hours"] * xs, "r-",
        label=f"unadjusted OLS (slope {res['naive_estimate']:.3f})")
sm = d["skill_index"].mean()
ax.plot(xs, adj.params["Intercept"] + adj.params["training_hours"] * xs + adj.params["skill_index"] * sm, "b-",
        label=f"adjusted for skill_index at mean (slope {res['effect_estimate']:.3f})")
ax.set_xlabel("training hours per employee"); ax.set_ylabel("productivity growth (%)")
ax.set_title("Productivity growth vs training hours (n=%d)" % res["n_analysed"]); ax.legend()
fig.tight_layout(); fig.savefig(f"{OUT}/figure1.png", dpi=120)
print(json.dumps({k: res[k] for k in ["naive_estimate", "effect_estimate", "effect_se", "conclusion", "n_analysed"]}, indent=1))