#! /usr/bin/env python3
"""
Checks the outputs of a workflow run
"""
import argparse
from argparse import RawTextHelpFormatter
from losoto.h5parm import h5parm
import os
import sys
from filecmp import dircmp
import numpy as np


def main(results_path, control_path):
    """
    Checks the outputs of an end-to-end workflow job

    Parameters
    ----------
    results_path : str
        Path to output results directory of job
    control_path : str
        Path to control results directory
    """
    # Check that all expected output files are present
    dcmp = dircmp(results_path, control_path)
    if len(dcmp.left_only) > 0:
        print('ERROR: The following files are present in the output but not in the control: {}'.format(dcmp.left_only))
        sys.exit(1)
    if len(dcmp.right_only) > 0:
        print('ERROR: The following files are present in the control but not in the output: {}'.format(dcmp.right_only))
        sys.exit(1)

    # Check that the calibration solutions match the control ones
    check_h5parm = os.path.join(results_path, 'cal_values', 'cal_solutions.h5')
    control_h5parm = os.path.join(control_path, 'cal_values', 'cal_solutions.h5')
    check_h5 = h5parm(check_h5parm, readonly=True)
    control_h5 = h5parm(control_h5parm, readonly=True)
    for check_solset in check_h5.getSolsets():
        control_solset = control_h5.getSolset(check_solset.name)
        for check_soltab in check_solset.getSoltabs():
            control_soltab = control_solset.getSoltab(check_soltab.name)
            if not np.allclose(check_soltab.val, control_soltab.val, rtol=1e-03,
                               atol=1e-03, equal_nan=True):
                print('ERROR: Val array of soltab {} of solset {} does not match '
                      'the control'.format(check_soltab.anme, check_solset.name))
                sys.exit(1)


if __name__ == '__main__':
    descriptiontext = "Checks the ouput of a workflow run.\n"

    parser = argparse.ArgumentParser(description=descriptiontext, formatter_class=RawTextHelpFormatter)
    parser.add_argument('results_path', help='Path to output results')
    parser.add_argument('control_path', help='Path to control results')
    args = parser.parse_args()

    main(args.results_path, args.control_path)