import json, math, random
from pathlib import Path
import numpy as np
import torch

SEED = 1234
K = 15
STEPS = 1800
LR = 0.03
N_COL = 96
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)

# Formal-series recurrence for u' = 1 + u^2, u(0)=0.
# (n+1)c[n+1] = 1_{n=0} + sum_{i=0}^n c[i]c[n-i].
def formal_support(K):
    c = np.zeros(K + 1, dtype=np.float64)
    for n in range(K):
        rhs = (1.0 if n == 0 else 0.0) + sum(c[i] * c[n-i] for i in range(n + 1))
        c[n + 1] = rhs / (n + 1)
    return [i for i, x in enumerate(c) if abs(x) > 1e-14], c

def polynomial_value(coef, t, exponents):
    # Horner is unnecessary at this small K and this makes the exponent mask explicit.
    powers = torch.stack([t ** int(i) for i in exponents], dim=1)
    return (powers * coef.unsqueeze(0)).sum(dim=1)

def train(exponents, seed):
    seed_all(seed)
    t = torch.linspace(0.0, 1.0, N_COL, device=DEVICE)
    # Coefficients are the minimal coefficient-output layer described in the idea.
    coef = torch.nn.Parameter(0.01 * torch.randn(len(exponents), device=DEVICE))
    opt = torch.optim.Adam([coef], lr=LR)
    history = []
    ex = [int(i) for i in exponents]
    for step in range(STEPS):
        opt.zero_grad(set_to_none=True)
        u = polynomial_value(coef, t, ex)
        du = torch.zeros_like(t)
        for j, i in enumerate(ex):
            if i > 0:
                du = du + coef[j] * i * t ** (i - 1)
        residual = du - 1.0 - u * u
        loss = torch.mean(residual ** 2) + 10.0 * coef[ex.index(0)] ** 2 if 0 in ex else torch.mean(residual ** 2)
        loss.backward(); opt.step()
        if step in (0, 99, 299, 599, 1199, STEPS - 1): history.append(float(loss.detach().cpu()))
    with torch.no_grad():
        te = torch.linspace(0.0, 1.0, 401, device=DEVICE)
        ue = polynomial_value(coef, te, ex)
        exact = torch.tan(te)
        rel = torch.linalg.vector_norm(ue - exact) / torch.linalg.vector_norm(exact)
        u = polynomial_value(coef, t, ex)
        du = torch.zeros_like(t)
        for j, i in enumerate(ex):
            if i > 0: du += coef[j] * i * t ** (i - 1)
        residual = du - 1.0 - u * u
        return {"final_loss": float((torch.mean(residual**2)).cpu()),
                "relative_error": float(rel.cpu()),
                "coeff_l2": float(torch.linalg.vector_norm(coef).cpu()),
                "history": history, "parameters": len(ex)}

def main():
    global DEVICE
    try:
        seed_all(SEED)
        _ = torch.zeros(1, device=DEVICE)
    except Exception:
        DEVICE = "cpu"
        seed_all(SEED)
    support, coeffs = formal_support(K)
    # Independent direct check: all even coefficients are forced to zero by recurrence.
    even_abs = max(abs(coeffs[i]) for i in range(0, K + 1, 2))
    odd_nonzero = [i for i in range(1, K + 1, 2) if abs(coeffs[i]) > 1e-14]
    results = {"device": DEVICE, "K": K, "collocation": N_COL, "steps": STEPS,
               "formal_support": support, "expected_odd_support": odd_nonzero,
               "max_forced_even_coefficient": even_abs, "math_check_pass": bool(even_abs < 1e-14)}
    runs = {"dense": [], "restricted": []}
    for seed in (11, 22, 33):
        for name, ex in (("dense", list(range(K + 1))), ("restricted", support)):
            try:
                runs[name].append(train(ex, seed))
            except Exception as e:
                if DEVICE != "cpu":
                    DEVICE = "cpu"
                    runs[name].append(train(ex, seed))
                else:
                    raise
    results["runs"] = runs
    for name in runs:
        for metric in ("final_loss", "relative_error", "parameters"):
            vals = [x[metric] for x in runs[name]]
            results.setdefault(name, {})[metric + "_mean"] = float(np.mean(vals))
            results[name][metric + "_std"] = float(np.std(vals))
    Path("results.json").write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))

if __name__ == "__main__": main()