"""
Module containing the basic class for handling particles properties.
"""

from copy import deepcopy
from numpy import arange, ceil, empty, floor, int64
from numpy import load as np_load
from numpy import loadtxt, meshgrid, ndarray, pi, sqrt, triu_indices, zeros
from numpy.random import Generator, PCG64
from os.path import join
from scipy.linalg import norm
from scipy.spatial.distance import pdist
from warnings import warn

from .utilities.exceptions import ParticlesError, ParticlesWarning

class Particles:
    """
    Class handling particles' properties.

    Attributes
    ----------
    kB : float
        Boltzmann constant.

    fourpie0: float
        Electrostatic constant :math:`4\\pi \\epsilon_0`.

    pos : numpy.ndarray
        Particles' positions.

    vel : numpy.ndarray
        Particles' velocities.

    acc : numpy.ndarray
        Particles' accelerations.

    box_lengths : numpy.ndarray
        Box sides' lengths.

    pbox_lengths : numpy.ndarray
        Initial particle box sides' lengths.

    masses : numpy.ndarray
        Mass of each particle. Shape = (attr:`sarkas.core.Parameters.total_num_ptcls`).

    charges : numpy.ndarray
        Charge of each particle. Shape = (attr:`sarkas.core.Parameters.total_num_ptcls`).

    id : numpy.ndarray,
        Species identifier. Shape = (attr:`sarkas.core.Parameters.total_num_ptcls`).

    names : numpy.ndarray
        Species' names. (attr:`sarkas.core.Parameters.total_num_ptcls`).

    rdf_nbins : int
        Number of bins for radial pair distribution.

    no_grs : int
        Number of independent :math:`g_{ij}(r)`.

    rdf_hist : numpy.ndarray
        Histogram array for the radial pair distribution function.

    prod_dump_dir : str
        Directory name where to store production phase's simulation's checkpoints. Default = 'dumps'.

    eq_dump_dir : str
        Directory name where to store equilibration phase's simulation's checkpoints. Default = 'dumps'.

    total_num_ptcls : int
        Total number of simulation's particles.

    num_species : int
        Number of species.

    species_num : numpy.ndarray
        Number of particles of each species. Shape = (attr:`sarkas.particles.Particles.num_species`).

    dimensions : int
        Number of non-zero dimensions. Default = 3.

    potential_energy : float
        Instantaneous value of the potential energy.

    rnd_gen : numpy.random.Generator
        Random number generator.

    """

    def __init__(self):
        self.mag_dump_dir = None
        self.rdf_nbins = None
        self.potential_energy = 0.0
        self.kB = None
        self.fourpie0 = None
        self.prod_dump_dir = None
        self.eq_dump_dir = None
        self.box_lengths = None
        self.pbox_lengths = None
        self.total_num_ptcls = None
        self.num_species = 1
        self.species_num = None
        self.dimensions = None
        self.rnd_gen = None

        self.pos = None
        self.vel = None
        self.acc = None

        self.virial = None
        self.pbc_cntr = None

        self.names = None
        self.id = None

        self.species_initial_velocity = None
        self.species_thermal_velocity = None

        self.masses = None
        self.charges = None
        self.cyclotron_frequencies = None

        self.no_grs = None
        self.rdf_hist = None

    def __repr__(self):
        sortedDict = dict(sorted(self.__dict__.items(), key=lambda x: x[0].lower()))
        disp = "Particles( \n"
        for key, value in sortedDict.items():
            disp += "\t{} : {}\n".format(key, value)
        disp += ")"
        return disp

    def __copy__(self):
        """
        Make a shallow copy of the object using copy by creating a new instance of the object and copying its __dict__."""
        # Create a new object
        _copy = type(self)()
        # copy the dictionary
        _copy.__dict__.update(self.__dict__)
        return _copy

    def __deepcopy__(self, memodict: dict = {}):
        """Make a deepcopy of the object.

        Parameters
        ----------
        memodict: dict
            Dictionary of id's to copies

        Returns
        -------
        _copy: :class:`sarkas.particles.Particles`
            A new Particles class.
        """
        id_self = id(self)  # memorization avoids unnecessary recursion
        _copy = memodict.get(id_self)
        if _copy is None:

            # Make a shallow copy of all attributes
            _copy = type(self)()
            # Make a deepcopy of the mutable arrays using numpy copy function
            for k, v in self.__dict__.items():
                if isinstance(v, ndarray):
                    _copy.__dict__[k] = v.copy()
                else:
                    _copy.__dict__[k] = deepcopy(v, memodict)

        return _copy

    def __getstate__(self):
        """Copy the object's state from self.__dict__ which contains all our instance attributes.
        Always use the dict.copy() method to avoid modifying the original state.
        Reference: https://docs.python.org/3/library/pickle.html#handling-stateful-objects
        """

        state = self.__dict__.copy()
        # Remove the data that is stored already
        del state["pos"]
        del state["vel"]
        del state["acc"]
        del state["id"]
        del state["names"]
        del state["pbc_cntr"]
        del state["rdf_hist"]
        del state["virial"]

        return state

    def __setstate__(self, state):
        # Restore instance attributes.
        self.__dict__.update(state)
        # Initialize arrays
        self.pos = zeros((self.__dict__["total_num_ptcls"], 3))
        self.vel = zeros((self.__dict__["total_num_ptcls"], 3))
        self.acc = zeros((self.__dict__["total_num_ptcls"], 3))
        self.id = zeros(self.__dict__["total_num_ptcls"])
        self.names = zeros(self.__dict__["total_num_ptcls"])
        self.pbc_cntr = zeros((self.__dict__["total_num_ptcls"], 3))
        self.rdf_hist = zeros((self.__dict__["rdf_nbins"], self.__dict__["num_species"], self.__dict__["num_species"]))
        self.virial = zeros((self.__dict__["dimensions"], self.__dict__["dimensions"], self.__dict__["total_num_ptcls"]))

    def copy_params(self, params):
        """
        Copy necessary parameters.

        Parameters
        ----------
        params: :class:`sarkas.core.Parameters`
            Simulation's parameters.

        """

        self.kB = params.kB
        self.fourpie0 = params.fourpie0
        self.prod_dump_dir = params.prod_dump_dir
        self.eq_dump_dir = params.eq_dump_dir
        self.box_lengths = params.box_lengths.copy()
        self.pbox_lengths = params.pbox_lengths.copy()
        self.total_num_ptcls = params.total_num_ptcls
        self.total_num_density = params.total_num_density
        self.num_species = params.num_species
        self.species_num = params.species_num.copy()
        self.dimensions = params.dimensions
        self.load_method = params.load_method
        self.restart_step = params.restart_step
        self.particles_input_file = params.particles_input_file
        self.load_perturb = params.load_perturb
        self.load_rejection_radius = params.load_rejection_radius
        self.load_halton_bases = params.load_halton_bases

        if hasattr(params, "np_per_side"):
            self.np_per_side = params.np_per_side

        if hasattr(params, "initial_lattice_config"):
            self.lattice_type = params.initial_lattice_config

        if hasattr(params, "load_gauss_sigma"):
            self.load_gauss_sigma = params.load_gauss_sigma.copy()

        self.species_names = params.species_names

        if hasattr(params, "rdf_nbins"):
            self.rdf_nbins = params.rdf_nbins
        else:
            # nbins = 5% of the number of particles.
            self.rdf_nbins = int(0.05 * params.total_num_ptcls)
            params.rdf_nbins = self.rdf_nbins

    def gaussian(self, mean, sigma, size):
        """
        Initialize particles' velocities according to a normalized Maxwell-Boltzmann (Normal) distribution.
        It calls ``numpy.random.Generator.normal``

        Parameters
        ----------
        size : tuple
            Size of the array to initialize. (no. of particles, dimensions).

        mean : float
            Center of the normal distribution.

        sigma : float
            Scale of the normal distribution.

        Returns
        -------
         : numpy.ndarray
            Particles property distributed according to a Normal probability density function.

        """
        return self.rnd_gen.normal(mean, sigma, size)

    def halton_reject(self, bases, r_reject):
        """
        Place particles according to a Halton sequence from 0 to LP (the initial particle box length)
        and uses a rejection radius to avoid placing particles to close to each other.

        Parameters
        ----------
        bases : numpy.ndarray
            Array of 3 ints each of which is a base for the Halton sequence.
            Defualt: bases = array([2,3,5])

        r_reject : float
            Value of rejection radius.

        """

        # Get bases
        b1, b2, b3 = bases

        # Allocate space and store first value from Halton
        x = zeros(self.total_num_ptcls)
        y = zeros(self.total_num_ptcls)
        z = zeros(self.total_num_ptcls)

        # Initialize particle counter and Halton counter
        i = 1
        k = 1

        # Loop over all particles
        while i < self.total_num_ptcls:

            # Increment particle counter
            n = k
            m = k
            p = k

            # Determine x coordinate
            f1 = 1
            r1 = 0
            while n > 0:
                f1 /= b1
                r1 += f1 * (n % int(b1))
                n = floor(n / b1)
            x_new = self.pbox_lengths[0] * r1  # new x value

            # Determine y coordinate
            f2 = 1
            r2 = 0
            while m > 0:
                f2 /= b2
                r2 += f2 * (m % int(b2))
                m = floor(m / b2)
            y_new = self.pbox_lengths[1] * r2  # new y value

            # Determine z coordinate
            f3 = 1
            r3 = 0
            while p > 0:
                f3 /= b3
                r3 += f3 * (p % int(b3))
                p = floor(p / b3)
            z_new = self.pbox_lengths[2] * r3  # new z value

            # Check if particle was place too close relative to all other current particles
            for j in range(len(x)):

                # Flag for if particle is outside of cutoff radius (1 -> not inside rejection radius)
                flag = 1

                # Compute distance b/t particles for initial placement
                x_diff = x_new - x[j]
                y_diff = y_new - y[j]
                z_diff = z_new - z[j]

                # Periodic condition applied for minimum image
                if x_diff < -self.pbox_lengths[0] / 2:
                    x_diff = x_diff + self.pbox_lengths[0]
                if x_diff > self.pbox_lengths[0] / 2:
                    x_diff = x_diff - self.pbox_lengths[0]

                if y_diff < -self.pbox_lengths[1] / 2:
                    y_diff = y_diff + self.pbox_lengths[1]
                if y_diff > self.pbox_lengths[1] / 2:
                    y_diff = y_diff - self.pbox_lengths[1]

                if z_diff < -self.pbox_lengths[2] / 2:
                    z_diff = z_diff + self.pbox_lengths[2]
                if z_diff > self.pbox_lengths[2] / 2:
                    z_diff = z_diff - self.pbox_lengths[2]

                # Compute distance
                r = sqrt(x_diff**2 + y_diff**2 + z_diff**2)

                # Check if new particle is below rejection radius. If not, break out and try again
                if r <= r_reject:
                    k += 1  # Increment Halton counter
                    flag = 0  # New position not added (0 -> no longer outside reject r)
                    break

            # If flag true add new position
            if flag == 1:
                # Add new positions to arrays
                x[i] = x_new
                y[i] = y_new
                z[i] = z_new

                k += 1  # Increment Halton counter
                i += 1  # Increment particle number

        self.pos[:, 0] = x + self.box_lengths[0] / 2 - self.pbox_lengths[0] / 2
        self.pos[:, 1] = y + self.box_lengths[1] / 2 - self.pbox_lengths[1] / 2
        self.pos[:, 2] = z + self.box_lengths[2] / 2 - self.pbox_lengths[2] / 2

    def initialize_accelerations(self):
        """
        Initialize particles' accelerations.
        """
        self.acc = zeros((self.total_num_ptcls, 3))

    def initialize_arrays(self):
        """Initialize the needed arrays"""
        self.pos = zeros((self.total_num_ptcls, 3))
        self.vel = zeros((self.total_num_ptcls, 3))
        self.acc = zeros((self.total_num_ptcls, 3))

        self.pbc_cntr = zeros((self.total_num_ptcls, 3))
        self.virial = zeros((3, 3, self.total_num_ptcls))

        self.names = empty(self.total_num_ptcls, dtype=self.species_names.dtype)
        self.id = zeros(self.total_num_ptcls, dtype=int64)

        self.species_initial_velocity = zeros((self.num_species, 3))
        self.species_thermal_velocity = zeros((self.num_species, 3))

        self.masses = zeros(self.total_num_ptcls)  # mass of each particle
        self.charges = zeros(self.total_num_ptcls)  # charge of each particle
        self.cyclotron_frequencies = zeros(self.total_num_ptcls)

        # No. of independent rdf
        self.no_grs = int(self.num_species * (self.num_species + 1) / 2)

        self.rdf_hist = zeros((self.rdf_nbins, self.num_species, self.num_species))

    def initialize_positions(self):
        """
        Initialize particles' positions based on the load method.
        """
        # Particles Position Initialization
        if self.load_method in [
            "equilibration_restart",
            "eq_restart",
            "magnetization_restart",
            "mag_restart",
            "production_restart",
            "prod_restart",
        ]:
            # checks
            if self.restart_step is None:
                raise AttributeError("Restart step not defined." "Please define Parameters.restart_step.")

            if type(self.restart_step) is not int:
                self.restart_step = int(self.restart_step)

            if self.load_method[:2] == "eq":
                self.load_from_restart("equilibration", self.restart_step)
            elif self.load_method[:2] == "pr":
                self.load_from_restart("production", self.restart_step)
            elif self.load_method[:2] == "ma":
                self.load_from_restart("magnetization", self.restart_step)

        elif self.load_method == "file":
            # check
            if not hasattr(self, "particles_input_file"):
                raise AttributeError("Input file not defined." "Please define Parameters.particles_input_file.")
            self.load_from_file(self.particles_input_file)

        # position distribution.
        elif self.load_method == "lattice":
            self.lattice(self.load_perturb)

        elif self.load_method == "random_reject":
            # check
            if not hasattr(self, "load_rejection_radius"):
                raise AttributeError("Rejection radius not defined. " "Please define Parameters.load_rejection_radius.")
            self.random_reject(self.load_rejection_radius)

        elif self.load_method == "halton_reject":
            # check
            if not hasattr(self, "load_rejection_radius"):
                raise AttributeError("Rejection radius not defined. " "Please define Parameters.load_rejection_radius.")
            self.halton_reject(self.load_halton_bases, self.load_rejection_radius)

        elif self.load_method in ["uniform", "random_no_reject"]:
            self.pos = self.uniform_no_reject(
                0.5 * self.box_lengths - 0.5 * self.pbox_lengths, 0.5 * self.box_lengths + 0.5 * self.pbox_lengths
            )

        elif self.load_method == "gaussian":
            sp_start = 0
            sp_end = 0
            for sp, sp_num in enumerate(self.species_num):
                sp_end += sp_num
                self.pos[sp_start:sp_end, :] = self.gaussian(
                    self.box_lengths[0] / 2.0, self.load_gauss_sigma[sp], (sp_num, 3)
                )
                sp_start += sp_num
        else:
            raise AttributeError("Incorrect particle placement scheme specified.")

    def initialize_velocities(self, species):
        """
        Initialize particles' velocities based on the species input values. The velocities can be initialized from a
        Maxwell-Boltzmann distribution or from a monochromatic distribution.

        Parameters
        ----------
        species: list
            List of :class:`sarkas.core.Spcies`.

        """
        species_end = 0
        species_start = 0
        for ic, sp in enumerate(species):
            if sp.name != "electron_background":
                species_end += sp.num
                self.species_initial_velocity[ic, :] = sp.initial_velocity

                if sp.initial_velocity_distribution == "boltzmann":
                    if isinstance(sp.temperature, (int, float)):
                        sp_temperature = zeros(3)
                        for d in range(self.dimensions):
                            sp_temperature[d] = sp.temperature

                    self.species_thermal_velocity[ic] = sqrt(self.dimensions * self.kB * sp_temperature / (2.0 * sp.mass))
                    # Note gaussian(0.0, 0.0, N) = array of zeros
                    self.vel[species_start:species_end, :] = self.gaussian(
                        sp.initial_velocity, self.species_thermal_velocity[ic], (sp.num, 3)
                    )

                elif sp.initial_velocity_distribution == "monochromatic":
                    vrms = sqrt(self.dimensions * self.kB * sp.temperature / sp.mass)
                    self.vel[species_start:species_end, :] = vrms * self.random_unit_vectors(sp.num, self.dimensions)

                species_start += sp.num

    def kinetic_temperature(self):
        """
        Calculate the kinetic energy and temperature of each species.

        Returns
        -------
        K : numpy.ndarray
            Kinetic energy of each species. Shape=(``num_species``).

        T : numpy.ndarray
            Temperature of each species. Shape=(``num_species``).

        """
        K = zeros(self.num_species)
        T = zeros(self.num_species)
        const = 2.0 / (self.kB * self.species_num * self.dimensions)
        kinetic = 0.5 * self.masses * (self.vel * self.vel).transpose()

        species_start = 0
        species_end = 0
        for i, num in enumerate(self.species_num):
            species_end += num
            K[i] = kinetic[:, species_start:species_end].sum()
            T[i] = const[i] * K[i]
            species_start = species_end

        return K, T

    def lattice(self, perturb):
        """
        Place particles in a simple cubic lattice with a slight perturbation ranging
        from 0 to 0.5 times the lattice spacing.

        Parameters
        ----------
        perturb : float
            Value of perturbation, p, such that 0 <= p <= 1.

        """

        # Check if perturbation is below maximum allowed. If not, default to maximum perturbation.
        if perturb > 1:
            warn("Random perturbation must not exceed 1. Setting perturb = 1.", category=ParticlesWarning)

        if self.lattice_type == "simple_cubic":
            # Determining number of particles per side of simple cubic lattice
            part_per_side = self.total_num_ptcls ** (1.0 / 3.0)  # Number of particles per side of cubic lattice

            # Check if total number of particles is a perfect cube, if not, place more than the requested amount
            if round(part_per_side) ** 3 != self.total_num_ptcls:
                part_per_side = ceil(self.total_num_ptcls ** (1.0 / 3.0))
                raise ParticlesError(
                    f"N = {self.total_num_ptcls} cannot be placed in a simple cubic lattice. "
                    f"Use {int(part_per_side ** 3)} particles instead."
                )

            dx_lattice = self.pbox_lengths[0] / (self.total_num_ptcls ** (1.0 / 3.0))  # Lattice spacing
            dy_lattice = self.pbox_lengths[1] / (self.total_num_ptcls ** (1.0 / 3.0))  # Lattice spacing
            dz_lattice = self.pbox_lengths[2] / (self.total_num_ptcls ** (1.0 / 3.0))  # Lattice spacing

            # Create x, y, and z position arrays
            x = arange(0, self.pbox_lengths[0], dx_lattice) + 0.5 * dx_lattice
            y = arange(0, self.pbox_lengths[1], dy_lattice) + 0.5 * dy_lattice
            z = arange(0, self.pbox_lengths[2], dz_lattice) + 0.5 * dz_lattice

            # Create a lattice with appropriate x, y, and z values based on arange
            X, Y, Z = meshgrid(x, y, z)

            # Perturb lattice
            X += self.rnd_gen.uniform(-0.5, 0.5, X.shape) * perturb * dx_lattice
            Y += self.rnd_gen.uniform(-0.5, 0.5, Y.shape) * perturb * dy_lattice
            Z += self.rnd_gen.uniform(-0.5, 0.5, Z.shape) * perturb * dz_lattice

            # Flatten the meshgrid values for plotting and computation
            self.pos[:, 0] = X.ravel() + self.box_lengths[0] / 2 - self.pbox_lengths[0] / 2
            self.pos[:, 1] = Y.ravel() + self.box_lengths[1] / 2 - self.pbox_lengths[1] / 2
            self.pos[:, 2] = Z.ravel() + self.box_lengths[2] / 2 - self.pbox_lengths[2] / 2

        elif self.lattice_type in ["square", "tetragonal_2D"]:
            # Determining number of particles per side of simple cubic lattice
            part_per_side = round(sqrt(self.total_num_ptcls))  # Number of particles per side of a square lattice

            # Check if total number of particles is a perfect cube, if not, place more than the requested amount
            if part_per_side**2 != self.total_num_ptcls:
                raise ParticlesError(
                    f"N = {self.total_num_ptcls} cannot be placed in a square lattice. "
                    f"Use {int(part_per_side ** 2)} particles instead."
                )

            dx_lattice = self.pbox_lengths[0] / sqrt(self.total_num_ptcls)  # Lattice spacing
            dy_lattice = self.pbox_lengths[1] / sqrt(self.total_num_ptcls)  # Lattice spacing

            # Create x, y, and z position arrays
            x = arange(0, self.pbox_lengths[0], dx_lattice) + 0.5 * dx_lattice
            y = arange(0, self.pbox_lengths[1], dy_lattice) + 0.5 * dy_lattice

            # Create a lattice with appropriate x, y, and z values based on arange
            X, Y = meshgrid(x, y)

            # Perturb lattice
            X += self.rnd_gen.uniform(-0.5, 0.5, X.shape) * perturb * dx_lattice
            Y += self.rnd_gen.uniform(-0.5, 0.5, Y.shape) * perturb * dy_lattice

            # Flatten the meshgrid values for plotting and computation
            self.pos[:, 0] = X.ravel() + self.box_lengths[0] / 2 - self.pbox_lengths[0] / 2
            self.pos[:, 1] = Y.ravel() + self.box_lengths[1] / 2 - self.pbox_lengths[1] / 2
            self.pos[:, 2] = 0.0

        elif self.lattice_type in ["hexagonal", "triangular"]:

            # Determining number of particles per side of simple cubic lattice
            part_per_side = round(sqrt(self.total_num_ptcls))  # Number of particles per side of cubic lattice

            # Check if total number of particles is a perfect cube, if not, place more than the requested amount
            if self.np_per_side[:2].prod() != part_per_side * (part_per_side + 1):
                raise ParticlesError(
                    f"N = {self.total_num_ptcls} cannot be placed in an hexagonal lattice. "
                    f"Use Nx = {part_per_side} and Ny = {part_per_side + 1} particles instead."
                )

            dx_lattice = self.pbox_lengths[0] / (self.np_per_side[0])  # Lattice spacing
            dy_lattice = self.pbox_lengths[1] / (self.np_per_side[1])  # Lattice spacing

            if self.np_per_side[0] > self.np_per_side[1]:
                # Create x, y, and z position arrays
                x = arange(0, self.pbox_lengths[0], dx_lattice)
                y = arange(0, self.pbox_lengths[1], dy_lattice) + 0.5 * dy_lattice

                # Create a lattice with appropriate x, y, and z values based on arange
                X, Y = meshgrid(x, y)
                # Shift the Y axis of every other row of particles
                X[:, ::2] += 0.5 * dx_lattice

            else:
                # Create x, y, and z position arrays
                x = arange(0, self.pbox_lengths[0], dx_lattice) + 0.5 * dx_lattice
                y = arange(0, self.pbox_lengths[1], dy_lattice)

                # Create a lattice with appropriate x, y, and z values based on arange
                X, Y = meshgrid(x, y)
                # Shift the Y axis of every other row of particles
                Y[:, ::2] += 0.5 * dy_lattice

            # Perturb lattice
            X += self.rnd_gen.uniform(-0.5, 0.5, X.shape) * perturb * dx_lattice
            Y += self.rnd_gen.uniform(-0.5, 0.5, Y.shape) * perturb * dy_lattice

            # Flatten the meshgrid values for plotting and computation
            self.pos[:, 0] = X.ravel() + self.box_lengths[0] / 2 - self.pbox_lengths[0] / 2
            self.pos[:, 1] = Y.ravel() + self.box_lengths[1] / 2 - self.pbox_lengths[1] / 2
            self.pos[:, 2] = 0.0

    def load(self):
        """
        Initialize particles' positions and velocities.
        Positions are initialized based on the load method while velocities are chosen
        from a Maxwell-Boltzmann distribution.

        """

        warn(
            "Deprecated feature. It will be removed in the v2.0.0 release. \n"
            "Use parameters.calc_electron_properties(species). You need to pass the species list.",
            category=DeprecationWarning,
        )

        self.initialize_positions()

    def load_from_file(self, f_name):
        """
        Load particles' data from a specific file.

        Parameters
        ----------
        f_name : str
            Filename
        """
        pv_data = loadtxt(f_name)
        if not (pv_data.shape[0] == self.total_num_ptcls):
            msg = (
                f"Number of particles is not same between input file and initial p & v data file. \n "
                f"Input file: N = {self.total_num_ptcls}, load data: N = {pv_data.shape[0]}"
            )
            raise ParticlesError(msg)

        self.pos[:, 0] = pv_data[:, 0]
        self.pos[:, 1] = pv_data[:, 1]
        self.pos[:, 2] = pv_data[:, 2]

        self.vel[:, 0] = pv_data[:, 3]
        self.vel[:, 1] = pv_data[:, 4]
        self.vel[:, 2] = pv_data[:, 5]

    def load_from_restart(self, phase, it):
        """
        Load particles' data from a checkpoint of a previous run

        Parameters
        ----------
        it : int
            Timestep.

        phase: str
            Restart phase.

        """
        if phase == "equilibration":
            file_name = join(self.eq_dump_dir, "checkpoint_" + str(it) + ".npz")
            data = np_load(file_name, allow_pickle=True)
            self.id = data["id"]
            self.names = data["names"]
            self.pos = data["pos"]
            self.vel = data["vel"]
            self.acc = data["acc"]

        elif phase == "production":
            file_name = join(self.prod_dump_dir, "checkpoint_" + str(it) + ".npz")
            data = np_load(file_name, allow_pickle=True)
            self.id = data["id"]
            self.names = data["names"]
            self.pos = data["pos"]
            self.vel = data["vel"]
            self.acc = data["acc"]
            self.pbc_cntr = data["cntr"]
            self.rdf_hist = data["rdf_hist"]

        elif phase == "magnetization":
            file_name = join(self.mag_dump_dir, "checkpoint_" + str(it) + ".npz")
            data = np_load(file_name, allow_pickle=True)
            self.id = data["id"]
            self.names = data["names"]
            self.pos = data["pos"]
            self.vel = data["vel"]
            self.acc = data["acc"]
            self.pbc_cntr = data["cntr"]
            self.rdf_hist = data["rdf_hist"]

    def potential_energies(self):
        """
        Calculate the potential energies of each species.

        Returns
        -------
        P : numpy.ndarray
            Potential energy of each species. Shape=(``num_species``).

        """
        P = zeros(self.num_species)

        species_start = 0
        species_end = 0
        for i, num in enumerate(self.species_num):
            species_end += num

            # TODO: Consider writing a numba function speedup in distance calculation
            species_charges = self.charges[species_start:species_end]
            uti = triu_indices(species_charges.size, k=1)
            species_charge2 = species_charges[uti[0]] * species_charges[uti[1]]
            species_distances = pdist(self.pos[species_start:species_end, :])
            potential = species_charge2 / self.fourpie0 / species_distances
            P[i] = potential.sum()

            species_start = species_end

        return P

    def random_reject(self, r_reject):
        """
        Place particles by sampling a uniform distribution from 0 to LP (the initial particle box length)
        and uses a rejection radius to avoid placing particles to close to each other.

        Parameters
        ----------
        r_reject : float
            Value of rejection radius.
        """

        # Initialize Arrays
        x = zeros(self.total_num_ptcls)
        y = zeros(self.total_num_ptcls)
        z = zeros(self.total_num_ptcls)

        # Set first x, y, and z positions
        x_new = self.rnd_gen.uniform(0, self.pbox_lengths[0])
        y_new = self.rnd_gen.uniform(0, self.pbox_lengths[1])
        z_new = self.rnd_gen.uniform(0, self.pbox_lengths[2])

        # Append to arrays
        x[0] = x_new
        y[0] = y_new
        z[0] = z_new

        # Particle counter
        i = 1

        cntr_reject = 0
        cntr_total = 0
        # Loop to place particles
        while i < self.total_num_ptcls:

            # Set x, y, and z positions
            x_new = self.rnd_gen.uniform(0.0, self.pbox_lengths[0])
            y_new = self.rnd_gen.uniform(0.0, self.pbox_lengths[1])
            z_new = self.rnd_gen.uniform(0.0, self.pbox_lengths[2])

            # Check if particle was place too close relative to all other current particles
            for j in range(len(x)):

                # Flag for if particle is outside of cutoff radius (True -> not inside rejection radius)
                flag = 1

                # Compute distance b/t particles for initial placement
                x_diff = x_new - x[j]
                y_diff = y_new - y[j]
                z_diff = z_new - z[j]

                # periodic condition applied for minimum image
                if x_diff < -self.pbox_lengths[0] / 2:
                    x_diff += self.pbox_lengths[0]
                if x_diff > self.pbox_lengths[0] / 2:
                    x_diff -= self.pbox_lengths[0]

                if y_diff < -self.pbox_lengths[1] / 2:
                    y_diff += self.pbox_lengths[1]
                if y_diff > self.pbox_lengths[1] / 2:
                    y_diff -= self.pbox_lengths[1]

                if z_diff < -self.pbox_lengths[2] / 2:
                    z_diff += self.pbox_lengths[2]
                if z_diff > self.pbox_lengths[2] / 2:
                    z_diff -= self.pbox_lengths[2]

                # Compute distance
                r = sqrt(x_diff**2 + y_diff**2 + z_diff**2)

                # Check if new particle is below rejection radius. If not, break out and try again
                if r <= r_reject:
                    flag = 0  # new position not added (False -> no longer outside reject r)
                    cntr_reject += 1
                    cntr_total += 1
                    break

            # If flag true add new position
            if flag == 1:
                x[i] = x_new
                y[i] = y_new
                z[i] = z_new

                # Increment particle number
                i += 1
                cntr_total += 1

        self.pos[:, 0] = x + self.box_lengths[0] / 2 - self.pbox_lengths[0] / 2
        self.pos[:, 1] = y + self.box_lengths[1] / 2 - self.pbox_lengths[1] / 2
        self.pos[:, 2] = z + self.box_lengths[2] / 2 - self.pbox_lengths[2] / 2

    def random_unit_vectors(self, num_ptcls, dimensions):
        """
        Initialize random unit vectors for particles' velocities (e.g. for monochromatic energies but random velocities)
        It calls ``numpy.random.Generator.normal``

        Parameters
        ----------
        num_ptcls : int
            Number of particles to initialize.

        dimensions : int
            Number of non-zero dimensions.

        Returns
        -------
        uvec : numpy.ndarray
            Random unit vectors of specified dimensions for all particles

        """

        uvec = self.rnd_gen.normal(size=(num_ptcls, dimensions))
        # Broadcasting
        uvec /= norm(uvec).reshape(num_ptcls, 1)

        return uvec

    def remove_drift(self):
        """
        Enforce conservation of total linear momentum. Updates particles velocities
        """
        species_start = 0
        species_end = 0
        momentum = self.masses * self.vel.transpose()
        for ic, nums in enumerate(self.species_num):
            species_end += nums
            P = momentum[:, species_start:species_end].sum(axis=1)
            self.vel[species_start:species_end, :] -= P / (nums * self.masses[species_end - 1])
            species_start = species_end

    def setup(self, params, species):
        """
        Initialize class' attributes

        Parameters
        ----------
        params: :class:`sarkas.core.Parameters`
            Simulation's parameters.

        species : list
            List of :meth:`sarkas.plasma.Species` objects.

        """

        if hasattr(params, "rand_seed"):
            self.rand_seed = params.rand_seed
            self.rnd_gen = Generator(PCG64(params.rand_seed))
        else:
            self.rnd_gen = Generator(PCG64())

        self.copy_params(params)
        self.initialize_arrays()
        self.update_attributes(species)
        self.initialize_positions()
        self.initialize_velocities(species)
        self.initialize_accelerations()

    def uniform_no_reject(self, mins, maxs):
        """
        Randomly distribute particles along each direction.

        Parameters
        ----------
        mins : float
            Minimum value of the range of a uniform distribution.

        maxs : float
            Maximum value of the range of a uniform distribution.

        Returns
        -------
         : numpy.ndarray
            Particles' property, e.g. pos, vel. Shape = (``total_num_ptcls``, 3).

        """

        return self.rnd_gen.uniform(mins, maxs, (self.total_num_ptcls, 3))

    def update_attributes(self, species):
        """
        Assign particles attributes.

        Parameters
        ----------
        species : list
            List of :class:`sarkas.plasma.Species` objects.

        """
        species_end = 0
        species_start = 0

        for ic, sp in enumerate(species):
            if sp.name != "electron_background":
                species_end += sp.num

                self.names[species_start:species_end] = sp.name
                self.masses[species_start:species_end] = sp.mass

                if hasattr(sp, "charge"):
                    self.charges[species_start:species_end] = sp.charge
                else:
                    self.charges[species_start:species_end] = 1.0

                if hasattr(sp, "cyclotron_frequency"):
                    self.cyclotron_frequencies[species_start:species_end] = sp.cyclotron_frequency

                self.id[species_start:species_end] = ic
                species_start += sp.num


