import torch
from random_attracting_rnn import GeometricallyAttractingRandomRNN

device = 'cuda' if torch.cuda.is_available() else 'cpu'
try:
    torch.manual_seed(2803)
    m = GeometricallyAttractingRandomRNN(5, 7, candidates=3, target_rho=.9).to(device)
    x = torch.randn(11, 4, 5, device=device)
    states, routes = m(x, stochastic=False)
    loss = states.square().mean() + m.contraction_penalty(routes.mean(0))
    loss.backward()
    assert states.shape == (11, 4, 7)
    assert routes.shape == (11, 4, 3)
    assert torch.isfinite(loss) and m.W.grad is not None
    # Make the response of the regularizer explicit.
    with torch.no_grad():
        m.W.mul_(.5)
    low = float(m.contraction_penalty(torch.full((4, 3), 1/3, device=device)))
    with torch.no_grad():
        m.W.mul_(3.)
    high = float(m.contraction_penalty(torch.full((4, 3), 1/3, device=device)))
    print({'device': device, 'loss': float(loss), 'gain_norms': m.gains().detach().cpu().tolist(),
           'low_gain_penalty': low, 'high_gain_penalty': high})
except Exception as e:
    if device == 'cuda':
        print('CUDA failed, retrying CPU:', repr(e))
        m = GeometricallyAttractingRandomRNN(5, 7, candidates=3)
        x = torch.randn(11, 4, 5)
        states, routes = m(x, stochastic=False)
        print({'device':'cpu', 'states':tuple(states.shape), 'routes':tuple(routes.shape)})
    else:
        raise