"""Mechanisms to load and balance cells across the computing resources."""

import abc
import hashlib
import logging  # active only in rank 0 (init)
import os
import weakref
from contextlib import contextmanager
from io import StringIO
from pathlib import Path

import numpy as np

from .connection_manager import ConnectionManagerBase
from .core import (
    MPI,
    MComplexLoadBalancer,
    NeuronWrapper as Nd,
    ProgressBarRank0 as ProgressBar,
    mpi_no_errors,
    run_only_rank0,
)
from .core.configuration import (
    ConfigurationError,
    GlobalConfig,
    LoadBalanceMode,
    LogLevel,
    SimConfig,
)
from .core.nodeset import SelectionNodeSet
from .io import cell_readers
from .lfp_manager import LFPManager
from .metype import Cell_V6, EmptyCell, PointCell
from .target_manager import TargetSpec
from .utils import compat
from .utils.logging import log_all, log_verbose
from .utils.memory import DryRunStats, get_mem_usage_kb


class VirtualCellPopulation:
    """A virtual cell population offers a compatible interface with Cell Manager,
    however it doesnt instantiate cells.
    It is mostly used as source of projections
    """

    def __init__(self, population_name, gids=None, circuit_target=None):
        """Initializes a VirtualCellPopulation

        A virtual manager will have minimal set of attributes, namely
        the population name, the target name and the SelectionNodeSet
        """
        self.population_name = population_name
        self.circuit_target = circuit_target
        self.local_nodes = SelectionNodeSet(gids).register_global(population_name)

    is_virtual = property(lambda _self: True)

    def __str__(self):
        return f"([VIRT] {self.population_name:s})"


class _CellManager(abc.ABC):
    @abc.abstractmethod
    def get_final_gids(self): ...

    @abc.abstractmethod
    def get_cell(self, gid): ...

    def get_cellref(self, gid):
        """Retrieve a cell object given its gid.
        Note that this function handles multisplit cases incl converting to an
        spgid automatically \n
        Returns: Cell object
        """
        if self._binfo:
            # are we in load balance mode? raw gids are in the binfo
            gid_offset = self._local_nodes.offset
            gid = self._binfo.thishost_gid(gid - gid_offset) + gid_offset
        return self._pc.gid2obj(gid)

    # Methods for compat with hoc

    @abc.abstractmethod
    def getGidListForProcessor(self): ...

    def getCell(self, gid):
        return self.get_cellref(gid)


