import os
from pathlib import Path
import sys

os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

import schemdraw
import schemdraw.elements as elm
from schemdraw.segments import Segment, SegmentCircle, SegmentText

FONT_NAME = "STIXGeneral"
BG_COLOR = "white"
DEFAULT_OUTPUT = Path(__file__).resolve().parent.parent / "image" / "thevenin_equiv.svg"


def add_text(
    drawing,
    position,
    text,
    *,
    fontsize=18,
    rotation=0,
    align=("center", "center"),
    color="black",
):
    element = elm.Element().at(position).anchor("center").hold()
    element.anchors["center"] = (0, 0)
    element.segments.append(
        SegmentText(
            (0, 0),
            text,
            rotation=rotation,
            align=align,
            fontsize=fontsize,
            font=FONT_NAME,
            color=color,
        )
    )
    drawing.add(element)


def add_circle_symbol(
    drawing,
    center,
    *,
    radius=0.42,
    text=None,
    fontsize=16,
    horizontal_line=False,
):
    element = elm.Element().at(center).anchor("center").hold()
    element.anchors["center"] = (0, 0)
    element.segments.append(SegmentCircle((0, 0), radius))
    if horizontal_line:
        element.segments.append(Segment([(-radius, 0), (radius, 0)]))
    if text:
        element.segments.append(
            SegmentText((0, 0), text, fontsize=fontsize, font=FONT_NAME)
        )
    drawing.add(element)

def draw_circuit(save_path=DEFAULT_OUTPUT):
    save_path = Path(save_path)
    save_path.parent.mkdir(parents=True, exist_ok=True)

    with schemdraw.Drawing(show=False) as d:
        d.config(
            unit=1.0,
            fontsize=18,
            font=FONT_NAME,
            lw=2.0,
            margin=0.2,
            bgcolor=BG_COLOR,
        )

        # 总线节点
        top_y = 5.8
        bottom_y = 1.4
        left_bus_x = 3.0
        meter_x = 6.45
        ammeter_x = 8.05
        load_x = 10.2
        circle_r = 0.48

        # 左侧戴维南等效支路
        branch_x = 2.1
        source_top_y = 4.95
        source_bottom_y = 3.75
        resistor_top_y = 3.05
        resistor_bottom_y = 1.95

        d.add(elm.Line().at((left_bus_x, top_y)).to((branch_x, top_y)).hold())
        source_center_y = (source_top_y + source_bottom_y) / 2
        d.add(elm.Line().at((branch_x, top_y)).down(top_y - (source_center_y + circle_r)).hold())
        add_circle_symbol(d, (branch_x, source_center_y), radius=circle_r, horizontal_line=True)
        d.add(elm.Line().at((branch_x, source_center_y - circle_r)).down((source_center_y - circle_r) - resistor_top_y).hold())
        d.add(elm.ResistorIEC().at((branch_x, resistor_top_y)).to((branch_x, resistor_bottom_y)).hold())
        d.add(elm.Line().at((branch_x, resistor_bottom_y)).down(resistor_bottom_y - bottom_y).hold())
        d.add(elm.Line().at((branch_x, bottom_y)).to((left_bus_x, bottom_y)).hold())

        # 中间电压表和右侧负载支路
        d.add(elm.Line().at((left_bus_x, top_y)).right((ammeter_x - circle_r) - left_bus_x).hold())
        d.add(elm.Line().at((left_bus_x, bottom_y)).right(load_x - left_bus_x).hold())

        add_circle_symbol(d, (ammeter_x, top_y), radius=circle_r, text="A", fontsize=15)
        d.add(elm.Line().at((ammeter_x + circle_r, top_y)).right(load_x - (ammeter_x + circle_r)).hold())

        d.add(elm.Line().at((meter_x, top_y)).down(top_y - (3.6 + circle_r)).hold())
        add_circle_symbol(d, (meter_x, 3.6), radius=circle_r, text="V", fontsize=15)
        d.add(elm.Line().at((meter_x, 3.6 - circle_r)).down((3.6 - circle_r) - bottom_y).hold())

        d.add(elm.ResistorVarIEC().at((load_x, bottom_y + 0.7)).up(2.0).hold())

        # 右侧支路收口
        d.add(elm.Line().at((load_x, top_y)).down(top_y - (bottom_y + 2.7)).hold())
        d.add(elm.Line().at((load_x, bottom_y + 0.7)).down(0.7).hold())

        # 关键节点
        d.add(elm.Dot(radius=0.075).at((left_bus_x, top_y)).hold())
        d.add(elm.Dot(radius=0.075).at((left_bus_x, bottom_y)).hold())

        # 标注
        add_text(d, (3.2, 4.32), r"$U_{oc}$", fontsize=21)
        add_text(d, (3.0, 2.72), r"$R_0$", fontsize=21)

        # 负载标注
        add_text(d, (10.95, 3.02), r"$R_L$", fontsize=21)

        d.save(str(save_path), transparent=False)

    print(f"saved to: {save_path}")


if __name__ == "__main__":
    output_path = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_OUTPUT
    draw_circuit(output_path)