from collections import defaultdict
import torch
import copy

def pack_state_dict(module: torch.nn.Module, optimizer: torch.optim.Optimizer) -> dict:
    """
    整合模型和优化器状态到一个字典中
    """
    if optimizer is None:
        return {
            "__metadata__": {
                "param_groups": [],
            }
        }
    
    # 构建参数到参数名的映射
    param_name_map = {param: name for name, param in module.named_parameters()}
    
    # 构建参数组索引映射
    param_group_map = {}
    for group_idx, group in enumerate(optimizer.param_groups):
        for param in group["params"]:
            if param in param_name_map:
                param_group_map[param_name_map[param]] = group_idx
    
    # 构建整合后的状态字典
    packed_state = {}
    for name, param in module.named_parameters():

        assert name in param_group_map, f"Parameter {name} not found in optimizer"
            
        packed_state[name] = {
            "param": param.detach().clone(),
            "param_group": param_group_map[name],
            "state": copy.deepcopy(optimizer.state.get(param, {}))
        }
    
    # 保存参数组元数据
    packed_state["__metadata__"] = {
        "param_groups": [
            {k: v for k, v in group.items() if k not in ("params", "param_names")}
            for group in optimizer.param_groups
        ]
    }
    
    # from abc4d.utils.pprint import pprint
    # pprint(packed_state)
    return packed_state


def load_packed_state_dict(
    module: torch.nn.Module,
    packed_state: dict,
    optimizer_cls: type,
    optimizer_args: dict = None, # type: ignore
    device: str = None # type: ignore
) -> torch.optim.Optimizer:
    """
    从整合状态字典恢复模型和优化器
    """
    # 1. 恢复模型参数
    model_state = {name: data["param"] for name, data in packed_state.items() 
                   if name != "__metadata__"}
    module.load_state_dict(model_state, strict=True)
    
    # 2. 准备优化器参数组结构
    metadata = packed_state.get("__metadata__", {})
    param_groups = metadata.get("param_groups", [])
    
    # 按参数组索引分组参数
    group_params = defaultdict(list)
    for name, data in packed_state.items():
        if name == "__metadata__":
            continue
        group_idx = data["param_group"]
        group_params[group_idx].append(module.get_parameter(name))
    
    # 3. 构建完整参数组
    full_param_groups = []
    for group_idx, group_meta in enumerate(param_groups):
        assert group_idx in group_params, f"Parameter group {group_idx} not found in packed state"
        group = group_meta.copy()
        group["params"] = group_params[group_idx]
        full_param_groups.append(group)
    
    # 4. 创建并初始化优化器
    if optimizer_cls is None:
        return None
    
    optimizer_args = optimizer_args or {}
    optimizer = optimizer_cls(full_param_groups, **optimizer_args)
    
    # 5. 恢复优化器状态
    state_map = {}
    for name, data in packed_state.items():
        if name == "__metadata__":
            continue
            
        param = module.get_parameter(name)
        # 将状态转移到适当设备
        state = data["state"]
        for k, v in state.items():
            if isinstance(v, torch.Tensor) and device:
                state[k] = v.to(device=device)
        state_map[param] = state
    
    optimizer.state = defaultdict(dict, state_map)
    
    return optimizer