"""Fase 0: resuelve la discrepancia de nombres reales de proyecciones de Gated DeltaNet
(in_proj_qkvz/in_proj_ba fusionadas vs in_proj_qkv/in_proj_z/in_proj_a/in_proj_b separadas)
inspeccionando la arquitectura real (device_map="meta", sin cargar pesos) de una capa
linear_attention (layer 0) y una full_attention (layer 3).
"""
import re
import sys

import torch
from transformers import AutoConfig, AutoModelForCausalLM

MODEL_PATH = sys.argv[1] if len(sys.argv) > 1 else "/workspace/ft-models/Qwen--Qwen3.6-35B-A3B"


def main():
    config = AutoConfig.from_pretrained(MODEL_PATH)
    layer_types = config.text_config.layer_types
    print(f"num_hidden_layers: {len(layer_types)}")
    print(f"layer_types[:8]: {layer_types[:8]}")

    linear_idx = layer_types.index("linear_attention")
    full_idx = layer_types.index("full_attention")
    print(f"\nUsando layer {linear_idx} (linear_attention) y layer {full_idx} (full_attention)\n")

    with torch.device("meta"):
        model = AutoModelForCausalLM.from_config(config)

    for label, idx in [("linear_attention", linear_idx), ("full_attention", full_idx)]:
        print(f"=== capa {idx} ({label}) ===")
        pattern = re.compile(rf"\.layers\.{idx}\.")
        found = False
        for name, param in model.named_parameters():
            if pattern.search(name):
                found = True
                print(f"  {name}  {tuple(param.shape)}")
        if not found:
            print(f"  (no se encontraron params para layers.{idx} — revisar prefijo real del modelo)")
        print()

    print("=== nombres de nn.Module (no solo parámetros) para la capa linear_attention ===")
    pattern = re.compile(rf"\.layers\.{linear_idx}\.")
    for name, module in model.named_modules():
        if pattern.search(name) and list(module.children()) == []:
            print(f"  {name}  ({type(module).__name__})")


if __name__ == "__main__":
    main()