import torch
from torch import nn


class StableNNPhys(nn.Module):
    def __init__(self, hidden_size=32, input_size=71, output_size=68):
        super().__init__()
        self.hidden_size = hidden_size
        self.input_size = input_size
        self.output_size = output_size
        self.hidden = nn.Linear(input_size, hidden_size)
        self.output = nn.Linear(hidden_size, output_size)
        self.bypass = nn.Linear(input_size, output_size)
        nn.init.zeros_(self.output.weight)
        nn.init.zeros_(self.output.bias)
        nn.init.zeros_(self.bypass.weight)
        nn.init.zeros_(self.bypass.bias)

    def forward(self, x):
        return self.output(torch.relu(self.hidden(x))) + self.bypass(x)


def rollout(model, initial_state, surface, advection, state_mean, state_std,
            tendency_mean, tendency_std, dt_seconds=10800.0):
    """Integrate advection trapezoidally, then neural physics with Euler."""
    states = [initial_state]
    physics = []
    state = initial_state
    for step in range(surface.shape[1]):
        adv_now = advection[:, step]
        adv_next = advection[:, min(step + 1, advection.shape[1] - 1)]
        forced = state + 0.5 * dt_seconds * (adv_now + adv_next)
        surface_scaled = surface[:, step] / surface.new_tensor([100.0, 100.0, 1000.0])
        features = torch.cat(((forced - state_mean) / state_std, surface_scaled), dim=-1)
        tendency = model(features) * tendency_std + tendency_mean
        state = forced + dt_seconds * tendency
        physics.append(tendency)
        states.append(state)
    return torch.stack(states, dim=1), torch.stack(physics, dim=1)


def rollout_loss(prediction, target, layer_mass, mode="paper"):
    error = torch.abs(prediction[:, 1:] - target[:, 1:])
    if mode == "paper":
        weights = torch.cat((layer_mass, layer_mass), dim=-1)
        weights = weights / weights.mean(dim=-1, keepdim=True)
        return (error * weights[:, None]).mean()
    if mode == "official_v0_3":
        scale = target[:, 1:].std(dim=(0, 1), unbiased=False).clamp_min(1e-6)
        return (error / scale).mean()
    raise ValueError(f"Unknown loss mode: {mode}")