import logging
import os
from typing import Any, TypeAlias

import torch
import torch.nn.functional as F
import torchaudio
from coqpit import Coqpit

from TTS.vc.configs.knnvc_config import KNNVCConfig
from TTS.vc.layers.freevc.wavlm import get_wavlm
from TTS.vc.models.base_vc import BaseVC

logger = logging.getLogger(__name__)

PathOrTensor: TypeAlias = str | os.PathLike[Any] | torch.Tensor


class KNNVC(BaseVC):
    """
    Paper::
        https://arxiv.org/abs/2305.18975

    Paper Abstract::
        Any-to-any voice conversion aims to transform source speech
        into a target voice with just a few examples of the target speaker as a
        reference. Recent methods produce convincing conversions, but at the cost of
        increased complexity -- making results difficult to reproduce and build on.
        Instead, we keep it simple. We propose k-nearest neighbors voice conversion
        (kNN-VC): a straightforward yet effective method for any-to-any conversion.
        First, we extract self-supervised representations of the source and reference
        speech. To convert to the target speaker, we replace each frame of the source
        representation with its nearest neighbor in the reference. Finally, a pretrained
        vocoder synthesizes audio from the converted representation. Objective and
        subjective evaluations show that kNN-VC improves speaker similarity with similar
        intelligibility scores to existing methods.

    Samples::
        https://bshall.github.io/knn-vc

    Original code::
        https://github.com/bshall/knn-vc

    Examples:
        >>> from TTS.vc.configs.knnvc_config import KNNVCConfig
        >>> from TTS.vc.models.knnvc import KNNVC
        >>> config = KNNVCConfig()
        >>> model = KNNVC(config)
    """

    def __init__(self, config: Coqpit):
        super().__init__(config)
        self.ssl_dim = self.args.ssl_dim
        self.wavlm = get_wavlm()

    @staticmethod
    def init_from_config(config: KNNVCConfig) -> "KNNVC":
        return KNNVC(config)

    @torch.inference_mode()
    def get_features(self, audio: PathOrTensor, vad_trigger_level=0) -> torch.Tensor:
        """Return features for the given waveform with output shape (seq_len, dim).

        Optionally perform VAD trimming on start/end with `vad_trigger_level`.
        """
        # load audio
        if isinstance(audio, torch.Tensor):
            x: torch.Tensor = audio
            sr = self.config.audio.sample_rate
            if x.dim() == 1:
                x = x[None]
        else:
            x, sr = torchaudio.load(audio, normalize=True)

        if not sr == self.config.audio.sample_rate:
            logger.info("Resampling %d to %d in %s", sr, self.config.audio.sample_rate, audio)
            x = torchaudio.functional.resample(x, orig_freq=sr, new_freq=self.config.audio.sample_rate)
            sr = self.config.audio.sample_rate

        # trim silence from front and back
        if vad_trigger_level > 1e-3:
            transform = torchaudio.transforms.Vad(sample_rate=sr, trigger_level=vad_trigger_level)
            x_front_trim = transform(x)
            waveform_reversed = torch.flip(x_front_trim, (-1,))
            waveform_reversed_front_trim = transform(waveform_reversed)
            x = torch.flip(waveform_reversed_front_trim, (-1,))

        # extract the representation of each layer
        wav_input_16khz = x.to(self.device)
        features = self.wavlm.extract_features(
            wav_input_16khz, output_layer=self.config.wavlm_layer, ret_layer_results=False
        )[0]
        return features.squeeze(0)

    def get_matching_set(self, wavs: list[PathOrTensor], vad_trigger_level=7) -> torch.Tensor:
        """Get concatenated wavlm features for the matching set using all waveforms in `wavs`.

        Wavs are specified as either a list of paths or list of loaded waveform tensors of
        shape (channels, T), assumed to be of 16kHz sample rate.
        """
        feats = []
        for p in wavs:
            feats.append(self.get_features(p, vad_trigger_level=vad_trigger_level))

        feats = torch.concat(feats, dim=0).cpu()
        return feats

    @staticmethod
    def fast_cosine_dist(source_feats: torch.Tensor, matching_pool: torch.Tensor) -> torch.Tensor:
        """Like torch.cdist, but fixed dim=-1 and for cosine distance."""
        source_norms = torch.norm(source_feats, p=2, dim=-1)
        matching_norms = torch.norm(matching_pool, p=2, dim=-1)
        dotprod = (
            -(torch.cdist(source_feats[None], matching_pool[None], p=2)[0] ** 2)
            + source_norms[:, None] ** 2
            + matching_norms[None] ** 2
        )
        dotprod /= 2

        dists = 1 - (dotprod / (source_norms[:, None] * matching_norms[None]))
        return dists

    @torch.inference_mode()
    def match(
        self,
        query_seq: torch.Tensor,
        matching_set: torch.Tensor,
        synth_set: torch.Tensor | None = None,
        topk: int | None = None,
        target_duration: float | None = None,
    ) -> torch.Tensor:
        """Given `query_seq`, `matching_set`, and `synth_set` tensors of shape (N, dim), perform kNN regression matching
        with k=`topk`.

        Args:
            `query_seq`: Tensor (N1, dim) of the input/source query features.
            `matching_set`: Tensor (N2, dim) of the matching set used as the 'training set' for the kNN algorithm.
            `synth_set`: optional Tensor (N2, dim) corresponding to the matching set. We use the matching set to assign
                         each query vector to a vector in the matching set, and then use the corresponding vector from
                         the synth set during HiFiGAN synthesis.
                         By default, and for best performance, this should be identical to the matching set.
            `topk`: k in the kNN -- the number of nearest neighbors to average over.
            `target_duration`: if set to a float, interpolate waveform duration to be equal to this value in seconds.

        Returns:
            - converted features (1, N, dim)
        """
        if topk is None:
            topk = self.config.topk
        synth_set = matching_set.to(self.device) if synth_set is None else synth_set.to(self.device)
        matching_set = matching_set.to(self.device)
        query_seq = query_seq.to(self.device)

        if target_duration is not None:
            target_samples = int(target_duration * self.config.audio.sample_rate)
            scale_factor = (target_samples / self.hop_length) / query_seq.shape[0]  # n_targ_feats / n_input_feats
            query_seq = F.interpolate(query_seq.T[None], scale_factor=scale_factor, mode="linear")[0].T

        dists = self.fast_cosine_dist(query_seq, matching_set)
        best = dists.topk(k=topk, largest=False, dim=-1)
        out_feats = synth_set[best.indices].mean(dim=1)
        return out_feats.unsqueeze(0)

    def load_checkpoint(self, vc_config: KNNVCConfig, _vc_checkpoint: str | os.PathLike[Any]) -> None:
        """kNN-VC does not use checkpoints."""

    def forward(self) -> None: ...
    def inference(self) -> None: ...

    @torch.inference_mode()
    def voice_conversion(
        self,
        source: PathOrTensor,
        target: list[PathOrTensor],
        topk: int | None = None,
    ) -> torch.Tensor:
        if not isinstance(target, list):
            target = [target]
        source_features = self.get_features(source)
        matching_set = self.get_matching_set(target)
        return self.match(source_features, matching_set, topk=topk)
