'''
===============================================================================
MOOSE Herd Class

Authors: Lloyd Fletcher, Rory Spencer
===============================================================================
'''
import os
import time
import multiprocessing as mp
from pathlib import Path
from multiprocessing.pool import Pool

from mooseherder.directorymanager import DirectoryManager
from mooseherder.simrunner import SimRunner
from mooseherder.inputmodifier import InputModifier


class MooseHerdError(Exception):
    """MooseHerdError: custom error class for flagging errors with the moose
    herd.
    """

class MooseHerd:
    """ MooseHerd class that can run parametric sweeps of simulation chains in
    parallel with configurable parallelisation options. Takes a list of
    SimRunner objects and a corresponding list of InputModifiers to insert the
    variables into the input scripts for the SimRunners. When calling run_* the
    herd will first call all InputModifiers in the specified order and then
    call run on all the SimRunners in order. Uses the DirectoryManager class to
    create/clear and log the directories in which each parallel worker is
    creating input files and running simulations.
    """
    def __init__(self, sim_runners: list[SimRunner],
                 input_mods: list[InputModifier],
                 dir_manager: DirectoryManager) -> None:
        """__init__

        Args:
            sim_runners (list[SimRunner]): list of objects that inherit from
                the SimRunner ABC in the order they need to be run. The mesh
                needs to be created before runnning moose so a common chain
                would be [GmshRunner,MooseRunner].
            input_mods (list[InputModifier]): list of InputModifiers to create
                the required input scripts for the SimRunners.
            dir_manager (DirectoryManager): used to control how many and
                which directories are used to runthe simulations.
        """
        self._runners = sim_runners
        self._modifiers = input_mods
        self._dir_manager = dir_manager

        if len(self._runners) != len(self._modifiers):
            raise MooseHerdError('The sim runner list and the input modifier '+
                                 'list must be the same length')

        self._n_para_sims = 2

        self._input_names = [f'sim-{ii+1}' for ii,_ in enumerate(sim_runners)]

        self._keep_all = True

        self._var_sweep = list([])

        self._sweep_iter = 0
        self._sweep_run_time = -1.0

        self._sim_iter = 0
        self._iter_run_time = -1.0


    def set_input_copy_names(self, input_names: list[str] | None = None) -> None:
        """set_input_copy_name: sets the name that will be used when copying
        input files to the working directories for the sweep. The defualt name
        is 'sim-i' so the first combination of variables in the simulation chain
        will be called 'sim-1-1'.

        Args:
            input_names (list[str] | None, optional): List of name prefixes to
                be used for the simulation files. Defaults to None.

        Raises:
            MooseHerdError: The lengths of the sim runner list and the input
                modifier lists are not the same.
        """
        if input_names is None:
            self._input_names = [f'sim-{ii+1}' for ii,_ in enumerate(self._runners)]
            return

        if len(input_names) != len(self._runners):
            raise MooseHerdError(f'The length of the input names ({len(input_names)})'
                                 'must match the length of the sim runners ' +
                                 f'and input modifiers ({len(self._runners)})')

        self._input_names = input_names


    def set_keep_flag(self, keep_all: bool = True) -> None:
        """set_keep_flag: flag used for allowing multiple calls to run_para or
        run_seq to keep everything or to overwrite with every call to run_*.

        Args:
            keep_all (bool, optional): True = keep all inputs and outputs with
                multiple calls to run_*. False = overwrite inputs and outputs
                with multiple calls to run_*. Defaults to True.
        """
        self._keep_all = keep_all


    def set_num_para_sims(self, n_para: int = 1) -> None:
        """set_num_para_sims: sets the number of simulation chains to run in
        parallel. Limits the number

        Args:
            n_para (int, optional): _description_. Defaults to 1.
        """
        n_para = int(n_para)
        n_cpus = os.cpu_count()

        if n_cpus is None:
            pass
        elif n_para <= 0:
            n_para = 1
        elif n_para > n_cpus:
            n_para = n_cpus

        if self._n_para_sims != n_para:
            self._n_para_sims = n_para


    def get_sim_iter(self) -> int:
        """get_sim_iter: returns the current simulation iteration corresponding
        to the combination of variables being analysed. This number will
        accumulate with multiple calls to run_* is keep_all=true.

        Returns:
            int: current simulation iteration number.
        """
        return self._sim_iter


    def get_sweep_iter(self) -> int:
        """get_sweep_iter: gets the current sweep iteration. The sweep
        iteration is incremented with every call to run_* if keep_all = true.
        If keep_all = false then sweep_iter is held at 1.

        Returns:
            int: current sweep iteration number.
        """
        return self._sweep_iter


    def reset_iter_counts(self) -> None:
        """reset_iter_counts: resets the simulation iteration and the sweep
        iteration counters to zero.
        """
        self._sim_iter = 0
        self._sweep_iter = 0


    def _get_process_name(self) -> str:
        """_get_process_name: only here for monkey patching with pytest.

        Returns:
            str: the process name string.
        """
        return mp.current_process().name


    def _get_worker_num(self) -> str:
        """_get_worker_num: helper function to get the worker number for the
        current sub-process.

        Returns:
            str: number string taken from the process name. If this is the main
                process returns '1'.
        """
        name = self._get_process_name()

        if name == 'MainProcess':
            worker_num = '1'
        else:
            worker_num = name.split('-',1)[1]

        if int(worker_num) > self._n_para_sims:
            worker_num = str((int(worker_num) % self._n_para_sims)+1)

        return worker_num


    def _get_run_num(self, sim_iter: int, worker_num: str) -> str:
        """_get_run_num: helper function to get the run directory number string

        Args:
            sim_iter (int): the current simulation iteration.
            worker_num (str): the worker number extracted from the current
            process number as a string.

        Returns:
            str: the string specifying the run directory number for this
                simulation iteration.
        """
        if self._keep_all:
            run_num = str(sim_iter+1)
        else:
            run_num = worker_num

        return run_num


    def _mod_input(self,
                   modifier: InputModifier,
                   mod_vars: dict | None,
                   save_file: Path) -> None:
        """_mod_input: helper function that uses the input modifier to write
        new variables to the input file and save it to the specified path.

        Args:
            modifier (InputModifier): input modifier for the specified type of
                input file.
            mod_vars (dict | None): dictionary of variables to write to the
                input file, if None just copy the input file.
            save_file (Path): path with file name and extension to output the
                modified input file.
        """
        if mod_vars is not None:
            modifier.update_vars(mod_vars)
        modifier.write_file(save_file)


    def _run(self, runner: SimRunner, run_file: Path) -> Path | None:
        """_run: helper function to call the SimRunner and get the path to the
        output file.

        Args:
            runner (SimRunner): for running the simulation, must be a class
                that implements the SimRunner ABC.
            run_file (Path): path to the input file to run with SimRunner.

        Returns:
            Path | None: _description_
        """
        runner.run(run_file)
        return runner.get_output_path()


    def run_once(self, sim_iter: int, var_list: list[dict | None]
                 ) -> list[Path | None]:
        """run_once: runs a specific simulation chain with the given variable
        list once and returns a list of paths to the output files. Used by
        run_seq and run_para for parallelisation.

        Args:
            sim_iter (int): current simulation iteration which is the index of
                the var_list from the var_sweep.
            var_list (list[dict  |  None]): list of dictionaries that contain
                the variables that will be run for this iteration.

        Returns:
            list[Path | None]: list of paths to the simulation output. If there
                is no useful output from the runner in the simulation chain it
                returns None in the list.
        """
        iter_start_time = time.perf_counter()

        worker_num = self._get_worker_num()
        run_dir = self._dir_manager.get_run_dir(int(worker_num)-1)
        run_num = self._get_run_num(sim_iter,worker_num)

        run_files = list([])
        for ii,mm in enumerate(self._modifiers):
            ext = mm.get_input_file().suffix
            run_files.append(run_dir / (self._input_names[ii] +'-'+run_num+ext))
            self._mod_input(mm,var_list[ii],run_files[ii])

        output_list = list([])
        for ii,rr in enumerate(self._runners):
            output_list.append(self._run(rr,run_files[ii]))

        self._iter_run_time = time.perf_counter() - iter_start_time

        return output_list


    def _start_sweep(self, var_sweep: list[list[dict | None]]) -> float:
        """_start_sweep: helper function used at the start of a variable sweep
        in either run_seq or run_para. Sets the var_sweep attribute, deals with
        the management of directories and starts the performance counter.

        Args:
            var_sweep (list[list[dict  |  None]]): as passed to run_seq/para

        Returns:
            float: performance timer start value.
        """
        self._var_sweep = var_sweep

        if not self._keep_all:
            self.reset_iter_counts()
            self._dir_manager.clear_dirs()
            self._dir_manager.create_dirs()

        return time.perf_counter()


    def _end_sweep(self, start_sweep_time: float,
                   output_files: list[list[Path | None]]) -> None:
        """_end_sweep: helper function called at the end of runseq/para.
        Reacords the sweep run time. Increments the iteration counters. and
        writes the output key and sweep variables to the first workers
        directory.

        Args:
            start_sweep_time (float): the sweep start time taken from the
                _start_sweep() function.
            output_files (list[list[Path]]): list of list of paths to the
                simulation chain output files.
        """
        self._sweep_run_time = time.perf_counter() - start_sweep_time

        self._sweep_iter += 1
        self._sim_iter += len(self._var_sweep)

        self._dir_manager.set_output_paths(output_files)
        self._dir_manager.write_output_key(self._sweep_iter)
        self._dir_manager.write_sweep_vars(self._var_sweep,self._sweep_iter)


    def run_sequential(self, var_sweep: list[list[dict | None]]
                       ) -> list[list[Path | None]]:
        """run_sequential: runs the variable sweep given in var_sweep
        sequentially and returns the paths to the simulation outputs.

        Args:
            var_sweep (list[list[dict | None]]): outer list is the simulation
                iteration, inner list is the position in the simulation chain
                that the variable dictionary corresponds to. The dictionary
                contains the variables that will be inserted into the input
                file before calling run on the SimRunner. If None instead of
                a dictionary then the input file is copied with no modification

        Returns:
            list[list[Path | None]]: outer list is the simulation iteration and
                the inner list corresponds to the position of the SimRunner in
                the cimulation chain. Gives the path to the simulation output
                or None if no useful output is produced.
        """
        start_sweep_time = self._start_sweep(var_sweep)

        output_files = list([])

        ii = self._sim_iter
        for vv in var_sweep:
            output_files.append(self.run_once(ii,vv))
            ii += 1


        self._end_sweep(start_sweep_time,output_files)

        return output_files


    def run_para(self, var_sweep: list[list[dict | None]]
                 ) -> list[list[Path | None]]:
        """run_para: runs the variable sweep with the simulation chain in
        parallel.

        Args:
            var_sweep (list[list[dict | None]]): outer list is the simulation
                iteration, inner list is the position in the simulation chain
                that the variable dictionary corresponds to. The dictionary
                contains the variables that will be inserted into the input
                file before calling run on the SimRunner. If None instead of
                a dictionary then the input file is copied with no modification

        Returns:
            list[list[Path | None]]: outer list is the simulation iteration and
                the inner list corresponds to the position of the SimRunner in
                the cimulation chain. Gives the path to the simulation output
                or None if no useful output is produced.
        """
        sweep_start_time = self._start_sweep(var_sweep)

        with Pool(self._n_para_sims) as pool:
            processes = list([])

            ii = self._sim_iter
            for vv in var_sweep:
                processes.append(pool.apply_async(self.run_once, args=(ii,vv)))
                ii += 1

            output_files = [pp.get() for pp in processes]

        self._end_sweep(sweep_start_time, output_files)

        return output_files


    def get_sweep_time(self) -> float:
        """get_sweep_time

        Returns:
            float: the time taken for the variable sweep to run based on the
                performance counter.
        """
        return self._sweep_run_time


    def get_iter_time(self) -> float:
        """get_iter_time

        Returns
            float: the time taken for the current simulation iteration to run.
        """
        return self._iter_run_time

