import os
from pathlib import Path
import sys

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

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import AutoMinorLocator

FONT_NAME = "LXGW WenKai"
BG_COLOR = "white"
GRID_MAJOR = "#111111"
GRID_MINOR = "#8a8a8a"
LINE_COLOR = "#111111"
POINT_COLOR = "#111111"

CURRENT_MA = np.array([11.30, 12.71, 13.82, 15.47, 16.66], dtype=float)
VOLTAGE_V = np.array([11.39, 10.69, 9.97, 9.21, 8.58], dtype=float)

DEFAULT_OUTPUT = (
    Path(__file__).resolve().parent.parent
    / "image"
    / "exp4"
    / "vi_characteristic_curve_exp4.svg"
)


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

    plt.rcParams["font.family"] = FONT_NAME
    plt.rcParams["axes.unicode_minus"] = False

    slope, intercept = np.polyfit(CURRENT_MA, VOLTAGE_V, 1)
    fit_x = np.linspace(CURRENT_MA.min() - 0.4, CURRENT_MA.max() + 0.4, 200)
    fit_y = slope * fit_x + intercept

    fig, ax = plt.subplots(figsize=(8.2, 6.0), dpi=160)
    fig.patch.set_facecolor(BG_COLOR)
    ax.set_facecolor(BG_COLOR)

    ax.plot(
        fit_x,
        fit_y,
        color=LINE_COLOR,
        linewidth=2.2,
        zorder=2,
    )
    ax.scatter(
        CURRENT_MA,
        VOLTAGE_V,
        s=60,
        color=POINT_COLOR,
        edgecolors="white",
        linewidths=0.8,
        zorder=3,
    )

    ax.set_xlabel("端口电流 I / mA", fontsize=13)
    ax.set_ylabel("端口电压 U / V", fontsize=13)

    ax.set_xlim(11.0, 17.1)
    ax.set_ylim(8.5, 11.5)
    ax.set_xticks(np.arange(11, 18, 1))
    ax.set_yticks(np.arange(8.5, 11.6, 0.5))
    ax.xaxis.set_minor_locator(AutoMinorLocator(2))
    ax.yaxis.set_minor_locator(AutoMinorLocator(2))
    ax.set_box_aspect(6 / 5)

    ax.grid(True, which="major", color=GRID_MAJOR, linewidth=0.95, alpha=0.9)
    ax.grid(True, which="minor", color=GRID_MINOR, linewidth=0.55, alpha=0.8)

    for spine in ax.spines.values():
        spine.set_linewidth(1.0)
        spine.set_color("#111111")

    fig.tight_layout()
    fig.savefig(save_path, bbox_inches="tight")
    plt.close(fig)
    print(f"saved to: {save_path}")


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