from __future__ import annotations

import torch
import torch.nn as nn


class ResidualBlock(nn.Module):
    def __init__(self, dim: int, dropout: float = 0.0):
        super().__init__()
        self.net = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, dim),
            nn.GELU(),
            nn.Dropout(dropout) if dropout > 0 else nn.Identity(),
            nn.Linear(dim, dim),
        )
        self.act = nn.GELU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.act(x + self.net(x))


class TimeConditionedSurrogate(nn.Module):
    """Point-wise forward surrogate.

    f(params, full_schedule, loglog_time_point) -> [log_pressure, log_derivative]
    """

    def __init__(
        self,
        param_dim: int,
        schedule_dim: int,
        time_dim: int,
        hidden_dim: int = 256,
        n_blocks: int = 4,
        dropout: float = 0.05,
        use_schedule: bool = True,
    ):
        super().__init__()
        self.use_schedule = bool(use_schedule)

        self.param_encoder = nn.Sequential(
            nn.Linear(param_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
        )
        self.time_encoder = nn.Sequential(
            nn.Linear(time_dim, hidden_dim // 2),
            nn.LayerNorm(hidden_dim // 2),
            nn.GELU(),
        )
        if self.use_schedule:
            self.schedule_encoder = nn.Sequential(
                nn.Linear(schedule_dim, hidden_dim),
                nn.LayerNorm(hidden_dim),
                nn.GELU(),
            )
            fusion_dim = hidden_dim * 2 + hidden_dim // 2
        else:
            self.schedule_encoder = None
            fusion_dim = hidden_dim + hidden_dim // 2

        self.input_proj = nn.Sequential(
            nn.Linear(fusion_dim, hidden_dim),
            nn.GELU(),
        )
        self.blocks = nn.Sequential(*[ResidualBlock(hidden_dim, dropout=dropout) for _ in range(int(n_blocks))])
        self.head = nn.Sequential(
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.GELU(),
            nn.Linear(hidden_dim // 2, 2),
        )

    def forward(
        self,
        params_x: torch.Tensor,
        time_x: torch.Tensor,
        schedule_x: torch.Tensor | None = None,
    ) -> torch.Tensor:
        p = self.param_encoder(params_x)
        t = self.time_encoder(time_x)
        if self.use_schedule:
            if schedule_x is None:
                raise ValueError("use_schedule=True but schedule_x is None")
            s = self.schedule_encoder(schedule_x)
            x = torch.cat([p, s, t], dim=-1)
        else:
            x = torch.cat([p, t], dim=-1)
        x = self.input_proj(x)
        x = self.blocks(x)
        return self.head(x)