import numpy as np
import matplotlib.pyplot as plt
from scipy import signal


def create_filter_coeffs(filter_type, fc, gain_db, q, fs=48000):
    w0 = 2 * np.pi * fc / fs
    A = 10 ** (gain_db / 40)
    alpha = np.sin(w0) / (2 * q)

    if filter_type in ["PK"]:  # Peaking EQ
        b0 = 1 + alpha * A
        b1 = -2 * np.cos(w0)
        b2 = 1 - alpha * A
        a0 = 1 + alpha / A
        a1 = -2 * np.cos(w0)
        a2 = 1 - alpha / A

    elif filter_type in ["LSC", "LS"]:  # Low Shelf
        b0 = A * ((A + 1) - (A - 1) * np.cos(w0) + 2 * np.sqrt(A) * alpha)
        b1 = 2 * A * ((A - 1) - (A + 1) * np.cos(w0))
        b2 = A * ((A + 1) - (A - 1) * np.cos(w0) - 2 * np.sqrt(A) * alpha)
        a0 = (A + 1) + (A - 1) * np.cos(w0) + 2 * np.sqrt(A) * alpha
        a1 = -2 * ((A - 1) + (A + 1) * np.cos(w0))
        a2 = (A + 1) + (A - 1) * np.cos(w0) - 2 * np.sqrt(A) * alpha

    elif filter_type in ["HSC", "HS"]:  # High Shelf
        b0 = A * ((A + 1) + (A - 1) * np.cos(w0) + 2 * np.sqrt(A) * alpha)
        b1 = -2 * A * ((A - 1) + (A + 1) * np.cos(w0))
        b2 = A * ((A + 1) + (A - 1) * np.cos(w0) - 2 * np.sqrt(A) * alpha)
        a0 = (A + 1) - (A - 1) * np.cos(w0) + 2 * np.sqrt(A) * alpha
        a1 = 2 * ((A - 1) - (A + 1) * np.cos(w0))
        a2 = (A + 1) - (A - 1) * np.cos(w0) - 2 * np.sqrt(A) * alpha
    else:
        raise ValueError(f"Unknown filter type: {filter_type}")

    return [b0 / a0, b1 / a0, b2 / a0], [1.0, a1 / a0, a2 / a0]


def parse_filter_config(config_text):
    filters = []
    for line in config_text.strip().split("\n"):
        if line.startswith("Filter: ON"):
            parts = line.split()
            ftype = parts[2]
            fc = float(parts[4])
            gain = float(parts[7])
            q = float(parts[10]) if len(parts) > 10 else 0.707
            filters.append((ftype, fc, gain, q))
    return filters


def plot_eq_response(config_text):
    filters = parse_filter_config(config_text)
    fs = 48000
    freq = np.logspace(np.log10(20), np.log10(20000), 1000)
    w = 2 * np.pi * freq / fs

    total_response = np.ones_like(freq)

    plt.figure(figsize=(15, 8))
    for ftype, fc, gain, q in filters:
        b, a = create_filter_coeffs(ftype, fc, gain, q, fs)
        w, h = signal.freqz(b, a, worN=freq, fs=fs)  # type: ignore
        response_db = 20 * np.log10(np.abs(h))
        total_response += response_db
        plt.semilogx(freq, response_db, alpha=0.1, color="blue")

    plt.semilogx(freq, total_response, "r-", linewidth=2, label="Combined Response")
    plt.grid(True)
    plt.xlabel("Frequency [Hz]")
    plt.ylabel("Amplitude [dB]")
    plt.title("Parametric EQ Frequency Response")
    plt.ylim(-30, 20)
    plt.axhline(y=0, color="k", linestyle="-", alpha=0.3)
    plt.legend()
    plt.show()


def merge_filters(filters):
    """Merge multiple filters into single equivalent filter"""
    # Average frequency weighted by gain magnitude
    weights = [abs(f[2]) for f in filters]
    fc = np.average([f[1] for f in filters], weights=weights)

    # Sum gains
    gain = sum(f[2] for f in filters)

    # Average Q weighted by gain
    q = np.average([f[3] for f in filters], weights=weights)

    # Use most common filter type
    ftype = max(set(f[0] for f in filters), key=lambda x: filters.count((x)))

    return (ftype, fc, gain, q)


