"""Plot actual Bambu G-code extrusion paths, separating model and supports.

python3 script/plot_sliced_stand.py plate_1.gcode output_directory
This is inspection only: no G-code is modified or sent to a printer.
"""
import math
import os
import re
import sys

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from mpl_toolkits.mplot3d.art3d import Line3DCollection
import numpy as np

source, output = sys.argv[1:]
os.makedirs(output, exist_ok=True)
position = {a: 0.0 for a in "XYZE"}
relative_e, relative_xyz = False, False
feature = "Custom"
segments = {"model": [], "support": []}
with open(source) as stream:
    for line in stream:
        if line.startswith("; FEATURE:"):
            feature = line.partition(":")[2].strip()
        code = line.partition(";")[0].strip()
        if not code:
            continue
        command = code.split()[0]
        values = {a: float(v) for a, v in re.findall(r"([XYZEIJ])(-?(?:\d+(?:\.\d*)?|\.\d+))", code)}
        if command == "M83":
            relative_e = True
        elif command == "M82":
            relative_e = False
        elif command == "G91":
            relative_xyz = True
        elif command == "G90":
            relative_xyz = False
        elif command == "G92":
            position.update({a: v for a, v in values.items() if a in position})
        elif command in ("G0", "G1", "G2", "G3"):
            old = position.copy()
            for a in "XYZE":
                if a in values:
                    relative = relative_e if a == "E" else relative_xyz
                    position[a] = values[a] + (old[a] if relative else 0.0)
            if position["E"] <= old["E"] or feature in ("Custom", "Flush", "Prime tower", "Skirt"):
                continue
            start, end = [old[a] for a in "XYZ"], [position[a] for a in "XYZ"]
            points = [start, end]
            if command in ("G2", "G3") and ("I" in values or "J" in values):
                cx, cy = start[0] + values.get("I", 0), start[1] + values.get("J", 0)
                radius = math.hypot(start[0]-cx, start[1]-cy)
                a = math.atan2(start[1]-cy, start[0]-cx)
                b = math.atan2(end[1]-cy, end[0]-cx)
                sweep = (b-a) % (2*math.pi) if command == "G3" else -((a-b) % (2*math.pi))
                count = max(2, math.ceil(abs(sweep)*radius / 0.3))
                points = [[cx+radius*math.cos(a+sweep*t/count), cy+radius*math.sin(a+sweep*t/count),
                           start[2]+(end[2]-start[2])*t/count] for t in range(count+1)]
            group = "support" if feature.startswith("Support") else "model"
            segments[group].extend(zip(points, points[1:]))

paths = {name: np.asarray(items).reshape(-1, 2, 3) for name, items in segments.items()}
all_points = np.concatenate(list(paths.values())).reshape(-1, 3)
low, high = all_points.min(axis=0), all_points.max(axis=0)
colors = {"model": "#586879", "support": "#de790e"}
for name, elevation, azimuth in [("iso", 25, -65), ("front", 0, -90), ("top", 90, -90)]:
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(projection="3d")
    for group, paths_for_group in paths.items():
        ax.add_collection3d(Line3DCollection(paths_for_group, colors=colors[group], linewidths=0.3, alpha=0.85))
    ax.set(xlim=(low[0], high[0]), ylim=(low[1], high[1]), zlim=(0, high[2]))
    ax.set_box_aspect(high-low)
    ax.view_init(elev=elevation, azim=azimuth)
    ax.set_title("Sliced stand — gray: model; orange: supports")
    ax.set_axis_off()
    fig.savefig(os.path.join(output, name+".png"), dpi=150, bbox_inches="tight")
    plt.close(fig)

fig, axes = plt.subplots(2, 3, figsize=(14, 9))
for ax, z in zip(axes.flat, [0.2, 5.0, 15.0, 24.0, 30.0, 35.0]):
    for group, items in paths.items():
        subset = items[np.abs(items[:, 1, 2]-z) < 0.025]
        ax.add_collection(LineCollection(subset[:, :, :2], colors=colors[group], linewidths=0.5))
    ax.set(xlim=(low[0], high[0]), ylim=(low[1], high[1]), title=f"Z = {z:.1f} mm", aspect="equal")
fig.suptitle("Actual sliced layers — gray: model; orange: supports")
fig.tight_layout()
fig.savefig(os.path.join(output, "layers.png"), dpi=150)
print(f"Plotted {sum(len(v) for v in segments.values())} extrusion segments in {output}")