"""The Tutorial object: paths, configs, forcing and initial-condition copies, running CREDIT, and loading output."""

import os
import shutil
from importlib.metadata import PackageNotFoundError, version
from pathlib import Path

import cftime
import numpy as np
import pandas as pd
import torch
import xarray as xr
import yaml

from . import diagnostics as dg
from . import edits
from .runner import run_configs

# Dotted keys into the CAMulator rollout config.
DYN_FORCING_KEY = "data.source.CESM.variables.dynamic_forcing.path"
IC_KEYS = ("data.source.CESM.variables.prognostic.path", "data.source.CESM.variables.diagnostic.path")
OUTPUT_KEY = "inference.save_forecast"
FCST_START_KEY = "inference.single_forecast.start_datetime"
FCST_LEN_KEY = "inference.single_forecast.forecast_length"
OUTPUT_VARS_KEY = "inference.output.variables"
STATICS_PHYSICS_KEY = "postblocks.per_step.mass_fixer.args.save_loc_physics"
STATIC_KEY = "data.source.CESM.variables.static.path"
SCALER_KEY = "preblocks.per_step.scaler.args.scaler_path"

# Inputs each section needs in the data directory.
DATA_FILES = {
    "all sections": ["statics_b_credit_runs_f32_02.nc",
                     "b.e21.CREDIT_climate.statics_1.0deg_32levs_latlon_F32_hyai_fixed.nc"],
    "section 1 (2013 waves)": ["aew_initial_conditions_2013.nc", "aew_forcing_2013.nc", "truth_aew_2013.nc"],
    "sections 3 and 5": ["initial_conditions_1981.nc", "forcing_climo.nc"],
}


def get_by_path(d, dotted):
    for k in dotted.split("."):
        d = d[k]
    return d


def set_by_path(d, dotted, value):
    keys = dotted.split(".")
    for k in keys[:-1]:
        d = d[k]
    d[keys[-1]] = value


# Everyday units: (factor, units, new name or None). CAM's PRECT and QFLX are metres of water per 6-hour step, its
# surface heat fluxes are J/m² per step with negative = upward, and its TAUX is the stress on the atmosphere (positive
# under easterly winds), so the wind's eastward push on the surface is -TAUX.
EVERYDAY_UNITS = {
    "PRECT": (4000.0, "mm/day", None),
    "QFLX": (-4000.0, "mm/day (evaporation, positive upward)", None),
    "PS": (0.01, "hPa", None),
    "PSL": (0.01, "hPa", None),
    "LHFLX": (-1 / 21600, "W/m² (positive upward)", None),
    "SHFLX": (-1 / 21600, "W/m² (positive upward)", None),
    "TAUX": (-1000.0, "mN/m² (positive = westerly wind)", "westerly_stress"),
}


def everyday_units(ds):
    """Convert CAM's storage units to everyday ones (see EVERYDAY_UNITS); other variables are unchanged."""
    for var, (factor, units, new) in EVERYDAY_UNITS.items():
        if var in ds:
            da = (ds[var] * factor).assign_attrs(ds[var].attrs, units=units)
            ds = ds.drop_vars(var).assign({new or var: da})
    return ds


def _days(length):
    return pd.Timedelta(str(length).replace("d", "D"))


