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

import torch
import config
from model import HeatPINN, heat_pinn_loss


# ── HeatPINN.forward ─────────────────────────────────────────────────────────

def test_forward_output_shape(pinn_model, device):
    xt = torch.zeros(64, 2, device=device)
    xt[:, 0] = torch.rand(64) * config.L
    xt[:, 1] = torch.rand(64) * config.T_END
    assert pinn_model(xt).shape == (64, 1)


def test_forward_finite(pinn_model, device):
    xt = torch.zeros(100, 2, device=device)
    xt[:, 0] = torch.rand(100) * config.L
    xt[:, 1] = torch.rand(100) * config.T_END
    assert torch.isfinite(pinn_model(xt)).all()


def test_forward_zero_weights_returns_t_amb(device):
    """Con pesi nulli net(x,t)=0 ⇒ forward restituisce T_AMB per ogni input."""
    model = HeatPINN().to(device)
    for p in model.parameters():
        p.data.zero_()
    xt = torch.zeros(20, 2, device=device)
    xt[:, 0] = torch.rand(20) * config.L
    xt[:, 1] = torch.rand(20) * config.T_END
    out = model(xt)
    torch.testing.assert_close(out, torch.full_like(out, config.T_AMB), atol=1e-5, rtol=0.0)


def test_forward_t_normalization(device):
    """t viene normalizzato a [0,1]: il modello deve restituire output finiti
    anche a t=T_END senza saturazione di Tanh."""
    model = HeatPINN().to(device)
    torch.nn.init.xavier_uniform_(model.net[0].weight)
    xt = torch.tensor([[0.5, 0.0], [0.5, config.T_END]], device=device)
    out = model(xt)
    assert out.shape == (2, 1)
    assert torch.isfinite(out).all()


# ── heat_pinn_loss ────────────────────────────────────────────────────────────

def _dummy_inputs(device, n_f=100, n_ic=50, n_bc=50):
    x_f  = torch.rand(n_f,  device=device) * config.L
    t_f  = torch.rand(n_f,  device=device) * config.T_END
    x_ic = torch.rand(n_ic, device=device) * config.L
    t_bc = torch.rand(n_bc, device=device) * config.T_END
    return x_f, t_f, x_ic, t_bc


def test_loss_returns_four_values(pinn_model, device):
    result = heat_pinn_loss(pinn_model, *_dummy_inputs(device))
    assert len(result) == 4


def test_loss_components_non_negative(pinn_model, device):
    total, L_pde, L_ic, L_bc = heat_pinn_loss(pinn_model, *_dummy_inputs(device))
    assert total.item() >= 0.0
    assert L_pde.item() >= 0.0
    assert L_ic.item() >= 0.0
    assert L_bc.item() >= 0.0


def test_loss_finite(pinn_model, device):
    for v in heat_pinn_loss(pinn_model, *_dummy_inputs(device)):
        assert torch.isfinite(v), f"loss non finita: {v}"


def test_loss_weight_doubles_pde_contribution(pinn_model, device):
    """Raddoppiare w_pde con w_ic=w_bc=0 deve raddoppiare il totale."""
    inputs = _dummy_inputs(device)
    total1, L_pde1, _, _ = heat_pinn_loss(pinn_model, *inputs, w_pde=1.0, w_ic=0.0, w_bc=0.0)
    total2, L_pde2, _, _ = heat_pinn_loss(pinn_model, *inputs, w_pde=2.0, w_ic=0.0, w_bc=0.0)
    # L_pde deve essere identico tra le due chiamate (stesso modello, stessi dati)
    torch.testing.assert_close(L_pde1, L_pde2, atol=1e-5, rtol=1e-4)
    torch.testing.assert_close(total2, 2.0 * total1, atol=1e-5, rtol=1e-4)


def test_ic_loss_zero_when_net_is_zero(device):
    """Con net=0 ⇒ T = T_AMB = T0 ⇒ L_ic = 0."""
    model = HeatPINN().to(device)
    for p in model.parameters():
        p.data.zero_()
    _, _, L_ic, _ = heat_pinn_loss(model, *_dummy_inputs(device))
    assert L_ic.item() < 1e-8


def test_bc_loss_zero_when_net_is_zero(device):
    """Con net=0 ⇒ T = T_AMB e dT/dx = 0 ⇒ Robin BC soddisfatta ⇒ L_bc = 0."""
    model = HeatPINN().to(device)
    for p in model.parameters():
        p.data.zero_()
    _, _, _, L_bc = heat_pinn_loss(model, *_dummy_inputs(device))
    assert L_bc.item() < 1e-8