# spin_system.py
# Simon Hulse
# simon.hulse@chem.ox.ac.uk
# Last Edited: Fri 13 May 2022 10:55:09 BST

r"""This module provides the :py:class:`SpinSystem` class, allowing specification
of a specific spin system according to the nuclear idenitites, isotropic chemical
shifts, and J-couplings of each spin, along with the temperature and magnetic
field they are subjected to.

The ``SpinSystem`` class is used by all experiment simulation functions to derive
the equilibrium operator and Hamiltonians required.
"""

from __future__ import annotations
import random
from typing import Dict, Iterable, Optional, Tuple, Union

import networkx as nx
import numpy as np

from nmr_sims.nuclei import Nucleus
from nmr_sims._operators import CartesianBasis, Operator
from nmr_sims import _sanity

def triangle(n: int) -> int:
    return int(0.5 * n * (n - 1))

class SpinSystem(CartesianBasis):
    """Object representing a particular spin system.

    This represents an ensemble of many identical spin systems subject to a
    specified temperature and magnetic field strength.
    """
    kB = 1.380649e-23
    hbar = 1.054571817e-34

    def __init__(
        self,
        spins: Dict[int, Dict],
        default_nucleus: Union[Nucleus, str] = "1H",
        field: Union[int, float, str] = "500MHz",
        temperature: Union[int, float, str] = "298K",
    ) -> None:
        r"""Create an instance ``SpinSystem``.

        Parameters
        ----------

        spins
            A dictionary with information on each spin that makes up the system.
            Below is an example of a valid format. This specifies a three proton
            (AMX) system, with the following parameters:

            .. math::

                \delta_1 = 5.4\ \mathrm{ppm}

                \delta_2 = 2.8\ \mathrm{ppm}

                \delta_3 = 0.5\ \mathrm{ppm}

                J_{12} = 4.6\ \mathrm{Hz}

                J_{13} = 7.4\ \mathrm{Hz}

                J_{23} = 2.7\ \mathrm{Hz}

            .. code:: python3

                spins = {
                    1: {
                        "shift": 5.4,
                        "couplings": {
                            2: 4.6,
                            3: 7.4,
                        },
                    },
                    2: {
                        "shift": 2.8,
                        "couplings": {
                            3: 2.7,
                        },
                    },
                    3: {
                        "shift": 0.5,
                    },
                }

            Note that there is no need to repeat coupling information. As the
            coupling between spins 1 and 2 has been specified within the entry
            for spin 1, there is no need to repeat this for spin 2. Similary,
            spin 3 has been given no explicit coupling information as this has
            already been provided in the specification of spins 1 and 2.

            Each entry in the ``spins`` dictionary should have an ``int`` as a key
            and a ``dict`` as its value. Within the each ``dict``, the following
            key-value pairs are permitted:

            * ``"shift"`` (necessary) - The isotropic chemical shift of the spin
              in ppm.
            * ``"couplings"`` (optional) - The scalar couplings between the spin
              and its coupling partners. This should be denoted by a ``dict`` with
              an ``int`` key specifying the coupling partner, and a ``float`` value,
              specifying the J-coupling in Hz.
            * ``"nucleus"`` (optional) - The identity of nucelus giving rise to
              the spin. If not specified, then ``default_nucleus`` will be used.
              Can be either an instance of :py:class:`nmr_sims.nuclei.Nucleus`
              of a string corresponding to one of the pre-defined nuclei.

        default_nucleus
            The nucleus identity of any spin in the ``spins`` dictionary that is not
            given an explicit ``"nucleus"`` key. By default, all unspecified nuclei
            will correspond to proton (¹H).

        field
            The magnetic field strength. The following inputs are valid:

            * A positive number (``int`` or ``float``). The field will be taken to be
              in Telsa.
            * A ``str`` satisfying the regex ``r"\d+T"``. Again, the field will be
              taken to be in Tesla
            * A ``str`` satisfying the regex ``r"\d+MHz"``. This corresponds to the
              field which induces the specified Larmor frequency for proton.

            By default, this is set as ``"500MHz"``, correponding to approximately
            11.74T.

        temperature
            The temperature. The following inputs are valid:

            * A positive number (``int`` or ``float``). The field will be taken to be
              in Kelvin.
            * A ``str`` satisfying the regex ``r"\d+K"``. Again, the field will be
              taken to be in Kelvin.
            * A ``str`` satisfying the regex ``r"-?\d+C"``. The temperature will be
              taken to be in Celcius.

        .. warning::

            **Current bug**

            Negative Celcius specifications (``"-\d+C"``) are not supported currently,
            and will raise an error if given. PLease manually convert to Kelvin.
        """
        self.__couplings_on = True
        self.temperature = _sanity.process_temperature(temperature)
        self.field = _sanity.process_field(field)
        self.spin_dict = _sanity.process_spins(spins, default_nucleus)
        super().__init__(
            [(spin.nucleus.multiplicity - 1) / 2 for spin in self.spin_dict.values()]
        )
        if not all([I == 0.5 for I in self.spins]):
            raise ValueError("Only spin-1/2 nuclei are supported currently!")

    def exclude_couplings(self) -> None:
        """Spin system's couplings will be turned off.

        Can be reversed with :py:meth:`include_couplings`.
        """
        self.__couplings_on = False

    def include_couplings(self) -> None:
        """Spin system's couplings will be turned on.

        Can be reversed with :py:meth:`exclude_couplings`.
        """
        self.__couplings_on = True

    @classmethod
    def new_random(
        cls,
        nspins: Union[int, Tuple[int, int]] = 3,
        ncouplings: Optional[int] = None,
        max_couplings: Optional[int] = 3,
        shifts_range: Tuple[float, float] = (0.0, 10.0),
        couplings_range: Tuple[float, float] = (3.0, 15.0),
        field: Union[int, float, str] = "500MHz",
        temperature: Union[int, float, str] = "298K",
    ) -> SpinSystem:
        """Create a new instance with random chemical shifts and offsets.

        .. note::

            Only homonucler ¹H spin systems are available currently.

        Parameters
        ----------
        nspins
            The number of spins the system should comprise.

        ncouplings
            The total number of couplings featured in the spin system. This must be
            a value satisfying ``ncouplings <= int(0.5 * (nspins * (nspins - 1)))``
            There is a possibility that the final spin system will have fewer couplings
            than those stated if the random network of couplings does not satisfy
            the constraint placed by ``max_couplings``.

        max_couplings
            The largest possible number of coupling interactions any particular spin
            may have.

        shifts_range
            The smallest and greatest values of chemical shift any particular spin
            may have.

        couplings_range
            The smallest and greatest values of scalar coupling any particular
            interaction may have.

        field
            See :py:meth:`__init__`

        temperature
            See :py:meth:`__init__`
        """
        if not (isinstance(nspins, int) and nspins >= 1):
            raise ValueError("`nspins` is invalid")

        if not (isinstance(ncouplings, int) and 0 <= ncouplings <= triangle(nspins)):
            raise ValueError(
                "`ncouplings` should be a positive int no greater than "
                f"{triangle(nspins)}."
            )

        if max_couplings is None:
            # No chance of this being exceeded.
            max_couplings = nspins
        elif isinstance(max_couplings, int) and max_couplings >= 1:
            pass
        else:
            raise ValueError("`max_couplings` should be a positive int or None.")

        valid_range = lambda obj: (
            isinstance(shifts_range, (list, tuple)) and
            len(shifts_range) == 2 and
            all([isinstance(x, float) for x in shifts_range])
        )
        if valid_range(shifts_range):
            shifts_range = (min(shifts_range), max(shifts_range))
        else:
            raise ValueError("`shifts_range` is invalid.")

        if valid_range(couplings_range):
            couplings_range = (min(couplings_range), max(couplings_range))
        else:
            raise ValueError("`couplings_range` is invalid.")

        graph = nx.gnm_random_graph(nspins, ncouplings)
        overconnected_nodes = [
            node for node, degree in graph.degree if degree > max_couplings
        ]
        for node in overconnected_nodes:
            while graph.degree[node] > max_couplings:
                graph.remove_edge(node, random.choice([x for x in graph[node]]))

        spin_system = {}
        for i in range(nspins):
            spin_system[i + 1] = {
                "shift": np.round(
                    np.random.uniform(shifts_range[0], shifts_range[1]),
                    decimals=3,
                )
            }
            spin_system[i + 1]["couplings"] = {}

        for i, j in graph.edges:
            spin_system[i + 1]["couplings"][j + 1] = np.round(
                np.random.uniform(couplings_range[0], couplings_range[1]),
                decimals=3,
            )

        return cls(spin_system)

    @property
    def inverse_temperature(self) -> float:
        r"""Return the inverse temperature for the system.

        Given by :math:`\beta = \hbar / k_{\mathrm{B}} T`.
        """
        return self.hbar / (self.kB * self.temperature)

    @property
    def boltzmann_factor(self) -> Iterable[float]:
        r"""Return the Boltzmann factor for each spin in the spin system.

        This factor is given by :math:`2 \pi \beta \gamma B_0`. See
        :py:meth:`inverse_temperature` for the definition of :math:`\beta`.
        """
        return [
            2 * np.pi * self.inverse_temperature * spin.nucleus.gamma * self.field
            for spin in self.spin_dict.values()
        ]

    @property
    def basic_frequencies(self) -> np.ndarray:
        r"""Return the baisc (laboratory frame) frequencies of each spin.

        Given by :math:`-\gamma B_0 \left(1 + \delta \times 10^{-6}\right)`
        """
        return np.array([
            -spin.nucleus.gamma * self.field * (1 + (1e-6 * spin.shift))
            for spin in self.spin_dict.values()
        ])

    @property
    def rotframe_frequencies(self) -> np.ndarray:
        r"""Return the rotating frame frequencies for each spin.

        Given by :math:`-\gamma B_0 \left(\delta \times 10^{-6}\right)`
        """
        return np.array([
            -spin.nucleus.gamma * self.field * (1e-6 * spin.shift)
            for spin in self.spin_dict.values()
        ])

    @property
    def couplings(self) -> np.ndarray:
        couplings = np.zeros((self.nspins, self.nspins), dtype="float64")

        if self.__couplings_on:
            for i, spin in self.spin_dict.items():
                for j, coupling in spin.couplings.items():
                    if i < j:
                        couplings[i - 1, j - 1] = coupling
        else:
            couplings

        return couplings + couplings.T

    def pulse(
        self, nucleus: str, phase: float = 0., angle: float = np.pi / 2
    ) -> Operator:
        """Return the operator for a pulse targeting a specific nucleus, with
        specified phase and flip angle.

        Parameters
        ----------

        nucleus
            The identity of the nucleus to be targeted. Should match the
            ``nucleus.name``.

        phase
            Desired phase of the pulse in radians.

        Angle
            Desired flip angle in radians.
        """
        operator = self.zero
        for i, spin in self.spin_dict.items():
            if spin.nucleus.name == nucleus:
                operator += (
                    np.cos(phase) * self.get(f"{i}x") +
                    np.sin(phase) * self.get(f"{i}y")
                )
        return operator.rotation_operator(angle)

    @property
    def equilibrium_operator(self) -> Operator:
        r"""Return the equilibrium operator of the spin system.

        Given by:

        .. math::

            \hat{\rho}_{\mathrm{eq}} =
            \frac{1}{N} \left(\hat{E} + 2 \pi \beta B_0 \sum_{i=1}^N
            \gamma_i \hat{I}_{iz}\right)

        Where :math:`N` is the number of spins, and :math:`\hat{E}` is the identity
        operator.
        """
        return (1 / self.nspins) * (
            self.identity + sum(
                [b * self.get(f"{i}z") for i, b in
                 enumerate(self.boltzmann_factor, start=1)],
                self.zero,
            )
        )

    def hamiltonian(self, offsets: Union[dict, None] = None, decouple=None) -> Operator:
        r"""Return the Hamiltonian for the spin system.

        Given by:

        .. math::

            \hat{H} = -B_0 \sum_{i=1}^N (\delta_i \times 10^{-6}) \gamma_i
            \hat{I}_{iz}
            + 2 \pi \sum_{i=1}^{N-1} \sum_{j=i+1}^N
            J_{ij} \left(
            \hat{I}_{ix} \hat{I}_{jx} +
            \hat{I}_{iy} \hat{I}_{jy} +
            \hat{I}_{iz} \hat{I}_{jz} \right)

        Parameters
        ----------

        offsets
            Specification of transmitter offsets for given nuclei, in units of Hz.
            As an example, the argument ``{"1H": 5000}`` would inorporate a
            transmtter offset of 5000Hz for proton.
        """
        H = self.zero
        frequencies = self.rotframe_frequencies
        couplings = self.couplings
        isotopes = [self.spin_dict[i].nucleus.name for i in range(1, self.nspins + 1)]
        for i, (isotope1, freq) in enumerate(zip(isotopes, frequencies), start=1):
            H += freq * self.get(f"{i}z")
            if i == self.nspins:
                break
            for j, (isotope2, coupling) in enumerate(zip(
                isotopes[i:],
                couplings[i:, i - 1]
            ), start=i + 1):
                if isotope1 == isotope2:
                    H += np.pi * coupling * (
                        self.get(f"{i}x{j}x") +
                        self.get(f"{i}y{j}y") +
                        self.get(f"{i}z{j}z")
                    )
                elif decouple not in (isotope1, isotope2):
                    H += np.pi * coupling * self.get(f"{i}z{j}z")
                # Decouple: don't include scalar coupling

        if offsets is not None:
            for nuc, off in offsets.items():
                H += 2 * np.pi * off * sum(
                    [self.get(f"{i}z")
                     for i, spin in self.spin_dict.items()
                     if spin.nucleus.name == nuc],
                    self.zero
                )

        return H

    def _get_sum(self, coord: str, nucleus: str) -> Operator:
        if nucleus is None:
            labels = list(range(1, self.nspins + 1))
        else:
            labels = [i for i, spin in self.spin_dict.items()
                      if spin.nucleus.name == nucleus]

        return sum([self.get(f"{i}{coord}") for i in labels], self.zero)

    def Ix(self, nucleus: Union[str, None] = None) -> Operator:
        r"""Return the :math:`\hat{I}_x` operator correpsonding to the system,
        with the option of specifying the nucelus to target.

        For a given nucleus :math:`n`, this is given by:

        .. math::

            \hat{I}_{x, n} = \sum_{i} \hat{I}_{ix},

        :math:`\forall i \in \{1, \cdots, N\}` satisfying the requirement that
        spin :math:`i` corresponds to nucelus :math:`n`.

        Parameters
        ----------

        nucleus
            The identity of the nucelus to target. If ``None``, no constraint will
            be put on the nucleus.
        """
        return self._get_sum("x", nucleus)

    def Iy(self, nucleus: Union[str, None] = None) -> Operator:
        r"""Return the :math:`\hat{I}_y` operator correpsonding to the system,
        with the option of specifying the nucelus to target.

        For a given nucleus :math:`n`, this is given by:

        .. math::

            \hat{I}_{y, n} = \sum_{i} \hat{I}_{iy},

        :math:`\forall i \in \{1, \cdots, N\}` satisfying the requirement that
        spin :math:`i` corresponds to nucelus :math:`n`.

        Parameters
        ----------

        nucleus
            The identity of the nucelus to target. If ``None``, no constraint will
            be put on the nucleus.
        """
        return self._get_sum("y", nucleus)

    def Iz(self, nucleus: Union[str, None] = None) -> Operator:
        r"""Return the :math:`\hat{I}_z` operator correpsonding to the system,
        with the option of specifying the nucelus to target.

        For a given nucleus :math:`n`, this is given by:

        .. math::

            \hat{I}_{z, n} = \sum_{i} \hat{I}_{iz},

        :math:`\forall i \in \{1, \cdots, N\}` satisfying the requirement that
        spin :math:`i` corresponds to nucelus :math:`n`.

        Parameters
        ----------

        nucleus
            The identity of the nucelus to target. If ``None``, no constraint will
            be put on the nucleus.
        """
        return self._get_sum("z", nucleus)

