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

SEED = 2345
np.set_printoptions(precision=6, suppress=True)

def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)


def normalize(A, eps=1e-8):
    mu = A.mean(axis=0, keepdims=True)
    sig = A.std(axis=0, keepdims=True)
    return (A-mu)/(sig+eps), mu, sig


def toy_verification():
    # Two equally-sized return groups with controlled separation d. Every timestep
    # has the same return ordering, so the conditional-margin prediction is exact.
    B, T = 4000, 5
    rng = np.random.default_rng(SEED)
    feasible = np.zeros(B, dtype=bool); feasible[:B//2] = True
    rng.shuffle(feasible)
    noise = rng.normal(size=(B,T))
    results = {}

    # Prediction 1: normalized variance is (sigma/(sigma+eps))^2, approximately one.
    var_rows=[]
    for scale in [0.05, 0.2, 1.0, 5.0, 20.0]:
        A = scale * noise
        At, _, sig = normalize(A)
        observed = float(At.var(axis=0).mean())
        predicted = float(np.mean((sig/(sig+1e-8))**2))
        var_rows.append({'scale':scale, 'observed_variance':observed, 'predicted':predicted})

    # Prediction 2/3: with equal groups and feasible-return mean higher by d,
    # E[At|F]-E[At|V] = d/sigma; weighted margin is (wF+wV)d/(2 sigma).
    margin_rows=[]
    base_noise = rng.normal(scale=0.25, size=(B,T))
    for d in [0.0, 0.2, 0.5, 1.0, 2.0]:
        G = base_noise.copy()
        G[feasible] += d/2; G[~feasible] -= d/2
        At, _, sig = normalize(G)
        mf = float(At[feasible].mean()); mv = float(At[~feasible].mean())
        sigma = float(sig.mean())
        for ratio in [1.0, 2.0, 4.0]:
            wf, wv = ratio, 1.0
            observed = wf*mf - wv*mv
            predicted = (wf+wv)*d/(2*sigma) if sigma > 0 else 0.0
            margin_rows.append({'separation':d, 'wF':wf, 'wV':wv,
                                'observed_margin':observed, 'predicted_margin':predicted})
    # Fit observed margin slope at wF=2,wV=1, compare formula slope.
    selected=[r for r in margin_rows if r['wF']==2.0]
    x=np.array([r['separation'] for r in selected]); y=np.array([r['observed_margin'] for r in selected])
    slope=float(np.polyfit(x,y,1)[0])
    sigma_ref=float(np.mean(np.abs(base_noise))) # only informational
    return {'normalization_sweep':var_rows, 'margin_sweep':margin_rows,
            'margin_observed_slope_wF2_wV1':slope,
            'margin_prediction_slope_note':'(wF+wV)/(2*sigma), with sigma measured per sweep'}


def rollout(theta, B=64, T=5, device='cpu'):
    # action 0 is feasible/safe and pays .8; action 1 is risky and pays 1.2.
    p=torch.sigmoid(theta)
    acts=torch.bernoulli(p.expand(B,T))
    rewards=torch.where(acts < .5, torch.tensor(.8,device=device), torch.tensor(1.2,device=device))
    feasible=(acts.sum(dim=1)==0)
    G=rewards.sum(dim=1)
    return acts, rewards, feasible, G


def update(theta, old_logp, acts, G, feasible, ranked, clip_eps=.2, wf=3., wv=.25):
    # Group-relative trajectory advantage, repeated across time, then timestep norm.
    raw=(G-G.mean()).unsqueeze(1).expand_as(acts)
    mu=raw.mean(dim=0,keepdim=True); sd=raw.std(dim=0,keepdim=True,unbiased=False)
    A=(raw-mu)/(sd+1e-8)
    p=torch.sigmoid(theta)
    logp=acts*torch.log(p+1e-8)+(1-acts)*torch.log(1-p+1e-8)
    ratio=torch.exp(logp-old_logp)
    clipped=torch.clamp(ratio,1-clip_eps,1+clip_eps)
    weights=torch.where(feasible, torch.tensor(wf), torch.tensor(wv)) if ranked else torch.ones_like(feasible,dtype=torch.float)
    obj=torch.minimum(ratio*A,clipped*A)*weights.unsqueeze(1)
    loss=-obj.mean()
    loss.backward()
    return float(loss.detach()), float(A.var(dim=0,unbiased=False).mean()), float(feasible.float().mean())


def mini_experiment():
    # Same sampled-rollout setup and seeds for standard normalized GRPO vs ranking.
    device='cuda' if torch.cuda.is_available() else 'cpu'
    rows=[]
    for ranked in [False, True]:
        finals=[]; rewards=[]; curves=[]
        for s in range(8):
            seed_all(SEED+s)
            theta=torch.tensor(0.0,device=device,requires_grad=True)
            opt=torch.optim.SGD([theta],lr=.18)
            hist=[]
            for it in range(180):
                opt.zero_grad()
                acts,rews,feas,G=rollout(theta,device=device)
                p=torch.sigmoid(theta)
                oldlp=acts*torch.log(p+1e-8)+(1-acts)*torch.log(1-p+1e-8)
                loss,var,fb=update(theta,oldlp.detach(),acts,G,feas,ranked)
                opt.step()
                hist.append((fb,float(rews.mean())))
            with torch.no_grad():
                p=float(torch.sigmoid(theta)); final_feas=(1-p)**5
                finals.append(final_feas); rewards.append(.8*5*(1-p)+1.2*5*p)
                curves.append(hist)
        rows.append({'method':'ranked' if ranked else 'standard',
                     'final_feasibility_mean':float(np.mean(finals)),
                     'final_feasibility_std':float(np.std(finals)),
                     'policy_risky_probability_mean':float(1-np.mean(np.array(finals)**(1/5))),
                     'expected_reward_mean':float(np.mean(rewards)),
                     'training_rollout_feasibility_last20':float(np.mean([np.mean(c[-20:],axis=0)[0] for c in curves]))})
    return {'device':device,'rows':rows}


def main():
    seed_all(SEED)
    try:
        mini = mini_experiment()
    except Exception as exc:
        print('CUDA/accelerator run failed; retrying on CPU:', repr(exc))
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        old = torch.cuda.is_available
        torch.cuda.is_available = lambda: False
        try:
            mini = mini_experiment()
        finally:
            torch.cuda.is_available = old
    out={'toy_verification':toy_verification(),'mini_experiment':mini}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()