#!/usr/bin/env python3
"""Generate seed cabinet IR (Impulse Response) .wav files for the Pi Multi-FX Pedal.

Each IR models a different speaker cabinet type using DSP techniques:
- Speaker low-pass roll-off (cone diameter determines cutoff)
- Cabinet resonance peaks (Helmholtz resonance)
- Cone breakup modes (notch filters for mechanical resonances)
- Exponential decay envelope (room-dependent)
- Microphone proximity effect (slight high-end roll-off)

Output: ~/.pedal/irs/ directory or a bundled factory-irs path.

Usage:
    python3 generate_seed_irs.py [--dest DIR] [--mono]
"""

from __future__ import annotations

import argparse
import logging
import sys
from pathlib import Path

import numpy as np
from scipy.io import wavfile
from scipy.signal import butter, lfilter, freqz, sosfilt

logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger(__name__)

SAMPLE_RATE = 48000  # Fixed 48kHz for the pedal


# ── Filter helpers ───────────────────────────────────────────────────


def _lowpass_sos(cutoff_hz: float, order: int = 2) -> np.ndarray:
    """2nd-order Butterworth low-pass filter as SOS array."""
    sos = butter(order, cutoff_hz, btype="low", fs=SAMPLE_RATE, output="sos")
    return sos


def _highpass_sos(cutoff_hz: float, order: int = 2) -> np.ndarray:
    """2nd-order Butterworth high-pass filter as SOS array."""
    sos = butter(order, cutoff_hz, btype="high", fs=SAMPLE_RATE, output="sos")
    return sos


def _peak_sos(freq_hz: float, q: float, gain_db: float) -> np.ndarray:
    """Parametric peak EQ as SOS array (biquad)."""
    A = 10 ** (gain_db / 40.0)
    omega = 2.0 * np.pi * freq_hz / SAMPLE_RATE
    alpha = np.sin(omega) / (2.0 * q)
    b0 = 1.0 + alpha * A
    b1 = -2.0 * np.cos(omega)
    b2 = 1.0 - alpha * A
    a0 = 1.0 + alpha / A
    a1 = -2.0 * np.cos(omega)
    a2 = 1.0 - alpha / A
    # Normalize
    b = np.array([b0, b1, b2]) / a0
    a = np.array([1.0, a1 / a0, a2 / a0])
    # Convert to SOS
    sos = np.array([[b[0], b[1], b[2], a[0], a[1], a[2]]])
    return sos


def _notch_sos(freq_hz: float, q: float) -> np.ndarray:
    """Notch filter as SOS array."""
    omega = 2.0 * np.pi * freq_hz / SAMPLE_RATE
    alpha = np.sin(omega) / (2.0 * q)
    b = np.array([1.0, -2.0 * np.cos(omega), 1.0])
    a = np.array([1.0 + alpha, -2.0 * np.cos(omega), 1.0 - alpha])
    sos = np.array([[b[0] / a[0], b[1] / a[0], b[2] / a[0], 1.0, a[1] / a[0], a[2] / a[0]]])
    return sos


# ── Cabinet IR model ─────────────────────────────────────────────────


def _generate_cabinet_ir(
    num_taps: int,
    lowpass_cutoff: float,
    lowpass_order: int,
    highpass_cutoff: float,
    resonance_peaks: list[tuple[float, float, float]],  # (freq_hz, q, gain_db)
    notch_freqs: list[tuple[float, float]],  # (freq_hz, q)
    decay_time_ms: float,
) -> np.ndarray:
    """Generate a synthetic cabinet IR using cascaded filters + decay envelope.

    Args:
        num_taps: Length of the IR in samples.
        lowpass_cutoff: Low-pass filter cutoff (speaker roll-off) in Hz.
        lowpass_order: Order of the low-pass filter (2 or 4).
        highpass_cutoff: High-pass filter cutoff (cabinet resonance) in Hz.
        resonance_peaks: List of (freq_Hz, Q, gain_dB) peak filters.
        notch_freqs: List of (freq_Hz, Q) notch filters for cone breakup.
        decay_time_ms: T60-like decay time in milliseconds.
    
    Returns:
        float32 numpy array of the IR, shape (num_taps,).
    """
    # Start with a Dirac impulse
    ir = np.zeros(num_taps, dtype=np.float64)
    ir[0] = 1.0

    # Apply filters in series
    # 1. Low-pass (speaker roll-off)
    sos = _lowpass_sos(lowpass_cutoff, lowpass_order)
    ir = sosfilt(sos, ir)

    # 2. High-pass (cabinet resonance - prevents subsonic rumble)
    sos = _highpass_sos(highpass_cutoff, 1)
    ir = sosfilt(sos, ir)

    # 3. Resonance peaks (cabinet Helmholtz + mic position)
    for freq, q, gain_db in resonance_peaks:
        sos = _peak_sos(freq, q, gain_db)
        ir = sosfilt(sos, ir)

    # 4. Notch filters (cone breakup / standing waves)
    for freq, q in notch_freqs:
        sos = _notch_sos(freq, q)
        ir = sosfilt(sos, ir)

    # 5. Decay envelope (exponential, with slight room tail)
    decay_samples = int(decay_time_ms * SAMPLE_RATE / 1000.0)
    envelope = np.exp(-np.arange(num_taps) / max(1, decay_samples))
    ir *= envelope

    # Normalize to peak = 0.95 (headroom)
    peak = np.max(np.abs(ir))
    if peak > 0:
        ir /= peak * 1.0526  # normalize to ~0.95

    return ir.astype(np.float32)


