"""F5-conf-04: VFI vs closed form, optimal growth with log utility.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
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, beta = float(p["alpha"]), float(p["beta"])
y_min, y_max, N = float(p["y_min"]), float(p["y_max"]), int(float(p["grid_points"]))
spacing, tol = str(p["grid_spacing"]).strip(), float(p["vfi_tolerance"])
eval_y = [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
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) * np.asarray(y)
V_closed = lambda y: A + B * np.log(np.asarray(y))
savings_rate = ab  # (y - c)/y

# ---------- grid ----------
grid = np.exp(np.linspace(np.log(y_min), np.log(y_max), N))
lng = np.log(grid)

# ---------- VFI: choose next-period output y' on the grid ----------
# k = y'^(1/alpha) is the investment needed; c = y - k must be > 0.
k_needed = grid ** (1.0 / alpha)                 # indexed by j (y')
C = grid[:, None] - k_needed[None, :]            # rows i: current y, cols j: next y'
U = np.where(C > 0, np.log(np.where(C > 0, C, 1.0)), -np.inf)
V = np.zeros(N)
n_iter, sup = 0, np.inf
while sup >= tol:
    Vn = np.max(U + beta * V[None, :], axis=1)
    sup = float(np.max(np.abs(Vn - V)))
    V = Vn
    n_iter += 1
    if n_iter > 100000:
        raise RuntimeError("VFI did not converge")
jstar = np.argmax(U + beta * V[None, :], axis=1)
c_pol = grid - k_needed[jstar]
binding_low = int(np.sum(jstar == 0)); binding_high = int(np.sum(jstar == N - 1))
if binding_low or binding_high:
    data_issues.append({"issue": f"policy hits grid bound at {binding_low} (low) / {binding_high} (high) nodes",
                        "handling": "reported; grid kept as prescribed"})

def interp(vals, y):  # linear interpolation in ln y
    return float(np.interp(np.log(y), lng, vals))

on_node = [bool(np.min(np.abs(grid - y)) < 1e-12) for y in eval_y]
c_num = [interp(c_pol, y) for y in eval_y]
c_cl = [float(c_closed(y)) for y in eval_y]
errs = [abs(a - b) for a, b in zip(c_num, c_cl)]

# diagnostics (not required keys)
yprime_cf = (ab * grid) ** alpha
diag = {
    "eval_points_are_grid_nodes": on_node,
    "abs_error_policy_each_eval_point": errs,
    "max_abs_error_policy_all_grid_nodes": float(np.max(np.abs(c_pol - c_closed(grid)))),
    "max_abs_error_value_all_grid_nodes": float(np.max(np.abs(V - V_closed(grid)))),
    "vfi_truncation_bound_beta_over_1mbeta_times_last_change": beta / (1 - beta) * sup,
    "log_grid_step": float(lng[1] - lng[0]),
    "closed_form_yprime_range_on_grid": [float(yprime_cf.min()), float(yprime_cf.max())],
    "policy_nodes_at_lower_bound": binding_low, "policy_nodes_at_upper_bound": binding_high,
    "closed_form_A": float(A), "closed_form_B": float(B),
    "steady_state_y_closed_form": float(ab ** (alpha / (1 - alpha))),
    "interpolation": "linear in ln y between grid nodes",
    "vfi_choice_variable": "next-period output y' restricted to grid nodes (discrete VFI), initial V=0",
}

# refinement diagnostic (sensitivity only; headline results use the prescribed N)
def vfi_on(Nr):
    g = np.exp(np.linspace(np.log(y_min), np.log(y_max), Nr)); kn = g ** (1.0 / alpha)
    Cr = g[:, None] - kn[None, :]
    Ur = np.where(Cr > 0, np.log(np.where(Cr > 0, Cr, 1.0)), -np.inf)
    Vr, s = np.zeros(Nr), np.inf
    while s >= tol:
        Vq = np.max(Ur + beta * Vr[None, :], axis=1); s = np.max(np.abs(Vq - Vr)); Vr = Vq
    cr = g - kn[np.argmax(Ur + beta * Vr[None, :], axis=1)]
    return max(abs(float(np.interp(np.log(y), np.log(g), cr)) - float(c_closed(y))) for y in eval_y)
diag["sensitivity_max_abs_error_policy_by_grid_points"] = {str(n): vfi_on(n) for n in (200, 800)}

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": float(max(errs)),
    "value_closed_y1": float(V_closed(1.0)),
    "value_num_y1": interp(V, 1.0),
    "savings_rate_closed_form": float(savings_rate),
    "baseline_savings_rate_closed_form": float(savings_rate),
    "contraction_modulus": float(beta),
    "grid_points_used": int(N),
    "n_iterations": int(n_iter),
    "final_sup_change": float(sup),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": data_issues,
    "n_input_rows": {"model.csv": int(len(df))},
    "n_analysed": int(N),
    "random_seed": None,
    "diagnostics": diag,
}
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps({k: v for k, v in res.items() if k != "diagnostics"}, indent=2))