import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
from engine import prepare_data
from model import HeatPINN, heat_pinn_loss


def test_data_to_model_forward():
    """prepare_data → forward: shape e device coerenti, nessun NaN."""
    data = prepare_data(N_f=100, N_ic=50, N_bc=50)
    model = HeatPINN().to(data['device'])
    xt = torch.stack([data['x_f'], data['t_f']], dim=1)
    out = model(xt)
    assert out.shape == (data['x_f'].shape[0], 1)
    assert out.device.type == data['device'].type
    assert torch.isfinite(out).all()


def test_full_loss_pipeline():
    """prepare_data → heat_pinn_loss: tutti i componenti finiti e non-negativi."""
    data = prepare_data(N_f=100, N_ic=50, N_bc=50)
    model = HeatPINN().to(data['device'])
    total, L_pde, L_ic, L_bc = heat_pinn_loss(
        model, data['x_f'], data['t_f'], data['x_ic'], data['t_bc']
    )
    for name, v in [('total', total), ('L_pde', L_pde), ('L_ic', L_ic), ('L_bc', L_bc)]:
        assert torch.isfinite(v), f"{name} non è finita"
        assert v.item() >= 0.0, f"{name} è negativa"


def test_backward_gradients_finite():
    """Il backward della loss non produce NaN/Inf nei parametri."""
    data = prepare_data(N_f=100, N_ic=50, N_bc=50)
    model = HeatPINN().to(data['device'])
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    optimizer.zero_grad()
    total, _, _, _ = heat_pinn_loss(
        model, data['x_f'], data['t_f'], data['x_ic'], data['t_bc']
    )
    total.backward()
    for p in model.parameters():
        assert p.grad is not None
        assert torch.isfinite(p.grad).all(), "gradiente NaN/Inf"


def test_training_loop_stable():
    """20 step di Adam non producono NaN/Inf nei parametri né nella loss."""
    data = prepare_data(N_f=200, N_ic=100, N_bc=100)
    model = HeatPINN().to(data['device'])
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    args = (model, data['x_f'], data['t_f'], data['x_ic'], data['t_bc'])

    for _ in range(20):
        optimizer.zero_grad()
        loss, _, _, _ = heat_pinn_loss(*args)
        loss.backward()
        optimizer.step()

    for p in model.parameters():
        assert torch.isfinite(p).all(), "parametro NaN/Inf dopo training"

    total, _, _, _ = heat_pinn_loss(*args)
    assert torch.isfinite(total)