from dataclasses import dataclass
from typing import Optional, Literal, List

import torch
import torch.nn as nn
import torch.nn.functional as F

# 来自你的代码库（可直接使用）
from hnet.modules.dc import RoutingModule, ChunkLayer
from hnet.modules.isotropic import Isotropic
from hnet.models.config_hnet import HNetConfig, SSMConfig, AttnConfig

# -------------------- 辅助 --------------------
def create_isotropic_encoder(d_model, arch="m", height=4, device=None, dtype=None):
    """创建简化的Isotropic编码器"""
    factory_kwargs = {"device": device, "dtype": dtype}
    
    # 创建HNetConfig，确保list字段有足够的元素
    config = HNetConfig(
        arch_layout=[f"{arch}{height}"],
        d_model=[d_model], 
        d_intermediate=[d_model * 2],
        ssm_cfg=SSMConfig(
            d_conv=4,
            expand=2,
            d_state=128,
            chunk_size=256
        ),
        attn_cfg=AttnConfig(
            num_heads=[8],          # 确保有至少一个元素
            rotary_emb_dim=[0],     # 确保有至少一个元素  
            window_size=[-1]        # 确保有至少一个元素
        )
    )
    
    return Isotropic(
        config=config,
        pos_idx=0,
        stage_idx=0,
        **factory_kwargs
    )

def ratio_loss(boundary_mask: torch.Tensor, boundary_prob: torch.Tensor, target_N: int) -> torch.Tensor:
    F_act = boundary_mask.float().mean(dim=1)      # (B,)
    G_prob = boundary_prob[..., 1].mean(dim=1)     # (B,)
    N = float(target_N)
    loss = N / (N - 1.0) * (((N - 1.0) * F_act) + (1.0 - F_act) * (1.0 - G_prob))
    return loss.mean()

