"""
Fine-tune VTXAI/tts (VITS-architecture) on Hindi speech.

This is a REAL training script:
  - real audio -> mel-spectrogram target (via torchaudio, no torchcodec needed)
  - real VITS forward pass (text encoder, posterior encoder, flow, duration predictor, decoder)
  - real losses: mel L1, KL divergence, duration loss, adversarial (LSGAN) + feature matching
  - MultiPeriodDiscriminator trained jointly (adversarial training is required for VITS to
    produce intelligible/natural audio -- without it you just get a blurry mean waveform)
  - proper batching/padding via a collate function, not batch-size-1 streaming
  - checkpointing so you can resume 1000+ hours of training across sessions

Usage:
    Paste this entire file into one Jupyter/Colab cell and run it. Edit the
    CONFIG block near the top (repo, dataset, batch size, etc.) instead of
    passing CLI flags.

Requires: torch, torchaudio, datasets, safetensors, huggingface_hub
    pip install -q torch torchaudio datasets safetensors huggingface_hub
Do NOT need torchcodec: we force decode=False and decode audio ourselves with torchaudio,
which is more memory-stable for long streaming runs anyway.
"""

import os
import io
import json
import math
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
import torchaudio
from torch.utils.data import IterableDataset, DataLoader
import sys
import importlib.util

from huggingface_hub import snapshot_download

# ==============================================================================
# ⚙️ CONFIG — edit these values directly, then just run the cell
# ==============================================================================
HF_REPO        = "VTXAI/tts"            # base model repo to download+finetune
BASE_MODEL_DIR = "./vtx_base_model"     # where the base snapshot gets cached
DATASET_NAME   = "SPRINGLab/IndicTTS-Hindi"
DATASET_SPLIT  = "train"
OUT_DIR        = "./ckpt"               # finetune checkpoints saved here
RESUME_PATH    = ""                     # set to a ./ckpt/*.pth path to resume, else ""

BATCH_SIZE   = 16
EPOCHS       = 1
MAX_STEPS    = 0        # 0 = run full epochs, otherwise stop after this many new steps
LR           = 2e-4
SAVE_EVERY   = 1000
LOG_EVERY    = 50
NUM_WORKERS  = 2
# ==============================================================================


def load_modeling(repo_id="VTXAI/tts", local_dir="./vtx_base_model"):
    """
    Ensures modeling.py is available from local repo directory or downloaded snapshot
    and imports required model classes/functions from it.
    """
    local_modeling = Path("modeling.py")
    if not local_modeling.exists():
        snapshot_dir = Path(snapshot_download(repo_id=repo_id, local_dir=local_dir))
        local_modeling = snapshot_dir / "modeling.py"

    spec = importlib.util.spec_from_file_location("vtx_modeling", local_modeling)
    vtx_modeling = importlib.util.module_from_spec(spec)
    sys.modules["vtx_modeling"] = vtx_modeling
    spec.loader.exec_module(vtx_modeling)

    return (
        vtx_modeling.SynthesizerTrn,
        vtx_modeling.MultiPeriodDiscriminator,
        vtx_modeling.kl_divergence,
        vtx_modeling.slice_segments,
    )

SynthesizerTrn, MultiPeriodDiscriminator, kl_divergence, slice_segments = load_modeling(HF_REPO, BASE_MODEL_DIR)
def fetch_base_model(repo_id="VTXAI/tts", local_dir="./vtx_base_model"):
    """
    Downloads the full VTXAI/tts repo snapshot and locates config.json + weights
    within it. Returns (config_path, weights_path).
    """
    local_dir = Path(local_dir)
    print(f"Downloading base model from https://huggingface.co/{repo_id} ...")
    repo_path = Path(snapshot_download(repo_id=repo_id, local_dir=str(local_dir)))
    print(f"Downloaded to {repo_path}")

    config_path = None
    p = repo_path / "config.json"
    if p.exists():
        config_path = p
    else:
        found = list(repo_path.rglob("config.json"))
        if not found:
            raise FileNotFoundError(f"No config.json found in {repo_id} snapshot.")
        config_path = found[0]

    weights_path = None
    for cand in ["model.safetensors", "G.safetensors", "model.pth", "G.pth"]:
        p = repo_path / cand
        if p.exists():
            weights_path = p
            break
    if weights_path is None:
        for pattern in ["*.safetensors", "*.pth"]:
            found = list(repo_path.rglob(pattern))
            if found:
                weights_path = found[0]
                break
    if weights_path is None:
        raise FileNotFoundError(f"No weights (.safetensors/.pth) found in {repo_id} snapshot.")

    print(f"Using config: {config_path}")
    print(f"Using weights: {weights_path}")
    return config_path, weights_path

