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

SEED = 924
rng = np.random.default_rng(SEED)

# A binary-state finite-horizon process. At each step, next state is sampled,
# reward has conditional mean r[state,next] and Gaussian conditional noise.
def recursion(h, p, gamma, r, reward_var, terminal_mean=None, terminal_var=None):
    # m[t,x], v[t,x] describe return from state x with t steps remaining.
    m = np.zeros((h + 1, 2)); v = np.zeros((h + 1, 2))
    if terminal_mean is not None: m[0] = terminal_mean
    if terminal_var is not None: v[0] = terminal_var
    P = np.array([[1-p, p], [p, 1-p]], float)
    for t in range(1, h + 1):
        for x in range(2):
            z = r[x] + gamma * m[t-1] # reward mean depends only on current state
            m[t,x] = z @ P[x]
            # conditional reward noise plus propagated child uncertainty and selection variance
            v[t,x] = np.sum(P[x] * (reward_var[x] + gamma**2 * v[t-1] + (z - m[t,x])**2))
    return m, v

def sample_returns(h, n, p, gamma, r, sigma, start=0):
    x = np.full(n, start, dtype=np.int64); out = np.zeros(n)
    for t in range(h):
        nx = np.where(rng.random(n) < p, 1-x, x)
        out += (gamma**t) * (r[x] + sigma[x] * rng.normal(size=n))
        x = nx
    return out

def main():
    gamma=.9; r=np.array([-1., 1.]); sig=np.array([.35,.35]); rv=sig**2
    rows=[]
    # Prediction 1: recursive variance equals held-out Monte Carlo variance.
    for h in [1,2,4,8,12]:
        m,v=recursion(h,.35,gamma,r,rv)
        g=sample_returns(h, 300000, .35,gamma,r,sig)
        rows.append({'h':h,'recursive_var':float(v[h,0]),'mc_var':float(g.var()),
                     'relative_error':float(abs(g.var()-v[h,0])/v[h,0])})
    # Prediction 2: one-step transition-selection term is p(1-p)*(delta child mean)^2,
    # hence zero at deterministic transitions and maximal at p=.5.
    sel=[]
    child=np.array([-2., 3.]); delta=child[1]-child[0]
    for p in [0.,.1,.25,.5,.75,.9,1.]:
        P=np.array([1-p,p]); mean=P@child
        measured=np.sum(P*(child-mean)**2)
        predicted=p*(1-p)*delta**2
        sel.append({'p':p,'predicted':float(predicted),'measured':float(measured)})
    # Prediction 3: with independently redrawn state each step, variance is the
    # discounted sum of per-step variances; removing transition randomness removes
    # the state-selection contribution and leaves reward-noise variance only.
    horizon=[]
    for transition_var in [0.,1.]:
        # transition_var=1 means iid +/-1 state each step, 0 means fixed state +1.
        vals=[]
        for h in [1,2,4,8,16]:
            state_var=transition_var # Var +/-1 with equal probability
            predicted=sum((gamma**(2*t))*(state_var + .35**2) for t in range(h))
            g=np.zeros(250000)
            for t in range(h):
                states=(rng.choice([-1.,1.],size=g.size) if transition_var else np.ones(g.size))
                g += gamma**t*(states + .35*rng.normal(size=g.size))
            vals.append({'h':h,'predicted':float(predicted),'mc_var':float(g.var())})
        horizon.append({'transition_randomness':transition_var,'points':vals})
    # Small baseline-vs-idea regression: contexts have heteroscedastic known recursive
    # targets. MSE is standard; weighted NLL uses the detached recursive variance.
    try:
        import torch
        torch.manual_seed(SEED)
        device='cuda' if torch.cuda.is_available() else 'cpu'
        try:
            x=torch.rand(12000,1,device=device)
            # x controls branch probability; target is child +/-1 plus reward noise.
            pp=.05+.9*x[:,0]
            means=2*pp-1
            q=.25 + 4*pp*(1-pp) # reward variance + transition selection variance
            y=means + torch.sqrt(torch.tensor(.25,device=device))*torch.randn_like(means)
            # split is fixed and models are intentionally tiny
            def train(kind):
                net=torch.nn.Sequential(torch.nn.Linear(1,24),torch.nn.Tanh(),torch.nn.Linear(24,2)).to(device)
                opt=torch.optim.Adam(net.parameters(),lr=.01)
                for _ in range(180):
                    out=net(x[:9000]); mu=out[:,0]; logv=out[:,1]
                    if kind=='mse': loss=((mu-y[:9000])**2).mean()
                    else:
                        # recursive q is detached, as prescribed; mean is trained with
                        # heteroscedastic Gaussian NLL and the variance head is fitted too.
                        loss=.5*(((y[:9000]-mu)**2)/q[:9000].detach()+torch.log(q[:9000].detach())).mean()
                        loss += .05*.5*(logv-q[:9000].detach().log()).pow(2).mean()
                    opt.zero_grad(); loss.backward(); opt.step()
                with torch.no_grad():
                    o=net(x[9000:]); err=(o[:,0]-y[9000:])**2
                    cal=(err/q[9000:]).mean().item()
                    mse=err.mean().item()
                return mse,cal
            mse,cal_mse=train('mse'); mse_i,cal_i=train('idea')
            neural={'device':device,'baseline_mse':mse,'baseline_calibration_ratio':cal_mse,
                    'idea_mse':mse_i,'idea_calibration_ratio':cal_i}
        except Exception as e:
            neural={'fallback_error':str(e)}
    except Exception as e: neural={'torch_error':str(e)}
    result={'seed':SEED,'parameters':{'gamma':gamma,'reward_sigma':.35},
            'prediction_1_recursive_vs_mc':rows,'prediction_2_selection_term':sel,
            'prediction_3_horizon':horizon,'mini_experiment':neural}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()