#!/usr/bin/env python3
"""
Sallen-Key low-pass: plot magnitude + phase from actual R/C values,
and optionally do Monte-Carlo tolerance sweeps to see "wonkiness".

Model (standard Sallen-Key LPF, non-inverting gain K):
  H(s) = K / ( s^2*R1*R2*C1*C2 + s*(R1*C1 + (R1+R2)*C2 - K*R1*C2) + 1 )

Usage examples:
  python3 sallen_key_plot.py --R1 22000 --R2 22000 --C1 22e-9 --C2 10e-9
  python3 sallen_key_plot.py --R1 22k --R2 22k --C1 22n --C2 10n --K 1.0 --fc_guess 700
  python3 sallen_key_plot.py --R1 22k --R2 22k --C1 22n --C2 10n --mc 200 --tolR 0.01 --tolC 0.05

Notes:
- Needs numpy + matplotlib. On Ubuntu: sudo apt install python3-numpy python3-matplotlib
"""

import argparse
import math
import re
import numpy as np
import matplotlib.pyplot as plt


SI = {
    "p": 1e-12, "n": 1e-9, "u": 1e-6, "µ": 1e-6, "m": 1e-3,
    "k": 1e3, "K": 1e3, "M": 1e6, "G": 1e9,
}

def parse_si(s: str) -> float:
    """Parse numbers like '22k', '10n', '1.5M', '22e-9'."""
    s = s.strip()
    # plain float
    try:
        return float(s)
    except ValueError:
        pass
    m = re.fullmatch(r"([0-9]*\.?[0-9]+)\s*([pnumµmkKMG])", s)
    if not m:
        raise ValueError(f"Can't parse value '{s}'. Use e.g. 22000, 22k, 10n, 22e-9")
    val = float(m.group(1))
    suf = m.group(2)
    return val * SI[suf]

def H_sallen_key(jw: np.ndarray, R1: float, R2: float, C1: float, C2: float, K: float) -> np.ndarray:
    s = 1j * jw
    a2 = R1 * R2 * C1 * C2
    a1 = (R1 * C1) + ((R1 + R2) * C2) - (K * R1 * C2)
    den = (a2 * s * s) + (a1 * s) + 1.0
    return K / den

def estimate_fc(R1, R2, C1, C2) -> float:
    """Rough natural frequency estimate for guidance (not exact fc at -3 dB)."""
    w0 = 1.0 / math.sqrt(R1 * R2 * C1 * C2)
    return w0 / (2.0 * math.pi)

def main():
    ap = argparse.ArgumentParser(description="Sallen-Key LPF plotter from actual R/C values (+ Monte-Carlo tolerances).")
    ap.add_argument("--R1", required=True, type=str, help="Ohms (e.g. 22000 or 22k)")
    ap.add_argument("--R2", required=True, type=str, help="Ohms (e.g. 22000 or 22k)")
    ap.add_argument("--C1", required=True, type=str, help="Farads (e.g. 22e-9 or 22n)")
    ap.add_argument("--C2", required=True, type=str, help="Farads (e.g. 10e-9 or 10n)")
    ap.add_argument("--K", default="1.0", type=str, help="Non-inverting gain (default 1.0)")

    ap.add_argument("--fmin", type=float, default=None, help="Plot min freq Hz (default: fc_est/100)")
    ap.add_argument("--fmax", type=float, default=None, help="Plot max freq Hz (default: fc_est*100)")
    ap.add_argument("--points", type=int, default=2000, help="Plot points (default 2000)")

    ap.add_argument("--mc", type=int, default=0, help="Monte-Carlo runs (default 0)")
    ap.add_argument("--tolR", type=float, default=0.0, help="R tolerance as fraction (e.g. 0.01 for 1%%)")
    ap.add_argument("--tolC", type=float, default=0.0, help="C tolerance as fraction (e.g. 0.05 for 5%%)")
    ap.add_argument("--seed", type=int, default=None, help="Random seed for Monte-Carlo")

    ap.add_argument("--save", type=str, default=None, help="Save PNG to this path instead of showing")
    args = ap.parse_args()

    R1 = parse_si(args.R1)
    R2 = parse_si(args.R2)
    C1 = parse_si(args.C1)
    C2 = parse_si(args.C2)
    K  = float(parse_si(args.K))

    fc_est = estimate_fc(R1, R2, C1, C2)
    fmin = args.fmin if args.fmin is not None else max(fc_est / 100.0, 0.1)
    fmax = args.fmax if args.fmax is not None else fc_est * 100.0

    f = np.logspace(np.log10(fmin), np.log10(fmax), args.points)
    w = 2.0 * np.pi * f

    # Nominal response
    H0 = H_sallen_key(w, R1, R2, C1, C2, K)
    mag0 = 20.0 * np.log10(np.abs(H0))
    ph0  = np.unwrap(np.angle(H0)) * 180.0 / np.pi

    # Figure
    fig, (axm, axp) = plt.subplots(2, 1, figsize=(10, 7), sharex=True)

    # Monte-Carlo sweeps (optional)
    if args.mc and (args.tolR > 0 or args.tolC > 0):
        rng = np.random.default_rng(args.seed)
        for _ in range(args.mc):
            # uniform tolerance: value*(1 + u), u in [-tol, +tol]
            r1 = R1 * (1.0 + rng.uniform(-args.tolR, args.tolR))
            r2 = R2 * (1.0 + rng.uniform(-args.tolR, args.tolR))
            c1 = C1 * (1.0 + rng.uniform(-args.tolC, args.tolC))
            c2 = C2 * (1.0 + rng.uniform(-args.tolC, args.tolC))
            Hi = H_sallen_key(w, r1, r2, c1, c2, K)
            axm.semilogx(f, 20.0*np.log10(np.abs(Hi)), linewidth=0.8, alpha=0.15)
            axp.semilogx(f, np.unwrap(np.angle(Hi))*180.0/np.pi, linewidth=0.8, alpha=0.15)

    # Nominal on top
    axm.semilogx(f, mag0, linewidth=2.0)
    axp.semilogx(f, ph0,  linewidth=2.0)

    axm.set_ylabel("Magnitude (dB)")
    axp.set_ylabel("Phase (deg)")
    axp.set_xlabel("Frequency (Hz)")

    axm.grid(True, which="both", linestyle=":")
    axp.grid(True, which="both", linestyle=":")

    title = (
        f"Sallen-Key LPF | R1={R1:g}Ω R2={R2:g}Ω C1={C1:g}F C2={C2:g}F K={K:g}\n"
        f"w0 est: {2*math.pi*fc_est:.3g} rad/s,  f0 est: {fc_est:.3g} Hz"
    )
    if args.mc and (args.tolR > 0 or args.tolC > 0):
        title += f" | Monte-Carlo: {args.mc} runs (tolR={args.tolR:.3g}, tolC={args.tolC:.3g})"
    fig.suptitle(title)
    fig.tight_layout()

    if args.save:
        plt.savefig(args.save, dpi=200)
        print(f"Saved plot to {args.save}")
    else:
        plt.show()


if __name__ == "__main__":
    main()