# ----------------------------------------------------------------------------
# 0. maximum_path fallback
# ----------------------------------------------------------------------------
# The reference VITS repo ships a compiled `monotonic_align` extension exposing
# maximum_path(neg_cent, mask). modeling.py calls it but doesn't define it, so we
# provide a numba-free, pure-torch fallback here (slower, but correct and dependency-light).
def maximum_path(neg_cent, mask):
    """
    neg_cent: [b, t_t, t_s]  mask: [b, t_t, t_s] -> returns path [b, t_t, t_s]
    Pure-python dynamic programming, batched via a simple loop (fine at these batch sizes).
    """
    device = neg_cent.device
    dtype = neg_cent.dtype
    neg_cent = neg_cent.detach().cpu().numpy().astype(np.float64)
    mask = mask.detach().cpu().numpy().astype(np.bool_)
    b, t_t, t_s = neg_cent.shape
    path = np.zeros_like(neg_cent, dtype=np.float64)

    for i in range(b):
        t_y = int(mask[i, :, 0].sum())
        t_x = int(mask[i, 0, :].sum())
        if t_y == 0 or t_x == 0:
            continue
        value = neg_cent[i, :t_y, :t_x]
        v = np.full((t_y, t_x), -1e9, dtype=np.float64)
        v[0, 0] = value[0, 0]
        for y in range(1, t_y):
            v[y, 0] = v[y - 1, 0] + value[y, 0]
        for x in range(1, t_x):
            v[0, x] = -1e9  # can't start later in x at t_y=0
        for y in range(1, t_y):
            for x in range(1, min(t_x, y + 1)):
                v[y, x] = value[y, x] + max(v[y - 1, x], v[y - 1, x - 1])
        # backtrack
        index = t_x - 1
        for y in range(t_y - 1, -1, -1):
            path[i, y, index] = 1.0
            if index > 0 and (y == 0 or v[y - 1, index - 1] > v[y - 1, index]):
                index -= 1
    return torch.from_numpy(path).to(device=device, dtype=dtype)


# ----------------------------------------------------------------------------
# 1. Text symbols / tokenizer (Hindi, matches your existing pipeline)
# ----------------------------------------------------------------------------
_pad = "_"
_punctuation = ';:,.!?¡¿—…"«»“” '
_hindi_letters = "अआइईउऊऋएऐओऔकखगघङचछजझञटठडढणतथदधनपफबभमयरलवशषसहक्षत्रज्ञ़ािीुूृेैोौॉ्िंः"
_letters_ipa = "ɑɐɒæɓɓβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞"

hindi_symbols = [_pad] + list(_punctuation) + list(_hindi_letters) + list(_letters_ipa)
hindi_symbol_to_id = {s: i for i, s in enumerate(hindi_symbols)}


def hindi_text_to_sequence(text: str) -> list:
    seq = []
    for ch in text:
        if ch in hindi_symbol_to_id:
            seq.append(hindi_symbol_to_id[ch])
        elif ch == " ":
            seq.append(1)
    if not seq:
        seq = [1, 2, 3]
    interspersed = [0] * (len(seq) * 2 + 1)
    interspersed[1::2] = seq
    return interspersed