import datetime
import logging
import os
import numpy as np
from multiprocessing import Pool
from functools import partial
import us
import pickle
import simplejson as json
import copy
from collections import defaultdict
from pyseir.models.seir_model import SEIRModel
from pyseir.parameters.parameter_ensemble_generator import ParameterEnsembleGenerator
from pyseir.models.suppression_policies import generate_empirical_distancing_policy, generate_covidactnow_scenarios
from pyseir import load_data
from pyseir.reports.county_report import CountyReport
from pyseir.utils import get_run_artifact_path, RunArtifact, RunMode
from pyseir.load_data import FAULTY_HOSPITAL_DATA_STATES
from libs.datasets.dataset_utils import AggregationLevel
from libs.datasets import CovidTrackingDataSource
from libs.datasets import JHUDataset

_logger = logging.getLogger(__name__)

compartment_to_capacity_attr_map = {
    'HGen': 'beds_general',
    'HICU': 'beds_ICU',
    'HVent': 'ventilators'
}

class EnsembleRunner:
    """
    The EnsembleRunner executes a collection of N_samples simulations based on
    priors defined in the ParameterEnsembleGenerator.

    Parameters
    ----------
    fips: str
        County or state fips code
    n_years: int
        Number of years to simulate
    n_samples: int
        Ensemble size to run for each suppression policy.
    suppression_policy: list(float or str)
        List of suppression policies to apply.
    output_percentiles: list
        List of output percentiles desired. These will be computed for each
        compartment.
    run_mode: str
        Individual parameters can be overridden here.
    min_hospitalization_threshold: int
        Require this number of hospitalizations before initializing based on
        observations. Fallback to cases otherwise.
    hospitalization_to_confirmed_case_ratio: float
        When hospitalization data is not available directly, this fraction of
        confirmed cases defines the initial number of hospitalizations.
    covid_timeseries: NoneType or DataSet
        Can be optionally passed in to prevent reloading.
    """
    def __init__(self, fips, n_years=.5, n_samples=250,
                 suppression_policy=(0.35, 0.5, 0.75, 1),
                 skip_plots=False,
                 output_percentiles=(5, 25, 32, 50, 75, 68, 95),
                 generate_report=True,
                 run_mode=RunMode.DEFAULT,
                 min_hospitalization_threshold=5,
                 hospitalization_to_confirmed_case_ratio=1 / 4,
                 covid_timeseries=None):

        self.fips = fips
        self.agg_level = AggregationLevel.COUNTY if len(fips) == 5 else AggregationLevel.STATE

        self.t_list = np.linspace(0, int(365 * n_years), int(365 * n_years))
        self.skip_plots = skip_plots
        self.run_mode = RunMode(run_mode)
        self.hospitalizations_for_state = None
        self.min_hospitalization_threshold = min_hospitalization_threshold
        self.hospitalization_to_confirmed_case_ratio = hospitalization_to_confirmed_case_ratio

        if self.agg_level is AggregationLevel.COUNTY:
            self.county_metadata = load_data.load_county_metadata_by_fips(fips)
            self.state_abbr = us.states.lookup(self.county_metadata['state']).abbr
            self.state_name = us.states.lookup(self.county_metadata['state']).name

            self.output_file_report = get_run_artifact_path(self.fips, RunArtifact.ENSEMBLE_REPORT)
            self.output_file_data = get_run_artifact_path(self.fips, RunArtifact.ENSEMBLE_RESULT)

        else:
            self.state_abbr = us.states.lookup(self.fips).abbr
            self.state_name = us.states.lookup(self.fips).name

            self.output_file_report = None
            self.output_file_data = get_run_artifact_path(self.fips, RunArtifact.ENSEMBLE_RESULT)

        county_fips = None if self.agg_level is AggregationLevel.STATE else self.fips

        if not covid_timeseries:
            covid_timeseries = JHUDataset.local().timeseries()
        else:
            covid_timeseries = covid_timeseries.timeseries()

        self.covid_data = covid_timeseries\
            .get_subset(self.agg_level, country='USA', state=self.state_abbr) \
            .get_data(country='USA', state=self.state_abbr, fips=county_fips) \
            .sort_values('date')

        os.makedirs(os.path.dirname(self.output_file_data), exist_ok=True)
        if self.output_file_report:
            os.makedirs(os.path.dirname(self.output_file_report), exist_ok=True)

        self.output_percentiles = output_percentiles
        self.n_samples = n_samples
        self.n_years = n_years
        # TODO: Will be soon replaced with loaders for all the inferred params.
        # self.t0 = fit_results.load_t0(fips)
        self.date_generated = datetime.datetime.utcnow().isoformat()
        self.suppression_policy = suppression_policy
        self.summary = copy.deepcopy(self.__dict__)
        self.summary.pop('t_list')
        self.generate_report = generate_report

        self.suppression_policies = None
        self.override_params = dict()
        self.init_run_mode()

        self.all_outputs = {}

    def get_initial_hospitalizations(self, use_cases=False):
        """
        Attempt a two level hierarchy of lookups for hospitalizations.

        1. Direct hospitalizations if available
        2. Inferred from covid case data.

        Returns
        -------
        latest_date: date
            Last date of data available.
        hospitalizations_total: int
            Estimated number of current hospitalizations total.
        """
        fips = None if self.agg_level is AggregationLevel.STATE else self.fips

        hospitalization_data = CovidTrackingDataSource.local()\
            .timeseries()\
            .get_subset(self.agg_level, country='USA', state=self.state_abbr)\
            .get_data(state=self.state_abbr, country='USA', fips=fips)\
            .sort_values('date')

        # If there are enough hospitalizations, use those to define initial conditions.
        if not use_cases and len(hospitalization_data) > 0 and self.state_abbr not in FAULTY_HOSPITAL_DATA_STATES:
            latest_date = hospitalization_data.iloc[-1]['date'].date()
            n_current = hospitalization_data.iloc[-1]['current_hospitalized']
            if n_current > self.min_hospitalization_threshold and not np.isnan(n_current):
                hospitalizations_total = n_current

            # TODO: We will need a better estimator for current hospitalizations
            # in cases where cumulative is not available. Punting on this until
            # post-release.
            else:
                hospitalizations_total = hospitalization_data.iloc[-1]['cumulative_hospitalized']

        # Fallback to case data if not.
        else:
            latest_date = self.covid_data.date.max()
            hospitalizations_total = self.covid_data.cases.max() * self.hospitalization_to_confirmed_case_ratio
        return latest_date, hospitalizations_total

    def init_run_mode(self):
        """
        Based on the run mode, generate suppression policies and ensemble
        parameters.  This enables different model combinations and project
        phases.
        """
        self.suppression_policies = dict()

        if self.run_mode is RunMode.CAN_BEFORE_HOSPITALIZATION:
            self.n_samples = 1

            for scenario in ['no_intervention', 'flatten_the_curve', 'full_containment', 'social_distancing']:
                R0 = 3.6
                self.override_params['R0'] = R0
                policy = generate_covidactnow_scenarios(t_list=self.t_list, R0=R0, t0=datetime.datetime.today(), scenario=scenario)
                self.suppression_policies[f'suppression_policy__{scenario}'] = policy
                self.override_params = ParameterEnsembleGenerator(
                    self.fips, N_samples=500, t_list=self.t_list, suppression_policy=policy).get_average_seir_parameters()

            self.override_params['mortality_rate_no_general_beds'] = 0.0
            self.override_params['mortality_rate_from_hospital'] = 0.0
            self.override_params['mortality_rate_from_ICU'] = 0.40
            self.override_params['mortality_rate_no_ICU_beds'] = 1.0

            self.override_params['hospitalization_length_of_stay_general'] = 6
            self.override_params['hospitalization_length_of_stay_icu'] = 13
            self.override_params['hospitalization_length_of_stay_icu_and_ventilator'] = 14

            self.override_params['hospitalization_rate_general'] = 0.0727
            self.override_params['hospitalization_rate_icu'] = 0.13 * self.override_params['hospitalization_rate_general']
            self.override_params['beds_ICU'] = 0
            self.override_params['symptoms_to_hospital_days'] = 6

            if len(self.covid_data) > 0 and self.covid_data.cases.max() > 0:
                self.t0 = self.covid_data.date.max()
                self.t0, hospitalizations_total = self.get_initial_hospitalizations()

                self.override_params['HGen_initial'] = hospitalizations_total * (1 - self.override_params['hospitalization_rate_icu'] / self.override_params['hospitalization_rate_general'])
                self.override_params['HICU_initial'] = hospitalizations_total * self.override_params['hospitalization_rate_icu']/ self.override_params['hospitalization_rate_general']
                self.override_params['HICUVent_initial'] = self.override_params['HICU_initial'] * self.override_params['fraction_icu_requiring_ventilator']
                self.override_params['I_initial'] = hospitalizations_total / self.override_params['hospitalization_rate_general']

                # The following two params disable the asymptomatic compartment.
                self.override_params['A_initial'] = 0
                self.override_params['gamma'] = 1   # 100% of Exposed go to the infected bucket.

                # 0.6 is a ~ steady state for the exposed bucket initialization at Reff ~ 1.2
                self.override_params['E_initial'] = 0.6 * (self.override_params['I_initial'] + self.override_params['A_initial'])
                self.override_params['D_initial'] = self.covid_data.deaths.max()

            else:
                self.t0 = datetime.datetime.today()
                self.override_params['I_initial'] = 1
                self.override_params['A_initial'] = 0
                self.override_params['gamma'] = 1  # 100% of Exposed go to the infected bucket.

        elif self.run_mode is RunMode.CAN_BEFORE_HOSPITALIZATION_NEW_PARAMS:
            self.n_samples = 1

            for scenario in ['no_intervention', 'flatten_the_curve', 'inferred', 'social_distancing']:
                R0 = 3.6
                self.override_params['R0'] = R0
                if scenario != 'inferred':
                    policy = generate_covidactnow_scenarios(t_list=self.t_list, R0=R0, t0=datetime.datetime.today(), scenario=scenario)
                else:
                    policy = None
                self.suppression_policies[f'suppression_policy__{scenario}'] = policy
                self.override_params = ParameterEnsembleGenerator(
                    self.fips, N_samples=500, t_list=self.t_list, suppression_policy=policy).get_average_seir_parameters()

            if len(self.covid_data) > 0 and self.covid_data.cases.max() > 0:
                self.t0 = self.covid_data.date.max()
                self.t0, hospitalizations_total = self.get_initial_hospitalizations()

                self.override_params['HGen_initial'] = hospitalizations_total * (1 - self.override_params['hospitalization_rate_icu'])
                self.override_params['HICU_initial'] = hospitalizations_total * self.override_params['hospitalization_rate_icu']
                self.override_params['HICUVent_initial'] = self.override_params['HICU_initial'] * self.override_params['fraction_icu_requiring_ventilator']
                self.override_params['I_initial'] = hospitalizations_total / self.override_params['hospitalization_rate_general']

                # The following two params disable the asymptomatic compartment.
                self.override_params['A_initial'] = 0
                self.override_params['gamma'] = 1   # 100% of Exposed go to the infected bucket.

                # 1.2 is a ~ steady state for the exposed bucket initialization.
                self.override_params['E_initial'] = 0.6 * (self.override_params['I_initial'] + self.override_params['A_initial'])
                self.override_params['D_initial'] = self.covid_data.deaths.max()

        elif self.run_mode is RunMode.DEFAULT:
            for suppression_policy in self.suppression_policy:
                self.suppression_policies[f'suppression_policy__{suppression_policy}']= generate_empirical_distancing_policy(
                    t_list=self.t_list, fips=self.fips, future_suppression=suppression_policy)
            self.override_params = dict()
        else:
            raise ValueError('Invalid run mode.')

    @staticmethod
    def _run_single_simulation(parameter_set):
        """
        Run a single simulation instance.

        Parameters
        ----------
        parameter_set: dict
            Params passed to the SEIR model

        Returns
        -------
        model: SEIRModel
            Executed model.
        """
        model = SEIRModel(**parameter_set)
        model.run()
        return model

    def run_ensemble(self):
        """
        Run an ensemble of models for each suppression policy nad generate the
        output report / results dataset.
        """
        for suppression_policy_name, suppression_policy in self.suppression_policies.items():

            logging.info(f'Running simulation ensemble for {self.state_name} {self.fips} {suppression_policy_name}')

            if suppression_policy_name == 'suppression_policy__inferred':

                artifact_path = get_run_artifact_path(self.fips, RunArtifact.MLE_FIT_MODEL)
                if os.path.exists(artifact_path):
                    with open(artifact_path, 'rb') as f:
                        model_ensemble = [pickle.load(f)]
                else:
                    logging.warning(f'No MLE model found for {self.state_name}: {self.fips}. Skipping.')
            else:
                parameter_sampler = ParameterEnsembleGenerator(
                    fips=self.fips,
                    N_samples=self.n_samples,
                    t_list=self.t_list,
                    suppression_policy=suppression_policy)
                parameter_ensemble = parameter_sampler.sample_seir_parameters(override_params=self.override_params)
                model_ensemble = list(map(self._run_single_simulation, parameter_ensemble))

            if self.agg_level is AggregationLevel.COUNTY:
                self.all_outputs['county_metadata'] = self.county_metadata
                self.all_outputs['county_metadata']['age_distribution'] = list(self.all_outputs['county_metadata']['age_distribution'])
                self.all_outputs['county_metadata']['age_bins'] = list(self.all_outputs['county_metadata']['age_distribution'])

            self.all_outputs[f'{suppression_policy_name}'] = self._generate_output_for_suppression_policy(model_ensemble)

        if self.generate_report and self.output_file_report:
            report = CountyReport(self.fips,
                                  model_ensemble=model_ensemble,
                                  county_outputs=self.all_outputs,
                                  filename=self.output_file_report,
                                  summary=self.summary)
            report.generate_and_save()

        with open(self.output_file_data, 'w') as f:
            json.dump(self.all_outputs, f)

    @staticmethod
    def _generate_compartment_arrays(model_ensemble):
        """
        Given a collection of SEIR models, convert these to numpy arrays for
        each compartment, with axis 0 being the model index and axis 1 being the
        timestep.

        Parameters
        ----------
        model_ensemble: list(SEIRModel)

        Returns
        -------
        value_stack: array[n_samples, time steps]
            Array with the stacked model output results.
        """
        compartments = {key: [] for key in model_ensemble[0].results.keys() if key not in ('t_list', 'county_metadata')}
        for model in model_ensemble:
            for key in compartments:
                compartments[key].append(model.results[key])

        return {key: np.vstack(value_stack) for key, value_stack in compartments.items()}

    @staticmethod
    def _get_surge_window(model_ensemble, compartment):
        """
        Calculate the list of surge window starts and ends for an ensemble.

        Parameters
        ----------
        model_ensemble: list(SEIRModel)
            List of models to compute the surge windows for.
        compartment: str
            Compartment to calculate the surge window over.

        Returns
        -------
        surge_start: np.array
            For each model, the surge start window time (since beginning of
            simulation). NaN implies no surge occurred.
        surge_end: np.array
            For each model, the surge end window time (since beginning of
            simulation). NaN implies no surge occurred.
        """
        surge_start = []
        surge_end = []
        for m in model_ensemble:
            # Find the first t where overcapacity occurs
            surge_start_idx = np.argwhere(m.results[compartment] > getattr(m, compartment_to_capacity_attr_map[compartment]))
            surge_start.append(m.t_list[surge_start_idx[0][0]] if len(surge_start_idx) > 0 else float('NaN'))

            # Reverse the t-list and capacity and do the same.
            surge_end_idx = np.argwhere(m.results[compartment][::-1] > getattr(m, compartment_to_capacity_attr_map[compartment]))
            surge_end.append(m.t_list[::-1][surge_end_idx[0][0]] if len(surge_end_idx) > 0 else float('NaN'))

        return surge_start, surge_end

    def _detect_peak_time_and_value(self, value_stack, t_list):
        """
        Compute the peak times for each compartment by finding the arg
        max, and selecting the corresponding time.

        Parameters
        ----------
        value_stack: array[n_samples, time steps]
            Array with the stacked model output results.
        t_list: array
            Array of timesteps.

        Returns
        -------
        peak_data: dict
            For each confidence interval, produce key, value pairs for e.g.
                - peak_time_cl50
                - peak_value_cl50
            Also add peak_value_mean.
        """
        peak_indices = value_stack.argmax(axis=1)
        peak_times = [t_list[peak_index] for peak_index in peak_indices]
        values_at_peak_index = [val[idx] for val, idx in zip(value_stack, peak_indices)]

        peak_data = dict()
        for percentile in self.output_percentiles:
            peak_data['peak_value_ci%i' % percentile] = np.percentile(values_at_peak_index, percentile).tolist()
            peak_data['peak_time_ci%i' % percentile] = np.percentile(peak_times, percentile).tolist()

        peak_data['peak_value_mean'] = np.mean(values_at_peak_index).tolist()
        return peak_data

    def _generate_output_for_suppression_policy(self, model_ensemble):
        """
        Generate output data for a given suppression policy.

        Parameters
        ----------
        model_ensemble: list(SEIRModel)
            List of models to compute the surge windows for.

        Returns
        -------
        outputs: dict
            Output data for this suppression policc ensemble.
        """
        outputs = defaultdict(dict)
        outputs['t_list'] = model_ensemble[0].t_list.tolist()

        # ------------------------------------------
        # Calculate Confidence Intervals and Peaks
        # ------------------------------------------
        for compartment, value_stack in self._generate_compartment_arrays(model_ensemble).items():
            compartment_output = dict()

            # Compute percentiles over the ensemble
            for percentile in self.output_percentiles:
                outputs[compartment]['ci_%i' % percentile] = np.percentile(value_stack, percentile, axis=0).tolist()

            if compartment in compartment_to_capacity_attr_map:
                compartment_output['surge_start'], compartment_output['surge_start'] = self._get_surge_window(model_ensemble, compartment)
                compartment_output['capacity'] = [getattr(m, compartment_to_capacity_attr_map[compartment]) for m in model_ensemble]

            compartment_output.update(self._detect_peak_time_and_value(value_stack, outputs['t_list']))

            # Merge this dictionary into the suppression level one.
            outputs[compartment].update(compartment_output)

        return outputs

