# This file is part of QuTiP: Quantum Toolbox in Python.
#
#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
#       of its contributors may be used to endorse or promote products derived
#       from this software without specific prior written permission.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################
import warnings

import numpy as np

from qutip.operators import tensor, identity, destroy, sigmax, sigmaz
from qutip.states import basis
from qutip.qip.circuit import QubitCircuit, Gate
from qutip.qip.device.processor import Processor
from qutip.qip.device.modelprocessor import ModelProcessor
from qutip.qip.operations import expand_operator
from qutip.qobj import Qobj
from qutip.qobjevo import QobjEvo
from qutip.qip.pulse import Pulse
from qutip.qip.compiler.gatecompiler import GateCompiler
from qutip.qip.compiler import CavityQEDCompiler

__all__ = ['DispersiveCavityQED']

class DispersiveCavityQED(ModelProcessor):
    """
    The processor based on the physical implementation of
    a dispersive cavity QED system.
    The available Hamiltonian of the system is predefined.
    For a given pulse amplitude matrix, the processor can
    calculate the state evolution under the given control pulse,
    either analytically or numerically.
    (Only additional attributes are documented here, for others please
    refer to the parent class :class:`qutip.qip.device.ModelProcessor`)

    Parameters
    ----------
    N: int
        The number of qubits in the system.

    correct_global_phase: float, optional
        Save the global phase, the analytical solution
        will track the global phase.
        It has no effect on the numerical solution.

    num_levels: int, optional
        The number of energy levels in the resonator.

    deltamax: int or list, optional
        The sigma-x paraicient for each of the qubits in the system.

    epsmax: int or list, optional
        The sigma-z paraicient for each of the qubits in the system.

    w0: int, optional
        The base frequency of the resonator.

    eps: int or list, optional
        The epsilon for each of the qubits in the system.

    delta: int or list, optional
        The epsilon for each of the qubits in the system.

    g: int or list, optional
        The interaction strength for each of the qubit with the resonator.

    t1: list or float
        Characterize the decoherence of amplitude damping for
        each qubit. A list of size `N` or a float for all qubits.

    t2: list of float
        Characterize the decoherence of dephasing for
        each qubit. A list of size `N` or a float for all qubits.

    Attributes
    ----------
    sx_ops: list
        A list of sigmax Hamiltonians for each qubit.

    sz_ops: list
        A list of sigmaz Hamiltonians for each qubit.

    cavityqubit_ops: list
        A list of interacting Hamiltonians between cavity and each qubit.

    sx_u: array_like
        Pulse matrix for sigmax Hamiltonians.

    sz_u: array_like
        Pulse matrix for sigmaz Hamiltonians.

    g_u: array_like
        Pulse matrix for interacting Hamiltonians
        between cavity and each qubit.

    wq: list of float
        The frequency of the qubits calculated from
        eps and delta for each qubit.

    Delta: list of float
        The detuning with repect to w0 calculated
        from wq and w0 for each qubit.
    """

    def __init__(self, N, correct_global_phase=True,
                 num_levels=10, deltamax=1.0,
                 epsmax=9.5, w0=10., wq=None, eps=9.5,
                 delta=0.0, g=0.01, t1=None, t2=None):
        super(DispersiveCavityQED, self).__init__(
            N, correct_global_phase=correct_global_phase,
            t1=t1, t2=t2)
        self.correct_global_phase = correct_global_phase
        self.spline_kind = "step_func"
        self.num_levels = num_levels
        self._paras = {}
        self.set_up_params(
            N=N, num_levels=num_levels, deltamax=deltamax,
            epsmax=epsmax, w0=w0, wq=wq, eps=eps,
            delta=delta, g=g)
        self.set_up_ops(N)
        self.dims = [num_levels] + [2] * N

    @property
    def ctrls(self):
        result = []
        for pulse in self.pulses:
            result.append(pulse.get_ideal_qobj(self.dims))
        return result

    def set_up_ops(self, N):
        """
        Generate the Hamiltonians for the spinchain model and save them in the
        attribute `ctrls`.

        Parameters
        ----------
        N: int
            The number of qubits in the system.
        """
        # single qubit terms
        self.a = tensor(destroy(self.num_levels))
        self.pulses.append(
            Pulse(self.a.dag() * self.a, [0], spline_kind=self.spline_kind))
        for m in range(N):
            self.pulses.append(
                Pulse(sigmax(), [m+1], spline_kind=self.spline_kind))
        for m in range(N):
            self.pulses.append(
                Pulse(sigmaz(), [m+1], spline_kind=self.spline_kind))
        # interaction terms
        a_full = tensor([destroy(self.num_levels)] +
                        [identity(2) for n in range(N)])
        for n in range(N):
            sm = tensor([identity(self.num_levels)] +
                        [destroy(2) if m == n else identity(2)
                         for m in range(N)])
            self.pulses.append(
                Pulse(a_full.dag() * sm + a_full * sm.dag(),
                      list(range(N+1)), spline_kind=self.spline_kind))

        self.psi_proj = tensor([basis(self.num_levels, 0)] +
                               [identity(2) for n in range(N)])

    def set_up_params(
            self, N, num_levels, deltamax,
            epsmax, w0, wq, eps, delta, g):
        """
        Save the parameters in the attribute `params` and check the validity.

        Parameters
        ----------
        N: int
            The number of qubits in the system.

        num_levels: int
            The number of energy levels in the resonator.

        deltamax: list
            The sigma-x paraicient for each of the qubits in the system.

        epsmax: list
            The sigma-z paraicient for each of the qubits in the system.

        wo: int
            The base frequency of the resonator.

        wq: list
            The frequency of the qubits.

        eps: list
            The epsilon for each of the qubits in the system.

        delta: list
            The delta for each of the qubits in the system.

        g: list
            The interaction strength for each of the qubit with the resonator.

        Notes
        -----
        All parameters will be multiplied by 2*pi for simplicity
        """
        sx_para = super(DispersiveCavityQED, self)._para_list(deltamax, N)
        self._paras["sx"] = sx_para
        sz_para = super(DispersiveCavityQED, self)._para_list(epsmax, N)
        self._paras["sz"] = sz_para
        w0 = w0 * 2 * np.pi
        self._paras["w0"] = w0
        eps = super(DispersiveCavityQED, self)._para_list(eps, N)
        self._paras["eps"] = eps
        delta = super(DispersiveCavityQED, self)._para_list(delta, N)
        self._paras["delta"] = delta
        g = super(DispersiveCavityQED, self)._para_list(g, N)
        self._paras["g"] = g

        # computed
        self.wq = [np.sqrt(eps[i]**2 + delta[i]**2) for i in range(N)]
        self.Delta = [self.wq[i] - w0 for i in range(N)]

        # rwa/dispersive regime tests
        if any([g[i] / (w0 - self.wq[i]) > 0.05 for i in range(N)]):
            warnings.warn("Not in the dispersive regime")

        if any([(w0 - self.wq[i])/(w0 + self.wq[i]) > 0.05 for i in range(N)]):
            warnings.warn(
                "The rotating-wave approximation might not be valid.")

    @property
    def sx_ops(self):
        return self.ctrls[1: self.N + 1]

    @property
    def sz_ops(self):
        return self.ctrls[self.N + 1: 2*self.N + 1]

    @property
    def cavityqubit_ops(self):
        return self.ctrls[2*self.N + 1: 3*self.N + 1]

    @property
    def sx_u(self):
        return self.coeffs[1: self.N + 1]

    @property
    def sz_u(self):
        return self.coeffs[self.N + 1: 2*self.N + 1]

    @property
    def g_u(self):
        return self.coeffs[2*self.N + 1: 3*self.N + 1]

    def get_ops_labels(self):
        """
        Get the labels for each Hamiltonian.
        """
        return ([r"$a^\dagger a$"] +
                [r"$\sigma_x^%d$" % n for n in range(self.N)] +
                [r"$\sigma_z^%d$" % n for n in range(self.N)] +
                [r"$g_{%d}$" % (n) for n in range(self.N)])

    def optimize_circuit(self, qc):
        """
        Take a quantum circuit/algorithm and convert it into the
        optimal form/basis for the desired physical system.

        Parameters
        ----------
        qc: :class:`qutip.QubitCircuit`
            Takes the quantum circuit to be implemented.

        Returns
        -------
        qc: :class:`qutip.QubitCircuit`
            The circuit representation with elementary gates
            that can be implemented in this model.
        """
        self.qc0 = qc
        self.qc1 = self.qc0.resolve_gates(
            basis=["SQRTISWAP", "ISWAP", "RX", "RZ"])
        return self.qc1

    def eliminate_auxillary_modes(self, U):
        """
        Eliminate the auxillary modes like the cavity modes in cqed.
        """
        return self.psi_proj.dag() * U * self.psi_proj

    def load_circuit(self, qc):
        """
        Decompose a :class:`qutip.QubitCircuit` in to the control
        amplitude generating the corresponding evolution.

        Parameters
        ----------
        qc: :class:`qutip.QubitCircuit`
            Takes the quantum circuit to be implemented.

        Returns
        -------
        tlist: array_like
            A NumPy array specifies the time of each coefficient

        coeffs: array_like
            A 2d NumPy array of the shape (len(ctrls), len(tlist)). Each
            row corresponds to the control pulse sequence for
            one Hamiltonian.
        """
        gates = self.optimize_circuit(qc).gates

        dec = CavityQEDCompiler(
            self.N, self._paras, self.wq, self.Delta,
            global_phase=0., num_ops=len(self.ctrls))
        tlist, self.coeffs, self.global_phase = dec.decompose(gates)
        for i in range(len(self.pulses)):
            self.pulses[i].tlist = tlist
        # TODO The amplitude of the first control a.dag()*a
        # was set to zero before I made this refactoring.
        # It is probably due to the fact that
        # it contributes only a constant (N) and can be neglected.
        # but change the below line to np.ones leads to test error.
        self.coeffs[0] = self._paras["w0"] * np.zeros(len(tlist))
        return tlist, self.coeffs


#!/usr/bin/env python

"""
    The purpose of this script is to edit the topology file of a system containing molecules which have a benzene ring
    in order to create your choice of two things:
    (1) An artificial dipole which will act as electron clouds participating in a pi bond. The dipole is
        created by centering two virtual sites above and below the plane of the benzene ring and assigning them
        appropriate charges values.
    (2) Add position restraints with a given force constant to chosen atoms w.r.t. to a specified axis or axes
"""
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
import argparse
import numpy as np
import warnings
import os
from LLC_Membranes.llclib import file_rw, topology
import mdtraj as md

location = os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file__)))  # Location of this script

