import typing
from typing import Any, Dict, Union, cast

import torch
import torch.optim
from torch import nn

from zyplib.utils.time import now

from .shapes import output_shape

__all__ = [
    'output_shape',
    'weight_initializer',
    'save_checkpoint',
    'load_checkpoint',
    'freeze',
    'unfreeze',
]


def weight_initializer(m: nn.Module):
    """神经网络权重初始化器

    对不同类型的神经网络层应用合适的权重初始化策略。

    Parameters
    ----------
    - `m` : `nn.Module`
        - 需要初始化的网络层模块
    """
    if isinstance(m, (nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.Linear)):
        nn.init.kaiming_normal_(m.weight.data, mode='fan_out', nonlinearity='relu')
        if m.bias is not None:
            nn.init.constant_(m.bias.data, 0)
    elif isinstance(m, (nn.ConvTranspose1d, nn.ConvTranspose2d, nn.ConvTranspose3d)):
        # 转置卷积层使用 Kaiming 初始化，适用于生成模型等场景
        nn.init.kaiming_normal_(m.weight.data, mode='fan_in', nonlinearity='relu')
        if m.bias is not None:
            nn.init.constant_(m.bias.data, 0)
    elif isinstance(
        m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.GroupNorm, nn.LayerNorm)
    ):
        # 归一化层的权重设为1，偏置设为0
        if hasattr(m, 'weight') and m.weight is not None:
            nn.init.constant_(m.weight.data, 1)
        if hasattr(m, 'bias') and m.bias is not None:
            nn.init.constant_(m.bias.data, 0)
    elif isinstance(m, nn.Embedding):
        # 嵌入层使用正态分布初始化
        nn.init.normal_(m.weight.data, mean=0, std=0.1)
    elif isinstance(m, nn.LSTM):
        for name, param in m.named_parameters():
            if 'weight' in name:
                nn.init.orthogonal_(param.data)
            elif 'bias' in name:
                nn.init.constant_(param.data, 0)
    elif isinstance(m, nn.GRU):
        for name, param in m.named_parameters():
            if 'weight' in name:
                nn.init.xavier_uniform_(param.data)
            elif 'bias' in name:
                nn.init.constant_(param.data, 0)
    elif isinstance(m, nn.MultiheadAttention):
        for param in m.parameters():
            if param.dim() > 1:
                nn.init.xavier_uniform_(param)
    elif isinstance(m, nn.Transformer):
        # Transformer 层的权重初始化
        for param in m.parameters():
            if param.dim() > 1:
                nn.init.xavier_uniform_(param)


def save_checkpoint(
    models: Union[typing.Mapping[str, nn.Module], nn.Module],
    path: str,
    optimizers: Union[typing.Mapping[str, Any], Any] | None = None,
    epoch=None,
    loss=None,
    **kwargs,
):
    if isinstance(loss, torch.Tensor):
        loss = loss.item()

    normalized_models: typing.Mapping[str, nn.Module]
    if not isinstance(models, dict):
        normalized_models = cast(typing.Mapping[str, nn.Module], {'model': models})
    else:
        normalized_models = models
    normalized_optimizers: typing.Mapping[str, Any] | None
    if optimizers is not None and not isinstance(optimizers, dict):
        normalized_optimizers = {'optimizer': optimizers}
    else:
        normalized_optimizers = optimizers

    checkpoint = {
        'models_state_dict': {
            name: model.state_dict() for name, model in normalized_models.items()
        },
        'optimizers_state_dict': {
            name: opt.state_dict() for name, opt in normalized_optimizers.items()
        }
        if normalized_optimizers
        else None,
        'epoch': epoch,
        'loss': loss,
        'save_time': now(),
        **kwargs,
    }
    checkpoint = {k: v for k, v in checkpoint.items() if v is not None}
    torch.save(checkpoint, path)


def load_checkpoint(
    models: Union[typing.Mapping[str, nn.Module], nn.Module],
    path: str,
    optimizers: Union[typing.Mapping[str, Any], Any] | None = None,
    device: Union[str, torch.device] | None = None,
) -> Dict[str, Any]:
    checkpoint = torch.load(path, map_location=device, weights_only=True)

    normalized_models: typing.Mapping[str, nn.Module]
    if not isinstance(models, dict):
        normalized_models = cast(typing.Mapping[str, nn.Module], {'model': models})
    else:
        normalized_models = models
    normalized_optimizers: typing.Mapping[str, Any] | None
    if optimizers is not None and not isinstance(optimizers, dict):
        normalized_optimizers = {'optimizer': optimizers}
    else:
        normalized_optimizers = optimizers

    for name, model in normalized_models.items():
        model.load_state_dict(checkpoint['models_state_dict'][name])

    if normalized_optimizers and 'optimizers_state_dict' in checkpoint:
        for name, optimizer in normalized_optimizers.items():
            if name in checkpoint['optimizers_state_dict']:
                optimizer.load_state_dict(checkpoint['optimizers_state_dict'][name])

    return checkpoint


def freeze(model: nn.Module):
    for param in model.parameters():
        param.requires_grad = False


def unfreeze(model: nn.Module):
    for param in model.parameters():
        param.requires_grad = True
