# removed torchvision, torchaudio, safetensors, yaml, math, json, sys, io, warnings,
# random, T5Tokenizer, T5EncoderModel, tqdm, MelSpectrogram. not needed for inference.

import librosa
import torch
import accelerate
import numpy as np

from librosa.feature import chroma_stft

from anyaccomp.fmt_model import FlowMatchingTransformerConcat
from models.codec.amphion_codec.vocos import Vocos
from models.codec.coco.rep_coco_model import CocoStyle

from utils.util import load_config


class Sing2SongInferencePipeline:
    """
    Wraps the three model components needed for inference:
      1. CocoStyle  – encodes the vocal chromagram into discrete style tokens
      2. FlowMatchingTransformerConcat – diffuses those tokens into a mel spectrogram
      3. Vocos  – decodes the mel spectrogram into a waveform
    """

    def __init__(
        self,
        checkpoint_path,
        cfg_path,
        vocoder_checkpoint_path,
        vocoder_cfg_path,
        device="cuda",
    ):
        self.cfg = load_config(cfg_path)
        self.device = device

        self.checkpoint_path = checkpoint_path
        self._load_model(checkpoint_path)       # flow matching transformer

        self._build_input_model()               # chromagram encoder (CocoStyle)
        self.vocoder_checkpoint_path = vocoder_checkpoint_path
        self.vocoder_cfg = load_config(vocoder_cfg_path)
        self._build_output_model()              # vocoder (Vocos)
        print("Output model built")

    def _load_model(self, checkpoint_path):
        self.model = FlowMatchingTransformerConcat(
            cfg=self.cfg.model.flow_matching_transformer
        )

        accelerate.load_checkpoint_and_dispatch(self.model, checkpoint_path)
        self.model.eval().to(self.device)
        print(
            f"model Params: {round(sum(p.numel() for p in self.model.parameters() if p.requires_grad)/1e6, 2)}M"
        )
        print(f"Loaded model from {checkpoint_path}")

    def _build_input_model(self):
        # construct_only_for_quantizer=True skips building the decoder half of CocoStyle —
        # we only need the encoder + quantizer for inference.
        self.coco_model = CocoStyle(
            cfg=self.cfg.model.coco, construct_only_for_quantizer=True
        )
        self.coco_model.eval()
        self.coco_model.to(self.device)
        accelerate.load_checkpoint_and_dispatch(
            self.coco_model, self.cfg.model.coco.pretrained_path
        )

    def _build_output_model(self):
        self.vocoder = Vocos(cfg=self.vocoder_cfg.model.vocos)
        accelerate.load_checkpoint_and_dispatch(
            self.vocoder, self.vocoder_checkpoint_path
        )
        self.vocoder = self.vocoder.eval().to(self.device)

    @torch.no_grad()
    @torch.amp.autocast("cuda", dtype=torch.bfloat16)
    def _extract_coco_codec(self, speech):
        """Compute a chromagram from the waveform, then quantize it into discrete tokens."""
        target_chroma_dim = self.cfg.model.coco.chromagram_dim

        speech = speech.cpu().numpy().squeeze()

        # librosa returns [n_chroma, T]; transpose to [T, n_chroma] for the model.
        chromagram = chroma_stft(
            y=speech,
            sr=self.cfg.preprocess.chromagram.sample_rate,
            n_fft=self.cfg.preprocess.chromagram.n_fft,
            hop_length=self.cfg.preprocess.chromagram.hop_size,
            win_length=self.cfg.preprocess.chromagram.win_size,
            n_chroma=target_chroma_dim,
        ).T
        chromagram_feats = torch.tensor(chromagram).unsqueeze(0).to(self.device)
        codecs, _ = self.coco_model.quantize(chromagram_feats)
        return codecs

    @torch.no_grad()
    def encode_vocal(self, speech):  # (B, T)
        speech = speech.to(self.device)
        return self._extract_coco_codec(speech)

    @torch.no_grad()
    def _generate_audio(self, mel):
        # mel is [B, T, C]; vocoder expects [B, C, T].
        return (self.vocoder(mel.transpose(1, 2)).detach().cpu())[0]