"""
完整流水线 Demo: 端到端网络结构可视化
=====================================
集成所有子网络，展示从输入到输出的完整数据流。

运行模式:
  python 08_full_pipeline_demo.py --mode bev    # 仅BEV分支
  python 08_full_pipeline_demo.py --mode img    # 仅图像分支
  python 08_full_pipeline_demo.py --mode fusion # 完整融合模式
"""

import torch
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')

import sys
import os
import argparse
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from net import Fusion, BEVHead, ImgHead, FusionHead
from BEVNet import RICNN
from ALIKE.alnet import ALNet
from netvlad import NetVLAD
from uot import UOTHead

OUTPUT_DIR = os.path.join(os.path.dirname(__file__), 'output')
os.makedirs(OUTPUT_DIR, exist_ok=True)


def create_dummy_batch_dict(mode='fusion'):
    """创建模拟的batch_dict"""
    B = 2  # batch中1对 (query + positive)
    batch_dict = {
        'batch_size': 2 * B,
    }

    if mode in ('fusion', 'bev'):
        batch_dict['bev'] = torch.randn(2 * B, 7, 320, 320)
        batch_dict['bev'][:, :3] = torch.sigmoid(batch_dict['bev'][:, :3])  # 可视通道
        batch_dict['bev'][:, 2:3] = (batch_dict['bev'][:, 2:3] > 0.3).float()  # guider mask

    if mode in ('fusion', 'img'):
        batch_dict['img'] = torch.randint(0, 256, (2 * B, 5, 192, 576)).float()

    if mode == 'fusion':
        # 模拟 relation: (B, max_len, K, 2)
        max_len, K = 200, 11  # K=1+10: last dim is bev coord
        batch_dict['relation'] = torch.zeros(2 * B, max_len, K, 2, dtype=torch.long)
        for i in range(2 * B):
            n_valid = 150
            batch_dict['relation'][i, :n_valid, :K - 1, 0] = torch.randint(0, 576, (n_valid, K - 1))
            batch_dict['relation'][i, :n_valid, :K - 1, 1] = torch.randint(0, 192, (n_valid, K - 1))
            batch_dict['relation'][i, :n_valid, K - 1, 0] = torch.randint(0, 320, (n_valid,))
            batch_dict['relation'][i, :n_valid, K - 1, 1] = torch.randint(0, 320, (n_valid,))

        # pose_to_frame (训练时需要)
        angle = 0.3
        pose = torch.eye(4).unsqueeze(0).repeat(B, 1, 1)
        pose[:, 0, 0] = torch.cos(torch.tensor(angle))
        pose[:, 0, 1] = -torch.sin(torch.tensor(angle))
        pose[:, 1, 0] = torch.sin(torch.tensor(angle))
        pose[:, 1, 1] = torch.cos(torch.tensor(angle))
        pose[:, 0, 3] = 2.0
        pose[:, 1, 3] = -1.0
        batch_dict['pose_to_frame'] = pose.clone()

        batch_dict['pose_query'] = torch.eye(4).unsqueeze(0).repeat(B, 1, 1)
        batch_dict['pose_positive'] = torch.eye(4).unsqueeze(0).repeat(B, 1, 1)

        batch_dict['label_score'] = torch.zeros(B, 320, 320, 2)
        batch_dict['id_query'] = torch.arange(B)
        batch_dict['id_positive'] = torch.arange(B)
        batch_dict['sequence'] = torch.zeros(B, dtype=torch.long)

    return batch_dict