if __name__ == "__main__":
    ss = SpinSystem.new_random(nspins=6, ncouplings=6, max_couplings=2)
    from nmr_sims.experiments.pa import PulseAcquireSimulation

    sim = PulseAcquireSimulation(ss, 4096, "10ppm", "5ppm")
    sim.simulate()
    sh, sp, lab = sim.spectrum(zf_factor=4)
    import matplotlib as mpl
    mpl.use("tkAgg")
    import matplotlib.pyplot as plt
    fig = plt.figure()
    ax = fig.add_subplot()
    ax.plot(sh, sp)
    ax.set_xlim(reversed(ax.get_xlim()))
    ax.set_xlabel(lab)
    plt.show()


"""

                  extra classes used in findR

                  frame    -> store a frame
                  species  -> store a species in a frame (or just as such)
                  fcompare -> frame comparer
                  revent   -> reaction event

##################### NOTES /TO BE IMPLEMENTED

add xyz coord reference to frame object for mol object generation
-> not completely implemented

add pmol also to species
move molg generation to the frame and species objects

"""

import numpy as np
import os

import molsys
from molsys.util import rotations

from graph_tool import Graph, GraphView
import graph_tool.topology as gtt
import graph_tool.util as gtu
import graph_tool.draw as gtd

import uuid
import graph_tool

