"""
morph.py — core RAVE latent interpolation engine

Encodes two audio files into RAVE latent space,
interpolates between them in N steps, decodes each step
back to audio, and returns the audio arrays + a latent
path figure.
"""

import copy
import os
import tempfile
import numpy as np
import torch
import librosa
import soundfile as sf
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA


# ── Model loading ─────────────────────────────────────────────

_model_cache = {}

def load_model(model_path: str):
    """Load and cache a TorchScript RAVE model."""
    if model_path not in _model_cache:
        model = torch.jit.load(model_path, map_location="cpu")
        model.eval()
        _model_cache[model_path] = model
    return _model_cache[model_path]


# ── Audio loading / encoding ───────────────────────────────────

MIN_SAMPLES = 4096


def load_audio_segment(
    audio_path: str,
    target_sr: int = 48000,
    duration_sec: float | None = None,
) -> np.ndarray:
    """Load mono audio, trim silence, optionally cap length."""
    audio, _ = librosa.load(audio_path, sr=target_sr, mono=True)
    audio, _ = librosa.effects.trim(audio, top_db=30)
    if duration_sec is not None:
        audio = audio[: int(duration_sec * target_sr)]
    return audio


def prepare_audio_pair(
    audio_a_path: str,
    audio_b_path: str,
    target_sr: int = 48000,
    duration_sec: float | None = 8.0,
) -> tuple[np.ndarray, np.ndarray, float]:
    """
    Load two files and crop them to the same length so the morph
    isn't silently truncated to a tiny hit from the shorter input.
    """
    audio_a = load_audio_segment(audio_a_path, target_sr, duration_sec)
    audio_b = load_audio_segment(audio_b_path, target_sr, duration_sec)
    n = min(len(audio_a), len(audio_b))
    if n < MIN_SAMPLES:
        raise ValueError(
            f"Aligned audio is only {n / target_sr:.2f}s after trimming. "
            "Use longer clips (5–15 s each) with similar length."
        )
    return audio_a[:n], audio_b[:n], n / target_sr


def encode_waveform(model, audio: np.ndarray) -> torch.Tensor:
    """Encode a mono waveform array with RAVE. Returns [1, D, T']."""
    if len(audio) < MIN_SAMPLES:
        audio = np.pad(audio, (0, MIN_SAMPLES - len(audio)))

    x = torch.from_numpy(audio).float().unsqueeze(0).unsqueeze(0)
    with torch.no_grad():
        z = model.encode(x)
    return z


def encode_audio(model, audio_path: str, target_sr: int = 48000) -> torch.Tensor:
    """Load a WAV/MP3 and encode with RAVE."""
    audio, _ = librosa.load(audio_path, sr=target_sr, mono=True)
    audio, _ = librosa.effects.trim(audio, top_db=30)
    return encode_waveform(model, audio)


# ── Interpolation ─────────────────────────────────────────────

def slerp(z1: torch.Tensor, z2: torch.Tensor, alpha: float) -> torch.Tensor:
    """
    Spherical interpolation per latent time frame, preserving magnitude.
  """
    frames = []
    for t in range(z1.shape[-1]):
        v1 = z1[0, :, t]
        v2 = z2[0, :, t]
        n1 = v1.norm()
        n2 = v2.norm()
        if n1 < 1e-8 or n2 < 1e-8:
            frames.append(((1 - alpha) * v1 + alpha * v2).unsqueeze(-1))
            continue

        u1 = v1 / n1
        u2 = v2 / n2
        dot = (u1 * u2).sum().clamp(-1 + 1e-6, 1 - 1e-6)
        theta = torch.acos(dot)
        sin_theta = torch.sin(theta)
        if sin_theta.abs() < 1e-6:
            direction = (1 - alpha) * u1 + alpha * u2
        else:
            direction = (
                (torch.sin((1 - alpha) * theta) / sin_theta) * u1
                + (torch.sin(alpha * theta) / sin_theta) * u2
            )
        mag = (1 - alpha) * n1 + alpha * n2
        frames.append((direction / (direction.norm() + 1e-8) * mag).unsqueeze(-1))

    return torch.cat(frames, dim=-1).unsqueeze(0)