def masked_mean(x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    mask_f = mask.float().unsqueeze(-1)     # (B, L, 1)
    s = (x * mask_f).sum(dim=1)             # (B, D)
    denom = mask_f.sum(dim=1).clamp_min(1.0)
    return s / denom

# -------------------- 多层Encoder（金字塔）：每层Mamba2 + 路由下采样，只有最终有主网络 --------------------
class PyramidEncoders_NoDechunk(nn.Module):
    """
    层级结构（仅编码器逐层压缩；主网络只在最终一层）：
      输入 x0: (B, L0, 1)
      - 线性升维 -> D0
      For s = 0..S-1:
        Es(Mamba2, D_s) -> h_s (B, L_s, D_s)
        路由 + 下采样 -> x_{s+1} (B, L_{s+1}, D_s), mask_{s+1}
        维度扩展 D_s -> D_{s+1}（拼接共享向量）
      最终 x_S: (B, L_S, D_S) 送入单一主网络 M (Transformer/Mamba)
    跨尺度融合（不去分块）：融合 E^0 的 pooled_enc0 与 主网络 pooled_main
    """
    def __init__(
        self,
        d_models: List[int],                    # [D0, D1, ..., D_S] 单调非降
        encoder_cfg_per_stage: List[dict],      # S个编码器配置（必须 arch='m'/'M'）
        main_cfg: dict,                         # 单一主网络配置（在最压缩序列上工作）
        fusion_dropout: float = 0.1,
        dtype: Optional[torch.dtype] = None,
        device: Optional[torch.device] = None,
    ):
        super().__init__()
        factory_kwargs = {"device": device, "dtype": dtype}

        assert len(d_models) >= 1
        S = len(d_models) - 1
        assert S == len(encoder_cfg_per_stage), "stage数等于encoder配置数"
        for i in range(S):
            assert d_models[i+1] >= d_models[i], "需满足 D_s <= D_{s+1}（宽度单调增加）"
            assert encoder_cfg_per_stage[i].get("arch", "m") in ("m", "M"), "Encoder必须为Mamba2"

        self.S = S
        self.d_models = d_models

        # 输入升维到 D0
        self.input_proj = nn.Linear(1, d_models[0], **factory_kwargs)

        # 每层编码器 + 路由 + 下采样 + 扩宽参数
        self.encoders = nn.ModuleList()
        self.routers = nn.ModuleList()
        self.chunks = nn.ModuleList()
        self.pad_vectors = nn.ParameterList()
        for s in range(S):
            self.encoders.append(
                create_isotropic_encoder(
                    d_model=d_models[s], 
                    **{k: v for k, v in encoder_cfg_per_stage[s].items() if k != "d_model"}, 
                    **factory_kwargs
                )
            )
            self.routers.append(RoutingModule(d_models[s], **factory_kwargs))
            self.chunks.append(ChunkLayer())
            delta = d_models[s+1] - d_models[s]
            self.pad_vectors.append(nn.Parameter(torch.zeros(delta, **factory_kwargs)) if delta > 0 else nn.Parameter(torch.empty(0, **factory_kwargs)))

        # 最终唯一的主网络：在 D_S & L_S 上运行
        self.main_network = create_isotropic_encoder(
            d_model=d_models[-1], 
            **{k: v for k, v in main_cfg.items() if k != "d_model"}, 
            **factory_kwargs
        )

        # 跨尺度融合：将 pooled_enc0(D0) 投到 D_S 并与 pooled_main(D_S) 融合 -> D_S
        self.proj_enc0_to_DS = nn.Linear(d_models[0], d_models[-1], **factory_kwargs)
        self.fusion_head = nn.Sequential(
            nn.Linear(d_models[-1] + d_models[-1], d_models[-1], **factory_kwargs),
            nn.GELU(),
            nn.Dropout(fusion_dropout),
            nn.Linear(d_models[-1], d_models[-1], **factory_kwargs),
        )

    def _expand_width(self, x: torch.Tensor, pad_vec: nn.Parameter) -> torch.Tensor:
        if pad_vec.numel() == 0:
            return x
        early = x.shape[:-1]
        return torch.cat([x, pad_vec.expand(*early, -1)], dim=-1)

    def forward(self, x_scalar: torch.Tensor, mask: Optional[torch.Tensor] = None, return_seq: bool = False):
        """
        x_scalar: (B, L) 或 (B, L, 1)
        mask: (B, L) bool
        返回:
          fused_vec: (B, D_S)
          debug: 可选
          aux: 包含各层路由信息（供ratio loss）
        """
        if x_scalar.dim() == 2:
            x_scalar = x_scalar.unsqueeze(-1)  # (B, L, 1)
        B, L, _ = x_scalar.shape
        device = x_scalar.device
        if mask is None:
            mask = torch.ones(B, L, dtype=torch.bool, device=device)

        # 初始升维到 D0
        x = self.input_proj(x_scalar)  # (B, L0, D0)
        cur_mask = mask

        pooled_enc0 = None
        aux_per_stage = []
        seq_debug = [] if return_seq else None

        # 逐层：Encoder(Mamba2)->Routing->Chunk->Expand D
        for s in range(self.S):
            d_in = self.d_models[s]
            # 细粒度编码（未压缩序列）
            h_enc = self.encoders[s](x, mask=cur_mask)  # (B, L_s, D_s)

            if s == 0:
                pooled_enc0 = masked_mean(h_enc, cur_mask)  # (B, D0)

            # 路由 + 下采样（得到更短序列）
            bpred = self.routers[s](h_enc, mask=cur_mask)
            x_next, _, _, mask_next = self.chunks[s](h_enc, bpred.boundary_mask, mask=cur_mask)  # (B, L_{s+1}, D_s)

            # 扩展宽度 D_s -> D_{s+1}
            x_next = self._expand_width(x_next, self.pad_vectors[s])  # (B, L_{s+1}, D_{s+1})

            # 推进到下一层
            x, cur_mask = x_next, mask_next

            aux_per_stage.append({
                "boundary_mask": bpred.boundary_mask,
                "boundary_prob": bpred.boundary_prob,
                "selected_probs": bpred.selected_probs,
            })
            if return_seq:
                seq_debug.append({"stage": s, "seq": x, "mask": cur_mask})

        # 现在 x: (B, L_S, D_S), cur_mask: (B, L_S)
        # 最终单一主网络在最压缩序列上
        h_main = self.main_network(x, mask=cur_mask)  # (B, L_S, D_S)

        # 主网络池化
        if cur_mask is None:
            pooled_main = h_main.mean(dim=1)  # (B, D_S)
        else:
            pooled_main = (h_main * cur_mask.float().unsqueeze(-1)).sum(dim=1) / \
                          cur_mask.float().sum(dim=1, keepdim=True).clamp_min(1.0)

        # 跨尺度融合：E^0 全局池化 与 主网络池化
        pooled_enc0_in_DS = self.proj_enc0_to_DS(pooled_enc0)  # (B, D_S)
        fused = torch.cat([pooled_enc0_in_DS, pooled_main], dim=-1)  # (B, 2*D_S)
        fused = self.fusion_head(fused)  # (B, D_S)

        aux = {"per_stage": aux_per_stage}
        if return_seq:
            return fused, {"stages": seq_debug, "main_seq": h_main, "main_mask": cur_mask}, aux
        else:
            return fused, None, aux

# -------------------- 顶层：多通道融合 + 分类头（仅一个主网络） --------------------
@dataclass
class HierEncodersSingleMainConfig:
    num_channels: int
    d_models: List[int]                      # [D0, D1, ..., D_S] 单调非降
    num_classes: int
    encoder_cfg_per_stage: List[dict]        # S个编码器配置（均为Mamba2, height≈4）
    main_cfg: dict                           # 单一主网络配置（Transformer或Mamba2），d_model自动用D_S
    target_compression_N_per_stage: List[int]
    share_channel: bool = True
    fusion_across_channels: Literal["mean", "concat"] = "mean"
    dropout: float = 0.1

class HierEncodersSingleMainClassifier(nn.Module):
    def __init__(self, cfg: HierEncodersSingleMainConfig, dtype=None, device=None):
        super().__init__()
        self.cfg = cfg
        factory_kwargs = {"dtype": dtype, "device": device}

        S = len(cfg.d_models) - 1
        assert S == len(cfg.encoder_cfg_per_stage) == len(cfg.target_compression_N_per_stage), "stage数不一致"

        if cfg.share_channel:
            self.channel_encoder = PyramidEncoders_NoDechunk(
                d_models=cfg.d_models,
                encoder_cfg_per_stage=cfg.encoder_cfg_per_stage,
                main_cfg=cfg.main_cfg,
                **factory_kwargs,
            )
        else:
            self.channel_encoder = nn.ModuleList([
                PyramidEncoders_NoDechunk(
                    d_models=cfg.d_models,
                    encoder_cfg_per_stage=cfg.encoder_cfg_per_stage,
                    main_cfg=cfg.main_cfg,
                    **factory_kwargs,
                )
                for _ in range(cfg.num_channels)
            ])

        fusion_dim = (cfg.num_channels * cfg.d_models[-1]) if cfg.fusion_across_channels == "concat" \
                     else cfg.d_models[-1]
        self.dropout = nn.Dropout(cfg.dropout)
        self.head = nn.Linear(fusion_dim, cfg.num_classes, **factory_kwargs)

    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None, return_seq: bool = False):
        """
        x: (B, L, N)  多通道输入
        mask: (B, L)  时序mask
        """
        B, L, N = x.shape
        assert N == self.cfg.num_channels

        channel_vecs: List[torch.Tensor] = []
        ratio_losses = []
        seq_dbg_all = [] if return_seq else None

        for c in range(N):
            x_c = x[..., c]  # (B, L)
            if self.cfg.share_channel:
                vec, seq_dbg, aux = self.channel_encoder(x_c, mask=mask, return_seq=return_seq)
            else:
                vec, seq_dbg, aux = self.channel_encoder[c](x_c, mask=mask, return_seq=return_seq)

            # ratio loss 累加（每个encoder stage一项）
            total_rl = 0.0
            for s, aux_s in enumerate(aux["per_stage"]):
                rl = ratio_loss(aux_s["boundary_mask"], aux_s["boundary_prob"], self.cfg.target_compression_N_per_stage[s])
                total_rl = total_rl + rl
            ratio_losses.append(total_rl)

            channel_vecs.append(vec)
            if return_seq:
                seq_dbg_all.append(seq_dbg)

        if self.cfg.fusion_across_channels == "concat":
            fused = torch.cat(channel_vecs, dim=-1)  # (B, N*D_S)
        else:
            fused = torch.stack(channel_vecs, dim=1).mean(dim=1)  # (B, D_S)

        fused = self.dropout(fused)
        logits = self.head(fused)

        aux_all = {"ratio_loss": torch.stack(ratio_losses).mean()}
        if return_seq:
            return logits, seq_dbg_all, aux_all
        else:
            return logits, None, aux_all