#####################  FRAME CLASS ########################################################################################

class frame:
    """container to store info on a frame

    the main attributes are:
    - fid: frame id
    - xyz:     numpy array (N,3) with the xyz coords of all atoms in this frame
    - elems:   list of elems
    - bondtab: bond table (M,2) with m bonds
    - bondord: bon orders (M)
    """

    def __init__(self, fid, xyz, mol, bondtab, bondord, cutoff=0.5, min_atom=6, min_elems={"c":2}):
        self.fid = fid
        self.xyz = xyz
        self.mol = mol
        self.natoms = self.mol.get_natoms()
        self.bondord = bondord
        self.bondtab = bondtab
        self.maxbond = len(bondord)
        self.specs = {}
        # generate the molecular graph
        molg = Graph(directed=False)
        molg.add_vertex(n=self.natoms)
        molg.vp.aid  = molg.new_vertex_property("int")
        molg.vp.aid.a[:] = np.arange(self.natoms, dtype="int32")
        molg.vp.filt = molg.new_vertex_property("bool")
        molg.vp.mid  = molg.new_vertex_property("int")
        molg.ep.bord = molg.new_edge_property("float")
        molg.ep.filt = molg.new_edge_property("bool")
        for j in range(self.maxbond):
            o = bondord[j]
            # invalid entries are marked with a bondord = -1.0
            if o < 0.0:
                break
            e = bondtab[j]-1
            # TO BE REMOVED ... this is a legacy check for the old incorrect mfp5 files
            if (e[0] < 0 or e[0] >= self.natoms or e[1] < 0 or e[1] >= self.natoms):
                break
            # END TO BE REMOVED
            newb = molg.add_edge(e[0], e[1])
            isbond = o > cutoff
            molg.ep.bord[newb] = o
            molg.ep.filt[newb] = isbond
        # apply edge filter
        molg.set_edge_filter(molg.ep.filt)
        self.molg = molg
        self.mid, hist = gtt.label_components(molg, vprop=molg.vp.mid)
        nspecies_all = len(hist)
        nspecies = []
        # collect number of critical elements in the species
        elem_spec = {}
        for e in min_elems.keys():
            spec_ecount = np.zeros(nspecies_all, dtype="int32")
            for i in range(self.natoms):
                if self.mol.elems[i] == e:
                    spec_ecount[self.mid[i]] += 1
            elem_spec[e] = spec_ecount
        for i in range(nspecies_all):
            track = True
            if hist[i] < min_atom:
                track = False
            for e in min_elems.keys():
                if elem_spec[e][i] < min_elems[e]:
                    track = False
            if track:
                nspecies.append(i)
        # append the species
        for s in nspecies:
            self.add_species(s)
        return

    def add_species(self, mid, tracked=True):
        # make all frame species (observed) with a local graph representation
        self.specs[mid] = species(self.fid, mid, self.molg, make_graph=True, tracked=tracked)
        return

    def make_species(self, mid):
        # return a species object (without a local graph)
        return species(self.fid, mid, self.molg, make_graph=False, tracked=False)

    @property
    def nspecies(self):
        return len(self.specs)

    def plot(self, selection=False):
        if selection:
            print ("plotting species %s of frame %d" % (str(selection), self.fid))
            for s in selection:
                gfname = "frame_%d_species_%d.png" % (self.fid, s)
                g = self.specs[s].graph
                pos = gtd.arf_layout(g, max_iter=0)
                gtd.graph_draw(g, pos=pos, vertex_text=g.vp.aid, vertex_font_size=12, vertex_size=8, \
                                           edge_text=g.ep.bord, edge_font_size=10,\
                output_size=(800, 800), output=gfname, bg_color=[1,1,1,1])
        else:
            print ("plotting all species of frame %d" % self.fid)
            for s in self.specs.keys():
                gfname = "frame_%d_species_%d.png" % (self.fid, s)
                g = self.specs[s].graph
                pos = gtd.arf_layout(g, max_iter=0)
                gtd.graph_draw(g, pos=pos, vertex_text=g.vp.aid, vertex_font_size=12, vertex_size=8, \
                output_size=(200, 200), output=gfname, bg_color=[1,1,1,1])
        return

    def get_border(self, b):
        """get the bond order for a bond b

        currently no test is made if the edge exists ... you need to check

        Args:
            b (tuple of indices): atom indices of bond
        """
        e = self.molg.edge(b[0], b[1])
        border = self.molg.ep.bord[e]
        return border

    def make_mol(self, species):
        """generate a single mol object from a list of species for a TS

        NOTE: the additional species added to the frame are considered to be non tracked

        """
        aids = []
        for mid in species:
            if mid not in self.specs:
                self.add_species(mid, tracked=False)
            s = self.specs[mid]
            aids += list(s.aids)
        aids.sort() # This is necessary to ensure same ordering in mol object
        mol = molsys.mol.from_array(self.xyz[aids])
        mol.set_cell(self.mol.get_cell())
        mol.set_elems([self.mol.elems[i] for i in aids])
        mol.set_real_mass()
        mol.center_com(check_periodic=False)
        mol.apply_pbc()
        mol.make_nonperiodic()
        # rotate into principal axes
        xyz = mol.get_xyz()
        mol.set_xyz(rotations.align_pax(xyz, masses = mol.get_mass()))
        mol.set_atypes(mol.get_elems())
        # add connectivity
        conn = []
        for i in aids:
            v = self.molg.vertex(i)
            conn_i = []
            for j in v.all_neighbors():
                assert int(j) in aids
                conn_i.append(aids.index(int(j)))
            conn.append(conn_i)
        mol.set_conn(conn)
        return mol, aids

    def make_species_mol(self):
        """get a list of mol objects for all species in the frame

        NOTE: we use here the global xyz and pmol -> should be delegated to the species object which needs refactoring
        - we should attacg
        """
        mols = {}
        for s in self.specs:
            sp = self.specs[s]
            sxyz = self.xyz[list(sp.aids)]
            mols[s] = self.specs[s].make_mol(sxyz, self.mol)
        return mols

    ### some DEBUG methods ###

    def write_species(self):
        foldername = "frame_%d_species" % self.fid
        os.mkdir(foldername)
        os.chdir(foldername)
        mols = self.make_species_mol()
        for s in mols:
            m = mols[s]
            m.write("spec_%d.mfpx" % s)
        os.chdir("..")

    def get_main_species_formula(self):
        """get the sumformula of the tracked species
        """
        sp = list(self.specs.keys())
        sp.sort()
        sumforms = []
        for s in sp:
            aids = list(self.specs[s].aids)
            elems = [self.mol.elems[i] for i in aids]
            cont  = list(set(elems))
            cont.sort()
            sumform = ""
            for e in cont:
                sumform += "%s%d " % (e, elems.count(e))
            sumforms.append(sumform)
        return sumforms

    def DEBUG_write_as_xyz(self, fname):
        f = open(fname, "w")
        f.write("%d\n\n" % self.natoms)
        for i in range(self.natoms):
            x, y, z = self.xyz[i]
            f.write("%3s %12.6f %12.6f %12.6f\n" % (self.mol.elems[i], x, y, z))
        f.close()
        return

