"""
贝叶斯神经网络系列算法
实现 MC Dropout 和 Deep Ensemble 两种不确定性估计方法

理论基础:
- MC Dropout: 通过测试时启用 Dropout 进行 Monte Carlo 采样估计不确定性
- Deep Ensemble: 通过训练多个独立模型的集成来估计不确定性

参考:
- Gal, Y., & Ghahramani, Z. (2016). Dropout as a bayesian approximation: Representing model uncertainty in deep learning.
- Lakshminarayanan, B., Pritzel, A., & Blundell, C. (2017). Simple and scalable predictive uncertainty estimation using deep ensembles.
"""

import typing

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    List,
    Literal,
    Optional,
    Tuple,
    Union,
)

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from zyplib.nn.modules import BaseModule, get_activation
from zyplib.utils.print import warning

if TYPE_CHECKING:
    from zyplib.train.common import FabricTrainingContext


@dataclass(frozen=True)
class BNNOutput:
    """贝叶斯神经网络的不确定性输出

    理论精确的设计：不强行区分不确定性类型，专注于实际可测量的量
    """

    predictive_mean: Tensor  # 预测均值 E[y|x,D]
    predictive_variance: Tensor  # 预测方差 Var[y|x,D]
    prediction_samples: Tensor  # 预测样本 [num_samples, batch_size, ...]


class MCDropoutLayer(nn.Module):
    """Monte Carlo Dropout 层

    可复用的 Dropout 层，在训练和推理时都保持激活状态
    """

    def __init__(self, p: float = 0.5, inplace: bool = False):
        super().__init__()
        if p < 0 or p > 1:
            raise ValueError(f'Dropout 概率必须在 [0, 1] 范围内，得到 {p}')
        self.p = p
        self.inplace = inplace

    def forward(self, x: Tensor) -> Tensor:
        """前向传播，始终应用 Dropout"""
        # 关键：无论 self.training 状态如何，都强制应用 dropout
        return F.dropout(x, p=self.p, training=True, inplace=self.inplace)

    def extra_repr(self) -> str:
        return f'p={self.p}, inplace={self.inplace}'