def _run_county(fips, ensemble_kwargs):
    """
    Execute the ensemble runner for a specific county.

    Parameters
    ----------
    fips: str
        County fips.
    ensemble_kwargs: dict
        Kwargs passed to the EnsembleRunner object.
    """
    runner = EnsembleRunner(fips=fips, **ensemble_kwargs)
    runner.run_ensemble()

def run_state(state, ensemble_kwargs, states_only=False):
    """
    Run the EnsembleRunner for each county in a state.

    Parameters
    ----------
    state: str
        State to run against.
    ensemble_kwargs: dict
        Kwargs passed to the EnsembleRunner object.
    states_only: bool
        If True only run the state level.
    """
    # Run the state level
    runner = EnsembleRunner(fips=us.states.lookup(state).fips, **ensemble_kwargs)
    runner.run_ensemble()

    if not states_only:
        # Run county level
        df = load_data.load_county_metadata()
        all_fips = df[df['state'].str.lower() == state.lower()].fips
        p = Pool()
        f = partial(_run_county, ensemble_kwargs=ensemble_kwargs)
        p.map(f, all_fips)
        p.close()


# Copyright 2019 The Cirq Developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import List, Optional

import numpy as np
import sympy

import cirq

SQRT_ISWAP = cirq.ISWAP ** 0.5
SQRT_ISWAP_INV = cirq.ISWAP ** -0.5

