import logging
import time

import numpy as np
import numpy.typing as npt
import onnxruntime as ort
from pyannote.core import SlidingWindow, SlidingWindowFeature

from onnx_diarization.powerset import PowersetConverter

logger = logging.getLogger(__name__)


class PyannnoteSegmentation:
    """
    Pyannote segmentation model for speaker diarization.

    Args:
        session: ONNX Runtime session or InferenceSessionProxy for the segmentation model
        duration: Duration of segmentation chunks in seconds
        step: Step size between chunks in seconds (defaults to duration * 0.1)
        inference_batch_size: Maximum number of chunks to process in a single inference call
    """

    def __init__(
        self,
        session: ort.InferenceSession,
        duration: float = 10.0,
        step: float | None = None,
        inference_batch_size: int = 32,
    ) -> None:
        self.session = session
        self.duration = duration
        self.step = step if step is not None else duration * 0.1
        self.sample_rate = 16000
        self.inference_batch_size = inference_batch_size

        self.powerset_converter = PowersetConverter(num_classes=3, max_set_size=2)

    def get_num_samples(self) -> int:
        return int(self.duration * self.sample_rate)

    @property
    def frames(self) -> SlidingWindow:
        return SlidingWindow(duration=0.0170068, step=0.0170068)

    def __call__(self, audio_data: npt.NDArray[np.float32]) -> SlidingWindowFeature:
        t_start = time.perf_counter()

        num_samples = self.get_num_samples()
        step_samples = int(self.step * self.sample_rate)
        total_samples = len(audio_data)

        chunks = []
        chunk_starts = []

        start = 0
        while start < total_samples:
            end = min(start + num_samples, total_samples)
            chunk = audio_data[start:end]

            if len(chunk) < num_samples:
                chunk = np.pad(chunk, (0, num_samples - len(chunk)), mode="constant")

            chunks.append(chunk)
            chunk_starts.append(start / self.sample_rate)

            start += step_samples

            if end >= total_samples:
                break

        t_prep = time.perf_counter() - t_start

        chunks_array = np.array(chunks)[:, np.newaxis, :]

        t_infer_start = time.perf_counter()
        powerset_logprobs_list = []

        num_chunks_total = len(chunks_array)
        for batch_start in range(0, num_chunks_total, self.inference_batch_size):
            batch_end = min(batch_start + self.inference_batch_size, num_chunks_total)
            batch = chunks_array[batch_start:batch_end]

            onnx_input = {"input": batch.astype(np.float32)}
            batch_logprobs = self.session.run(None, onnx_input)[0]
            powerset_logprobs_list.append(batch_logprobs)

            logger.debug(
                f"Processed batch {batch_start // self.inference_batch_size + 1}/{(num_chunks_total + self.inference_batch_size - 1) // self.inference_batch_size}: "
                f"chunks {batch_start}-{batch_end - 1}"
            )

        powerset_logprobs = np.concatenate(powerset_logprobs_list, axis=0)
        t_infer = time.perf_counter() - t_infer_start

        logger.debug(
            f"Segmentation inference: {t_infer:.3f}s for {len(chunks)} chunks (throughput: {len(chunks) / t_infer:.1f} chunks/s)"
        )

        logger.info(f"Powerset (log probs): min={np.min(powerset_logprobs):.4f}, max={np.max(powerset_logprobs):.4f}")
        powerset_probs = np.exp(powerset_logprobs)
        logger.info(
            f"Powerset (probs): min={np.min(powerset_probs):.4f}, max={np.max(powerset_probs):.4f}, sum_across_classes={np.mean(np.sum(powerset_probs, axis=-1)):.4f}"
        )

        segmentations = self.powerset_converter.to_multilabel(powerset_logprobs, soft=False)
        logger.info(
            f"Multilabel (after conversion): shape={segmentations.shape}, min={np.min(segmentations):.4f}, max={np.max(segmentations):.4f}"
        )

        _num_chunks, _num_frames, _num_classes = segmentations.shape

        chunks_window = SlidingWindow(
            start=0.0,
            duration=self.duration,
            step=self.step,
        )

        t_total = time.perf_counter() - t_start
        logger.debug(
            f"Segmentation total: {t_total:.3f}s (prep: {t_prep / t_total * 100:.1f}%, inference: {t_infer / t_total * 100:.1f}%)"
        )

        return SlidingWindowFeature(segmentations, chunks_window, labels=None)