#####################  SPECIES CLASS ########################################################################################

class species:
    """container class to keep species info (per frame!)
    """

    def __init__(self, fid, mid, molg, make_graph=False, tracked = True):
        """init species

        Args:
            fid (int): frame number
            mid (int): molecule id ("name" of species from label components)
            molg (graph): molg of the frame
        """
        self.fid = fid
        self.mid = mid
        self.molg = molg # parent molgraph
        # find all vertices in molg that belong to this species mid
        vs = gtu.find_vertex(molg, molg.vp.mid, mid)
        self.aids = set([int(v) for v in vs]) # atomids -> TBI do we need to sort? seems like they are sorted as they come
        self.graph = None
        if make_graph:
            # now make a view of molg for this species
            self.make_graph()
        self.tracked = tracked
        return

    def make_graph(self):
        vfilt = self.molg.new_vertex_property("bool")
        vfilt.a[:] = False
        for v in self.aids:
            vfilt[v] = True
        self.graph = GraphView(self.molg, vfilt=vfilt)
        return

    @property
    def natoms(self):
        return len(self.aids)

    def make_mol(self, xyz, pmol):
        """generate a mol object for the species from a frame

        we get the current coordinates as xyz and the parent mol object pmol from the mfp5 file

        Args:
            xyz (numpy): coordinates of the atoms
            pmol (molsys object): parent mol object from mfp5 file
        """
        aids = list(self.aids) # in order to map backwards
        self.mol = molsys.mol.from_array(xyz)
        if pmol.get_cell() is not None:
            self.mol.set_cell(pmol.get_cell())
        self.mol.set_elems([pmol.elems[i] for i in aids])
        self.mol.set_real_mass()
        self.mol.center_com(check_periodic=False)
        self.mol.apply_pbc()
        self.mol.make_nonperiodic()
        # rotate into principal axes
        xyz = self.mol.get_xyz()
        self.mol.set_xyz(rotations.align_pax(xyz, masses = self.mol.get_mass()))
        self.mol.set_atypes(self.mol.get_elems())
        # add connectivity
        if self.graph is None:
            self.make_graph()
        ctab = []
        for e in self.graph.edges():
            i = aids.index(int(e.source()))
            j = aids.index(int(e.target()))
            ctab.append([i, j])
        self.mol.set_ctab(ctab, conn_flag=True)
        return self.mol

    def __eq__(self, other):
        """compare two species whether they are equal

        Compares "exactly" .. atom indices (wrt frame) and bonding must match exactly.
        This can not be used to find similar species (same bonding etc but different atom numbers)
        Note: The conformation of the species can be entirely differnt -> we do not compare postions!

        Args:
            other (species object): the other species object

        Returns:
            bool: if they are equal or not equal
        """
        if not type(self) is type(other):
            return NotImplemented
        # check if atom IDs are equal
        if self.aids != other.aids:
            return False
        # generate a set of bonds for each species
        bself  = set([(int(e.source()),int(e.target())) for e in self.graph.edges()])
        bother = set([(int(e.source()),int(e.target())) for e in other.graph.edges()])
        if bself != bother:
            return False

        # if we are here the species are equal
        return True

    def __repr__(self):
        return "SPECIES{frame %5d / molid %3d (%4d atoms)}" % (self.fid, self.mid, self.natoms)