def calculate_response(filters, freq, fs):
    """Calculate combined frequency response for set of filters"""
    total_response = np.zeros_like(freq)
    for ftype, fc, gain, q in filters:
        b, a = create_filter_coeffs(ftype, fc, gain, q, fs)
        w, h = signal.freqz(b, a, worN=freq, fs=fs)  # type: ignore
        response_db = 20 * np.log10(np.abs(h))
        total_response += response_db
    return total_response


def optimize_eq_config(config_text):
    # Parse original filters
    original_filters = parse_filter_config(config_text)
    fs = 48000
    freq = np.logspace(np.log10(20), np.log10(20000), 1000)

    # Get original response
    orig_response = calculate_response(original_filters, freq, fs)

    # Initial filter grouping
    grouped_filters = []
    current_group = []

    for i, (ftype, fc, gain, q) in enumerate(original_filters):
        if not current_group:
            current_group.append((ftype, fc, gain, q))
            continue

        prev_fc = current_group[-1][1]
        if fc / prev_fc < 1.2:  # Group filters within 20% frequency range
            current_group.append((ftype, fc, gain, q))
        else:
            # Merge group into single filter
            if len(current_group) > 1:
                merged = merge_filters(current_group)
                grouped_filters.append(merged)
            else:
                grouped_filters.extend(current_group)
            current_group = [(ftype, fc, gain, q)]

    # Add remaining group
    if current_group:
        if len(current_group) > 1:
            grouped_filters.append(merge_filters(current_group))
        else:
            grouped_filters.extend(current_group)

    # Generate new config text
    optimized_config = ""
    for ftype, fc, gain, q in grouped_filters:
        optimized_config += (
            f"Filter: ON {ftype} Fc {fc:.0f} Hz Gain {gain:.1f} dB Q {q:.1f}\n"
        )

    # Plot comparison
    plt.figure(figsize=(15, 8))
    plt.semilogx(freq, orig_response, "b-", label="Original", alpha=0.5)
    opt_response = calculate_response(grouped_filters, freq, fs)
    plt.semilogx(freq, opt_response, "r--", label="Optimized")
    plt.grid(True)
    plt.legend()
    plt.xlabel("Frequency [Hz]")
    plt.ylabel("Amplitude [dB]")
    plt.title(
        f"Original ({len(original_filters)} filters) vs Optimized ({len(grouped_filters)} filters)"
    )
    plt.show()

    return optimized_config


def optimize_eq_config_v2(config_text, max_filters=15, tolerance_db=0.5):
    original_filters = parse_filter_config(config_text)
    fs = 48000
    freq = np.logspace(np.log10(20), np.log10(20000), 1000)
    orig_response = calculate_response(original_filters, freq, fs)

    grouped_filters = []
    current_group = []

    # Sort filters by frequency for better grouping
    sorted_filters = sorted(original_filters, key=lambda x: x[1])

    for ftype, fc, gain, q in sorted_filters:
        if not current_group:
            current_group.append((ftype, fc, gain, q))
            continue

        prev_fc = current_group[-1][1]
        if fc / prev_fc < 1.2 and len(grouped_filters) < max_filters:
            current_group.append((ftype, fc, gain, q))
        else:
            if len(current_group) > 1:
                merged = merge_filters(current_group)
                grouped_filters.append(merged)
            else:
                grouped_filters.extend(current_group)
            current_group = [(ftype, fc, gain, q)]

    # Add remaining group
    if current_group:
        if len(current_group) > 1:
            grouped_filters.append(merge_filters(current_group))
        else:
            grouped_filters.extend(current_group)

    # Iteratively remove least significant filters until error exceeds tolerance
    while len(grouped_filters) > max_filters:
        min_error = float("inf")
        best_filters = None

        # Try removing each filter
        for i in range(len(grouped_filters)):
            test_filters = grouped_filters[:i] + grouped_filters[i + 1 :]
            test_response = calculate_response(test_filters, freq, fs)
            error = np.max(np.abs(test_response - orig_response))

            if error < min_error:
                min_error = error
                best_filters = test_filters

        # Stop if error exceeds tolerance
        if min_error > tolerance_db:
            break

        grouped_filters = best_filters

    # Generate optimized config
    optimized_config = ""
    for ftype, fc, gain, q in grouped_filters:
        optimized_config += (
            f"Filter: ON {ftype} Fc {fc:.0f} Hz Gain {gain:.1f} dB Q {q:.1f}\n"
        )

    # Plot comparison
    plt.figure(figsize=(15, 8))
    plt.semilogx(
        freq,
        orig_response,
        "b-",
        label=f"Original ({len(original_filters)} filters)",
        alpha=0.5,
    )
    opt_response = calculate_response(grouped_filters, freq, fs)
    plt.semilogx(
        freq, opt_response, "r--", label=f"Optimized ({len(grouped_filters)} filters)"
    )
    plt.grid(True)
    plt.legend()
    plt.xlabel("Frequency [Hz]")
    plt.ylabel("Amplitude [dB]")
    plt.title(f"Max Error: {np.max(np.abs(opt_response - orig_response)):.1f} dB")
    plt.show()

    return optimized_config


