#!/usr/bin/env python3
"""
Capture a regression baseline for display_netcdf.

Sweeps every variable of every file in sample_files/ and records, per
(file, variable): the exit status, the plot classification, the dimensions left
after selection, the figure's title and axis labels, and a hash of the arrays
actually handed to matplotlib.

The hash is taken on the plotted data, not on the PNG, so the baseline survives
matplotlib version changes; only a change in what we plot moves it.

Usage:
    python tests/baseline_capture.py --out tests/baseline/plot_matrix.json
    python tests/baseline_capture.py --check tests/baseline/plot_matrix.json
"""

import os

os.environ.setdefault('MPLBACKEND', 'Agg')

import io as _io
import sys
import json
import glob
import hashlib
import argparse
import contextlib
import importlib
import tempfile
import traceback

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from netCDF4 import Dataset

_HERE = os.path.dirname(os.path.abspath(__file__))
_ROOT = os.path.dirname(_HERE)
EXAMPLES = os.path.join(_ROOT, 'sample_files')


def hash_array(arr):
    """
    Stable hash of a numeric array: NaN-normalized, C-contiguous, float64.
    Returns None for anything non-numeric so the harness never crashes on a
    stray string argument.
    """
    try:
        a = np.asarray(arr)
        if hasattr(a, 'mask'):
            a = np.where(np.ma.getmaskarray(a), np.nan, np.ma.getdata(a))
        a = np.ascontiguousarray(np.asarray(a, dtype=np.float64))
    except (TypeError, ValueError):
        return None
    # Canonicalize every NaN payload to one bit pattern and -0.0 to 0.0
    a = np.where(np.isnan(a), np.nan, a) + 0.0
    h = hashlib.sha256()
    h.update(str(a.shape).encode())
    h.update(a.tobytes())
    return h.hexdigest()[:16]


class Recorder:
    """
    Wraps the matplotlib Axes methods the script draws with, recording the data
    passed to each call, then delegating to the original implementation.
    """

    METHODS = ('pcolormesh', 'plot', 'contour', 'contourf', 'quiver', 'streamplot')

    def __init__(self):
        self.calls = []
        self._saved = {}

    def __enter__(self):
        for name in self.METHODS:
            original = getattr(matplotlib.axes.Axes, name, None)
            if original is None:
                continue
            self._saved[name] = original
            setattr(matplotlib.axes.Axes, name, self._wrap(name, original))
        # `describe_figures` reads the titles and labels off figures that are
        # still open, and every sweep run passes -o. A renderer that closes what
        # it has just saved - which is the right thing for a tool whose
        # interactive loop would otherwise collect windows until matplotlib
        # complains - would leave nothing here to read, and every record's
        # `axes` would silently go empty. Closing is the caller's business;
        # inspecting afterwards is the harness's.
        self._close = plt.close
        plt.close = lambda *args, **kwargs: None
        return self

    def __exit__(self, *exc):
        for name, original in self._saved.items():
            setattr(matplotlib.axes.Axes, name, original)
        self._saved.clear()
        plt.close = self._close
        return False

    def _wrap(self, name, original):
        calls = self.calls

        def wrapper(ax_self, *args, **kwargs):
            calls.append({
                'method': name,
                'args': [hash_array(a) for a in args],
            })
            return original(ax_self, *args, **kwargs)

        return wrapper


def describe_figures():
    """
    Snapshot the title and axis labels of every open figure, in creation order.

    Two things have to be looked for that `fig.axes` alone does not hold, and
    both of them are how a map describes itself:

    - the title is on the *figure*, not the axes. A cartopy GeoAxes resolves an
      axes title's position to NaN and never rasterizes it, so geomap puts it on
      `fig.suptitle` - which `ax.get_title()` cannot see.
    - the colour bars are inset axes. They are insets because a GeoAxes with
      labelled gridlines and a colorbar taking its room out of the axes collapses
      to nothing under `bbox_inches='tight'`; but an inset is a child of its
      parent axes and is not registered in `fig.axes`.

    Between them, every map was recording one axes with an empty title and empty
    labels - which is to say, nothing at all. This is the record that is supposed
    to catch a renamed axis or a lost colour bar.
    """
    out = []
    for num in plt.get_fignums():
        fig = plt.figure(num)
        out.append({'title': fig.get_suptitle(), 'xlabel': '', 'ylabel': ''})
        for ax in fig.axes:
            for target in (ax, *ax.child_axes):
                out.append({
                    'title': target.get_title(),
                    'xlabel': target.get_xlabel(),
                    'ylabel': target.get_ylabel(),
                })
    return out


