"""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"]))
tol = float(p["vfi_tolerance"])
evals = [float(p["eval_y_1"]), float(p["eval_y_2"]), float(p["eval_y_3"])]
assert p["grid_spacing"].strip().lower() == "log"
data_issues = []

# ---- Closed form (guess-and-verify V = A + B ln y) ----
ab = alpha * beta
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) * np.asarray(y)
V_closed = lambda y: A + B * np.log(np.asarray(y))
savings_rate = ab  # baseline quantity (y-c)/y

# ---- VFI on the prescribed log grid; choice = next-period output y' on the same grid ----
y = np.exp(np.linspace(np.log(y_min), np.log(y_max), N))
k_needed = y ** (1.0 / alpha)                 # capital (savings) needed to reach y'_j
C = y[:, None] - k_needed[None, :]            # consumption for state i, choice j
feasible = C > 0
U = np.where(feasible, np.log(np.where(feasible, C, 1.0)), -np.inf)
V = np.zeros(N)
n_iter, sup = 0, np.inf
while sup >= tol:
    TV = np.max(U + beta * V[None, :], axis=1)
    sup = float(np.max(np.abs(TV - V)))
    V = TV
    n_iter += 1
    if n_iter > 100000:
        raise RuntimeError("VFI did not converge")
pol = np.argmax(U + beta * V[None, :], axis=1)   # policy for the converged V
c_grid = y - k_needed[pol]
hits_bound = bool(np.any(pol == 0) or np.any(pol == N - 1))
if hits_bound:
    data_issues.append({"issue": "VFI policy hits grid boundary", "handling": "reported; results may be truncated"})

c_num = [float(np.interp(e, y, c_grid)) for e in evals]          # linear in y
V_num1 = float(np.interp(np.log(evals[1]), np.log(y), V))       # linear in ln y
cc = [float(c_closed(e)) for e in evals]
data_issues.append({"issue": "Evaluation points y=0.5,1,2 are not grid nodes",
                    "handling": "VFI policy linearly interpolated in y, value function linearly interpolated in ln y between adjacent nodes (allowed by TASK.md)"})

res = {
    "task_id": "F5-conf-04",
    "c_closed_y05": cc[0], "c_closed_y1": cc[1], "c_closed_y2": cc[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(abs(a - b) for a, b in zip(c_num, cc))),
    "value_closed_y1": float(V_closed(1.0)),
    "value_num_y1": V_num1,
    "savings_rate_closed_form": 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(len(df))},
    "n_analysed": int(N),
    "random_seed": None,
    # supplementary diagnostics (not required keys)
    "diag_closed_form_A": float(A), "diag_closed_form_B": float(B),
    "diag_max_abs_policy_error_grid": float(np.max(np.abs(c_grid - c_closed(y)))),
    "diag_max_abs_value_error_grid": float(np.max(np.abs(V - V_closed(y)))),
    "diag_vfi_error_bound_tol": float(beta * tol / (1 - beta)),
    "diag_log_grid_step": float(np.log(y[1] / y[0])),
    "diag_policy_hits_grid_boundary": hits_bound,
    "diag_policy_index_range": [int(pol.min()), int(pol.max())],
    "diag_eval_points": evals,
}
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    json.dump(res, f, indent=2)
for k in ["c_closed_y05","c_num_y05","c_closed_y1","c_num_y1","c_closed_y2","c_num_y2",
          "max_abs_error_policy","value_closed_y1","value_num_y1","n_iterations","final_sup_change"]:
    print(k, res[k])