#!/usr/bin/env python3
"""Shared rendering helpers: member polygons, frames from *_frames.bin, load-displacement curves."""
import json, os, numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import PolyCollection, LineCollection
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
RUNS = os.path.join(ROOT, 'data', 'runs')
LEVEL_COLORS = ['#1f3a93', '#2a7ab0', '#2aa198', '#5fb760', '#b0c23a', '#e0a020', '#e06020', '#c02040', '#8a2080', '#503090', '#304090', '#203070']

def load_run(id_):
    with open(os.path.join(RUNS, id_ + '.json')) as f: return json.load(f)

def load_frames(run):
    if not run.get('framesFile'): return None
    idx = json.load(open(os.path.join(RUNS, run['id'] + '_frames.json')))
    N, E, n = idx['N'], idx['E'], idx['nFrames']; per = 2 * N + 2 * E
    buf = np.fromfile(os.path.join(RUNS, run['framesFile']), dtype='<f4').reshape(n, per)
    return {'index': idx['frames'], 'ux': buf[:, :N], 'uy': buf[:, N:2*N], 'r': buf[:, 2*N:2*N+E], 'Nax': buf[:, 2*N+E:], 'N': N, 'E': E, 'n': n}

def geometry_arrays(run):
    g = run['geometry']; nodes = np.array([[n[0], n[1]] for n in g['nodes']]); mem = g['members']
    n1 = np.array([m[0] for m in mem]); n2 = np.array([m[1] for m in mem]); lev = np.array([m[2] for m in mem]); w = np.array([m[3] for m in mem])
    return nodes, n1, n2, lev, w

def member_polys(P1, P2, w):
    d = P2 - P1; L = np.linalg.norm(d, axis=1); L[L == 0] = 1e-12; nx = -d[:, 1] / L * w / 2; ny = d[:, 0] / L * w / 2
    return np.stack([np.stack([P1[:, 0] + nx, P1[:, 1] + ny], 1), np.stack([P2[:, 0] + nx, P2[:, 1] + ny], 1), np.stack([P2[:, 0] - nx, P2[:, 1] - ny], 1), np.stack([P1[:, 0] - nx, P1[:, 1] - ny], 1)], 1)

def draw_structure(ax, run, frame=None, frames=None, color='level', deform=1.0, show_failed=True, cmap='viridis', rmax=1.0, lw_failed=0.5, title=None, grips=True):
    nodes, n1, n2, lev, w = geometry_arrays(run); W, H = run['W'], run['H']
    P = nodes.copy()
    failed = np.zeros(len(n1), bool)
    if frame is not None and frames is not None:
        P = P + deform * np.stack([frames['ux'][frame], frames['uy'][frame]], 1)
        nE = frames['index'][frame]['nEvents']; fm = [e['member'] for e in run['events'][:nE]]; failed[fm] = True
    polys = member_polys(P[n1], P[n2], w)
    if color == 'level': cols = np.array([plt.matplotlib.colors.to_rgba(LEVEL_COLORS[(l - 1) % len(LEVEL_COLORS)]) for l in lev])
    elif color == 'r' and frame is not None: cm = plt.get_cmap(cmap); cols = cm(np.clip(frames['r'][frame] / rmax, 0, 1))
    elif color == 'width': cm = plt.get_cmap(cmap); cols = cm((w - w.min()) / max(1e-9, w.max() - w.min()))
    else: cols = np.tile(np.array([[0.35, 0.4, 0.5, 1.0]]), (len(n1), 1))
    ax.add_collection(PolyCollection(polys[~failed], facecolors=cols[~failed], edgecolors='none'))
    if show_failed and failed.any():
        segs = np.stack([P[n1[failed]], P[n2[failed]]], 1); ax.add_collection(LineCollection(segs, colors='#e5484d', linewidths=lw_failed, alpha=0.9))
    if grips:
        gh = 0.05 * H; dtop = deform * (frames['index'][frame]['d'] if (frame is not None and frames is not None) else 0)
        ax.add_patch(plt.Rectangle((-0.02 * W, -gh), 1.04 * W, gh, color='#666'))
        ax.add_patch(plt.Rectangle((-0.02 * W, H + dtop), 1.04 * W, gh, color='#666'))
    ax.set_xlim(-0.04 * W, 1.04 * W); ax.set_ylim(-0.08 * H, 1.08 * H + (deform * run['metrics']['dFinal'] if frames is not None and frame is not None else 0)); ax.set_aspect('equal')
    ax.set_xticks([]); ax.set_yticks([]); [s.set_visible(False) for s in ax.spines.values()]
    if title: ax.set_title(title, fontsize=8)

def curve_xy(run):
    c = run['curve']; return np.array([p['d'] for p in c]), np.array([p['F'] for p in c])

def frame_at_displacement(frames, d):
    ds = np.array([f['d'] for f in frames['index']]); return int(np.argmin(np.abs(ds - d)))