# TODO: Combine this with the equivalent functions in google/gate_set.py
# Or better yet, write a proper gate set so we don't need this in two places.
# Github issue: https://github.com/quantumlib/Cirq/issues/2970
def _near_mod_n(e, t, n, atol=1e-8):
    return abs((e - t + 1) % n - 1) <= atol

def _near_mod_2pi(e, t, atol=1e-8):
    return _near_mod_n(e, t, 2 * np.pi, atol=atol)

class ConvertToSqrtIswapGates(cirq.PointOptimizer):
    """Attempts to convert gates into ISWAP**-0.5 gates.

    Since we have Z rotations and arbitrary XY rotations, we
    can rely on cirq decomposition for one qubit gates and
    need to only specify special decompositions for two qubit gates.

    Currently natively specified gates are CZPowGate, ISwapPowGate,
    and FSimGate.  This will also support gates that decompose into
    the above gates.
    """

    def __init__(self, ignore_failures=False) -> None:
        """
        Args:
            ignore_failures: If set, gates that fail to convert are forwarded
                unchanged. If not set, conversion failures raise a TypeError.
        """
        super().__init__()
        self.ignore_failures = ignore_failures

    def _convert_one(self, op: cirq.Operation) -> cirq.OP_TREE:
        """
        Decomposer intercept:  Let cirq decompose one-qubit gates,
        intercept on 2-qubit gates if they are known gates.
        """
        if isinstance(op, cirq.GlobalPhaseOperation):
            return []

        gate = op.gate

        if len(op.qubits) != 2:
            return NotImplemented

        q0, q1 = op.qubits

        if isinstance(gate, cirq.CZPowGate):
            if isinstance(gate.exponent, sympy.Basic):
                return cphase_symbols_to_sqrt_iswap(q0, q1, gate.exponent)
            else:
                return cphase_to_sqrt_iswap(q0, q1, gate.exponent)
        if isinstance(gate, cirq.SwapPowGate):
            return swap_to_sqrt_iswap(q0, q1, gate.exponent)
        if isinstance(gate, cirq.ISwapPowGate):
            return iswap_to_sqrt_iswap(q0, q1, gate.exponent)
        if isinstance(gate, cirq.FSimGate):
            return fsim_gate(q0, q1, gate.theta, gate.phi)

        return NotImplemented

    def _on_stuck_raise(self, bad):
        return TypeError(
            f"Don't know how to work with {bad}. "
            "It isn't a native sqrt ISWAP operation, "
            "a 1 or 2 qubit gate with a known unitary, "
            "or composite."
        )

    def convert(self, op: cirq.Operation) -> List[cirq.Operation]:
        return cirq.decompose(
            op,
            keep=is_sqrt_iswap_compatible,
            intercepting_decomposer=self._convert_one,
            on_stuck_raise=(None if self.ignore_failures else self._on_stuck_raise),
        )

    def optimization_at(
        self, circuit: cirq.Circuit, index: int, op: cirq.Operation
    ) -> Optional[cirq.PointOptimizationSummary]:
        if isinstance(op.gate, cirq.MatrixGate) and len(op.qubits) == 1:
            return None

        converted = self.convert(op)
        if len(converted) == 1 and converted[0] is op:
            return None

        return cirq.PointOptimizationSummary(
            clear_span=1, new_operations=converted, clear_qubits=op.qubits
        )