# ── Cabinet definitions ─────────────────────────────────────────────

# Each definition creates one IR file.
# Parameters derived from published measurements of real cabinets
# (frequency responses, not IRs themselves, so no copyright issue).

CABINET_SPECS: list[dict] = [
    {
        "name": "vintage-1x12",
        "display": "Vintage 1x12 — Fender-style Open Back",
        "description": "Warm, scooped cleans with bell-like top end. "
        "Models a Fender-style 1x12 open-back combo.",
        "num_taps": 2048,
        "lowpass_cutoff": 4800,
        "lowpass_order": 2,
        "highpass_cutoff": 75,
        "resonance_peaks": [
            (100, 2.0, 4.0),     # Cab resonance hump
            (800, 1.5, -3.0),    # Scooped mids
            (2800, 3.0, 2.0),   # Presence peak
        ],
        "notch_freqs": [
            (1200, 8.0),         # Minor cone mode
            (3500, 10.0),        # Cone edge resonance
        ],
        "decay_time_ms": 60,
    },
    {
        "name": "british-4x12",
        "display": "British 4x12 — Marshall-style Closed Back",
        "description": "Mid-forward, aggressive rock tones with tight low end. "
        "Models a Marshall 1960A-style 4x12 closed-back cabinet.",
        "num_taps": 4096,
        "lowpass_cutoff": 5200,
        "lowpass_order": 4,
        "highpass_cutoff": 80,
        "resonance_peaks": [
            (110, 1.8, 5.0),     # Big cab resonance
            (700, 1.2, 4.0),     # Mid-forward punch
            (1500, 2.0, 3.0),    # Upper mid grind
            (3200, 2.5, 2.5),    # Presence
        ],
        "notch_freqs": [
            (2500, 12.0),        # Cross-over null
            (4000, 8.0),         # Speaker breakup
        ],
        "decay_time_ms": 85,
    },
    {
        "name": "american-2x12",
        "display": "American 2x12 — Vox-style Open Back",
        "description": "Chimey, complex midrange with sparkling highs. "
        "Models a Vox AC30-style 2x12 open-back cabinet.",
        "num_taps": 2048,
        "lowpass_cutoff": 5800,
        "lowpass_order": 2,
        "highpass_cutoff": 70,
        "resonance_peaks": [
            (120, 1.5, 3.0),     # Cab resonance
            (600, 1.0, 2.0),     # Low-mid body
            (1300, 2.0, 5.0),    # Chime / complex mids
            (3500, 3.0, 3.0),    # Top-end sparkle
        ],
        "notch_freqs": [
            (1800, 10.0),        # Minor comb filter
            (4200, 12.0),        # Cone resonance
        ],
        "decay_time_ms": 55,
    },
    {
        "name": "modern-4x12",
        "display": "Modern 4x12 — Mesa/Boogie-style Closed Back",
        "description": "Tight low-end, aggressive mids, smooth highs. "
        "Models a Mesa Rectifier-style 4x12 closed-back cab.",
        "num_taps": 4096,
        "lowpass_cutoff": 4800,
        "lowpass_order": 4,
        "highpass_cutoff": 85,
        "resonance_peaks": [
            (95, 2.5, 6.0),      # Tight low-end thump
            (600, 1.5, 2.0),     # Low-mid body
            (1000, 1.0, 5.0),    # Aggressive mid bark
            (2200, 3.0, 1.0),    # Upper mid cut
            (3200, 2.0, -1.0),   # Slight high-end roll for smoothness
        ],
        "notch_freqs": [
            (450, 15.0),         # Subsonic resonance cleanup
            (2800, 10.0),        # Cone breakup suppression
        ],
        "decay_time_ms": 90,
    },
    {
        "name": "jazz-1x15",
        "display": "Jazz 1x15 — Deep Open Back",
        "description": "Deep, warm, scooped tone for jazz cleans. "
        "Models a 15-inch speaker in a large open-back cab.",
        "num_taps": 2048,
        "lowpass_cutoff": 4500,
        "lowpass_order": 2,
        "highpass_cutoff": 55,
        "resonance_peaks": [
            (80, 2.0, 8.0),      # Deep low-end warmth
            (500, 2.0, -4.0),    # Scooped mids
            (2500, 2.0, 3.0),    # Presence for articulation
        ],
        "notch_freqs": [
            (200, 8.0),          # Cab resonance smoothing
            (3800, 8.0),         # Cone edge roll-off
        ],
        "decay_time_ms": 70,
    },
    {
        "name": "boutique-1x12",
        "display": "Boutique 1x12 — Dumble-style Open Back",
        "description": "Smooth, rounded cleans with enhanced mid complexity. "
        "Models a Dumble-style 1x12 open-back combo.",
        "num_taps": 2048,
        "lowpass_cutoff": 5200,
        "lowpass_order": 2,
        "highpass_cutoff": 78,
        "resonance_peaks": [
            (105, 2.0, 4.0),     # Low-end body
            (800, 1.8, 5.0),     # Complex mid character
            (1800, 2.5, 2.0),    # Mid sparkle
            (3000, 2.0, 1.5),    # Top presence
        ],
        "notch_freqs": [
            (1500, 12.0),        # Smoothing
            (3500, 10.0),
        ],
        "decay_time_ms": 65,
    },
    {
        "name": "mini-1x8",
        "display": "Mini 1x8 — Small Practice Amp",
        "description": "Lo-fi, boxy tone for vintage radio-style sounds. "
        "Models a small 8-inch practice amp speaker.",
        "num_taps": 1024,
        "lowpass_cutoff": 3800,
        "lowpass_order": 2,
        "highpass_cutoff": 90,
        "resonance_peaks": [
            (150, 1.5, 6.0),     # Boxy resonance
            (1000, 1.0, 4.0),    # Nasally mids
            (2500, 1.5, -2.0),   # Rolled-off highs
        ],
        "notch_freqs": [
            (3000, 6.0),
        ],
        "decay_time_ms": 40,
    },
]


