#!/usr/bin/env python
"""
This script uses matplotlib to compare reference output file(s) with data in tabular form
with the corresponding file(s) generated by the automatic test.
Multiple files can be specified via the command line interface.

Example:

    pltdiff.py t88o_DS2_PHDOS

Usage: pltdiff.py file1 [file2, ...]
"""
from __future__ import unicode_literals, division, print_function, absolute_import

import sys
import os
import numpy as np


def find_top_build_tree(start_path, with_abinit=True, ntrials=10):
    """
    Returns the absolute path of the ABINIT build tree.
    Assume start_path is within the build tree.

    Raises:
        `RuntimeError` if build tree is not found after ntrials attempts.
    """
    abs_path = os.path.abspath(start_path)

    trial = 0
    while trial <= ntrials:
        config_h = os.path.join(abs_path, "config.h")
        abinit_bin = os.path.join(abs_path, "src", "98_main", "abinit")
        # Check if we are in the top of the ABINIT source tree
        if with_abinit:
            found = os.path.isfile(config_h) and os.path.isfile(abinit_bin)
        else:
            found = os.path.isfile(config_h)

        if found:
            return abs_path
        else:
            abs_path, tail = os.path.split(abs_path)
            trial += 1

    raise RuntimeError("Cannot find the ABINIT build tree after %s trials" % ntrials)


def get_array_from_path(path):
    """Parse text file with data in tabular form, return numpy array."""
    data = []
    with open(path, "rt") as fh:
        for l in fh:
            l = l.strip()
            if l.startswith("#") or not l: continue
            data.append(list(map(float, l.split())))

    return np.array(data).T.copy()


def get_axarray_fig_plt(ax_array, nrows=1, ncols=1, sharex=False, sharey=False,
                        squeeze=True, subplot_kw=None, gridspec_kw=None,
                        **fig_kw):
    """
    Helper function used in plot functions that accept an optional array of Axes
    as argument. If ax_array is None, we build the `matplotlib` figure and
    create the array of Axes by calling plt.subplots else we return the
    current active figure.

    Returns:
        ax: Array of :class:`Axes` objects
        figure: matplotlib figure
        plt: matplotlib pyplot module.
    """
    import matplotlib.pyplot as plt

    if ax_array is None:
        fig, ax_array = plt.subplots(nrows=nrows, ncols=ncols, sharex=sharex,
                                     sharey=sharey, squeeze=squeeze,
                                     subplot_kw=subplot_kw,
                                     gridspec_kw=gridspec_kw, **fig_kw)
    else:
        fig = plt.gcf()
        ax_array = np.reshape(np.array(ax_array), (nrows, ncols))
        if squeeze:
            if ax_array.size == 1:
                ax_array = ax_array[0]
            elif any(s == 1 for s in ax_array.shape):
                ax_array = ax_array.ravel()

    return ax_array, fig, plt


def compare_data(ref_data, out_data, title):
    print("ref_data.shape:", ref_data.shape, "out_data.shape:", out_data.shape)
    #if ref_data.shape == out_data.shape:
    #    print(ref_data - out_data)

    nrows = min(len(ref_data), len(out_data)) - 1
    axlist, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=1, sharex=True, sharey=False,
                                           squeeze=True, subplot_kw=None, gridspec_kw=None)
    axlist = np.array(axlist).ravel()

    # Assume file contains data in the form | x | f(x) | g(x) | ....
    ref_xs = ref_data[0]
    out_xs = out_data[0]

    for i, ax in enumerate(axlist):
        ax.plot(ref_xs, ref_data[i+1], "--", label="ref" if i == 0 else None)
        ax.plot(out_xs, out_data[i+1], ":", label="out" if i == 0 else None)
        ax.grid(True)
        if i == 0: ax.legend(loc="best", shadow=True)

    fig.suptitle(title)
    plt.show()


def main():
    if len(sys.argv) == 1 or "-h" in sys.argv or "--help" in sys.argv:
        print(__doc__)
        return 1

    out_paths = sys.argv[1:]
    out_paths = [os.path.abspath(p) for p in out_paths]

    find_refs = len(out_paths) == 1
    if find_refs:
        top = find_top_build_tree(".", with_abinit=False)
        # Need top level directory for ref files so check if tests/v1 exists else try ..
        v1refs = os.path.join(top, "tests", "v1", "Refs")
        print(v1refs)
        if not os.path.isdir(v1refs):
            top = os.path.join(top, "..")
            v1refs = os.path.join(top, "tests", "v1", "Refs")
            if not os.path.isdir(v1refs):
                raise RuntimeError("Cannot find top-level abinit directory containing ~abinit/v1/Refs")

        for out_path in out_paths:
            # Get the name of the substuite from dirname e.g. v7_t85-t86-t87-t88-t89
            suite_dir = os.path.basename(os.path.dirname(out_path)).split("_")[0]
            ref_path = os.path.join(top, "tests", suite_dir, "Refs", os.path.basename(out_path))
            print("Comparing:", os.path.relpath(out_path), "with", ref_path)
            if not os.path.isfile(ref_path):
                raise RuntimeError("Cannot find reference file: `%s`" % ref_path)
            ref_data = get_array_from_path(ref_path)
            out_data = get_array_from_path(out_path)
            compare_data(ref_data, out_data, os.path.basename(out_path))
    else:
        assert len(out_paths) == 2
        ref_data = get_array_from_path(out_paths[0])
        out_data = get_array_from_path(out_paths[1])
        compare_data(ref_data, out_data, os.path.basename(out_paths[0]))

    return 0


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