run:
    import xarray as xr
    ds = xr.open_dataset(input[0])
    mean = ds.mean(['x', 'time'])
    mean.to_netcdf(output[0])

shell: """
papermill -p run_path $PWD/{params.run} -p training_data $PWD/{TRAINING_DATA} \
         -p caseid {wildcards.id} \
        -p training_data_mean {TRAINING_MEAN} \
        --prepare-only {params.template} {params.ipynb}
jupyter nbconvert  --ExecutePreprocessor.timeout=600 \
                   --allow-errors \
                   --execute {params.ipynb}
# clean up the notebook
rm -f {params.ipynb}
"""

import xarray as xr
from uwnet.debias import insert_apparent_sources, LassoDebiasedModel
import torch


o = snakemake.output
p = snakemake.params

ds = xr.open_dataset(p.data)
model = torch.load(p.model)
ds = insert_apparent_sources(ds, prognostics=p.prognostics)

mapping = [
    ('QT', 'QT', 'QQT'),
    ('SLI', 'SLI', 'QSLI'),
]

debias = LassoDebiasedModel(model, mapping).fit(ds)
torch.save(debias, o[0])

import torch
from uwnet.loss import get_input_output
from uwnet.utils import mean_other_dims
import json
from tqdm import tqdm
import argparse

import uwnet.ml_models.nn.datasets_handler as d


def batch_to_residual(model, batch):
    from uwnet.timestepper import Batch

    batch = Batch(batch.float(), prognostics=["QT", "SLI"])
    with torch.no_grad():
        prediction, truth = get_input_output(model, 0.125, batch)
    return prediction - truth


def vertically_resolved_mse_from_residual(residual):
    return {k: mean_other_dims(residual[k] ** 2, 2).squeeze() for k in residual}


def batch_to_mse(model, batch):
    residual = batch_to_residual(model, batch)
    return vertically_resolved_mse_from_residual(residual)


def _parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("data", help="path to zarr reshaped training or test data")
    parser.add_argument("model", help="path to model")

    return parser.parse_args()


args = _parse_args()

model_path = args.model
path = args.data
prognostics = ["QT", "SLI"]

model = torch.load(model_path)
train_dataset = d.get_dataset(path, predict_radiation=False)
dl = d.get_data_loader(train_dataset, prognostics, batch_size=64)

total = {}
count = 0
for batch in tqdm(dl):
    mse = batch_to_mse(model, batch)
    for key in mse:
        count += 1
        alpha = 1 / count
        zeros = torch.zeros_like(mse[key])
        total[key] = total.get(key, zeros) * (1 - alpha) + mse[key] * alpha

total = {"mse_apparent_source_" + k.lower(): total[k].numpy().tolist() for k in total}
print(json.dumps(total))

import xarray as xr
import json
import argparse
from src.sam.case import get_ngqaua_ic
from uwnet.data.blur import blur_dataset
from uwnet.thermo import layer_mass
from src.data.ngaqua import NGAqua
from src.sam.process_ngaqua import run_sam_nsteps
from subprocess import run
import os
import shutil
from os.path import abspath


class Namespace:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)


def get_args_snakemake():
    print("Using snakemake")
    sigma = snakemake.params.sigma
    if sigma:
        sigma = float(sigma)

    return Namespace(
        time_step=int(snakemake.wildcards.step),
        sam_parameters=snakemake.input.sam_parameters,
        sam=snakemake.config.get('sam_path',  '/opt/sam'),
        ngaqua_root=snakemake.params.ngaqua_root,
        output=snakemake.output[0],
        sigma=sigma)


def get_args_argparse():

    parser = argparse.ArgumentParser(
        description='Pre-process a single time step')
    parser.add_argument('-n', '--ngaqua-root', type=str)
    parser.add_argument('-s', '--sam', type=str, default='/opt/sam')
    parser.add_argument('-t', '--time-step', type=int, default=0)
    parser.add_argument('-p', '--sam-parameters', type=str, default=0)
    parser.add_argument('--sigma', type=float, help='Radius for Gaussian blurring')

    parser.add_argument('output')
    return parser.parse_args()


def get_args():
    """Get the arguments needed to run this script

    This will function will behave differently if run from snakemake"""

    try:
        snakemake
    except NameError:
        return get_args_argparse()
    else:
        return get_args_snakemake()


args = get_args()


def get_extra_features(ngaqua: NGAqua, time_step):

    # compute layer_mass
    stat = ngaqua.stat
    rho = stat.RHO.isel(time=0).drop('time')
    w = layer_mass(rho)

    # get 2D variables
    time = ngaqua.times_3d[time_step]
    d2 = ngaqua.data_2d.sel(time=time)

    # add variables to three-d
    d2['RADTOA'] = d2.LWNT - d2.SWNT
    d2['RADSFC'] = d2.LWNS - d2.SWNS
    d2['layer_mass'] = w
    d2['rho'] = rho

    # 3d variables
    qrad = ngaqua.data_3d.QRAD.drop(['x', 'y'])
    d2['QRAD'] = qrad.sel(time=time)

    return d2


# get initial condition
ic = get_ngqaua_ic(args.ngaqua_root, args.time_step)

# get data
ngaqua = NGAqua(args.ngaqua_root)
features = get_extra_features(ngaqua, args.time_step)

if args.sigma:
    print(f"Blurring data with radius {args.sigma}")
    ic = blur_dataset(ic, args.sigma)
    features = blur_dataset(features, args.sigma)

# compute the forcings by running through sam
if args.sam_parameters:
    with open(args.sam_parameters) as f:
        prm = json.load(f)

path = run_sam_nsteps(ic, prm, sam_src=abspath(args.sam))
files = os.path.join(path, 'OUT_3D', '*.nc')
ds = xr.open_mfdataset(files)
shutil.rmtree(path)

for key in ['QT', 'SLI', 'U', 'V']:
    forcing_key = 'F' + key
    src = ds[key].diff('time') / ds.time.diff('time') / 86400
    src = src.isel(time=0)
    ds[forcing_key] = src

ds = ds.isel(time=0)

# forcing data
ic['x'] = ds.x
ic['y'] = ds.y
for key in ic.data_vars:
    if key not in ds:
        ds[key] = ic[key]

ds = ds.merge(features).expand_dims('time')
ds.attrs['sam_namelist'] = json.dumps(prm)
ds.to_netcdf(args.output, unlimited_dims=['time'], engine='h5netcdf')

import xarray as xr
import numpy as np

# arguments
input_file = snakemake.input[0]
output_file = snakemake.output[0]
variables = snakemake.params.variables
shuffle = snakemake.params.shuffle
train_or_test = snakemake.wildcards.train_or_test

chunk_size = 2**10

# open data
ds = xr.open_dataset(input_file)


if train_or_test == "train":
    ds = ds.isel(x=slice(0, 64))
elif train_or_test == "test":
    ds = ds.isel(x=slice(64, None))
else:
    raise NotImplementedError("{test_or_train} is not \"train\" or \"test\"")

# perform basic validation
assert ds['SLI'].dims == ('time', 'z', 'y', 'x')

# stack data
variables = list(variables) # needs to be a list for xarray
stacked = (ds[variables]
           .stack(sample=['y', 'x'])
           .drop('sample'))

# add needed variables
stacked['layer_mass'] = ds.layer_mass.isel(time=0)

# shuffle samples
if shuffle:
    n = len(stacked.sample)
    indices = np.random.choice(n, n, replace=False)
    stacked = stacked.isel(sample=indices)

chunked = stacked.chunk({'sample': chunk_size})

# save to disk
chunked.to_zarr(output_file)