"""End-to-End 两阶段推理脚本 (Text → Semantic DiT → VAE DiT → Audio)

推理管线:
  Text prompt
    ↓
  [Text Encoders] (PE + Flan-T5)
    pe_text_embeds: (B, 1024)
    flan_text_feature: (B, seq, 1024)
    ↓
  [Stage 1: Semantic DiT]  (train_semantic.py 训练的)
    FlowMatchingWrapper.sample(shape=(B,250,8/64), cfg_scale=...)
    noise → semantic latent z (B, 250, 8/64)
    ↓
  [Stage 2: VAE DiT]  (train_vae_scaled.py 训练的)
    vae_gen.dit(x, time, pe_audio_cond=z)
    Euler 采样 t=1→0, noise → DAC latent (B, 250, 128)
    注意: z 直接作为 pe_audio_cond，跳过 audio_projector
    ↓
  [DAC Decoder]
    (B, 128, 250) → audio waveform (48kHz) → .wav

用法:
    # 自定义文本推理
    python infer_semantic.py --prompts "guitar music" "a dog barking" --device cuda:4

    # 从 validation 数据取样（有 GT 对比）
    python infer_semantic.py --from_data --num_samples 10 --device cuda:4
"""

import argparse
import os
import sys
from pathlib import Path

import torch
import torchaudio
from tqdm import tqdm

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

PE_MODULE_PATH = "/data/peaudio/perception_models"


# ============================================================
# 模型加载
# ============================================================

def load_text_encoders(device):
    """加载 PE text encoder 和 Flan-T5 encoder（复用 infer_text.py）"""

    # ---- PE (pe-a-frame-large) ----
    print("加载 Perception Encoder (pe-a-frame-large)...")
    sys.path.insert(0, PE_MODULE_PATH)
    from core.audio_visual_encoder import PEAudioFrame, PEAudioFrameTransform
    pe_model = PEAudioFrame.from_config("pe-a-frame-large", pretrained=True).to(device)
    pe_model.eval()
    pe_transform = PEAudioFrameTransform.from_config("pe-a-frame-large")
    print("  PE loaded.")

    # ---- Flan-T5 ----
    print("加载 Flan-T5...")
    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
    flan_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large").to(device)
    flan_model.eval()
    flan_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
    print("  Flan-T5 loaded.")

    return pe_model, pe_transform, flan_model, flan_tokenizer


def load_semantic_dit(ckpt_path, pe_audio_out_dim, device):
    """加载 Semantic DiT 模型（train_semantic.py 训练的）

    从 checkpoint 中提取 model.* 权重（FlowMatchingWrapper 包裹的 DiT）。
    推理时 cfg_dropout=0.0。
    """
    from dit.dit import DiT
    from dit.config import TransformerConfig
    from flow_matching import FlowMatchingWrapper

    print(f"加载 Semantic DiT (out_dim={pe_audio_out_dim}): {ckpt_path}")
    checkpoint = torch.load(ckpt_path, map_location='cpu', weights_only=False)
    hparams = checkpoint.get('hyper_parameters', {})

    config = TransformerConfig(
        dim=hparams.get('dim', 1024),
        n_layers=hparams.get('n_layers', 16),
        n_heads=hparams.get('n_heads', 16),
        in_channels=hparams.get('in_channels', pe_audio_out_dim),
        out_channels=hparams.get('out_channels', pe_audio_out_dim),
        global_cond_dim=hparams.get('global_cond_dim', 1024),
        cross_cond_dim=hparams.get('cross_cond_dim', 1024),
        use_global_cond=True,
        use_cross_cond=True,
        max_positions=hparams.get('max_positions', 1024),
        cfg_dropout=0.0,  # 推理时不 dropout
    )

    dit_model = DiT(config)
    flow_wrapper = FlowMatchingWrapper(
        model=dit_model,
        inference_mode='euler',
        num_steps=50,
        reverse_flow=True,
    )

    # 提取 model.* 权重（train_semantic.py 中 self.model = FlowMatchingWrapper(...)）
    flow_wrapper_state = {}
    for k, v in checkpoint['state_dict'].items():
        if k.startswith('model.'):
            flow_wrapper_state[k[len('model.'):]] = v

    flow_wrapper.load_state_dict(flow_wrapper_state, strict=True)
    flow_wrapper = flow_wrapper.to(device).eval()
    print(f"  加载了 {len(flow_wrapper_state)} 个参数")
    return flow_wrapper