# ----------------------------------------------------------------------------
# 2. Mel-spectrogram (matches VITS training target: linear spec for enc_q,
#    mel spec for the reconstruction loss)
# ----------------------------------------------------------------------------
class AudioProcessor:
    def __init__(self, cfg):
        d = cfg["data"]
        self.sampling_rate = d["sampling_rate"]
        self.filter_length = d["filter_length"]
        self.hop_length = d["hop_length"]
        self.win_length = d["win_length"]
        self.n_mel_channels = d.get("n_mel_channels", 80)
        self.mel_fmin = d.get("mel_fmin", 0.0)
        self.mel_fmax = d.get("mel_fmax", None)

        self.spec_transform = torchaudio.transforms.Spectrogram(
            n_fft=self.filter_length,
            hop_length=self.hop_length,
            win_length=self.win_length,
            power=1.0,
            center=False,
            pad=(self.filter_length - self.hop_length) // 2,
        )
        self.mel_scale = torchaudio.transforms.MelScale(
            n_mels=self.n_mel_channels,
            sample_rate=self.sampling_rate,
            f_min=self.mel_fmin,
            f_max=self.mel_fmax,
            n_stft=self.filter_length // 2 + 1,
        )

    def to(self, device):
        self.spec_transform = self.spec_transform.to(device)
        self.mel_scale = self.mel_scale.to(device)
        return self

    def wav_to_spec(self, wav):
        # wav: [1, T] float32 in [-1, 1] -> spec [C, F]
        spec = self.spec_transform(wav)
        return spec.squeeze(0)

    def spec_to_mel(self, spec):
        if spec.device != self.mel_scale.fb.device:
            self.mel_scale = self.mel_scale.to(spec.device)
        return self.mel_scale(spec.unsqueeze(0)).squeeze(0)

    def mel_scale_transform(self, spec):
        if spec.device != self.mel_scale.fb.device:
            self.mel_scale = self.mel_scale.to(spec.device)
        return self.mel_scale(spec)

    def wav_to_mel(self, wav):
        return self.spec_to_mel(self.wav_to_spec(wav))


# ----------------------------------------------------------------------------
# 3. Streaming dataset over the HF Hindi corpus, with real audio decoding
# ----------------------------------------------------------------------------
class HindiTTSIterable(IterableDataset):
    """
    Wraps datasets.load_dataset(..., streaming=True). We force decode=False on the
    audio column and decode manually with torchaudio -- this avoids the torchcodec
    dependency entirely and is more robust for very long (1000+ hr) streaming runs.
    """

    def __init__(self, dataset_name, split, ap: AudioProcessor, max_audio_sec=12.0, text_key=None):
        super().__init__()
        self.dataset_name = dataset_name
        self.split = split
        self.ap = ap
        self.max_samples = int(max_audio_sec * ap.sampling_rate)
        self.text_key = text_key

    def _load_stream(self):
        from datasets import load_dataset, Audio
        ds = load_dataset(self.dataset_name, split=self.split, streaming=True)
        ds = ds.cast_column("audio", Audio(decode=False))
        return ds

    def __iter__(self):
        ds = self._load_stream()
        for item in ds:
            try:
                audio_field = item["audio"]
                raw_bytes = audio_field["bytes"] if isinstance(audio_field, dict) else None
                if raw_bytes is not None:
                    wav, sr = torchaudio.load(io.BytesIO(raw_bytes))
                elif isinstance(audio_field, dict) and audio_field.get("path"):
                    wav, sr = torchaudio.load(audio_field["path"])
                else:
                    continue

                if wav.shape[0] > 1:
                    wav = wav.mean(0, keepdim=True)
                if sr != self.ap.sampling_rate:
                    wav = torchaudio.functional.resample(wav, sr, self.ap.sampling_rate)
                if wav.shape[1] > self.max_samples:
                    continue  # skip overly long clips for stable batch memory
                if wav.shape[1] < self.ap.hop_length * 8:
                    continue

                text_key = self.text_key
                if text_key is None:
                    text_key = "normalized_text" if "normalized_text" in item else "text"
                text = item.get(text_key, "")
                if not text or not text.strip():
                    continue

                tok = torch.LongTensor(hindi_text_to_sequence(text))
                yield tok, wav.squeeze(0)
            except Exception:
                # skip malformed rows rather than crashing a 1000hr run
                continue


def collate_fn(batch):
    batch = [b for b in batch if b[0].numel() > 0 and b[1].numel() > 0]
    batch.sort(key=lambda x: x[1].shape[0], reverse=True)
    text_lens = torch.LongTensor([b[0].shape[0] for b in batch])
    wav_lens = torch.LongTensor([b[1].shape[0] for b in batch])

    text_pad = torch.zeros(len(batch), text_lens.max(), dtype=torch.long)
    wav_pad = torch.zeros(len(batch), wav_lens.max(), dtype=torch.float32)

    for i, (tok, wav) in enumerate(batch):
        text_pad[i, : tok.shape[0]] = tok
        wav_pad[i, : wav.shape[0]] = wav

    return text_pad, text_lens, wav_pad, wav_lens


