DEFAULT = {
    "model": "unet",
    "data": "data/Goop",
    "epochs": 100,
    "steps_per_call": 8,
    "autoregressive_steps": 2,
    "grid_size": 64,
    "batch_size": 128,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "architecture": {"hidden": [64, 64, 64]},
    "lr": 1e-3,
    "min_lr": 1e-6,
    "use_schedulers": True,
    "length": 0.015,
    "grid_noise": [0.001, 0.001, 0.001, 0.0],
    "particle_noise": 0.0003,
}

"""
FNO : architecture = [hidden_channels_list, modes]
                                             ^ None -> grid_size // 2

Types:                                      
Explicit: iterate over parameters
Combi: all combinations of parameters

If only one parameter, either [param] or param (param will be converted to [param])
"""

fno_vs_unet = {
    "exp_type": "explicit",
    "model": ("fno", "unet"),
    "architecture": (
        {"hidden": [16, 16, 16, 16], "modes": None},
        {"hidden": [64, 64, 64]},
    ),
}

"""
Test experience.
"""
dummy = {
    "exp_type": "explicit",
    "model": "fno",
    "run_name": "dummy2",
    "architecture": {"hidden": [8, 8], "modes": 4},
    "grid_size": 32,
    "batch_size": 256,
    "autoregressive_steps": 3,
    "epochs": 50,
    "grid_noise": [0.01, 0.001, 0.001, 0.0],
    "particle_noise": 0.003,
}

fno = {
    "exp_type": "explicit",
    "model": "fno",
    "architecture": (
        {"hidden": [H] * (i + 1), "modes": None, "use_mlp": use_mlp}
        for i in range(2, 5)
        for H in (32, 64)
        for use_mlp in (True, False)
    ),
}