def is_sqrt_iswap_compatible(op: cirq.Operation) -> bool:
    """Check if the given operation is compatible with the sqrt_iswap gateset
    gate set.

    Args:
        op: Input operation.

    Returns:
        True if the operation is native to the gate set, false otherwise.
    """
    return is_basic_gate(op.gate) or is_sqrt_iswap(op.gate)

def is_sqrt_iswap(gate: Optional[cirq.Gate]) -> bool:
    """Checks if this is a ± sqrt(iSWAP) gate specified using either
    ISwapPowGate or with the equivalent FSimGate.
    """
    if (
        isinstance(gate, cirq.FSimGate)
        and not isinstance(gate.theta, sympy.Basic)
        and _near_mod_2pi(abs(gate.theta), np.pi / 4)
        and _near_mod_2pi(gate.phi, 0)
    ):
        return True
    return (
        isinstance(gate, cirq.ISwapPowGate)
        and not isinstance(gate.exponent, sympy.Basic)
        and _near_mod_n(abs(gate.exponent), 0.5, 4)
    )

def is_basic_gate(gate: Optional[cirq.Gate]) -> bool:
    """Check if a gate is a basic supported one-qubit gate.

    Args:
        gate: Input gate.

    Returns:
        True if the gate is native to the gate set, false otherwise.
    """
    return isinstance(
        gate,
        (
            cirq.MeasurementGate,
            cirq.PhasedXZGate,
            cirq.PhasedXPowGate,
            cirq.XPowGate,
            cirq.YPowGate,
            cirq.ZPowGate,
        ),
    )