# ----------------------------------------------------------------------------
# 4. Loss functions
# ----------------------------------------------------------------------------
def discriminator_loss(disc_real_outputs, disc_generated_outputs):
    loss = 0
    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
        r_loss = torch.mean((1 - dr) ** 2)
        g_loss = torch.mean(dg ** 2)
        loss += r_loss + g_loss
    return loss


def generator_adv_loss(disc_outputs):
    loss = 0
    for dg in disc_outputs:
        loss += torch.mean((1 - dg) ** 2)
    return loss


def feature_loss(fmap_r, fmap_g):
    loss = 0
    for dr, dg in zip(fmap_r, fmap_g):
        for rl, gl in zip(dr, dg):
            loss += torch.mean(torch.abs(rl.float() - gl.float()))
    return loss * 2


# ----------------------------------------------------------------------------
# 5. Checkpoint loading (keeps your existing "transfer weights, resize vocab" fix)
# ----------------------------------------------------------------------------
def load_pretrained_generator(model, ckpt_path):
    ckpt_path = Path(ckpt_path)
    if ckpt_path.suffix == ".safetensors":
        from safetensors.torch import load_file
        state_dict = load_file(str(ckpt_path))
    else:
        raw = torch.load(ckpt_path, map_location="cpu")
        state_dict = raw.get("model", raw)

    # Drop embedding weights: symbol vocab size for Hindi differs from source checkpoint
    filtered = {
        k: v for k, v in state_dict.items()
        if not (k.endswith("emb.weight") or k.startswith("enc_p.emb"))
    }
    model_dict = model.state_dict()
    compatible = {k: v for k, v in filtered.items() if k in model_dict and model_dict[k].shape == v.shape}
    skipped = [k for k in filtered if k not in compatible]
    model_dict.update(compatible)
    model.load_state_dict(model_dict, strict=False)
    print(f"Loaded {len(compatible)}/{len(state_dict)} tensors from checkpoint "
          f"(skipped {len(skipped)}, incl. resized text embedding).")