def load_vae_model(ckpt_path, pe_audio_out_dim, device):
    """加载 VAE Generator 模型（复用 infer_all.py 的实现）

    从 checkpoint 中提取 vae_generator.* 权重。
    """
    from generator.vaegen import VAEGenerator, VAEGeneratorConfig

    print(f"加载 VAE Generator (out_dim={pe_audio_out_dim}): {ckpt_path}")
    checkpoint = torch.load(ckpt_path, map_location='cpu', weights_only=False)

    vae_state = {}
    for k, v in checkpoint['state_dict'].items():
        if k.startswith('vae_generator.'):
            vae_state[k[len('vae_generator.'):]] = v

    hparams = checkpoint.get('hyper_parameters', {})
    config = VAEGeneratorConfig(
        model_dim=hparams.get('model_dim', 1024),
        n_layers=hparams.get('n_layers', 16),
        n_heads=hparams.get('n_heads', 16),
        pe_audio_in_dim=hparams.get('pe_audio_in_dim', 1024),
        pe_audio_out_dim=pe_audio_out_dim,
        dac_dim=hparams.get('dac_dim', 128),
        dropout=hparams.get('dropout', 0.1),
    )

    vae_gen = VAEGenerator(config)
    vae_gen.load_state_dict(vae_state, strict=True)
    vae_gen = vae_gen.to(device).eval()
    print(f"  加载了 {len(vae_state)} 个参数")
    return vae_gen


# ============================================================
# 采样函数
# ============================================================

def extract_text_features(prompts, pe_model, pe_transform, flan_model, flan_tokenizer, device):
    """从文本提取 pe_text_embeds 和 flan_text_feature（复用 infer_text.py）"""

    # ---- PE text embeds ----
    pe_inputs = pe_transform(text=prompts)
    pe_inputs = pe_inputs.to(device)

    with torch.inference_mode():
        text_model_output = pe_model._get_text_output(
            pe_inputs['input_ids'], pe_inputs['attention_mask']
        )
        pe_text_embeds = pe_model.text_head(text_model_output.pooler_output)
        # pe_text_embeds: (B, 1024)

    # ---- Flan-T5 text features ----
    flan_inputs = flan_tokenizer(
        prompts, return_tensors="pt", padding=True, truncation=True
    ).to(device)

    with torch.no_grad():
        flan_out = flan_model.encoder(**flan_inputs)
    flan_text_feature = flan_out.last_hidden_state  # (B, seq_len, 1024)
    flan_text_mask = flan_inputs['attention_mask']   # (B, seq_len)

    return pe_text_embeds, flan_text_feature, flan_text_mask


def sample_semantic(flow_wrapper, pe_text_embeds, flan_text_feature, flan_text_mask,
                    out_dim, num_steps=50, cfg_scale=3.0):
    """Stage 1: Semantic DiT 采样

    Text → semantic latent z (B, 250, out_dim)
    使用 FlowMatchingWrapper.sample() 的 CFG 支持。
    """
    B = pe_text_embeds.shape[0]
    shape = (B, 250, out_dim)

    generated = flow_wrapper.sample(
        shape=shape,
        global_cond=pe_text_embeds,
        cross_cond=flan_text_feature,
        memory_padding_mask=flan_text_mask,
        num_steps=num_steps,
        cfg_scale=cfg_scale,
    )
    return generated  # (B, 250, out_dim)


def sample_vae_from_semantic(vae_gen, semantic_z, num_steps=50):
    """Stage 2: VAE DiT 采样（以 semantic latent 为条件）

    semantic_z 直接作为 pe_audio_cond，跳过 audio_projector。
    手动 Euler 采样 t=1→0。

    Args:
        vae_gen: VAEGenerator 模型
        semantic_z: (B, 250, out_dim) — Stage 1 生成的 semantic latent
        num_steps: Euler 采样步数

    Returns:
        x: (B, 250, 128) — 生成的 DAC latent
    """
    B = semantic_z.shape[0]
    T_seq, dac_dim = 250, 128
    device = semantic_z.device
    dtype = semantic_z.dtype

    # 从随机噪声出发
    x = torch.randn(B, T_seq, dac_dim, device=device, dtype=dtype)
    dt = -1.0 / num_steps  # t: 1→0, 负方向

    for step_i in range(num_steps):
        t_val = 1.0 - step_i / num_steps
        t = torch.full((B,), t_val, device=device, dtype=x.dtype)
        # 直接传 semantic_z 作为 pe_audio_cond（跳过 audio_projector）
        v = vae_gen.dit(x=x, time=t, pe_audio_cond=semantic_z)
        x = x + v * dt

    return x  # (B, 250, 128)


