"""
Parametric EQ Engine — Applies biquad filters to audio streams.
"""

import numpy as np
import scipy.signal as signal
from dataclasses import dataclass
from typing import List, Optional

@dataclass
class EQBand:
    """Parameters for a single EQ band."""
    freq: float       # Center frequency in Hz
    gain: float       # Gain in dB (not used for HP/LP)
    q: float          # Quality factor
    filter_type: str  # 'peaking', 'lowshelf', 'highshelf', 'highpass', 'lowpass'
    enabled: bool = True

class EQEngine:
    """
    Processes audio through a chain of parametric EQ bands.
    Uses Second-Order Sections (SOS) for numerical stability.
    """
    def __init__(self, sample_rate: int = 44100):
        self.sample_rate = sample_rate
        self.bands: List[EQBand] = []
        self._sos = None
        self._zi = None  # Filter state for continuous processing
        
        # Default initialization: 1 low shelf, 10 peaking, 1 high shelf
        self.reset_to_default_bands()

    def reset_to_default_bands(self):
        """Setup a default 12-band EQ."""
        self.bands = [
            EQBand(freq=100.0, gain=0.0, q=0.707, filter_type='lowshelf'),
            EQBand(freq=250.0, gain=0.0, q=1.414, filter_type='peaking'),
            EQBand(freq=500.0, gain=0.0, q=1.414, filter_type='peaking'),
            EQBand(freq=1000.0, gain=0.0, q=1.414, filter_type='peaking'),
            EQBand(freq=2000.0, gain=0.0, q=1.414, filter_type='peaking'),
            EQBand(freq=4000.0, gain=0.0, q=1.414, filter_type='peaking'),
            EQBand(freq=8000.0, gain=0.0, q=1.414, filter_type='peaking'),
            EQBand(freq=12000.0, gain=0.0, q=0.707, filter_type='highshelf'),
        ]
        self.update_filters()

    def _get_biquad_coeffs(self, band: EQBand) -> tuple:
        """Calculate biquad coefficients based on Audio EQ Cookbook."""
        A = 10.0 ** (band.gain / 40.0)
        w0 = 2.0 * np.pi * band.freq / self.sample_rate
        alpha = np.sin(w0) / (2.0 * band.q)
        
        if band.filter_type == 'peaking':
            b0 = 1.0 + alpha * A
            b1 = -2.0 * np.cos(w0)
            b2 = 1.0 - alpha * A
            a0 = 1.0 + alpha / A
            a1 = -2.0 * np.cos(w0)
            a2 = 1.0 - alpha / A
        elif band.filter_type == 'lowshelf':
            sqrtA = np.sqrt(A)
            b0 = A * ((A + 1.0) - (A - 1.0) * np.cos(w0) + 2.0 * sqrtA * alpha)
            b1 = 2.0 * A * ((A - 1.0) - (A + 1.0) * np.cos(w0))
            b2 = A * ((A + 1.0) - (A - 1.0) * np.cos(w0) - 2.0 * sqrtA * alpha)
            a0 = (A + 1.0) + (A - 1.0) * np.cos(w0) + 2.0 * sqrtA * alpha
            a1 = -2.0 * ((A - 1.0) + (A + 1.0) * np.cos(w0))
            a2 = (A + 1.0) + (A - 1.0) * np.cos(w0) - 2.0 * sqrtA * alpha
        elif band.filter_type == 'highshelf':
            sqrtA = np.sqrt(A)
            b0 = A * ((A + 1.0) + (A - 1.0) * np.cos(w0) + 2.0 * sqrtA * alpha)
            b1 = -2.0 * A * ((A - 1.0) + (A + 1.0) * np.cos(w0))
            b2 = A * ((A + 1.0) + (A - 1.0) * np.cos(w0) - 2.0 * sqrtA * alpha)
            a0 = (A + 1.0) - (A - 1.0) * np.cos(w0) + 2.0 * sqrtA * alpha
            a1 = 2.0 * ((A - 1.0) - (A + 1.0) * np.cos(w0))
            a2 = (A + 1.0) - (A - 1.0) * np.cos(w0) - 2.0 * sqrtA * alpha
        elif band.filter_type == 'highpass':
            b0 = (1.0 + np.cos(w0)) / 2.0
            b1 = -(1.0 + np.cos(w0))
            b2 = (1.0 + np.cos(w0)) / 2.0
            a0 = 1.0 + alpha
            a1 = -2.0 * np.cos(w0)
            a2 = 1.0 - alpha
        elif band.filter_type == 'lowpass':
            b0 = (1.0 - np.cos(w0)) / 2.0
            b1 = 1.0 - np.cos(w0)
            b2 = (1.0 - np.cos(w0)) / 2.0
            a0 = 1.0 + alpha
            a1 = -2.0 * np.cos(w0)
            a2 = 1.0 - alpha
        else:
            raise ValueError(f"Unknown filter type: {band.filter_type}")
            
        return np.array([b0, b1, b2, a0, a1, a2])

    def update_filters(self):
        """Recompute SOS coefficients from all enabled bands."""
        sos_list = []
        for band in self.bands:
            if band.enabled:
                coeffs = self._get_biquad_coeffs(band)
                # Normalize by a0
                a0 = coeffs[3]
                sos_section = np.array([coeffs[0]/a0, coeffs[1]/a0, coeffs[2]/a0, 
                                        1.0, coeffs[4]/a0, coeffs[5]/a0])
                sos_list.append(sos_section)
                
        if not sos_list:
            self._sos = None
            self._zi = None
        else:
            self._sos = np.vstack(sos_list)
            # Reset zi to match the new filter order (number of sections)
            # zi shape needs to handle stereo if needed, but for now we'll 
            # initialize it on the first process() call based on input shape.
            self._zi = None

    def set_sample_rate(self, sr: int):
        if self.sample_rate != sr:
            self.sample_rate = sr
            self.update_filters()

    def process(self, chunk: np.ndarray) -> np.ndarray:
        """
        Process a chunk of audio data.
        chunk: shape (N_samples, channels)
        """
        if self._sos is None:
            return chunk

        # Initialize or reshape filter state zi if needed
        n_sections = self._sos.shape[0]
        n_channels = chunk.shape[1] if chunk.ndim > 1 else 1
        
        if self._zi is None or self._zi.shape[2] != n_channels or self._zi.shape[0] != n_sections:
             # zi needs shape (n_sections, 2, n_channels) for sosfilt
             self._zi = np.zeros((n_sections, 2, n_channels))

        # Apply filter and update state
        out, self._zi = signal.sosfilt(self._sos, chunk, axis=0, zi=self._zi)
        return out.astype(np.float32)

    def reset_state(self):
        """Clear filter state (call when stopping/starting playback to avoid glitches)."""
        self._zi = None

    def save_bands(self, filepath: str):
        import json
        data = []
        for b in self.bands:
            data.append({
                'freq': b.freq,
                'gain': b.gain,
                'q': b.q,
                'filter_type': b.filter_type,
                'enabled': b.enabled
            })
        with open(filepath, 'w') as f:
            json.dump(data, f, indent=4)
            
    def load_bands(self, filepath: str):
        import json
        with open(filepath, 'r') as f:
            data = json.load(f)
        
        new_bands = []
        for d in data:
            new_bands.append(EQBand(
                freq=d['freq'],
                gain=d['gain'],
                q=d['q'],
                filter_type=d['filter_type'],
                enabled=d.get('enabled', True)
            ))
        self.bands = new_bands
        self.update_filters()