"""Plot planning and synchronization gains without hiding invalid programs."""
from __future__ import annotations

import argparse
import csv
import math
from pathlib import Path

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


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--input', type=Path, required=True)
    parser.add_argument('--output-dir', type=Path, required=True)
    args = parser.parse_args()
    with args.input.open(newline='') as stream:
        rows = list(csv.DictReader(stream))
    if not rows:
        raise ValueError('The ablation CSV is empty')
    panels = [('plan', 'latency', 'latency_gain_percent', 'Planning: latency reduction (%)'),
              ('plan', 'throughput', 'throughput_gain_percent', 'Planning: throughput increase (%)'),
              ('sync', 'throughput', 'throughput_gain_percent', 'Synchronization: throughput increase (%)')]
    fig, axes = plt.subplots(3, 1, figsize=(11, 9), layout='constrained')
    for ax, (comparison, mode, field, label) in zip(axes, panels):
        selected = sorted((r for r in rows if r['comparison'] == comparison and r['mode'] == mode),
                          key=lambda r: (r['arch'], r['model'], int(r['raptor_pipeline'])))
        for x, row in enumerate(selected):
            try:
                value = float(row[field].rstrip('%'))
            except (ValueError, TypeError):
                value = math.nan
            valid = row['functional_validation'] == row['baseline_functional_validation'] == 'PASS'
            if math.isfinite(value):
                ax.bar(x, value, color='#B3D5CA' if valid else '#FF9494',
                       edgecolor='#222222', hatch='' if valid else '//')
                ax.annotate(f'{value:.1f}' + ('' if valid else ' INVALID'),
                            (x, value), xytext=(0, 4 if value >= 0 else -4),
                            textcoords='offset points', ha='center',
                            va='bottom' if value >= 0 else 'top', fontsize=7)
            else:
                ax.text(x, 0, 'NA', ha='center', va='bottom', fontsize=8)
        ax.set_xticks(range(len(selected)),
                      [f"{r['model']}\n{r['arch']} p{r['raptor_pipeline']}" for r in selected], fontsize=8)
        ax.set_ylabel(label)
        ax.axhline(0, color='#222222', linewidth=.7)
        ax.set_xlim(-.6, max(.6, len(selected) - .4))
        ax.margins(y=.3)
    fig.suptitle('Raptor ablation (red hatching: functional validation failed)')
    args.output_dir.mkdir(parents=True, exist_ok=True)
    for extension in ('pdf', 'svg', 'png'):
        fig.savefig(args.output_dir / f'ablation.{extension}', dpi=200)
    plt.close(fig)


if __name__ == '__main__':
    main()