import logging
import warnings

import numpy as np
from pyannote.core import SlidingWindow, SlidingWindowFeature

logger = logging.getLogger(__name__)


def trim(
    scores: SlidingWindowFeature,
    warm_up: tuple[float, float] = (0.1, 0.1),
) -> SlidingWindowFeature:
    """Trim left and right warm-up regions

    Parameters
    ----------
    scores : SlidingWindowFeature
        (num_chunks, num_frames, num_classes)-shaped scores.
    warm_up : (float, float) tuple
        Left/right warm up ratio of chunk duration.
        Defaults to (0.1, 0.1), i.e. 10% on both sides.

    Returns
    -------
    trimmed : SlidingWindowFeature
        (num_chunks, trimmed_num_frames, num_speakers)-shaped scores
    """

    assert scores.data.ndim == 3, "trim expects (num_chunks, num_frames, num_classes)-shaped scores"
    _, num_frames, _ = scores.data.shape

    chunks = scores.sliding_window

    num_frames_left = round(num_frames * warm_up[0])
    num_frames_right = round(num_frames * warm_up[1])

    num_frames_step = round(num_frames * chunks.step / chunks.duration)
    if num_frames - num_frames_left - num_frames_right < num_frames_step:
        warnings.warn(
            f"Total warm_up is so large ({sum(warm_up) * 100:g}% of each chunk) "
            f"that resulting trimmed scores does not cover a whole step ({chunks.step:g}s)",
            stacklevel=2,
        )
    new_data = scores.data[:, num_frames_left : num_frames - num_frames_right]

    new_chunks = SlidingWindow(
        start=chunks.start + warm_up[0] * chunks.duration,
        step=chunks.step,
        duration=(1 - warm_up[0] - warm_up[1]) * chunks.duration,
    )

    return SlidingWindowFeature(new_data, new_chunks)


def aggregate(
    scores: SlidingWindowFeature,
    frames: SlidingWindow,
    warm_up: tuple[float, float] = (0.0, 0.0),
    epsilon: float = 1e-12,
    hamming: bool = False,
    missing: float = np.nan,
    skip_average: bool = False,
) -> SlidingWindowFeature:
    """Aggregation

    Parameters
    ----------
    scores : SlidingWindowFeature
        Raw (unaggregated) scores. Shape is (num_chunks, num_frames_per_chunk, num_classes).
    frames : SlidingWindow
        Frames resolution.
    warm_up : (float, float) tuple, optional
        Left/right warm up duration (in seconds).
    missing : float, optional
        Value used to replace missing (ie all NaNs) values.
    skip_average : bool, optional
        Skip final averaging step.

    Returns
    -------
    aggregated_scores : SlidingWindowFeature
        Aggregated scores. Shape is (num_frames, num_classes)
    """

    num_chunks, num_frames_per_chunk, num_classes = scores.data.shape

    chunks = scores.sliding_window
    frames = SlidingWindow(
        start=chunks.start,
        duration=frames.duration,
        step=frames.step,
    )

    hamming_window = np.hamming(num_frames_per_chunk).reshape(-1, 1) if hamming else np.ones((num_frames_per_chunk, 1))

    warm_up_window = np.ones((num_frames_per_chunk, 1))
    warm_up_left = round(warm_up[0] / scores.sliding_window.duration * num_frames_per_chunk)
    warm_up_window[:warm_up_left] = epsilon
    warm_up_right = round(warm_up[1] / scores.sliding_window.duration * num_frames_per_chunk)
    warm_up_window[num_frames_per_chunk - warm_up_right :] = epsilon

    num_frames = (
        frames.closest_frame(
            scores.sliding_window.start
            + scores.sliding_window.duration
            + (num_chunks - 1) * scores.sliding_window.step
            + 0.5 * frames.duration
        )
        + 1
    )
    aggregated_output: np.ndarray = np.zeros((num_frames, num_classes), dtype=np.float32)
    overlapping_chunk_count: np.ndarray = np.zeros((num_frames, num_classes), dtype=np.float32)
    aggregated_mask: np.ndarray = np.zeros((num_frames, num_classes), dtype=np.float32)

    for chunk, score in scores:
        mask = 1 - np.isnan(score)
        np.nan_to_num(score, copy=False, nan=0.0)

        start_frame = frames.closest_frame(chunk.start + 0.5 * frames.duration)

        aggregated_output[start_frame : start_frame + num_frames_per_chunk] += (
            score * mask * hamming_window * warm_up_window
        )

        overlapping_chunk_count[start_frame : start_frame + num_frames_per_chunk] += (
            mask * hamming_window * warm_up_window
        )

        aggregated_mask[start_frame : start_frame + num_frames_per_chunk] = np.maximum(
            aggregated_mask[start_frame : start_frame + num_frames_per_chunk],
            mask,
        )

    average = aggregated_output if skip_average else aggregated_output / np.maximum(overlapping_chunk_count, epsilon)

    average[aggregated_mask == 0.0] = missing

    return SlidingWindowFeature(average, frames)
