"""ClimODE's official global neural transport model.

The equations and layer layout follow Aalto-QuML/ClimODE.  The small wrapper
at the bottom adds configuration-friendly construction and checkpoint loading;
the core ``ClimateEncoderFreeUncertain`` class intentionally keeps the
official state and tensor layout.
"""

from __future__ import annotations

import importlib
from pathlib import Path
import sys
from typing import Any, Sequence

import torch
import torch.nn as nn
import torch.nn.functional as F

try:
    from torchdiffeq import odeint as _torchdiffeq_odeint
except ImportError:  # pragma: no cover - exercised when optional dependency is absent
    _torchdiffeq_odeint = None


def _euler_odeint(func, y0: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """Equivalent fixed-step Euler solver used only when torchdiffeq is absent."""

    states = [y0]
    for index in range(1, len(t)):
        previous = states[-1]
        dt = t[index] - t[index - 1]
        states.append(previous + dt * func(t[index - 1], previous))
    return torch.stack(states, dim=0)


def odeint(func, y0: torch.Tensor, t: torch.Tensor, method: str = "euler", **kwargs):
    if _torchdiffeq_odeint is not None:
        return _torchdiffeq_odeint(func, y0, t, method=method, **kwargs)
    if method != "euler":
        raise ImportError(
            "torchdiffeq is required for solver=%r; install the official dependency."
            % method
        )
    return _euler_odeint(func, y0, t)


class OptimVelocity(nn.Module):
    """Learn the initial per-channel velocity used to start the ODE system."""

    def __init__(self, num_years: int, height: int, width: int, out_channels: int = 5):
        super().__init__()
        self.out_channels = out_channels
        self.v_x = nn.Parameter(
            torch.randn(num_years, 1, out_channels, height, width)
        )
        self.v_y = nn.Parameter(
            torch.randn(num_years, 1, out_channels, height, width)
        )

    def forward(self, data: torch.Tensor):
        u_y = torch.gradient(data, dim=3)[0]
        u_x = torch.gradient(data, dim=4)[0]
        divergence = torch.gradient(self.v_y, dim=3)[0] + torch.gradient(
            self.v_x, dim=4
        )[0]
        adv = self.v_x * u_x + self.v_y * u_y + data * divergence
        return adv, self.v_x, self.v_y


class BoundaryPad(nn.Module):
    """Reflect at the poles and wrap around the longitude seam."""

    def forward(self, value: torch.Tensor) -> torch.Tensor:
        return F.pad(F.pad(value, (0, 0, 1, 1), "reflect"), (1, 1, 0, 0), "circular")


class ResidualBlock(nn.Module):
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()
        self.activation = nn.LeakyReLU(0.3)
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=0)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=0)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.drop = nn.Dropout(p=0.1)
        self.shortcut = (
            nn.Conv2d(in_channels, out_channels, kernel_size=1)
            if in_channels != out_channels
            else nn.Identity()
        )

    def forward(self, value: torch.Tensor) -> torch.Tensor:
        value_padded = F.pad(
            F.pad(value, (0, 0, 1, 1), "reflect"), (1, 1, 0, 0), "circular"
        )
        hidden = self.activation(self.bn1(self.conv1(value_padded)))
        hidden = F.pad(
            F.pad(hidden, (0, 0, 1, 1), "reflect"), (1, 1, 0, 0), "circular"
        )
        hidden = self.activation(self.bn2(self.conv2(hidden)))
        return self.drop(hidden) + self.shortcut(value)


class ClimateResNet2D(nn.Module):
    def __init__(
        self,
        num_channels: int,
        layers: Sequence[int],
        hidden_size: Sequence[int],
    ):
        super().__init__()
        if len(layers) != len(hidden_size):
            raise ValueError("layers and hidden_size must have equal lengths")
        self.layer_cnn = nn.ModuleList(
            [nn.Sequential(*blocks_for_layer) for blocks_for_layer in self._split_blocks(layers, hidden_size, num_channels)]
        )

    @staticmethod
    def _split_blocks(
        layers: Sequence[int], hidden_size: Sequence[int], num_channels: int
    ) -> list[list[nn.Module]]:
        groups: list[list[nn.Module]] = []
        in_channels = num_channels
        for repetitions, out_channels in zip(layers, hidden_size):
            group = [ResidualBlock(in_channels, out_channels)]
            group.extend(
                ResidualBlock(out_channels, out_channels)
                for _ in range(1, int(repetitions))
            )
            groups.append(group)
            in_channels = out_channels
        return groups

    def forward(self, value: torch.Tensor) -> torch.Tensor:
        output = value.float()
        for layer in self.layer_cnn:
            output = layer(output)
        return output