class CellManagerBase(_CellManager):
    CellType = NotImplemented  # please override
    """The underlying Cell type class

    signature:
        __init__(self, gid, cell_info, circuit_conf)
    """

    _node_loader = None
    """Default function implementing the loading of nodes data

    signature:
        load(circuit_conf, gidvec, stride=1, stride_offset=0)
    """

    def __init__(self, circuit_conf, target_manager, _run_conf=None, **_kw):
        """Initializes CellDistributor

        Args:
            circuit_conf: The "Circuit" config block
            target_parser: the target parser hoc object, to retrieve target cells' info
            run_conf: Run configuration

        """
        self._circuit_conf = circuit_conf
        self._target_manager = target_manager
        self._target_spec = TargetSpec(circuit_conf.CircuitTarget)
        self._population_name = None
        self._local_nodes = None
        self._total_cells = 0  # total cells in target, being simulated
        self._gid2cell = {}

        self._ionchannel_seed = 0
        self._binfo = None
        self._pc = Nd.pc
        self._conn_managers_per_src_pop = weakref.WeakValueDictionary()

        if isinstance(circuit_conf.CellLibraryFile, str):
            self._init_config(circuit_conf, self._target_spec.population)
        else:
            logging.info(" => %s Circuit has been disabled", self.circuit_name or "(default)")

    # read-only properties
    target_manager = property(lambda self: self._target_manager)
    local_nodes = property(
        lambda self: self._local_nodes if self._local_nodes is not None else SelectionNodeSet()
    )
    total_cells = property(lambda self: self._total_cells)
    cells = property(lambda self: self._gid2cell.values())
    gid2cell = property(lambda self: self._gid2cell)
    pc = property(lambda self: self._pc)
    population_name = property(lambda self: self._population_name)
    circuit_target = property(lambda self: self._target_spec.name)
    circuit_name = property(lambda self: self._circuit_conf.name)
    is_virtual = property(lambda _self: False)
    connection_managers = property(lambda self: self._conn_managers_per_src_pop)

    def is_initialized(self):
        return self._local_nodes is not None

    def __str__(self):
        return f"({self.__class__.__name__}: {self._population_name!s})"

    # Compatibility with neurodamus-core (used by TargetManager, CompMapping)
    # Create hoc vector from numpy.array
    def getGidListForProcessor(self):
        return compat.hoc_vector(self.local_nodes.gids(raw_gids=False))

    def get_final_gids(self):
        return np.array(self.local_nodes.gids(raw_gids=False))

    def _init_config(self, circuit_conf, pop):
        if not pop:  # Last attempt to get pop name
            pop = self._get_sonata_population_name(circuit_conf.CellLibraryFile)
            logging.info(" -> Discovered node population name: %s", pop)
        self._population_name = pop
        self._local_nodes = SelectionNodeSet().register_global(pop)

    @staticmethod
    def _get_sonata_population_name(node_file):
        import libsonata  # only for SONATA

        pop_names = libsonata.NodeStorage(node_file).population_names
        if len(pop_names) != 1:
            raise ConfigurationError(
                "Could not determine population name from sonata file circuit."
            )
        return next(iter(pop_names), None)

    def load_nodes(self, load_balancer=None, *, _loader=None, loader_opts=None):
        """Top-level loader of nodes."""
        if self._local_nodes is None:
            return
        conf = self._circuit_conf
        _loader = _loader or self._node_loader
        cycle = loader_opts.pop("cycle_i", None)
        loader_f = (lambda *args: _loader(*args, **loader_opts)) if loader_opts else _loader

        logging.info("Reading Nodes (METype) info from '%s'", conf.CellLibraryFile)
        if (
            load_balancer
            and hasattr(load_balancer, "population")
            and load_balancer.population != self._target_spec.population
        ):
            log_verbose("Load balance object doesn't apply to '%s'", self._target_spec.population)
            load_balancer = None
        if not load_balancer or SimConfig.dry_run:
            gidvec, me_infos, *cell_counts = self._load_nodes(loader_f)
        elif load_balancer and SimConfig.loadbal_mode == LoadBalanceMode.Memory:
            gidvec, me_infos, *cell_counts = self._load_nodes_balance_mem(
                loader_f, load_balancer, cycle
            )
        else:
            gidvec, me_infos, *cell_counts = self._load_nodes_balance(loader_f, load_balancer)
        self._local_nodes.add_gids(gidvec, me_infos)
        self._total_cells = cell_counts[0]
        logging.info(" => Loaded info about %d target cells (out of %d)", *cell_counts)

    def _load_nodes(self, loader_f):
        """Base loader which handles targets"""
        conf = self._circuit_conf
        target_spec = self._target_spec

        if target_spec.name:
            logging.info(" -> Distributing '%s' target cells Round-Robin", target_spec)
            target_gids = self._target_manager.get_target(target_spec).gids(raw_gids=True)
            gidvec, me_infos, full_size = loader_f(conf, target_gids, MPI.size, MPI.rank)
            total_cells = len(target_gids)
        else:
            logging.info(" -> Distributing ALL cells Round-Robin (No target provided)")
            gidvec, me_infos, full_size = loader_f(conf, None, MPI.size, MPI.rank)
            total_cells = full_size
        return gidvec, me_infos, total_cells, full_size

    def _load_nodes_balance(self, loader_f, load_balancer):
        target_spec = self._target_spec
        if not load_balancer.valid_load_distribution(target_spec):
            raise RuntimeError(
                "No valid Load Balance info could be found or derived."
                "Please perform a full load balance."
            )

        logging.info(" -> Distributing target '%s' using Load-Balance", target_spec.name)
        self._binfo = load_balancer.load_balance_info(target_spec)
        # self._binfo has gidlist, but gids can appear multiple times
        all_gids = np.unique(self._binfo.gids.as_numpy().astype("uint32"))
        total_cells = len(all_gids)
        gidvec, me_infos, full_size = loader_f(self._circuit_conf, all_gids)
        return gidvec, me_infos, total_cells, full_size

    def _load_nodes_balance_mem(self, loader_f, load_balancer, cycle_i):
        targetspec: TargetSpec = self._target_spec

        population = targetspec.population
        if population in load_balancer:
            all_gids = load_balancer.get(population).get((MPI.rank, cycle_i), [])
            all_gids = np.array(all_gids, dtype="uint32")
            total_cells = len(all_gids)
            logging.debug("Loading %d cells in rank %d", total_cells, MPI.rank)
            if total_cells == 0:
                return [], [], 0, 0
            gidvec, me_infos, full_size = loader_f(self._circuit_conf, all_gids)
            return gidvec, me_infos, total_cells, full_size
        return self._load_nodes(loader_f)

    # -
    def finalize(self, **opts):
        """Instantiates cells and initializes the network in the simulator.

        Note: it should be called after all cell distributors have done load_nodes()
            so gids offsets are final.
        """
        if self._local_nodes is None:
            return
        logging.info("Finalizing cells... Gid offset: %d", self._local_nodes.offset)
        self._instantiate_cells(**opts)
        self._update_targets_local_gids()
        self._init_cell_network()
        self._local_nodes.clear_cell_info()

    @mpi_no_errors
    def _instantiate_cells(self, cell_type=None, **_opts):
        cell_type = cell_type or self.CellType
        if SimConfig.crash_test_mode:
            cell_type = PointCell

        assert cell_type is not None, "Undefined cell_type in Manager"
        Nd.execute("xopen_broadcast_ = 0")

        logging.info(" > Instantiating cells... (%d in Rank 0)", len(self._local_nodes))
        cell_offset = self._local_nodes.offset

        if GlobalConfig.verbosity >= LogLevel.DEBUG:
            gid_info_iter = self._local_nodes.iter_cell_info()
        else:
            gid_info_iter = ProgressBar.iter(
                self._local_nodes.iter_cell_info(), len(self._local_nodes)
            )

        for gid, cell_info in gid_info_iter:
            cell = cell_type(gid, cell_info, self._circuit_conf)
            self._store_cell(gid + cell_offset, cell)

    @mpi_no_errors
    def _instantiate_cells_dry(self, cell_type, skip_metypes, **_opts):
        """Instantiates the subset of selected cells while measuring memory taken by each metype

        Args:
            cell_type: The cell type class
            full_memory_counter: The memory counter to be updated for each metype
        """
        assert cell_type is not None, "Undefined cell_type in Manager"
        Nd.execute("xopen_broadcast_ = 0")

        logging.info(" > Dry run on cells... (%d in Rank 0)", len(self._local_nodes))
        cell_offset = self._local_nodes.offset

        prev_metype = None
        prev_memory = get_mem_usage_kb()
        metype_n_cells = 0
        memory_dict = {}
        MAX_CELLS = 50

        def store_metype_stats(metype, n_cells):
            nonlocal prev_memory
            end_memory = get_mem_usage_kb()
            memory_allocated = end_memory - prev_memory
            log_all(
                logging.DEBUG,
                " * METype %s: %.1f KiB averaged over %d cells",
                metype,
                memory_allocated / n_cells,
                n_cells,
            )
            memory_dict[metype] = max(0, memory_allocated / n_cells)
            prev_memory = end_memory

        for gid, cell_info in self._local_nodes.iter_cell_info():
            if cell_info is None:
                continue
            metype = f"{cell_info.mtype}-{cell_info.etype}"
            if metype in skip_metypes:
                continue
            if prev_metype is not None and metype != prev_metype:
                store_metype_stats(prev_metype, metype_n_cells)
                metype_n_cells = 0
            if metype_n_cells >= MAX_CELLS:
                continue
            cell = cell_type(gid, cell_info, self._circuit_conf)
            self._store_cell(gid + cell_offset, cell)
            prev_metype = metype
            metype_n_cells += 1

        if prev_metype is not None and metype_n_cells > 0:
            store_metype_stats(prev_metype, metype_n_cells)

        return memory_dict

    def _update_targets_local_gids(self):
        logging.info(" > Updating targets")
        # Add local gids to matching targets
        self._target_manager.register_local_nodes(self._local_nodes)

    def _init_cell_network(self):
        """Init global gids for cell networking"""
        logging.info(" > Initializing cell network")
        self._init_rng()
        pc = self._pc

        for final_gid, cell in self._gid2cell.items():
            cell.re_init_rng(self._ionchannel_seed)
            nc = cell.connect2target(None)  # Netcon doesnt require being stored
            raw_gid = final_gid - self._local_nodes.offset
            if self._binfo:
                gid_i = int(self._binfo.gids.indwhere("==", raw_gid))
                cb = self._binfo.bilist.object(self._binfo.cbindex.x[gid_i])
                # multisplit cells call cb.multisplit() instead
                if cb.subtrees.count() > 0:
                    cb.multisplit(nc, self._binfo.msgid, pc, pc.id())
                    cell.gid = final_gid
                    cell.raw_gid = raw_gid
                    continue

            pc.set_gid2node(final_gid, pc.id())
            pc.cell(final_gid, nc)
            cell.gid = final_gid  # update the cell.gid last (RNGs had to use the base gid)
            cell.raw_gid = raw_gid

        pc.multisplit()

    def load_artificial_cell(self, gid, artificial_cell):
        logging.info(" > Adding Artificial cell for CoreNeuron")
        cell = EmptyCell(gid, artificial_cell)
        self._store_cell(gid, cell)
        nc = cell.connect2target(None)
        self._pc.set_gid2node(cell.gid, self._pc.id())
        self._pc.cell(cell.gid, nc)

    def _init_rng(self):
        rng_info = Nd.RNGSettings()
        self._ionchannel_seed = rng_info.getIonChannelSeed()
        return rng_info

    def _store_cell(self, gid, cell):
        self._gid2cell[gid] = cell

    def get_cell(self, gid):
        return self._gid2cell[gid]

    def get_cellref(self, gid):
        return self._gid2cell[gid]._cellref

    def record_spikes(self, gids=None, append_spike_vecs=None):
        """Setup recording of spike events (crossing of threshold) for cells on this node"""
        if not self._local_nodes:
            return None
        spikevec, idvec = append_spike_vecs or (Nd.Vector(), Nd.Vector())
        if gids is None:
            gids = self._local_nodes.gids(raw_gids=False)
            gid_offset = self._local_nodes.offset

        for gid in gids:
            # only want to collect spikes of cell pieces with the soma (i.e. the real gid)
            if not self._binfo or self._binfo.thishost_gid(gid - gid_offset) + gid_offset == gid:
                self._pc.spike_record(gid, spikevec, idvec)
        return spikevec, idvec

    def register_connection_manager(self, conn_manager: ConnectionManagerBase):
        src_population = conn_manager.src_node_population
        if src_population in self._conn_managers_per_src_pop:
            logging.warning("Skip registering %s as a second pop source", conn_manager)
        else:
            self._conn_managers_per_src_pop[src_population] = conn_manager

    def pre_stdinit(self):
        """Post stdinit actions"""
        if self._circuit_conf.DetailedAxon:
            log_verbose("Now deleting the axon! 🤩")
            for c in self.cells:
                c.delete_axon()


