#!/usr/bin/env python3
"""
Standalone inference script for generating songs from lyrics using Suno GPT model.
Extracted from inference_notebook_original.ipynb
"""

import os
import torch
from pathlib import Path

# Set GPU
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

# For reproducibility
torch.manual_seed(42)

# 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 (Block-based Conditioning Transformer) 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"  # DPO model
    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."""
    
    # Example lyrics with timing annotations
    lyrics = """    
{start_offset:125}


[verse]
[5.0]oh, my love[8.0]
[8.1]My friend you know[15.0]
[15.1]it's been a while[20.0]
[20.1]Without thinking of you[25.0]
[25.1]but the thought makes me smile[30.0]

[chorus]
[40.0]I'm so tired of wanting
[40.1]wanting more than this[45.0]
[45.1]i know it but what am i to do[50.0]
[50.1]i need some space to breathe,[55.0]
[55.1]so give me some room[60.0]

[versef]
[80.0]oh, my love
you have a heart of stone
cause since i've come home
i've never felt so alone
but the thought makes me smile
"""
    
    # 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...")
    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
    output_path = "generated_song.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!")
    return audio


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