#! /usr/bin/env python
#
# This script compares RFMIP results from RTE+RRTMGP against a benchmark
#
import argparse
import os
import sys
import urllib.request

import numpy as np
import xarray as xr

#
# Comparing reference and test results
#
if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Compares all-sky example output to file in reference "
                    "directory")
    parser.add_argument("--allsky_file", type=str, default="rrtmgp-allsky.nc",
                        dest="file",
                        help="Name of file inputs and outputs for "
                             "all-sky problem (same for test and reference)")
    parser.add_argument("--ref_dir", type=str, default="ref",
                        help="Directory where reference values are")
    parser.add_argument("--download_reference", action="store_true",
                        help="Download reference files even if they exist")
    parser.add_argument("--report_threshold", type=float, default=0.,
                        help="Threshold for reporting differences")
    parser.add_argument("--failure_threshold", type=float, default=1.e-5,
                        help="Threshold at which differences cause failure "
                             "(for continuous integration)")
    args = parser.parse_args()

    tst_file = args.file
    ref_file = os.path.join(args.ref_dir, tst_file)
    # If a version of the file exists in the reference directory,
    # no need to download (can be over-ridden)
    if args.download_reference or not os.path.exists(ref_file):
        os.makedirs(args.ref_dir, exist_ok=True)
        urllib.request.urlretrieve(
            "ftp://ftp.ldeo.columbia.edu/pub/robertp/rte-rrtmgp/continuous-integration/rrtmgp-allsky.nc",
            ref_file)
    tst = xr.open_dataset(tst_file)
    ref = xr.open_dataset(ref_file)

    failed = False
    for v in ['lw_flux_up', 'lw_flux_dn', 'sw_flux_up', 'sw_flux_dn',
              'sw_flux_dir']:
        if np.all(np.isnan(tst.variables[v].values)):
            raise Exception("All test values are missing. Were the tests run?")
        if np.any(np.isnan(tst.variables[v].values)):
            raise Exception(
                "Some test values are missing. Now that is strange.")

        # express as (tst-ref).variables[v].values when replacing reference file
        # to have same number of columns
        diff = abs((tst.variables[v] - ref.variables[v]).values)
        avg = 0.5 * (tst.variables[v] + ref.variables[v]).values
        # Division raises a runtime warning when we divide by zero even if the
        # values in those locations will be ignored.
        with np.errstate(divide='ignore', invalid='ignore'):
            frac_diff = np.abs(
                np.where((avg > 2. * np.finfo(float).eps), diff / avg, 0))

        if diff.max() > args.report_threshold:
            print(
                'Variable %s differs (max abs difference: %e; '
                'max percent difference: %e%%)' % (
                    v, diff.max(), 100.0 * frac_diff.max()))
        else:
            print('Variable %s: No diffs' % v)

        if diff.max() > args.failure_threshold:
            failed = True

    sys.exit(1) if failed else sys.exit(0)