#!/usr/bin/env python3
"""Render an HFX structure (geometry JSON from HFX.geometryJSON) with matplotlib.
Usage as module: render(geom, ax, u=None, scale=1, color_by='level', state=None, util=None, N=None, fail_order=None, ...)
"""
import json, sys, math
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib import cm, colors

LEVEL_COLORS = ['#d62728', '#ff7f0e', '#bcbd22', '#2ca02c', '#17becf', '#1f77b4', '#9467bd', '#e377c2', '#8c564b', '#7f7f7f', '#171717']

def element_colors(geom, color_by, state=None, util=None, N=None, fail_order=None, n_failed=0, disp=None):
    els = geom['elements']
    out = []
    Nmax = max(1e-300, max(abs(x) for x in N)) if N is not None else 1
    dmax = max(1e-300, max(disp)) if disp is not None else 1
    tmin = geom['stats']['tMin']; tmax = geom['stats']['tMax']
    for i, e in enumerate(els):
        st = state[i] if state is not None else 0
        if st == 2 and color_by != 'failOrder':
            out.append((0.75, 0.75, 0.75, 1.0)); continue
        if color_by == 'level':
            out.append(colors.to_rgba(LEVEL_COLORS[min(e['level'], len(LEVEL_COLORS) - 1)]))
        elif color_by == 'interface':
            out.append(colors.to_rgba('#d62728' if e['iface'] else ('#ff7f0e' if e['kind'] == 'brace' else '#505050')))
        elif color_by == 'util':
            u = util[i] if util is not None else 0
            out.append(cm.viridis(min(1.0, max(0.0, u))))
        elif color_by == 'N':
            v = (N[i] / Nmax) if N is not None else 0
            out.append(cm.coolwarm(0.5 + 0.5 * max(-1, min(1, v))))
        elif color_by == 'disp':
            out.append(cm.magma(min(1.0, (disp[i] / dmax) if disp is not None else 0)))
        elif color_by == 'failOrder':
            if fail_order is None or fail_order[i] < 0: out.append((0.82, 0.82, 0.82, 1.0))
            else: out.append(cm.viridis(fail_order[i] / max(1, n_failed - 1)))
        elif color_by == 'thickness':
            out.append(cm.magma(math.log(e['t'] / tmin) / max(1e-9, math.log(tmax / tmin))))
        else:
            out.append((0.2, 0.2, 0.2, 1.0))
    return out

def render(geom, ax, u=None, scale=1.0, color_by='level', state=None, util=None, N=None, fail_order=None, n_failed=0,
           show_failed='ghost', line_scale=1.0, min_lw=0.3, title=None, show_cells=False, lw_units_per_inch=None, frame=True):
    nodes = np.array([[n[0], n[1]] for n in geom['nodes']], dtype=float)
    if u is not None:
        uu = np.asarray(u, dtype=float).reshape(-1, 3)
        nodes = nodes + scale * uu[:, :2]
    els = geom['elements']
    W, H = geom['W'], geom['H']
    # line width in points: thickness (normalised) * axis width in points
    fig = ax.figure
    bbox = ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
    pts_per_unit = bbox.width * 72.0 / (W * 1.08)
    disp = None
    if u is not None and color_by == 'disp':
        uu = np.asarray(u, dtype=float).reshape(-1, 3)
        disp = [0.5 * (np.hypot(*uu[e['n1'], :2]) + np.hypot(*uu[e['n2'], :2])) for e in els]
    cols = element_colors(geom, color_by, state, util, N, fail_order, n_failed, disp)
    order = sorted(range(len(els)), key=lambda i: els[i]['level'])
    segs, widths, cc, ghost_segs, ghost_w = [], [], [], [], []
    for i in order:
        e = els[i]
        if e.get('absent'): continue
        st = state[i] if state is not None else 0
        seg = [nodes[e['n1']], nodes[e['n2']]]
        lw = max(min_lw, e['t'] * pts_per_unit * line_scale)
        if st == 2:
            if show_failed == 'hide': continue
            if show_failed == 'ghost': ghost_segs.append(seg); ghost_w.append(lw * 0.6); continue
        segs.append(seg); widths.append(lw); cc.append(cols[i])
    if show_cells and geom.get('leafCells'):
        for c in geom['leafCells']:
            poly = np.array(c + [c[0]])
            ax.plot(poly[:, 0], poly[:, 1], color='#e6e6e6', lw=0.4, zorder=0)
    if ghost_segs:
        ax.add_collection(LineCollection(ghost_segs, linewidths=ghost_w, colors='#c8c8c8', linestyles=(0, (2, 2)), capstyle='round', zorder=1))
    ax.add_collection(LineCollection(segs, linewidths=widths, colors=cc, capstyle='round', zorder=2))
    ax.set_xlim(-0.04 * W, 1.04 * W); ax.set_ylim(-0.04 * H, 1.04 * H)
    ax.set_aspect('equal'); ax.set_xticks([]); ax.set_yticks([])
    if not frame:
        for s in ax.spines.values(): s.set_visible(False)
    if title: ax.set_title(title, fontsize=9)
    return ax

if __name__ == '__main__':
    geom = json.load(open(sys.argv[1]))
    out = sys.argv[2]
    color_by = sys.argv[3] if len(sys.argv) > 3 else 'level'
    fig, ax = plt.subplots(figsize=(5, 5 * geom['H'] / geom['W']), dpi=150)
    render(geom, ax, color_by=color_by, title=sys.argv[4] if len(sys.argv) > 4 else None)
    fig.tight_layout(); fig.savefig(out, dpi=150); print('wrote', out)