#####################  FRAME COMPARE CLASS ########################################################################################

class fcompare:
    """This class compares two frames at different levels of resolution whether species are different
       All info collected during the compairson is kept in this class in order to be exploited later
    """

    def __init__(self, f1, f2):
        """generate comparer

        Args:
            f1 (frame object): first frame object
            f2 (frame object): second frame object (to be compared with f1)
        """
        self.f1 = f1
        self.f2 = f2
        # defaults
        self.compare_level = 0 # 0: no comparison, 1: atom id level, 2: connectivity level
        self.umatch_f1 = [s for s in self.f1.specs if self.f1.specs[s].tracked]
        self.umatch_f2 = [s for s in self.f2.specs if self.f2.specs[s].tracked]
        self.aids_match = []
        self.bond_match = []
        self.aids_analyzed = False
        self.bonds_analyzed = False
        self.reacs = []
        self.broken_bonds = []
        self.formed_bonds = []
        self.nreacs = 0 # number of independent reactions for this pair of frames .. should always be 1 (??)
        return

    def report(self, all = False):
        """this method is just to implement all levels of comparison

           it does not return anything and just reports ... meant for debugging
        """
        print("##################################################")
        print("FRAMES        %5d         %5d" % (self.f1.fid, self.f2.fid))
        mf = self.check_aids()
        if all:
            for m in self.aids_match:
                print("          species %5d == species %5d" % m)
                print(" natoms:          %5d            %5d" % (self.f1.specs[m[0]].natoms, self.f2.specs[m[1]].natoms))
        if mf > 0:
            self.analyse_aids()
            for r in self.reacs:
                print ("educts  : %s" % str(list(r[0].keys())))
                print ("products: %s" % str(list(r[1].keys())))
        return

    def check_aids(self, verbose = False):
        """check on level 1 for atom id matches

        this method implements a rather complex logic:
        which species form the educts and products to define a complete reaction event
        """
        if self.compare_level < 1:
            # we need to compare
            for sk1 in self.f1.specs:
                s1 = self.f1.specs[sk1]
                # find a corresponding species in frame 2
                for sk2 in self.umatch_f2:
                    s2 = self.f2.specs[sk2]
                    if s1.aids == s2.aids:      # NOTE: this works because the vertices are always properly sorted
                        # s1 and s2 match -> add to match and to remove
                        self.aids_match.append((sk1, sk2))
                        self.umatch_f1.remove(sk1)
                        self.umatch_f2.remove(sk2)
                        break
            # now matches are found -> set compare level
            self.compare_level = 1
        match_flag = 0 # all is equal on level 1
        if len(self.umatch_f1) > 0 or len(self.umatch_f2) > 0:
            match_flag = 1 # unmatched species!!
        if verbose:
            print ("species in f1   : %s" % str(self.f1.specs.keys()))
            print ("species in f2   : %s" % str(self.f2.specs.keys()))
            print ("unmatched in f1 : %s" % str(self.umatch_f1))
            print ("unmatched in f2 : %s" % str(self.umatch_f2))

        return match_flag

    def analyse_aids(self):
        if self.aids_analyzed:
            return
        if self.compare_level == 0:
            self.check_aids()
        if len(self.umatch_f1)==0 and len(self.umatch_f2)==0:
            # there is nothing to do
            return
        # we have some unmatched species -> there is one (or more) reaction(s) between these frames
        # find groups of species that define a reaction
        #    all atom ids in the union of the educts sets must be also in the products set
        for sk1 in self.umatch_f1:
            # first search for atoms in the unmatched f2 species
            s1 = self.f1.specs[sk1]
            educts = {sk1: s1} # dictionary of species key/species for this reaction
            educt_aids = s1.aids.copy() # set of atomids
            products = {}
            product_aids = set()
            for sk2 in self.umatch_f2:
                s2 = self.f2.specs[sk2]
                common_aids = s1.aids & s2.aids
                if len(common_aids)>0:
                    products[sk2] = s2
                    product_aids |= s2.aids
            # do the following until educt_aids and product_aids match
            while not (educt_aids == product_aids):
                # which atoms are in products that are not in the educt species? --> add them
                for a in product_aids - educt_aids:
                    # to which species in frame 1 does this atom belong to?
                    esk = self.f1.molg.vp.mid[a]
                    # is this already in educts?
                    if esk not in educts:
                        # we need to make a new species object and add it (as non-tracked)
                        educts[esk] = self.f1.make_species(esk)
                        educt_aids |= educts[esk].aids

                        # avoid that found educt is in first list. In that case you will find a reaction twice
                        if esk in self.umatch_f1:
                            self.umatch_f1.remove(esk)

                # which atoms are in the educts that are not in the product species? add them
                for a in educt_aids - product_aids:
                    # to which species in frame 2 does this atom belong to?
                    psk = self.f2.molg.vp.mid[a]
                    # is this already in educts?
                    if psk not in products:
                        # we need to make a new species object and add it (as non-tracked)
                        products[psk] = self.f2.make_species(psk)
                        product_aids |= products[psk].aids

                        # GS: unsure if this is required. Should avoid to assign a product twice
                        if psk in self.umatch_f2:
                            self.umatch_f2.remove(psk)

            self.reacs.append((educts, products))
        # the above will not work if there is no tracked species in umatch_f1 (but in umatch_f2)
        # ... in other words a species has "appeared" or formed by merging two or more untracked species
        if len(self.umatch_f1)==0:
            for sk2 in self.umatch_f2:
                s2 = self.f2.specs[sk2]
                products = {sk2: s2}
                product_aids = s2.aids.copy()
                # now find all the species in frame 1 to match the atoms
                educts = {}
                educt_aids = set()
                for a in product_aids:
                    esk = self.f1.molg.vp.mid[a]
                    if esk not in educts:
                        educts[esk] = self.f1.make_species(esk)
                        educt_aids |= educts[esk].aids
                # do the following until educt_aids and product_aids match
                while not (educt_aids == product_aids):
                    # which atoms are in products that are not in the educt species? --> add them
                    for a in product_aids - educt_aids:
                        # to which species in frame 1 does this atom belong to?
                        esk = self.f1.molg.vp.mid[a]
                        # is this already in educts?
                        if esk not in educts:
                            # we need to make a new species object and add it (as non-tracked)
                            educts[esk] = self.f1.make_species(esk)
                            educt_aids |= educts[esk].aids

                    # which atoms are in the educts that are not in the product species? add them
                    for a in educt_aids - product_aids:
                        # to which species in frame 2 does this atom belong to?
                        psk = self.f2.molg.vp.mid[a]
                        # is this already in educts?
                        if psk not in products:
                            # we need to make a new species object and add it (as non-tracked)
                            products[psk] = self.f2.make_species(psk)
                            product_aids |= products[psk].aids
                # now add the final results to the reacs list
                self.reacs.append((educts, products))
        self.nreacs = len(self.reacs)
        self.aids_analyzed = True
        return

    def find_react_bond(self):
        """find a reactive bond in a bimolecular reaction

        Args:
            r (int): index in self.reac to analyse
        """
        assert self.aids_analyzed
        g1 = self.f1.molg
        g2 = self.f2.molg
        for r in range(self.nreacs):
            broken_bonds = []
            formed_bonds = []
            educts, products = self.reacs[r]
            # get all involved atoms
            aids = set()
            for s in educts:
                aids |= educts[s].aids
            for a in aids:
                bonds1 = []
                v1 = g1.vertex(a)
                for e in v1.out_edges():
                    bonds1.append(int(e.target()))
                bonds2 = []
                v2 = g2.vertex(a)
                for e in v2.out_edges():
                    bonds2.append(int(e.target()))
                bs = set(bonds1) - set(bonds2)
                fs = set(bonds2) - set(bonds1)
                for b in bs:
                    if b > a:
                        broken_bonds.append((a,b))
                for b in fs:
                    if b > a:
                        formed_bonds.append((a,b))
            self.broken_bonds.append(broken_bonds)
            self.formed_bonds.append(formed_bonds)
        return

    def check_bonds(self, verbose = False):
        """check on level 2 for identical bonds

        we use the existing pairs of species in self.aids_match to identify species that have identical aids
        => now we test if they have the same bonding pattern
        """
        assert self.compare_level > 0
        self.missmatch = []
        for p in self.aids_match:
            # get the species of the pair
            s1 = self.f1.specs[p[0]]
            s2 = self.f2.specs[p[1]]
            # # TBI: this might not be enough
            # #      do we need elements as vertex properties to be considered?
            # #      what about a tautomerism when the initial and fianl state are symmetric?
            # f = gtt.isomorphism(s1.graph, s2.graph)
            # INDEED: isomoprhism is not strict enough
            # print ("%d %d %d %d %s" % (self.f1.fid, self.f2.fid, p[0], p[1], f))
            if  s1 == s2:
                self.bond_match.append(p)
            else:
                self.missmatch.append(p)
        self.compare_level = 2
        if len(self.missmatch) > 0:
            # unmatched species on level 2
            return 2
        else:
            return 0 # all equal

    def analyse_bonds(self):
        if self.bonds_analyzed:
            return
        if self.compare_level < 2:
            self.check_bonds()
        if len(self.missmatch) == 0:
            return
        # now analyse all the pairs in self.missmatch: they have identical aids but a diffrent bond graph
        self.nreacs = len(self.missmatch)
        for p in self.missmatch:
            # get the species of the pair
            s1 = self.f1.specs[p[0]]
            s2 = self.f2.specs[p[1]]
            self.reacs.append(({p[0]: s1},{p[1]: s2}))
            # get all edges as vertex id tuples (int tuples)
            bs1 = set([(int(e.source()),int(e.target())) for e in s1.graph.edges()])
            bs2 = set([(int(e.source()),int(e.target())) for e in s2.graph.edges()])
            self.broken_bonds.append(list(bs1-bs2)) # broken: bond in frame1 but not in frame2
            self.formed_bonds.append(list(bs2-bs1)) # formed: bond in frame2 but not in frame1
        self.bonds_analyzed = True
        return

    def check(self, verbose = False):
        """check identity of species on all levels (1 and 2)
        """
        mf = self.check_aids(verbose=verbose)
        if mf > 0:
            return mf
        # aids are equal -> chek bonds
        return self.check_bonds(verbose=verbose)

