import json
import numpy as np
import torch
from order_adaptive import verify, OrderAdaptiveIntegral


def torch_smoke(seed=11, steps=160):
    torch.manual_seed(seed)
    x = torch.randn(96, 2)
    y = (2.0 * x[:, :1] - 0.2 * x[:, 1:]).tanh()

    def train(adaptive):
        torch.manual_seed(seed)
        model = torch.nn.Sequential(
            torch.nn.Linear(2, 8), torch.nn.Tanh(), torch.nn.Linear(8, 1)
        )
        if adaptive:
            opt = OrderAdaptiveIntegral(
                model.parameters(), lr=0.04, max_order=2, beta=0.95,
                rho=0.98, decision_interval=10, patience=2,
                ramp_steps=20, gains=(0.05, 0.0005)
            )
        else:
            opt = torch.optim.SGD(model.parameters(), lr=0.04)
        losses = []
        for _ in range(steps):
            opt.zero_grad()
            loss = ((model(x) - y) ** 2).mean()
            loss.backward()
            opt.step()
            losses.append(float(loss))
        return {
            'final_loss': losses[-1], 'loss_40': losses[39],
            'loss_100': losses[99],
            'activations': getattr(opt, 'activations', [])
        }

    return {'sgd': train(False), 'adaptive': train(True)}


if __name__ == '__main__':
    out = {'quadratic_verification': verify(), 'torch_smoke': torch_smoke()}
    with open('results.json', 'w') as f:
        json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))