#!/usr/bin/env python3
"""
Test inference script for FA3 environment
Generates song with custom postfix
"""

import os
import sys
import torch
from pathlib import Path

# Set GPU
os.environ["CUDA_VISIBLE_DEVICES"] = "2"  # Use different GPU from FA2 test

# For reproducibility
torch.manual_seed(42)

# Test Flash Attention v3 version
try:
    import flash_attn_interface
    print(f"Flash Attention v3 successfully imported from: {flash_attn_interface.__file__}")
except ImportError as e:
    print(f"Flash Attention v3 not available: {e}")
    # Try fallback to v2
    try:
        import flash_attn
        print(f"Flash Attention v2 fallback: {flash_attn.__version__}")
    except ImportError:
        print("No Flash Attention available")

# Import Suno utilities
from suno_utils.gpt.generation import load_model, GPT
from suno_utils.utils.s3 import download_s3_file_if_needed
from suno_utils.audio import Audio
from suno_utils.diffusion import generation as diffusion_gen
from suno_utils.tasks.upsample_engine import UpsampleEngine, Request
from suno_utils.tasks.dac_vae_fixed_25hz import preload_models as preload_codec_models, decode

# BCT imports
from suno_utils.gpt.bct.bct_generation_simple import BCTGenerationConfig, BlockSequence, generate_block
from suno_utils.gpt.bct.bct import Block, BlockType, TensorDict


def setup_models():
    """Load all necessary models for inference."""
    print("Loading GPT model...")
    
    # Model paths
    gpt_model_path = "/app2/suno/checkpoints/2025-08-24_08-31-23/last_ckpt_infer.pt"  
    tokenizer_path = "s3://suno-data/georg/models/tokenizers/tokenizer_60k.json"
    
    # Load GPT model
    model_container = load_model(
        ckpt_path=download_s3_file_if_needed(gpt_model_path),
        tokenizer_path=download_s3_file_if_needed(tokenizer_path),
    )
    model = model_container["model"]
    assert isinstance(model, GPT)
    cfg = model.config
    print(f"Model loaded: {cfg.n_layer} layers, {cfg.n_head} heads, {cfg.n_embd} dimensions")
    
    # Load diffusion model
    print("Loading diffusion model...")
    dit_model_filepath = "/app/suno/checkpoints/2025-02-17_16-54-01_s7787/last_ckpt.pt"
    diffusion_gen.preload_models(dit_model_filepath=dit_model_filepath)
    
    # Load codec model
    print("Loading codec model...")
    preload_codec_models("s3://suno-data/minz/models/dac_vae_tuned_25hz.pth")
    
    # Create diffusion engine
    diffusion_engine = UpsampleEngine(compile=False)
    print("✅ All models loaded successfully")
    
    return model, model_container, cfg, diffusion_engine


def setup_block_types():
    """Define block types for BCT generation."""
    TextBlockType = BlockType(name="text", is_causal=True)
    CausalSemanticBlockType = BlockType(name="semantic", is_causal=True)
    return TextBlockType, CausalSemanticBlockType


def make_text_block(text: str, model_container, cfg, TextBlockType):
    """Create text conditioning block."""
    text_tokens = model_container["tokenizer"].encode(text) + [cfg.text_infer_token]
    return Block(
        TextBlockType,
        inputs=TensorDict({
            "text_input": torch.tensor(text_tokens).reshape(1, 1, -1),
        }),
    )


def run_diffusion(codes, diffusion_engine, lyrics="", cfg_scale=2.0):
    """Convert semantic codes to audio via diffusion."""
    gen_cfg = diffusion_gen.DiffusionGenerationConfig(
        lyrics=lyrics,
        text_cfg_coef=cfg_scale,
        ctx_cfg_coef=1.0,
        steps=12,
        codec_scale_factor=0.4,
        scale_ctx_vector=True,
    )
    request = Request(
        id="lyrics2song",
        generation_config=gen_cfg,
        tokens=codes.cpu(),
        input_tokens_finished=True,
    )
    result = diffusion_engine.run_request(request)
    vae_latents = torch.concat(result.vae_latents)
    audio = decode(vae_latents)
    return audio


