"""
Real-time pitch detection module using autocorrelation and FFT methods.
"""

import numpy as np
import scipy.signal
from typing import Optional, Tuple


class PitchDetector:
    """Real-time pitch detector for audio signals."""
    
    def __init__(self, sample_rate: int = 44100, frame_size: int = 4096):
        """
        Initialize the pitch detector.
        
        Args:
            sample_rate: Audio sample rate in Hz
            frame_size: Size of audio frame to analyze
        """
        self.sample_rate = sample_rate
        self.frame_size = frame_size
        self.min_freq = 80.0  # Minimum frequency to detect (around E2)
        self.max_freq = 2000.0  # Maximum frequency to detect (around B6)
        
        # Create note mapping
        self.note_names = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
        
    def detect_pitch(self, audio_data: np.ndarray) -> Optional[float]:
        """
        Detect the fundamental frequency of audio data using autocorrelation.
        
        Args:
            audio_data: Array of audio samples
            
        Returns:
            Detected frequency in Hz, or None if no clear pitch detected
        """
        # Apply window to reduce spectral leakage
        windowed = audio_data * np.hanning(len(audio_data))
        
        # Remove DC component
        windowed = windowed - np.mean(windowed)
        
        # Autocorrelation method
        correlation = np.correlate(windowed, windowed, mode='full')
        correlation = correlation[len(correlation)//2:]
        
        # Find the minimum and maximum lag based on frequency range
        min_lag = int(self.sample_rate / self.max_freq)
        max_lag = int(self.sample_rate / self.min_freq)
        
        if max_lag >= len(correlation):
            return None
            
        # Find peaks in autocorrelation
        correlation_segment = correlation[min_lag:max_lag]
        
        if len(correlation_segment) == 0:
            return None
            
        # Find the highest peak
        peak_idx = np.argmax(correlation_segment)
        lag = peak_idx + min_lag
        
        # Check if the peak is significant enough
        if correlation[lag] < 0.3 * correlation[0]:
            return None
            
        # Convert lag to frequency
        frequency = self.sample_rate / lag
        
        # Validate frequency is in expected range
        if self.min_freq <= frequency <= self.max_freq:
            return frequency
        
        return None
    
    def detect_pitch_fft(self, audio_data: np.ndarray) -> Optional[float]:
        """
        Alternative pitch detection using FFT method.
        
        Args:
            audio_data: Array of audio samples
            
        Returns:
            Detected frequency in Hz, or None if no clear pitch detected
        """
        # Apply window
        windowed = audio_data * np.hanning(len(audio_data))
        
        # Compute FFT
        fft = np.fft.rfft(windowed)
        magnitude = np.abs(fft)
        
        # Create frequency bins
        freqs = np.fft.rfftfreq(len(windowed), 1/self.sample_rate)
        
        # Filter to frequency range of interest
        valid_idx = (freqs >= self.min_freq) & (freqs <= self.max_freq)
        if not np.any(valid_idx):
            return None
            
        # Find peak frequency
        peak_idx = np.argmax(magnitude[valid_idx])
        peak_freq = freqs[valid_idx][peak_idx]
        
        # Check if peak is significant
        peak_magnitude = magnitude[valid_idx][peak_idx]
        avg_magnitude = np.mean(magnitude[valid_idx])
        
        if peak_magnitude < 3 * avg_magnitude:  # Peak should be at least 3x average
            return None
            
        return peak_freq
    
    def frequency_to_note(self, frequency: float) -> Tuple[str, int, float]:
        """
        Convert frequency to musical note name and octave.
        
        Args:
            frequency: Frequency in Hz
            
        Returns:
            Tuple of (note_name, octave, cents_offset)
        """
        # A4 = 440 Hz
        A4_freq = 440.0
        A4_note_number = 69  # MIDI note number for A4
        
        # Calculate semitone distance from A4
        note_number = 12 * np.log2(frequency / A4_freq) + A4_note_number
        
        # Round to nearest semitone
        rounded_note_number = round(note_number)
        
        # Calculate cents offset (100 cents = 1 semitone)
        cents_offset = (note_number - rounded_note_number) * 100
        
        # Get note name and octave
        note_index = rounded_note_number % 12
        octave = rounded_note_number // 12 - 1
        note_name = self.note_names[note_index]
        
        return note_name, octave, cents_offset
    
    def format_note_info(self, frequency: float) -> str:
        """
        Format pitch detection result as a readable string.
        
        Args:
            frequency: Detected frequency in Hz
            
        Returns:
            Formatted string with note information
        """
        note_name, octave, cents_offset = self.frequency_to_note(frequency)
        
        # Format cents with direction
        if abs(cents_offset) < 5:
            cents_str = "in tune"
        elif cents_offset > 0:
            cents_str = f"+{cents_offset:.0f} cents"
        else:
            cents_str = f"{cents_offset:.0f} cents"
            
        # Bold the note name using ANSI escape codes
        bold_note = f"\033[1m{note_name}{octave}\033[0m"
        
        return f"{bold_note} ({frequency:.1f} Hz) - {cents_str}" 