# 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/>.


"""Legacy energy-based transcription backend.

Maintains backward compatibility with the original NaSong note detection
logic, which uses energy thresholds for onsets and FFT peak detection
for pitch estimation.
"""

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

#
import numpy as np

#
from .base import NoteDetector


#
class LegacyDetector(NoteDetector):
    """Legacy note detection using energy peaks and FFT analysis.

    Simple, lightweight, and deterministic. Good for clear onset sounds but
    less robust than modern neural or DSP methods.
    """

    def detect(self, audio_data: np.ndarray, sample_rate: int) -> list[dict[str, Any]]:
        """Detects notes using energy-based onset detection.

        Args:
            audio_data (np.ndarray): Audio data.
            sample_rate (int): Sampling rate.

        Returns:
            list[dict[str, Any]]: Detected notes.
        """
        # Unpack config
        num_notes = self.config.get("legacy_num_notes", 1)
        use_onset_detection = self.config.get("legacy_use_onset", True)

        total_duration = len(audio_data) / sample_rate
        notes = []

        if use_onset_detection:
            # Simple energy-based onset detection
            # Calculate short-time energy
            frame_len_sec = self.config.get("legacy_frame_len", 0.02)
            hop_len_sec = self.config.get("legacy_hop_len", 0.01)

            frame_length = int(frame_len_sec * sample_rate)
            hop_length = int(hop_len_sec * sample_rate)

            energy = []
            for i in range(0, len(audio_data) - frame_length, hop_length):
                frame = audio_data[i : i + frame_length]
                energy.append(np.sum(frame**2))

            energy_arr = np.array(energy)

            # Find peaks in energy (onsets)
            threshold_ratio = self.config.get("legacy_onset_threshold", 0.3)
            threshold = threshold_ratio * np.max(energy_arr)
            onsets = []
            for i in range(1, len(energy_arr) - 1):
                if (
                    energy_arr[i] > threshold
                    and energy_arr[i] > energy_arr[i - 1]
                    and energy_arr[i] > energy_arr[i + 1]
                ):
                    onset_time = i * hop_length / sample_rate
                    onsets.append(onset_time)

            # Limit to num_notes
            onsets = onsets[:num_notes]

            # If no onsets detected, fall back to simple division
            if len(onsets) == 0:
                onsets = [i * total_duration / num_notes for i in range(num_notes)]
        else:
            # Simple even division
            onsets = [i * total_duration / num_notes for i in range(num_notes)]

        # Extract pitch and duration for each onset
        for i, start_time in enumerate(onsets):
            # Determine duration (to next onset or end)
            if i < len(onsets) - 1:
                duration = onsets[i + 1] - start_time
            else:
                duration = total_duration - start_time

            # Extract segment for pitch detection
            start_idx = int(start_time * sample_rate)
            # Use first 100ms or duration, whichever is shorter
            analysis_duration = min(0.1, duration)
            end_idx = int((start_time + analysis_duration) * sample_rate)
            note_segment = audio_data[start_idx:end_idx]

            # Detect multiple pitches (chord detection)
            max_pitches = self.config.get("legacy_max_pitches", 3)
            min_freq = self.config.get("legacy_min_freq", 50.0)
            max_freq = self.config.get("legacy_max_freq", 4000.0)

            frequencies = self._detect_pitches_fft(
                note_segment,
                sample_rate,
                max_pitches=max_pitches,
                min_freq=min_freq,
                max_freq=max_freq,
            )

            notes.append(
                {
                    "frequencies": frequencies,  # List of frequencies (chord)
                    "start_time": start_time,
                    "duration": duration * 0.9,  # Leave small gap
                    "amplitude": float(np.sqrt(np.mean(note_segment**2)))
                    if len(note_segment) > 0
                    else 0.0,
                }
            )

        return notes

    def _detect_pitches_fft(
        self,
        audio_data: np.ndarray,
        sample_rate: int,
        max_pitches: int = 3,
        min_freq: float = 50.0,
        max_freq: float = 4000.0,
    ) -> list[float]:
        """
        Detect multiple pitches using FFT peak detection.
        """

        if len(audio_data) == 0:
            return [220.0]  # Default A3

        # Apply window
        windowed = audio_data * np.hanning(len(audio_data))

        # FFT
        fft = np.fft.rfft(windowed)
        freqs = np.fft.rfftfreq(len(windowed), 1 / sample_rate)
        magnitudes = np.abs(fft)

        # Focus on frequency range
        freq_mask = (freqs >= min_freq) & (freqs <= max_freq)
        freqs_filtered = freqs[freq_mask]
        mags_filtered = magnitudes[freq_mask]

        if len(mags_filtered) == 0:
            return [220.0]

        # Find peaks
        detected_freqs = []
        threshold = 0.1 * np.max(mags_filtered)

        # Find local maxima
        for i in range(1, len(mags_filtered) - 1):
            if (
                mags_filtered[i] > threshold
                and mags_filtered[i] > mags_filtered[i - 1]
                and mags_filtered[i] > mags_filtered[i + 1]
            ):
                detected_freqs.append((freqs_filtered[i], mags_filtered[i]))

        # Sort by magnitude and take top max_pitches
        detected_freqs.sort(key=lambda x: x[1], reverse=True)
        detected_freqs = detected_freqs[:max_pitches]

        # Extract just frequencies, sorted by frequency
        result = [f for f, _ in detected_freqs]
        result.sort()

        if len(result) == 0:
            return [220.0]

        return result
