import yaml
import numpy as np
import xarray as xr
from pathlib import Path
from strataer.cmip7.convert_cmip7 import (
    interpolate_ext,
    interpolate_ssa,
    interpolate_asy,
    planck_weights,
)
from strataer.utils import wavenumber_to_wavelength
from functools import partial
import os


def cli_run(config):
    """Interpolate the CMIP7 stratospheric aerosol forcing to new wavelength bands."""

    config = yaml.safe_load(open(config, "r"))
    folder = Path(config["input"]["folder"])

    for key, band in config["output"]["bands"].items():
        output_wavel_bnds = np.array(band["bounds"])
        units = band["units"]
        base_filename = config["input"]["base_filename"]
        if "model" in config["output"]:
            model = f'{config["output"]["model"]}_{key}'
        else:
            model = f"INTERPOLATED_{key}"

        if units == "cm-1":
            output_wavel_bnds = wavenumber_to_wavelength(output_wavel_bnds) * 1e-9
        if units == "nm":
            output_wavel_bnds = output_wavel_bnds * 1e-9

        source = band["source"]
        function_name = list(source.keys())[0]
        if function_name.lower() == "planck":
            weight_function = partial(planck_weights, **source[function_name])
        else:
            raise ValueError("only planck source function weighting is currently supported")

        wavelength = [(b[1] + b[0]) / 2 for b in output_wavel_bnds]
        output_wavel_bnds = xr.DataArray(
            np.array(output_wavel_bnds),
            dims=["wavelength", "bnds"],
            coords=[wavelength, [0, 1]],
        )

        output_folder = Path(config["output"]["folder"])
        os.makedirs(output_folder, exist_ok=True)
        new_base = base_filename.replace("UOEXETER", model)

        print(f"converting extinction for {key}")
        ext = interpolate_ext(
            folder, output_wavel_bnds, base_filename=base_filename, weight_function=weight_function
        )
        ext.to_netcdf(output_folder / f"ext_{new_base}")

        print(f"converting assymetry for {key}")
        asy = interpolate_asy(
            folder, output_wavel_bnds, base_filename=base_filename, weight_function=weight_function
        )
        asy.to_netcdf(output_folder / f"asy_{new_base}")

        print(f"converting single scattering albedo for {key}")
        ssa = interpolate_ssa(
            folder, output_wavel_bnds, base_filename=base_filename, weight_function=weight_function
        )
        ssa.to_netcdf(output_folder / f"ssa_{new_base}")
