import copy
import json

from dataclasses import dataclass
from typing import Tuple

@dataclass
class Configuration:
    nodes: int
    tasks: int
    gpus_per_node: int
    blockcount: Tuple[int, int, int]
    blocksize: Tuple[int, int, int]

    def get_domain_size(self) -> int:
        return self.blockcount[0] * self.blocksize[0] * self.blockcount[1] * self.blocksize[1] * self.blockcount[2] * self.blocksize[2]

configurations = [
    Configuration( 1,   12,   1, ( 1,  4, 3), (400, 100, 134)),
    Configuration( 1,   24,   2, ( 2,  4, 3), (200, 100, 134)),
    Configuration( 1,   48,   4, ( 4,  4, 3), (100, 100, 134)),
    Configuration( 2,   96,   4, ( 4,  4, 6), (100, 100,  67)),
    Configuration( 4,  192,   4, ( 4,  8, 6), (100,  50,  67)),
    Configuration( 8,  384,   4, ( 8,  8, 6), ( 50,  50,  67)),
    Configuration(16,  768,   4, ( 8, 16, 6), ( 50,  25,  67)),
    Configuration(32, 1536,   4, (16, 16, 6), ( 25,  25,  67))
]

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 = f"{c.nodes:02}"

    if c.gpus_per_node < 4:
        label += f"g{c.gpus_per_node}"

    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=06:00: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 "${{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_{c.nodes:02}.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)