def run_inference(blocks, model, cfg, max_steps=1000, text_cfg_boost=0.3, return_raw=False):
    """Run inference with BCT generation."""
    if text_cfg_boost == 0.0 or blocks[0].spec.name != "text":
        prompts = [(1, blocks)]
    else:
        no_text_blocks = BlockSequence(blocks[1:])
        prompts = [(1 + text_cfg_boost, blocks), (-text_cfg_boost, no_text_blocks)]
    
    gconf = BCTGenerationConfig(
        prompts,
        max_autoregressive_steps=max_steps,
        eos_token=cfg.semantic_pad_token,
        temperature=0.9,
        compile=False,
    )
    
    block = generate_block(model, gconf)
    all_sem_codes = block.inputs["semantic_input"][0, 0, 1 : len(block)]
    
    if return_raw:
        return all_sem_codes
    
    return all_sem_codes


def generate_song_from_lyrics(lyrics, model, model_container, cfg, diffusion_engine, 
                              TextBlockType, CausalSemanticBlockType,
                              max_steps=1000, text_cfg_boost=0.3):
    """Generate a song from lyrics text."""
    
    # Create text block from lyrics
    text_block = make_text_block(lyrics, model_container, cfg, TextBlockType)
    
    # Create semantic block
    sem_block = Block(
        CausalSemanticBlockType,
        inputs=TensorDict({
            "semantic_input": torch.full((1, 1, 1), cfg.semantic_infer_token),
        }),
    )
    
    # Create block sequence
    blocks = BlockSequence([text_block, sem_block])
    
    # Run inference to get semantic codes
    print("Generating semantic codes...")
    codes = run_inference(blocks, model, cfg, max_steps=max_steps, 
                         text_cfg_boost=text_cfg_boost, return_raw=True)
    
    # Run diffusion to generate audio
    print("Running diffusion to generate audio...")
    audio = run_diffusion(codes, diffusion_engine, lyrics=lyrics)
    
    return audio


def main():
    """Main function to generate song from lyrics."""
    
    print("=" * 60)
    print("Testing Flash Attention v3 Environment")
    print("=" * 60)
    
    # Example lyrics with timing annotations - different lyrics for FA3 test
    lyrics = """    
{start_offset:125}


[verse]
[5.0]In the neon lights[8.0]
[8.1]Where dreams take flight[15.0]
[15.1]Electric pulses flow[20.0]
[20.1]Through circuits we don't know[25.0]
[25.1]Digital hearts beating slow[30.0]

[chorus]
[40.0]We're running on H100 power
[40.1]Flash attention every hour[45.0]
[45.1]Version three is what we need[50.0]
[50.1]Maximum performance speed[55.0]
[55.1]Future tech is here indeed[60.0]

[versef]
[80.0]In the silicon sky
where algorithms fly
processing thoughts divine
in parallel design
optimized and fine
"""
    
    # Setup models
    model, model_container, cfg, diffusion_engine = setup_models()
    
    # Setup block types
    TextBlockType, CausalSemanticBlockType = setup_block_types()
    
    # Generate song
    print("\nGenerating song from lyrics (FA3 environment)...")
    print("=" * 50)
    print("Lyrics preview:")
    print(lyrics[:200] + "...")
    print("=" * 50)
    
    audio = generate_song_from_lyrics(
        lyrics=lyrics,
        model=model,
        model_container=model_container,
        cfg=cfg,
        diffusion_engine=diffusion_engine,
        TextBlockType=TextBlockType,
        CausalSemanticBlockType=CausalSemanticBlockType,
        max_steps=1000,
        text_cfg_boost=0.3
    )
    
    # Save the generated audio with custom postfix for FA3
    output_path = "generated_song_fa3.mp3"
    print(f"\nSaving generated song to {output_path}")
    audio.write_mp3(output_path)
    
    # Play the audio (if in an environment that supports it)
    try:
        audio.play()
    except:
        print("Audio playback not available in this environment")
    
    print("\n✅ Song generation complete with FA3 environment!")
    return audio


if __name__ == "__main__":
    audio = main()