#!/usr/bin/env python

import argparse
import warnings
import sys
import xarray as xr
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import momgrid

warnings.filterwarnings(
    "ignore",
    message="The input coordinates to pcolormesh are interpreted as cell centers, but are not monotonically increasing or decreasing.*",
)


def main(args):
    resolution = 1.0
    plot_type = "yx"
    level = 0.0
    region = None
    tdim = "time"
    projection = None
    # projection = "Robinson"

    var = args.varname[0]
    varlist = args.varname

    ds1 = momgrid.Gridset(args.ds1)
    ds2 = momgrid.Gridset(args.ds2)

    ds1.subset(varlist)
    ds2.subset(varlist)

    # If an xy plot is requested, extract the desired depth level
    if plot_type == "yx":
        try:
            ds1.data = momgrid.util.extract_level(ds1.data, level)
            ds2.data = momgrid.util.extract_level(ds2.data, level)
        except:
            pass

    if tdim in ds1.data.dims:
        ds1.data = ds1.data.mean(tdim, keep_attrs=True)

    if tdim in ds2.data.dims:
        ds2.data = ds2.data.mean(tdim, keep_attrs=True)

    # Check if grids are identical and regrid if necessary
    if ds1.model == ds2.model:
        ds1 = ds1.data
        ds2 = ds2.data
    else:
        ds1 = ds1.regrid(resolution=resolution)
        ds2 = ds2.regrid(resolution=resolution)

    if plot_type == "yz":
        ds1 = momgrid.util.x_average_dataset(ds1, region=region)
        ds2 = momgrid.util.x_average_dataset(ds2, region=region)

    projection = ccrs.__dict__[projection]() if projection is not None else None

    _ = momgrid.plot.compare_2d(
        ds1[var],
        ds2[var],
        dpi=90,
        singlepanel=False,
        plot_type=plot_type,
        projection=projection,
    )

    plt.show()


def list_vars(args):
    ds1 = xr.open_mfdataset(args.ds1, decode_times=False)
    ds2 = xr.open_mfdataset(args.ds2, decode_times=False)
    varlist = set(ds1.keys()).intersection(ds2.keys())
    for x in sorted(varlist):
        _longname = (
            ds1[x].attrs["long_name"] if "long_name" in ds1[x].attrs.keys() else ""
        )
        print(f" *  {x}  : {_longname}")


if __name__ == "__main__":
    # Create the parser
    parser = argparse.ArgumentParser(
        description="Compare two datasets generated by the MOM ocean model."
    )

    # Add arguments
    parser.add_argument(
        "ds1", type=str, help="Comma-separated list of filenames for the first dataset."
    )
    parser.add_argument(
        "ds2",
        type=str,
        help="Comma-separated list of filenames for the second dataset.",
    )
    parser.add_argument(
        "varname",
        type=str,
        nargs="?",
        default=None,
        help="Variable name to compare. If omitted, a list of available variables will be offered",
    )

    # Parse arguments
    args = parser.parse_args()

    args.ds1 = args.ds1.split(",")
    args.ds2 = args.ds2.split(",")

    if args.varname is None:
        list_vars(args)
        sys.exit(1)
    else:
        args.varname = args.varname.split(",")

    main(args)

    sys.exit(0)
