"""Generate the PSU schematic (SVG + PNG) with schemdraw.

    python -m pip install schemdraw matplotlib
    python hardware/psu_schematic.py
"""
from pathlib import Path

import schemdraw
import schemdraw.elements as elm

OUT = Path(__file__).parent
schemdraw.use("matplotlib")


def gnd(d, at):
    """Ground symbol pointing down (for parts hanging below a node)."""
    d.add(elm.Ground().theta(0).at(at))


def gnd_up(d, at):
    """Ground symbol pointing up (for parts standing on a negative rail)."""
    d.add(elm.Ground().theta(180).at(at))


def vdd(d, at, name):
    d.add(elm.Vdd().theta(0).at(at).label(name))


def ic(left, right, label, bottom=(), top=()):
    pins = [elm.IcPin(name=n, side="left") for n in left]
    pins += [elm.IcPin(name=n, side="right") for n in right]
    pins += [elm.IcPin(name=n, side="bottom") for n in bottom]
    pins += [elm.IcPin(name=n, side="top") for n in top]
    return elm.Ic(pins=pins, pinspacing=1.2, edgepadH=0.6, edgepadW=1.2).right().label(label, "top")


def title(d, at, text):
    d.add(elm.Label().at(at).label(text, fontsize=12, halign="left", valign="bottom"))


def note(d, at, text, size=9, color="black"):
    d.add(elm.Label().at(at).label(text, fontsize=size, halign="left", valign="top", color=color))


# For vertical parts schemdraw's "left"/"right" label locations are the part's ends;
# these give the visual side instead. A .reverse()d part swaps its sides again.
R_DOWN, L_DOWN, R_UP, L_UP = "bottom", "top", "bottom", "top"


def shunt_down(d, at, part, label):
    """part is a factory: elements auto-add when created inside a Drawing."""
    d.add(elm.Dot().at(at))
    d.add(part().at(at).down().label(label, R_DOWN))
    gnd(d, d.here)


def shunt_up(d, at, part, label):
    d.add(elm.Dot().at(at))
    d.add(part().at(at).up().label(label, R_UP))
    gnd_up(d, d.here)


# ---------------------------------------------------------------------------
def block_input(d, ox, oy):
    title(d, (ox - 3, oy + 3.5), "1. 12 V input and +5 V")
    d.add(elm.Dot(open=True).at((ox, oy)).label("J1 +12V in\n(certified brick)", "left"))
    d.add(elm.Fuse().right().label("F1 3A T"))
    n1 = d.here
    d.add(elm.Line().right().length(3))
    n2 = d.here
    d.add(elm.Line().right().length(3))
    n3 = d.here
    d.add(elm.Line().right().length(2))
    d.add(elm.Dot())
    d.add(elm.DiodeTVS().at(n1).down().reverse().label("D1\nP6KE15A", L_DOWN))
    gnd(d, d.here)
    d.add(elm.Dot().at(n1))
    shunt_down(d, n2, lambda: elm.Capacitor(polar=True), "C1\n470µF 25V")
    d.add(elm.Dot().at(n3))
    d.add(elm.Line().at(n3).up().length(1.2))
    vdd(d, d.here, "+12V")
    d.add(elm.Dot(open=True).at((ox, oy - 2.5)).label("J1 GND", "left"))
    gnd(d, (ox, oy - 2.5))

    u1 = d.add(ic(["VIN"], ["VOUT"], "U1 Mini560\n5V fixed buck", bottom=["GND"])
               .at((n3.x + 2, oy)).anchor("VIN"))
    gnd(d, u1.GND)
    d.add(elm.Line().at(u1.VOUT).right().length(2))
    n4 = d.here
    d.add(elm.Line().right().length(2))
    vdd(d, d.here, "+5V")
    shunt_down(d, n4, lambda: elm.Capacitor(polar=True), "C2\n100µF 10V")


