class Parameters(object):
    """An enumeration of parameters which can be indexed
    like a dictionary."""

    def __index__(self, s):
        return getattr(self, s)

    def get_params(self):
        """Return the set of parameters and their default values
        as a python dictionary."""
        return {k: v for k, v in self.__class__.__dict__.items()
                if "__" not in k}

class TBParameters(Parameters):
    """The ThunderBoltz simulation settings and their
    default values."""
    #: (float) Cell length (m), default is ``1e-6``.
    L=1e-6
    #: (int) Number of species, default is ``2``.
    SP=2
    #: (list[int]) Number of particles for each species,
    #: default is ``[10000, 1000]``.
    NP=[10000, 1000]
    #: (list[float]) Temperature (eV) of each particle species, default
    #: is ``[0.0, 0.0259]``.
    TP=[0.0, 0.0259]
    #: (list[float]) Flow velocity for each particle, default is ``[0.0, 0.0]``.
    VV=[0.0, 0.0]
    #: (list[int]) Charge (elementary units) of each particle species,
    #: default is ``[-1.0, 0.0]``.
    QP=[-1.0, 0.0]
    #: (list[float]) Mass of each particle species (amu),
    #: default is ``[5.4857e-4, 28.0]``.
    MP=[5.4857e-4, 28.0]
    #: (int) Number of time steps, default is ``1000001``.
    NS=1000001
    #: (int) Number of random samples used to find
    #: :math:`\max_\epsilon (v\sigma(\epsilon))` for each process,
    #: default is ``1000``.
    VS=1000
    #: (float) Time increment interval (s), default is ``5e-12``.
    DT=5.0e-12
    #: (float) Electric field in z-direction (V/m), default is
    #: ``-24640.0``.
    E=-24640.0
    #: (float) E-field oscillation frequency (Hz), default is ``0``.
    ET=0.
    #: (list[int]) Magnetic field vector (Tesla),
    #: default is ``[0.0, 0.0, 0.0]``.
    B=[0.0, 0.0, 0.0]
    #: (list[int]) Output velocity dump settings [start, stride, species ID],
    #: default is ``[1000, 1000000, 0]``.
    FV=[1000, 1000000, 0]
    #: (int) Cross section extrapolation — options are ``0`` (extrapolated
    #: cross sections are set to :math:`0` m :math:`^2`), or
    #: ``1`` (linearly extrapolated from last two points), default is ``0``.
    EX=0
    #: (float) Request memory (GB) for particle arrays. If 0 is passed,
    #: then enough space for :math:`4N_{\rm init}` will be
    #: allocated, default is ``0``.
    MEM=0.
    #: (int) When using a SLURM manager on HPC, auto dump particle velocity
    #: data before job allocation runs out —
    #: options are ``0`` (don't auto dump) | ``1`` (dump using SLURM setup).
    SE=0
    #: (int) Time step stride for output parameters, default is ``100``.
    OS=100
    #: (int) If ``1``, then the remainder of :math:`N_{\rm pairs}` is
    #: carried into the next :math:`N_{\rm pairs}` evaluation of the
    #: same process, default is ``1``.
    CR=1
    #: (list[str,int]) Optionally load particle velocities from a comma separated
    #: text file, default is ``None``; specify the name of the file at index
    #: ``0`` and the particle species index it applies to at index ``1``.
    LV=None

