import logging

from einops import rearrange
import numpy as np
import numpy.typing as npt
from pyannote.core import SlidingWindowFeature
from scipy.cluster.hierarchy import fcluster, linkage
from scipy.optimize import linear_sum_assignment
from scipy.spatial.distance import cdist
from scipy.special import logsumexp, softmax
from sklearn.cluster import KMeans

from onnx_diarization.clustering.plda import PLDATransform

logger = logging.getLogger(__name__)


# Ah, If only I knew what the fuck is going on here


def vbx_algorithm(
    X: npt.NDArray[np.float64],  # noqa: N803
    Phi: npt.NDArray[np.float64],  # noqa: N803
    Fa: float = 1.0,  # noqa: N803
    Fb: float = 1.0,  # noqa: N803
    pi: int | npt.NDArray[np.float64] = 10,
    gamma: npt.NDArray[np.float64] | None = None,
    maxIters: int = 10,  # noqa: N803
    epsilon: float = 1e-4,
    alphaQInit: float = 1.0,  # noqa: N803
) -> tuple[
    npt.NDArray[np.float64],
    npt.NDArray[np.float64],
    list[list[float]],
    npt.NDArray[np.float64],
    npt.NDArray[np.float64],
]:
    D = X.shape[1]  # noqa: N806

    if isinstance(pi, int):
        pi = np.ones(pi) / pi

    if gamma is None:
        gamma = np.random.default_rng().gamma(alphaQInit, size=(X.shape[0], len(pi)))
        gamma = gamma / gamma.sum(1, keepdims=True)

    assert gamma is not None
    assert gamma.shape[1] == len(pi) and gamma.shape[0] == X.shape[0]

    G = -0.5 * (np.sum(X**2, axis=1, keepdims=True) + D * np.log(2 * np.pi))  # noqa: N806
    V = np.sqrt(Phi)  # noqa: N806
    rho = X * V
    Li: list[list[float]] = []  # noqa: N806
    ELBO = None  # noqa: N806
    alpha = None
    invL = None  # noqa: N806

    for ii in range(maxIters):
        if ii > 0 or alpha is None or invL is None:
            invL = 1.0 / (1 + Fa / Fb * gamma.sum(axis=0, keepdims=True).T * Phi)  # noqa: N806
            alpha = Fa / Fb * invL * gamma.T.dot(rho)

        log_p_ = Fa * (rho.dot(alpha.T) - 0.5 * (invL + alpha**2).dot(Phi) + G)

        eps = 1e-8
        lpi = np.log(pi + eps)
        log_p_x = logsumexp(log_p_ + lpi, axis=-1)
        log_pX_ = np.sum(log_p_x, axis=0)  # noqa: N806

        gamma = np.exp(log_p_ + lpi - log_p_x[:, None])
        pi = np.sum(gamma, axis=0)
        pi = pi / pi.sum()

        ELBO = log_pX_ + Fb * 0.5 * np.sum(np.log(invL) - invL - alpha**2 + 1)  # noqa: N806
        Li.append([ELBO])

        if ii > 0 and ELBO - Li[-2][0] < epsilon:
            if ELBO - Li[-2][0] < 0:
                logger.warning("Value of auxiliary function decreased in VBx")
            break

    return gamma, pi, Li, alpha, invL