def block_enable(d, ox, oy):
    title(d, (ox - 3, oy + 3.5), "2. HV enable relay + hardware watchdog")
    # charge pump
    d.add(elm.Dot(open=True).at((ox, oy)).label("ESP32 GPIO21\n(toggles every\nscanned row)", "left"))
    d.add(elm.Capacitor().right().label("C3 100nF"))
    a = d.here
    d.add(elm.Diode().right().label("D2 1N4148"))
    b = d.here
    d.add(elm.Line().right().length(2.5))
    c = d.here
    d.add(elm.Line().right().length(2.5))
    q1 = d.add(elm.NFet().theta(0).reverse().anchor("gate").label("Q1\n2N7000", "right"))
    d.add(elm.Dot().at(a))
    d.add(elm.Diode().at(a).down().reverse().label("D3\n1N4148", L_DOWN))
    gnd(d, d.here)
    shunt_down(d, b, lambda: elm.Capacitor(polar=True), "C4\n10µF")
    shunt_down(d, c, lambda: elm.Resistor(), "R1\n100k")
    gnd(d, q1.source)

    # relay module
    k1 = d.add(ic(["IN", "VCC"], ["NO", "COM"], "K1 relay module\n5V, opto, active-low IN",
                  bottom=["GND"]).at((q1.drain.x + 5, q1.drain.y + 2.5)).anchor("IN"))
    d.add(elm.Wire("|-").at(q1.drain).to(k1.IN))
    d.add(elm.Line().at(k1.VCC).left().length(1.5))
    vdd(d, d.here, "+5V")
    gnd(d, k1.GND)
    d.add(elm.Line().at(k1.COM).right().length(2))
    vdd(d, d.here, "+12V")
    d.add(elm.Line().at(k1.NO).right().length(3).label("+12V_HV", "right"))

    note(d, (ox - 6, oy - 4.5),
         "Static GPIO (high, low, reset, crash) → no charge on C4 → Q1 off → relay open → HV off.\n"
         "Drop-out ≈ 1 s (C4 × R1). Check your relay module's IN polarity.")


def block_pos(d, ox, oy):
    title(d, (ox - 3, oy + 4.5), "3. Positive HV: boost module + original regulator → +65 V")
    d.add(elm.Dot(open=True).at((ox, oy)).label("+12V_HV", "left"))
    d.add(elm.Fuse().right().label("F2 1A T"))
    d.add(elm.Line().right().length(1))
    u2 = d.add(ic(["VIN"], ["HV+"], "U2 NCH6100HV\nset 90V", bottom=["GND"]).anchor("VIN"))
    gnd(d, u2.GND)
    y0 = oy
    d.add(elm.Line().at(u2.absanchors["HV+"]).right().length(1))
    d.add(elm.Fuse().right().label("F4 100mA F"))
    x = d.here.x
    d.add(elm.Label().at((x + 0.3, y0 + 0.2)).label("+90V raw", "top", fontsize=9, halign="left"))

    # raw rail: bulk cap, bleeder, HV-present LED
    shunt_down(d, (x + 1.5, y0), lambda: elm.Capacitor(polar=True), "C5\n22µF\n250V")
    shunt_down(d, (x + 4.5, y0), lambda: elm.Resistor(), "R2\n330k")
    d.add(elm.Dot().at((x + 7.5, y0)))
    d.add(elm.Resistor().at((x + 7.5, y0)).down().label("R3\n100k 1W", R_DOWN))
    d.add(elm.LED().down().label("LED1", R_DOWN))
    gnd(d, d.here)

    # regulator (original Q1/Q3 section)
    xa, xb, xd = x + 12, x + 16, x + 21
    d.add(elm.Resistor().at((xb, y0)).down().label("R4\n1.5k", L_DOWN))
    qa = d.add(elm.BjtPnp(circle=True).theta(0).anchor("emitter").label("Q2\nMPSA92", "right"))
    by = qa.base.y
    d.add(elm.Dot().at((xa, y0)))
    d.add(elm.Zener().endpoints((xa, y0), (xa, by)).reverse().label("D4\n3.9V\n1N5228", R_DOWN))
    d.add(elm.Dot().at((xa, by)))
    d.add(elm.Line().at((xa, by)).to(qa.base))
    d.add(elm.Resistor().at((xa, by)).down().label("R5\n47k 1W", L_DOWN))
    gnd(d, d.here)

    qby = qa.collector.y - 1.2
    d.add(elm.Line().at(qa.collector).toy(qby))
    nqb = d.add(elm.Dot()).center
    d.add(elm.Zener().at(nqb).down().reverse().label("D5\n68V\n1N4760A", R_DOWN))
    gnd(d, d.here)
    q3 = d.add(elm.BjtNpn(circle=True).theta(0).anchor("base").at((xd, qby))
               .label("Q3\nMJE15030\n(heatsink)", "right"))
    d.add(elm.Line().at(nqb).to(q3.base))
    d.add(elm.Line().at(q3.collector).toy(y0))
    d.add(elm.Line().at((x, y0)).tox(q3.collector.x))

    oy2 = q3.emitter.y - 0.8
    d.add(elm.Line().at(q3.emitter).toy(oy2))
    d.add(elm.Line().right().length(2.5))
    o1 = d.here
    d.add(elm.Line().right().length(3))
    o2 = d.here
    d.add(elm.Line().right().length(3))
    d.add(elm.Dot(open=True).label("+65V to panel", "right"))
    shunt_down(d, o1, lambda: elm.Capacitor(), "C6\n0.1µF\n500V")
    shunt_down(d, o2, lambda: elm.Resistor(), "R6\n330k")