grid_noise = {
    "exp_type": "explicit",
    "model": "unet",
    "run_name": (
        "grid_0",
        "grid_0-001",
        "grid_0-005",
    ),
    "architecture": {"hidden": [64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": (
        [0.0, 0.0, 0.0, 0.0],
        [0.001, 0.001, 0.001, 0.0],
        [0.005, 0.005, 0.005, 0.0],
    ),
    "particle_noise": (0.0003,),
}


sanity = {
    "exp_type": "explicit",
    "model": "unet",
    "run_name": ("sanity_1",),
    "architecture": {"hidden": [64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 100,
    "steps_per_call": 8,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": ([0.001, 0.001, 0.001, 0.0],),
    "particle_noise": (0.0003,),
}

particles_noise = {
    "exp_type": "explicit",
    "model": "unet",
    "run_name": (
        "particles_0",
        "particles_0-0003",
        "particles_0-0006",
        "particles_0-001",
    ),
    "architecture": {"hidden": [64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": ([0.001, 0.001, 0.001, 0.0],),
    "particle_noise": (0.0, 0.0003, 0.0006, 0.001),
}


autoregressive = {
    "exp_type": "explicit",
    "model": "unet",
    "run_name": (
        "auto_1",
        "auto_2",
        "auto_3",
        "auto_4",
    ),
    "architecture": {"hidden": [64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": (1, 2, 3, 4),
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": ([0.001, 0.001, 0.001, 0.0],),
    "particle_noise": (0.0003,),
}

datasets = {
    "exp_type": "explicit",
    "project:": "test_datasets",
    "data": (
        "data/WaterRamps",
        "data/Goop",
        "data/SandRamps",
    ),
}


autoregressive_no_noise = {
    "project": "hyperparam_tunning",
    "exp_type": "explicit",
    "model": "unet",
    "run_name": (
        "nauto_1",
        "nauto_2",
        "nauto_3",
        "nauto_4",
    ),
    "architecture": {"hidden": [64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": (1, 2, 3, 4),
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": ([0.0, 0.0, 0.0, 0.0],),
    "particle_noise": (0.0,),
}

width = {
    "project": "hyperparam_tunning",
    "exp_type": "explicit",
    "model": "unet",
    "run_name": (
        "w32",
        "w64",
        "w128",
    ),
    "architecture": (
        {"hidden": [32, 32, 32]},
        {"hidden": [64, 64, 64]},
        {"hidden": [128, 128, 128]},
    ),
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
}


depth = {
    "project": "hyperparam_tunning",
    "exp_type": "explicit",
    "model": "unet",
    "run_name": (
        "d2",
        "d3",
        "d4",
        "d5",
    ),
    "architecture": (
        {"hidden": [64, 64]},
        {"hidden": [64, 64, 64]},
        {"hidden": [64, 64, 64, 64]},
        {"hidden": [64, 64, 64, 64, 64]},
    ),
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
}


grid_size = {
    "project": "hyperparam_tunning",
    "exp_type": "explicit",
    "model": "unet",
    "run_name": (
        "g32",
        "g64",
        "g128",
    ),
    "architecture": {"hidden": [64, 64, 64]},
    "grid_size": (32, 64, 128),
    "batch_size": 64,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 8,
}


bundle = {
    "data": "data/Goop",
    "project": "hyperparam_tunning",
    "exp_type": "explicit",
    "model": "unet",
    "run_name": (
        "b1",
        "b2",
        "b4",
        "b8",
        "b16",
        "b32",
    ),
    "architecture": {"hidden": [64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": (1, 2, 4, 8, 16, 32),
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
}


timee_bundle = {
    "data": "data/Goop",
    "project": "hyperparam_tunning",
    "exp_type": "explicit",
    "model": "unet",
    "run_name": (
        "lb1",
        "lb2",
        "lb4",
        "lb8",
        "lb16",
        "lb32",
    ),
    "architecture": {"hidden": [64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": (1, 2, 4, 8, 16, 32),
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
}

"""
Demo experiment for a UNO run.
"""
demo_uno = {
    "exp_type": "explicit",
    "model": "uno",
    "run_name": "demo_uno_v2",
    "architecture": {
        "hidden_channels": 64,
        "uno_out_channels": [64, 128, 128, 128, 128, 128, 64],
        "n_modes": [32, 32, 16, 8, 8, 16, 32],
        "scalings": [1.0, 0.5, 0.5, 1.0, 2.0, 2.0, 1.0],
    },
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 500,
    "steps_per_call": 8,
    "autoregressive_steps": 2,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": [0.001, 0.001, 0.001, 0.0],
    "particle_noise": 0.0003,
}

SOTA = {
    "project": "SOTA",
    "exp_type": "explicit",
    "model": "unet",
    "data": "data/Goop",
    "run_name": ("SOTA_Goop",),
    "architecture": {"hidden": [64, 64, 64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 4,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": ([0.0, 0.0, 0.0, 0.0],),
    "particle_noise": (0.0,),
}

SOTAW = {
    "project": "SOTA",
    "exp_type": "explicit",
    "model": "unet",
    "data": "data/WaterRamps",
    "run_name": ("SOTA_WaterRamps",),
    "architecture": {"hidden": [64, 64, 64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 4,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": ([0.0, 0.0, 0.0, 0.0],),
    "particle_noise": (0.0,),
}

SOTASv2 = {
    "project": "SOTA",
    "exp_type": "explicit",
    "model": "unet",
    "data": "data/SandRamps",
    "run_name": ("SOTA_SandRamps2",),
    "architecture": {"hidden": [64, 64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": ([0.0, 0.0, 0.0, 0.0],),
    "particle_noise": (0.0,),
}


SOTAWBC = {
    "project": "SOTA",
    "exp_type": "explicit",
    "model": "unet",
    "data": "data/WBC-SPH",
    "run_name": ("SOTA_WBC4", "SOTA_WBC3"),
    "architecture": ({"hidden": [64, 64, 64, 64]}, {"hidden": [64, 64, 64, 64]}),
    "grid_size": 64,
    "batch_size": 64,
    "epochs": 1000,
    "steps_per_call": 16,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 4,
    "grid_noise": ([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
    "particle_noise": 0.0,
}

DEFAULT_MODELS_NOISE = {
    "project": "default_models",
    "exp_type": "explicit",
    "model": "unet",
    "data": ("data/Goop", "data/WaterRamps", "data/SandRamps"),
    "run_name": ("Goop, WaterRamps, SandRamps",),
    "architecture": {"hidden": [64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": ([0.001, 0.001, 0.001, 0.0]),
    "particle_noise": (0.001,),
}


DEFAULT_MODELS_NO_NOISE = {
    "project": "default_models_non_noise",
    "exp_type": "explicit",
    "model": "unet",
    "data": ("data/Goop", "data/WaterRamps", "data/SandRamps"),
    "run_name": ("Goop, WaterRamps, SandRamps",),
    "architecture": {"hidden": [64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": ([0.0, 0.0, 0.0, 0.0],),
    "particle_noise": (0.0,),
}


SOTA_auto_4 = {
    "project": "SOTA",
    "exp_type": "explicit",
    "model": "unet",
    "data": ("data/Goop", "data/WaterRamps", "data/SandRamps"),
    "run_name": ("SOTA_Goop_auto_4", "SOTA_Water_auto_4", "SOTA_Sand_auto_4"),
    "architecture": {"hidden": [64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 4,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": ([0.0, 0.0, 0.0, 0.0],),
    "particle_noise": (0.0,),
}

SOTA_depth_4 = {
    "project": "SOTA",
    "exp_type": "explicit",
    "model": "unet",
    "data": ("data/Goop", "data/WaterRamps", "data/SandRamps"),
    "run_name": ("SOTA_Goop_depth_4", "SOTA_Water_depth_4", "SOTA_Sand_depth_4"),
    "architecture": {"hidden": [64, 64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": ([0.0, 0.0, 0.0, 0.0],),
    "particle_noise": (0.0,),
}


SOTA_bundle_3 = {
    "project": "SOTA",
    "exp_type": "explicit",
    "model": "unet",
    "data": ("data/Goop", "data/WaterRamps", "data/SandRamps"),
    "run_name": ("SOTA_Goop_bundle_16", "SOTA_Water_bundle_16", "SOTA_Sand_bundle_16"),
    "architecture": {"hidden": [64, 64, 64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 16,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": ([0.0, 0.0, 0.0, 0.0],),
    "particle_noise": (0.0,),
}

lgrid_size = {
    "data": "data/Goop",
    "project": "hyperparam_tunning",
    "exp_type": "explicit",
    "model": "unet",
    "run_name": (
        "lg32",
        "lg64",
        "lg128",
    ),
    "architecture": {"hidden": [64, 64, 64]},
    "grid_size": (32, 64, 128),
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
}


improved_lgrid_size = {
    "data": "data/Goop",
    "project": "hyperparam_tunning",
    "exp_type": "explicit",
    "model": "unet",
    "run_name": (
        "ilg32",
        "ilg64",
        "ilg128",
    ),
    "architecture": {"hidden": [64, 64, 64]},
    "grid_size": (32, 64, 128),
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
}


WBC = {
    "project": "WBC",
    "exp_type": "explicit",
    "model": "unet",
    "data": "data/WBC-SPH",
    "run_name": ("WBC1", "WBC2", "WBC3", "WBC4", "WBC5", "WBC6"),
    "architecture": (
        {"hidden": [128, 128, 128]},
        {"hidden": [128, 128, 128, 128]},
        {"hidden": [128, 128, 128]},
        {"hidden": [128, 128, 128, 128]},
        {"hidden": [128, 128, 128]},
        {"hidden": [128, 128, 128, 128]},
    ),
    "grid_size": (64, 64, 128, 128, 128, 128),
    "batch_size": (64, 64, 64, 64, 32, 32),
    "epochs": 1000,
    "steps_per_call": (16, 32, 16, 16, 32, 32),
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 4,
    "grid_noise": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    "particle_noise": 0.0,
}


tiny_WBC = {
    "project": "experiments",
    "exp_type": "explicit",
    "model": "unet",
    "data": "data/WBC-SPH",
    "run_name": ("WBC_batch_64", "WBC_batch_32", "WBC_batch_16", "WBC_batch_8"),
    "architecture": ({"hidden": [128, 128, 128, 128]},),
    "grid_size": 128,
    "batch_size": (64, 32, 16, 8),
    "epochs": 1000,
    "steps_per_call": 16,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 4,
    "grid_noise": ([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
    "particle_noise": 0.0,
}


repeat_bundle = {
    "data": "data/Goop",
    "project": "hyperparam_tunning",
    "exp_type": "explicit",
    "model": "unet",
    "run_name": (
        "rb8",
        "rb16",
        "rb32",
    ),
    "architecture": {"hidden": [64, 64, 64]},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": (1000, 1500, 2000),
    "steps_per_call": (8, 16, 32),
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
}


repeat_waterramps = {
    "project": "WaterRamps",
    "exp_type": "explicit",
    "model": "unet",
    "data": "data/WaterRamps",
    "run_name": ("WR1", "WR2", "WR3", "WR4",),
    "architecture": (
        {"hidden": [64, 64, 64]},
        {"hidden": [64, 64, 64, 64]},
        {"hidden": [64, 64, 64]},
        {"hidden": [64, 64, 64, 64]},
    ),
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": (8, 8, 16, 16),
    "autoregressive_steps": (3, 3, 4, 4),
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": (
        [0.001, 0.001, 0.001, 0.0],
        [0.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 0.0],
    ),
    "particle_noise": (0.0003, 0.0, 0.0, 0.0),
}


repeat_WBC = {
    "project": "WBC",
    "exp_type": "explicit",
    "model": "unet",
    "data": "data/WBC-SPH",
    "run_name": ("WBC1", "WBC2", "WBC3", "WBC4", "WBC5", "WBC6"),
    "architecture": (
        {"hidden": [128, 128, 128]},
        {"hidden": [128, 128, 128, 128]},
        {"hidden": [128, 128, 128]},
        {"hidden": [128, 128, 128, 128]},
        {"hidden": [128, 128, 128]},
        {"hidden": [128, 128, 128, 128]},
    ),
    "grid_size": (64, 64, 128, 128, 128, 128),
    "batch_size": (32, 32, 32, 32, 32, 32),
    "epochs": 1500,
    "steps_per_call": (8, 16, 8, 8, 16, 32),
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 4,
    "grid_noise": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    "particle_noise": 0.0,
    "lr": 1e-4,
    "min_lr": 1e-6,
}



fno_noise = {
    "project": "fno",
    "exp_type": "explicit",
    "model": "fno",
    "data": ("data/Goop", "data/WaterRamps", "data/SandRamps"),
    "run_name": ("Goop_noise", "WaterRamps_noise", "SandRamps_noise",),
    "architecture": {"hidden": [64, 64, 64, 64], "modes": None, "use_mlp": True},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 4,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": ([0.001, 0.001, 0.001, 0.0]),
    "particle_noise": (0.001,),
}


fno_no_noise = {
    "project": "fno",
    "exp_type": "explicit",
    "model": "fno",
    "data": ("data/Goop", "data/WaterRamps", "data/SandRamps"),
    "run_name": ("Goop", "WaterRamps", "SandRamps",),
    "architecture": {"hidden": [64, 64, 64, 64], "modes": None, "use_mlp": True},
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 1000,
    "steps_per_call": 8,
    "autoregressive_steps": 4,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": ([0.0, 0.0, 0.0, 0.0],),
    "particle_noise": (0.0,),
}


WR_A5000_1 = {
    "project": "WaterRamps",
    "exp_type": "explicit",
    "model": "unet",
    "data": "data/WaterRamps",
    "run_name": "WRA5000_1",
    "architecture": (
        {"hidden": [64, 64, 64, 64]},
    ),
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 2000,
    "steps_per_call": 8,
    "autoregressive_steps": 4,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": (
        [0.0, 0.0, 0.0, 0.0],
    ),
    "particle_noise": 0.0,
}

WR_A5000_2 = {
    "project": "WaterRamps",
    "exp_type": "explicit",
    "model": "unet",
    "data": "data/WaterRamps",
    "run_name": "WRA5000_2",
    "architecture": (
        {"hidden": [64, 64, 64, 64]},

    ),
    "grid_size": 64,
    "batch_size": 128,
    "epochs": 2000,
    "steps_per_call": 8,
    "autoregressive_steps": 5,
    "passes_over_buffer": 20,
    "sims_in_memory": 16,
    "grid_noise": (
        [0.0, 0.0, 0.0, 0.0],
    ),
    "particle_noise": 0.0,
}

test_WBC = {
    "project": "experiments",
    "exp_type": "explicit",
    "model": "unet",
    "data": "data/WBC-SPH",
    "run_name": "test_WBC",
    "architecture": (
        {"hidden": [128, 128, 128, 128]},
    ),
    "grid_size": 128,
    "batch_size": 32,
    "epochs": 1500,
    "steps_per_call": 8,
    "autoregressive_steps": 3,
    "passes_over_buffer": 20,
    "sims_in_memory": 8,
    "grid_noise": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    "particle_noise": 0.,
    "lr": 1e-4,
    "min_lr": 1e-6,
}


another_WBC = {
    "project": "new_WBC",
    "exp_type": "explicit",
    "model": "unet",
    "data": "data/WBC-SPH",
    "run_name": ("A1", "A2", "A3", "A4", "A5"),
    "architecture": (
        {"hidden": [128, 128, 128, 128]},
        {"hidden": [128, 128, 128, 128]},
        {"hidden": [128, 128, 128, 128]},
        {"hidden": [128, 128, 128, 128, 128]},
        {"hidden": [128, 128, 128, 128, 128]},
    ),
    "grid_size": 128,
    "batch_size": (64, 64, 64, 32, 32),
    "epochs": 5000,
    "steps_per_call": (8, 8, 32, 16, 16),
    "autoregressive_steps": (3, 5, 3, 4, 5),
    "passes_over_buffer": 20,
    "sims_in_memory": 8,
    "grid_noise": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    "particle_noise": (0.0005, 0., 0., 0.,0.),
    "lr": 1e-4,
    "min_lr": 1e-6,
}

debug_exp = {
    "model": "unet",
    "data": "data/WaterRamps",
    "epochs": 10,
    "steps_per_call": 2,
    "autoregressive_steps": 1,
    "grid_size": 64,
    "batch_size": 16,
    "passes_over_buffer": 2,
    "sims_in_memory": 2,
    "architecture": {"hidden": [8, 8]},
    "lr": 1e-3,
    "min_lr": 1e-6,
    "use_schedulers": True,
    "length": 0.015,
    "grid_noise": [0.0] * 4,
    "particle_noise": 0.0,
}