# Copyright (C) 2026 Nathan Cerisara <https://github.com/nath54/nasong>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.


"""Note detection backend using the TorchCrepe implementation.

Leverages PyTorch-based CREPE for state-of-the-art pitch tracking and
Viterbi decoding for temporal consistency. Best for high-accuracy
monophonic transcription when GPU acceleration is available.
"""

#
### Import Modules. ###
#
from typing import Any

#
import numpy as np

#
from .base import NoteDetector

#
### Try importing torch, handle failure gracefully ###
#
try:
    import torch

    HAS_TORCH = True
except (ImportError, OSError):
    HAS_TORCH = False
    torch = None  # type: ignore

#
### Try importing torchcrepe, handle failure gracefully ###
#
try:
    import torchcrepe  # type: ignore
except ImportError:
    torchcrepe = None  # type: ignore


class TorchCrepeDetector(NoteDetector):
    """SOTA pitch detection using PyTorch implementation of CREPE.

    Utilizes Viterbi decoding for smooth pitch contours and a thresholded
    segmentation approach to identify note boundaries.
    """

    def detect(self, audio_data: np.ndarray, sample_rate: int) -> list[dict[str, Any]]:
        """Detects notes using TorchCrepe inference and Viterbi decoding.

        Args:
            audio_data (np.ndarray): The audio data.
            sample_rate (int): Audio sampling rate.

        Returns:
            list[dict[str, Any]]: List of detected note events.
        """
        if not HAS_TORCH:
            raise ImportError("PyTorch is not installed. Cannot use TorchCrepe.")

        if torchcrepe is None:
            raise ImportError(
                "TorchCrepe is not installed. Please install 'torchcrepe'."
            )

        device = "cuda" if torch.cuda.is_available() else "cpu"  # pylint: disable=no-member

        # Prepare audio tensor
        # Crepe expects shape (batch, time)
        audio_tensor = torch.tensor(  # pylint: disable=no-member
            audio_data,
            dtype=torch.float32,
            device=device,
        ).unsqueeze(0)

        step_size_ms = self.config.get("crepe_step_size", 10)
        hop_length = int(step_size_ms * sample_rate / 1000)

        fmin = 50.0
        fmax = 2000.0
        model_size = self.config.get("crepe_model", "medium")
        conf_thresh = self.config.get("crepe_confidence_threshold", 0.8)

        # Predict
        try:
            # Note: predict returns pitch in Hz and periodicity (confidence)
            pitch, periodicity = torchcrepe.predict(
                audio_tensor,
                sample_rate,
                hop_length=hop_length,
                fmin=fmin,
                fmax=fmax,
                model=model_size,
                decoder=torchcrepe.decode.viterbi,
                return_periodicity=True,
                device=device,
                batch_size=2048,
            )
        except Exception as e:  # pylint: disable=broad-except
            raise RuntimeError(f"TorchCrepe prediction failed: {e}") from e

        # Move to CPU for processing
        pitch = pitch.squeeze(0).cpu().numpy()
        periodicity = periodicity.squeeze(0).cpu().numpy()

        # Segmentation logic
        # Find continuous regions where periodicity > threshold
        is_voiced = periodicity > conf_thresh

        notes: list[dict[str, Any]] = []
        hop_s = step_size_ms / 1000.0

        current_start_idx = None
        current_pitches = []
        current_confs = []

        for i, voiced in enumerate(is_voiced):
            if voiced:
                if current_start_idx is None:
                    current_start_idx = i
                current_pitches.append(pitch[i])
                current_confs.append(periodicity[i])
            else:
                if current_start_idx is not None:
                    # End of note
                    self._add_note(
                        notes,
                        current_start_idx,
                        i,
                        current_pitches,
                        current_confs,
                        hop_s,
                        audio_data,
                        sample_rate,
                    )
                    current_start_idx = None
                    current_pitches = []
                    current_confs = []

        # Check last note
        if current_start_idx is not None:
            self._add_note(
                notes,
                current_start_idx,
                len(is_voiced),
                current_pitches,
                current_confs,
                hop_s,
                audio_data,
                sample_rate,
            )

        return notes

    def _add_note(
        self,
        notes,
        start_idx,
        end_idx,
        pitches,
        confs,
        hop_s,
        audio_data,
        sample_rate,
    ):
        duration = (end_idx - start_idx) * hop_s
        # Min duration, e.g. 50ms
        if duration < 0.05:
            return

        start_time = start_idx * hop_s
        end_time = end_idx * hop_s

        # Calculate amplitude (RMS)
        start_sample = int(start_time * sample_rate)
        end_sample = int(end_time * sample_rate)

        # Clamp indices
        start_sample = max(0, start_sample)
        end_sample = min(len(audio_data), end_sample)

        if end_sample > start_sample:
            segment = audio_data[start_sample:end_sample]
            amplitude = float(np.sqrt(np.mean(segment**2)))
        else:
            amplitude = 0.0

        # Median pitch to robustly handle fluctuations
        median_pitch = np.median(pitches)
        mean_conf = np.mean(confs)

        notes.append(
            {
                "start_time": float(start_time),
                "duration": float(duration),
                "frequencies": [float(median_pitch)],
                "confidence": float(mean_conf),
                "amplitude": amplitude,
            }
        )