def generate_all_irs(dest_dir: Path, verify: bool = True) -> list[Path]:
    """Generate all seed IR files into dest_dir.

    Returns:
        List of created file paths (sorted by name).
    """
    dest_dir.mkdir(parents=True, exist_ok=True)
    created: list[Path] = []

    for spec in CABINET_SPECS:
        name = spec["name"]
        path = dest_dir / f"{name}.wav"

        logger.info(f"  Generating {name}.wav ...")

        ir = _generate_cabinet_ir(
            num_taps=spec["num_taps"],
            lowpass_cutoff=spec["lowpass_cutoff"],
            lowpass_order=spec["lowpass_order"],
            highpass_cutoff=spec["highpass_cutoff"],
            resonance_peaks=spec["resonance_peaks"],
            notch_freqs=spec["notch_freqs"],
            decay_time_ms=spec["decay_time_ms"],
        )

        # Write as 16-bit WAV (smaller, compatible with scipy)
        int16_data = (ir * 32767).astype(np.int16)
        wavfile.write(str(path), SAMPLE_RATE, int16_data)

        # Quick verification
        file_size = path.stat().st_size
        duration_ms = (spec["num_taps"] / SAMPLE_RATE) * 1000

        logger.info(
            f"    → {path.name}  ({spec['num_taps']} taps, "
            f"{duration_ms:.1f}ms, {file_size // 1024}KB)"
        )

        if verify:
            # Read back and check
            sr, data = wavfile.read(str(path))
            assert sr == SAMPLE_RATE, f"Sample rate mismatch: {sr}"
            assert len(data) == spec["num_taps"], f"Length mismatch: {len(data)}"
            peak = np.max(np.abs(data.astype(np.float64))) / 32767.0
            assert peak > 0.1, f"IR {name} peak too low: {peak:.3f}"

        created.append(path)

    return created