def cluster_vbx(
    ahc_init: npt.NDArray[np.int64],
    fea: npt.NDArray[np.float64],
    Phi: npt.NDArray[np.float64],  # noqa: N803
    Fa: float,  # noqa: N803
    Fb: float,  # noqa: N803
    maxIters: int = 20,  # noqa: N803
    init_smoothing: float = 7.0,
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
    qinit = np.zeros((len(ahc_init), ahc_init.max() + 1))
    qinit[range(len(ahc_init)), ahc_init.astype(int)] = 1.0
    qinit = qinit if init_smoothing < 0 else softmax(qinit * init_smoothing, axis=1)

    gamma, pi, _, _, _ = vbx_algorithm(
        fea,
        Phi,
        Fa=Fa,
        Fb=Fb,
        pi=qinit.shape[1],
        gamma=qinit,
        maxIters=maxIters,
    )

    return gamma, pi


class VBxClustering:
    def __init__(
        self,
        plda: PLDATransform,
        metric: str = "cosine",
        constrained_assignment: bool = True,
        threshold: float = 0.7,
        Fa: float = 0.07,  # noqa: N803
        Fb: float = 0.8,  # noqa: N803
    ) -> None:
        self.plda = plda
        self.metric = metric
        self.constrained_assignment = constrained_assignment
        self.threshold = threshold
        self.Fa = Fa
        self.Fb = Fb

    def filter_embeddings(
        self,
        embeddings: npt.NDArray[np.float32],
        segmentations: SlidingWindowFeature | None = None,
        min_active_ratio: float = 0.2,
    ) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.int64], npt.NDArray[np.int64]]:
        num_chunks, num_speakers, _dimension = embeddings.shape

        if segmentations is None:
            active = np.ones((num_chunks, num_speakers), dtype=bool)
            train_embeddings = embeddings[active]
            train_chunk_idx, train_speaker_idx = np.where(active)
            return train_embeddings, train_chunk_idx, train_speaker_idx

        _, num_frames, _ = segmentations.data.shape

        single_active_mask = np.sum(segmentations.data, axis=2, keepdims=True) == 1
        num_clean_frames = np.sum(segmentations.data * single_active_mask, axis=1)
        active = num_clean_frames >= min_active_ratio * num_frames
        valid = ~np.any(np.isnan(embeddings), axis=2)

        chunk_idx, speaker_idx = np.where(active * valid)

        return embeddings[chunk_idx, speaker_idx], chunk_idx, speaker_idx

    def constrained_argmax(self, soft_clusters: npt.NDArray[np.float64]) -> npt.NDArray[np.int32]:
        num_chunks, num_speakers, _num_clusters = soft_clusters.shape

        hard_clusters = np.zeros((num_chunks, num_speakers), dtype=np.int32)

        for c in range(num_chunks):
            cost_matrix = -soft_clusters[c]
            cost_matrix = np.nan_to_num(cost_matrix, nan=1e10, posinf=1e10, neginf=-1e10)
            row_ind, col_ind = linear_sum_assignment(cost_matrix)
            hard_clusters[c, row_ind] = col_ind

        return hard_clusters

    def __call__(
        self,
        embeddings: npt.NDArray[np.float32],
        segmentations: SlidingWindowFeature | None = None,
        num_clusters: int | None = None,
        min_clusters: int | None = None,
        max_clusters: int | None = None,
    ) -> tuple[npt.NDArray[np.int32], npt.NDArray[np.float64], npt.NDArray[np.float32]]:
        constrained_assignment = self.constrained_assignment

        train_embeddings, _, _ = self.filter_embeddings(embeddings, segmentations=segmentations)

        if train_embeddings.shape[0] < 2:
            num_chunks, num_speakers, _ = embeddings.shape
            hard_clusters = np.zeros((num_chunks, num_speakers), dtype=np.int8)
            soft_clusters = np.ones((num_chunks, num_speakers, 1))
            centroids = np.mean(train_embeddings, axis=0, keepdims=True)
            return hard_clusters, soft_clusters, centroids

        if not np.all(np.isfinite(train_embeddings)):
            logger.warning(
                f"Non-finite values in raw train embeddings: {np.sum(~np.isfinite(train_embeddings))} out of {train_embeddings.size}"
            )
            train_embeddings = np.nan_to_num(train_embeddings, nan=0.0, posinf=1.0, neginf=-1.0)

        norms = np.linalg.norm(train_embeddings, axis=1, keepdims=True)
        norms = np.maximum(norms, 1e-8)
        train_embeddings_normed = train_embeddings / norms

        if not np.all(np.isfinite(train_embeddings_normed)):
            logger.warning(
                f"Non-finite values in normalized embeddings: {np.sum(~np.isfinite(train_embeddings_normed))} out of {train_embeddings_normed.size}"
            )
            train_embeddings_normed = np.nan_to_num(train_embeddings_normed, nan=0.0, posinf=1.0, neginf=-1.0)

        dendrogram = linkage(train_embeddings_normed, method="centroid", metric="euclidean")
        ahc_clusters = fcluster(dendrogram, self.threshold, criterion="distance") - 1
        _, ahc_clusters = np.unique(ahc_clusters, return_inverse=True)

        fea = self.plda(train_embeddings)
        q, sp = cluster_vbx(
            ahc_clusters,
            fea,
            self.plda.phi,
            Fa=self.Fa,
            Fb=self.Fb,
            maxIters=20,
        )

        num_chunks, num_speakers, dimension = embeddings.shape
        W = q[:, sp > 1e-7]  # noqa: N806
        centroids = W.T @ train_embeddings.reshape(-1, dimension) / W.sum(0, keepdims=True).T

        auto_num_clusters, _ = centroids.shape
        if min_clusters is not None and auto_num_clusters < min_clusters:
            num_clusters = min_clusters
        elif max_clusters is not None and auto_num_clusters > max_clusters:
            num_clusters = max_clusters

        if num_clusters and num_clusters != auto_num_clusters:
            constrained_assignment = False
            kmeans_clusters = KMeans(n_clusters=num_clusters, n_init=3, random_state=42, copy_x=False).fit_predict(
                train_embeddings_normed
            )
            centroids = np.vstack(
                [np.mean(train_embeddings[kmeans_clusters == k], axis=0) for k in range(num_clusters)]
            )

        e2k_distance = rearrange(
            cdist(
                rearrange(embeddings, "c s d -> (c s) d"),
                centroids,
                metric=self.metric,
            ),
            "(c s) k -> c s k",
            c=num_chunks,
            s=num_speakers,
        )
        soft_clusters = 2 - e2k_distance

        if constrained_assignment:
            const = soft_clusters.min() - 1.0
            if segmentations is not None:
                soft_clusters[segmentations.data.sum(1) == 0] = const
            hard_clusters = self.constrained_argmax(soft_clusters)
        else:
            hard_clusters = np.argmax(soft_clusters, axis=2)

        hard_clusters = hard_clusters.reshape(num_chunks, num_speakers)

        return hard_clusters, soft_clusters, centroids
