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" / "exp4" / "norton_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_outline(drawing, center, *, radius=0.48, vertical_line=False):
    element = elm.Element().at(center).anchor("center").hold()
    element.anchors["center"] = (0, 0)
    element.segments.append(SegmentCircle((0, 0), radius))
    if vertical_line:
        element.segments.append(Segment([(0, -radius), (0, radius)]))
    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

        source_x = 2.0
        resistor_x = 3.45
        node_x = 4.35
        meter_x = 6.15
        ammeter_x = 7.95
        load_x = 10.1
        circle_r = 0.48

        source_center_y = 3.6
        resistor_top_y = 4.9
        resistor_bottom_y = 2.3

        # 诺顿等效支路：电流源与电阻并联
        d.add(elm.Line().at((source_x, top_y)).to((node_x, top_y)).hold())
        d.add(elm.Line().at((source_x, bottom_y)).to((node_x, bottom_y)).hold())

        d.add(
            elm.Line()
            .at((source_x, top_y))
            .down(top_y - (source_center_y + circle_r))
            .hold()
        )
        add_circle_outline(
            d,
            (source_x, source_center_y),
            radius=circle_r,
            vertical_line=True,
        )
        d.add(
            elm.Line()
            .at((source_x, source_center_y - circle_r))
            .down((source_center_y - circle_r) - bottom_y)
            .hold()
        )

        d.add(elm.ResistorIEC().at((resistor_x, resistor_top_y)).to((resistor_x, resistor_bottom_y)).hold())
        d.add(elm.Line().at((resistor_x, top_y)).down(top_y - resistor_top_y).hold())
        d.add(elm.Line().at((resistor_x, resistor_bottom_y)).down(resistor_bottom_y - bottom_y).hold())

        # 右侧测量与负载
        d.add(elm.Line().at((node_x, top_y)).right((ammeter_x - circle_r) - node_x).hold())
        d.add(elm.Line().at((node_x, bottom_y)).right(load_x - node_x).hold())

        element_a = elm.Element().at((ammeter_x, top_y)).anchor("center").hold()
        element_a.anchors["center"] = (0, 0)
        element_a.segments.append(SegmentCircle((0, 0), circle_r))
        element_a.segments.append(SegmentText((0, 0), "A", fontsize=15, font=FONT_NAME))
        d.add(element_a)
        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())
        element_v = elm.Element().at((meter_x, 3.6)).anchor("center").hold()
        element_v.anchors["center"] = (0, 0)
        element_v.segments.append(SegmentCircle((0, 0), circle_r))
        element_v.segments.append(SegmentText((0, 0), "V", fontsize=15, font=FONT_NAME))
        d.add(element_v)
        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((node_x, top_y)).hold())
        d.add(elm.Dot(radius=0.075).at((node_x, bottom_y)).hold())

        # 标注
        add_text(d, (1.2, 3.42), r"$I_{sc}$", fontsize=21, align=("right", "center"))
        add_text(d, (4.15, 3.55), r"$R_0$", fontsize=21)
        add_text(d, (10.85, 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)