class MCDropoutNN(BaseModule):
    """Monte Carlo Dropout 神经网络

    通过在推理时保持 Dropout 激活并进行多次前向传播来估计不确定性。
    理论上等价于对深度高斯过程的变分推断。

    Parameters
    ----------
    - `backbone` : `nn.Module`
        - 主干网络
    - `dropout_rate` : `float`, optional
        - Dropout 概率, by default 0.1
    - `auto_mcdp` : `Literal[None, 'insert', 'convert']`, optional
        - MC Dropout 处理模式, by default None
        - None: 不进行任何自动处理，假设模型已包含 MCDropoutLayer
        - 'insert': 在线性层和卷积层后自动插入 MCDropoutLayer
        - 'convert': 将现有的 Dropout 层转换为 MCDropoutLayer
    - `input_shape` : `Tuple[int, ...]`, optional
        - 输入形状，用于模型验证, by default None
    - `validate_model` : `bool`, optional
        - 是否验证模型, by default True

    Examples
    ----------
    ```python
    # 模式1: 使用已包含 MCDropoutLayer 的模型
    backbone = nn.Sequential(
        nn.Linear(784, 128),
        nn.ReLU(),
        MCDropoutLayer(0.1),
        nn.Linear(128, 10)
    )
    mc_model = MCDropoutNN(backbone, auto_mcdp=None)

    # 模式2: 自动插入 MCDropoutLayer
    simple_backbone = nn.Sequential(
        nn.Linear(784, 128),
        nn.ReLU(),
        nn.Linear(128, 10)
    )
    mc_model = MCDropoutNN(simple_backbone, dropout_rate=0.2, auto_mcdp='insert')

    # 模式3: 转换现有 Dropout 为 MCDropoutLayer
    dropout_backbone = nn.Sequential(
        nn.Linear(784, 128),
        nn.Dropout(0.1),
        nn.ReLU(),
        nn.Linear(128, 10),
        nn.Dropout(0.2)
    )
    mc_model = MCDropoutNN(dropout_backbone, auto_mcdp='convert')

    # 使用
    uncertainty_output = mc_model.predict_bnn(x, num_samples=50)
    ```
    """

    BNN_TYPE = 'MCDropoutNN'

    def __init__(
        self,
        backbone: nn.Module,
        dropout_rate: float = 0.1,
        auto_mcdp: Optional[Literal['insert', 'convert']] = None,
        activation: Optional[
            Literal[
                'relu', 'leaky_relu', 'prelu', 'gelu', 'silu', 'tanh', 'sigmoid', 'swish'
            ]
        ]
        | None = None,
    ):
        # 处理 MC Dropout 自动化
        if auto_mcdp == 'insert':
            backbone = self._insert_mc_dropout_layers(backbone, dropout_rate)
        elif auto_mcdp == 'convert':
            converted_count = self._convert_dropout_layers(backbone)
            if converted_count == 0:
                warning('模型中未找到任何 Dropout 层进行转换')

        super().__init__()
        self.backbone = backbone
        self.dropout_rate = dropout_rate
        self.auto_mcdp = auto_mcdp
        self.activation = get_activation(activation, return_type='function') or (
            lambda x: x
        )

    def _insert_mc_dropout_layers(
        self, model: nn.Module, dropout_rate: float
    ) -> nn.Module:
        """在线性层后自动插入 MCDropoutLayer"""
        layers = []

        for name, module in model.named_children():
            layers.append(module)

            # 在线性层和卷积层后插入 Dropout
            if isinstance(module, (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d)):
                layers.append(MCDropoutLayer(p=dropout_rate))

        return nn.Sequential(*layers)

    def _convert_dropout_layers(self, model: nn.Module) -> int:
        """将现有 Dropout 层转换为 MCDropoutLayer"""
        return MCDropoutNN.apply_mcdropout(model)

    @staticmethod
    def apply_mcdropout(model: nn.Module) -> int:
        """静态方法：将模型中的 Dropout 层转换为 MCDropoutLayer

        借鉴案例代码的 apply_mcdropout 实现，提供独立的转换功能

        Parameters
        ----------
        - `model` : `nn.Module`
            - 要转换的模型

        Returns
        ----------
        - `int`
            - 转换的 Dropout 层数量

        Examples
        ----------
        ```python
        # 直接转换现有模型
        model = EEGNet(channels=8, samples=1000, n_class=2)
        converted_count = MCDropoutNN.apply_mcdropout(model)
        print(f"转换了 {converted_count} 个 Dropout 层")

        # 然后可以直接使用转换后的模型
        mc_model = MCDropoutNN(model, auto_mcdp=None)
        ```
        """
        converted_count = 0

        for name, module in model.named_modules():
            if isinstance(module, (nn.Dropout, nn.Dropout1d, nn.Dropout2d, nn.Dropout3d)):
                parent_name = '.'.join(name.split('.')[:-1])
                layer_name = name.split('.')[-1]

                try:
                    # 递归查找父模块
                    parent = model
                    if parent_name:
                        for part in parent_name.split('.'):
                            parent = getattr(parent, part)

                    # 创建 MCDropoutLayer 替换原 Dropout
                    mc_dropout = MCDropoutLayer(
                        p=module.p, inplace=getattr(module, 'inplace', False)
                    )
                    setattr(parent, layer_name, mc_dropout)
                    converted_count += 1

                except AttributeError as e:
                    warning(f"无法转换 Dropout 层 '{name}': {e}")
                    continue

        return converted_count

    def forward(self, x: Tensor) -> Tensor:
        """标准前向传播"""
        logits = self.backbone(x)
        return self.activation(logits)

    def predict_bnn(
        self, x: Tensor, num_samples: int = 100, *, CHUNCK_SIZE=20
    ) -> BNNOutput:
        # self.eval()

        # 批量化处理：复制输入以便并行采样
        batch_size = x.shape[0]

        if num_samples <= CHUNCK_SIZE:  # 小样本数时批量处理
            x_expanded = x.unsqueeze(0).repeat(num_samples, 1, *[1] * (x.dim() - 1))
            x_batch = x_expanded.view(-1, *x.shape[1:])

            with torch.no_grad():
                predictions_batch = self.backbone(x_batch)
                predictions_batch = self.activation(predictions_batch)

            prediction_samples = predictions_batch.view(num_samples, batch_size, -1)
        else:  # 大样本数时分批处理
            samples = []
            chunk_size = min(CHUNCK_SIZE, num_samples)

            for i in range(0, num_samples, chunk_size):
                current_chunk = min(chunk_size, num_samples - i)
                x_chunk = x.unsqueeze(0).repeat(current_chunk, 1, *[1] * (x.dim() - 1))
                x_batch = x_chunk.view(-1, *x.shape[1:])

                with torch.no_grad():
                    chunk_pred = self.backbone(x_batch)
                    chunk_pred = self.activation(chunk_pred)

                samples.append(chunk_pred.view(current_chunk, batch_size, -1))

            prediction_samples = torch.cat(samples, dim=0)

        # 计算统计量
        predictive_mean = torch.mean(prediction_samples, dim=0)
        predictive_variance = torch.var(prediction_samples, dim=0, unbiased=True)

        return BNNOutput(
            predictive_mean=predictive_mean,
            predictive_variance=predictive_variance,
            prediction_samples=prediction_samples,
        )


