"""F5-conf-04: VFI vs closed-form policy, optimal growth with log utility.
Reads only data/model.csv; writes out/results.json (deterministic, no randomness)."""
import json, os
import numpy as np
import pandas as pd

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
df = pd.read_csv(os.path.join(ROOT, "data", "model.csv"))
p = dict(zip(df["parameter"], df["value"]))
alpha = float(p["alpha"]); beta = float(p["beta"])
y_min = float(p["y_min"]); y_max = float(p["y_max"])
N = int(float(p["grid_points"])); spacing = str(p["grid_spacing"]).strip()
tol = float(p["vfi_tolerance"])
evals = [float(p["eval_y_1"]), float(p["eval_y_2"]), float(p["eval_y_3"])]
assert spacing == "log"

data_issues = []

# ---- Closed form (guess V = A + B ln y, verified in report) ----
ab = alpha * beta
s_closed = ab                                   # savings rate (y-c)/y
B = 1.0 / (1.0 - ab)
A = (np.log(1.0 - ab) + (ab / (1.0 - ab)) * np.log(ab)) / (1.0 - beta)
c_closed = lambda y: (1.0 - ab) * y
V_closed = lambda y: A + B * np.log(y)

# ---- VFI on the prescribed log grid; choice of next-period output y' on the same grid ----
y = np.exp(np.linspace(np.log(y_min), np.log(y_max), N))
k_next = y ** (1.0 / alpha)                     # capital needed to produce y' = k^alpha
C = y[:, None] - k_next[None, :]                # consumption for (y, y')
feasible = C > 0
U = np.full(C.shape, -np.inf)
U[feasible] = np.log(C[feasible])
V = np.zeros(N)
n_iter = 0; sup_change = np.inf
while sup_change >= tol:
    RHS = U + beta * V[None, :]
    V_new = RHS.max(axis=1)
    sup_change = float(np.max(np.abs(V_new - V)))
    V = V_new; n_iter += 1
    if n_iter > 100000:
        raise RuntimeError("VFI did not converge")
pol_idx = (U + beta * V[None, :]).argmax(axis=1)
c_grid = y - k_next[pol_idx]
# check that the upper/lower bound of the choice grid never binds
n_bind_upper = int(np.sum(pol_idx == N - 1)); n_bind_lower = int(np.sum(pol_idx == 0))
if n_bind_upper or n_bind_lower:
    data_issues.append({"issue": f"choice grid bound binds at {n_bind_lower} (lower) / {n_bind_upper} (upper) grid points",
                        "handling": "reported; affects policy near grid edges only"})

lny = np.log(y)
interp = lambda f, x: float(np.interp(np.log(x), lny, f))   # linear in ln y
on_grid = [bool(np.any(np.isclose(y, e, rtol=0, atol=1e-12))) for e in evals]
off = [e for e, g in zip(evals, on_grid) if not g]
if off:
    data_issues.append({"issue": f"evaluation points {off} are not nodes of the prescribed log grid",
                        "handling": "VFI policy and value interpolated linearly in ln y between adjacent grid nodes (as allowed by TASK.md)"})

c_num = [interp(c_grid, e) for e in evals]
c_cl = [float(c_closed(e)) for e in evals]
max_err = float(max(abs(a - b) for a, b in zip(c_num, c_cl)))

res = {
    "task_id": "F5-conf-04",
    "c_closed_y05": c_cl[0], "c_closed_y1": c_cl[1], "c_closed_y2": c_cl[2],
    "c_num_y05": c_num[0], "c_num_y1": c_num[1], "c_num_y2": c_num[2],
    "max_abs_error_policy": max_err,
    "value_closed_y1": float(V_closed(1.0)),
    "value_num_y1": interp(V, 1.0),
    "savings_rate_closed_form": float(s_closed),
    "contraction_modulus": beta,
    "grid_points_used": N,
    "n_iterations": n_iter,
    "final_sup_change": sup_change,
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": data_issues,
    "n_input_rows": {"model.csv": int(len(df))},
    "n_analysed": N,
    "random_seed": None,
    # supplementary diagnostics
    "baseline_savings_rate_closed_form": float(s_closed),
    "diag_closed_form_A": float(A), "diag_closed_form_B": float(B),
    "diag_grid_log_step": float(lny[1] - lny[0]),
    "diag_max_abs_value_error_grid": float(np.max(np.abs(V - V_closed(y)))),
    "diag_max_abs_policy_error_grid": float(np.max(np.abs(c_grid - c_closed(y)))),
    "diag_mean_savings_rate_num_grid": float(np.mean(1 - c_grid / y)),
    "diag_stopping_error_bound": float(beta * sup_change / (1 - beta)),
    "diag_choice_bound_binds": n_bind_lower + n_bind_upper,
}
os.makedirs(os.path.join(ROOT, "out"), exist_ok=True)
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    json.dump(res, f, indent=2)
pd.DataFrame([{"step": "none", "rule": "no filtering applied; all 10 parameter rows used",
               "n_before": len(df), "n_after": len(df), "n_removed": 0,
               "rationale": "inputs are model parameters, not observations"}]).to_csv(
    os.path.join(ROOT, "out", "exclusions.csv"), index=False)
print(json.dumps(res, indent=2))