def block_neg(d, ox, oy):
    title(d, (ox - 3, oy + 4.5),
          "4. Negative HV: isolated converter + floating boost module + original regulator → −112 V / −100 V")
    d.add(elm.Dot(open=True).at((ox, oy)).label("+12V_HV", "left"))
    d.add(elm.Fuse().right().label("F3 1A T"))
    d.add(elm.Line().right().length(1))
    u3 = d.add(ic(["VIN", "GND"], ["−VO", "+VO"], "U3 isolated DC-DC\n12V→12V 6–10W").anchor("VIN"))
    d.add(elm.Line().at(u3.GND).left().length(1))
    gnd(d, d.here)
    d.add(elm.Line().at(u3.absanchors["+VO"]).right().length(2.5))
    u4 = d.add(ic(["VIN"], ["HV+"], "U4 NCH6100HV\nset 130V", bottom=["GND"]).anchor("VIN"))
    jy = u4.GND.y - 1.5
    d.add(elm.Line().at(u3.absanchors["−VO"]).right().length(1.2))
    d.add(elm.Line().toy(jy))
    d.add(elm.Line().tox(u4.GND.x))
    d.add(elm.Dot())
    hv = u4.absanchors["HV+"]
    d.add(elm.Line().at(hv).right().length(1.5))
    gnd(d, d.here)

    yb = u4.GND.y - 6                 # −130 V raw rail
    d.add(elm.Line().at(u4.GND).toy(yb))
    d.add(elm.Fuse().right().label("F5 100mA F", "bottom"))
    x = d.here.x
    d.add(elm.Label().at((x + 0.3, yb - 0.3)).label("−130V raw", "bottom", fontsize=9, halign="left"))

    d.add(elm.EncircleBox([u4], padx=1.0, pady=1.2).linestyle("--").color("red"))
    note(d, (u3.center.x - 3, yb - 0.8),
         "⚠ U3's output side, U4 and this whole\nnet sit at −130 V: insulate, never touch.", 8, "red")

    shunt_up(d, (x + 1.5, yb), lambda: elm.Capacitor(polar=True).reverse(), "C7\n22µF\n250V")
    shunt_up(d, (x + 4.5, yb), lambda: elm.Resistor(), "R7\n330k")
    d.add(elm.Dot().at((x + 7.5, yb)))
    d.add(elm.Resistor().at((x + 7.5, yb)).up().label("R8\n100k 1W", R_UP))
    d.add(elm.LED().up().reverse().label("LED2", L_UP))
    gnd_up(d, d.here)

    # regulator (original Q2/Q4/Q5 section)
    xa, xb, xz, xd = x + 12, x + 16, x + 19, x + 23
    d.add(elm.Resistor().at((xb, yb)).up().label("R9\n1.5k", L_UP))
    qa = d.add(elm.BjtNpn(circle=True).theta(0).anchor("emitter").label("Q4\nMPSA42", "right"))
    by = qa.base.y
    d.add(elm.Dot().at((xa, yb)))
    d.add(elm.Zener().endpoints((xa, yb), (xa, by)).label("D6\n3.9V\n1N5228", L_UP))
    d.add(elm.Dot().at((xa, by)))
    d.add(elm.Line().at((xa, by)).to(qa.base))
    d.add(elm.Resistor().at((xa, by)).up().label("R10\n47k 1W", L_UP))
    gnd_up(d, d.here)

    qby = qa.collector.y + 0.6
    d.add(elm.Line().at(qa.collector).toy(qby))
    d.add(elm.Line().tox(xz))
    nqb = d.add(elm.Dot()).center
    d.add(elm.Zener().at(nqb).up().label("D7\n13V\n1N4743", R_UP))
    m = d.add(elm.Dot()).center
    d.add(elm.Zener().up().label("D8\n100V\n1N4764A", R_UP))
    gnd_up(d, d.here)

    q4 = d.add(elm.BjtPnp(circle=True).theta(0).anchor("base").at((xd, qby))
               .label("Q5\nMJE15031\n(heatsink)", "right"))
    d.add(elm.Line().at(nqb).to(q4.base))
    d.add(elm.Line().at(q4.collector).toy(yb))
    d.add(elm.Line().at((x, yb)).tox(q4.collector.x))

    # −112 V output
    out_x = xd + 17
    y112 = q4.emitter.y + 0.6
    d.add(elm.Line().at(q4.emitter).toy(y112))
    d.add(elm.Line().right().length(6))
    o1 = d.here
    d.add(elm.Line().right().length(3))
    o2 = d.here
    d.add(elm.Line().tox(out_x))
    d.add(elm.Dot(open=True).label("−112V to panel", "right"))
    shunt_up(d, o1, lambda: elm.Capacitor(), "C8\n0.1µF\n500V")
    shunt_up(d, o2, lambda: elm.Resistor(), "R11\n330k")

    # −100 V: Q5 follower driven from node M (D6/D7 junction)
    yq5 = m.y + 5.5
    d.add(elm.Line().at(m).tox(m.x + 1.5))
    d.add(elm.Line().toy(yq5))
    d.add(elm.Resistor().right().tox(xd + 4.5).label("R12 1.5k", "top"))
    q5 = d.add(elm.BjtNpn(circle=True).theta(0).anchor("base").label("Q6\nMJE15030\n(heatsink)", "right"))
    d.add(elm.Resistor().at(q5.collector).up().label("R13\n2k 5W", R_UP))
    gnd_up(d, d.here)
    y100 = q5.emitter.y - 0.6
    d.add(elm.Line().at(q5.emitter).toy(y100))
    d.add(elm.Line().right().length(6))
    o3 = d.here
    d.add(elm.Line().tox(out_x))
    d.add(elm.Dot(open=True).label("−100V to panel", "right"))
    shunt_up(d, o3, lambda: elm.Capacitor(), "C9\n0.1µF\n500V")