def interpolate_latents(
    z1: torch.Tensor,
    z2: torch.Tensor,
    steps: int = 12,
    use_slerp: bool = False
) -> list[torch.Tensor]:
    """
    Return a list of `steps` latent tensors from z1 → z2.
    Handles mismatched time dimensions by trimming to shorter.
    """
    # Align time dimension (T') — trim to shortest
    t = min(z1.shape[-1], z2.shape[-1])
    z1 = z1[..., :t]
    z2 = z2[..., :t]

    alphas = np.linspace(0.0, 1.0, steps)
    latents = []
    for a in alphas:
        if use_slerp:
            z = slerp(z1, z2, float(a))
        else:
            z = (1 - a) * z1 + a * z2
        latents.append(z)
    return latents


# ── Decoding ──────────────────────────────────────────────────

def decode_latents(model, latents: list[torch.Tensor]) -> list[np.ndarray]:
    """
    Decode each latent tensor back to audio.

    RAVE's exported decoder keeps streaming state in cached conv layers,
    so we deep-copy the model before each decode to avoid silent/NaN output.
    """
    audio_arrays = []
    with torch.no_grad():
        for z in latents:
            decoder = copy.deepcopy(model)
            decoder.eval()
            audio = decoder.decode(z)
            arr = audio.squeeze().cpu().numpy()
            if np.isnan(arr).any():
                raise RuntimeError(
                    "Decode produced invalid audio. Try linear interpolation "
                    "or a different model for these inputs."
                )
            peak = np.abs(arr).max()
            if peak > 0:
                arr = arr / peak * 0.9
            audio_arrays.append(arr)
    return audio_arrays


def run_morph(
    model_a_path: str,
    model_b_path: str,
    audio_a_path: str,
    audio_b_path: str,
    steps: int = 12,
    use_slerp: bool = False,
    duration_sec: float | None = 8.0,
    target_sr: int = 48000,
) -> tuple[list[np.ndarray], list[torch.Tensor], float, int]:
    """
    Full encode → interpolate → decode pipeline.

    Sound A is encoded with model_a, Sound B with model_b.
    Interpolated latents are decoded with model_b (toward B's space).

    Returns (audio_arrays, latents, segment_duration_sec, target_sr).
    """
    model_a = load_model(model_a_path)
    model_b = load_model(model_b_path)

    audio_a, audio_b, segment_sec = prepare_audio_pair(
        audio_a_path, audio_b_path, target_sr, duration_sec
    )

    with tempfile.TemporaryDirectory() as tmp:
        aligned_a = os.path.join(tmp, "aligned_a.wav")
        aligned_b = os.path.join(tmp, "aligned_b.wav")
        sf.write(aligned_a, audio_a, target_sr)
        sf.write(aligned_b, audio_b, target_sr)
        z1 = encode_audio(model_a, aligned_a, target_sr)
        z2 = encode_audio(model_b, aligned_b, target_sr)

    latents = interpolate_latents(z1, z2, steps=steps, use_slerp=use_slerp)
    arrays = decode_latents(model_b, latents)
    return arrays, latents, segment_sec, target_sr


# ── Save outputs ──────────────────────────────────────────────

def save_morphed_audio(
    arrays: list[np.ndarray],
    sr: int,
    out_dir: str
) -> list[str]:
    """Save each decoded array as a numbered WAV. Returns paths."""
    os.makedirs(out_dir, exist_ok=True)
    paths = []
    for i, arr in enumerate(arrays):
        path = os.path.join(out_dir, f"morph_step_{i:02d}.wav")
        sf.write(path, arr, sr)
        paths.append(path)
    return paths


# ── Latent path visualisation ─────────────────────────────────

