"""
Analyzer — Performs frequency response measurement.

Compares a recorded signal against a reference signal to determine
the transfer function (frequency response) of the tape deck.
"""

import numpy as np
from scipy import signal
from typing import Tuple, Dict

def align_signals(reference: np.ndarray, recorded: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """
    Align recorded signal to the reference signal using cross-correlation.
    Accounts for latency introduced by the audio interface and tape path.
    """
    # Use cross-correlation to find the time delay
    # We correlate a portion of the signals for speed if they are long
    max_len = min(len(reference), len(recorded), 44100 * 2) # use max 2 seconds for alignment
    
    ref_slice = reference[:max_len]
    rec_slice = recorded[:max_len]
    
    # correlation
    corr = signal.correlate(rec_slice, ref_slice, mode='full')
    delay = np.argmax(corr) - len(ref_slice) + 1
    
    if delay > 0:
        # Recorded is delayed relative to reference
        aligned_rec = recorded[delay:]
        aligned_ref = reference[:len(aligned_rec)]
    elif delay < 0:
        # Should not happen physically, but handle just in case
        aligned_ref = reference[-delay:]
        aligned_rec = recorded[:len(aligned_ref)]
    else:
        aligned_ref = reference
        aligned_rec = recorded
        
    # Trim to same length
    min_len = min(len(aligned_ref), len(aligned_rec))
    
    return aligned_ref[:min_len], aligned_rec[:min_len]

def calculate_transfer_function(
    reference: np.ndarray, 
    recorded: np.ndarray, 
    sample_rate: int = 44100,
    nperseg: int = 4096
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Calculate the transfer function H(f) = cross_power(ref, rec) / auto_power(ref).
    Using Welch's method for better noise reduction.
    
    Returns:
        freqs: Frequency bins (Hz)
        response_db: Magnitude response in dB
    """
    # Ensure aligned
    ref_aligned, rec_aligned = align_signals(reference, recorded)
    
    if len(ref_aligned) < nperseg:
        # Fallback if signal is too short
        nperseg = len(ref_aligned)
        
    # Calculate cross power spectral density (Pxy) and power spectral density of reference (Pxx)
    freqs, Pxy = signal.csd(ref_aligned, rec_aligned, fs=sample_rate, nperseg=nperseg)
    _, Pxx = signal.welch(ref_aligned, fs=sample_rate, nperseg=nperseg)
    
    # Transfer function H = Pxy / Pxx
    # Add small epsilon to avoid division by zero
    H = Pxy / (Pxx + 1e-10)
    
    # Magnitude response
    magnitude = np.abs(H)
    
    # Convert to dB
    response_db = 20 * np.log10(magnitude + 1e-10)
    
    return freqs, response_db

def smooth_curve(freqs: np.ndarray, response_db: np.ndarray, fraction: float = 1/6) -> np.ndarray:
    """
    Smooth the frequency response curve using fractional octave smoothing.
    fraction = 1/3 for third-octave, 1/6 for sixth-octave.
    """
    smoothed = np.zeros_like(response_db)
    
    for i, f in enumerate(freqs):
        if f == 0:
            smoothed[i] = response_db[i]
            continue
            
        # Octave bandwidth around frequency f
        f_lower = f * (2.0 ** (-fraction / 2.0))
        f_upper = f * (2.0 ** (fraction / 2.0))
        
        # Find indices within this band
        idx = np.where((freqs >= f_lower) & (freqs <= f_upper))[0]
        
        if len(idx) > 0:
            smoothed[i] = np.mean(response_db[idx])
        else:
            smoothed[i] = response_db[i]
            
    return smoothed