#!/usr/bin/env python3
"""Plot observed nonrotating surface asphericity at saved radial nodes."""
import argparse
import csv
import os
from pathlib import Path
for name in ('OPENBLAS_NUM_THREADS', 'OMP_NUM_THREADS', 'VECLIB_MAXIMUM_THREADS'):
    os.environ[name] = '1'
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('directories', type=Path, nargs=2)
    parser.add_argument('--output', type=Path, required=True)
    args = parser.parse_args()
    if args.output.exists():
        parser.error('Use a fresh output file.')
    data = []
    for directory, expected_order in zip(args.directories, (3, 4)):
        metadata = dict(line.split("=", 1) for line in (directory / "metadata.txt").read_text().splitlines() if "=" in line)
        with (directory / "weak_shape_physical_metadata.csv").open() as stream:
            physical_metadata = {r["key"]: r["value"] for r in csv.DictReader(stream)}
        if float(metadata["n"]) != 1.5 or float(metadata["J"]) != 0 or int(physical_metadata["displacement_family_order"]) != expected_order:
            raise ValueError("This comparison requires n=1.5, J=0, and Q3 followed by Q4 displacement.")
        with (directory / 'weak_shape_physical_surface_nodes.csv').open() as stream:
            rows = list(csv.DictReader(stream))
        with (directory / 'newton.csv').open() as stream:
            accepted = sum(int(row['accepted']) for row in csv.DictReader(stream))
        points = np.array([[float(row['unit_radial_'+a]) for a in 'xyz'] for row in rows])
        radii = np.array([float(row['physical_radius']) for row in rows])
        if not np.isfinite(points).all() or not np.isfinite(radii).all() or np.any(radii <= 0):
            raise ValueError('Invalid physical surface samples.')
        latitude = np.arcsin(np.clip(points[:, 2], -1, 1))
        longitude = np.arctan2(points[:, 1], points[:, 0])
        deviation = 100 * (radii / np.mean(radii) - 1)
        data.append((longitude, latitude, deviation, len(rows), accepted))
    limit = max(float(np.max(np.abs(row[2]))) for row in data)
    if limit == 0:
        limit = 1e-12
    fig, axes = plt.subplots(1, 2, figsize=(12, 5), subplot_kw={'projection': 'mollweide'}, layout='constrained')
    fig.get_layout_engine().set(rect=(0, .07, 1, .93))
    for axis, (longitude, latitude, deviation, count, accepted), title in zip(axes, data, ('Q3 displacement / surface', 'Q4 displacement / surface')):
        picture = axis.scatter(longitude, latitude, c=deviation, s=14, cmap='RdBu_r', vmin=-limit, vmax=limit, linewidths=0)
        axis.set_title(f'{title}\n{accepted} accepted updates; {count} radial nodes')
        axis.grid(True, alpha=.25)
        axis.tick_params(labelsize=8)
    fig.colorbar(picture, ax=list(axes), orientation='horizontal', shrink=.65, pad=.08,
                 label='Radius relative to its unweighted nodal mean (%)')
    fig.suptitle('Nonrotating n = 1.5: surface samples after the bounded order trial', fontsize=14)
    fig.text(.5, .01, 'Common color scale; different nodal sampling. Neither endpoint is certified by this plot.', ha='center', fontsize=9)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(args.output, dpi=200)
    plt.close(fig)


if __name__ == '__main__':
    main()