class DeepEnsemble(BaseModule):
    """Deep Ensemble 不确定性估计

    通过训练多个独立初始化的神经网络并对其预测进行集成来估计不确定性。
    每个模型代表后验分布的一个近似样本。
    """

    BNN_TYPE = 'DeepEnsemble'

    def __init__(
        self,
        models: typing.Sequence[nn.Module],
        activation: Optional[
            Literal[
                'relu', 'leaky_relu', 'prelu', 'gelu', 'silu', 'tanh', 'sigmoid', 'swish'
            ]
        ]
        | None = None,
    ):
        super().__init__()
        # 使用 ModuleList 管理多个模型
        self.backbone = nn.ModuleList(models)

        self.num_models = len(models)
        self.models = self.backbone  # 别名，便于访问
        self.activation = get_activation(activation, return_type='function') or (
            lambda x: x
        )

    def forward(self, x: Tensor) -> Tensor:
        """集成预测（所有模型的平均）"""
        predictions = []

        for model in self.models:
            pred = model(x)
            pred = self.activation(pred)  # 应用激活函数
            predictions.append(pred)

        # 返回平均预测
        ensemble_pred = torch.stack(predictions, dim=0)
        return torch.mean(ensemble_pred, dim=0)

    def predict_bnn(
        self,
        x: Tensor,
        # num_samples: Optional[int] = None,  # Deep Ensemble 不需要额外采样
    ) -> BNNOutput:
        """Deep Ensemble 不确定性估计"""
        predictions = []

        # 获取每个模型的预测
        for model in self.models:
            model.eval()
            with torch.no_grad():
                pred = model(x)
                pred = self.activation(pred)
                predictions.append(pred)

        # 堆叠预测 [num_models, batch_size, ...]
        prediction_samples = torch.stack(predictions, dim=0)

        # 计算统计量
        predictive_mean = torch.mean(prediction_samples, dim=0)
        predictive_variance = torch.var(prediction_samples, dim=0)

        return BNNOutput(
            predictive_mean=predictive_mean,
            predictive_variance=predictive_variance,
            prediction_samples=prediction_samples,
        )

    def get_model(self, index: int) -> nn.Module:
        """获取指定索引的模型"""
        if index >= self.num_models:
            raise IndexError(f'模型索引 {index} 超出范围 [0, {self.num_models-1}]')
        return self.models[index]