def initialize():

    parser = argparse.ArgumentParser(description='Duplicate points periodically in the x-y directions')

    parser.add_argument('-g', '--gro', default='initial.gro', type=str, help='Coordinate file')
    parser.add_argument('-o', '--out', default='dipole.itp', type=str, help='Name of output topology file')
    parser.add_argument('-f', '--f_const', nargs='+', default=[1000, 1000, 1000], type=float, help='Force constant')
    parser.add_argument('-a', '--atoms', nargs='+', default=['C', 'C1', 'C2', 'C3', 'C4', 'C5'], type=str, help='Name of carbons in ring')
    parser.add_argument('-d', '--distance', default=0.1, help='Distance to offset dipole from ring (Angstroms)')
    parser.add_argument('-m', '--monomer', default='NAcarb11V', help='Which monomer topology is being used')
    parser.add_argument('-A', '--axis', default='xy', help='Axis to restrain along with position restraints')
    ######## These parameters are not implemented yet but left here since they will have some use in the future ########
    parser.add_argument('--xlink', help='Specify this flag if the system is being crosslinked while restrained',
                        action="store_true")
    parser.add_argument('--append', help='Specify this to prevent the topology from being re-written, and instead, add '
                                         'restraints to the topology for other atoms', action="store_true")
    parser.add_argument('-i', '--input', type=str, default='dipole.itp', help='Name of topology file to be edited if '
                                                                              'you use the option --append')
    ####################################################################################################################
    parser.add_argument('-dr', '--dihedral_restraints', action='append', nargs='+',
                        help='Specify atom names of dihedral to be restrained, followed by angle at which to restrain'
                             'them, the deviation from that angle allowed and the force constant to apply. For example:'
                             '"restrain.py -dr C1 C C6 O4 90 0 1000" means to keep the angle between the planes formed'
                             'by C1-C-C6 and C-C6-O4 90 degrees apart with 0 degrees of leeway and a force constant of '
                             '1000')
    parser.add_argument('-com', '--center_of_mass', action="store_true", help="Add position restraints at the center of"
                                                                              "mass of args.atoms")
    parser.add_argument('-v', '--virtual_site_parameters', nargs='+', default=['3fd', 'C', 'C2', 'C4', '.5', '.14'],
                        help='A list in the following order : virtual site construction type, atoms to use to build '
                             'virtual site, required length parameters (i.e. a, b, c or d) in the order specified in'
                             'GROMACS documentation')
    parser.add_argument('-b', '--bond_restraints', default=[.37, 1000], nargs='+', help='Bond restraint pararmeters. A'
                        'list where the first entry is the equilibrium distance (nm) and the second entry is the force'
                        'constant for a harmonic potential (kJ/mol/nm^2)')

    args = parser.parse_args()

    return args

warnings.filterwarnings("error")  # This makes it so numpy warnings are treated as real errors

def virtual_sites(all_coords, monomers, valence, a, b, c, funct):

    n_atoms = np.shape(all_coords)[1] - (1 / valence) * monomers
    atoms_per_molecule = n_atoms / monomers  # subtract valence to exclude counterion

    n_vsites = monomers * 2  # number of virtual sites need to create a dipole (2 per ring)
    vsites = np.zeros([8, n_vsites])
    vsites[4, :] = funct
    vsites[5, :] = a  # all a and b values are the same
    vsites[6, :] = b

    for i in range(monomers):
        for j in range(2):
            vsites[0, i * 2 + j] = n_atoms + i * 2 + j + 1  # new atoms were placed at the end of the .gro file
            vsites[1:4, i * 2 + j] = [i * atoms_per_molecule + 1, i * atoms_per_molecule + 3,
                              i * atoms_per_molecule + 5]
            vsites[7, i * 2 + j] = (-1)**j * c

    return vsites

def exclusions(coord_file, monomers, valence, toplines, atoms, n_atoms, vsites):
    """
    :param coord_file: the original .gro file stored in a list
    :param monomers: number of monomers
    :param valence: charge on ions
    :param atoms: the names of the atoms which should be excluded
    :param n_atoms: the number of atoms total
    :return: a list of exclusions
    """
    n_atoms = np.shape(all_coords)[1] - (1 / valence) * monomers
    atoms_per_molecule = n_atoms / monomers  # subtract valence to exclude counterion (includes new PI atoms)

    n_excluded = len(atoms) + 1  # number of atoms excluded. +1 because it will be excluded from it complementary vsite

    n_exclusions = monomers * 2  # number of exclusions to be specified (one for each vsite)
    exclusions = np.zeros([n_excluded + 1, n_exclusions])  # the first entry is the virtual site itself

    for i in range(np.shape(vsites)[1]):
        exclusions[0, i] = vsites[0, i]
        # add exclusion from the complementary vsite. This is pretty specific to the format and will likely need to be
        # re-written eventually
        if i % 2 == 0:
            exclusions[1, i] = vsites[0, i + 1]
        elif i % 2 == 1:
            exclusions[1, i] = vsites[0, i - 1]

    x = 0
    for i in range(monomers):
        a = 2
        for j in range(atoms_per_molecule):
            line = i*atoms_per_molecule + j + toplines
            if str.strip(coord_file[line][10:15]) in atoms:
                exclusions[a, x] = int(coord_file[line][15:20])
                exclusions[a, x + 1] = int(coord_file[line][15:20])
                a += 1
        x += 2

    return exclusions