#####################  REACTIVE EVENT CLASS ########################################################################################
# this class is instantiated with a comparer (between two frames)
# and stores a reactive event
# it knows the TS_fid and if the event is bi- or unimolecular

class revent:

    def __init__(self, comparer, fR, unimol= False, ireac=0):
        """generate a reaction event object

        TBI: what to do if there are more then one reaction event (in two diffrent tracked species)
             at the same time ... this is properly tracked in the comparer (nreacs>1)
             but this means there are more revents to be generated. should we do this recursively?
             how to store?

        Args:
            comparer (fcompare object): comparer that gave a reactive event
            fR (parent findR object): to access process_frame
            unimol (bool, optional): is unimolecular. Defaults to False.
            ireac (int,optional): specifies the reaction event
        """
        self.unimol = unimol
        self.fR = fR
        self.comparer = comparer
        # currently we allow only for single reaction events per frame
        # this would change if there is more than one tracked species ....
        #assert comparer.nreacs == 1 , "Currently only single reaction events are processed. Error occured for frames %s and %s" % (comparer.f1.fid, comparer.f2.fid)
        #r = 0 # pick first reaction (the only one)
        assert comparer.nreacs > ireac, "Too many reaction events requested..."
        r = ireac
        print("ireac " + str(ireac) + " frames " + str(comparer.f1.fid) + " and " + str(comparer.f2.fid) )
        educts, products = comparer.reacs[r]
        self.broken_bonds     = comparer.broken_bonds[r]
        self.formed_bonds     = comparer.formed_bonds[r]
        f1               = comparer.f1
        f2               = comparer.f2
        # choose which frame we use as TS
        # everything is referenced to f1 which is 0 (f2 is +1)
        # find avereage bond order of reactive bonds at f1/f2
        f1_averborder = 0.0
        if len(self.broken_bonds) >0:
            for b in self.broken_bonds:
                f1_averborder += f1.molg.ep.bord[f1.molg.edge(b[0], b[1])]
            f1_averborder/len(self.broken_bonds)
        f2_averborder = 0.0
        if len(self.formed_bonds) >0:
            for b in self.formed_bonds:
                f2_averborder += f2.molg.ep.bord[f2.molg.edge(b[0], b[1])]
            f2_averborder/len(self.formed_bonds)
        if f1_averborder == 0.0:
            TS_rfid = 1
        elif f2_averborder == 0.0:
            TS_rfid = 0
        else:
            if abs(f1_averborder-0.5) < abs(f2_averborder-0.5):
                # f1 closer to TS
                TS_rfid = 0
            else:
                TS_rfid = 1
        # now store data depending on relative frame id (rfid) of the TS
        if TS_rfid == 0:
            self.TS = f1
            self.PR = f2
            self.ED = self.fR.process_frame(self.TS.fid-1)
            # get corresponding species numbers
            self.TS_spec = educts
            self.PR_spec = products
            loccomp = fcompare(self.ED, self.TS)
            if loccomp.check_aids() == 0:
                # no change in atom ids .. we can use TS species for ED as well
                self.ED_spec = {}
                for e in educts:
                    if e in self.ED.specs:
                        self.ED_spec[e] = self.ED.specs[e]
                    else:
                        self.ED_spec[e] = self.ED.make_species(e)
            else:
                print ("Houston we have a problem!!! species changed between ED and TS")
        else:
            self.ED = f1
            self.TS = f2
            if self.fR.nframes > self.TS.fid+1:
                self.PR = self.fR.process_frame(self.TS.fid+1)
                self.ED_spec = educts
                self.TS_spec = products
                loccomp = fcompare(self.TS, self.PR)
                if loccomp.check_aids() == 0:
                    # no change in atom ids .. we can use TS species for PR as well
                    self.PR_spec = {}
                    for e in products:
                        if e in self.PR.specs:
                            self.PR_spec[e] = self.PR.specs[e]
                        else:
                            self.PR_spec[e] = self.PR.make_species(e)
                else:
                    print ("Houston we have a problem!!! species changed between TS and PR")
            else:
                print("Warning! We run out of frames to process to identify TS and PR.")
        # get TS_fid for ease
        self.TS_fid = self.TS.fid

        return