def verify_quality(ir_dir: Path) -> None:
    """Verify that generated IRs have reasonable frequency response."""
    import matplotlib
    matplotlib.use("Agg")  # headless
    import matplotlib.pyplot as plt

    fig, axes = plt.subplots(len(CABINET_SPECS), 1, figsize=(10, 2 * len(CABINET_SPECS)))
    if len(CABINET_SPECS) == 1:
        axes = [axes]

    for ax, spec in zip(axes, CABINET_SPECS):
        path = ir_dir / f"{spec['name']}.wav"
        if not path.exists():
            continue
        sr, data = wavfile.read(str(path))
        # Normalize
        data_f = data.astype(np.float64) / 32767.0
        # FFT for frequency response
        fft = np.abs(np.fft.rfft(data_f, n=4096))
        freqs = np.fft.rfftfreq(4096, d=1.0 / sr)
        # dB scale, normalized
        fft_db = 20 * np.log10(fft / np.max(fft) + 1e-10)
        ax.semilogx(freqs[1:], fft_db[1:], linewidth=0.8)
        ax.axhline(-3, color="gray", linestyle=":", linewidth=0.5)
        ax.axhline(-12, color="gray", linestyle=":", linewidth=0.5)
        ax.set_xlim(20, sr / 2)
        ax.set_ylim(-48, 3)
        ax.set_title(f"{spec['display']} — Freq Response")
        ax.set_xlabel("Frequency (Hz)")
        ax.set_ylabel("dB")
        ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plot_path = ir_dir / "frequency_response.png"
    fig.savefig(plot_path, dpi=150)
    plt.close(fig)
    logger.info(f"Frequency response plot: {plot_path}")


def create_manifest(ir_dir: Path, created: list[Path]) -> Path:
    """Create a README manifest for the IR files."""
    manifest = ir_dir / "README.md"
    lines = [
        "# Pi Multi-FX Pedal — Default Cabinet IR Files",
        "",
        "Synthetic impulse responses (.wav) for guitar cabinet simulation.",
        "Generated using DSP filter cascades (Butterworth LP/HP + parametric",
        "peak/notch filters + exponential decay envelope).",
        "",
        "**All IRs are 48kHz, 16-bit mono WAV files.**",
        "",
        "| File | Cabinet Model | Taps | Length | Description |",
        "|------|--------------|------|--------|-------------|",
    ]

    for spec in CABINET_SPECS:
        path = ir_dir / f"{spec['name']}.wav"
        if path.exists():
            num_taps = spec["num_taps"]
            length_ms = (num_taps / SAMPLE_RATE) * 1000
            lines.append(
                f"| `{spec['name']}.wav` | {spec['display']} | "
                f"{num_taps} | {length_ms:.0f}ms | {spec['description']} |"
            )

    lines.extend([
        "",
        "## Usage",
        "",
        "These IRs are loaded by the `IRLoader` class through the pedal's cab simulation.",
        "Place them in `~/.pedal/irs/` or reference them by absolute path in preset chain blocks.",
        "",
        "## Replacement",
        "",
        "Replace any `.wav` file with a real captured IR (48kHz, 16-bit mono) to",
        "upgrade from synthetic to authentic cabinet tone. The pedal treats all",
        "`.wav` files in the IR directory identically.",
        "",
        "## Cabinet Type Guide",
        "",
        "- **Vintage 1x12** — Fender Deluxe / Princeton-style cleans. Scooped mids, warm lows.",
        "- **British 4x12** — Marshall 1960 / JCM-style crunch and rock. Mid-forward, aggressive.",
        "- **American 2x12** — Vox AC30-style chime. Complex mids, sparkling treble.",
        "- **Modern 4x12** — Mesa Rectifier-style high gain. Tight lows, smooth highs.",
        "- **Jazz 1x15** — Polytone / Henriksen-style jazz. Deep lows, scooped mids.",
        "- **Boutique 1x12** — Dumble / Two-Rock-style. Smooth, complex midrange.",
        "- **Mini 1x8** — Small practice amp. Boxy, lo-fi, vintage radio tone.",
    ])

    manifest.write_text("\n".join(lines) + "\n")
    logger.info(f"Manifest written: {manifest}")
    return manifest


# ── Main ─────────────────────────────────────────────────────────────


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate seed cabinet IR files")
    parser.add_argument(
        "--dest",
        default=str(Path.home() / ".pedal" / "irs"),
        help="Output directory (default: ~/.pedal/irs)",
    )
    parser.add_argument(
        "--plot",
        action="store_true",
        help="Generate frequency response verification plot",
    )
    args = parser.parse_args()

    dest = Path(args.dest)
    logger.info(f"Generating {len(CABINET_SPECS)} seed cabinet IRs → {dest}")
    logger.info("")

    created = generate_all_irs(dest, verify=True)

    logger.info("")
    logger.info(f"Generated {len(created)} IR files successfully.")

    if args.plot:
        verify_quality(dest)

    create_manifest(dest, created)

    logger.info(f"\nAll IRs in: {dest.resolve()}")


if __name__ == "__main__":
    main()