class SelfAttnConv(nn.Module):
    """Official key-query-value attention convolution."""

    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()
        self.query = self._conv(in_channels, in_channels // 8, stride=1)
        self.key = self._key_conv(in_channels, in_channels // 8, stride=2)
        self.value = self._key_conv(in_channels, out_channels, stride=2)
        self.post_map = nn.Sequential(
            nn.Conv2d(out_channels, out_channels, kernel_size=1, stride=1, padding=0)
        )
        self.out_ch = out_channels

    @staticmethod
    def _conv(n_in: int, n_out: int, stride: int) -> nn.Sequential:
        return nn.Sequential(
            BoundaryPad(),
            nn.Conv2d(n_in, n_in // 2, kernel_size=3, stride=stride, padding=0),
            nn.LeakyReLU(0.3),
            BoundaryPad(),
            nn.Conv2d(n_in // 2, n_out, kernel_size=3, stride=stride, padding=0),
            nn.LeakyReLU(0.3),
            BoundaryPad(),
            nn.Conv2d(n_out, n_out, kernel_size=3, stride=stride, padding=0),
        )

    @staticmethod
    def _key_conv(n_in: int, n_out: int, stride: int) -> nn.Sequential:
        return nn.Sequential(
            nn.Conv2d(n_in, n_in // 2, kernel_size=3, stride=stride, padding=0),
            nn.LeakyReLU(0.3),
            nn.Conv2d(n_in // 2, n_out, kernel_size=3, stride=stride, padding=0),
            nn.LeakyReLU(0.3),
            nn.Conv2d(n_out, n_out, kernel_size=3, stride=1, padding=0),
        )

    def forward(self, value: torch.Tensor) -> torch.Tensor:
        size = value.size()
        value = value.float()
        query = self.query(value).flatten(-2, -1)
        key = self.key(value).flatten(-2, -1)
        val = self.value(value).flatten(-2, -1)
        beta = F.softmax(torch.bmm(query.transpose(1, 2), key), dim=1)
        output = torch.bmm(val, beta.transpose(1, 2))
        output = output.view(-1, self.out_ch, size[-2], size[-1]).contiguous()
        return self.post_map(output)


class ClimateEncoderFreeUncertain(nn.Module):
    """Official global ``Climate_encoder_free_uncertain`` implementation."""

    def __init__(
        self,
        num_channels: int = 5,
        const_channels: int = 2,
        out_types: int = 5,
        method: str = "euler",
        use_att: bool = True,
        use_err: bool = True,
        use_pos: bool = False,
    ):
        super().__init__()
        self.layers = [5, 3, 2]
        self.hidden = [128, 64, 2 * out_types]
        input_channels = 30 + out_types * int(use_pos) + 34 * (1 - int(use_pos))
        self.vel_f = ClimateResNet2D(input_channels, self.layers, self.hidden)
        if use_att:
            self.vel_att = SelfAttnConv(input_channels, 10)
            self.gamma = nn.Parameter(torch.tensor([0.1]))

        self.scales = num_channels
        self.const_channel = const_channels
        self.out_ch = out_types
        self.past_samples: torch.Tensor | int = 0
        self.const_info: torch.Tensor | int = 0
        self.lat_map: torch.Tensor | int = 0
        self.lon_map: torch.Tensor | int = 0
        self.method = method
        err_in = 9 + out_types * int(use_pos) + 34 * (1 - int(use_pos))
        if use_err:
            self.noise_net = ClimateResNet2D(err_in, [3, 2, 2], [128, 64, 2 * out_types])
        if use_pos:
            self.pos_enc = ClimateResNet2D(4, [2, 1, 1], [32, 16, out_types])
        self.att = use_att
        self.err = use_err
        self.pos = use_pos
        self.pos_feat: torch.Tensor | int = 0
        self.lsm: torch.Tensor | int = 0
        self.oro: torch.Tensor | int = 0

    def update_param(self, params: Sequence[torch.Tensor]) -> None:
        if len(params) != 4:
            raise ValueError("update_param expects past_samples, constants, latitude, longitude")
        self.past_samples, self.const_info, self.lat_map, self.lon_map = params

    def _time_features(
        self, t: torch.Tensor, height: int, width: int, batch_size: int
    ) -> tuple[torch.Tensor, ...]:
        t_emb = ((t * 100) % 24).view(1, 1, 1, 1).expand(batch_size, 1, height, width)
        sin_t = torch.sin(torch.pi * t_emb / 12 - torch.pi / 2)
        cos_t = torch.cos(torch.pi * t_emb / 12 - torch.pi / 2)
        sin_season = torch.sin(torch.pi * t_emb / (12 * 365) - torch.pi / 2)
        cos_season = torch.cos(torch.pi * t_emb / (12 * 365) - torch.pi / 2)
        return t_emb, torch.cat([sin_t, cos_t], dim=1), torch.cat(
            [sin_season, cos_season], dim=1
        )

    def pde(self, t: torch.Tensor, state: torch.Tensor) -> torch.Tensor:
        height, width = state.shape[-2:]
        ds = state[:, -self.out_ch :].view(-1, self.out_ch, height, width).float()
        velocity = state[:, : 2 * self.out_ch].view(
            -1, 2 * self.out_ch, height, width
        ).float()
        t_emb, day_emb, season_emb = self._time_features(
            t, height, width, ds.shape[0]
        )
        ds_grad_x = torch.gradient(ds, dim=3)[0]
        ds_grad_y = torch.gradient(ds, dim=2)[0]
        nabla_u = torch.cat([ds_grad_x, ds_grad_y], dim=1)

        if self.pos:
            combined = torch.cat(
                [t_emb / 24, day_emb, season_emb, nabla_u, velocity, ds, self.pos_feat],
                dim=1,
            )
        else:
            cos_lat, sin_lat = torch.cos(self.new_lat_map), torch.sin(self.new_lat_map)
            cos_lon, sin_lon = torch.cos(self.new_lon_map), torch.sin(self.new_lon_map)
            time_cyclic = torch.cat([day_emb, season_emb], dim=1)
            pos_feats = torch.cat(
                [
                    cos_lat,
                    cos_lon,
                    sin_lat,
                    sin_lon,
                    sin_lat * cos_lon,
                    sin_lat * sin_lon,
                ],
                dim=1,
            )
            pos_time = self.get_time_pos_embedding(time_cyclic, pos_feats)
            combined = torch.cat(
                [
                    t_emb / 24,
                    day_emb,
                    season_emb,
                    nabla_u,
                    velocity,
                    ds,
                    self.new_lat_map,
                    self.new_lon_map,
                    self.lsm,
                    self.oro,
                    pos_feats,
                    pos_time,
                ],
                dim=1,
            )

        dv = self.vel_f(combined)
        if self.att:
            dv = dv + self.gamma * self.vel_att(combined)
        v_x = velocity[:, : self.out_ch]
        v_y = velocity[:, self.out_ch :]
        advection = v_x * ds_grad_x + v_y * ds_grad_y
        advection = advection + ds * (
            torch.gradient(v_x, dim=3)[0] + torch.gradient(v_y, dim=2)[0]
        )
        return torch.cat([dv, advection], dim=1)

    @staticmethod
    def get_time_pos_embedding(
        time_features: torch.Tensor, position_features: torch.Tensor
    ) -> torch.Tensor:
        outputs = [feature.unsqueeze(1) * position_features for feature in time_features.unbind(1)]
        return torch.cat(outputs, dim=1)

    def noise_net_contrib(
        self,
        time: torch.Tensor,
        pos_enc: torch.Tensor,
        s_final: torch.Tensor,
        height: int,
        width: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        t_emb = (time % 24).view(-1, 1, 1, 1, 1)
        sin_t = torch.sin(torch.pi * t_emb / 12 - torch.pi / 2).expand(
            len(s_final), s_final.shape[1], 1, height, width
        )
        cos_t = torch.cos(torch.pi * t_emb / 12 - torch.pi / 2).expand(
            len(s_final), s_final.shape[1], 1, height, width
        )
        sin_season = torch.sin(torch.pi * t_emb / (12 * 365) - torch.pi / 2).expand(
            len(s_final), s_final.shape[1], 1, height, width
        )
        cos_season = torch.cos(torch.pi * t_emb / (12 * 365) - torch.pi / 2).expand(
            len(s_final), s_final.shape[1], 1, height, width
        )
        pos_rep = pos_enc.expand(len(s_final), s_final.shape[1], -1, height, width)
        pos_rep = pos_rep.flatten(start_dim=0, end_dim=1)
        time_cyclic = torch.cat([sin_t, cos_t, sin_season, cos_season], dim=2)
        time_cyclic = time_cyclic.flatten(start_dim=0, end_dim=1)
        pos_time = self.get_time_pos_embedding(time_cyclic, pos_rep[:, 2:-2])
        combined = torch.cat(
            [time_cyclic, s_final.flatten(start_dim=0, end_dim=1), pos_rep, pos_time],
            dim=1,
        )
        final_out = self.noise_net(combined).view(
            len(time), -1, 2 * self.out_ch, height, width
        )
        mean = s_final + final_out[:, :, : self.out_ch]
        std = F.softplus(final_out[:, :, self.out_ch :])
        return mean, std

    def forward(
        self,
        time_steps: torch.Tensor,
        data: torch.Tensor,
        atol: float = 0.1,
        rtol: float = 0.1,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if not isinstance(self.past_samples, torch.Tensor):
            raise RuntimeError("Call update_param before forward")
        height, width = self.past_samples.shape[-2:]
        values = data.float().view(-1, self.out_ch, height, width)
        final_data = torch.cat([self.past_samples, values], dim=1)
        init_time = time_steps[0].item() * 6
        final_time = time_steps[-1].item() * 6
        steps_val = final_time - init_time

        if self.pos:
            lat_map = self.lat_map.unsqueeze(0) * torch.pi / 180
            lon_map = self.lon_map.unsqueeze(0) * torch.pi / 180
            pos_rep = torch.cat([lat_map.unsqueeze(0), lon_map.unsqueeze(0), self.const_info], dim=1)
            self.pos_feat = self.pos_enc(pos_rep).expand(
                values.shape[0], -1, values.shape[-2], values.shape[-1]
            )
            final_pos_enc = self.pos_feat
        else:
            self.oro = self.const_info[0, 0]
            self.lsm = self.const_info[0, 1]
            self.lsm = self.lsm.unsqueeze(0).expand(values.shape[0], -1, height, width)
            self.oro = F.normalize(self.oro).unsqueeze(0).expand(values.shape[0], -1, height, width)
            self.new_lat_map = self.lat_map.expand(values.shape[0], 1, height, width) * torch.pi / 180
            self.new_lon_map = self.lon_map.expand(values.shape[0], 1, height, width) * torch.pi / 180
            cos_lat, sin_lat = torch.cos(self.new_lat_map), torch.sin(self.new_lat_map)
            cos_lon, sin_lon = torch.cos(self.new_lon_map), torch.sin(self.new_lon_map)
            pos_feats = torch.cat(
                [cos_lat, cos_lon, sin_lat, sin_lon, sin_lat * cos_lon, sin_lat * sin_lon],
                dim=1,
            )
            final_pos_enc = torch.cat(
                [self.new_lat_map, self.new_lon_map, pos_feats, self.lsm, self.oro], dim=1
            )

        integration_steps = max(int(steps_val) + 1, 1)
        new_time_steps = torch.linspace(
            init_time, final_time, steps=integration_steps, device=values.device
        )
        ode_time = 0.01 * new_time_steps.float()
        final_result = odeint(
            self.pde,
            final_data,
            ode_time,
            method=self.method,
            atol=atol,
            rtol=rtol,
        )
        s_final = final_result[:, :, -self.out_ch :].view(
            len(ode_time), -1, self.out_ch, height, width
        )
        sampled = s_final[0 : len(s_final) : 6]
        if self.err:
            mean, std = self.noise_net_contrib(
                time_steps, final_pos_enc, sampled, height, width
            )
            return mean, std, sampled
        return sampled, torch.zeros_like(sampled), sampled


class ClimODE(ClimateEncoderFreeUncertain):
    """Configuration-friendly public model name."""

    def __init__(
        self,
        num_channels: int = 5,
        const_channels: int = 2,
        out_types: int = 5,
        method: str = "euler",
        use_attention: bool = True,
        use_uncertainty: bool = True,
        use_positional_encoder: bool = False,
    ):
        super().__init__(
            num_channels=num_channels,
            const_channels=const_channels,
            out_types=out_types,
            method=method,
            use_att=use_attention,
            use_err=use_uncertainty,
            use_pos=use_positional_encoder,
        )


def _register_checkpoint_compat_modules() -> None:
    """Expose legacy top-level module names used by official pickle files."""

    for legacy_name, current_name in (
        ("model_function", "model.model_function"),
        ("model_utils", "model.model_utils"),
    ):
        module = importlib.import_module(current_name)
        sys.modules.setdefault(legacy_name, module)


def load_checkpoint(path: str | Path, map_location: str | torch.device = "cpu") -> nn.Module:
    """Load an official full-object checkpoint or a state-dict checkpoint."""

    checkpoint_path = Path(path)
    if not checkpoint_path.is_file():
        raise FileNotFoundError(checkpoint_path)
    _register_checkpoint_compat_modules()
    try:
        checkpoint = torch.load(checkpoint_path, map_location=map_location, weights_only=False)
    except TypeError:  # Older PyTorch does not expose weights_only.
        checkpoint = torch.load(checkpoint_path, map_location=map_location)
    if isinstance(checkpoint, nn.Module):
        return checkpoint
    model = ClimODE()
    if isinstance(checkpoint, dict):
        state_dict = checkpoint.get("state_dict", checkpoint.get("model", checkpoint))
    else:
        state_dict = checkpoint
    model.load_state_dict(state_dict)
    return model


# Names kept for the official checkpoint's pickle module/class references.
Climate_encoder_free_uncertain = ClimateEncoderFreeUncertain
Climate_ResNet_2D = ClimateResNet2D
Self_attn_conv = SelfAttnConv
boundarypad = BoundaryPad