from __future__ import annotations

import torch
import torch.nn as nn


"""用于数值试井双对数曲线预测的正演代理模型。

模型将变换后的物理参数和编码后的流量制度映射为 log_pressure /
log_derivative 响应。它不是反演模型：PSO 最终的 pbest/gbest 仍由
数值求解器的真实评价结果更新。
"""


def build_mlp(
    in_dim: int,
    hidden_dims: list[int],
    out_dim: int,
    dropout: float = 0.0,
) -> nn.Sequential:
    layers: list[nn.Module] = []
    prev = in_dim
    for h in hidden_dims:
        layers.append(nn.Linear(prev, h))
        layers.append(nn.ReLU())
        if dropout > 0:
            layers.append(nn.Dropout(dropout))
        prev = h
    layers.append(nn.Linear(prev, out_dim))
    return nn.Sequential(*layers)


class ScheduleEncoder(nn.Module):
    """编码固定长度的流量制度向量。"""

    def __init__(self, schedule_dim: int, hidden_dim: int, dropout: float = 0.0):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(schedule_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout) if dropout > 0 else nn.Identity(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
        )

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


class ParamEncoder(nn.Module):
    """编码经过变换的物理参数，例如 log10(k) 和 asinh(skin)。"""

    def __init__(self, param_dim: int, hidden_dim: int, dropout: float = 0.0):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(param_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout) if dropout > 0 else nn.Identity(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
        )

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


class ForwardSurrogate(nn.Module):
    """双分支正演代理模型，并分别设置压力和导数输出头。"""

    def __init__(
        self,
        param_dim: int,
        schedule_dim: int,
        curve_dim: int,
        hidden_dim: int = 128,
        fusion_hidden_dims: list[int] | None = None,
        dropout: float = 0.0,
        use_schedule: bool = True,
    ):
        super().__init__()

        if curve_dim % 3 != 0:
            raise ValueError(f"curve_dim={curve_dim} 不能被 3 整除；期望为 pressure/derivative/slope 三段")

        if fusion_hidden_dims is None:
            fusion_hidden_dims = [256, 256]

        self.curve_dim = curve_dim
        self.part_dim = curve_dim // 3
        self.use_schedule = bool(use_schedule)

        # 参数和流量制度的物理含义与尺度差异较大，因此采用两个分支分别编码。
        self.param_encoder = ParamEncoder(param_dim, hidden_dim, dropout=dropout)

        if self.use_schedule:
            self.schedule_encoder = ScheduleEncoder(schedule_dim, hidden_dim, dropout=dropout)
            trunk_in_dim = hidden_dim * 2
        else:
            self.schedule_encoder = None
            trunk_in_dim = hidden_dim

        trunk_out_dim = fusion_hidden_dims[-1]
        self.trunk = build_mlp(
            in_dim=trunk_in_dim,
            hidden_dims=fusion_hidden_dims,
            out_dim=trunk_out_dim,
            dropout=dropout,
        )

        # 压力曲线拆成 level + centered shape：
        # level 学习整体纵向偏移，shape 学习局部曲线形态。
        self.pressure_level_head = build_mlp(
            in_dim=trunk_out_dim,
            hidden_dims=[128],
            out_dim=1,
            dropout=dropout,
        )
        self.pressure_shape_head = build_mlp(
            in_dim=trunk_out_dim,
            hidden_dims=[128],
            out_dim=self.part_dim,
            dropout=dropout,
        )

        # 导数曲线同样拆分为 level + shape，因为平台、谷值和过渡段
        # 对自动拟合筛选非常重要。
        self.derivative_level_head = build_mlp(
            in_dim=trunk_out_dim,
            hidden_dims=[128],
            out_dim=1,
            dropout=dropout,
        )
        self.derivative_shape_head = build_mlp(
            in_dim=trunk_out_dim,
            hidden_dims=[128],
            out_dim=self.part_dim,
            dropout=dropout,
        )

        # slope 是辅助输出，主要用于保持数据布局兼容。
        self.slope_head = build_mlp(
            in_dim=trunk_out_dim,
            hidden_dims=[128],
            out_dim=self.part_dim,
            dropout=dropout,
        )

    @staticmethod
    def center_shape(x: torch.Tensor) -> torch.Tensor:
        """去除每个样本的均值，避免 shape 分支重复学习整体 level。"""
        return x - x.mean(dim=1, keepdim=True)

    def forward(self, params_x: torch.Tensor, schedule_x: torch.Tensor | None = None) -> torch.Tensor:
        p = self.param_encoder(params_x)

        if self.use_schedule:
            if schedule_x is None:
                raise ValueError("use_schedule=True，但 forward 没有传入 schedule_x")
            s = self.schedule_encoder(schedule_x)
            fused = torch.cat([p, s], dim=-1)
        else:
            fused = p

        trunk_feat = self.trunk(fused)

        pressure_level = self.pressure_level_head(trunk_feat)      # [B, 1]
        pressure_shape = self.pressure_shape_head(trunk_feat)      # [B, T]
        pressure_shape = self.center_shape(pressure_shape)
        pressure_pred = pressure_level + pressure_shape

        derivative_level = self.derivative_level_head(trunk_feat)  # [B, 1]
        derivative_shape = self.derivative_shape_head(trunk_feat)  # [B, T]
        derivative_shape = self.center_shape(derivative_shape)
        derivative_pred = derivative_level + derivative_shape

        slope_pred = self.slope_head(trunk_feat)                   # [B, T]

        curve_pred = torch.cat([pressure_pred, derivative_pred, slope_pred], dim=1)
        return curve_pred