#!/usr/bin/env python

#Title:         Module to facilitate basic reaction calculations for Sub-Chandra research
#Author:        Adam Jacobs
#Creation Date: Nov. 1, 2013

#Usage: load as a module

#Description: This module contains code to facilitate carrying out some basic reaction calculations
#             for analysis in the Sub-Chandra II paper.
#
#Revision History
#Programmer            Date                    Change
# ------------------------------------------------------------------------------
#Adam Jacobs         11/01/2013              Code created

#TODO:
# 1)

#Notes on conventions, programming style:
#  1) All classes are in CamelCase with the first letter capitalized.  Class methods
#     are also in CamelCase with the first letter lower-case, e.g. myMethod().
#  2) Non-class functions and all variables are lowercase with underscores (_)
#     acting as delimiters when needed/wanted.
#  3) An underscore prefix, e.g. _variable, means the named item is intended to be
#     private (Python doesn't provide for easily enforcing data hiding, it's all by
#     convention).
#  4) Names that are in ALL_CAPS are intended as constants.

######################
### Global Imports ###
######################
import sys

#################################
### Global Data and Constants ###
#################################
#_EXAMPLE = 'string'

###############
### Classes ###
###############

#################
### Functions ###
#################
def plot_logenuc():
  import numpy as np
  import matplotlib
  #If 'inline' backend, we're running within IPython so don't change anything.
  #Otherwise, we're running as script to generate eps, so use PS backend
  if 'inline' not in matplotlib.get_backend():
    matplotlib.use('PS')
  import matplotlib.pyplot as plt

  matplotlib.rcParams['text.usetex'] = 'True'
  # Make data based on MESA calcs in mesa_subch/
  rho_exp = np.linspace(5.0, 7.0, 11)
  rho     = np.array([10**rexp for rexp in rho_exp])
  #rho_nco = np.array([10**rexp for rexp in rho_exp[-3:]])

  # 200 MK rates, OLD, version 5XXX
  #enuc_cagob2  = np.array(
  #    [10**17.581540508,
  #     10**17.800549914,
  #     10**18.024481339,
  #     10**18.254609218,
  #     10**18.492973480,
  #     10**18.738377513,
  #     10**18.986558564,
  #     10**19.234718720,
  #     10**19.477401243,
  #     10**19.704819234,
  #     10**19.897451724])
  #enuc_tralf2  = np.array(
  #    [10**11.188457831,  #rho = 10**5.0
  #     10**11.626476644,  #rho = 10**5.2
  #     10**12.074446783,  #rho = 10**5.4
  #     10**12.531401511,  #rho = 10**5.6
  #     10**12.996679600,  #rho = 10**5.8
  #     10**13.469768554,  #rho = 10**6.0
  #     10**13.948445742,  #rho = 10**6.2
  #     10**14.432637076,  #rho = 10**6.4
  #     10**14.910185702,  #rho = 10**6.6
  #     10**15.450434567,  #rho = 10**6.8
  #     10**16.013534830]) #rho = 10**7.0
  #enuc_nagf2   = np.array(
  #    [10**10.356064284,
  #     10**10.596664766,
  #     10**10.840074663,
  #     10**11.083643491,
  #     10**11.322076922,
  #     10**11.545046637,
  #     10**11.796568378,
  #     10**12.066216636,
  #     10**12.348913239,
  #     10**12.646625087,
  #     10**12.961628265])
  #enuc_cago2   = np.array(
  #    [10**7.066187025,
  #     10**7.303537281,
  #     10**7.544452034,
  #     10**7.787414219,
  #     10**8.029050316,
  #     10**8.262882790,
  #     10**8.482066479,
  #     10**8.743164789,
  #     10**9.015982418,
  #     10**9.302270992,
  #     10**9.604056810])
  #enuc_nco2    = np.array(
  #    [10**7.621642798,
  #     10**7.907935225,
  #     10**8.209727138])

  # 200 MK rates, MESA version 7503, see mesa_subch/200MKSeries.out
  enuc_cagob2  = np.array(
      [3.88724408e+17,     #rho = 10**5.0
       6.43552100e+17,     #rho = 10**5.2
       1.07753194e+18,     #rho = 10**5.4
       1.82999980e+18,     #rho = 10**5.6
       3.16711338e+18,     #rho = 10**5.8
       5.57201399e+18,     #rho = 10**6.0
       9.86657427e+18,     #rho = 10**6.2
       1.74726877e+19,     #rho = 10**6.4
       3.05623900e+19,     #rho = 10**6.6
       5.16352101e+19,     #rho = 10**6.8
       8.07526663e+19])    #rho = 10**7.0
  enuc_tralf2  = np.array(
      [1.61236410e+11,
       4.41923321e+11,
       1.23924575e+12,
       3.54801266e+12,
       1.03546483e+13,
       3.07691004e+13,
       9.26265273e+13,
       2.82426937e+14,
       8.46686308e+14,
       2.93507889e+15,
       1.07243491e+16])
  enuc_nagf2   = np.array(
      [2.39846615e+10,
       4.17329568e+10,
       7.30904519e+10,
       1.28073324e+11,
       2.21835352e+11,
       3.71300418e+11,
       6.60773488e+11,
       1.22889032e+12,
       2.35501530e+12,
       4.67150539e+12,
       9.64213931e+12])
  enuc_cago2   = np.array(
      [1.21008577e+07,
       2.08980546e+07,
       3.63892899e+07,
       6.36683332e+07,
       1.11075656e+08,
       1.90389587e+08,
       3.15051582e+08,
       5.74527718e+08,
       1.07630179e+09,
       2.07969722e+09,
       4.16416868e+09 ])
  #
  # 300 MK rates, OLD, version 5XXX
  #
  enuc_cagob3  = np.array(
      [10**19.030792008,
       10**19.241106991,
       10**19.454092787,
       10**19.670440935,
       10**19.891022033,
       10**20.116932101,
       10**20.349550944,
       10**20.591053956,
       10**20.837613055,
       10**21.086204225,
       10**21.333254625])
  enuc_tralf3  = np.array(
      [10**13.785826231,  #rho = 10**5.0
       10**14.206456199,  #rho = 10**5.2
       10**14.632427790,  #rho = 10**5.4
       10**15.065124085,  #rho = 10**5.6
       10**15.506286282,  #rho = 10**5.8
       10**15.957694147,  #rho = 10**6.0
       10**16.417507200,  #rho = 10**6.2
       10**16.885601880,  #rho = 10**6.4
       10**17.360983029,  #rho = 10**6.6
       10**17.842218167,  #rho = 10**6.8
       10**18.324423763]) #rho = 10**7.0
  enuc_nagf3   = np.array(
      [10**13.685591364,
       10**13.909659660,
       10**14.141720810,
       10**14.379825361,
       10**14.621578864,
       10**14.865447555,
       10**15.108015882,
       10**15.342720865,
       10**15.567331414,
       10**15.829352083,
       10**16.103271007])
  enuc_cago3   = np.array(
      [10**9.738896065,
       10**9.959526033,
       10**10.185497624,
       10**10.419619014,
       10**10.658290572,
       10**10.900155663,
       10**11.143164162,
       10**11.383040770,
       10**11.611826212,
       10**11.842872776,
       10**12.107808383])
  #enuc_nco3    = np.array(
  #    [10**10.973769648,
  #     10**11.204818056,
  #     10**11.469756036])

  #make the plots
  #fig, ax = plt.subplots(nrows=2, ncols=1)
  fig, ax = plt.subplots(nrows=1, ncols=1)

  ###plot the 200MK data
  ax.plot(rho,     enuc_cagob2, label=r'CagO-by')
  ax.plot(rho,     enuc_tralf2, label=r'3$\alpha$')
  ax.plot(rho,     enuc_nagf2,  label=r'NagF')
  ax.plot(rho,     enuc_cago2,  label=r'CagO')
  #ax.plot(rho_nco, enuc_nco2,   label=r'NCO')

  #ax[0].plot(rho,     enuc_tralf3, 'b--')
  #ax[0].plot(rho,     enuc_cagob3, 'g--')
  #ax[0].plot(rho,     enuc_cago3,  'r--')
  #ax[0].plot(rho_nco, enuc_nco3,   'c--')

  #tune plot settings
  ax.set_yscale('log')
  ax.set_xscale('log')
  ax.set_ylabel(r'$\dot{\epsilon}~[\rm{erg}~\rm{g}^{-1}~\rm{s}^{-1}]$', fontsize='x-large')
  ax.set_ylim(1.e6, 1.e22)

  ax.set_xlabel(r'$\rho~[\rm{g}~\rm{cm}^{-3}]$', fontsize='x-large')

  ax.tick_params(labelsize='x-large')
  ax.legend(loc=2, fontsize='medium')
  ax.set_title(r'$T = 200~\rm{MK}$', fontsize='x-large')

  ###plot the 300MK data
  #ax[1].plot(rho,     enuc_cagob3, label=r'CagO-by')
  #ax[1].plot(rho,     enuc_tralf3, label=r'3$\alpha$')
  #ax[1].plot(rho,     enuc_nagf3,  label=r'NagF')
  #ax[1].plot(rho,     enuc_cago3,  label=r'CagO')
  #ax[1].plot(rho_nco, enuc_nco3,   label=r'NCO')

  ##tune plot settings
  #ax[1].set_yscale('log')
  #ax[1].set_xscale('log')
  #ax[1].set_ylabel(r'$\dot{\epsilon}~[\rm{erg}~\rm{g}^{-1}~\rm{s}^{-1}]$', fontsize='xx-large')
  #ax[1].set_ylim(1.e6, 1.e22)

  #ax[1].set_xlabel(r'$\rho~[\rm{g}~\rm{cm}^{-3}]$', fontsize='xx-large')
  #
  #ax[1].tick_params(labelsize='x-large')
  #ax[1].legend(loc=2, fontsize='large')
  #ax[1].set_title(r'$T = 300 ~\rm{MK}$', fontsize='xx-large')

  fig.set_size_inches(5.0*1.25, 5.0)
  #fig.tight_layout()
  #This fixes a bug in mpl's eps generation
  matplotlib.rc('ps', usedistiller='xpdf')
  fig.savefig('rxns.eps', bbox_inches='tight')

  return

#################
### Execution ###
#################
#This is only for testing. rxncalcs.py is intended to be used as a module.
if __name__== "__main__":
  if(len(sys.argv) <= 1):
    #TODO: Add arg checks
    pass

  #Put tests here
  plot_logenuc()