import matplotlib.pyplot as plt
from pathlib import Path
import xarray as xr
import numpy as np


plt.style.use(
    r"/space/hall5/sitestore/eccc/crd/ccrn/users/rlr001/software/ht-paper/ht-paper.mplstyle"
)

folder = Path(
    r"/space/hall5/sitestore/eccc/crd/ccrn/users/rlr001/data/cmip/input/strat_aerosol/v1.3"
)
model = "UKESM"

fields = {
    "ext": "Extinction [m$^{-1}$]",
    "asy": "Asymmetry Factor",
    "ssa": "Single Scattering Albedo",
}

for var in fields:
    files = {
        "original": [
            folder
            / f"{var}_input4MIPs_aerosolProperties_CMIP_UOEXETER-CMIP-1-3-0_gnz_175001-202312.nc"
        ],
        "weighted": list(
            (folder / "UKESM" / "weighted").glob(
                f"{var}_input4MIPs_aerosolProperties_CMIP_{model}_*-CMIP-1-3-0_gnz_175001-202312.nc"
            )
        ),
        "interpolated": list((folder / "UKESM" / "linear").glob(f"{var}_*_CMIP7_interpolated.nc")),
    }

    time = "2010-01-01"
    height = 20000.0
    lat = 0

    data = {}
    for key, file in files.items():
        # need to manually concatenate non-monotonic wavelength datasets
        tmp = xr.concat(
            [
                xr.open_mfdataset(f)
                .sel(height=slice(0, 40000))
                .sel(lat=slice(-90, 90))
                .sel(time=slice("1850-01-01", "2023-12-31"))
                .mean(dim="time")
                .mean(dim="height")
                .sortby("wavelength")
                for f in file
            ],
            dim="wavelength",
        ).sortby("wavelength")
        data[key] = (tmp * np.cos(tmp.lat * np.pi / 180)).sum(dim="lat") / np.sum(
            np.cos(tmp.lat.values * np.pi / 180)
        )

    fig, ax = plt.subplots(2, 1, figsize=(4, 4), sharex=True)
    fig.subplots_adjust(left=0.14, bottom=0.09, top=0.96, right=0.97, hspace=0.02)
    for key, ds in data.items():
        ax[0].plot(ds.wavelength.values, ds[var].values, marker=".", label=key)

    percent = False
    if var == "ext":
        diff = (data["weighted"][var] - data["interpolated"][var]) / data["weighted"][var] * 100
        percent = True
    else:
        diff = data["weighted"][var] - data["interpolated"][var]

    ax[0].set_title(fields[var])
    ax[1].plot(diff.wavelength.values, diff.values, marker=".", label=key)
    ax[0].set_ylabel(fields[var])
    ax[1].set_ylabel(f"Difference (weighted - interp){' [%]' if percent else ''}")
    ax[0].set_xscale("log")
    if percent:
        ax[0].set_yscale("log")
    else:
        ax[0].set_yscale("linear")
    # ax[1].set_yscale('log')
    ax[1].set_xlabel("Wavelength [nm]")
    ax[0].legend()
    fig.savefig(f"compare_methods_{var}-{model}.png", dpi=450)
