from abc import ABCMeta
from typing import Any, Callable, Dict, Literal, Protocol, Union, cast

import torch
from torch import nn
from torch.nn import Module

from zyplib.nn import utils as nn_utils
from zyplib.utils.print import print_info
from zyplib.utils.time import now


class _WeightedModule(Protocol):
    weight: torch.Tensor


class MaxNormConstraint(nn.Module):
    def __init__(
        self, module: nn.Module, max_norm: float = 1.0, p: int = 2, dim: int = 0
    ):
        """
        一个通用的模块，用于对任何具有`.weight`属性的模块应用max-norm约束。

        Note
        ----------
        在每次前向时，把权重张量沿着某一维度切片（slice），对每一个 slice 做 p-范数截断（超过 max_norm 的部分缩放回去）。对于大多数 PyTorch 自带的线性／卷积层，这个“切片”的维度恰好是第 0 维：

        1. nn.Linear
            - `weight.shape = (out_features, in_features)`
            - `weight.select(0, i)` 就是第 i 个“输出神经元”对应的输入权重向量，长度 = in_features。
            - 对它们作 L2（p=2）范数约束，自然要设 dim=0。

        2. nn.ConvNd
            - `weight.shape = (out_channels, in_channels, …kernel…)`
            - weight.select(0, i) 是第 i 个输出通道上的整个卷积核（形状 = (in_channels, …kernel…)）。
            - 同理，对每个输出通道的 kernel 做 L2 范数约束，也选 dim=0。
        """
        super().__init__()
        self.module = module
        self.max_norm = max_norm
        self.p = p
        self.dim = dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            cast(_WeightedModule, self.module).weight.renorm_(
                p=self.p, dim=self.dim, maxnorm=self.max_norm
            )
        return self.module(x)


class BaseModule(Module):
    __metaclass__ = ABCMeta

    def __init__(self):
        super().__init__()

    def __hash__(self) -> int:
        return hash(str(self))

    def init_weights(self, weight_init_func=None):
        if weight_init_func is None:
            self.apply(nn_utils.weight_initializer)
        else:
            self.apply(weight_init_func)

    def count_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

    def freeze(self):
        nn_utils.freeze(self)

    def unfreeze(self):
        nn_utils.unfreeze(self)

    def save_checkpoint(self, path: str, optimizer=None, epoch=None, loss=None, **kwargs):
        nn_utils.save_checkpoint(self, path, optimizer, epoch, loss, **kwargs)

    def load_checkpoint(self, path: str, optimizer=None, device=None) -> Dict[str, Any]:
        return nn_utils.load_checkpoint(self, path, optimizer, device)

    def summary(self):
        print_info(f'Model Architecture:\n{self}')
        print_info(f'Total trainable parameters: {self.count_parameters():,}')


class LambdaModule(BaseModule):
    def __init__(self, func: Callable[[Any], torch.Tensor]):
        super().__init__()
        self.func = func

    def forward(self, *args: Any) -> torch.Tensor:
        return self.func(*args)
