#!/usr/bin/env python
from abc import ABC, abstractmethod
from dataclasses import dataclass
from parsl.config import Config
from parsl.providers import PBSProProvider
from parsl.executors import HighThroughputExecutor
from parsl.launchers import MpiExecLauncher, GnuParallelLauncher
from parsl.addresses import address_by_interface
from parsl.utils import get_all_checkpoints
from pathlib import Path
from typing import Union

PathLike = Union[str, Path]

@dataclass
class BaseComputeSettings(ABC):
    """Compute settings (HPC platform, number of GPUs, etc)."""

    @abstractmethod
    def config_factory(self, run_dir: PathLike) -> Config:
        """
        Create new Parsl configuration.
        """

class PolarisSettings(BaseComputeSettings):
    num_nodes: int = 1
    worker_init: str = ''
    scheduler_options: str = ''
    account: str
    queue: str
    walltime: str
    cpus_per_node: int = 64
    strategy: str = 'simple'

    def config_factory(self, run_dir: PathLike) -> Config:
        """Create a configuration suitable for running all tasks on single nodes of Polaris
        We will launch 4 workers per node, each pinned to a different GPU
        Args:
            num_nodes: Number of nodes to use for the MPI parallel tasks
            user_options: Options for which account to use, location of environment files, etc
            run_dir: Directory in which to store Parsl run files. Default: `runinfo`
        """
        return Config(
            retries=1,  # Allows restarts if jobs are killed by the end of a job
            executors=[
                HighThroughputExecutor(
                    label=self.label,
                    heartbeat_period=15,
                    heartbeat_threshold=120,
                    worker_debug=True,
                    available_accelerators=4,  # Ensures one worker per accelerator
                    address=address_by_hostname(),
                    cpu_affinity="alternating",
                    prefetch_capacity=0,  # Increase if you have many more tasks than workers
                    start_method="spawn",
                    provider=PBSProProvider(  # type: ignore[no-untyped-call]
                        launcher=MpiExecLauncher(
                            bind_cmd="--cpu-bind", overrides="--depth=64 --ppn 1"
                        ),  # Updates to the mpiexec command
                        account=self.account,
                        queue=self.queue,
                        select_options="ngpus=4",
                        # PBS directives (header lines): for array jobs pass '-J' option
                        scheduler_options=self.scheduler_options,
                        worker_init=self.worker_init,
                        nodes_per_block=self.num_nodes,
                        init_blocks=1,
                        min_blocks=0,
                        max_blocks=1,  # Can increase more to have more parallel jobs
                        cpus_per_node=self.cpus_per_node,
                        walltime=self.walltime,
                    ),
                ),
            ],
            run_dir=str(run_dir),
            strategy=self.strategy,
            app_cache=True,
        )

class AuroraSettings(BaseComputeSettings):
    worker_init: str = ""
    num_nodes: int = 1
    scheduler_options: str = ""
    account: str
    queue: str
    walltime: str
    retries: int = 0
    cpus_per_node: int = 48 # only 4 cpus per OpenMM job
    strategy: str = "simple"

    def config_factory(self, run_dir: PathLike) -> Config:
        """Create a Parsl configuration for running on Aurora."""
        accel_ids = [str(i) for i in range(12)]
        return Config(
            executors=[
                HighThroughputExecutor(
                    label=self.label,
                    available_accelerators=accel_ids,  # Ensures one worker per accelerator
                    cpu_affinity="block",  # Assigns cpus in sequential order
                    prefetch_capacity=0,
                    max_workers=12,
                    cores_per_worker=16,
                    heartbeat_period=30,
                    heartbeat_threshold=300,
                    worker_debug=False,
                    provider=PBSProProvider(
                        launcher=MpiExecLauncher(
                            bind_cmd="--cpu-bind",
                            overrides=f"--depth=208 --ppn 1"
                        ),  # Ensures 1 manger per node and allows it to divide work among all 208 threads
                        worker_init=self.worker_init,
                        nodes_per_block=self.num_nodes,
                        account=self.account,
                        queue=self.queue,
                        walltime=self.walltime,

                    ),
                ),
            ],
            run_dir=str(run_dir),
            checkpoint_mode='task_exit',
            retries=self.retries,
            app_cache=True,
        )


@dataclass
class AuroraCPUSettings:
    worker_init: str = ""
    num_nodes: int = 1
    scheduler_options: str = ""
    account: str
    queue: str
    walltime: str
    retries: int = 0
    cpus_per_node: int = 208
    strategy: str = "simple"

    def config_factory(self, run_dir: PathLike) -> Config:
        """Create a Parsl configuration for running on CPU only on Aurora."""
        return Config(
            executors=[
                HighThroughputExecutor(
                    label='Aurora_CPU',
                    available_accelerators=None,
                    cpu_affinity="block",  # Assigns cpus in sequential order
                    prefetch_capacity=0,
                    max_workers=208,
                    cores_per_worker=1,
                    heartbeat_period=30,
                    heartbeat_threshold=300,
                    worker_debug=False,
                    provider=PBSProProvider(
                        launcher=MpiExecLauncher(
                            bind_cmd="--cpu-bind",
                            overrides=f"--depth=208 --ppn 1"
                        ),  # Ensures 1 manger per node and allows it to divide work among all 208 threads
                        worker_init=self.worker_init,
                        nodes_per_block=self.num_nodes,
                        account=self.account,
                        queue=self.queue,
                        walltime=self.walltime,

                    ),
                ),
            ],
            run_dir=str(run_dir),
            checkpoint_mode='task_exit',
            retries=self.retries,
            app_cache=True,
        )
