"""
Calibrator — Auto-fits EQ bands to compensate for measured frequency response.
"""

import numpy as np
from scipy.optimize import minimize
from typing import List, Tuple, Dict
from .eq_engine import EQBand, EQEngine

TAPE_PROFILES = {
    "Type I (Ferric)": {"low_cut": 50.0, "high_cut": 15000.0, "eq_us": 120},
    "Type II (Chrome)": {"low_cut": 45.0, "high_cut": 16000.0, "eq_us": 70},
    "Type IV (Metal)": {"low_cut": 40.0, "high_cut": 17000.0, "eq_us": 70},
}

def calculate_correction_curve(response_db: np.ndarray, max_correction: float = 15.0) -> np.ndarray:
    """
    Invert the measured response to get the correction curve.
    Limits the maximum boost/cut to prevent extreme EQ values.
    """
    # Invert the response (if tape cuts 5dB, we boost 5dB)
    correction = -response_db
    
    # Normalize to 0dB mean around 1kHz (typical reference)
    # Assume response_db is already somewhat normalized, or just mean shift
    mean_level = np.mean(correction)
    correction -= mean_level
    
    # Clip to max correction
    correction = np.clip(correction, -max_correction, max_correction)
    
    return correction

def _evaluate_eq_curve(freqs: np.ndarray, bands: List[EQBand], sample_rate: int) -> np.ndarray:
    """
    Calculate the combined frequency response of the EQ bands.
    """
    # Create temporary EQ engine to calculate response
    engine = EQEngine(sample_rate=sample_rate)
    engine.bands = bands
    engine.update_filters()
    
    # If no filters, return flat response
    if engine._sos is None:
        return np.zeros_like(freqs)
        
    # Calculate frequency response using sosfreqz
    # freqs normalized to nyquist
    w = 2 * np.pi * freqs / sample_rate
    
    from scipy.signal import sosfreqz
    _, h = sosfreqz(engine._sos, worN=w)
    
    return 20 * np.log10(np.abs(h) + 1e-10)

def fit_eq_bands(
    freqs: np.ndarray, 
    correction_curve: np.ndarray, 
    num_peaking_bands: int = 10,
    sample_rate: int = 44100,
    tape_type: str = "Type I (Ferric)"
) -> List[EQBand]:
    """
    Fit parametric EQ bands to the target correction curve, constrained by tape type.
    """
    profile = TAPE_PROFILES.get(tape_type, TAPE_PROFILES["Type I (Ferric)"])
    low_cut = profile["low_cut"]
    high_cut = profile["high_cut"]
    
    # Initialize bands
    bands = []
    
    # Instead of shelves that fight the natural roll-off, use peaking bands 
    # distributed strictly within the usable frequency range of the tape.
    # We add 2 extra bands instead of the shelves to give 12 peaking bands total.
    total_bands = num_peaking_bands + 2
    
    # Log-spaced center frequencies
    f_centers = np.logspace(np.log10(low_cut * 1.5), np.log10(high_cut * 0.8), total_bands)
    for f in f_centers:
        bands.append(EQBand(freq=float(f), gain=0.0, q=1.414, filter_type='peaking'))
    
    # Optimization target function
    def error_func(params):
        # params is a flattened array of [gain, q] for each band
        # Keep frequencies fixed for stability
        temp_bands = []
        for i, b in enumerate(bands):
            gain = params[i*2]
            q = max(0.1, min(10.0, params[i*2+1])) # Constrain Q
            temp_bands.append(EQBand(
                freq=b.freq,
                gain=gain,
                q=q,
                filter_type=b.filter_type
            ))
            
        eq_curve = _evaluate_eq_curve(freqs, temp_bands, sample_rate)
        
        # Weighted error: ignore extremes outside tape limits, focus on mids
        weights = np.zeros_like(freqs)
        # Weight 1.0 for usable range
        valid_mask = (freqs >= low_cut) & (freqs <= high_cut)
        weights[valid_mask] = 1.0
        # Extra weight for critical mid-range (300Hz - 5kHz)
        mid_mask = (freqs >= 300) & (freqs <= 5000)
        weights[mid_mask] = 2.0
        
        # Mean squared error
        mse = np.average((eq_curve - correction_curve)**2, weights=weights)
        return mse
        
    # Initial guess: flat 0dB, default Q
    initial_guess = np.zeros(len(bands) * 2)
    for i in range(len(bands)):
        initial_guess[i*2+1] = 1.414 if bands[i].filter_type == 'peaking' else 0.707
        
    # Bounds: Gain [-15, 15], Q [0.1, 10]
    bounds = []
    for _ in range(len(bands)):
        bounds.append((-15.0, 15.0)) # Gain
        bounds.append((0.1, 10.0))   # Q
        
    # Run optimization
    result = minimize(
        error_func, 
        initial_guess, 
        bounds=bounds,
        method='L-BFGS-B',
        options={'maxiter': 100}
    )
    
    # Apply optimized parameters
    optimized_bands = []
    for i, b in enumerate(bands):
        gain = float(result.x[i*2])
        q = float(result.x[i*2+1])
        optimized_bands.append(EQBand(
            freq=b.freq,
            gain=gain,
            q=q,
            filter_type=b.filter_type
        ))
        
    return optimized_bands