import json
import random
import time
import numpy as np
import torch
from torch import nn

SEED = 19
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
except Exception:
    device = torch.device('cpu')

L, D, H, N = 30, 4, 8, 2400
rng = np.random.default_rng(SEED)
x = rng.normal(size=(N, L, D)).astype('float32')
y = (x.sum(axis=(1, 2)) > 0).astype('int64')
xt = torch.tensor(x); yt = torch.tensor(y)

class StandardRNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.inp = nn.Linear(D, H)
        self.rec = nn.Linear(H, H)
        self.out = nn.Linear(H, 2)
    def forward(self, x):
        h = torch.zeros(x.size(0), H, device=x.device)
        for t in range(L):
            h = torch.tanh(self.inp(x[:, t]) + self.rec(h))
        return self.out(h)

class ContractingState(nn.Module):
    # A=rI, with a bounded driver and learned forcing phi(x_t,u_t).
    def __init__(self, r):
        super().__init__()
        self.r = r
        self.wx = nn.Linear(D, 1)
        self.force = nn.Sequential(nn.Linear(D + 1, 16), nn.Tanh(), nn.Linear(16, H))
        self.out = nn.Linear(H, 2)
    def forward(self, x):
        h = torch.zeros(x.size(0), H, device=x.device)
        for t in range(L):
            driver = torch.sigmoid(self.wx(x[:, t]))
            h = self.r * h + self.force(torch.cat([driver, x[:, t]], dim=1))
        return self.out(h)

def run(model, epochs=18):
    model = model.to(device)
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    lossfn = nn.CrossEntropyLoss()
    tic = time.time()
    for _ in range(epochs):
        perm = torch.randperm(1800, device='cpu')
        for start in range(0, 1800, 64):
            ind = perm[start:start+64]
            loss = lossfn(model(xt[ind].to(device)), yt[ind].to(device))
            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
    with torch.no_grad():
        pred = model(xt[1800:].to(device)).argmax(1).cpu()
        acc = float((pred == yt[1800:]).float().mean())
        test_loss = float(lossfn(model(xt[1800:].to(device)), yt[1800:].to(device)))
    return {'test_accuracy': acc, 'test_loss': test_loss, 'seconds': time.time()-tic,
            'parameters': sum(p.numel() for p in model.parameters())}

def main():
    global device
    results = {'seed': SEED, 'device': str(device), 'task': 'sign of sequence-wide sum'}
    try:
        results['baseline'] = run(StandardRNN())
        results['idea_r095'] = run(ContractingState(.95))
    except Exception as e:
        # CUDA allocation/runtime errors are handled by rerunning on CPU.
        if device.type == 'cuda':
            device = torch.device('cpu')
            results['device_fallback'] = str(e)
            results['baseline'] = run(StandardRNN())
            results['idea_r095'] = run(ContractingState(.95))
        else:
            raise
    with open('mini_results.json', 'w') as f: json.dump(results, f, indent=2)
    print(json.dumps(results, indent=2))

if __name__ == '__main__': main()