class WrapParameters(Parameters):
    """Additional Python interface settings and
    their defaults. An asterisk (*) at the beginning
    of the description indicates a parameter that is
    specific to the built-in He model.
    """
    #: (float) Specify reduced electric field (Td),
    #: default is ``None``.
    Ered=None
    #: (list[float]) Specify reduced magnetic field (Hx),
    #: default is ``None``.
    Bred=None
    #: (int) The index of the neutral gas species (if applicable),
    #: default is ``1``.
    gas_index=1
    #: (float) Run simulation until duration (s) is complete,
    #: default is ``None``.
    duration=None
    #: (int) Number of background neutral gas particles
    #: (assumed gas species is of index 1), default is ``None``.
    NN=None
    #: (float) Set the ratio :math:`\frac{N_{\rm e}}{N_{\rm gas}}`,
    #: default is ``None``.
    pct_ion=None
    #: (float) Minimum number of pseudo pairs to be generated by the
    #: smallest cross section of interest, default is ``1.0``.
    Nmin=1.0
    #: (float) The initial energy (eV) of a hypothetical electron per time step,
    #: default is ``10``.
    EP_0=10.
    #: (float) The maximum change in energy (eV) of a hypothetical electron per time
    #: step, default is ``0.1``.
    DE=0.1
    #: (bool) Flag to calculate ``DT`` / ``NP`` / ``E`` from ``Ered`` / ``L`` /
    #: ``pct_ion`` / ``DE`` / ``EP_0``, default is ``False``.
    autostep=False
    #: (dict[str->float]) Multiply any property involved in the :math:`N_{\rm min}`
    #: constraint by a constant after the constraint is imposed; useful for
    #: convergence testing, default is ``{}`` (do nothing).
    pc_scale={}
    #: (callable or str) Function for auto-generating indeck and CS object or
    #: string path to directory with CS data and indeck,
    #: default is ``None``.
    indeck=None
    #: (list[str or int or 2D-array, int]) Initialize particles with velocity data —
    #: a string at index ``0`` will read velocity data from that file path;
    #: an int at index ``0`` will attempt to reinitialize a previous calculation
    #: from a that time step if that step dump file is available;
    #: an array or DataFrame of shape (NP,3) at index ``0``
    #: will create a velocity initialization file with the provided data,
    #: the value at index 1 should represent the species type.
    vdf_init=None
    #: (bool) If specified with ``vdf_init``, truncate the init file
    #: such that NP is satisfied, default is ``False``.
    downsample=False
    # Cross section model
    #: *(int) For the built-in ``thunderboltz.input.He_TB`` Helium model, include CCC excitation
    #: processes from the ground state to (up to and including) states with
    #: principle quantum number ``n``.
    n=3
    #: *(str or bool) For the built-in ``thunderboltz.input.He_TB`` Helium model, use
    #: either tabulated data, analytic fits, or a mix of both,
    #: options are ``False`` | ``True`` | ``"mixed"``.
    analytic_cs="mixed"
    #: *(int) For the built-in ``thunderboltz.input.He_TB``, this specifies
    #: the number of tabulated cross section values for analytic sampling.
    nsamples=250
    #: *(float) If ``analytic_cs`` is ``"mixed"``, use numerical data
    #: at energies lower than this threshold value (in eV), and use analytic
    #: data at higher energies.
    mix_thresh=300
    #: (bool) Flag to append "FixedParticle2" to each of the reaction types in the indeck
    fixed_background=True
    #: (str) Elastic angular distribution function model —
    #: options are ``"default"``, or ``"He_Park"``.
    eadf="default"
    #: *(str or None) The total elastic cross section model for the He built-in —
    #: options are ``"ICS"`` or ``"MTCS"``, default is ``"ICS"`` if an anisotropic
    #: angular distribution function is used and ``"MTCS"`` if an isotropic
    #: angular distribution function is used.
    ECS=None
    # Ionization model
    egen=False #: *(bool) Allow secondary electron generation for the ionization model.
    #: (str) Electron energy sharing distribution model — options are
    #: ``"default"`` (one takes all) | ``"equal"`` | ``"uniform"``.
    eesd="default"