def cphase_to_sqrt_iswap(a, b, turns):
    """Implement a C-Phase gate using two sqrt ISWAP gates and single-qubit
    operations. The circuit is equivalent to cirq.CZPowGate(exponent=turns).

    Output unitary:
    [1   0   0   0],
    [0   1   0   0],
    [0   0   1   0],
    [0   0   0   e^{i turns pi}].

    Args:
        a: the first qubit
        b: the second qubit
        turns: Exponent specifying the evolution time in number of rotations.
    """
    theta = (turns % 2) * np.pi
    if 0 <= theta <= np.pi:
        sign = 1.0
        theta_prime = theta
    elif np.pi < theta < 2 * np.pi:
        sign = -1.0
        theta_prime = 2 * np.pi - theta

    if np.isclose(theta, np.pi):
        # If we are close to pi, just set values manually to avoid possible
        # numerical errors with arcsin of greater than 1.0 (Ahem, Windows).
        phi = np.pi / 2
        xi = np.pi / 2
    else:
        phi = np.arcsin(np.sqrt(2) * np.sin(theta_prime / 4))
        xi = np.arctan(np.tan(phi) / np.sqrt(2))

    yield cirq.rz(sign * 0.5 * theta_prime).on(a)
    yield cirq.rz(sign * 0.5 * theta_prime).on(b)
    yield cirq.rx(xi).on(a)
    yield cirq.X(b) ** (-sign * 0.5)
    yield SQRT_ISWAP_INV(a, b)
    yield cirq.rx(-2 * phi).on(a)
    yield SQRT_ISWAP(a, b)

    yield cirq.rx(xi).on(a)
    yield cirq.X(b) ** (sign * 0.5)
    # Corrects global phase
    yield cirq.GlobalPhaseOperation(np.exp(sign * theta_prime * 0.25j))