class Tutorial:
    """Everything the notebook needs to set up, run and read back CAMulator experiments.

    model_dir:  checkpoint, scaler, metadata and the base config (camulator_inference_tutorial.yml)
    data_dir:   initial conditions, forcing, statics, truth
    work_dir:   written by the tutorial: configs, forcing and initial-condition copies, rollouts
    """

    def __init__(self, model_dir, data_dir, work_dir, checkpoint="checkpoint.pt00069.pt",
                 config_name="camulator_inference_tutorial.yml", data_repo="charlie-becker/CAMulator_IC_and_forcing",
                 everyday_units=False):
        self.model_dir, self.data_dir, self.work_dir = Path(model_dir), Path(data_dir), Path(work_dir)
        self.data_repo = data_repo  # Hugging Face dataset repo with the per-year initial conditions and forcing
        self.everyday_units = everyday_units  # load() and truth() convert to mm/day, hPa, W/m² (see EVERYDAY_UNITS)
        # The config's paths are $CAMULATOR_MODEL / $CAMULATOR_DATA / $CAMULATOR_WORK templates; the credit subprocess
        # inherits these.
        os.environ["CAMULATOR_MODEL"] = str(self.model_dir)
        os.environ["CAMULATOR_DATA"] = str(self.data_dir)
        os.environ["CAMULATOR_WORK"] = str(self.work_dir)
        self.base_config = self.model_dir / config_name
        self.checkpoint = self.model_dir / checkpoint
        missing = [str(p) for p in (self.base_config, self.checkpoint) if not p.exists()]
        if missing:
            raise FileNotFoundError("Missing model files:\n  " + "\n  ".join(missing))

        self.conf = self.load_config()
        self.timestep = pd.Timedelta(self.conf["data"]["timestep"])
        self.rollout_root = self.work_dir / "rollout"
        self.forcing_root = self.work_dir / "forcing"
        self.ic_root = self.work_dir / "initial_conditions"
        for p in (self.rollout_root, self.forcing_root, self.ic_root):
            p.mkdir(parents=True, exist_ok=True)

        # CREDIT loads <save_loc>/checkpoint.pt and writes its grid schema there: link the checkpoint into the work dir.
        save_loc = Path(self.conf["save_loc"])
        save_loc.mkdir(parents=True, exist_ok=True)
        link = save_loc / "checkpoint.pt"
        if link.is_symlink() or link.exists():
            link.unlink()
        link.symlink_to(self.checkpoint)

        self.statics_physics = Path(get_by_path(self.conf, STATICS_PHYSICS_KEY))
        if self.statics_physics.exists():
            with xr.open_dataset(self.statics_physics) as st:
                self.hyai, self.hybi = st["hyai"].load(), st["hybi"].load()  # hyai in Pa in this file
                self.phis = st["PHIS"].load()  # surface geopotential (m²/s²)
                self.land = st["LANDFRAC"].load() > 0.5
            assert float(self.hyai.max()) > 1, "expected hyai in Pa"

    # --- environment ---------------------------------------------------------------------------------------------------
    def check(self):
        """Print the compute environment and any missing inputs."""
        try:
            credit_version = version("miles-credit")
        except PackageNotFoundError:
            credit_version = "not installed"
        if torch.cuda.is_available():
            props = torch.cuda.get_device_properties(0)
            device = f"cuda: {props.name}, {props.total_memory / 2**30:.0f} GiB"
        elif torch.backends.mps.is_available():
            device = "mps (Apple GPU; no TF32, so FAST_MATH has no effect)"
        else:
            device = "cpu (very slow)"
        free = shutil.disk_usage(self.work_dir).free / 1e9
        import credit.applications.rollout_gen2 as rollout

        print(f"Device         : {device}")
        print(f"CREDIT         : miles-credit {credit_version}"
              + ("" if hasattr(rollout, "load_model_for_inference") else "  (WARNING: ensemble launcher unsupported)"))
        print(f"Free disk      : {free:.0f} GB in {self.work_dir}" + ("  (WARNING: under 30 GB)" if free < 30 else ""))
        print(f"Model          : {self.checkpoint}")
        for section, files in DATA_FILES.items():
            gone = [f for f in files if not (self.data_dir / f).exists()]
            print(f"Data, {section:20s}: " + ("ok" if not gone else "MISSING " + ", ".join(gone)))

    def load_config(self):
        # Expand $CAMULATOR_* before parsing, so every path in the config is absolute.
        return yaml.safe_load(os.path.expandvars(self.base_config.read_text()))

    def variables(self):
        """Names of the variables a run can save, by kind: {'prognostic 3d': [...], 'diagnostic 2d': [...], ...}."""
        out = {}
        for ftype in ("prognostic", "diagnostic"):
            spec = self.conf["data"]["source"]["CESM"]["variables"][ftype]
            for dims, tag in (("vars_3D", "3d"), ("vars_2D", "2d")):
                if spec.get(dims):
                    out[f"{ftype} {tag}"] = list(spec[dims])
        return out

    def fetch_year(self, start, length):
        """Download (once) what a run with the actual forcing needs from the Hugging Face dataset repo: the start year's
        initial_conditions_<year>.nc and forcing_real_<year>.nc for every year the run reads."""
        from huggingface_hub import hf_hub_download

        names = [f"initial_conditions_{pd.Timestamp(start).year}.nc"]
        names += [f"forcing_real_{y}.nc" for y in self.years(start, length)]
        for name in names:
            if not (self.data_dir / name).exists():
                print(f"  downloading {name} ...", flush=True)
                hf_hub_download(self.data_repo, name, repo_type="dataset", local_dir=self.data_dir)
        print("  have: " + ", ".join(names))

    def var_key(self, name):
        """Config key of a variable, e.g. 'PS' -> 'CESM/prognostic/2d/PS'."""
        for ftype, spec in self.conf["data"]["source"]["CESM"]["variables"].items():
            for dims, tag in (("vars_3D", "3d"), ("vars_2D", "2d")):
                if name in (spec.get(dims) or []):
                    return f"CESM/{ftype}/{tag}/{name}"
        raise KeyError(f"{name} is not a CAMulator variable")

    # --- forcing -------------------------------------------------------------------------------------------------------
    def years(self, start, length):
        """Years of forcing a run reads: the last step reads forcing one timestep before the final valid time."""
        start = pd.Timestamp(start)
        end = start + _days(length) - self.timestep
        return list(range(start.year, end.year + 1))

    def forcing(self, kind, start, length):
        """%Y path template of the base forcing for a run: 'real' (forcing_real_<year>.nc, the actual year) or 'climo'
        (the 1980-2014 climatology in forcing_climo.nc, copied with its time stamps set to each year the run reads,
        because CREDIT finds forcing by year)."""
        years = self.years(start, length)
        if kind == "real":
            template = str(self.data_dir / "forcing_real_%Y.nc")
            gone = [template.replace("%Y", str(y)) for y in years if not Path(template.replace("%Y", str(y))).exists()]
            assert not gone, "Missing forcing: " + ", ".join(gone)
            return template
        assert kind == "climo", "forcing kind must be 'real' or 'climo'"
        src = self.data_dir / "forcing_climo.nc"
        out_dir = self.forcing_root / "climo"
        out_dir.mkdir(parents=True, exist_ok=True)
        for y in years:
            out = out_dir / f"forcing_climo_{y}.nc"
            if out.exists() and out.stat().st_mtime > src.stat().st_mtime:
                continue  # already made from this climatology
            with xr.open_dataset(src) as f:
                clim = f.load()
            calendar = clim["time"].encoding.get("calendar", "noleap")
            assert len({t.year for t in clim["time"].values}) == 1, f"{src} must hold exactly one year"
            ds = clim.assign_coords(time=[cftime.datetime(y, t.month, t.day, t.hour, calendar=calendar)
                                          for t in clim["time"].values])
            ds["time"].encoding = {"units": f"days since {y}-01-01", "calendar": calendar}
            self._write_forcing(ds, out)
        return str(out_dir / "forcing_climo_%Y.nc")

    def perturb_forcing(self, base_template, name, start, length, edit_fn):
        """Write edit_fn(forcing) for every year the run reads to work_dir/forcing/<name>/forcing_<year>.nc and return
        the %Y template. Only forcing changes; the initial state is untouched."""
        out_dir = self.forcing_root / name
        out_dir.mkdir(parents=True, exist_ok=True)
        for y in self.years(start, length):
            with xr.open_dataset(base_template.replace("%Y", str(y))) as src:
                ds = edit_fn(src.load())
            self._write_forcing(ds, out_dir / f"forcing_{y}.nc")
        return str(out_dir / "forcing_%Y.nc")

    @staticmethod
    def _write_forcing(ds, out):
        # One time step per chunk keeps the rollout's per-step reads cheap.
        enc = {v: {"chunksizes": (1, *ds[v].shape[1:]), "zlib": True, "complevel": 1} for v in ds.data_vars}
        ds.to_netcdf(out, encoding=enc)
        print(f"  wrote {out}")

    # --- initial conditions --------------------------------------------------------------------------------------------
    def initial_state(self, start, ic_template=None):
        """The initial state at `start` (00Z dates in the file), time dim of length 1 kept."""
        start = pd.Timestamp(start)
        path = (ic_template or get_by_path(self.conf, IC_KEYS[0])).replace("%Y", str(start.year))
        with xr.open_dataset(path) as src:
            hit = [k for k, t in enumerate(src["time"].values)
                   if (t.year, t.month, t.day, t.hour) == (start.year, start.month, start.day, start.hour)]
            assert hit, f"no initial condition at {start} in {path}"
            return src.isel(time=hit).load()

    def perturb_initial_state(self, start, edit_fn, name, ic_template=None):
        """Write edit_fn(initial state at start) to work_dir/initial_conditions/<name>/initial_conditions_<year>.nc and
        return the %Y template (set it on the run's initial-condition keys)."""
        out_dir = self.ic_root / name
        out_dir.mkdir(parents=True, exist_ok=True)
        ds = edit_fn(self.initial_state(start, ic_template))
        out = out_dir / f"initial_conditions_{pd.Timestamp(start).year}.nc"
        ds.to_netcdf(out, encoding={v: {"zlib": True, "complevel": 1} for v in ds.data_vars})
        return str(out_dir / "initial_conditions_%Y.nc")

    # --- configs and runs ----------------------------------------------------------------------------------------------
    def run_dir(self, name):
        return self.rollout_root / name

    def write_config(self, name, start, length, forcing, ic=None, outputs=None, output_interval=None):
        """Write work_dir/rollout/<name>/config.yml for one run and return its path. Clears that run's old output.

        start: '1981-02-25' (00Z) or '1981-02-25T06'; length: '14d'; forcing: %Y template (see forcing());
        ic: %Y initial-condition template (default: the base config's); outputs: variable names to save (default: the
        base config's list); output_interval: e.g. '24h' to save once a day (default: every 6-h step)."""
        conf = self.load_config()
        start = pd.Timestamp(start)
        run_dir = self.run_dir(name)
        set_by_path(conf, FCST_START_KEY, start.strftime("%Y-%m-%dT%H:%MZ"))
        set_by_path(conf, FCST_LEN_KEY, str(length))
        conf["data"]["start_datetime"] = start.strftime("%Y-%m-%d %H:%M:%S")
        conf["data"]["end_datetime"] = (start + _days(length)).strftime("%Y-%m-%d %H:%M:%S")
        set_by_path(conf, DYN_FORCING_KEY, forcing)
        for key in IC_KEYS if ic else ():
            set_by_path(conf, key, ic)
        if outputs is not None:
            set_by_path(conf, OUTPUT_VARS_KEY, [{"name": self.var_key(v)} for v in outputs])
        conf["inference"]["output"]["output_interval"] = output_interval
        set_by_path(conf, OUTPUT_KEY, str(run_dir) + "/")
        if run_dir.exists():  # old forecasts (one sub-directory per start time) would mix with the new one
            for old in run_dir.glob("*_*Z"):
                shutil.rmtree(old)
        run_dir.mkdir(parents=True, exist_ok=True)
        path = run_dir / "config.yml"
        with open(path, "w") as f:
            yaml.safe_dump(conf, f, sort_keys=False)
        return path

    def member(self, name, start, length, forcing, amplitude=0.0, seed=0, ic_template=None, **config_kw):
        """Config for one run from the initial state at `start`, with T noise of std `amplitude` K (pattern set by
        `seed`) added when amplitude > 0. Returns the config path."""
        ic = ic_template
        if amplitude > 0:
            ic = self.perturb_initial_state(start, edits.add_T_noise(amplitude, seed), name, ic_template)
        return self.write_config(name, start, length, forcing, ic=ic, **config_kw)

    def ensemble(self, name, n_members, amplitude, start, length, forcing, ic_template=None, **config_kw):
        """Configs for members <name>/m00 ... : member 0 starts from the unperturbed state, member k adds T noise with
        seed k. Two ensembles with the same n_members and amplitude get the same perturbations member by member."""
        return [self.member(f"{name}/m{k:02d}", start, length, forcing, amplitude=amplitude if k else 0.0, seed=k,
                            ic_template=ic_template, **config_kw) for k in range(n_members)]

    def run(self, configs, fast_math=False):
        """Run configs in one CREDIT process (the model is loaded once). fast_math: TF32 (see credit_launch.py)."""
        run_configs(configs, fast_math=fast_math)

    # --- output --------------------------------------------------------------------------------------------------------
    def load(self, name, variables=None):
        """A run's output as one lazy dataset, with a 'lead' coordinate (days since the start)."""
        run_dir = self.run_dir(name)
        files = sorted(run_dir.glob("*_*Z/*.nc"))
        if not files:
            raise FileNotFoundError(f"No output under {run_dir}")
        ds = xr.open_mfdataset(files, combine="by_coords")
        if variables is not None:
            ds = ds[[v for v in variables if v in ds]]
        conf = yaml.safe_load((run_dir / "config.yml").read_text())
        t0 = pd.Timestamp(get_by_path(conf, FCST_START_KEY).replace("Z", ""))
        t0 = cftime.datetime(t0.year, t0.month, t0.day, t0.hour, calendar="noleap")
        lead = np.array([(t - t0).total_seconds() for t in ds["time"].values]) / 86400.0
        ds = ds.assign_coords(lead=("time", lead))
        return everyday_units(ds) if self.everyday_units else ds

    def load_ensemble(self, name, n_members):
        """The members <name>/m00, m01, ... of an ensemble, as a list of datasets."""
        return [self.load(f"{name}/m{k:02d}") for k in range(n_members)]

    def truth(self, name="truth_cyclone_1981.nc"):
        ds = xr.open_dataset(self.data_dir / name)
        return everyday_units(ds) if self.everyday_units else ds

    def first_forcing_step(self, template, start):
        """The first time step of a forcing file (for previews of SST changes)."""
        with xr.open_dataset(template.replace("%Y", str(pd.Timestamp(start).year))) as f:
            return f.isel(time=[0]).load()

    @staticmethod
    def _ps_pa(ds):
        ps = ds["PS"]
        return ps * 100.0 if ps.attrs.get("units") == "hPa" else ps

    def mslp_hpa(self, ds):
        """Sea-level pressure (hPa) of a dataset with PS and TREFHT (CREDIT's Trenberth formula, PHIS from statics)."""
        return dg.mslp(self._ps_pa(ds).load(), ds["TREFHT"].load(), self.phis) / 100.0

    def land_mask(self, like):
        """Land (LANDFRAC > 0.5) as a boolean DataArray on like's latitude/longitude coordinates."""
        return xr.DataArray(self.land.values, dims=(dg.LAT, dg.LON),
                            coords={dg.LAT: like[dg.LAT].values, dg.LON: like[dg.LON].values})

    def precipitable_water(self, ds):
        """Column water vapour (mm) from Qtot and surface pressure."""
        return dg.precipitable_water(ds.assign(PS=self._ps_pa(ds)), self.hyai, self.hybi)