USDA-ARS-NWRC/awsm

#!/usr/bin/env python3

#
# Script to plot difference in gold files between current branch and master
# Must be ran from the project root directory
#

import argparse
from pathlib import PurePath

from goldmeister.compare import GoldGitBranchCompare

import awsm

files = [
    'ipysnobal',
]


def argument_parser():
    parser = argparse.ArgumentParser(
        description="Plot the differences in the gold files "
        "for two different branches. This is meant to document "
        "changes in the gold files"
    )

    parser.add_argument(
        '--gold_dir', '-g',
        metavar='GOLD_DIR',
        type=str,
        default='gold',
        choices=['gold', 'gold_hrrr'],
        help='Gold directory (gold)'
    )

    parser.add_argument(
        '--basin', '-b',
        metavar='BASIN',
        type=str,
        default='RME',
        choices=['RME', 'Lakes'],
        help='Test basin name (RME or Lakes)'
    )

    parser.add_argument(
        '--old_branch', '-o',
        metavar='OLD_BRANCH',
        type=str,
        default='main',
        help='Old branch to compare against. Default: main'
    )

    parser.add_argument(
        '--new_branch', '-n',
        metavar='NEW_BRANCH',
        type=str,
        help='New branch to compare against'
    )

    parser.add_argument(
        '--show-plots',
        action='store_true',
        help='Show the plots before storing to the output folder'
    )

    return parser


def main(args):
    source_path = PurePath(awsm.__file__).parent
    repo_path = source_path.parent
    tests_dir = source_path.joinpath('tests', 'basins')

    gold_files = []
    for gf in files:
        gold_files.append(
            tests_dir.joinpath(
                args.basin,
                args.gold_dir,
                "{}.nc".format(gf)
            ).as_posix()
        )

    gc = GoldGitBranchCompare(
        repo_path=repo_path,
        gold_files=gold_files,
        file_type='netcdf',
        old_branch=args.old_branch,
        new_branch=args.new_branch)

    results = gc.compare()

    gc.plot_results(
        results,
        plot_original_data=False,
        include_hist=True,
        show_plots=args.show_plots,
    )


if __name__ == '__main__':
    script_arguments = argument_parser().parse_args()
    main(script_arguments)