import logging
from pathlib import Path

import numpy as np
import numpy.typing as npt
from scipy.linalg import eigh

logger = logging.getLogger(__name__)

_MODULE_DIR = Path(__file__).parent
XVEC_TRANSFORM_PATH = _MODULE_DIR / "data" / "xvec_transform.npz"
PLDA_PATH = _MODULE_DIR / "data" / "plda.npz"


def load_xvex_and_plda_data(
    xvec_transform_path: Path = XVEC_TRANSFORM_PATH,
    plda_path: Path = PLDA_PATH,
) -> tuple[dict[str, npt.NDArray], dict[str, npt.NDArray]]:
    xvec_data = np.load(xvec_transform_path)
    plda_data = np.load(plda_path)
    return xvec_data, plda_data


def l2_norm(vec_or_matrix: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
    if len(vec_or_matrix.shape) == 1:
        return vec_or_matrix / np.linalg.norm(vec_or_matrix)
    elif len(vec_or_matrix.shape) == 2:
        return vec_or_matrix / np.linalg.norm(vec_or_matrix, axis=1, ord=2)[:, np.newaxis]
    else:
        raise ValueError(f"Wrong number of dimensions, 1 or 2 is supported, not {len(vec_or_matrix.shape)}.")


class PLDATransform:
    def __init__(
        self,
        xvec_data: dict[str, npt.NDArray],
        plda_data: dict[str, npt.NDArray],
        lda_dimension: int = 128,
    ) -> None:
        self.mean1 = xvec_data["mean1"]
        self.mean2 = xvec_data["mean2"]
        self.lda = xvec_data["lda"]

        plda_mu = plda_data["mu"]
        plda_tr = plda_data["tr"]
        plda_psi = plda_data["psi"]

        w = np.linalg.inv(plda_tr.T.dot(plda_tr))
        b = np.linalg.inv((plda_tr.T / plda_psi).dot(plda_tr))

        acvar, wccn = eigh(b, w)
        self.plda_psi = acvar[::-1]
        self.plda_tr = wccn.T[::-1]
        self.plda_mu = plda_mu

        self.lda_dimension = lda_dimension

    @property
    def phi(self) -> np.ndarray:
        return self.plda_psi[: self.lda_dimension]

    def xvec_tf(self, x: npt.NDArray[np.float32]) -> npt.NDArray[np.float64]:
        x = np.sqrt(self.lda.shape[1]) * l2_norm(
            self.lda.T.dot(np.sqrt(self.lda.shape[0]) * l2_norm(x - self.mean1).T).T - self.mean2
        )
        return x

    def plda_tf(self, x: npt.NDArray[np.float64], lda_dim: int | None = None) -> npt.NDArray[np.float64]:
        if lda_dim is None:
            lda_dim = self.lda_dimension
        return (x - self.plda_mu).dot(self.plda_tr.T)[:, :lda_dim]

    def __call__(self, embeddings: npt.NDArray[np.float32]) -> npt.NDArray[np.float64]:
        return self.plda_tf(self.xvec_tf(embeddings), lda_dim=self.lda_dimension)
