import json, math, random, time
import numpy as np
import torch
from fractional_optimizer import FractionalMemory, LateTimeOrderEstimator, AdamW

SEED = 1729

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


def get_device():
    try:
        d = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        if d.type == "cuda": torch.zeros(1, device=d).sum().item()
        return d
    except Exception:
        return torch.device("cpu")


def verify_power_law():
    # M(t)=C t^(-m alpha), so the estimator should recover alpha for m=1.
    true_alpha, m, rho = 0.68, 1, 1.8
    est = LateTimeOrderEstimator(alpha0=.5, amin=.2, amax=.95,
                                 beta=0.0, warmup=10, lag=10,
                                 check_every=1, monotone_checks=2)
    vals = []
    for k in range(1, 401):
        # Positive, monotonically decreasing late-time signal with tiny noise.
        t = float(k + 1)
        loss = t ** (-m * true_alpha) * (1 + 0.0002 * math.sin(k))
        vals.append(est.update(k, loss))
    raw = -math.log((401.0/391.0)**(-m*true_alpha)) / (m*math.log(401.0/391.0))
    return {"true_alpha": true_alpha, "raw_asymptotic_estimate": raw,
            "final_estimate": vals[-1], "absolute_error": abs(vals[-1]-true_alpha),
            "passed": abs(vals[-1]-true_alpha) < .06}


def make_problem(n=24, d=20, device=None):
    g = torch.Generator(device="cpu").manual_seed(SEED)
    X = torch.randn(n, d, generator=g).to(device)
    true_w = torch.randn(d, 1, generator=g).to(device)
    y = X @ true_w + .05 * torch.randn(n, 1, generator=g).to(device)
    return X, y


def train(kind, device, steps=180):
    seed_all(SEED)
    X, y = make_problem(device=device)
    w = torch.nn.Parameter(torch.zeros(20, 1, device=device))
    if kind == "adamw":
        opt = AdamW([w], lr=.08, weight_decay=0.0)
    elif kind.startswith("fixed"):
        alpha = float(kind.split("_")[1])
        opt = FractionalMemory([w], lr=.08, alpha=alpha, history=32)
    else:
        # The estimator intentionally starts conservatively and adapts only
        # after warmup; loss itself is the positive scalar observation.
        e = LateTimeOrderEstimator(alpha0=.5, amin=.2, amax=.95, beta=.8,
                                   warmup=25, lag=12, check_every=3,
                                   monotone_checks=2)
        opt = FractionalMemory([w], lr=.08, alpha=.5, history=32,
                               adaptive=True, estimator=e)
    losses, alphas = [], []
    start = time.perf_counter()
    for step in range(1, steps + 1):
        pred = X @ w
        loss = ((pred-y)**2).mean()
        loss.backward()
        if kind == "adamw":
            opt.step()
        else:
            opt.step(loss_value=float(loss.detach()))
        opt.zero_grad(set_to_none=True)
        losses.append(float(loss.detach()))
        if kind == "adaptive": alphas.append(opt.estimator.alpha)
    elapsed = time.perf_counter() - start
    tail = float(np.mean(losses[-20:]))
    return {"final_loss": losses[-1], "tail20_loss": tail,
            "loss_at_40": losses[39], "seconds": elapsed,
            "alpha_final": (alphas[-1] if alphas else None),
            "alpha_min": (min(alphas) if alphas else None),
            "alpha_max": (max(alphas) if alphas else None)}


def main():
    device = get_device()
    check = verify_power_law()
    results = {"device": str(device), "math_check": check, "runs": {}}
    for kind in ["adamw", "fixed_0.3", "fixed_0.5", "fixed_0.7", "fixed_0.9", "adaptive"]:
        results["runs"][kind] = train(kind, device)
    print(json.dumps(results, indent=2))

if __name__ == "__main__": main()