def run_bev_only():
    """仅BEV分支"""
    print('\n' + '=' * 60)
    print('模式: BEV Only (仅点云分支)')
    print('=' * 60)

    cfg = {
        'flag': 'bev',
        'kpts_number_bev': 150,
        'kpts_number_img': 150,
        'cluster_num_bev': 16,
        'cluster_num_img': 16,
        'cluster_num_fusion': 16,
        'sinkhorn_iter': 5,
        'vlad_size': 256,
    }

    model = Fusion(cfg)
    model.eval()
    total_params = sum(p.numel() for p in model.parameters())
    print(f'模型参数量: {total_params:,} ({total_params / 1e6:.2f}M)')

    batch_dict = create_dummy_batch_dict('bev')

    with torch.no_grad():
        output = model(batch_dict)

    print('\n输出:')
    for k, v in output.items():
        if isinstance(v, torch.Tensor):
            print(f'  {k:30s}: {list(v.shape)}')
        else:
            print(f'  {k:30s}: {v}')

    # 可视化BEV分支数据流
    fig, axes = plt.subplots(2, 4, figsize=(18, 9))

    # BEV输入 (3个可视通道)
    if 'bev' in output or 'bev' in batch_dict:
        bev_in = batch_dict['bev'][0, :3].permute(1, 2, 0).numpy()
        axes[0, 0].imshow(bev_in)
        axes[0, 0].set_title('BEV输入 (3通道)')
        axes[0, 0].axis('off')

    # Score Map
    if 'score_bev' in output:
        axes[0, 1].imshow(output['score_bev'][0].numpy(), cmap='hot')
        axes[0, 1].set_title('BEV Score Map')
        axes[0, 1].axis('off')

    # 关键点位置
    if 'key_points' in output and 'pixels_kpt' in output:
        bev_show = batch_dict['bev'][0, :3].permute(1, 2, 0).numpy()
        axes[0, 2].imshow(bev_show)
        kpt = output['pixels_kpt'][0].numpy()
        axes[0, 2].scatter(kpt[:, 1], kpt[:, 0], c='red', s=5, alpha=0.8)
        axes[0, 2].set_title(f'BEV Top-{len(kpt)} 关键点')
        axes[0, 2].axis('off')

    # Descriptor Map (第一通道)
    if 'fea_bev' in output:
        axes[0, 3].imshow(output['fea_bev'][0, 0].numpy(), cmap='viridis')
        axes[0, 3].set_title('BEV Descriptor ch0')
        axes[0, 3].axis('off')

    # 关键点特征相似度
    if 'fea_kpt_original' in output:
        fea = output['fea_kpt_original']
        # query vs positive 的相似度
        B = fea.shape[0] // 2
        sim = torch.nn.functional.cosine_similarity(
            fea[:B].permute(0, 2, 1).unsqueeze(-1),
            fea[B:].permute(0, 2, 1).unsqueeze(-2),
            dim=1
        )[0]
        im = axes[1, 0].imshow(sim.numpy(), cmap='RdYlBu_r', vmin=-1, vmax=1)
        axes[1, 0].set_title('Query-Positive 特征相似度')
        axes[1, 0].set_xlabel('Positive'); axes[1, 0].set_ylabel('Query')
        plt.colorbar(im, ax=axes[1, 0])

    # VLAD
    if 'vlads' in output:
        vlad = output['vlads'][0].view(16, 128).numpy()
        im = axes[1, 1].imshow(vlad, cmap='RdBu_r', aspect='auto')
        axes[1, 1].set_title('VLAD描述子 (16×128)')
        axes[1, 1].set_xlabel('Feature Dim'); axes[1, 1].set_ylabel('Cluster')
        plt.colorbar(im, ax=axes[1, 1])

    # 数据流图
    axes[1, 2].set_title('BEV分支数据流')
    flow = [
        'bev (7,320,320)',
        '→ x = bev[:3] (可视BEV)',
        '→ points = bev[3:7] (坐标)',
        '→ RICNN前向',
        '→ score_bev (1,320,320)',
        '→ fea_bev (128,320,320)',
        '→ NMS + Top-K(150)',
        '→ key_points (150,4)',
        '→ fea_kpt (128,150)',
        '→ EncodePosition',
        '→ NetVLAD → vlad_bev (2048)',
    ]
    for i, f in enumerate(flow):
        axes[1, 2].text(0.1, 0.95 - i * 0.1, f, transform=axes[1, 2].transAxes,
                        fontsize=9, family='monospace')
    axes[1, 2].axis('off')

    # 参数量饼图
    axes[1, 3].set_title('BEV分支参数分布')
    modules = dict(model.bev.feature_extractor.named_children())
    sizes = []
    labels = []
    for name, mod in modules.items():
        p = sum(pm.numel() for pm in mod.parameters())
        if p > 0:
            sizes.append(p)
            labels.append(f'{name}\n({p/1e3:.0f}K)')
    axes[1, 3].pie(sizes, labels=labels, autopct='%1.1f%%', textprops={'fontsize': 8})

    plt.suptitle('BEV Only 模式: 点云分支可视化', fontsize=14, fontweight='bold')
    plt.tight_layout()
    path = os.path.join(OUTPUT_DIR, 'full_pipeline_bev.png')
    plt.savefig(path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f'[保存] {path}')


