#!/usr/bin/env python3
"""Render a structure JSON (nodes, members with widths and levels) to PNG with true member widths."""
import json, sys
import numpy as np
import matplotlib; matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.collections import PolyCollection

def member_polys(nodes, members, widths=None):
    polys = []
    for k, m in enumerate(members):
        a = nodes[m['n1']]; b = nodes[m['n2']]
        w = widths[k] if widths is not None else m['w']
        dx, dy = b['x']-a['x'], b['y']-a['y']; L = np.hypot(dx, dy) or 1e-12
        nx, ny = -dy/L*w/2, dx/L*w/2
        polys.append([(a['x']+nx, a['y']+ny), (b['x']+nx, b['y']+ny), (b['x']-nx, b['y']-ny), (a['x']-nx, a['y']-ny)])
    return polys

def render(struct, out, title=None, color_by='level'):
    nodes, members = struct['nodes'], struct['members']
    W, H = struct['W'], struct['H']
    fig, ax = plt.subplots(figsize=(6, 6*H/W))
    polys = member_polys(nodes, members)
    lv = np.array([m['level'] for m in members])
    cmap = plt.get_cmap('viridis')
    D = max(1, lv.max())
    cols = cmap((lv-1)/max(1, D-1)) if D > 1 else np.tile([[0.2,0.4,0.7,1]], (len(members),1))
    ax.add_collection(PolyCollection(polys, facecolors=cols, edgecolors='none'))
    ax.set_xlim(-0.02*W, 1.02*W); ax.set_ylim(-0.02*H, 1.02*H); ax.set_aspect('equal')
    ax.set_xlabel('x [mm]'); ax.set_ylabel('y [mm]')
    if title: ax.set_title(title, fontsize=9)
    fig.tight_layout(); fig.savefig(out, dpi=150); plt.close(fig)

if __name__ == '__main__':
    src, out = sys.argv[1], sys.argv[2]
    with open(src) as f: s = json.load(f)
    render(s, out, title=sys.argv[3] if len(sys.argv) > 3 else None)