def decode_dac(dac_latent, dac_model, device):
    """DAC latent -> audio waveform（复用 infer_text.py）"""
    if dac_latent.dim() == 3 and dac_latent.shape[-1] == 128:
        dac_latent = dac_latent.transpose(1, 2)
    dac_latent = dac_latent.to(device).float()
    with torch.no_grad():
        audio = dac_model.decode(dac_latent)
    return audio


def save_audio(audio_tensor, path, sample_rate):
    """保存音频，归一化到 [-1, 1]"""
    wav = audio_tensor.squeeze(0).cpu().float()
    wav = wav / (wav.abs().max() + 1e-8)
    torchaudio.save(str(path), wav, sample_rate)


# ============================================================
# Main
# ============================================================

def main():
    parser = argparse.ArgumentParser(
        description='End-to-End 两阶段推理: Text → Semantic DiT → VAE DiT → Audio'
    )

    # 使用模式（二选一）
    mode_group = parser.add_mutually_exclusive_group(required=True)
    mode_group.add_argument('--prompts', nargs='+',
                            help='自定义文本 prompt 列表，如 "guitar music" "piano solo"')
    mode_group.add_argument('--from_data', action='store_true',
                            help='从 validation 数据取样（有 GT 对比）')

    # 数据参数（from_data 模式）
    parser.add_argument('--num_samples', type=int, default=10,
                        help='from_data 模式下取样数量')
    parser.add_argument('--data_dir', type=str,
                        default='/data/data/parquet_features/')
    parser.add_argument('--data_pattern', type=str, default='audiocaps_validation_*.parquet')

    # Checkpoint 路径（自动检测）
    parser.add_argument('--ckpt_semantic_8', type=str, default=None,
                        help='Semantic DiT 8dim checkpoint')
    parser.add_argument('--ckpt_semantic_64', type=str, default=None,
                        help='Semantic DiT 64dim checkpoint')
    parser.add_argument('--ckpt_semantic_64_large', type=str, default=None,
                        help='Semantic DiT 64dim large (d28 h1152) checkpoint')
    parser.add_argument('--ckpt_vae_8', type=str, default=None,
                        help='VAE DiT 8dim checkpoint')
    parser.add_argument('--ckpt_vae_64', type=str, default=None,
                        help='VAE DiT 64dim checkpoint')

    # 通用 checkpoint（支持任意 dim）
    parser.add_argument('--ckpt_semantic', type=str, default=None,
                        help='通用 Semantic DiT checkpoint（配合 --out_dim 使用）')
    parser.add_argument('--ckpt_vae', type=str, default=None,
                        help='通用 VAE DiT checkpoint（配合 --out_dim 使用）')
    parser.add_argument('--out_dim', type=int, default=64,
                        help='通用 pipeline 的 out_dim（如 32, 64, 128）')
    parser.add_argument('--pipeline_name', type=str, default=None,
                        help='通用 pipeline 的名称（默认自动生成）')

    # 采样参数
    parser.add_argument('--cfg_scale', type=float, default=3.0,
                        help='Semantic DiT 的 CFG 引导强度')
    parser.add_argument('--num_steps_semantic', type=int, default=50,
                        help='Stage 1 (Semantic DiT) 采样步数')
    parser.add_argument('--num_steps_vae', type=int, default=50,
                        help='Stage 2 (VAE DiT) 采样步数')

    # 输出
    parser.add_argument('--output_dir', type=str,
                        default='/data/semaudio/infer_semantic_output')

    # 硬件
    parser.add_argument('--device', type=str, default='cuda:0')

    args = parser.parse_args()

    # ---- Checkpoint 自动检测 ----
    log_base = '/data/semaudio/logs'
    auto_ckpts = {
        'ckpt_semantic_8':        f'{log_base}/semantic_8dim/checkpoints/last.ckpt',
        'ckpt_semantic_64':       f'{log_base}/semantic_64dim/checkpoints/last.ckpt',
        'ckpt_semantic_64_large': f'{log_base}/semantic_64dim_large/checkpoints/last.ckpt',
        'ckpt_vae_8':             f'{log_base}/vae_8dim_kl/checkpoints/last.ckpt',
        'ckpt_vae_64':            f'{log_base}/vae_64dim_kl/checkpoints/last.ckpt',
    }
    for attr, path in auto_ckpts.items():
        if getattr(args, attr) is None and os.path.exists(path):
            setattr(args, attr, path)

    # ---- 配置 pipeline ----
    # 每条 pipeline = (semantic_ckpt, vae_ckpt, out_dim, name)
    # 只添加 checkpoint 文件实际存在的 pipeline
    def _ckpt_valid(p):
        return p is not None and os.path.isfile(p)

    pipelines = []
    if _ckpt_valid(args.ckpt_semantic_8) and _ckpt_valid(args.ckpt_vae_8):
        pipelines.append((args.ckpt_semantic_8, args.ckpt_vae_8, 8, 'semantic_8dim'))
    if _ckpt_valid(args.ckpt_semantic_64) and _ckpt_valid(args.ckpt_vae_64):
        pipelines.append((args.ckpt_semantic_64, args.ckpt_vae_64, 64, 'semantic_64dim'))
    if _ckpt_valid(args.ckpt_semantic_64_large) and _ckpt_valid(args.ckpt_vae_64):
        pipelines.append((args.ckpt_semantic_64_large, args.ckpt_vae_64, 64, 'semantic_64dim_large'))

    # 通用 pipeline（支持任意 dim）
    if _ckpt_valid(args.ckpt_semantic) and _ckpt_valid(args.ckpt_vae):
        name = args.pipeline_name or f'semantic_{args.out_dim}dim_custom'
        pipelines.append((args.ckpt_semantic, args.ckpt_vae, args.out_dim, name))

    if not pipelines:
        print("错误: 未找到可用的 checkpoint 组合。")
        print("需要 Semantic DiT + VAE DiT 的 checkpoint 配对。")
        print("检测路径:")
        for attr, path in auto_ckpts.items():
            exists = "✅" if os.path.exists(path) else "❌"
            print(f"  {exists} {attr}: {path}")
        return

    device = args.device
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # ============================================================
    # 加载模型
    # ============================================================

    # DAC
    print("\n加载 DAC 模型...")
    from dacvae import DACVAE
    dac_model = DACVAE.load("facebook/dacvae-watermarked").to(device).eval()
    sample_rate = dac_model.sample_rate
    print(f"DAC 加载完成, sample_rate={sample_rate}")

    # 加载各 pipeline 的模型
    print("\n" + "=" * 60)
    print("加载推理管线模型...")
    loaded_pipelines = {}  # name -> (semantic_flow_wrapper, vae_gen, out_dim)

    for sem_ckpt, vae_ckpt, out_dim, name in pipelines:
        print(f"\n--- Pipeline: {name} (out_dim={out_dim}) ---")
        sem_model = load_semantic_dit(sem_ckpt, out_dim, device)
        vae_model = load_vae_model(vae_ckpt, out_dim, device)
        loaded_pipelines[name] = (sem_model, vae_model, out_dim)

    print(f"\n共 {len(loaded_pipelines)} 条管线: {list(loaded_pipelines.keys())}")

    # ============================================================
    # Prompts 模式: 自定义文本 → 音频
    # ============================================================
    if args.prompts:
        # 加载文本编码器
        print("\n加载文本编码器...")
        pe_model, pe_transform, flan_model, flan_tokenizer = load_text_encoders(device)

        # 提取文本特征
        print(f"\n提取文本特征 ({len(args.prompts)} 个 prompt)...")
        pe_text_embeds, flan_text_feature, flan_text_mask = extract_text_features(
            args.prompts, pe_model, pe_transform, flan_model, flan_tokenizer, device
        )
        print(f"  pe_text_embeds: {pe_text_embeds.shape}")
        print(f"  flan_text_feature: {flan_text_feature.shape}")

        # 释放文本编码器显存
        del pe_model, flan_model
        torch.cuda.empty_cache()

        # 各 pipeline 推理
        print(f"\n开始推理 (cfg_scale={args.cfg_scale}, "
              f"steps_semantic={args.num_steps_semantic}, steps_vae={args.num_steps_vae})...")
        print("=" * 60)

        for name, (sem_model, vae_model, out_dim) in loaded_pipelines.items():
            print(f"\n[{name}] Stage 1: Semantic DiT 采样...")
            with torch.no_grad():
                semantic_z = sample_semantic(
                    sem_model, pe_text_embeds, flan_text_feature, flan_text_mask,
                    out_dim=out_dim,
                    num_steps=args.num_steps_semantic,
                    cfg_scale=args.cfg_scale,
                )
            print(f"  semantic_z: {semantic_z.shape}")

            print(f"[{name}] Stage 2: VAE DiT 采样...")
            with torch.no_grad():
                gen_dac = sample_vae_from_semantic(
                    vae_model, semantic_z,
                    num_steps=args.num_steps_vae,
                )
            print(f"  gen_dac: {gen_dac.shape}")

            print(f"[{name}] DAC 解码...")
            gen_audio = decode_dac(gen_dac, dac_model, device)

            # 保存（按 prompt 分目录）
            for i, prompt in enumerate(args.prompts):
                wav = gen_audio[i].cpu().float()
                wav = wav / (wav.abs().max() + 1e-8)
                dir_safe = "".join(
                    c if c.isalnum() or c in " _-" else "_" for c in prompt
                )[:80].strip()
                prompt_dir = output_dir / f"{i:02d}_{dir_safe}"
                prompt_dir.mkdir(parents=True, exist_ok=True)
                # 保存 caption
                (prompt_dir / "caption.txt").write_text(prompt)
                filepath = prompt_dir / f"{name}.wav"
                torchaudio.save(str(filepath), wav, sample_rate)
                print(f"  [{i+1}/{len(args.prompts)}] \"{prompt}\" -> {filepath}")

    # ============================================================
    # from_data 模式: 从 validation 数据取样（有 GT 对比）
    # ============================================================
    elif args.from_data:
        from data_loaders.scaled_dataset import ScaledFeatureDataset

        # 加载文本编码器（from_data 模式也需要，因为要提取文本特征用于 Semantic DiT）
        print("\n加载文本编码器...")
        pe_model, pe_transform, flan_model, flan_tokenizer = load_text_encoders(device)

        # 加载数据
        print(f"\n加载测试数据: {args.data_pattern}")
        dataset = ScaledFeatureDataset(args.data_dir, pattern=args.data_pattern)
        num_samples = min(args.num_samples, len(dataset))
        print(f"共 {len(dataset)} 样本，推理 {num_samples} 个")

        print(f"\n开始推理 (cfg_scale={args.cfg_scale}, "
              f"steps_semantic={args.num_steps_semantic}, steps_vae={args.num_steps_vae})...")
        print("=" * 60)

        for idx in tqdm(range(num_samples), desc="推理中"):
            sample = dataset[idx]

            sample_dir = output_dir / f"sample_{idx:04d}"
            sample_dir.mkdir(exist_ok=True)

            # 保存 caption
            caption = sample.get('caption', '')
            with open(sample_dir / "caption.txt", 'w') as f:
                f.write(caption)

            # GT 音频
            gt_dac = sample['dac_encoded_sampled'].unsqueeze(0).to(device)  # (1, 128, 250)
            gt_audio = decode_dac(gt_dac, dac_model, device)
            save_audio(gt_audio, sample_dir / "gt.wav", sample_rate)

            # 提取文本特征（从 caption）
            pe_text_embeds, flan_text_feature, flan_text_mask = extract_text_features(
                [caption], pe_model, pe_transform, flan_model, flan_tokenizer, device
            )

            # 各 pipeline 生成
            for name, (sem_model, vae_model, out_dim) in loaded_pipelines.items():
                with torch.no_grad():
                    # Stage 1: Text → semantic latent
                    semantic_z = sample_semantic(
                        sem_model, pe_text_embeds, flan_text_feature, flan_text_mask,
                        out_dim=out_dim,
                        num_steps=args.num_steps_semantic,
                        cfg_scale=args.cfg_scale,
                    )

                    # Stage 2: semantic latent → DAC latent
                    gen_dac = sample_vae_from_semantic(
                        vae_model, semantic_z,
                        num_steps=args.num_steps_vae,
                    )

                gen_audio = decode_dac(gen_dac, dac_model, device)
                save_audio(gen_audio, sample_dir / f"{name}.wav", sample_rate)

            print(f"  [{idx+1}/{num_samples}] caption: {caption[:80]}...")

        # 释放文本编码器
        del pe_model, flan_model
        torch.cuda.empty_cache()

    # ============================================================
    # 完成
    # ============================================================
    print(f"\n{'=' * 60}")
    print(f"✅ 推理完成！输出目录: {output_dir}")
    if args.prompts:
        print(f"生成了 {len(args.prompts)} 个 prompt × {len(loaded_pipelines)} 条管线")
    else:
        print(f"生成了 {num_samples} 个样本 × {len(loaded_pipelines)} 条管线")
        print(f"每个 sample 目录包含: caption.txt, gt.wav" +
              "".join(f", {name}.wav" for name in loaded_pipelines.keys()))


if __name__ == '__main__':
    main()