# ----------------------------------------------------------------------------
# 6. Training run (executes directly — paste this whole file into one notebook cell)
# ----------------------------------------------------------------------------
def main():
    out_dir = Path(OUT_DIR)
    out_dir.mkdir(parents=True, exist_ok=True)

    config_path, checkpoint_path = fetch_base_model(HF_REPO, BASE_MODEL_DIR)

    with open(config_path) as f:
        config = json.load(f)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    ap = AudioProcessor(config).to(device)

    net_g = SynthesizerTrn(
        len(hindi_symbols),
        config["data"]["filter_length"] // 2 + 1,
        config["train"]["segment_size"] // config["data"]["hop_length"],
        **config["model"],
    ).to(device)
    net_d = MultiPeriodDiscriminator(config["model"].get("use_spectral_norm", False)).to(device)

    start_step = 0
    if RESUME_PATH and Path(RESUME_PATH).exists():
        ckpt = torch.load(RESUME_PATH, map_location="cpu")
        net_g.load_state_dict(ckpt["net_g"])
        net_d.load_state_dict(ckpt["net_d"])
        start_step = ckpt.get("step", 0)
        print(f"Resumed from {RESUME_PATH} at step {start_step}")
    else:
        load_pretrained_generator(net_g, checkpoint_path)

    optim_g = torch.optim.AdamW(net_g.parameters(), lr=LR, betas=(0.8, 0.99), eps=1e-9)
    optim_d = torch.optim.AdamW(net_d.parameters(), lr=LR, betas=(0.8, 0.99), eps=1e-9)
    if RESUME_PATH and Path(RESUME_PATH).exists():
        optim_g.load_state_dict(ckpt["optim_g"])
        optim_d.load_state_dict(ckpt["optim_d"])

    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=0.999875)
    scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=0.999875)

    dataset = HindiTTSIterable(DATASET_NAME, DATASET_SPLIT, ap)
    loader = DataLoader(
        dataset, batch_size=BATCH_SIZE, collate_fn=collate_fn,
        num_workers=NUM_WORKERS, pin_memory=(device.type == "cuda"),
    )

    hop = config["data"]["hop_length"]
    seg_size = config["train"]["segment_size"]
    seg_frames = seg_size // hop
    c_mel = config["train"].get("c_mel", 45.0)
    c_kl = config["train"].get("c_kl", 1.0)

    step = start_step
    t0 = time.time()
    net_g.train()
    net_d.train()

    for epoch in range(EPOCHS):
        for text_pad, text_lens, wav_pad, wav_lens in loader:
            text_pad, text_lens = text_pad.to(device), text_lens.to(device)
            wav_pad, wav_lens = wav_pad.to(device), wav_lens.to(device)

            # Build linear spectrogram targets per-sample (variable length), then pad
            specs, spec_lens = [], []
            for i in range(wav_pad.shape[0]):
                w = wav_pad[i, : wav_lens[i]].unsqueeze(0).cpu()
                s = ap.wav_to_spec(w)
                specs.append(s)
                spec_lens.append(s.shape[-1])
            spec_lens = torch.LongTensor(spec_lens).to(device)
            spec_pad = torch.zeros(len(specs), specs[0].shape[0], spec_lens.max(), device=device)
            for i, s in enumerate(specs):
                spec_pad[i, :, : s.shape[-1]] = s.to(device)

            # ---- Generator forward ----
            (y_hat, l_length, attn, ids_slice, x_mask, y_mask,
             (z, z_p, m_p, logs_p, m_q, logs_q)) = net_g(text_pad, text_lens, spec_pad, spec_lens)

            mel_pad = ap.mel_scale_transform(spec_pad)
            mel_slice = slice_segments(mel_pad, ids_slice, seg_frames)
            with torch.no_grad():
                wav_hat_mono = y_hat.squeeze(1)
            y_hat_mel = torch.stack([
                ap.spec_to_mel(ap.wav_to_spec(wav_hat_mono[i:i + 1].detach().cpu())).to(device)
                for i in range(wav_hat_mono.shape[0])
            ])
            y_slice = slice_segments(wav_pad.unsqueeze(1), ids_slice * hop, seg_size)

            # ---- Discriminator step ----
            y_d_hat_r, y_d_hat_g, _, _ = net_d(y_slice, y_hat.detach())
            loss_disc = discriminator_loss(y_d_hat_r, y_d_hat_g)
            optim_d.zero_grad()
            loss_disc.backward()
            torch.nn.utils.clip_grad_norm_(net_d.parameters(), 5.0)
            optim_d.step()

            # ---- Generator step ----
            y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(y_slice, y_hat)
            loss_mel = F.l1_loss(mel_slice, y_hat_mel) * c_mel
            loss_kl = torch.mean(kl_divergence(m_p, logs_p, m_q, logs_q)) * c_kl
            loss_dur = torch.sum(l_length.float())
            loss_fm = feature_loss(fmap_r, fmap_g)
            loss_adv = generator_adv_loss(y_d_hat_g)
            loss_gen_total = loss_adv + loss_fm + loss_mel + loss_kl + loss_dur

            optim_g.zero_grad()
            loss_gen_total.backward()
            torch.nn.utils.clip_grad_norm_(net_g.parameters(), 5.0)
            optim_g.step()

            step += 1
            if step % LOG_EVERY == 0:
                elapsed = time.time() - t0
                print(f"[step {step}] gen={loss_gen_total.item():.4f} "
                      f"mel={loss_mel.item():.4f} kl={loss_kl.item():.4f} "
                      f"dur={loss_dur.item():.4f} adv={loss_adv.item():.4f} "
                      f"fm={loss_fm.item():.4f} disc={loss_disc.item():.4f} "
                      f"({elapsed:.1f}s elapsed)")

            if step % SAVE_EVERY == 0:
                ckpt_path = out_dir / f"vtx_hindi_step{step}.pth"
                torch.save({
                    "net_g": net_g.state_dict(),
                    "net_d": net_d.state_dict(),
                    "optim_g": optim_g.state_dict(),
                    "optim_d": optim_d.state_dict(),
                    "step": step,
                    "symbols": hindi_symbols,
                }, ckpt_path)
                print(f"Saved checkpoint: {ckpt_path}")

            if MAX_STEPS and step >= start_step + MAX_STEPS:
                break
        scheduler_g.step()
        scheduler_d.step()
        if MAX_STEPS and step >= start_step + MAX_STEPS:
            break

    final_path = out_dir / "vtx_hindi_final.pth"
    torch.save({"net_g": net_g.state_dict(), "net_d": net_d.state_dict(), "step": step,
                "symbols": hindi_symbols}, final_path)
    print(f"Done. Final checkpoint: {final_path}")


if __name__ == "__main__":
    main()