def cphase_symbols_to_sqrt_iswap(a, b, turns):
    """Version of cphase_to_sqrt_iswap that works with symbols.

    Note that the formulae contained below will need to be flattened
    into a sweep before serializing.
    """
    theta = sympy.Mod(turns, 2.0) * sympy.pi

    # -1 if theta > pi.  Adds a hacky fudge factor so theta=pi is not 0
    sign = sympy.sign(sympy.pi - theta + 1e-9)

    # For sign = 1: theta. For sign = -1, 2pi-theta
    theta_prime = (sympy.pi - sign * sympy.pi) + sign * theta

    phi = sympy.asin(np.sqrt(2) * sympy.sin(theta_prime / 4))
    xi = sympy.atan(sympy.tan(phi) / np.sqrt(2))

    yield cirq.rz(sign * 0.5 * theta_prime).on(a)
    yield cirq.rz(sign * 0.5 * theta_prime).on(b)
    yield cirq.rx(xi).on(a)
    yield cirq.X(b) ** (-sign * 0.5)
    yield SQRT_ISWAP_INV(a, b)
    yield cirq.rx(-2 * phi).on(a)
    yield SQRT_ISWAP(a, b)
    yield cirq.rx(xi).on(a)
    yield cirq.X(b) ** (sign * 0.5)