def sweep_variable(main, nc_path, varname, dims, shape, tmpdir):
    """
    Run one variable through the CLI in-process and return its record.

    Dimensions beyond the last two are pinned to index 0 so that most variables
    reach a real plotting path instead of being refused for having too many
    dimensions. This mirrors the sweep the README describes.
    """
    extra = {d: 0 for d in dims[:-2]} if len(dims) > 2 else {}
    out_png = os.path.join(tmpdir, 'out.png')
    argv = [nc_path, '-v', varname, '-o', out_png]
    if extra:
        argv += ['-e', json.dumps(extra)]

    record = {'argv': argv[1:], 'extra_indices': extra}
    stdout = _io.StringIO()
    old_argv = sys.argv
    sys.argv = ['display_netcdf.py'] + argv
    try:
        with Recorder() as rec, contextlib.redirect_stdout(stdout), \
                contextlib.redirect_stderr(stdout):
            try:
                status = main()
            except SystemExit as exc:  # argparse
                status = exc.code
            except Exception:
                status = 'EXC'
                stdout.write(traceback.format_exc())
            record['axes'] = describe_figures()
            record['draw_calls'] = rec.calls
    finally:
        sys.argv = old_argv
        plt.close('all')

    record['status'] = status
    # The temporary directory name is different on every run, so scrub it out;
    # otherwise every "Saved to ..." line reports a spurious difference.
    text = stdout.getvalue().strip().replace(tmpdir, '<TMP>')
    record['stdout'] = text.splitlines()
    return record


def capture(entry):
    """
    Import the entry point and sweep every variable of every sample file.
    `entry` is a "module:function" string.
    """
    mod_name, func_name = entry.split(':')
    sys.path.insert(0, _ROOT)
    main = getattr(importlib.import_module(mod_name), func_name)

    result = {'entry': entry, 'files': {}}
    with tempfile.TemporaryDirectory() as tmpdir:
        for nc_path in sorted(glob.glob(os.path.join(EXAMPLES, '*.nc'))):
            name = os.path.basename(nc_path)
            ds = Dataset(nc_path, 'r')
            try:
                variables = [(v, list(ds.variables[v].dimensions),
                              list(ds.variables[v].shape)) for v in ds.variables]
            finally:
                ds.close()

            entries = {}
            for varname, dims, shape in variables:
                print(f"  {name}:{varname}", file=sys.stderr)
                entries[varname] = sweep_variable(main, nc_path, varname,
                                                  dims, shape, tmpdir)
                entries[varname]['dims'] = dims
                entries[varname]['shape'] = shape
            result['files'][name] = entries
    return result


def summarize(result):
    total = ok = 0
    for entries in result['files'].values():
        for rec in entries.values():
            total += 1
            if rec['status'] == 0:
                ok += 1
    return total, ok


def diff(old, new):
    """
    Return a list of human-readable differences between two baselines.
    The `entry` key is ignored: comparing the old script to the new package is
    the whole point.

    A baseline only describes the sample set it was taken on. When that set has
    moved, say so once and compare the files the two have in common, rather than
    reporting every added or removed file as a regression.
    """
    problems = []
    gone = sorted(set(old['files']) - set(new['files']))
    added = sorted(set(new['files']) - set(old['files']))
    if gone or added:
        lines = ["the sample set has changed since this baseline was taken"]
        if gone:
            lines.append(f"    no longer present: {', '.join(gone)}")
        if added:
            lines.append(f"    new: {', '.join(added)}")
        lines.append("    only the files in both are compared; recapture with --out")
        problems.append('\n'.join(lines))

    for name in sorted(set(old['files']) & set(new['files'])):
        a, b = old['files'][name], new['files'][name]
        for var in sorted(set(a) | set(b)):
            ra, rb = a.get(var), b.get(var)
            if ra is None or rb is None:
                problems.append(f"{name}:{var}: present in only one baseline")
                continue
            for key in ('status', 'dims', 'shape', 'axes', 'draw_calls', 'stdout'):
                if ra.get(key) != rb.get(key):
                    problems.append(
                        f"{name}:{var}: {key} differs\n"
                        f"    old: {ra.get(key)}\n"
                        f"    new: {rb.get(key)}"
                    )
    return problems


def main():
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('--entry', default='display_netcdf:main',
                        help='module:function to drive (default: display_netcdf:main)')
    parser.add_argument('--out', help='write the captured baseline to this path')
    parser.add_argument('--check', help='compare against this existing baseline')
    args = parser.parse_args()

    # The sample files are not versioned, and an empty glob would sweep nothing
    # and report a confident pass against a baseline of 194 variables
    if not glob.glob(os.path.join(EXAMPLES, '*.nc')):
        print(f"No sample files in {EXAMPLES}.")
        print("The sweep needs them; see README.md for where to get them.")
        return 2

    result = capture(args.entry)
    total, ok = summarize(result)
    print(f"\nSwept {total} variables, {ok} plotted successfully, {total - ok} refused.")

    if args.out:
        os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
        with open(args.out, 'w') as fh:
            json.dump(result, fh, indent=1, sort_keys=True)
        print(f"Baseline written to {args.out}")

    if args.check:
        with open(args.check) as fh:
            old = json.load(fh)
        problems = diff(old, result)
        if problems:
            print(f"\n{len(problems)} difference(s) against {args.check}:\n")
            for p in problems:
                print(f"  - {p}")
            return 1
        print(f"Identical to {args.check}.")
    return 0


if __name__ == '__main__':
    sys.exit(main())