class GlobalCellManager(_CellManager):
    """GlobalCellManager is a wrapper over all Cell Managers so that we can query
    any cell from its global gid
    """

    def __init__(self):
        self._cell_managers = []
        self._pc = Nd.pc
        self._lfp_manager = LFPManager()

    def register_manager(self, cell_manager):
        self._cell_managers.append(cell_manager)

    def finalize(self):
        self._cell_managers.sort(key=lambda x: x.local_nodes.offset)

    # Accessor methods (Keep CamelCase API for compatibility with existing hoc)
    # ----------------
    def getGidListForProcessor(self):
        def _hoc_append(vec_a, vec_b):
            return vec_a.append(vec_b)

        from functools import reduce

        return reduce(_hoc_append, (man.getGidListForProcessor() for man in self._cell_managers))

    def get_final_gids(self):
        return np.concatenate([man.get_final_gids() for man in self._cell_managers])

    def _find_manager(self, gid):
        cell_managers_iter = iter(self._cell_managers)
        prev_manager = next(cell_managers_iter)  # base cell manager
        for manager in cell_managers_iter:
            if manager.local_nodes.offset > gid:
                break
            prev_manager = manager
        return prev_manager

    def get_cell(self, gid):
        manager = self._find_manager(gid)
        return manager.get_cell(gid)

    def get_cellref(self, gid):
        """Retrieve a cell object given its gid.
        Note that this function handles multisplit cases incl converting to an
        spgid automatically \n
        Returns: Cell object
        """
        manager = self._find_manager(gid)
        if manager._binfo:
            # are we in load balance mode? raw gids are in the binfo
            gid_offset = manager.local_nodes.offset
            gid = manager._binfo.thishost_gid(gid - gid_offset) + gid_offset
        return self._pc.gid2obj(gid)

    def getSpGid(self, gid):
        """Retrieve the spgid from a gid (provided we are using loadbalancing)

        Args:
            gid: The base gid (as read from start.ncs)

        Returns: The gid as it appears on this cpu (if this is the same as the base gid,
        then that is the soma piece)
        """
        manager = self._find_manager(gid)
        if manager._binfo:
            gid_offset = manager.local_nodes.offset
            return manager._binfo.thishost_gid(gid - gid_offset) + gid_offset
        return gid

    def getPopulationInfo(self, gid):
        manager = self._find_manager(gid)
        return manager.population_name, manager.local_nodes.offset


