from __future__ import annotations

from functools import lru_cache
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio.functional as AF
from huggingface_hub import hf_hub_download
from transformers import AutoModel, Wav2Vec2FeatureExtractor


SAMPLE_RATE = 24_000
CLIP_SECONDS = 10
NUM_SAMPLES = SAMPLE_RATE * CLIP_SECONDS
EXTRACT_LAYERS = (3, 4, 5, 6, 23)

MERT_REPO = "m-a-p/MERT-v1-330M"
MERT_REVISION = "5240c2708a5acaee1007f43fb9735c7dcd0b78c9"
MERIT_REPO = "amaai-lab/merit"
MERIT_REVISION = "a85df30eca1ba112eb594285f3ba1d96488e7883"
HEAD_FILES = {
    "melody": "head_mel/best_head.pt",
    "rhythm": "head_rhy/best_head.pt",
    "timbre": "head_tim/best_head.pt",
}


def _select_device() -> str:
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


class ProjectionHead(nn.Module):
    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, out_dim, bias=False),
        )

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        return F.normalize(self.net(features), dim=-1)


def _load_head(filename: str, device: torch.device) -> ProjectionHead:
    checkpoint_path = hf_hub_download(
        repo_id=MERIT_REPO,
        filename=filename,
        revision=MERIT_REVISION,
    )
    checkpoint = torch.load(
        checkpoint_path,
        map_location="cpu",
        weights_only=True,
    )
    head = ProjectionHead(
        in_dim=int(checkpoint["in_dim"]),
        hidden_dim=int(checkpoint["hidden_dim"]),
        out_dim=int(checkpoint["out_dim"]),
    )
    head.load_state_dict(checkpoint["state_dict"])
    return head.to(device).eval()


@lru_cache(maxsize=2)
def _load_models(device_name: str):
    device = torch.device(device_name)
    processor = Wav2Vec2FeatureExtractor.from_pretrained(
        MERT_REPO,
        revision=MERT_REVISION,
    )
    model = AutoModel.from_pretrained(
        MERT_REPO,
        revision=MERT_REVISION,
        trust_remote_code=True,
    ).to(device).eval()
    heads = {
        name: _load_head(filename, device)
        for name, filename in HEAD_FILES.items()
    }
    return processor, model, heads


def _load_audio(path: str) -> np.ndarray:
    audio, sample_rate = sf.read(
        Path(path),
        dtype="float32",
        always_2d=True,
    )
    waveform = torch.from_numpy(audio).mean(dim=1)

    if sample_rate != SAMPLE_RATE:
        waveform = AF.resample(waveform, sample_rate, SAMPLE_RATE)

    waveform = waveform[:NUM_SAMPLES]
    if waveform.numel() < NUM_SAMPLES:
        waveform = F.pad(waveform, (0, NUM_SAMPLES - waveform.numel()))
    return waveform.numpy()


@torch.inference_mode()
def compare_audio(
    reference_path: str,
    comparison_path: str,
) -> dict[str, float]:
    device_name = _select_device()
    processor, model, heads = _load_models(device_name)
    waveforms = [
        _load_audio(reference_path),
        _load_audio(comparison_path),
    ]
    inputs = processor(
        waveforms,
        sampling_rate=SAMPLE_RATE,
        return_tensors="pt",
        padding=True,
    )
    inputs = {
        name: value.to(device_name)
        for name, value in inputs.items()
    }
    output = model(
        **inputs,
        output_hidden_states=True,
    )
    backbone = torch.cat(
        [
            output.hidden_states[layer].mean(dim=1)
            for layer in EXTRACT_LAYERS
        ],
        dim=-1,
    )

    scores = {}
    for name, head in heads.items():
        embeddings = head(backbone)
        score = torch.sum(embeddings[0] * embeddings[1]).item()
        scores[name] = round(float(score), 6)
    return scores