"""Prescribed-performance hidden-state observer MVP.

Run: /home/maxwelhelp/main/bin/python3 ppo_observer_experiment.py
"""
import json, math, random
from pathlib import Path
import numpy as np

SEED = 7
np.random.seed(SEED); random.seed(SEED)


def envelope(k, r0=1.0, rinf=.08, lo=.035, dt=.02):
    return (r0-rinf)*np.exp(-lo*k*dt)+rinf


def ppo_update(hat, e, r, dt=.02, gamma=(2.0,), rinf=.08, eps=1e-6):
    """Euler PPO update, with safe transformed innovation clipping."""
    n = len(hat)
    xi = float(np.clip((e-hat[0])/r, -1+eps, 1-eps))
    T = math.log((1+xi)/(1-xi))
    old = hat.copy()
    for i in range(n-1):
        hat[i] += dt*(old[i+1] + gamma[i]*r*T/(rinf**(i+1)))
    hat[-1] += dt*gamma[-1]*r*T/(rinf**n)
    return hat, xi, T


def fixed_update(hat, e, gain=1.0, dt=.02):
    # A simple conventional fixed-gain first-order residual observer.
    hat[0] += dt*gain*(e-hat[0])
    return hat


def math_check():
    # Constant error: PPO should rapidly reduce normalized innovation while
    # retaining a shrinking prescribed envelope; compare a deliberately modest
    # fixed gain, which cannot track the shrinking envelope as well.
    steps=1800; e=.18
    hp=np.zeros(1); hf=np.zeros(1); pxis=[]; fxis=[]; violations=0; fviol=0
    for k in range(steps):
        r=envelope(k)
        hp, xi, _ = ppo_update(hp,e,r,gamma=(1.4,),rinf=.08)
        hf=fixed_update(hf,e,gain=1.0)
        pxis.append(abs((e-hp[0])/r)); fxis.append(abs((e-hf[0])/r))
        violations += pxis[-1] >= 1
        fviol += fxis[-1] >= 1
    return {
        'ppo_max_normalized_innovation': float(max(pxis)),
        'fixed_max_normalized_innovation': float(max(fxis)),
        'ppo_final_normalized_innovation': float(pxis[-1]),
        'fixed_final_normalized_innovation': float(fxis[-1]),
        'ppo_violations': int(violations), 'fixed_violations': int(fviol),
        'ppo_improves_over_fixed': float(pxis[-1]) < float(fxis[-1]),
    }


def gain_sweep():
    """Probe the Euler gain boundary claimed by the construction."""
    out = {}
    for g in (0.2, 1.4, 5.0, 20.0, 100.0, 500.0):
        h = np.zeros(1); max_xi = 0.0
        for k in range(600):
            h, xi, _ = ppo_update(h, .18, envelope(k), dt=.02,
                                   gamma=(g,), rinf=.08)
            max_xi = max(max_xi, abs(xi))
        out[str(g)] = {'max_abs_xi': float(max_xi),
                       'final_hat_e1': float(h[0])}
    return out


def make_data(n=2600):
    # Observable noisy sinusoid with slowly varying frequency/amplitude.
    t=np.arange(n+1,dtype=np.float64)
    clean=np.sin(.075*t)+.25*np.sin(.013*t+1.1)
    noisy=clean + .10*np.random.randn(n+1)
    return noisy.astype(np.float32), clean.astype(np.float32)


def train_gru(train, epochs=5, hidden=20):
    import torch
    from torch import nn
    torch.manual_seed(SEED)
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        model=nn.GRU(1,hidden,batch_first=True).to(device)
        head=nn.Linear(hidden,1).to(device)
        opt=torch.optim.Adam(list(model.parameters())+list(head.parameters()),lr=3e-3)
        # Short windows make the comparison cheap and reproducible.
        x=[]; y=[]; W=24
        for i in range(0,len(train)-W-1,2):
            x.append(train[i:i+W]); y.append(train[i+W])
        x=torch.tensor(np.asarray(x),device=device).unsqueeze(-1)
        y=torch.tensor(np.asarray(y),device=device).view(-1,1)
        for _ in range(epochs):
            perm=torch.randperm(len(x),device=device)
            for j in range(0,len(x),128):
                ix=perm[j:j+128]; out,_=model(x[ix]); pred=head(out[:,-1])
                loss=((pred-y[ix])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
        return model,head,device
    except Exception as ex:
        # CUDA OOM or another backend issue: retry on CPU as required.
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        device='cpu'; model=nn.GRU(1,hidden,batch_first=True).to(device); head=nn.Linear(hidden,1).to(device)
        opt=torch.optim.Adam(list(model.parameters())+list(head.parameters()),lr=3e-3)
        W=24; x=[]; y=[]
        for i in range(0,len(train)-W-1,2): x.append(train[i:i+W]); y.append(train[i+W])
        x=torch.tensor(np.asarray(x)).unsqueeze(-1); y=torch.tensor(np.asarray(y)).view(-1,1)
        for _ in range(epochs):
            for j in range(0,len(x),128):
                out,_=model(x[j:j+128]); pred=head(out[:,-1]); loss=((pred-y[j:j+128])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
        return model,head,device


def predict_compare(data, clean, model, head, device):
    import torch
    W=24; hp=np.zeros(1); hf=np.zeros(1); raw=[]; ppo=[]; fixed=[]; truth=[]; xis=[]
    with torch.no_grad():
        for k in range(W,len(data)-1):
            seq=torch.tensor(data[k-W:k],device=device).view(1,W,1)
            pred=float(head(model(seq)[0][:,-1]).cpu().item())
            # e is the currently observable residual; observer estimates its
            # slowly varying component and corrects the next output.
            e=float(data[k]-pred); r=envelope(k-W,r0=.55,rinf=.08,lo=.08)
            hp,xi,_=ppo_update(hp,e,r,dt=.05,gamma=(1.0,),rinf=.08)
            hf=fixed_update(hf,e,gain=1.2,dt=.05)
            raw.append(pred); ppo.append(pred+hp[0]); fixed.append(pred+hf[0]); truth.append(float(data[k+1])); xis.append(abs(xi))
    def rmse(a,b): return float(np.sqrt(np.mean((np.asarray(a)-np.asarray(b))**2)))
    return {'raw_gru_rmse':rmse(raw,truth),'ppo_corrected_rmse':rmse(ppo,truth),'fixed_observer_rmse':rmse(fixed,truth),
            'ppo_mean_abs_xi':float(np.mean(xis)),'ppo_max_abs_xi':float(np.max(xis)),
            'ppo_fraction_envelope_violations':float(np.mean(np.asarray(xis)>=1.0))}


def main():
    check=math_check()
    train,clean=make_data(); split=1900
    model,head,device=train_gru(train[:split],epochs=5)
    pred=predict_compare(train[split-24:],clean[split-24:],model,head,device)
    result={'seed':SEED,'device':device,'math_check':check,'gain_sweep':gain_sweep(),'prediction':pred}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()