import copy
import json

from typing import Tuple

from toolkit import Configuration

SIZE_X = 400
SIZE_Y = 400
SIZE_Z = 400

def make_config(gpus: int, blockcount: Tuple[int, int, int]) -> Configuration:
    assert gpus % 4 == 0
    assert SIZE_X % blockcount[0] == 0
    assert SIZE_Y % blockcount[1] == 0
    assert SIZE_Z % blockcount[2] == 0
    assert blockcount[0] * blockcount[1] * blockcount[2] == gpus

    return Configuration(
        gpus // 4,
        gpus,
        4,
        blockcount,
        (SIZE_X // blockcount[0], SIZE_Y // blockcount[1], SIZE_Z // blockcount[2])
    )

configurations = [
    Configuration(1, 1, 1, (1, 1, 1), (400, 400, 400)),
    Configuration(1, 2, 2, (1, 1, 2), (400, 400, 200)),
    Configuration(1, 2, 2, (1, 2, 1), (400, 200, 400)),
    Configuration(1, 2, 2, (2, 1, 1), (200, 400, 400)),
    make_config(4, (1, 1, 4)),
    make_config(4, (1, 4, 1)),
    make_config(4, (4, 1, 1)),
    make_config(4, (1, 2, 2)),
    make_config(4, (2, 1, 2)),
    make_config(4, (2, 2, 1)),
    make_config(8, (2, 2, 2)),
    make_config(8, (1, 2, 4)),
    make_config(16, (1, 4, 4)),
    make_config(16, (2, 2, 4)),
    make_config(32, (2, 4, 4)),
    make_config(64, (4, 4, 4)),
    make_config(128, (4, 4, 8)),
    make_config(256, (4, 8, 8)),
    make_config(512, (8, 8, 8)),
    make_config(1024, (8, 8, 16)),
    make_config(2048, (8, 16, 16))
]

with open("templates/spheroid.json") as template_file:
    template = json.load(template_file)

for c in configurations:
    print(c)
    assert(c.get_domain_size() == configurations[0].get_domain_size())

    nastja_config = copy.deepcopy(template)

    nastja_config["Geometry"]["blockcount"] = c.blockcount
    nastja_config["Geometry"]["blocksize"] = c.blocksize

    label = c.get_label()

    with open(f"configs/measurements/strong/spheroid_{label}.json", "w") as config_file:
        json.dump(nastja_config, config_file, indent=2)

    batch_config = f"""#!/usr/bin/env bash

#SBATCH --job-name=strong-{label}
#SBATCH --account=hkf6
#SBATCH --partition=booster
#SBATCH --nodes={c.nodes}
#SBATCH --ntasks={c.tasks}
# Counted per node
#SBATCH --gres=gpu:{c.gpus_per_node}
#SBATCH --time=00:30:00
#SBATCH --output=logs/strong-{label}-%A_%a.log
#SBATCH --error=logs/strong-{label}-%A_%a.log
#SBATCH --array=1-5

SOURCE_DIR=/p/project/cellsinsilico/paulslustigebude
OUTPUT_DIR="/p/scratch/cellsinsilico/paul/nastja-out/strong-{label}-${{SLURM_ARRAY_TASK_ID}}"

echo "outdir is ${{OUTPUT_DIR}}"

mkdir -p "${{OUTPUT_DIR}}"
source "${{SOURCE_DIR}}/activate-nastja-modules"

srun --unbuffered "${{SOURCE_DIR}}/nastja/build-cuda/nastja" \\
  -c "${{SOURCE_DIR}}/ma/experiments/configs/measurements/strong/spheroid_{label}.json" \\
  -o "${{OUTPUT_DIR}}"
"""

    with open(f"batch/measurements/strong/strong-{label}", "w", encoding="utf8") as batch_config_file:
        batch_config_file.write(batch_config)