"""Render retained simulation/reference data, with no synthetic values."""
import argparse
import json
import pathlib
import xml.etree.ElementTree as ET

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--run', required=True, type=pathlib.Path)
    args = parser.parse_args()
    run = args.run.resolve()
    data = json.loads((run / 'figure-data.json').read_text())
    cases = ET.parse(run / 'pytest.xml').findall('.//testcase')
    passed = sum(not any(case.find(tag) is not None for tag in ('failure', 'error', 'skipped')) for case in cases)
    violet, ink, muted, paper = '#6240ab', '#27232d', '#746d7a', '#fcfaf7'
    plt.rcParams.update({'font.family': 'STIXGeneral', 'font.size': 12,
                         'axes.spines.top': False, 'axes.spines.right': False,
                         'axes.labelcolor': ink, 'xtick.color': muted, 'ytick.color': muted,
                         'axes.edgecolor': '#c6bfca', 'axes.titleweight': 'normal',
                         'savefig.facecolor': paper})
    fig = plt.figure(figsize=(12.5, 9), dpi=160, facecolor=paper)
    gs = fig.add_gridspec(2, 2, left=.085, right=.96, bottom=.16, top=.735,
                         hspace=.52, wspace=.27, height_ratios=[1, 1.2])
    fig.text(.085, .93, 'felisyn  /  computational research', color=violet, fontsize=14)
    fig.text(.085, .868, 'The first cat-circuit regression', color=ink, fontsize=28)
    fig.text(.085, .815, f'{passed} of {len(cases)} upstream comparisons passed exact equality',
             color=ink, fontsize=16)
    ax = fig.add_subplot(gs[0, :], facecolor=paper)
    for label, color, width, style in [('reference', '#7b7980', 2.4, '-'), ('run', violet, 1.2, '--')]:
        d = data[label]
        ax.plot(d['time_ms'], d['voltage_mV'], color=color, lw=width, linestyle=style,
                label='Authors’ reference' if label == 'reference' else 'Our simulation')
    ax.set_title(f"Membrane voltage · {data['population']} · recorded neuron {data['run']['voltage_neuron_id']}", loc='left', pad=12)
    ax.set_xlabel('Time within stimulus (ms)')
    ax.set_ylabel('Voltage (mV)')
    ax.legend(frameon=False, loc='best', fontsize=11)
    all_ids = [row['neuron_id'] for label in ['reference', 'run'] for row in data[label]['spike_trains']]
    xmin = min(data[label]['time_ms'][0] for label in ['reference', 'run'])
    xmax = max(data[label]['time_ms'][-1] for label in ['reference', 'run'])
    for column, label in enumerate(['reference', 'run']):
        ax = fig.add_subplot(gs[1, column], facecolor=paper)
        rows = data[label]['spike_trains']
        times = [t for row in rows for t in row['times_ms']]
        neurons = [row['neuron_id'] for row in rows for _ in row['times_ms']]
        ax.scatter(times, neurons, s=4, color='#7b7980' if label == 'reference' else violet, linewidths=0)
        ax.set_xlim(xmin, xmax)
        if all_ids:
            ax.set_ylim(min(all_ids)-1, max(all_ids)+1)
        ax.set_xlabel('Time within stimulus (ms)')
        ax.set_ylabel('Recorded neuron ID')
        ax.set_title(('Authors’ reference' if label == 'reference' else 'Our simulation') + f' · {len(times)} spikes', loc='left', pad=12)
    fig.text(.085, .088, f"Example: {data['population']}, segment {data['run']['segment_identifier']}; {len(data['run']['spike_trains'])} recorded spike trains.", color=muted, fontsize=11)
    fig.text(.085, .062, 'Small LSV1M software test. These are simulated outputs; this is not independent biological validation.', color=muted, fontsize=11)
    fig.text(.085, .034, 'Source: Mozaik v0.4.0 reference data · model context: Antolík et al., PLOS Computational Biology (2024)', color=muted, fontsize=10)
    fig.savefig(run / 'regression-comparison.png', dpi=160)
    fig.savefig(run / 'regression-comparison.svg')
    plt.close(fig)
    print(run / 'regression-comparison.png')


if __name__ == '__main__':
    main()