def plot_latent_path(latents: list[torch.Tensor]) -> plt.Figure:
    """
    PCA-project all latents to 2D and draw the interpolation
    path. Dots go from blue (Sound A) to orange (Sound B).
    """
    # Flatten each latent to a 1D vector for PCA
    vecs = [z.squeeze().mean(dim=-1).cpu().numpy() for z in latents]
    vecs = np.stack(vecs)  # [steps, D]

    n = len(vecs)
    if vecs.shape[1] >= 2:
        pca = PCA(n_components=2)
        proj = pca.fit_transform(vecs)
    else:
        # Edge case: single latent dim — use step index as X
        proj = np.column_stack([np.arange(n), vecs[:, 0]])

    fig, ax = plt.subplots(figsize=(5, 4))
    fig.patch.set_facecolor("#f8f8f8")
    ax.set_facecolor("#f8f8f8")

    # Draw connecting line
    ax.plot(proj[:, 0], proj[:, 1], color="#cccccc", linewidth=1.5, zorder=1)

    # Color each dot along a blue → orange gradient
    colors = plt.cm.coolwarm(np.linspace(0, 1, n))
    for i, (x, y) in enumerate(proj):
        ax.scatter(x, y, color=colors[i], s=60, zorder=2, edgecolors="white", linewidths=0.5)

    # Label endpoints
    ax.annotate("A", proj[0], fontsize=11, fontweight="500",
                xytext=(-12, 4), textcoords="offset points", color="#2255aa")
    ax.annotate("B", proj[-1], fontsize=11, fontweight="500",
                xytext=(6, 4), textcoords="offset points", color="#cc4400")

    ax.set_title("Latent path (PCA)", fontsize=12, fontweight="500", pad=10)
    ax.set_xlabel("PC 1", fontsize=10, color="#666")
    ax.set_ylabel("PC 2", fontsize=10, color="#666")
    ax.tick_params(colors="#999", labelsize=9)
    for spine in ax.spines.values():
        spine.set_edgecolor("#dddddd")

    fig.tight_layout()
    return fig


# ── CLI entry point ───────────────────────────────────────────

if __name__ == "__main__":
    import sys

    if len(sys.argv) < 5:
        print(
            "Usage: python morph.py <model_a.ts> <model_b.ts> "
            "<audio_a.wav> <audio_b.wav> [steps] [duration_sec]\n"
            "\n"
            "Example:\n"
            "  python morph.py models/guitar.ts models/organ_archive.ts \\\n"
            "    samples/356390__mtg__violin-d-major.wav \\\n"
            "    samples/356148__mtg__violin-e6-bad-dynamics-tremolo.wav \\\n"
            "    12 8"
        )
        sys.exit(1)

    model_a_path = sys.argv[1]
    model_b_path = sys.argv[2]
    audio_a      = sys.argv[3]
    audio_b      = sys.argv[4]
    steps        = int(sys.argv[5]) if len(sys.argv) > 5 else 12
    duration     = float(sys.argv[6]) if len(sys.argv) > 6 else 8.0

    print(f"Loading model A: {model_a_path}")
    print(f"Loading model B: {model_b_path}")
    print(f"Encoding {audio_a} with model A...")
    print(f"Encoding {audio_b} with model B...")
    print(f"Decoding with model B...")
    print(f"Using {duration:.1f}s segment (capped to shorter input)...")

    arrays, latents, segment_sec, sr = run_morph(
        model_a_path,
        model_b_path,
        audio_a,
        audio_b,
        steps=steps,
        duration_sec=duration,
    )
    print(f"Segment length: {segment_sec:.2f}s")

    print(f"Interpolating {steps} steps...")
    print("Decoding...")

    out_dir = os.path.join(os.path.dirname(audio_a), "output")
    paths = save_morphed_audio(arrays, sr=sr, out_dir=out_dir)
    print(f"Saved {len(paths)} files to {out_dir}/")

    fig = plot_latent_path(latents)
    fig_path = os.path.join(out_dir, "latent_path.png")
    fig.savefig(fig_path, dpi=150)
    print(f"Latent path plot saved to {fig_path}")