class CellDistributor(CellManagerBase):
    """Manages a group of cells for BBP simulations, V5 and V6

    Instantiated cells are stored locally (.cells property)
    """

    _sonata_with_extra_attrs = True  # Enable search extra node attributes

    def _init_config(self, circuit_conf, _pop):
        if not circuit_conf.CellLibraryFile:
            raise ConfigurationError("CellLibraryFile not set")

        super()._init_config(circuit_conf, _pop)

    def load_nodes(self, load_balancer=None, **kw):
        """Gets gids from target, splits and returns a GidSet with all metadata"""
        loader_opts = kw.pop("loader_opts", {}).copy()
        all_cell_requirements = SimConfig.cell_requirements
        cell_requirements = all_cell_requirements.get(self._population_name)
        loader = cell_readers.load_sonata
        loader_opts["node_population"] = self._population_name  # mandatory in Sonata
        loader_opts["load_dynamic_props"] = cell_requirements
        loader_opts["has_extra_data"] = self._sonata_with_extra_attrs

        log_verbose("Nodes Format: SONATA , Loader: %s", loader.__name__)
        return super().load_nodes(load_balancer, _loader=loader, loader_opts=loader_opts)

    def _instantiate_cells(self, dry_run_stats_obj: DryRunStats = None, **opts):
        """Instantiates cells, honouring dry_run if provided

        Args:
            dry_run_store_stats: Do a dry-run and update the inner fields accordingly
        """
        if self.CellType is not NotImplemented:
            return super()._instantiate_cells(self.CellType)

        conf = self._circuit_conf
        cell_type = Cell_V6
        if conf.MorphologyType:
            cell_type.morpho_extension = conf.MorphologyType

        log_verbose("Loading metypes from: %s", conf.METypePath)
        log_verbose(
            "Loading '%s' morphologies from: %s", cell_type.morpho_extension, conf.MorphologyPath
        )
        if dry_run_stats_obj is None:
            super()._instantiate_cells(cell_type, **opts)
        else:
            cur_metypes_mem = dry_run_stats_obj.metype_memory
            memory_dict = self._instantiate_cells_dry(cell_type, cur_metypes_mem, **opts)
            log_verbose("Updating global dry-run memory counters with %d items", len(memory_dict))
            cur_metypes_mem.update(memory_dict)

        return None