def dihedral_restraints(file, atoms):
    """
    This function needs to be moved into the class
    """

    ndihedrals = len(atoms)

    all_restraints = np.zeros([0, 8])
    for n in range(ndihedrals):

        atom_numbers = []
        d = np.zeros([4])
        count = 0
        for line in file:
            atom = str.strip(line[10:15])  # name of atom at that line
            if atom in atoms[n]:
                d[atoms[n].index(atom)] = int(line[15:20])  # atom number placed in d in order that dihedral was passed
                if np.count_nonzero(d) == 4:  # and len(atom_numbers) % 8 == 0:
                    if count % 2 == 1:
                        for i in range(4):
                            atom_numbers.append(d[i])
                    d = np.zeros([4])
                    count += 1

        restraints = np.zeros([len(atom_numbers)//4, 8])
        for i in range(len(atom_numbers)//4):
            d = 4 * i
            restraints[i, :] = [atom_numbers[d], atom_numbers[d + 1], atom_numbers[d + 2], atom_numbers[d + 3], 1, atoms[n][4],
                                atoms[n][5], atoms[n][6]]

        all_restraints = np.concatenate((all_restraints, restraints))

    return all_restraints

class RestrainedTopology(object):

    def __init__(self, gro, res, atoms, name='restrained', com=False, xlink=False,
                 vparams=None):
        """ Write topology to restrain one or more residues with position restraints in GROMACS

        :param gro: coordinate file where restraints will be placed
        :param res: name of residue where position restraints are being added
        :param atoms: name of atoms to be restrained in res
        :param name: name of output topology file
        :param com: restrain center of mass of atoms instead of individual atoms
        :param xlink : whether or not the system is in the process of being crosslinked
        :param vparams: A list in the following order : virtual site construction type, atoms to use to build virtual
               site, required length parameters (i.e. a, b, c or d) in the order specified in GROMACS documentation
        """

        topology.fix_resnumbers(gro)
        t = md.load(gro)

        if type(res) is str:
            res = [res]
            atoms = [atoms]

        self.all_coords = t.xyz[0, :, :]  # all coordinates for system
        self.atom_numbers = []
        for i in range(len(atoms)):
            self.atom_numbers.append([a.index + 1 for a in t.topology.atoms if a.name in atoms[i] and a.residue.name ==
                                      res[i]])

        self.atoms = t.n_atoms  # number of atoms in full system
        self.nmon = [len(self.atom_numbers[i]) // len(atoms[i]) for i in range(len(atoms))]  # number of monomer residues
        self.name = name  # name of output files (.itp, .gro if you are using centers of masses)
        self.residue = res  # name of residue(s) to which position restraints are being applied
        self.LC = [topology.LC('%s' % r) for r in self.residue]  # everything we can know about the residue
        self.com = com

        # self.keep = np.array([a.index for a in t.topology.atoms if a.name in atoms])  # atoms to keep
        # self.atom_numbers = self.keep + 1  # numbers (not indices) of atoms to keep
        # self.coords = self.all_coords[self.keep, :]  # coordinates of atoms in keep

        if self.com:  # add center of mass virtual site

            # These things are only needed for center of mass virtual site construction
            self.ids = [a.name for a in t.topology.atoms]  # names of all atoms in system
            self.res = [a.residue.name for a in t.topology.atoms]  # residue names of all atoms in system
            self.vparams = vparams
            self.vatoms_numbers = [a.index + 1 for a in t.topology.atoms if a.name in vparams]
            self.box = t.unitcell_vectors[0, :, :]  # box vectors in mdtraj formate
            self.box_gromacs = [self.box[0, 0], self.box[1, 1], self.box[2, 2], self.box[0, 1], self.box[2, 0],
                                self.box[1, 0], self.box[0, 2], self.box[1, 2],
                                self.box[2, 0]]  # gromacs format box vects

            with open('%s/../top/Monomer_Tops/%s.itp' % (location, self.residue), 'r') as f:
                residue_top = []
                for line in f:
                    residue_top.append(line)

            atoms_index = 0
            while residue_top[atoms_index].count('[ atoms ]') == 0:
                atoms_index += 1

            while residue_top[atoms_index] != '\n':
                atoms_index += 1

            residue_top.insert(atoms_index, '{:>6d}{:>5s}{:>6d}{:>6s}{:>6s}{:>5d}{:>13.6f}{:>13.6f}\n'.format(
                self.LC.natoms + 1, 'hc_d', 1, self.LC.residues[0], 'HD', self.LC.natoms + 1, 0, 0))

            if self.vparams[0] == '3fd':
                # 'a' = 0.5 and 'b' = 0.14 (aromatic carbon bond length (nm)) puts a vsite in the middle of benzene
                # if the constructor atoms are 3 non-adjacent carbons from the ring
                residue_top.append('[ virtual_sites3 ]\n')
                residue_top.append('{:<6d}{:<6d}{:<6d}{:<6d}{:<6d}{:<8.4f}{:<8.4f}\n'.format(self.LC.natoms + 1,
                                    self.vatoms_numbers[0], self.vatoms_numbers[1], self.vatoms_numbers[2], 2,
                                    float(self.vparams[-2]), float(self.vparams[-1])))
            else:
                print('Your choice of virtual site has not yet been implemented')
                exit()

            # groups = np.reshape(self.coords, (len(self.keep) // len(atoms), len(atoms), 3))
            # centers_of_mass = np.mean(groups, axis=1)
            # self.coords = centers_of_mass  # redefine coordinates as centers of mass

            file_rw.write_assembly(residue_top, '%s.itp' % self.name, self.nmon, xlink=xlink)

            # now the dummies need to be added to the .gro file. They are placed at the end of the residue section
            # This loop works for a single virtual site per monomer. It will need to be modified if multiple sites
            # are to be constructed.
            insert_ndx = self.LC.natoms
            self.atom_numbers = []  # redefine this since everything is renumbered
            for i in range(self.nmon):
                ndx = (i + 1)*insert_ndx + i
                self.ids.insert(ndx, 'HD')
                self.res.insert(ndx, 'HII')  # should make this more general
                self.all_coords = np.insert(self.all_coords, ndx, np.array([0, 0, 0]), axis=0)
                self.atom_numbers.append(ndx + 1)

            file_rw.write_gro_pos(self.all_coords, '%s.gro' % self.name, ids=self.ids, res=self.res, box=self.box_gromacs)
        else:

            file_rw.write_assembly(res, '%s.itp' % self.name, self.nmon, xlink=xlink)

        with open('%s.itp' % self.name, 'r') as f:
            self.topology = []
            for line in f:
                self.topology.append(line)

    def add_position_restraints(self, axis, f_const):
        """
        Restrain the selected atoms in desired directions
        :param axis: which direction to restrain (xyz, xy, z, xz .. etc.)
        :param f_const: force constant in each direction. Order of force constants matches that of axis argument
        :return: an array of position restraints formatted for easy writing into the topology (.itp)
        """

        self.topology.append("\n[ position_restraints ]\n")

        fc = np.zeros([3])
        for i, a in enumerate(axis):
            if a == 'x':
                fc[0] = f_const[i]
            if a == 'y':
                fc[1] = f_const[i]
            if a == 'z':
                fc[2] = f_const[i]

        atom_numbers = []
        for i in self.atom_numbers:
            atom_numbers += i

        restraints = np.zeros([5, len(atom_numbers)])  # organize them into a list which can be translated to a topology
        for i in range(len(atom_numbers)):
            restraints[:, i] = [atom_numbers[i], 1, fc[0], fc[1], fc[2]]  # See: http://www.gromacs.org/Documentation/How-tos/Position_Restraints
            self.topology.append('{:6d}{:6d}{:1s}{:9f}{:1s}{:9f}{:1s}{:9f}\n'.format(int(restraints[0, i]),
                                int(restraints[1, i]),'', restraints[2, i], '', restraints[3, i], '', restraints[4, i]))

    def add_distance_restraint_columns(self, b0, kb, layers=20, pores=4):
        """
        Add distance constraints to centers of mass of monomer head groups. This is a function specialized for an
        HII system built with build.py (without the flag -columns).
        :param b0 : equilibrium distance
        :param kb : force constant for harmonic potential
        :param layers : layers per pore
        :param pores : number of pore columns
        """

        self.topology.append("\n[ bonds ]\n")
        mpl = int(self.nmon / 4 / layers)  # monomers per layer

        for p in range(pores):
            for l in range(layers):
                for m in range(mpl):
                    if l == (layers - 1):  # handles periodicity
                        self.topology.append('{:<6d}{:<6d}{:<6d}{:<6.1f}{:<6.1f}\n'.format(self.atom_numbers[layers*mpl*p
                                             + l*mpl + m], self.atom_numbers[layers*mpl*p + m], 6, b0, kb))
                    else:
                        self.topology.append('{:<6d}{:<6d}{:<6d}{:<6.1f}{:<6.1f}\n'.format(self.atom_numbers[layers*mpl*p
                                             + l*mpl + m], self.atom_numbers[layers*mpl*p + (l + 1)*mpl + m], 6, b0, kb))

        # Tether together columns (this is temporary so variables are hard coded)
        b0 = 2*0.6*np.sin(36*np.pi/180)
        kb = 1000
        for p in range(pores):
            for l in range(layers):
                for m in range(mpl):
                    if m == (mpl - 1):  # handles periodicity
                        self.topology.append('{:<6d}{:<6d}{:<6d}{:<6.3f}{:<6.1f}\n'.format(self.atom_numbers[layers*mpl*p
                                             + l*mpl + m], self.atom_numbers[layers*mpl*p + l*mpl], 6, b0, kb))
                    else:
                        self.topology.append('{:<6d}{:<6d}{:<6d}{:<6.3f}{:<6.1f}\n'.format(self.atom_numbers[layers*mpl*p
                                             + l*mpl + m], self.atom_numbers[layers*mpl*p + l*mpl + m + 1], 6, b0, kb))

    def write_topology(self):
        with open('%s.itp' % self.name, 'w') as f:
            for line in self.topology:
                f.write(line)

if __name__ == "__main__":

    args = initialize()

    top = RestrainedTopology(args.gro, args.monomer, args.atoms, com=args.center_of_mass,
                             vparams=args.virtual_site_parameters)
    # top.add_distance_restraint_columns(float(args.bond_restraints[0]), float(args.bond_restraints[1]))
    top.add_position_restraints(args.axis, args.f_const)
    top.write_topology()


""" Clustered Hierarchical Entropy-scaling Manifold Mapping.

# TODO: https://docs.python.org/3/whatsnew/3.8.html#f-strings-support-for-self-documenting-expressions-and-debugging
"""
import logging
import pickle
import random
from collections import deque
from operator import itemgetter
from queue import Queue
from threading import Thread
from typing import Set, Dict, Iterable, BinaryIO, List, Union

import numpy as np
from scipy.spatial.distance import pdist, cdist

from chess.types import Data, Radius, Vector, Metric

SUBSAMPLE_LIMIT = 100
BATCH_SIZE = 10_000
LOG_LEVEL = logging.INFO

logging.basicConfig(
    level=LOG_LEVEL,
    format="%(asctime)s:%(levelname)s:%(name)s:%(module)s.%(funcName)s:%(message)s"
)

class Cluster:
    """ A cluster of points.

    Clusters maintain references to their their children, the manifold to which they belong,
    the indices of the points they are responsible for, and neighbors (clusters with which they overlap).

    You can compare clusters, hash them, partition them, perform tree search, prune them, and more.
    In general, they implement methods that create and utilize the underlying tree structure used by Manifold.
    """

    def __init__(self, manifold: 'Manifold', argpoints: Vector, name: str, **kwargs):
        """
        A Cluster needs to know the manifold it belongs to and the indexes of the points it contains.
        The name of a Cluster indicated its position in the tree.

        :param manifold: The manifold to which the cluster belongs.
        :param argpoints: A list of indexes of the points that belong to the cluster.
        :param name: The name of the cluster indicating its position in the tree.
        """
        logging.debug(f"Cluster(name={name}, argpoints={argpoints})")
        self.manifold: 'Manifold' = manifold
        self.argpoints: Vector = argpoints
        self.name: str = name

        # TODO: Consider relying on Graph.edges instead of having neighbors be a member of Cluster.
        self.neighbors: Dict['Cluster', float] = dict()  # key is neighbor, value is distance to neighbor
        self.children: Set['Cluster'] = set()

        self.__dict__.update(**kwargs)

        # This is used during Cluster.from_json().
        if not argpoints and self.children:
            self.argpoints = [p for child in self.children for p in child.argpoints]
        elif not argpoints:
            raise ValueError(f"Cluster {name} needs argpoints.")
        return

    def __eq__(self, other: 'Cluster') -> bool:
        """ Two clusters are identical if they have the same name and the same set of points. """
        return all((
            self.name == other.name,
            set(self.argpoints) == set(other.argpoints),
        ))

    def __hash__(self):
        """ Be careful to use this only with other clusters. """
        return hash(self.name)

    def __str__(self) -> str:
        return self.name or 'root'

    def __repr__(self) -> str:
        return ','.join([self.name, ';'.join(map(str, self.argpoints))])

    def __len__(self) -> int:
        """ Returns cardinality of the set of points.
        TODO: Consider deprecating __len__ and providing Cluster().cardinality
        """
        return len(self.argpoints)

    def __iter__(self) -> Vector:
        # Iterates in batches, instead of by element.
        for i in range(0, len(self), BATCH_SIZE):
            yield self.argpoints[i:i + BATCH_SIZE]

    def __contains__(self, point: Data) -> bool:
        """ Check weather the given point could be inside this cluster. """
        return self.overlaps(point=point, radius=0.)

    @property
    def metric(self) -> str:
        """ The metric used in the manifold. """
        return self.manifold.metric

    @property
    def depth(self) -> int:
        """ The depth in the tree at which the cluster exists. """
        return len(self.name)

    @property
    def points(self) -> Data:
        """ An iterator, in batches, over the points in the Clusters. """
        for i in range(0, len(self), BATCH_SIZE):
            yield self.manifold.data[self.argpoints[i:i + BATCH_SIZE]]

    @property
    def samples(self) -> Data:
        """ Returns the samples from the cluster. Samples are used in computing approximate centers and poles.
        """
        return self.manifold.data[self.argsamples]

    @property
    def argsamples(self) -> Vector:
        """ Indices used to retrieve samples.

        Ensures that there are at least 2 different points in samples,
        otherwise returns a single sample that represents the entire cluster.
        i.e., if len(argsamples) == 1, the cluster contains only duplicates.
        """
        if '_argsamples' not in self.__dict__:
            logging.debug(f"building cache for {self}")
            if len(self) <= SUBSAMPLE_LIMIT:
                n = len(self.argpoints)
                indices = self.argpoints
            else:
                n = int(np.sqrt(len(self)))
                indices = list(np.random.choice(self.argpoints, n, replace=False))

            # Handle Duplicates.
            if pdist(self.manifold.data[indices], self.metric).max(initial=0.) == 0.:
                indices = np.unique(self.manifold.data[self.argpoints], return_index=True, axis=0)[1]
                indices = [self.argpoints[i] for i in indices][:n]

            # Cache it.
            self.__dict__['_argsamples'] = indices
        return self.__dict__['_argsamples']

    @property
    def nsamples(self) -> int:
        """ The number of samples for the cluster. """
        return len(self.argsamples)

    @property
    def centroid(self) -> Data:
        """ The Geometric Mean of the cluster. """
        return np.average(self.samples, axis=0)

    @property
    def medoid(self) -> Data:
        """ The Geometric Median of the cluster. """
        return self.manifold.data[self.argmedoid]

    @property
    def argmedoid(self) -> int:
        """ The index used to retrieve the medoid. """
        if '_argmedoid' not in self.__dict__:
            logging.debug(f"building cache for {self}")
            _argmedoid = np.argmin(cdist(self.samples, self.samples, self.metric).sum(axis=1))
            self.__dict__['_argmedoid'] = self.argsamples[int(_argmedoid)]
        return self.__dict__['_argmedoid']

    @property
    def radius(self) -> Radius:
        """ The radius of the cluster.

        Computed as distance from medoid to the farthest point in the cluster.
        """
        if '_min_radius' in self.__dict__:
            logging.debug(f'taking min_radius from {self}')
            return self.__dict__['_min_radius']
        elif '_radius' not in self.__dict__:
            logging.debug(f'building cache for {self}')
            _ = self.argradius
        return self.__dict__['_radius']

    @property
    def argradius(self) -> int:
        """ The index used to retrieve the point which is farthest from the medoid. """
        if ('_argradius' not in self.__dict__) or ('_radius' not in self.__dict__):
            logging.debug(f'building cache for {self}')

            def argmax_max(b):
                distances = self.distance(self.manifold.data[b])
                argmax = int(np.argmax(distances))
                return b[argmax], distances[argmax]

            argradii_radii = [argmax_max(batch) for batch in iter(self)]
            _argradius, _radius = max(argradii_radii, key=itemgetter(1))
            self.__dict__['_argradius'], self.__dict__['_radius'] = int(_argradius), float(_radius)
        return self.__dict__['_argradius']

    @property
    def local_fractal_dimension(self) -> float:
        """ The local fractal dimension of the cluster. """
        # TODO: Consider computing by using search.
        if '_local_fractal_dimension' not in self.__dict__:
            logging.debug(f'building cache for {self}')
            if self.nsamples == 1:
                return 0.
            count = [d <= (self.radius / 2)
                     for batch in self
                     for d in self.distance(self.manifold.data[batch])]
            count = np.sum(count)
            self.__dict__['_local_fractal_dimension'] = count if count == 0. else np.log2(len(self.argpoints) / count)
        return self.__dict__['_local_fractal_dimension']

    def clear_cache(self) -> None:
        """ Clears the cache for the cluster. """
        logging.debug(f'clearing cache for {self}')
        for prop in ['_argsamples', '_argmedoid', '_argradius', '_radius', '_local_fractal_dimension']:
            try:
                del self.__dict__[prop]
            except KeyError:
                pass

    def tree_search(self, point: Data, radius: Radius, depth: int) -> Dict['Cluster', Radius]:
        """ Searches down the tree for clusters that overlap point with radius at depth. """
        logging.debug(f'tree_search(point={point}, radius={radius}, depth={depth}')
        if depth == -1:
            depth = len(self.manifold.graphs)
        if depth < self.depth:
            raise ValueError('depth must not be less than cluster.depth')  # TODO: Cover

        results: Dict['Cluster', Radius] = dict()
        if self.depth == depth:
            results = {self: self.distance(np.asarray([point]))[0]}  # TODO: Cover
        elif self.overlaps(point, radius):
            results = self._tree_search(point, radius, depth)

        return results

    def _tree_search(self, point: Data, radius: Radius, depth: int) -> Dict['Cluster', Radius]:
        distance = self.distance(np.asarray([point]))[0]
        assert distance <= radius + self.radius, f'_tree_search was started with no overlap.'
        assert self.depth < depth, f'_tree_search needs to have depth ({depth}) > self.depth ({self.depth}). '

        # results and candidates ONLY contain clusters that have overlap with point
        results: Dict['Cluster', Radius] = dict()
        candidates: Dict['Cluster', Radius] = {self: distance}
        for d_ in range(self.depth, depth):
            # if cluster was not partitioned any further, add it to results.
            results.update({c: d for c, d in candidates.items() if len(c.children) < 2})

            # filter out only those candidates that were partitioned.
            candidates = {c: d for c, d in candidates.items() if len(c.children) > 1}

            # proceed down th tree
            children: List[Cluster] = [c for candidate in candidates.keys() for c in candidate.children]
            if len(children) == 0:
                break

            # filter out clusters that are too far away to possibly contain any hits.
            centers = np.asarray([c.medoid for c in children])
            distances = cdist(np.expand_dims(point, 0), centers, self.metric)[0]
            radii = [radius + c.radius for c in children]
            candidates = {c: d for c, d, r in zip(children, distances, radii) if d <= r}
            if len(candidates) == 0:
                break  # TODO: Cover

        assert all((depth >= r.depth for r in results))
        assert all((depth == c.depth for c in candidates))

        # put all potential clusters in one dictionary.
        results.update(candidates)
        return results

    def partition(self, *criterion) -> Iterable['Cluster']:
        """ Partitions the cluster into 1 or 2 children.

        2 children are produced if the cluster can be split, otherwise 1 child is produced.
        """
        if not all((
                len(self.argpoints) > 1,
                len(self.argsamples) > 1,
                *(c(self) for c in criterion),
        )):
            # TODO: Can this be made more efficient? In the context of the larger manifold and graph
            logging.debug(f'{self} did not partition.')
            self.children = {
                Cluster(
                    self.manifold,
                    self.argpoints,
                    self.name + '0',
                    _argsamples=self.argsamples,
                    _argmedoid=self.argmedoid,
                    _argradius=self.argradius,
                    _radius=self.radius,
                    _local_fractal_dimension=self.local_fractal_dimension,
                )
            }
            return self.children

        farthest = self.argsamples[int(np.argmax(cdist(
            np.expand_dims(self.manifold.data[self.argradius], 0),
            self.samples,
            self.metric,
        )[0]))]
        poles = np.stack([
            self.manifold.data[self.argradius],
            self.manifold.data[farthest],
        ])

        p1_idx, p2_idx = list(), list()
        [(p1_idx if p1 < p2 else p2_idx).append(i)
         for batch in iter(self)
         for i, p1, p2 in zip(batch, *cdist(poles, self.manifold.data[batch], self.metric))]

        # Ensure that p1 contains fewer points than p2
        p1_idx, p2_idx = (p1_idx, p2_idx) if len(p1_idx) < len(p2_idx) else (p2_idx, p1_idx)
        self.children = {
            Cluster(self.manifold, p1_idx, self.name + '1'),
            Cluster(self.manifold, p2_idx, self.name + '2'),
        }
        logging.debug(f'{self} was partitioned.')
        return self.children

    def distance(self, points: Data) -> List[Radius]:
        """ Returns the distance from self.medoid to every point in points. """
        return cdist(np.expand_dims(self.medoid, 0), points, self.metric)[0]

    def overlaps(self, point: Data, radius: Radius) -> bool:
        """ Checks if point is within radius + self.radius of cluster. """
        return self.distance(np.expand_dims(point, axis=0))[0] <= (self.radius + radius)

    def json(self):
        data = {
            'name': self.name,
            'argpoints': None,  # Do not save argpoints until at leaves.
            'children': [],
            'neighbors': {c.name: d for c, d in self.neighbors.items()},
            '_radius': self.radius,
            '_argradius': self.argradius,
            '_argsamples': self.argsamples,
            '_argmedoid': self.argmedoid,
            '_local_fractal_dimension': self.local_fractal_dimension,
        }
        if self.children:
            data['children'] = [c.json() for c in self.children]
        else:
            data['argpoints'] = self.argpoints
        return data

    @staticmethod
    def from_json(manifold, data):
        children = set([Cluster.from_json(manifold, c) for c in data.pop('children', [])])
        return Cluster(manifold, children=children, **data)

class Graph:
    """ A Graph is comprised of clusters. All constituent clusters must be at the same depth in the tree.

    Nodes in the Graph are Clusters. .Two clusters have an edge if they have overlapping volumes.
    The Graph class is responsible for handling operations that occur solely within a layer of Manifold.graphs.
    """

    def __init__(self, *clusters):
        logging.debug(f'Graph(clusters={[str(c) for c in clusters]})')
        assert all(isinstance(c, Cluster) for c in clusters)
        assert all([c.depth == clusters[0].depth for c in clusters[1:]])

        # self.clusters is a dictionary of the clusters in the graph and the connected component subgraph that the cluster belongs to.
        self.clusters: Dict[Cluster: 'Graph'] = {c: None for c in clusters}
        return

    def __eq__(self, other: 'Graph') -> bool:
        """ Two graphs are identical if they are composed of the same clusters. """
        return self.clusters.keys() == other.clusters.keys()

    def __iter__(self) -> Iterable[Cluster]:
        """ An iterator over the clusters in the graph. """
        yield from self.clusters.keys()

    def __len__(self) -> int:
        # TODO: Consider deprecating __len__ for Graph().cardinality
        return len(self.clusters.keys())

    def __str__(self) -> str:
        return ';'.join(sorted([str(c) for c in self.clusters.keys()]))

    def __repr__(self) -> str:
        return '\t'.join(sorted([repr(c) for c in self.clusters.keys()]))

    def __hash__(self):
        return hash(str(self))

    def __contains__(self, cluster: 'Cluster') -> bool:
        return cluster in self.clusters.keys()

    @property
    def manifold(self) -> 'Manifold':
        return next(iter(self.clusters.keys())).manifold

    @property
    def depth(self) -> int:
        return next(iter(self.clusters.keys())).depth

    @property
    def metric(self) -> Metric:
        return next(iter(self.clusters.keys())).metric

    def _build_edges_matrix(self) -> None:
        """ Calculates overlap for clusters in self in the naive way. """
        # TODO: Calculate memory cost of the distance matrix here.
        clusters: List[Cluster] = list(self.clusters.keys())

        centers = np.asarray([c.medoid for c in clusters], dtype=np.float64)
        radii = np.asarray([c.radius for c in clusters], dtype=np.float64)

        distances = cdist(centers, centers, self.metric)
        differences = (distances.T - radii).T - radii
        left, right = tuple(map(list, np.where(differences <= 0.)))

        [clusters[l_].neighbors.update({clusters[r_]: distances[l_, r_]}) for l_, r_ in zip(left, right) if l_ != r_]
        return

    def build_edges(self) -> None:
        """ Calculates edges for the Graph. """
        return self._build_edges_matrix()

    @property
    def edges(self) -> Dict[Set['Cluster'], float]:
        """ Returns all edges within the graph. """
        if '_edges' not in self.__dict__:
            logging.debug(f'building cache for {self}')
            self.__dict__['_edges'] = {frozenset([c, n]): d for c in self.clusters.keys() for n, d in c.neighbors.items()}
        return self.__dict__['_edges']

    @property
    def subgraphs(self) -> Set['Graph']:
        """ Returns all subgraphs within the graph. """
        if any((s is None for s in self.clusters.values())):
            unvisited = {c for c, s in self.clusters.items() if s is None}
            while unvisited:
                cluster = unvisited.pop()
                component = self.bft(cluster)
                unvisited -= component
                subgraph = Graph(*component)
                self.clusters.update({c: subgraph for c in subgraph})
        return set(self.clusters.values())

    def subgraph(self, cluster: 'Cluster') -> 'Graph':
        """ Returns the subgraph to which the cluster belongs. """
        if cluster not in self.clusters.keys():
            raise ValueError(f'Cluster {cluster} not a member of {self}')

        if self.clusters[cluster] is None:
            component = self.bft(cluster)
            subgraph = Graph(*component)
            self.clusters.update({c: subgraph for c in subgraph})

        return self.clusters[cluster]

    def clear_cache(self) -> None:
        """ Clears the cache of the graph. """
        for prop in ['_edges']:
            logging.debug(str(self.clusters))
            try:
                del self.__dict__[prop]
            except KeyError:
                pass
        # Clear all cached subgraphs.
        self.clusters = {c: None for c in self.clusters.keys()}
        return

    def random_walk(self, steps: int = 5, walks: int = 1) -> Dict[Cluster, int]:
        """ Performs a random walk, returning a modified graph instance.

        :param int steps: number of steps per walk
        :param int walks: number of walks to perform
        :returns a Dict of cluster names to visit counts
        """
        # TODO: Consider changing the type of parallelism here to not have to rely on lists.
        clusters = list(self.clusters.keys())
        results = {c: list() for c in clusters}

        def walk(cluster):
            for _ in range(steps):
                results[cluster].append(1)
                if not cluster.neighbors:
                    break  # TODO: Cover
                cluster = random.sample(cluster.neighbors.keys(), 1)[0]

        # Perform random walks in parallel.
        starts = random.sample(clusters, min(walks, len(clusters)))
        threads = [Thread(target=walk, args=(s,)) for s in starts]
        [t.start() for t in threads]
        [t.join() for t in threads]

        # Gather the results.
        results = {k: len(v) for k, v in results.items()}
        return results

    @staticmethod
    def bft(start: 'Cluster'):
        """ Breadth-First Traversal starting at start. """
        logging.debug(f'starting from {start}')
        visited = set()
        queue = deque([start])
        while queue:
            c = queue.popleft()
            if c not in visited:
                visited.add(c)
                [queue.append(neighbor) for neighbor in c.neighbors.keys()]
        return visited

    @staticmethod
    def dft(start: 'Cluster'):
        """ Depth-First Traversal starting at start. """
        logging.debug(f'starting from {start}')
        visited = set()
        stack: List[Cluster] = [start]
        while stack:
            c = stack.pop()
            if c not in visited:
                visited.add(c)
                stack.extend(c.neighbors.keys())
        return visited

class Manifold:
    """ Manifold of varying resolution.

    The Manifold class' main job is to organize the underlying Clusters ang Graphs.
    It does this by providing the ability to reset the build the Cluster-tree, and from them the Graph-stack.
    With this Cluster-tree and Graph-stack, Manifold provides utilities for rho-nearest neighbors search, k-nearest neighbors search.
    """
    # TODO: Bring in anomaly detection from experiments.

    def __init__(self, data: Data, metric: Metric, argpoints: Union[Vector, float] = None, **kwargs):
        """ A Manifold needs the data to learn the manifold for, and a distance metric to use while doing so.

        :param data: The data to learn. This could be a numpy.ndarray or a numpy.memmap.
        :param metric: The distance metric to use for the data. Any metric allowed by scipy.spatial.distance is allowed here.
        :param argpoints: Optional. List of indexes or portion of data to which to restrict Manifold.
        """
        logging.debug(f'Manifold(data={data.shape}, metric={metric}, argpoints={argpoints})')
        self.data: Data = data
        self.metric: Metric = metric

        if argpoints is None:
            self.argpoints = list(range(self.data.shape[0]))
        elif type(argpoints) is list:
            self.argpoints = list(map(int, argpoints))
        elif type(argpoints) is float:
            self.argpoints = np.random.choice(self.data.shape[0], int(self.data.shape[0] * argpoints), replace=False)
            self.argpoints = list(map(int, self.argpoints))
        else:
            raise ValueError(f"Invalid argument to argpoints. {argpoints}")

        self.graphs: List['Graph'] = [Graph(Cluster(self, self.argpoints, ''))]

        self.__dict__.update(**kwargs)
        return

    def __eq__(self, other: 'Manifold') -> bool:
        """ Two manifolds are identical if they have the same metric and the same leaf-clusters. """
        return all((
            self.metric == other.metric,
            self.graphs[-1] == other.graphs[-1],
        ))

    def __getitem__(self, depth: int) -> 'Graph':
        return self.graphs[depth]

    def __iter__(self) -> Iterable[Graph]:
        yield from self.graphs

    def __str__(self) -> str:
        return '\t'.join([self.metric, str(self.graphs[-1])])

    def __repr__(self) -> str:
        return '\n'.join([self.metric, repr(self.graphs[-1])])

    @property
    def depth(self) -> int:
        return len(self.graphs) - 1

    def find_points(self, point: Data, radius: Radius) -> Dict[int, Radius]:
        """ Returns all indices of points that are within radius of point. """
        # TODO: Need a default depth argument?
        # TODO: Consider returning results as a sorted list of tuples.
        candidates: List[int] = [p for c in self.find_clusters(point, radius, len(self.graphs)).keys() for p in c.argpoints]
        results: Dict[int, Radius] = dict()
        point = np.expand_dims(point, axis=0)
        for i in range(0, len(candidates), BATCH_SIZE):
            batch = candidates[i:i + BATCH_SIZE]
            distances = cdist(point, self.data[batch], self.metric)[0]
            results.update({p: d for p, d in zip(batch, distances) if d <= radius})
        return results

    def find_clusters(self, point: Data, radius: Radius, depth: int) -> Dict['Cluster', Radius]:
        """ Returns all clusters that contain points within radius of point at depth. """
        return {r: d for c in self.graphs[0] for r, d in c.tree_search(point, radius, depth).items()}

    def find_knn(self, point: Data, k: int) -> Dict[Data, Radius]:
        """ Finds and returns the k-nearest neighbors of point. """
        # TODO: Consider returning results as a sorted list of tuples.
        radius: Radius = np.float64(np.mean([c.radius for c in self.graphs[-1].clusters]))
        radius = np.float64(max(radius, 1e-16))
        results = self.find_points(point, radius)
        while len(results.keys()) < k:
            radius *= 2
            results = self.find_points(point, radius)

        sorted_results = sorted([(d, p) for p, d in results.items()])[:k]
        results = {p: d for d, p in sorted_results}
        return results

    def build(self, *criterion) -> 'Manifold':
        """ Rebuilds the Cluster-tree and the Graph-stack. """
        self.graphs = [Graph(Cluster(self, self.argpoints, ''))]
        self.build_tree(*criterion)
        self.build_graphs()
        return self

    def build_tree(self, *criterion) -> 'Manifold':
        """ Builds the Cluster-tree. """
        while True:
            logging.info(f'current depth: {len(self.graphs) - 1}')
            clusters = self._partition_threaded(criterion)
            if len(self.graphs[-1]) < len(clusters):
                g = Graph(*clusters)
                self.graphs.append(g)
            else:
                [c.children.clear() for c in self.graphs[-1]]
                break
        return self

    def build_graphs(self) -> 'Manifold':
        """ Builds the Graph-stack. """
        [g.build_edges() for g in self.graphs]
        return self

    def build_graph(self, depth: int) -> 'Manifold':
        """ Builds the graph at a given depth. """
        if depth > self.depth:
            raise ValueError(f'depth must not be greater than {self.depth}. Got {depth}.')
        self.graphs[depth].build_edges()
        return self

    def subgraph(self, cluster: Union[str, Cluster]) -> Graph:
        """ Returns the subgraph to which cluster belongs. """
        cluster = self.select(cluster) if type(cluster) is str else cluster
        return self.graphs[cluster.depth].subgraph(cluster)

    def graph(self, cluster: Union[str, Cluster]) -> Graph:
        """ Returns the graph to which cluster belongs. """
        cluster = self.select(cluster) if type(cluster) is str else cluster
        return self.graphs[cluster.depth]

    def _partition_single(self, criterion):
        return [child for cluster in self.graphs[-1] for child in cluster.partition(*criterion)]

    def _partition_threaded(self, criterion):
        queue = Queue()
        threads = [
            Thread(
                target=lambda cluster: [queue.put(c) for c in cluster.partition(*criterion)],
                args=(c,),
                name=c.name
            )
            for c in self.graphs[-1]]
        [t.start() for t in threads]
        [t.join() for t in threads]
        clusters = []
        while not queue.empty():
            clusters.append(queue.get())
        return clusters

    def select(self, name: str) -> Cluster:
        """ Returns the cluster with the given name. """
        if len(name) > self.depth:
            raise ValueError(f'depth of requested cluster must not be greater than depth of cluster-tree. Got {name}, max-depth: {self.depth}')

        # TODO: Consider how to change this for forests.
        cluster: Cluster = next(iter(self.graphs[0]))
        for depth in range(len(name) + 1):
            partial_name = name[:depth]
            for child in cluster.children:
                if child.name == partial_name:
                    cluster = child
                    break

        assert name == cluster.name, f'wanted {name} but got {cluster.name}.'
        return cluster

    def dump(self, fp: BinaryIO) -> None:  # TODO: Cover
        # TODO: Consider hoe to remove argpoints from this and just rebuild from leaves.
        pickle.dump({
            'metric': self.metric,
            'argpoints': self.argpoints,
            'root': [c.json() for c in self.graphs[0]],
        }, fp)
        return

    @staticmethod
    def load(fp: BinaryIO, data: Data) -> 'Manifold':
        d = pickle.load(fp)
        manifold = Manifold(data, metric=d['metric'], argpoints=d['argpoints'])
        graphs = [  # TODO: Cover
            Graph(*[Cluster.from_json(manifold, r) for r in d['root']])
        ]
        while True:
            layer = Graph(*(child for cluster in graphs[-1] for child in cluster.children))
            if not layer:
                break
            else:
                graphs.append(layer)

        manifold.graphs = graphs
        for graph in graphs:
            for cluster in graph.clusters.keys():
                cluster.neighbors = {manifold.select(n): d for n, d in cluster.__dict__['neighbors'].items()}
        return manifold


#! /usr/bin/env python
import sys
import os
import argparse
import numpy as np
import pandas as pd
import json
from ctdcal import process_ctd
import ctdcal.report_ctd as report_ctd
import ctdcal.fit_ctd as fit_ctd
import configparser
#import matplotlib.pyplot as plt
from scipy.optimize import leastsq
import gsw

DEBUG = False

#File extension to use for output files (csv-formatted)
FILE_EXT = 'csv'
PKL_EXT = 'pkl'

#File extension to use for output files (csv-formatted)
XML_EXT = 'XMLCON'

#File extension to use for output files (csv-formatted)
HEX_EXT = 'hex'

#File extension to use for raw output
BTL_SUFFIX = '_btl'

#File extension to use for raw output
FIT_SUFFIX = '_fit'

#File extension to use for raw output
MEAN_SUFFIX = '_mean'

#File extension to use for raw output
TIME_SUFFIX = '_time'

#File extension to use for output files (csv-formatted)
UPDT_SUFFIX = '_updt'

#File extension to use for converted output
CONVERTED_SUFFIX = '_cnv'

def debugPrint(*args, **kwargs):
    if DEBUG:
        errPrint(*args, **kwargs)

def errPrint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)

# -------------------------------------------------------------------------------------
# Main function of the script should it be run as a stand-alone utility.
# -------------------------------------------------------------------------------------
def main(argv):

    parser = argparse.ArgumentParser(description='Convert SBE raw data to a converted, csv-formatted text file')
    parser.add_argument('timeFile', metavar='time_file', help='the .csv data file to fit by bottle data')

    # debug messages
    parser.add_argument('-d', '--debug', action='store_true', help='display debug messages')

    # raw output
    parser.add_argument('-pres', '--pressure', action='store_true', help='Fit pressure data')

    # raw output
    parser.add_argument('-temp', '--temperature', action='store_true', help='Fit temperture data')

    # raw output
    parser.add_argument('-cond', '--conductivity', action='store_true', help='Fit conductivity data')

    # raw output
    parser.add_argument('-oxy', '--oxygen', type=argparse.FileType('r'), help='return the oxygen data file')

    # Process Command-line args
    args = parser.parse_args()
    if args.debug:
        global DEBUG
        DEBUG = True
        debugPrint("Running in debug mode")

    # Verify hex file exists
    if not os.path.isfile(args.timeFile):
        errPrint('ERROR: Input time dependent .csv file:', args.timeFile, 'not found\n')
        sys.exit(1)

    # Used later for building output file names
    filename_ext = os.path.basename(args.timeFile) # original filename with ext
    filename_base = os.path.splitext(filename_ext)[0] # original filename w/o ext

    if '_' in filename_base:
        filename_base = filename_base.split('_')[0]

    #Import Cruise Configuration File
    iniFile = 'data/ini-files/configuration.ini'
    config = configparser.RawConfigParser()
    config.read(iniFile)

    #Initialise Configuration Parameters
    expocode = config['cruise']['expocode']
    sectionID = config['cruise']['sectionid']
    raw_directory = config['ctd_processing']['raw_data_directory']
    time_directory = config['ctd_processing']['time_data_directory']
    pressure_directory = config['ctd_processing']['pressure_data_directory']
    oxygen_directory = config['ctd_processing']['oxygen_directory']
    btl_directory = config['ctd_processing']['bottle_directory']
    o2flask_file = config['ctd_processing']['o2flask_file']
    log_directory = config['ctd_processing']['log_directory']
    sample_rate = config['ctd_processing']['sample_rate']
    search_time = config['ctd_processing']['roll_filter_time']
    ctd = config['ctd_processing']['ctd_serial']

    time_zone = config['inputs']['time_zone']
    p_col = config['analytical_inputs']['p']
    p_btl_col = config['inputs']['p']
    t_col = config['analytical_inputs']['t']
    t_btl_col = config['inputs']['t']
    t1_col = config['analytical_inputs']['t1']
    t1_btl_col = config['inputs']['t1']
    t2_col = config['analytical_inputs']['t2']
    t2_btl_col = config['inputs']['t2']
    c_col = config['analytical_inputs']['c']
    c1_col = config['analytical_inputs']['c1']
    c1_btl_col = config['inputs']['c1']
    c2_col = config['analytical_inputs']['c2']
    c2_btl_col = config['inputs']['c2']
    sal_col = config['analytical_inputs']['salt']
    sal_btl_col = config['inputs']['salt']
    btl_sal_col = config['analytical_inputs']['btl_salt']
    dov_col = config['analytical_inputs']['dov']
    dov_btl_col = config['inputs']['dov']
    dopl_col = config['analytical_inputs']['dopl']
    dopl_btl_col = config['inputs']['dopl']
    dopkg_col = config['analytical_inputs']['dopkg']
    btl_oxy_col = config['analytical_inputs']['btl_oxy']
    xmis_col = config['analytical_inputs']['xmis']
    fluor_col = config['analytical_inputs']['fluor']
    timedate = config['analytical_inputs']['datetime']
    lat_col = config['analytical_inputs']['lat']
    lat_btl_col = config['inputs']['lat']
    lon_col = config['analytical_inputs']['lon']
    lon_btl_col = config['inputs']['lon']
    reft_col = config['inputs']['reft']
    btl_num_col = config['inputs']['btl_num']

    #time_column_data = config['time_series_output']['data_names'].split(',')
    time_column_data = config['time_series_output']['data_output']
    time_column_names = config['time_series_output']['column_name'].split(',')
    time_column_units = config['time_series_output']['column_units'].split(',')
    time_column_format = config['time_series_output']['format']

    #pressure_column_data = config['time_series_output']['data_names'].split(',')
    p_column_data = config['pressure_series_output']['data'].split(',')
    p_column_names = config['pressure_series_output']['column_name'].split(',')
    p_column_units = config['pressure_series_output']['column_units'].split(',')
    p_column_format = config['pressure_series_output']['format']
    p_column_qual = config['pressure_series_output']['qual_columns'].split(',')
    p_column_one = list(config['pressure_series_output']['q1_columns'].split(','))

    #bottle_data outputs
    btl_dtype = config['bottle_series_output']['dtype']

    hexfileName = str(filename_base + '.' + HEX_EXT)
    hexfilePath = os.path.join(raw_directory, hexfileName)

    xmlfileName = str(filename_base + '.' + XML_EXT)
    xmlfilePath = os.path.join(raw_directory, xmlfileName)

    outtimefileName = str(filename_base + TIME_SUFFIX + '.' + PKL_EXT)
    outtimefilePath = os.path.join(time_directory, outtimefileName)

    # pressfileName = str(filename_base + FIT_SUFFIX + '.' + FILE_EXT)
    # pressfilePath = os.path.join(pressure_directory, pressfileName)

    btlfileName = str(filename_base + BTL_SUFFIX + MEAN_SUFFIX + '.' + PKL_EXT)
    btlfilePath = os.path.join(btl_directory, btlfileName)

    # Get bottle data
    btl_data = process_ctd.dataToNDarray(btlfilePath,float,True,',',0)
    #import pdb; pdb.set_trace()
    #btl_data = btl_data[:][1:]

    # Get procesed time data
    time_data = process_ctd.dataToNDarray(args.timeFile,float,True,',',1)
    #btm = np.argmax(time_data[p_col][1:])
    #time_data = time_data[:][1:btm]
    time_data = pd.DataFrame.from_records(time_data)
    time_data = time_data.loc[:time_data['CTDPRS'].idxmax()]
    time_data = time_data.to_records(index=False)

    if args.pressure:
        print('In -pres flag fit condition')
        print(filename_base)
        pfileName = str('poffset' + '.' + FILE_EXT)
        pfilePath = os.path.join(log_directory, pfileName)
        poff_data = process_ctd.dataToNDarray(pfilePath,str,None,',',None)

        for line in poff_data:
            if filename_base in line[0]:
                for val in line:
                    if 'offset' in val:
                        offset = float(str.split(val, ':')[1])
            continue

        # Pressure offset
        btl_data[p_btl_col] = fit_ctd.offset(offset, btl_data[p_btl_col])
        time_data[p_col] = fit_ctd.offset(offset, time_data[p_col])
        # End pressure if condition

    if args.temperature:
        print('In -temp flag fit condition')
        print(filename_base)
        coef1 = []
        coef2 = []
        # Get descrete ref temp data
        t1fileName = str('fitting_t1' + '.' + FILE_EXT)
        t1filePath = os.path.join(log_directory, t1fileName)
        t1_coef = process_ctd.dataToNDarray(t1filePath,str,None,',',None)

        for line in t1_coef:
            if filename_base in line[0]:
                val = line[1:]
                for i in range(0,len(val)):
                    if 'coef' in val[i]:
                        coef1.append(float(str.split(val[i], ':')[1]))
                    else:
                        coef1.append(float(val[i]))
            continue
        btl_data[t1_btl_col] = fit_ctd.temperature_polyfit(coef1, btl_data[p_btl_col], btl_data[t1_btl_col])
        time_data[t1_col] = fit_ctd.temperature_polyfit(coef1, time_data[p_col], time_data[t1_col])

        t2fileName = str('fitting_t2' + '.' + FILE_EXT)
        t2filePath = os.path.join(log_directory, t2fileName)
        t2_coef = process_ctd.dataToNDarray(t2filePath,str,None,',',None)

        for line in t2_coef:
            if filename_base in line[0]:
                val = line[1:]
                for i in range(0,len(val)):
                    if 'coef' in val[i]:
                        coef2.append(float(str.split(val[i], ':')[1]))
                    else:
                        coef2.append(float(val[i]))
            continue
        btl_data[t2_btl_col] = fit_ctd.temperature_polyfit(coef2, btl_data[p_btl_col], btl_data[t2_btl_col])
        time_data[t2_col] = fit_ctd.temperature_polyfit(coef2, time_data[p_col], time_data[t2_col])

    if args.conductivity:
        print('In -cond flag fit condition')
        print(filename_base)
        coef1 = []
        coef2 = []

        # Get descrete cond data

        c1fileName = str('fitting_c1' + '.' + FILE_EXT)
        c1filePath = os.path.join(log_directory, c1fileName)
        if os.path.exists(c1filePath):
            c1_coef = process_ctd.dataToNDarray(c1filePath,str,None,',',None)

            for line in c1_coef:
                if filename_base in line[0]:
                    val = line[1:]
                    for i in range(0,len(val)):
                        if 'coef' in val[i]:
                            coef1.append(float(str.split(val[i], ':')[1]))
                        else:
                            coef1.append(float(val[i]))
                continue
            btl_data[c1_btl_col] = fit_ctd.conductivity_polyfit(coef1, btl_data[p_btl_col], btl_data[t1_btl_col], btl_data[c1_btl_col])
            time_data[c1_col] = fit_ctd.conductivity_polyfit(coef1, time_data[p_col], time_data[c1_col], time_data[c1_col])

        c2fileName = str('fitting_c2' + '.' + FILE_EXT)
        c2filePath = os.path.join(log_directory, c2fileName)
        if os.path.exists(c2filePath):
            c2_coef = process_ctd.dataToNDarray(c2filePath,str,None,',',None)

            for line in c2_coef:
                if filename_base in line[0]:
                    val = line[1:]
                    for i in range(0,len(val)):
                        if 'coef' in val[i]:
                            coef2.append(float(str.split(val[i], ':')[1]))
                        else:
                            coef2.append(float(val[i]))
                continue
            btl_data[c2_btl_col] = fit_ctd.conductivity_polyfit(coef2, btl_data[p_btl_col], btl_data[t2_btl_col], btl_data[c2_btl_col])
            time_data[c2_col] = fit_ctd.conductivity_polyfit(coef2, time_data[p_col], time_data[c2_col], time_data[c2_col])

        time_data[sal_col] = gsw.SP_from_C(time_data[c_col],time_data[t_col],time_data[p_col])

    if args.oxygen:
        print('In -oxy flag fit condition')
        print(filename_base)
        # Get Analytical Oxygen data
        o2pkg_btl, o2pl_btl = fit_ctd.o2_calc(o2flask_file,args.oxygen.name,btl_data[btl_num_col],btl_data[sal_btl_col])

        kelvin = []
        for i in range(0,len(time_data[t_col])):
            kelvin.append(time_data[t_col][i] + 273.15)

        # Find New Oxygen Coef
        oxy_coef = fit_ctd.find_oxy_coef(o2pl_btl['OXYGEN'], btl_data[p_btl_col], btl_data[t_btl_col], btl_data[sal_btl_col], btl_data[dov_btl_col], hexfilePath, xmlfilePath)

        # Convert CTD Oxygen Voltage Data with New DO Coef
        time_data[dopl_col] = fit_ctd.oxy_dict(oxy_coef, time_data[p_col], kelvin, time_data[t_col], time_data[sal_col], time_data[dov_col])
        # End oxygen flag fitting if condition

    # Find Isopycnal Down Trace Bottle Trip Equivalent
    # Write bottle data to file
    report_ctd.report_btl_data(btlfilePath, btl_dtype, btl_data)

    # Write time data to file
    report_ctd.report_time_series_data(filename_base, time_directory, expocode, time_column_names, time_column_units, time_column_names, time_column_format, time_data)

    # Pressure Sequence
    pressure_seq_data = process_ctd.pressure_sequence(filename_base, p_col, timedate, 2.0, -1.0, 0.0, 'down', int(sample_rate), int(search_time), time_data)

    # Convert dissolved oxygen from ml/l to umol/kg
    dopkg = process_ctd.o2pl2pkg(p_col, t1_col, sal_col, dopl_col, dopkg_col, lat_col, lon_col, pressure_seq_data)

    # Add quality codes to data
    qual_pseq_data = process_ctd.ctd_quality_codes(dopkg_col, None, None, True, p_column_qual, p_column_one, pressure_seq_data)

    # Collect Cast Details from Log
    logfileName = str('cast_details' + '.' + FILE_EXT)
    logfilePath = os.path.join(log_directory, logfileName)

    cast_details = process_ctd.dataToNDarray(logfilePath,str,None,',',0)
    for line in cast_details:
        if filename_base in line[0]:
            for val in line:
                if 'at_depth' in val: btime = float(str.split(val, ':')[1])
                if 'latitude' in val: btm_lat = float(str.split(val, ':')[1])
                if 'longitude' in val: btm_lon = float(str.split(val, ':')[1])
                if 'altimeter_bottom' in val: btm_alt = float(str.split(val, ':')[1])
            break

    # Write time data to file
    depth = -999
    #import pdb; pdb.set_trace()
    report_ctd.report_pressure_series_data(filename_base, expocode, sectionID, btime, btm_lat, btm_lon, depth, btm_alt, ctd, pressure_directory, p_column_names, p_column_units, p_column_data, qual_pseq_data, dopkg, pressure_seq_data)

    #plt.plot(o2pl_btl['OXYGEN'], btl_data[p_btl_col], color='b', marker='o')
    #plt.plot(tmpo2, time_data[p_col], color='g', label='raw')
    #plt.plot(time_data[dopl_col], time_data[p_col], color='b', label='raw')
    #plt.plot(pressure_seq_data[dopl_col], pressure_seq_data[p_col], color='r', label='raw')
    #plt.plot(pressure_seq_data[t2_col], pressure_seq_data[p_col], color='r', label='raw')
    #plt.plot(btl_data[t_btl_col], time_data[p_col], color='r', label='raw')
    #plt.plot(o2pl_btl['OXYGEN'], btl_data[p_btl_col], color='g', marker='o', label='raw')
    #plt.plot(pressure_seq_data[do_col], pressure_seq_data[p_col], color='b', label='raw')
    #plt.gca().invert_yaxis()
    #plt.axis()
    #plt.show()

    debugPrint('Done!')

# -------------------------------------------------------------------------------------
# Required python code for running the script as a stand-alone utility
# -------------------------------------------------------------------------------------
if __name__ == '__main__':
    main(sys.argv[1:])


import os
import numpy as np
import pandas as pd

from excalibur.constants import Ryd, u

def det_broad(input_directory):
    '''
    Determine the type of broadening that should be used in the case that the user specifies
    'default' broadening. Order of preference is: 1) H2-He, 2) air, 3) SB07

    Parameters
    ----------
    input_directory : String
        Local directory where the broadening file will be stored.

    Returns
    -------
    broadening : String
        The type of broadening being used.

    '''

    # Molecular hydrogen + helium broadening is first choice
    if 'H2.broad' in os.listdir(input_directory) and 'He.broad' in os.listdir(input_directory):
        broadening = 'H2-He'

    # If no H2 + He boradening files, search for an air boradening file
    elif 'air.broad' in os.listdir(input_directory):
        broadening = 'air'

    # If neither of the above are available (e.g. for most metal oxides), fall back to Sharp & Burrows (2007)
    else:
        broadening = 'SB07'
        if not 'SB07.broad' in os.listdir(input_directory):
            create_SB07(input_directory)

    return broadening

def create_SB07(input_directory):
    '''
    Create a broadening file according to Eq. 15 of Sharp & Burrows (2007),
    and add it to the input_directory.

    Note: S&B (2007) state Eq. 15 gives the FWHM. Personal communication from
    Richard Freedman indicates the equation actually gives the HWHM.

    Parameters
    ----------
    input_directory : String
         Local directory where the broadening file will be stored.

    Returns
    -------
    None.

    '''

    SB07_file = input_directory + 'SB07.broad'

    # Initialise total arrays
    J = np.arange(31.0)          # Total angular momentum
    gamma_L_0 = np.zeros(31)     # Lorentizian HWHM at P_ref and T_ref
    n_L = np.zeros(31)           # Temperature exponent

    # Implement Eq. 15 from S&B07 (wihout division by 2, as already HWHM)
    for i in range(len(J)):
        gamma_L_0[i] = (0.1 - min(J[i], 30) * 0.002) / 1.01325

    # Write broadening output file
    f_out = open(SB07_file, 'w')

    f_out.write('J | gamma_L_0 | n_L \n')

    for i in range(len(J)):
        f_out.write('%.1f %.4f %.3f \n' %(J[i], gamma_L_0[i], n_L[i]))

    f_out.close()

def read_H2_He(input_directory):
    '''
    Read the H2 and He broadening files from the input directory

    Parameters
    ----------
    input_directory : TYPE
        DESCRIPTION.

    Returns
    -------
    J_max : TYPE
        DESCRIPTION.
    gamma_0_H2 : TYPE
        DESCRIPTION.
    n_L_H2 : TYPE
        DESCRIPTION.
    gamma_0_He : TYPE
        DESCRIPTION.
    n_L_He : TYPE
        DESCRIPTION.

    '''

    # Read in H2 broadening file
    broad_file_H2 = pd.read_csv(input_directory + 'H2.broad',
                                sep = ' ', header=None, skiprows=1)
    J_max_H2 = int(np.max(np.array(broad_file_H2[0])))
    gamma_0_H2 = np.array(broad_file_H2[1])
    n_L_H2 = np.array(broad_file_H2[2])

    # Read in He broadening file
    broad_file_He = pd.read_csv(input_directory + 'He.broad',
                                sep = ' ', header=None, skiprows=1)
    J_max_He = int(np.max(np.array(broad_file_He[0])))
    gamma_0_He = np.array(broad_file_He[1])
    n_L_He = np.array(broad_file_He[2])

    # Take maximum J'' value for which broadening is a function of J to be lowest for which complete data available
    J_max = np.max(np.array([J_max_H2, J_max_He]))

    # If broadening files not of same length, extend shortest file to same length as longest
    if (J_max_H2 < J_max):

        for i in range (J_max_H2, J_max):

            gamma_0_H2 = np.append(gamma_0_H2, gamma_0_H2[-1])    # Extended values equal to final value
            n_L_H2 = np.append(n_L_H2, n_L_H2[-1])                # Extended values equal to final value

    if (J_max_He < J_max):

        for i in range (J_max_He, J_max):

            gamma_0_He = np.append(gamma_0_He, gamma_0_He[-1])    # Extended values equal to final value
            n_L_He = np.append(n_L_He, n_L_He[-1])                # Extended values equal to final value

    return J_max, gamma_0_H2, n_L_H2, gamma_0_He, n_L_He

def read_air(input_directory):
    '''
    Read the air broadening file from the input directory

    Parameters
    ----------
    input_directory : String
         Local directory where the broadening file will be stored.

    Returns
    -------
    J_max : TYPE
        DESCRIPTION.
    gamma_0_air : TYPE
        DESCRIPTION.
    n_L_air : TYPE
        DESCRIPTION.

    '''

    # Read in air broadening file
    broad_file_air = pd.read_csv(input_directory + 'air.broad',
                                 sep = ' ', header=None, skiprows = 1)
    J_max = int(np.max(np.array(broad_file_air[0])))
    gamma_0_air = np.array(broad_file_air[1])
    n_L_air = np.array(broad_file_air[2])

    return J_max, gamma_0_air, n_L_air

def read_SB07(input_directory):
    '''
    Read the Burrows broadening file from the input directory

    Parameters
    ----------
    input_directory : TYPE
        DESCRIPTION.

    Returns
    -------
    J_max : TYPE
        DESCRIPTION.
    gamma_0_SB07 : TYPE
        DESCRIPTION.

    '''

    # Read in Sharp & Burrows (2007) broadening file
    broad_file_SB07 = pd.read_csv(input_directory + 'SB07.broad',
                                           sep = ' ', header=None, skiprows=1)
    J_max = int(np.max(np.array(broad_file_SB07[0])))
    gamma_0_SB07 = np.array(broad_file_SB07[1])
    #n_L_SB07 = np.array(broad_file_SB07[2])       # Not really needed, as temperature exponent = 0 for all J''

    return J_max, gamma_0_SB07

def read_custom(input_directory):
    '''
    Read a user-provided broadening file from the input directory

    Parameters
    ----------
    input_directory : TYPE
        DESCRIPTION.

    Returns
    -------
    J_max : TYPE
        DESCRIPTION.
    gamma_0_air : TYPE
        DESCRIPTION.
    n_L_air : TYPE
        DESCRIPTION.

    '''

    # Read in custom broadening file
    broad_file_custom = pd.read_csv(input_directory + 'custom.broad',
                                    sep = ' ', header=None, skiprows = 1)
    J_max = int(np.max(np.array(broad_file_custom[0])))
    gamma_0_air = np.array(broad_file_custom[1])
    n_L_air = np.array(broad_file_custom[2])

    return J_max, gamma_0_air, n_L_air

def gamma_L_VALD(gamma_vdw, m_s, broadener):
    '''
    Computes Lorentzian HWHM at 296K and 1 bar for a given broadener from a
    tabulated VALD van der Waals broadening constant.

    Parameters
    ----------
    gamma_vdw : float
        van der Waals parameter from VALD.
    m_s : float
        mass of species whose spectral line is being broadened (u).
    broadener : str
        identity of broadening species (H2 or He).

    Returns
    -------
    gamma_L_0 : float
        Lorentzian HWHM at reference T (296K) and P (1 bar).
    n_L : float
        Temperature exponent (fixed to -0.7 for van der Waals theory).

    '''

    alpha_H = 0.666793  # Polarisability of atomic hydrogen (A^-3)
    m_H = 1.007825      # Mass of atomic hydrogen (u)

    if (broadener == 'H2'):
        alpha_p = 0.805000  # Polarisability of molecular hydrogen (A^-3)
        m_p = 2.01565       # Mass of molecular hydrogen (u)

    elif (broadener == 'He'):
        alpha_p = 0.205052  # Polarisability of helium (A^-3)
        m_p = 4.002603      # Mass of helium (u)

    else: print ("Invalid broadener for VALD!")

    # Compute Lorentzian HWHM
    gamma_L_0 = (2.2593427e7 * gamma_vdw * np.power(((m_H*(m_s+m_p))/(m_p*(m_s+m_H))), (3.0/10.0)) *
                                           np.power((alpha_p/alpha_H), (2.0/5.0)))

    # Temperature exponent
    n_L = 0.7

    return gamma_L_0, n_L

def gamma_L_impact(E_low, E_up, l_low, l_up, species, m_s, broadener):
    '''
    Computes Lorentzian HWHM at 296K and 1 bar for a given broadener using
    van der Waals impact theory.

    Parameters
    ----------
    E_low : float
        Lower level energy (cm^-1).
    E_up : float
        Upper level energy (cm^-1).
    l_low : int
        Lower level orbital angular momentum.
    l_up : int
        Upper level orbital angular momentum.
    species : str
        Identity of species whose spectral line is being broadened.
    m_s : float
        Mass of species whose spectral line is being broadened (u).
    broadener : str
        Identity of broadening species (H2 or He).

    Returns
    -------
    gamma_L_0 : float
        Lorentzian HWHM at reference T (296K) and P (1 bar).
    n_L : float
        Temperature exponent (fixed to -0.7 for van der Waals theory).

    '''

    alpha_H = 0.666793  # Polarisability of atomic hydrogen (A^-3)

    E_inf_eV = {'Li': 5.3917, 'Na': 5.1391, 'K': 4.3407, 'Rb': 4.1771, 'Cs': 3.8939}    # Ionisation energy (eV)
    E_inf = E_inf_eV[species] * 8065.547574991239  # convert from eV to cm^-1

    if (species in ['Li', 'Na','K', 'Cs', 'Rb']):
        Z = 0   # Ion charge

    if (broadener == 'H2'):
        alpha_p = 0.805000  # Polarisability of molecular hydrogen (A^-3)
        m_p = 2.01565       # Mass of molecular hydrogen (u)

    elif (broadener == 'He'):
        alpha_p = 0.205052  # Polarisability of helium (A^-3)
        m_p = 4.002603      # Mass of helium (u)

    else: print ("Invalid broadener for VALD!")

    # Evaluate effective principal quantum number for lower and upper levels
    n_low_sq = ((Ryd/100.0) * (Z + 1.0)**2)/(E_inf - E_low)
    n_up_sq =  ((Ryd/100.0) * (Z + 1.0)**2)/(E_inf - E_up)

    # Evaluate mean square orbital radius for lower and upper levels (in Bohr radii)
    r_low_sq = (n_low_sq/(2.0*(Z+1)**2)) * (5.0*n_low_sq + 1.0 - (3.0*l_low*(l_low + 1.0)))
    r_up_sq = (n_up_sq/(2.0*(Z+1)**2)) * (5.0*n_up_sq + 1.0 - (3.0*l_up*(l_up + 1.0)))

    # For lines where the Hydrogenic approximation breaks down, return zero vdw line width
    if (r_up_sq < r_low_sq):
        return 0.0, 0.7   # Reference HWHM | temperature exponent (dummy in this case)

    # Compute Lorentzian HWHM
    gamma_L_0 = (0.1972 * np.power(((m_s+m_p)/(m_s*m_p)), (3.0/10.0)) *
                          np.power((alpha_p/alpha_H), (2.0/5.0)) *
                          np.power((r_up_sq - r_low_sq), (2.0/5.0)))

    # Temperature exponent
    n_L = 0.7

    return gamma_L_0, n_L

def read_atom(species, nu_0, gf, E_low, E_up, J_low, l_low, l_up,
              gamma_nat, gamma_vdw, alkali, m):

    if alkali:  # Special treatments for alkali Van der Waals widths

        gamma_0_H2 = np.zeros(len(nu_0))
        gamma_0_He = np.zeros(len(nu_0))
        n_L_H2 = np.zeros(len(nu_0))
        n_L_He = np.zeros(len(nu_0))

        for i in range(len(nu_0)):

            if (gamma_vdw[i] != 0.0):  # For transitions with a VALD broadening value

                gamma_0_H2[i], n_L_H2[i] = gamma_L_VALD(gamma_vdw[i], (m/u), 'H2')
                gamma_0_He[i], n_L_He[i] = gamma_L_VALD(gamma_vdw[i], (m/u), 'He')

            elif (gamma_vdw[i] == 0.0):  # For transitions without a VALD broadening value

                gamma_0_H2[i], n_L_H2[i] = gamma_L_impact(E_low[i], E_up[i], l_low[i], l_up[i], species, (m/u), 'H2')
                gamma_0_He[i], n_L_He[i] = gamma_L_impact(E_low[i], E_up[i], l_low[i], l_up[i], species, (m/u), 'He')

    else:  # For non-alkali species

        gamma_0_H2, n_L_H2 = gamma_L_VALD(gamma_vdw, (m/u), 'H2')
        gamma_0_He, n_L_He = gamma_L_VALD(gamma_vdw, (m/u), 'He')

    return gamma_0_H2, n_L_H2, gamma_0_He, n_L_He

def compute_H2_He(gamma_0_H2, T_ref, T, n_L_H2, P, P_ref, X_H2, gamma_0_He, n_L_He, X_He):
    gamma = (gamma_0_H2 * np.power((T_ref/T), n_L_H2) * (P/P_ref) * X_H2 +   # H2+He Lorentzian HWHM for given T, P, and J (ang. mom.)
             gamma_0_He * np.power((T_ref/T), n_L_He) * (P/P_ref) * X_He)    # Note that these are only a function of J''

    return gamma

def compute_air(gamma_0_air, T_ref, T, n_L_air, P, P_ref):
    gamma = (gamma_0_air * np.power((T_ref/T), n_L_air) * (P/P_ref))      # Air-broadened Lorentzian HWHM for given T, P, and J (ang. mom.)

    return gamma

def compute_SB07(gamma_0_SB07, P, P_ref):
    gamma = (gamma_0_SB07 * (P/P_ref))      # Equation (15) in Sharp & Burrows (2007)

    return gamma


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

'''
QSDsan: Quantitative Sustainable Design for sanitation and resource recovery systems
Copyright (C) 2020, Quantitative Sustainable Design Group

This module is developed by:
    Lewis Rowles <stetsonsc@gmail.com>

This module is under the University of Illinois/NCSA Open Source License.
Please refer to https://github.com/QSD-Group/QSDsan/blob/master/LICENSE.txt
for license details.
'''
# %%

import numpy as np
from qsdsan import SanUnit, Construction
from qsdsan.utils.loading import load_data, data_path
import os
__all__ = ('StruvitePrecipitation',)

#path to csv with all the inputs
#data_path = '/Users/lewisrowles/opt/anaconda3/lib/python3.8/site-packages/exposan/biogenic_refinery/_struvite_precipitation.csv'
data_path += 'sanunit_data/_struvite_precipitation.tsv'

###
class StruvitePrecipitation(SanUnit):
    '''
    Stuvite Precipitation for P recovery from liquid stream. Solid stuvite is recovered.

    Reference documents
    -------------------
    N/A

    Parameters
    ----------
    ins : WasteStream (liquid), MagnesiumHydroxide, MagnesiumCarbonate, FilterBag

    outs : WasteStream, Struvite

    References
    ----------
    .. Lohman et al., Advancing Sustainable Sanitation and Agriculture
    through Investments in Human-Derived Nutrient Systems.
    Environ. Sci. Technol. 2020, 54, (15), 9217-9227.
    https://dx.doi.org/10.1021/acs.est.0c03764

    .. Tarpeh et al., Evaluating ion exchange for nitrogen recovery from
    source-separated urine in Nairobi, Kenya. Development Engineering. 2018,
    3, 188–195.
    https://doi.org/10.1016/j.deveng.2018.07.002

    '''

    def __init__(self, ID='', ins=None, outs=(), **kwargs):

        SanUnit.__init__(self, ID, ins, outs)

# load data from csv each name will be self.name
        data = load_data(path=data_path)
        for para in data.index:
            value = float(data.loc[para]['expected'])
            setattr(self, para, value)
        del data

        for attr, value in kwargs.items():
            setattr(self, attr, value)

# define the number of influent and effluent streams
    _N_ins = 4
    _N_outs = 2

# in _run: define influent and effluent streams and treatment processes
    def _run(self):
        waste, magnesium_hydroxide, magnesium_carbonate, bag_filter = self.ins
        treated, struvite = self.outs
        treated.copy_like(self.ins[0])
        magnesium_hydroxide.phase = 's'
        magnesium_carbonate.phase = 's'
        bag_filter = 's'
        struvite.phase = 's'

        # recovery of N, K, P
        N_recovered = (waste.imass['P'] * self.P_rec_1 / self.MW_P
                       * self.N_P_ratio_struvite * self.MW_N)  # (kg N/hr) total N recovered
        P_recovered = waste.imass['P'] * self.P_rec_1  # (kg P/hr) total P recovered
        K_recovered = waste.imass['K'] * self.K_rec_1 # (kg K/hr) total K recovered
        treated.imass['NH3'] =  waste.imass['NH3'] - N_recovered # kg N / hr
        treated.imass['P'] =  waste.imass['P'] - P_recovered # kg P / hr
        treated.imass['K'] =  waste.imass['K'] - K_recovered # kg K / hr

        # set values needed for _design and _cost as attributes
        self.volume_treated = waste.F_vol * 1000 * 24 # L liq / d
        self.quantity_tanks = np.ceil(self.volume_treated / self.cycles_per_day
                                      / self.reactor_volume) # number of tanks

        # !!! add MagnesiumHydroxide, MagnesiumCarbonate, FilterBag as influent streams and Struvite as effluent
        magnesium_hydroxide_demand_time = (waste.imass['P'] / self.MW_P / self.Mg_dose
                                           * self.Mg_MgOH2_ratio * self.MW_MgOH2) # (kg Mg(OH)2 per hr)
        magnesium_hydroxide.imass['MagnesiumHydroxide'] = magnesium_hydroxide_demand_time

        magnesium_carbonate_demand_time = (waste.imass['P'] / self.MW_P / self.Mg_dose
                                           * self.Mg_MgCO3_ratio * self.MW_MgCO3)  # (kg MgCO3 per hr)
        magnesium_carbonate.imass['MagnesiumCarbonate'] = magnesium_carbonate_demand_time

        struvite_production_time = P_recovered / self.MW_P * self.MW_struvite # kg (NH4)MgPO4•6(H2O) / hr
        struvite.imass['Struvite'] = struvite_production_time

        filter_bag_demand_time = self.quantity_tanks * self.cycles_per_day / self.filter_reuse / 24 # bags/hr
        self.ins[3].imass['FilterBag'] = filter_bag_demand_time # used in place of line below, not sure why its not working
        #bag_filter.imass['FilterBag'] = filter_bag_demand_time

    #_design will include all the construction or captial impacts
    def _design(self):
        design = self.design_results
        # defining the quantities of materials/items
        # note that these items to be to be in the _impacts_items.xlsx

        design['StainlessSteel'] = SS_quant = self.quantity_tanks * self.reactor_weight # kg SS
        design['PVC'] = PVC_quant = self.quantity_tanks * self.material_P_pipe * self.pvc_mass  # kg PVC

        self.construction = (
            Construction(item='StainlessSteel', quantity = SS_quant, quantity_unit = 'kg'),
            Construction(item='PVC', quantity = PVC_quant, quantity_unit = 'kg'),
            )
        self.add_construction()

    #_cost based on amount of steel and stainless plus individual components
    def _cost(self):
        #purchase_costs is used for capital costs
        #can use quantities from above (e.g., self.design_results['StainlessSteel'])
        #can be broken down as specific items within purchase_costs or grouped (e.g., 'Misc. parts')
        self.purchase_costs['Reactor'] = (self.quantity_tanks * self.cost_P_reactor)
        self.purchase_costs['Stirrer'] = (self.quantity_tanks * self.cost_P_stirrer)
        self.purchase_costs['PVC'] = (self.quantity_tanks * self.material_P_pipe * self.cost_P_pipe)

        self._BM = dict.fromkeys(self.purchase_costs.keys(), 1)

        #certain parts need to be replaced based on an expected lifefime
        #the cost of these parts is considered along with the cost of the labor to replace them
        struvite_replacement_parts_annual_cost = 0 # USD/yr only accounts for time running

        struvite_annual_maintenance = 0 #USD/yr only accounts for time running

        self.add_OPEX =  (struvite_replacement_parts_annual_cost + struvite_annual_maintenance) / (365 * 24) # USD/hr (all items are per hour)

        # costs associated with full time opperators can be added in the TEA as staff