def run_img_only():
    """仅图像分支"""
    print('\n' + '=' * 60)
    print('模式: Image Only (仅图像分支)')
    print('=' * 60)

    cfg = {
        'flag': 'img',
        'kpts_number_bev': 150,
        'kpts_number_img': 150,
        'cluster_num_bev': 16,
        'cluster_num_img': 16,
        'cluster_num_fusion': 16,
        'sinkhorn_iter': 5,
        'vlad_size': 256,
    }

    model = Fusion(cfg)
    model.eval()
    total_params = sum(p.numel() for p in model.parameters())
    print(f'模型参数量: {total_params:,} ({total_params / 1e6:.2f}M)')

    batch_dict = create_dummy_batch_dict('img')

    with torch.no_grad():
        output = model(batch_dict)

    print('\n输出:')
    for k, v in output.items():
        if isinstance(v, torch.Tensor):
            print(f'  {k:30s}: {list(v.shape)}')
        else:
            print(f'  {k:30s}: {v}')

    # 可视化
    fig, axes = plt.subplots(2, 4, figsize=(18, 9))

    # 输入图像
    img_in = batch_dict['img'][0, :3].permute(1, 2, 0).numpy().astype(np.uint8)
    axes[0, 0].imshow(img_in)
    axes[0, 0].set_title('图像输入 (192×576)')
    axes[0, 0].axis('off')

    # Score Map
    if 'score_img' in output:
        axes[0, 1].imshow(output['score_img'][0, 0].numpy(), cmap='hot')
        axes[0, 1].set_title('图像 Score Map')
        axes[0, 1].axis('off')

    # 关键点
    if 'key_pixels' in output:
        axes[0, 2].imshow(img_in)
        kpt = output['key_pixels'][0].numpy()
        axes[0, 2].scatter(kpt[:, 1], kpt[:, 0], c='red', s=5, alpha=0.8)
        axes[0, 2].set_title(f'Top-{len(kpt)} 关键点')
        axes[0, 2].axis('off')

    # Descriptor Map
    if 'fea_img' in output:
        axes[0, 3].imshow(output['fea_img'][0, 0].numpy(), cmap='viridis')
        axes[0, 3].set_title('图像 Descriptor ch0')
        axes[0, 3].axis('off')

    # 关键点特征相似度
    if 'fea_kpl' in output:
        fea = output['fea_kpl']
        B = fea.shape[0] // 2
        sim = torch.nn.functional.cosine_similarity(
            fea[:B].permute(0, 2, 1).unsqueeze(-1),
            fea[B:].permute(0, 2, 1).unsqueeze(-2),
            dim=1
        )[0]
        im = axes[1, 0].imshow(sim.numpy(), cmap='RdYlBu_r', vmin=-1, vmax=1)
        axes[1, 0].set_title('Query-Positive 特征相似度')
        plt.colorbar(im, ax=axes[1, 0])

    # 数据流图
    axes[1, 1].set_title('图像分支数据流')
    flow = [
        'img (5,192,576)',
        '→ x = img[:3]/255',
        '→ ALNet前向',
        '→ score_img (1,192,576)',
        '→ fea_img (128,192,576)',
        '→ NMS(2) + Top-K(150)',
        '→ key_pixels (150,2)',
        '→ fea_kpl (128,150)',
    ]
    for i, f in enumerate(flow):
        axes[1, 1].text(0.1, 0.95 - i * 0.11, f, transform=axes[1, 1].transAxes,
                        fontsize=9, family='monospace')
    axes[1, 1].axis('off')

    # 参数量饼图
    axes[1, 2].set_title('图像分支参数分布')
    modules = dict(model.img.feature_extractor.named_children())
    sizes = []
    labels = []
    for name, mod in modules.items():
        p = sum(pm.numel() for pm in mod.parameters())
        if p > 0:
            sizes.append(p)
            labels.append(f'{name}\n({p/1e3:.0f}K)')
    axes[1, 2].pie(sizes, labels=labels, autopct='%1.1f%%', textprops={'fontsize': 8})

    axes[1, 3].axis('off')

    plt.suptitle('Image Only 模式: 图像分支可视化', fontsize=14, fontweight='bold')
    plt.tight_layout()
    path = os.path.join(OUTPUT_DIR, 'full_pipeline_img.png')
    plt.savefig(path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f'[保存] {path}')