class LoadBalance:
    """Class handling the several types of load_balance info, including
    generating and loading the various files.

    LoadBalance instances target the current system (cpu count) and circuit
    BUT check/create load distribution for any given target.
    The circuit is identified by the nodes file AND population.

    NOTE: Given the heavy costs of computing load balance, some state files are created
    which allow the balance info to be reused. These are

     - cx_{TARGET}.dat: File with complexity information for the cells of a given target
     - cx_{TARGET}.{CPU_COUNT}.dat: The file assigning cells/pieces to individual CPUs ranks.

    For more information refer to the developer documentation.

    :param balance_mode: Mode of balancing the load, as defined by
        :class:`neurodamus.core.configuration.LoadBalanceMode`.
    :param nodes_path: Path to the nodes file.
    :param pop: Population identifier for the circuit.
    :param target_manager: Manages target information for load balancing.
    :param target_cpu_count: Optional number of CPUs for the target.
    """

    _base_output_dir = "sim_conf"
    _circuit_lb_dir_tpl = "_loadbal_%s.%s"  # Placeholders are (file_src_hash, population)
    _cx_filename_tpl = "cx_%s#.dat"  # use # to well delimiter the target name
    _cpu_assign_filename_tpl = "cx_%s#.%s.dat"  # prefix must be same (imposed by Neuron)

    def __init__(self, balance_mode, nodes_path, pop, target_manager, target_cpu_count=None):
        """Creates a new Load Balance object, associated with a given node file"""
        self.lb_mode = balance_mode
        self.target_cpu_count = target_cpu_count or MPI.size
        self._target_manager = target_manager
        self._valid_loadbalance = set()
        self.population = pop or ""
        self._lb_dir, self._cx_targets = self._get_circuit_loadbal_dir(nodes_path, self.population)
        log_verbose("Found existing targets with loadbal: %s", self._cx_targets)

    @classmethod
    @run_only_rank0
    def _get_circuit_loadbal_dir(cls, node_file, pop) -> tuple:
        """Ensure lbal dir exists. dir may be crated on rank 0"""
        lb_dir = cls._loadbal_dir(node_file, pop)
        if lb_dir.is_dir():
            return lb_dir, cls._get_lbdir_targets(lb_dir)

        logging.info("Creating load balance dir %s", lb_dir)
        os.makedirs(lb_dir)
        return lb_dir, set()

    @classmethod
    def _get_lbdir_targets(cls, lb_dir: Path) -> list:
        """Inspects the load-balance folder and detects which targets are load balanced"""
        prefix, suffix = cls._cx_filename_tpl.split("%s")
        return {
            fname.name[len(prefix) : -len(suffix)]
            for fname in lb_dir.glob(cls._cx_filename_tpl.replace("%s", "*"))
        }

    @run_only_rank0
    def valid_load_distribution(self, target_spec: TargetSpec) -> bool:
        """Checks whether we have valid load-balance files, attempting to
        derive from larger target distributions if possible.
        """
        if (target_spec.population) != self.population:
            logging.info(
                " => Load balance Population mismatch. Requested: %s, Existing: %s",
                target_spec.population,
                self.population,
            )
            return False

        target_name = target_spec.simple_name

        # Check cache
        if target_name in self._valid_loadbalance:
            logging.info(" => Cell distribution from Load Balance is valid")
            return True

        # A valid load-distribution requires BOTH the cx and the cpu_assign files.
        if self._cx_valid(target_spec):
            cpu_assign_filename = self._cpu_assign_filename(target_name)
            if cpu_assign_filename.is_file():
                logging.info(" => Found valid load balance: %s", cpu_assign_filename)
            else:
                # Still need to derive the CPU assignment
                logging.info(" => Found valid complexity file: %s", self._cx_filename(target_name))
                self._cpu_assign(target_name)
            self._valid_loadbalance.add(target_name)
            return True

        # If not found, attempt to find a cx file computed for a superset target
        # so we can derive a new cx file without instantiating cells
        if self._reuse_cell_complexity(target_spec):
            self._cpu_assign(target_name)
            self._valid_loadbalance.add(target_name)
            return True
        return False

    # -
    def _reuse_cell_complexity(self, target_spec: TargetSpec) -> bool:
        """Check if the complexities of all target gids were already calculated
        for another target.
        """
        # Abort if there are no cx files yet or in case now we request full circuit
        # since its impossible to have a superset of it
        if not target_spec.name or not self._cx_targets:
            logging.info(" => Target Cx reusing is not available.")
            return False

        logging.info("Attempt reusing cx files from other targets...")
        target_gids = self._get_target_raw_gids(target_spec)
        cx_other = {}

        for previous_target in self._cx_targets:
            log_verbose("Trying previous cx file on target %s", previous_target)
            ref_filename = self._cx_filename(previous_target)
            if self._cx_contains_gids(ref_filename, target_gids, cx_other):
                break  # done!
            log_verbose("  - Target is not a superset. Ignoring.")
        else:
            logging.info(" => Did not find any suitable target")
            return False

        new_cx_filename = self._cx_filename(target_spec.simple_name)
        logging.info(
            "Target %s is a subset of the target %s. Generating %s",
            target_spec.name,
            previous_target,
            new_cx_filename,
        )

        # Write the new cx file since Neuron needs it to do CPU assignment
        with open(new_cx_filename, "w", encoding="utf-8") as newfile:
            self._write_msdat_dict(newfile, cx_other, target_gids)
        # register
        self._cx_targets.add(target_spec.simple_name)
        return True

    # -
    def _cx_valid(self, target_spec) -> bool:
        """Determine if valid complexity files exist for the provided circuit and
        target spec, otherwise return False -> we need to regen load balance info
        """
        # _cx_targets populated at instantiation. If not there, no chance. If in list, verify
        if not self._cx_targets:
            logging.info(" => No complexity files for current circuit yet")
            return False

        target_name = target_spec.simple_name
        cx_filename = self._cx_filename(target_name)

        if target_name not in self._cx_targets:
            logging.info(" => No Cx files available for requested target")
            return False

        if target_spec:  # target provided, otherwise everything
            target_gids = self._get_target_raw_gids(target_spec)
            if not self._cx_contains_gids(cx_filename, target_gids):
                logging.warning(" => %s invalid: changed target definition!", cx_filename)
                return False
        return True

    @classmethod
    def _cx_contains_gids(cls, cxpath, target_gids, out_cx=None) -> bool:
        """Checks a cx file contains complexities for given gids"""
        if not cxpath.is_file():
            log_verbose("  - cxpath doesnt exist: %s", cxpath)
            return False
        with open(cxpath, encoding="utf-8") as f:
            cx_saved = cls._read_msdat(f)
        if not set(cx_saved.keys()) >= set(target_gids):
            log_verbose("  - Not all GIDs in target %s %s", set(cx_saved.keys()), set(target_gids))
            return False
        if out_cx is not None:
            out_cx.update(cx_saved)
        return True

    @contextmanager
    def generate_load_balance(self, target_spec, cell_distributor):
        """Context manager that creates load balance for the circuit instantiated within

        Args:
            target_str: a string representation of the target.
            cell_distributor: the cell distributor object to which we can query
                the cells to be load balanced
        """
        mcomplex = MComplexLoadBalancer()  # init mcomplex before building circuit
        yield
        target_str = target_spec.simple_name
        self._compute_save_complexities(target_str, mcomplex, cell_distributor)
        self._cpu_assign(target_str)
        self._valid_loadbalance.add(target_str)

    @mpi_no_errors
    def _compute_save_complexities(self, target_str, mcomplex, cell_distributor):
        msfactor = 1e6 if self.lb_mode == LoadBalanceMode.WholeCell else 0.8
        out_filename = self._cx_filename(target_str)

        cx_cells = self._compute_complexities(mcomplex, cell_distributor)
        total_cx, max_cx = self._cell_complexity_total_max(cx_cells)
        lcx = self._get_optimal_piece_complexity(total_cx, self.target_cpu_count, msfactor)
        logging.info(
            "LB Info: TC=%.3f MC=%.3f OptimalCx=%.3f FileName=%s",
            total_cx,
            max_cx,
            lcx,
            out_filename,
        )

        ms_list = []
        tmp = Nd.Vector()

        for cell in cell_distributor.cells:
            mcomplex.cell_complexity(cell.CellRef)
            mcomplex.multisplit(cell.raw_gid, lcx, tmp)
            ms_list.append(tmp.c())

        # To output build independently the contents of the file then append
        ostring = StringIO()
        for ms in ms_list:
            self._write_msdat(ostring, ms)

        all_ranks_cx = MPI.py_gather(ostring.getvalue(), 0)
        if MPI.rank == 0:
            with open(out_filename, "w", encoding="utf-8") as fp:
                fp.write(f"1\n{cell_distributor.total_cells}\n")
                fp.writelines(all_ranks_cx)

        # register
        self._cx_targets.add(target_str)

    @staticmethod
    def _cell_complexity_total_max(cx_cells):
        """Returns: Tuple of (TotalComplexity, max_complexity)"""
        local_max = max(cx_cells) if len(cx_cells) > 0 else 0.0
        local_sum = sum(cx_cells) if len(cx_cells) > 0 else 0.0

        global_total = MPI.allreduce(local_sum, MPI.SUM)
        global_max = MPI.allreduce(local_max, MPI.MAX)
        return global_total, global_max

    @classmethod
    def _compute_complexities(cls, mcomplex, cell_distributor):
        cx_cell = compat.Vector("f")
        pc = cell_distributor.pc
        for gid in cell_distributor.local_nodes.gids(raw_gids=False):
            cx_cell.append(mcomplex.cell_complexity(pc.gid2cell(gid)))
        return cx_cell

    @staticmethod
    def _get_optimal_piece_complexity(total_cx, nhost, msfactor):
        """Args:
        total_cx: Total complexity
        nhost: Prospective no of hosts
        """
        lps = total_cx * msfactor / nhost
        return int(lps + 1)

    @run_only_rank0
    def _cpu_assign(self, target_name):
        """Assigns cells to 'prospective_hosts' cpus using mymetis3.
        Results are written to file. basename.<NCPU>.dat
        """
        logging.info("Assigning Cells <-> %d CPUs [mymetis3]", self.target_cpu_count)
        base_filename = self._cx_filename(target_name, basename_str=True)
        Nd.mymetis3(base_filename, self.target_cpu_count)

    @staticmethod
    def _write_msdat(fp, ms):
        """Writes load balancing info to an output stream"""
        fp.write(str(int(ms.x[0])))  # gid
        fp.write(f" {ms.x[1]:g}")  # total complexity of cell
        piece_count = int(ms.x[2])
        fp.write(f" {piece_count}\n")
        i = 2

        for _ in range(piece_count):
            i += 1
            subtree_count = int(ms.x[i])
            fp.write(f"  {subtree_count}\n")
            for _ in range(subtree_count):
                i += 1
                cx = ms.x[i]  # subtree complexity
                i += 1
                children_count = int(ms.x[i])
                fp.write(f"   {cx:g} {children_count}\n")
                if children_count > 0:
                    fp.write("    ")
                for _ in range(children_count):
                    i += 1
                    elem_id = int(ms.x[i])  # at next child
                    fp.write(f" {elem_id}")
                if children_count > 0:
                    fp.write("\n")

    @staticmethod
    def _read_msdat(fp):
        """Read load balancing info from an input stream"""
        cx_saved = {}  # dict with key = gid, value = line content
        piece_count = 0
        gid = None
        next(fp)  # first line has 1. skip
        next(fp)  # ncells. skip

        for line in fp:
            if piece_count == 0:
                gid, _cx, piece_count = [int(float(x)) for x in line.split()]
                cx_saved[gid] = [line]
            else:  # Handle parts
                cx_saved[gid].append(line)
                for _ in range(2 * int(line)):  # each subtree has two lines
                    cx_saved[gid].append(next(fp))
                piece_count -= 1

        return cx_saved

    @staticmethod
    def _write_msdat_dict(fp, cx_dict, gids=None):
        """Write out selected gid cx lines from a cx_dict"""
        if gids is None:
            gids = cx_dict.keys()
        fp.write(f"1\n{len(gids)}\n")
        for gid in gids:
            for line in cx_dict[gid]:
                fp.write(line)  # raw lines, include \n

    # -
    def _get_target_raw_gids(self, target_spec) -> np.ndarray:
        return self._target_manager.get_target(target_spec).gids(raw_gids=True)

    def load_balance_info(self, target_spec):
        """Loads a load-balance info for a given target.
        NOTE: Please ensure the load balance exists or is derived before calling this function
        """
        bal_filename = self._cx_filename(target_spec.simple_name, basename_str=True)
        return Nd.BalanceInfo(bal_filename, MPI.rank, MPI.size)

    @classmethod
    def _loadbal_dir(cls, nodefile, population) -> Path:
        """Returns the dir where load balance files are stored for a given nodes file"""
        nodefile_hash = hashlib.md5(nodefile.encode()).hexdigest()[:10]
        return Path(cls._base_output_dir) / (cls._circuit_lb_dir_tpl % (nodefile_hash, population))

    def _cx_filename(self, target_str, basename_str=False) -> Path:
        """Gets the filename of a cell complexity file for a given target"""
        fname = self._lb_dir / (self._cx_filename_tpl % target_str)
        return str(fname)[:-4] if basename_str else fname

    def _cpu_assign_filename(self, target_str) -> Path:
        """Gets the CPU assignment filename for a given target, according to target CPU count"""
        return self._lb_dir / (self._cpu_assign_filename_tpl % (target_str, self.target_cpu_count))

    @staticmethod
    def select_lb_mode(sim_config, run_conf, target):
        """A method which selects the load balance mode according to run config"""
        # Check / set load balance mode
        lb_mode = sim_config.loadbal_mode
        if lb_mode == LoadBalanceMode.MultiSplit:
            if not sim_config.use_coreneuron:
                logging.info("Load Balancing ENABLED. Mode: MultiSplit")
            else:
                logging.warning("Load Balancing mode CHANGED to WholeCell for CoreNeuron")
                lb_mode = LoadBalanceMode.WholeCell

        elif lb_mode == LoadBalanceMode.WholeCell:
            logging.info("Load Balancing ENABLED. Mode: WholeCell")

        elif lb_mode is None:
            if target.is_void():
                lb_mode, reason = LoadBalanceMode.RoundRobin, "No target set, unknown cell count"
            else:
                lb_mode, reason = LoadBalanceMode.auto_select(
                    sim_config.use_neuron, target.gid_count(), run_conf["Duration"]
                )
            logging.warning("Load Balance AUTO-SELECTED: %s. Reason: %s", lb_mode.name, reason)

        return lb_mode