class OutputParameters(Parameters):
    """A listing of the main output parameters of the simulation, these keywords
    are the named columns of the time series and steady state data frames
    returned by :meth:`~.thunderboltz.ThunderBoltz.get_timeseries` and
    :meth:`~.thunderboltz.ThunderBoltz.get_ss_params` respectively. These data tables also
    include the :class:`ParticleParameters` of the species at index :math:`0`.
    The steady state parameters returned by :meth:`~.thunderboltz.ThunderBoltz.get_ss_params`
    will also include standard deviations for each parameter indicated by an
    added "_std" suffix.
    """

    #: (float) The mean energy (eV) of the species at index :math:`0` (usually electrons),
    #: computed as
    #: :math:`\langle\epsilon\rangle = \frac{m_0}{2N_0}\sum_{i=1}^{N_0}v_{0i}^2`
    #: where :math:`m_0` and :math:`N_0` are the mass and particle count
    #: of the :math:`0^{\rm th}` species, and :math:`v_{0i}` is the velocity
    #: vector of the :math:`i^{\rm th}` particle of species :math:`0`.
    MEe=0.

    #: (float) The reduced flux mobility :math:`(\text{V}^{-1}\text{m}^{-1}\text{s}^{-1})`
    #: of the species at index :math:`0`, computed as
    #: :math:`\mu^fn_{\rm gas}=\frac{n_{\rm gas}}{E}\frac{1}{N_0}\sum_{i=1}^{N_0}v_{\parallel,0i}`
    #: where :math:`N_0` is the particle count of the :math:`0^{\rm th}` species, :math:`E` is
    #: field magnitude, :math:`n_{\rm gas}` is the density of the background
    #: gas, and :math:`v_{\parallel,0i}` is the velocity component
    #: parallel to the E field vector of the :math:`i^{\rm th}` th particle of species :math:`0`.
    mobN=0.

    #: (float) The reduced bulk mobility :math:`(\text{V}^{-1}\text{m}^{-1}\text{s}^{-1})`
    #: of the species at index :math:`0`, computed as
    #: :math:`\mu^b n_{\rm gas}=\frac{n_{\rm gas}}{E}\frac{d}{dt}\left[\frac{1}{N_0}\sum_{i=1}^{N_0}r_{\parallel,0i}\right]`
    #: where :math:`N_0` is the particle count of the :math:`0^{\rm th}` species, :math:`E` is
    #: field magnitude, :math:`n_{\rm gas}` is the density of the background gas,
    #: and :math:`r_{\parallel,0i}` is displacement component
    #: parallel to the E field vector of the :math:`i^{\rm th}` th particle of species :math:`0`.
    mobN_bulk=0.

    #: (float) The reduced flux Townshend ionization coefficient :math:`(\text{m}^2)`
    #: of the species at index :math:`0`, computed as
    #: :math:`\frac{\alpha^f}{n_{\rm gas}} = \frac{1}{n_0 n_{\rm gas}}\frac{dC_{\rm ion}}{dt}\times \left( \frac{1}{N_0}\sum_{i=1}^{N_0}v_{\parallel,0i}\right)^{-1}`
    #: where :math:`N_0` is the particle count of the :math:`0^{\rm th}` species, :math:`C_{\rm ion}` is
    #: the count of ionization events, :math:`n_0` and :math:`n_{\rm gas}` are the
    #: :math:`0^{\rm th}` and background gas densities repectively, and :math:`v_{\parallel,0i}` is velocity component
    #: parallel to the E field vector of the :math:`i^{\rm th}` th particle of species :math:`0`.
    a_n=0.

    #: (float) The reduced bulk Townshend ionization coefficient :math:`(\text{m}^2)`
    #: of the species at index :math:`0`, computed as
    #: :math:`\frac{\alpha^b}{n_{\rm gas}} = \frac{1}{n_0 n_{\rm gas}}\frac{dC_{\rm ion}}{dt}\times\left( \frac{d}{dt} \left[\frac{1}{N_0}\sum_{i=1}^{N_0}r_{\parallel,0i}\right] \right)^{-1}`,
    #: where :math:`N_0` is the particle count of the :math:`0^{\rm th}` species, :math:`C_{\rm ion}` is
    #: the count of ionization events, :math:`n_0` and :math:`n_{\rm gas}` are the
    #: :math:`0^{\rm th}` and background gas densities respectively, and :math:`r_{\parallel,0i}` is displacement component
    #: parallel to the E field vector of the :math:`i^{\rm th}` th particle of species :math:`0`.
    a_n_bulk=0.

    #: (float) :math:`n_{\rm gas}`, the number density of the background gas.
    n_gas=0.

    #: (float) The electric field component (V/m) in the :math:`z` direction, which
    #: can change in AC scenarios.
    E=0.

    t=0. #: (float) The time (s) elapsed in the simulation.

    #: (int) The number of time steps elapsed in the simulation, with
    #: ``t`` :math:`=0` corresponding to ``step`` :math:`=0`, and with
    #: ``t`` = ``DT`` corresponding to ``step`` :math:`=1`.
    step=0

    #: (float) The ionization rate coefficient, computed as
    #: :math:`\frac{1}{n_0n_{\rm gas}} \frac{dC_{\rm ion}}{dt}`.
    k_ion=0.

    #: (float) The rate coefficient for the first reaction, computed as
    #: :math:`\frac{1}{n_0n_{\rm gas}} \frac{dC_{\rm 1}}{dt}`. If there are
    #: :math:`n` reactions included in the model, there will be the
    #: corresponding rate coefficients, :math:`k_1, k_2, ..., k_n` also
    #: included in the output tables.
    k_1=0.

class ParticleParameters(Parameters):
    """A listing of species dependent properties that can be accessed by
    :meth:`~.thunderboltz.ThunderBoltz.get_particle_tables`, which returns a
    list of data tables (one for each species) where each column of data is
    labeled with one of the following keywords."""
    Rxi=0. #: (float) The mean x component of all particle displacements (m).
    Ryi=0. #: (float) The mean y component of all particle displacements (m).
    Rzi=0. #: (float) The mean z component of all particle displacements (m).
    Vxi=0. #: (float) The mean x component velocity (m/s).
    Vyi=0. #: (float) The mean y component velocity (m/s).
    Vzi=0. #: (float) The mean z component velocity (m/s).
    Txi=0. #: (float) The mean x component temperature (eV).
    Tyi=0. #: (float) The mean y component temperature (eV).
    Tzi=0. #: (float) The mean z component temperature (eV).
    Ni=0.0 #: (float) The number density (m :math:`^{-3}`).
    Ki=0.0 #: (float) The total kinetic energy (eV).
    Mi=0.0 #: (float) The mean kinetic energy (eV).
    t=0. #: (float) The time (s) elapsed in the simulation.
    #: (int) The number of time steps elapsed in the simulation, with
    #: ``t`` :math:`=0` corresponding to ``step`` :math:`=0`.
    step=0

tb_parameters = TBParameters().get_params()
wrap_parameters = WrapParameters().get_params()
output_parameters = OutputParameters().get_params()
particle_parameters = ParticleParameters().get_params()