# Call the function with your config text
config_eq = """
Filter: ON LSC Fc 25 Hz Gain -6 dB Q 0.8
Filter: ON PK Fc 42 Hz Gain -1.1 dB Q 2
Filter: ON PK Fc 57 Hz Gain 1.2 dB Q 8
# ^ literally bassboost controller ^
Filter: ON PK Fc 69 Hz Gain 1.6 dB Q 7
Filter: ON PK Fc 75 Hz Gain 1.3 dB Q 2.7
Filter: ON PK Fc 84 Hz Gain -2 dB Q 3.4
Filter: ON PK Fc 115 Hz Gain -4 dB Q 10
Filter: ON LSC Fc 150 Hz Gain 1.5 dB Q 1.1
Filter: ON PK Fc 200 Hz Gain -7 dB Q 8
Filter: ON PK Fc 225 Hz Gain 0.5 dB Q 8
Filter: ON PK Fc 350 Hz Gain -2.5 dB Q 1.2
Filter: ON PK Fc 240 Hz Gain -2 dB Q 0.8
Filter: ON PK Fc 450 Hz Gain -4.5 dB Q 0.82
Filter: ON PK Fc 720 Hz Gain -2.8 dB Q 0.91
Filter: ON PK Fc 950 Hz Gain -1.2 dB Q 7
Filter: ON PK Fc 1900 Hz Gain -2 dB Q 9
Filter: ON PK Fc 1000 Hz Gain -3 dB Q 7
Filter: ON PK Fc 1500 Hz Gain 1.3 dB Q 0.8
Filter: ON PK Fc 2600 Hz Gain -1 dB Q 1.2
Filter: ON PK Fc 4300 Hz Gain -1.2 dB Q 7
Filter: ON PK Fc 4800 Hz Gain -4.2 dB Q 10
Filter: ON PK Fc 3200 Hz Gain 0.1 dB Q 2
Filter: ON PK Fc 5000 Hz Gain -2.3 dB Q 2
Filter: ON PK Fc 8200 Hz Gain -5 dB Q 1
Filter: ON HSC Fc 9600 Hz Gain -1.8 dB Q 1.6
Filter: ON PK Fc 9750 Hz Gain 1.4 dB Q 2.4
Filter: ON PK Fc 12000 Hz Gain 1.2 dB Q 1
Filter: ON PK Fc 13500 Hz Gain 0.3 dB Q 1.2498
Filter: ON PK Fc 14250 Hz Gain 0.2 dB Q 2.6
Filter: ON PK Fc 15000 Hz Gain -1.4 dB Q 1
Filter: ON HSC Fc 20000 Hz Gain -4.2 dB Q 0.8
Filter: ON LSC Fc 90 Hz Gain 2.5 dB Q 0.8
# Filter: ON HSC Fc 9500 Hz Gain -3 dB Q 0.5
Filter: ON PK Fc 136 Hz Gain -10 dB Q 12
Filter: ON PK Fc 180 Hz Gain -15 dB Q 15
Filter: ON PK Fc 1096 Hz Gain -4 dB Q 0.5
Filter: ON PK Fc 1500 Hz Gain -4.2 dB Q 2.8
Filter: ON LS Fc 120 Hz Gain -6.4 dB
Filter: ON HSC Fc 11600 Hz Gain -0.6 dB Q 4"""

# plot_eq_response(config_eq)
# optimized_config = optimize_eq_config(config_eq)
# print(optimized_config)

optimized_config_v2 = optimize_eq_config_v2(config_eq, max_filters=18)
print(optimized_config_v2)