# -------------------- 使用示例 --------------------
if __name__ == "__main__":
    """
    符合要求：
      - 多层仅增加编码器数量（每层Mamba2 + 动态分块），主网络只有最终一个
      - 序列长度逐层缩短（由DC决定），通道维度 d_model 单调增大（SpaceByte式共享向量拼接）
      - 不使用去分块（dechunk）；跨尺度融合用 E^0 的全局池化 + 最终主网络池化
    """
    B, L, N = 8, 1024, 6
    num_classes = 7
    d_models = [128, 256, 512]  # D0 <= D1 <= D2

    encoder_cfg_per_stage = [
        dict(arch="m", height=4, ssm_cfg=dict(), attn_cfg=dict()),  # stage 0 encoder (Mamba2)
        dict(arch="m", height=4, ssm_cfg=dict(), attn_cfg=dict()),  # stage 1 encoder (Mamba2)
    ]
    main_cfg = dict(
        arch="T", height=12, ssm_cfg=dict(), attn_cfg=dict(num_heads=8)  # 最终主网络（较重）
    )
    target_compression_N_per_stage = [4, 4]

    cfg = HierEncodersSingleMainConfig(
        num_channels=N,
        d_models=d_models,
        num_classes=num_classes,
        encoder_cfg_per_stage=encoder_cfg_per_stage,
        main_cfg=main_cfg,
        target_compression_N_per_stage=target_compression_N_per_stage,
        share_channel=True,
        fusion_across_channels="mean",
        dropout=0.1,
    )

    model = HierEncodersSingleMainClassifier(cfg).cuda().train()
    x = torch.randn(B, L, N, device="cuda")
    mask = torch.ones(B, L, dtype=torch.bool, device="cuda")

    logits, _, aux = model(x, mask=mask, return_seq=False)
    y = torch.randint(0, num_classes, (B,), device="cuda")
    cls_loss = F.cross_entropy(logits, y)
    ratio_reg = 0.03 * aux["ratio_loss"]
    loss = cls_loss + ratio_reg
    loss.backward()
    print("logits:", logits.shape, "loss:", float(loss))
