"""F5-conf-04: VFI vs closed form for Brock-Mirman growth with log utility.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Deterministic; no random numbers used."""
import json, os
import numpy as np
import pandas as pd

os.makedirs("out", exist_ok=True)
df = pd.read_csv("data/model.csv")
n_rows = len(df)
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" and 0 < alpha < 1 and 0 < beta < 1 and y_min < y_max
data_issues = []

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

# ---------- VFI on prescribed grid ----------
# State y on log grid; control = next-period output y' chosen on the same grid.
# Consumption implied: c = y - y'^(1/alpha) (capital k = y'^(1/alpha)).
grid = np.exp(np.linspace(np.log(y_min), np.log(y_max), N))
k_needed = grid ** (1.0 / alpha)                     # capital required for each y'
C = grid[:, None] - k_needed[None, :]                # rows: y, cols: y'
feasible = C > 0
assert feasible.any(axis=1).all(), "some grid node has no feasible choice"
U = np.full_like(C, -np.inf); U[feasible] = np.log(C[feasible])
V = np.zeros(N); n_iter = 0; sup = np.inf
while sup >= tol:
    V_new = np.max(U + beta * V[None, :], axis=1)
    sup = float(np.max(np.abs(V_new - V))); V = V_new; n_iter += 1
    if n_iter > 100000:
        raise RuntimeError("VFI did not converge")
idx = np.argmax(U + beta * V[None, :], axis=1)
c_pol = C[np.arange(N), idx]
lg = np.log(grid)
interp = lambda f, y: float(np.interp(np.log(y), lg, f))   # linear in ln y
c_num = [interp(c_pol, y) for y in evals]
c_cl = [float(c_closed(y)) for y in evals]
max_err = float(max(abs(a - b) for a, b in zip(c_num, c_cl)))
v_num_1 = interp(V, 1.0)
for y in evals:
    if not np.any(np.isclose(grid, y, rtol=0, atol=1e-12)):
        data_issues.append({"issue": f"evaluation point y={y} is not a grid node",
                            "handling": "VFI policy/value linearly interpolated in ln y between adjacent grid nodes"})
data_issues.append({"issue": "VFI control is discretised: next-period output y' restricted to the 400-point grid",
                    "handling": "standard discrete-choice VFI on the prescribed grid; no off-grid optimisation, so policy error reflects grid spacing"})

# diagnostics (not required keys)
diag = {
    "grid_log_step": float(lg[1] - lg[0]),
    "max_abs_error_policy_on_grid": float(np.max(np.abs(c_pol - c_closed(grid)))),
    "max_abs_error_value_on_grid": float(np.max(np.abs(V - V_closed(grid)))),
    "vfi_error_bound_beta_over_1mbeta_times_sup": float(beta / (1 - beta) * sup),
    "steady_state_y": float(ab ** (alpha / (1 - alpha))),
    "choice_index_min_max": [int(idx.min()), int(idx.max())],
    "closed_form_A": float(A), "closed_form_B": float(B),
}

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": v_num_1,
    "savings_rate_closed_form": float(savings_rate),
    "baseline_savings_rate": float(savings_rate),
    "contraction_modulus": beta,
    "grid_points_used": int(N), "n_iterations": int(n_iter), "final_sup_change": sup,
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": data_issues,
    "n_input_rows": {"model.csv": int(n_rows)},
    "n_analysed": int(N), "random_seed": None,
    "diagnostics": diag,
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
pd.DataFrame([{"step": "parse", "rule": "read all parameter rows", "n_before": n_rows,
               "n_after": n_rows, "n_removed": 0, "rationale": "no rows excluded"}]).to_csv("out/exclusions.csv", index=False)
print(json.dumps({k: v for k, v in res.items() if k not in ("data_issues", "diagnostics")}, indent=1))