def run_fusion():
    """完整融合模式"""
    print('\n' + '=' * 60)
    print('模式: Fusion (完整融合)')
    print('=' * 60)

    cfg = {
        'flag': 'fusion',
        'kpts_number_bev': 150,
        'kpts_number_img': 150,
        'cluster_num_bev': 16,
        'cluster_num_img': 16,
        'cluster_num_fusion': 16,
        'sinkhorn_iter': 5,
        'vlad_size': 256,
    }

    model = Fusion(cfg)
    model.eval()
    total_params = sum(p.numel() for p in model.parameters())
    print(f'模型参数量: {total_params:,} ({total_params / 1e6:.2f}M)')

    batch_dict = create_dummy_batch_dict('fusion')

    with torch.no_grad():
        output = model(batch_dict)

    print('\n输出:')
    for k, v in output.items():
        if isinstance(v, torch.Tensor):
            print(f'  {k:30s}: {list(v.shape)}')
        else:
            print(f'  {k:30s}: {v}')

    # 可视化融合数据流
    fig, axes = plt.subplots(3, 4, figsize=(22, 15))

    # BEV输入
    bev_in = batch_dict['bev'][0, :3].permute(1, 2, 0).numpy()
    axes[0, 0].imshow(bev_in)
    axes[0, 0].set_title('BEV 输入 (320×320)')
    axes[0, 0].axis('off')

    # 图像输入
    img_in = batch_dict['img'][0, :3].permute(1, 2, 0).numpy().astype(np.uint8)
    axes[0, 1].imshow(img_in)
    axes[0, 1].set_title('图像输入 (192×576)')
    axes[0, 1].axis('off')

    # Score maps
    if 'score_bev' in output:
        axes[0, 2].imshow(output['score_bev'][0].numpy(), cmap='hot')
        axes[0, 2].set_title('BEV Score')
        axes[0, 2].axis('off')
    if 'score_img' in output:
        axes[0, 3].imshow(output['score_img'][0, 0].numpy(), cmap='hot')
        axes[0, 3].set_title('Image Score')
        axes[0, 3].axis('off')

    # 融合特征空间中的相似度
    if 'fea_kpt_original' in output and 'fea_kpt_fusion' in output:
        fea_orig = output['fea_kpt_original']
        fea_fusion = output['fea_kpt_fusion']
        B = fea_orig.shape[0] // 2

        sim_orig = torch.nn.functional.cosine_similarity(
            fea_orig[:B].permute(0, 2, 1).unsqueeze(-1),
            fea_orig[B:].permute(0, 2, 1).unsqueeze(-2),
            dim=1
        )[0].numpy()

        sim_fusion = torch.nn.functional.cosine_similarity(
            fea_fusion[:B].permute(0, 2, 1).unsqueeze(-1),
            fea_fusion[B:].permute(0, 2, 1).unsqueeze(-2),
            dim=1
        )[0].numpy()

        im1 = axes[1, 0].imshow(sim_orig, cmap='RdYlBu_r', vmin=-1, vmax=1)
        axes[1, 0].set_title('原始特征 相似度 (150×150)')
        plt.colorbar(im1, ax=axes[1, 0])

        im2 = axes[1, 1].imshow(sim_fusion, cmap='RdYlBu_r', vmin=-1, vmax=1)
        axes[1, 1].set_title('融合特征 相似度 (150×150)')
        plt.colorbar(im2, ax=axes[1, 1])

        axes[1, 2].imshow(np.abs(sim_orig - sim_fusion), cmap='YlOrRd')
        axes[1, 2].set_title('相似度变化 |差异|')
        plt.colorbar(im2, ax=axes[1, 2])

    # VLAD
    if 'vlads' in output:
        vlad = output['vlads'][0].view(16, 128).numpy()
        im = axes[1, 3].imshow(vlad, cmap='RdBu_r', aspect='auto')
        axes[1, 3].set_title('VLAD 融合 (16×128)')
        plt.colorbar(im, ax=axes[1, 3])

    # 整体架构图
    axes[2, 0].set_title('完整架构')
    arch = [
        '┌─ BEVHead ─────────────┐',
        '│ RICNN + EncodePos     │',
        '│ → fea_kpt_original    │',
        '│ → vlad_bev            │',
        '└───────────────────────┘',
        '┌─ ImgHead ─────────────┐',
        '│ ALNet + NMS           │',
        '│ → fea_kpl             │',
        '│ → fea_img             │',
        '└───────────────────────┘',
        '┌─ FusionHead ──────────┐',
        '│ LocalPool + Converter │',
        '│ Generator + FusionHead│',
        '│ → fea_kpt_fusion      │',
        '└───────────────────────────────────────────────────────┘',
        '  VLAD = w·vlad_fusion + (1-w)·vlad_bev'
    ]
    for i, a in enumerate(arch):
        axes[2, 0].text(0.05, 0.98 - i * 0.075, a, transform=axes[2, 0].transAxes,
                        fontsize=7.5, family='monospace')
    axes[2, 0].axis('off')

    # 模块参数对比
    axes[2, 1].set_title('各模块参数量')
    module_names = []
    module_params = []
    for name, mod in model.named_children():
        p = sum(pm.numel() for pm in mod.parameters())
        if p > 0:
            module_names.append(name)
            module_params.append(p)
    colors = plt.cm.Set3(np.linspace(0, 1, len(module_names)))
    axes[2, 1].barh(range(len(module_names)), module_params, color=colors)
    axes[2, 1].set_yticks(range(len(module_names)))
    axes[2, 1].set_yticklabels(module_names, fontsize=8)
    for i, p in enumerate(module_params):
        axes[2, 1].text(p, i, f' {p/1e3:.0f}K', va='center', fontsize=8)

    # 数据流汇总
    axes[2, 2].set_title('融合模式数据流')
    flow = [
        'img, bev, relation 输入',
        '├─ ImgHead → ALNet',
        '│   ├─ score_img',
        '│   ├─ fea_img (密集描述子)',
        '│   └─ fea_kpl (关键点)',
        '├─ BEVHead → RICNN',
        '│   ├─ score_bev',
        '│   ├─ fea_bev (密集描述子)',
        '│   ├─ fea_kpt_original',
        '│   └─ vlad_bev',
        '└─ FusionHead',
        '    ├─ GridSample → fea_pl_dual, fea_pt_dual',
        '    ├─ Converters → 跨模态转换',
        '    ├─ Generator → 全景特征',
        '    ├─ FusionHead → 融合特征',
        '    └─ NetVLAD → vlad_fusion',
        '最终: vlads = w·vlad_fusion + (1-w)·vlad_bev',
        '     UOT: → transformation (位姿)',
    ]
    for i, f in enumerate(flow):
        axes[2, 2].text(0.05, 0.98 - i * 0.06, f, transform=axes[2, 2].transAxes,
                        fontsize=7.5, family='monospace')
    axes[2, 2].axis('off')

    axes[2, 3].axis('off')

    plt.suptitle('Fusion 模式: 完整跨模态融合可视化', fontsize=14, fontweight='bold')
    plt.tight_layout()
    path = os.path.join(OUTPUT_DIR, 'full_pipeline_fusion.png')
    plt.savefig(path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f'[保存] {path}')


def main():
    parser = argparse.ArgumentParser(description='全流水线可视化')
    parser.add_argument('--mode', type=str, default='all',
                        choices=['all', 'bev', 'img', 'fusion'],
                        help='运行模式')
    args = parser.parse_args()

    if args.mode in ('all', 'bev'):
        run_bev_only()
    if args.mode in ('all', 'img'):
        run_img_only()
    if args.mode in ('all', 'fusion'):
        run_fusion()

    print(f'\n所有可视化结果保存在: {OUTPUT_DIR}')


if __name__ == '__main__':
    main()