def block_out(d, ox, oy):
    title(d, (ox - 4, oy + 3.5), "5. Panel power connector")
    names = ["+5V", "+12V", "+65V", "GND", "−100V", "−112V"]
    j2 = d.add(ic(list(reversed(names)), [], "J2 to panel\n(keyed, ≥300V rated)").at((ox, oy)).anchor("+5V"))
    for n in names:
        if n == "+12V":
            d.add(elm.Switch().at(j2.absanchors[n]).left().length(3).label("JP1", "top"))
            d.add(elm.Label().at(d.here).label(n, "left"))
        else:
            d.add(elm.Line().at(j2.absanchors[n]).left().length(3).label(n, "left"))
    note(d, (ox - 4, j2.absanchors["−112V"].y - 1.5),
         "Pin order is a placeholder: map it to YOUR panel's\nconnector before wiring. Fit JP1 only if the\npanel uses +12V.")


def block_notes(d, ox, oy):
    title(d, (ox, oy), "Notes")
    note(d, (ox, oy - 0.4),
         "• Panel voltages are set by zeners D5 (68V → about +65V out) and D7+D8 (113V → about −112V / −100V out), copied from Stern 520-5138-00.\n"
         "   Module trim pots only set the raw rails, so a mis-set module can't over-volt the panel.\n"
         "• Original board refs → this design: Q1→Q2, Q2→Q4, Q3→Q3, Q4→Q5, Q5→Q6, D3→D4, D4→D6, D5→D5, D6→D8, D7→D7, R2→R5, R3→R10, R4→R4, R5→R9,\n"
         "   R6→R6, R7→R11, R8→R13, R10→R12, C3→C6, C4→C8, C7→C9. Designators match the KiCad project.\n"
         "• Set U2/U4 with no load, then lock the trim pots. First power-up with the panel: 1 kΩ in series with each HV output, measure current.\n"
         "• Q3/Q4/Q5 on small heatsinks. HV transistors and zeners from a real distributor. Raw caps ≥ 250 V, output caps 500 V, ≥ 3 mm HV clearance.\n"
         "• Bleeders discharge all rails at power-off; LEDs are a hint only. Always measure before touching.",
         size=9)


with schemdraw.Drawing(show=False) as d:
    d.config(fontsize=9, unit=2.5)
    title(d, (-8, 8), "Plasma DMD — 12 V DC input power supply  (rev 0.1, unverified)")
    block_input(d, 0, 0)
    block_enable(d, 42, 0)
    block_pos(d, 0, -17)
    block_out(d, 58, -14)
    block_neg(d, 0, -36)
    block_notes(d, -8, -56)

d.save(str(OUT / "psu-schematic.svg"))
d.save(str(OUT / "psu-schematic.png"), dpi=150)
print("saved")