def iswap_to_sqrt_iswap(a, b, turns):
    """Implement the evolution of the hopping term using two sqrt_iswap gates
     and single-qubit operations. Output unitary:
    [1   0   0   0],
    [0   c  is   0],
    [0  is   c   0],
    [0   0   0   1],
    where c = cos(t * np.pi / 2) and s = sin(t * np.pi / 2).

    Args:
        a: the first qubit
        b: the second qubit
        t: Exponent that specifies the evolution time in number of rotations.
    """
    yield cirq.Z(a) ** 0.75
    yield cirq.Z(b) ** 0.25
    yield SQRT_ISWAP_INV(a, b)
    yield cirq.Z(a) ** (-turns / 2 + 1)
    yield cirq.Z(b) ** (turns / 2)
    yield SQRT_ISWAP_INV(a, b)
    yield cirq.Z(a) ** 0.25
    yield cirq.Z(b) ** -0.25

def swap_to_sqrt_iswap(a, b, turns):
    """Implement the evolution of the hopping term using two sqrt_iswap gates
     and single-qubit operations. Output unitary:
    [[1, 0,        0,     0],
     [0, g·c,    -i·g·s,  0],
     [0, -i·g·s,  g·c,    0],
     [0,   0,      0,     1]]
     where c = cos(theta) and s = sin(theta).
        Args:
            a: the first qubit
            b: the second qubit
            theta: The rotational angle that specifies the gate, where
            c = cos(π·t/2), s = sin(π·t/2), g = exp(i·π·t/2).
    """
    if not isinstance(turns, sympy.Basic) and _near_mod_n(turns, 1.0, 2):
        # Decomposition for cirq.SWAP
        yield cirq.Y(a) ** 0.5
        yield cirq.Y(b) ** 0.5
        yield SQRT_ISWAP(a, b)
        yield cirq.Y(a) ** -0.5
        yield cirq.Y(b) ** -0.5
        yield SQRT_ISWAP(a, b)
        yield cirq.X(a) ** -0.5
        yield cirq.X(b) ** -0.5
        yield SQRT_ISWAP(a, b)
        yield cirq.X(a) ** 0.5
        yield cirq.X(b) ** 0.5
        return

    yield cirq.Z(a) ** 1.25
    yield cirq.Z(b) ** -0.25
    yield cirq.ISWAP(a, b) ** -0.5
    yield cirq.Z(a) ** (-turns / 2 + 1)
    yield cirq.Z(b) ** (turns / 2)
    yield cirq.ISWAP(a, b) ** -0.5
    yield cirq.Z(a) ** (turns / 2 - 0.25)
    yield cirq.Z(b) ** (turns / 2 + 0.25)
    yield cirq.CZ.on(a, b) ** (-turns)

def fsim_gate(a, b, theta, phi):
    """FSimGate has a default decomposition in cirq to XXPowGate and YYPowGate,
    which is an awkward decomposition for this gate set.
    Decompose into ISWAP and CZ instead."""
    if theta != 0.0:
        yield cirq.ISWAP(a, b) ** (-2 * theta / np.pi)
    if phi != 0.0:
        yield cirq.CZPowGate(exponent=-phi / np.pi)(a, b)