from typing import Optional, List, Iterable, Tuple

import numpy as np
import torch
from pyannote.core import SlidingWindowFeature

from ..mapping import SpeakerMap, SpeakerMapBuilder


class OnlineSpeakerClustering:
    """Implements constrained incremental online clustering of speakers and manages cluster centers.

    Parameters
    ----------
    tau_active:float
        Threshold for detecting active speakers. This threshold is applied on the maximum value of per-speaker output
        activation of the local segmentation model.
    rho_update: float
        Threshold for considering the extracted embedding when updating the centroid of the local speaker.
        The centroid to which a local speaker is mapped is only updated if the ratio of speech/chunk duration
        of a given local speaker is greater than this threshold.
    delta_new: float
        Threshold on the distance between a speaker embedding and a centroid. If the distance between a local speaker and all
        centroids is larger than delta_new, then a new centroid is created for the current speaker.
    metric: str. Defaults to "cosine".
        The distance metric to use.
    max_speakers: int
        Maximum number of global speakers to track through a conversation. Defaults to 20.
    """

    def __init__(
        self,
        tau_active: float,
        rho_update: float,
        delta_new: float,
        metric: Optional[str] = "cosine",
        max_speakers: int = 20,
    ):
        self.tau_active = tau_active
        self.rho_update = rho_update
        self.delta_new = delta_new
        self.metric = metric
        self.max_speakers = max_speakers
        self.centers: Optional[np.ndarray] = None
        self.active_centers = set()
        self.blocked_centers = set()

    @property
    def num_free_centers(self) -> int:
        return self.max_speakers - self.num_known_speakers - self.num_blocked_speakers

    @property
    def num_known_speakers(self) -> int:
        return len(self.active_centers)

    @property
    def num_blocked_speakers(self) -> int:
        return len(self.blocked_centers)

    @property
    def inactive_centers(self) -> List[int]:
        return [
            c
            for c in range(self.max_speakers)
            if c not in self.active_centers or c in self.blocked_centers
        ]

    def get_next_center_position(self) -> Optional[int]:
        for center in range(self.max_speakers):
            if center not in self.active_centers and center not in self.blocked_centers:
                return center

    def init_centers(self, dimension: int):
        """Initializes the speaker centroid matrix

        Parameters
        ----------
        dimension: int
            Dimension of embeddings used for representing a speaker.
        """
        self.centers = np.zeros((self.max_speakers, dimension))
        self.active_centers = set()
        self.blocked_centers = set()

    def update(self, assignments: Iterable[Tuple[int, int]], embeddings: np.ndarray):
        """Updates the speaker centroids given a list of assignments and local speaker embeddings

        Parameters
        ----------
        assignments: Iterable[Tuple[int, int]])
            An iterable of tuples with two elements having the first element as the source speaker
            and the second element as the target speaker.
        embeddings: np.ndarray, shape (local_speakers, embedding_dim)
            Matrix containing embeddings for all local speakers.
        """
        if self.centers is not None:
            for l_spk, g_spk in assignments:
                assert g_spk in self.active_centers, "Cannot update unknown centers"
                self.centers[g_spk] += embeddings[l_spk]

    def add_center(self, embedding: np.ndarray) -> int:
        """Add a new speaker centroid initialized to a given embedding

        Parameters
        ----------
        embedding: np.ndarray
            Embedding vector of some local speaker

        Returns
        -------
            center_index: int
                Index of the created center
        """
        center = self.get_next_center_position()
        self.centers[center] = embedding
        self.active_centers.add(center)
        return center

    def identify(
        self, segmentation: SlidingWindowFeature, embeddings: torch.Tensor
    ) -> SpeakerMap:
        """Identify the centroids to which the input speaker embeddings belong.

        Parameters
        ----------
        segmentation: np.ndarray, shape (frames, local_speakers)
            Matrix of segmentation outputs
        embeddings: np.ndarray, shape (local_speakers, embedding_dim)
            Matrix of embeddings

        Returns
        -------
            speaker_map: SpeakerMap
                A mapping from local speakers to global speakers.
        """
        embeddings = embeddings.detach().cpu().numpy()
        active_speakers = np.where(
            np.max(segmentation.data, axis=0) >= self.tau_active
        )[0]
        long_speakers = np.where(np.mean(segmentation.data, axis=0) >= self.rho_update)[
            0
        ]
        num_local_speakers = segmentation.data.shape[1]

        if self.centers is None:
            self.init_centers(embeddings.shape[1])
            assignments = [
                (spk, self.add_center(embeddings[spk])) for spk in active_speakers
            ]
            return SpeakerMapBuilder.hard_map(
                shape=(num_local_speakers, self.max_speakers),
                assignments=assignments,
                maximize=False,
            )

        # Obtain a mapping based on distances between embeddings and centers
        dist_map = SpeakerMapBuilder.dist(embeddings, self.centers, self.metric)
        # Remove any assignments containing invalid speakers
        inactive_speakers = np.array(
            [spk for spk in range(num_local_speakers) if spk not in active_speakers]
        )
        dist_map = dist_map.unmap_speakers(inactive_speakers, self.inactive_centers)
        # Keep assignments under the distance threshold
        valid_map = dist_map.unmap_threshold(self.delta_new)

        # Some speakers might be unidentified
        missed_speakers = [
            s for s in active_speakers if not valid_map.is_source_speaker_mapped(s)
        ]

        # Add assignments to new centers if possible
        new_center_speakers = []
        for spk in missed_speakers:
            has_space = len(new_center_speakers) < self.num_free_centers
            if has_space and spk in long_speakers:
                # Flag as a new center
                new_center_speakers.append(spk)
            else:
                # Cannot create a new center
                # Get global speakers in order of preference
                preferences = np.argsort(dist_map.mapping_matrix[spk, :])
                preferences = [
                    g_spk for g_spk in preferences if g_spk in self.active_centers
                ]
                # Get the free global speakers among the preferences
                _, g_assigned = valid_map.valid_assignments()
                free = [g_spk for g_spk in preferences if g_spk not in g_assigned]
                if free:
                    # The best global speaker is the closest free one
                    valid_map = valid_map.set_source_speaker(spk, free[0])

        # Update known centers
        to_update = [
            (ls, gs)
            for ls, gs in zip(*valid_map.valid_assignments())
            if ls not in missed_speakers and ls in long_speakers
        ]
        self.update(to_update, embeddings)

        # Add new centers
        for spk in new_center_speakers:
            valid_map = valid_map.set_source_speaker(
                spk, self.add_center(embeddings[spk])
            )

        return valid_map

    def __call__(
        self, segmentation: SlidingWindowFeature, embeddings: torch.Tensor
    ) -> SlidingWindowFeature:
        return SlidingWindowFeature(
            self.identify(segmentation, embeddings).apply(segmentation.data),
            segmentation.sliding_window,
        )
