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


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/"
)

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": (folder / "canesm").glob(
            f"{var}_input4MIPs_aerosolProperties_CMIP_CanESM_*-CMIP-1-3-0_gnz_175001-202312.nc"
        ),
        "interpolated": (folder / "official").glob(f"{var}_*_CMIP7_interpolated.nc"),
    }

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

    data = {}
    for key, file in files.items():
        data[key] = (
            xr.open_mfdataset(file).sel(time=time).sel(height=height, lat=lat, method="nearest")
        )

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

    diff = (data["weighted"][var] - data["interpolated"][var]) / data["weighted"][var] * 100
    ax[0].set_title(fields[var])
    ax[1].plot(diff.wavelength.values * 1e9, diff.values, marker=".", label=key)
    ax[0].set_ylabel(fields[var])
    ax[1].set_ylabel("Difference (weighted - interp) [%]")
    ax[0].set_xscale("log")
    ax[0].set_yscale("log")
    ax[1].set_xlabel("Wavelength [nm]")
    ax[0].legend()
    fig.savefig(f"compare_methods_{var}.png", dpi=450)
