import torch
import torch.nn as nn
import config


class HeatPINN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, 128), nn.Tanh(),
            nn.Linear(128, 128), nn.Tanh(),
            nn.Linear(128, 128), nn.Tanh(),
            nn.Linear(128, 128), nn.Tanh(),
            nn.Linear(128, 1),
        )

    def forward(self, x):
        # Output scaled to physical range: T_AMB + (Q*L/K) * net
        # net learns dimensionless perturbation in [0,1] range
        T_scale = config.T_AMB + (config.Q_VAL * config.L / config.K) * self.net(x)
        return T_scale


def heat_pinn_loss(model, x_f, t_f, x_ic, t_bc, w_pde=1.0, w_ic=1.0, w_bc=10.0):
    # Characteristic scales for normalization
    T_char = config.Q_VAL * config.L / config.K      # ~50 °C — temperature scale
    grad_char = (config.Q_VAL / config.K) ** 2        # ~2500 — gradient scale²

    # PDE residual: dT/dt - alpha * d2T/dx2 - source(x,t) = 0  (normalized by T_char/t_char)
    x_f = x_f.detach().requires_grad_(True)
    t_f = t_f.detach().requires_grad_(True)
    T_f = model(torch.stack([x_f, t_f], dim=1))
    dT_dt = torch.autograd.grad(T_f.sum(), t_f, create_graph=True)[0]
    dT_dx = torch.autograd.grad(T_f.sum(), x_f, create_graph=True)[0]
    d2T_dx2 = torch.autograd.grad(dT_dx.sum(), x_f, create_graph=True)[0]
    sigma = 0.02
    Q_t_f = torch.where(t_f >= config.T_STEP,
                        torch.tensor(config.Q_VAL, device=t_f.device, dtype=t_f.dtype),
                        torch.tensor(0.0,          device=t_f.device, dtype=t_f.dtype))
    gauss = torch.exp(-0.5 * ((x_f - config.X_SRC) / sigma) ** 2) / (sigma * (2 * torch.pi) ** 0.5)
    source_term = (config.ALPHA / config.K) * Q_t_f * gauss
    pde_scale = (T_char / config.T_END) ** 2 + 1e-8
    L_pde = ((dT_dt - config.ALPHA * d2T_dx2 - source_term) ** 2).mean() / pde_scale

    # IC: T(x, 0) = T0 — normalized by T_char²
    T_ic_pred = model(torch.stack([x_ic, torch.zeros_like(x_ic)], dim=1))
    T_ic_true = torch.full_like(T_ic_pred, config.T0)
    L_ic = ((T_ic_pred - T_ic_true) ** 2).mean() / (T_char ** 2 + 1e-8)

    # BC x=0: Robin — dT/dx + H_CONV/K * (T(0,t) - T_AMB) = 0
    x_left = torch.zeros(t_bc.shape[0], device=t_bc.device).requires_grad_(True)
    T_left = model(torch.stack([x_left, t_bc.detach()], dim=1))
    dT_dx_left = torch.autograd.grad(T_left.sum(), x_left, create_graph=True)[0]
    L_bc_left = ((dT_dx_left + (config.H_CONV / config.K) * (T_left.squeeze() - config.T_AMB)) ** 2).mean() / grad_char

    # BC x=L: Robin — dT/dx + H_CONV/K * (T(L,t) - T_AMB) = 0
    x_right = torch.full((t_bc.shape[0],), config.L, device=t_bc.device).requires_grad_(True)
    T_right = model(torch.stack([x_right, t_bc.detach()], dim=1))
    dT_dx_right = torch.autograd.grad(T_right.sum(), x_right, create_graph=True)[0]
    L_bc_right = ((dT_dx_right + (config.H_CONV / config.K) * (T_right.squeeze() - config.T_AMB)) ** 2).mean() / grad_char

    L_bc = L_bc_left + L_bc_right
    total = w_pde * L_pde + w_ic * L_ic + w_bc * L_bc
    return total, L_pde, L_ic, L_bc