#A library of code to examine properties of bulk water and near solutes
#
#Should eventually be able to handle local densities and fluctuations,
#solute-water and water-water energies, 3-body angles, hydrogen bonds,
#energy densities, and all of this as a function of space. Additionally,
#should also be able to compute interfaces, such as Willard-Chandler
#instantaneous interface, or vdW surface, SASA and volume of solute.
#
#Will work with pytraj interface for trajectory analysis, since this
#should later allow easier energy decomposition?
#If doesn't work out, will go back to sim package with netcdf plugin.
#
#Also, should have test script and some test system where know answers
#

import sys, os
import numpy as np
import scipy.optimize as optimize
from scipy.special import sph_harm
import waterlib as wl

#Define constants and unit conversions

#conversion for surface tension
kBJ = 1.38064852*(10**(-23))
temp = 300.0
tomJm2 = kBJ*temp*1000.0*(10**20) #converts kBT/Angstrom^2 to mJ/m^2

#Convert potential energy to kBT
kBTkcal = 0.0019858775*300.0

#Water density
watdens = 0.033456 # molecules or oxygens per Angstrom ^ 3 near 300 K

#Define library of useful functions

def SASAperAtom(pos, radii, radius=1.4, nPoints = 1000, nExpose = 10):
  """Inputs:
     pos - Nx3 array of atomic positions
     radii - N array of atomic radii
     radius - solvent radius to "roll" over surface
     nPoints - number points on each sphere
     nExpose - number exposed points on atom (sphere) to be considered on surface
     Outputs:
     SASAper - SASA for each atom
     surfAtoms - array of 1 for solvent exposed, 0 for not on surface
  """

  points = wl.spherepoints(nPoints)
  SASAper, surfAtoms = wl.spheresurfaceareas(pos, radii+radius, points, nExpose)

  return SASAper, surfAtoms

def PepWatHBonds(allPos, pepAccInds, pepDonInds, watInds, distCut = 2.1, angCut = 30.0):
  """Currently kind of wack (does acceptor to hydrogen distance). Also, calculating
H-bonds geometrically seems less useful.
     Inputs:
     allPos - full position array for trajectory frame (all atoms included)
     pepAccInds - global indices of peptide acceptors
     pepDonInds - global indices of peptide donors
     watInds - global indices of water atoms in selected hydration shell(s)
     distCut(=2.1) - distance cutoff for H-bond detection
     angCut(=30.0) - angle cutoff for H-bond detection
     Outputs:
     NBonds - number of detected H-bonds
     bondsPer - number H-bonds for each water molecule with peptide
     donors - indices of donors (H atoms only) as string
     acceptors - indices of acceptors as string
  """

  #Get H-bond info
  NBonds, watAcc, watDon, pepAcc, pepDon = wl.findhbonds(
                             allPos[pepAccInds], allPos[pepDonInds], allPos[watInds], distCut, angCut)

  #And sort nicely into just acceptors and donors
  acceptorsList = []
  donorsList = []
  bondsWat = np.zeros(int(len(watInds)/3))

  for (j, val) in enumerate(pepAcc):
    acceptorsList = acceptorsList + (val*[pepAccInds[j]])

  for (j, val) in enumerate(pepDon):
    donorsList = donorsList + (val*[pepDonInds[j]])

  for (j, val) in enumerate(watAcc):
    acceptorsList = acceptorsList + (val*[watInds[j]])
    bondsWat[int(j/3)] = bondsWat[int(j/3)] + val

  for (j, val) in enumerate(watDon):
    donorsList = donorsList + (val*[watInds[j]])
    bondsWat[int(j/3)] = bondsWat[int(j/3)] + val

  #Above uses properties of python lists to add each index the number of H-bonds it participates in

  bondsPer = bondsWat

  #For easy file writing, make donors and acceptors into strings of indices
  #Remember that the sim package indexes at zero!
  donors = ''.join(str(e)+"|" for e in donorsList)
  acceptors = ''.join(str(e)+"|" for e in acceptorsList)

  return NBonds, bondsPer, acceptors, donors

def BBHBonds(allPos, pepAccInds, pepDonInds, distCut = 2.1, angCut = 30.0):
  """Finds H bonds between two list of acceptors and donors. Intended for just peptide backbone.
     Inputs:
     allPos - full position array for trajectory frame
     pepAccInds - global indics of peptide acceptors
     pepDonInds - global indices of peptide doneors
     distCut(=2.1) - distance cutoff for H-bond detection
     angCut(=30.0) - angle cutoff for H-bond detection
     Outputs:
     NBonds - number of detected H-bonds
     donors - indices of donors as string
     acceptors - indices of acceptors as string
  """

  #Get H-bonds
  NBonds, pepAcc, pepDon = wl.bbhbonds(allPos[pepAccInds], allPos[pepDonInds], distCut, angCut)

  #Sort nicely
  acceptorsList = []
  donorsList = []

  for (j, val) in enumerate(pepAcc):
    acceptorsList = acceptorsList + (val*[pepAccInds[j]])

  for (j, val) in enumerate(pepDon):
    donorsList = donorsList + (val*[pepDonInds[j]])

  #set lists to strings and return
  donors = ''.join(str(e)+"|" for e in donorsList)
  acceptors = ''.join(str(e)+"|" for e in acceptorsList)

  return NBonds, acceptors, donors

def WatHBonds(allPos, watInds, allWatInds, BoxDims, distCut = 2.1, angCut = 30.0):
  """Also kind of wack, but keeping since used in peptide-surface pulling analysis.
     For a better, more general algorithm, use HBondsGeneral.
     Inputs:
     allPos - full position array for trajectory frame (all atoms included)
     watInds - global indices of water atoms in selected hydration shell(s)
     allWatInds - global indices of ALL water atoms
     BoxDims - dimensions of box to account for periodic BCs (to turn off, set to zero)
     distCut(=2.1) - distance cutoff for H-bond detection
     angCut(=30.0) - angle cutoff for H-bond detection
     Outputs:
     NBonds - number of detected H-bonds
     bondsPer - number of detected H-bonds for each water molecule in selection
     acceptors - indices of acceptors as string
     donors - indices of donors (H atoms only) as string
  """

  #Get H-bond info
  NBonds, watAcc, watDon = wl.wathbonds(allPos[watInds], allPos[allWatInds], BoxDims, distCut, angCut)

  #And sort nicely into just acceptors and donors
  #Also count number of H-bonds for each water to get estimate of average per water
  acceptorsList = []
  donorsList = []
  bondsWat = np.zeros(int(len(watInds)/3))

  for (j, val) in enumerate(watAcc):
    acceptorsList = acceptorsList + (val*[watInds[j]])
    bondsWat[int(j/3)] = bondsWat[int(j/3)] + val

  for (j, val) in enumerate(watDon):
    donorsList = donorsList + (val*[watInds[j]])
    bondsWat[int(j/3)] = bondsWat[int(j/3)] + val

  #Above uses properties of python lists to add each index the number of H-bonds it participates in

  #print bondsWat
  #bondsPer = np.average(bondsWat)
  bondsPer = bondsWat

  #For easy file writing, make donors and acceptors into strings of indices
  #Remember that the sim package indexes at zero!
  donors = ''.join(str(e)+"|" for e in donorsList)
  acceptors = ''.join(str(e)+"|" for e in acceptorsList)

  return NBonds, bondsPer, acceptors, donors

def getCosAngs(subPos, Pos, BoxDims, lowCut=0.0, highCut=3.413):
  """This is called getCosAngs, but actually just returns the angles themselves (faster to convert
     from cos(theta) to theta in Fortran)
     Inputs:
     subPos - positions of set of atoms to measure tetrahedrality of (may be different, subset, or same as Pos)
     Pos - positions of ALL atoms that can make tetrahedral configurations (needed if subPos not same as Pos)
     BoxDims - current box dimensions to account for periodicity
     lowCut - lower cutoff for nearest-neighbor shell (default 0.0)
     highCut - higher cutoff for nearest-neighbor shell (default 3.413 - see Chaimovich, 2014, but should really
               change to reflect first peak in g(r) for the chosen water model)
     Outputs:
     angVals - all angle values for current configuration of positions supplied
     numAngs - number of angles for each central oxygen atom (i.e. number neighbors factorial)
               This is useful for finding which angles belong to which central oxygens
               This return value was added on 07/09/2017, so any code using this function
               before then will break, unfortunately, but the fix is easy.
  """

  #Set-up array to hold angle results and stack as go... list increases in size!
  angVals = np.array([])
  numAngs = np.zeros(len(subPos))

  #Find nearest neighbors for ALL atoms in subPos
  #But make sure using efficient algorithm...
  #If subPos is same as Pos, use allnearneighbors instead
  if np.array_equal(subPos, Pos):
    nearNeighbs = wl.allnearneighbors(Pos, BoxDims, lowCut, highCut).astype(bool)
  else:
    nearNeighbs = wl.nearneighbors(subPos, Pos, BoxDims, lowCut, highCut).astype(bool)

  #Loop over each position in subPos, finding angle made with all neighbor pairs
  for (i, apos) in enumerate(subPos):
    #Make sure have nearest neighbors...
    if len(Pos[nearNeighbs[i]]) > 0:
      #below returns symmetric, square array (zero diagonal)
      tempAng = wl.tetracosang(apos, Pos[nearNeighbs[i]], BoxDims)
      #Only want half of array, flattened
      angVals = np.hstack((angVals, tempAng[np.triu_indices(len(tempAng),k=1)].tolist()))
      numAngs[i] = tempAng.shape[0]

  return angVals, numAngs

def tetrahedralMetrics(angVals, nBins=500, binRange=[0.0, 180.0]):
  """Inputs:
     angVals - all angle values sampled
     nBins - number histogram bins to use
     binRange - histogram bin range to apply
     Outputs:
     angDist - distribution of angle
     bins - bins used in histogramming
     fracTet - fraction of distribution that is tetrahedral (integrate cosDist from -0.75 to 0.25 - see Chaimovich, 2014)
     avgCos - average Cos(angle) within tetrahedral peak
     stdCos - second moment of Cos(angle) within tetrahedral peak
  """

  #Histogram the data - note that density set so just returns number of counts, not normalized
  angDist, bins = np.histogram(angVals, bins=nBins, range=binRange, density=False)

  #Take index before since want histogram bin containing this value
  startTet = np.argmax(bins>np.arccos(0.25)*180.0/np.pi) - 1
  endTet = np.argmax(bins>np.arccos(-0.75)*180.0/np.pi) - 1

  fracTet = np.sum(angDist[startTet:endTet]) / np.sum(angDist)

  #Take average and second moment within peak
  avgCos = 0.0
  stdCos = 0.0
  angCount = 0

  for ang in angVals:
    if (ang >= np.arccos(0.25)*180.0/np.pi) and (ang <= np.arccos(-0.75)*180.0/np.pi):
      avgCos = avgCos + np.cos(ang*np.pi/180.0)
      stdCos = stdCos + np.cos(ang*np.pi/180.0)**2
      angCount += 1

  avgCos = avgCos / angCount
  stdCos = stdCos / angCount

  return angDist, bins, fracTet, avgCos, stdCos

def getOrderParamq(subPos, Pos, BoxDims, lowCut=0.0, highCut=8.0):
  """Finds angles for 4 nearest neighbors of each water and returns for all waters the
tetrahedral order parameter, q, used by Errington and Debenedetti (2001).
     Inputs:
     subPos - positions of set of atoms to measure tetrahedrality of (may be different, subset, or same as Pos)
     Pos - positions of ALL atoms that can make tetrahedral configurations (needed if subPos not same as Pos)
     BoxDims - current box dimensions to account for periodicity
     lowCut - lower cutoff for nearest-neighbor shell (default 0.0)
     highCut - higher cutoff for nearest-neighbor shell used to find 4 nearest neighbors
     Outputs:
     qVals - returns an order parameter value for each water
     distNeighbs - returns distances from central oxygen to 4 nearest neighbors
  """

  #Set-up array to hold results
  qVals = np.zeros(len(subPos))
  distNeighbs = np.zeros((len(subPos), 4))

  #Find nearest neighbors for ALL atoms in subPos
  #But make sure using efficient algorithm...
  #If subPos is same as Pos, use allnearneighbors instead
  if np.array_equal(subPos, Pos):
    nearNeighbs = wl.allnearneighbors(Pos, BoxDims, lowCut, highCut).astype(bool)
  else:
    nearNeighbs = wl.nearneighbors(subPos, Pos, BoxDims, lowCut, highCut).astype(bool)

  #Loop over each position in subPos, finding angle made with the closest 4 neighbors, then q
  for (i, apos) in enumerate(subPos):
    #Make sure have nearest neighbors...
    if np.sum(nearNeighbs[i]) > 0:
      thisPos = wl.reimage(Pos[nearNeighbs[i]], apos, BoxDims)
      thisDists = np.linalg.norm(thisPos - apos, axis=1)
      sortInds = np.argsort(thisDists)
      newPos = thisPos[sortInds][:4]
      distNeighbs[i,:] = thisDists[sortInds][:4]
      #below returns symmetric, square array (zero diagonal)
      tempAng = wl.tetracosang(apos, newPos, BoxDims)
      #Only want half of array, flattened
      angVals = tempAng[np.triu_indices(len(tempAng),k=1)]
      #Now compute q for this set of angles
      qVals[i] = 1.0 - (3.0/8.0)*np.sum((np.cos(angVals*np.pi/180.0) + (1.0/3.0))**2)

  #Return all of the order parameter values
  return qVals, distNeighbs

def findSineCoeffs(allangs, Norder=180, doNormalize=False):
  """Given an array of angles, computes the sine coefficients to the given order.
Note that to get right coefficients, will need to divide by total number of angles.
This is not done by default, assuming that angles provided are for each frame.
     Inputs:
             allangs - array or list of angles
             Norder - (default 180) number of terms in sine series to use (excludes k=0)
             doNormalize - (default False) if true, divides by number of samples to correctly normalize
     Outputs:
             coeffs - Norder x 2 array; 1st column is k, second column is coefficient
                      Comes from fact that period is zero to Pi, so only keep sin(k*angle) in series
  """
  #Check if angles in radians - if any values are greater than Pi, assume in degrees
  if np.max(allangs) > np.pi:
    allangs = allangs * np.pi / 180.0
  coeffs = np.zeros((Norder,2))
  for k in range(Norder):
    coeffs[k,0] = k+1
    coeffs[k,1] = np.sqrt(2.0/np.pi)*np.sum(np.sin((k+1)*allangs))
  if doNormalize:
    coeffs = coeffs / len(allangs)
  return coeffs

def distFromCoeffs(coeffs, angvals=None, Norder=60):
  """Given an array of coefficients for a sine series, compute the distribution.
     Inputs:
             coeffs - coefficients for each term in a sine series
                      assuming that for sin(k*angle) form, this array is sorted from small to large k
             angvals - (default 0.0 to 180.0 by 0.01) angle values in degrees at which distribution
                       should be evaluated - normalization will be done for PDF along degrees
             Norder - (default 60) number of terms in the series to use (i.e. number of coeffs)
     Outputs:
             adist - returns a normalized distribution
  """
  if angvals is None:
    angvals = np.arange(0.0, 180.0, 0.01)
  #Also define in radians
  radvals = angvals * np.pi / 180.0
  adist = np.zeros(len(angvals))
  normfac = 0.0
  for k in range(Norder):
    adist += coeffs[k]*np.sin((k+1)*radvals)
    if (k+1)%2 != 0:
      normfac += coeffs[k]*2.0/(k+1)
  adist = adist / (normfac*(angvals[1]-angvals[0]))
  return adist

def fitDist(refDists, Dist, bruteNs=200):
  """Given a set of reference distributions, as a numpy array with each distribution as a row,
fits the current distribution using a linear combination of the reference distributions.
  Inputs:
  refDists - array with each reference distribution as a row
  Dist - (3-body angle) distribution to fit as linear combination of references with
         the fitting parameters (linear coefficients) summing to one
  bruteNs - number of discrete bins to use along each parameter when searching for brute minimum
  Outputs:
  fitParams - fit parameters resulting from fitting
  resSq - sum of squared residuals for fit
  resSigned - signed residuals at each point of fit
  """

  #Define tolerance
  tolf = 1.0e-12
  tolx = 1.0e-12

  #Initialize parameters to seek for - start in 4 ways and take minimum of these
  initParams = np.eye(refDists.shape[0])
  initParams = np.vstack((initParams, np.ones(refDists.shape[0]) * (1.0/refDists.shape[0])))

  #Define an objective function to be minimized
  def funcMin(vals, *withcon):
    #Give it parameter values, returns squared residuals
    func = np.sum((np.dot(vals, refDists) - Dist)**2)
    if withcon:
      func = func + (np.sum(vals) - 1.0)**2
    return func

  def jacFunc(vals):
    #Returns the Jacobian of the function to minimize
    func = np.dot(refDists, 2.0*(np.dot(vals, refDists) - Dist))
    return func

  def funcSquares(vals):
    #Gives vector of squared residuals to see where best/worst parts of fit are
    func = (np.dot(vals, refDists) - Dist)**2
    return func

  #Define constraints... for now say that all parms must sum to one
  cons = ({'type' : 'eq',
           'fun' : lambda x: np.sum(x) - 1.0,
           'jac' : lambda x: np.ones(len(x))})

  #And define bounds to keep all params between 0 and 1
  bnds = [(0.0,1.0)]*refDists.shape[0]

  #For each set of starting conditions, do minimization, then pick global min
  globMinInfo = None

  #And will store squared residuals at found mins as go
  #Checks if one part of curve fits better than another
  resSq = np.zeros((refDists.shape[1], initParams.shape[0]))

  for (i, params) in enumerate(initParams):

    #If only one distribution given, don't use constraint
    if refDists.shape[0] == 1:
      mininfo = optimize.minimize(funcMin, params, jac=jacFunc, method='SLSQP',
                                                bounds=bnds, options={'ftol':tolf})
    else:
      mininfo = optimize.minimize(funcMin, params, jac=jacFunc, method='SLSQP',
                                constraints=cons, bounds=bnds, options={'ftol':tolf})

    #print "Minimum sum of squares: %e  at values "%mininfo.fun+str(mininfo.x)

    if globMinInfo != None:
      if mininfo.fun < globMinInfo.fun:
        globMinInfo = mininfo
    else:
      globMinInfo = mininfo

    resSq[:,i] = funcSquares(mininfo.x)

  #Compare to global min with brute force
  if refDists.shape[0] == 1:
    (bruteMinInfo) = optimize.brute(funcMin, tuple(bnds), Ns=bruteNs, finish=None, full_output=True, disp=False)
  else:
    (bruteMinInfo) = optimize.brute(funcMin, tuple(bnds), args=(1,), Ns=bruteNs, finish=None, full_output=True, disp=False)

  fitParams = bruteMinInfo[0]

  #print "Brute force finds minima at "+str(fitParams)

  #Also compute regular residuals, not squared
  resSigned = np.dot(fitParams, refDists) - Dist

  #print "Best fit found at:"
  #print [float(q) for q in fitParams]
  #print "And with parameters summing to %f" % np.sum(fitParams)
  return fitParams, resSq, resSigned

def waterOrientationBinZ(Opos, Hpos, boxDim, refVec=[0.0, 0.0, 1.0], refBins=None, angBins=None):
  """Determines the angle between a reference vector and the dipoles and plane-normal vector
of all water molecule positions provided.
     Inputs:
     Opos - all water oxygen positions
     Hpos - all water hydrogen positions
     boxDim - box dimensions for imaging
     refVec - the reference vector for water orientation, default
              is the z-direction [1, 0, 0]
     refBins - bins along the direction of refVec that the waters
               should be placed into (default is min and max along refVec)
     angBins - bins for calculated angles, default 500 bins from 0 to 180
     Outputs:
     plane2Dhist - 2D histogram with angle bins varying over rows and
                   refVec bins varying over rows (for the water plane vector angles)
     dip2Dhist - 2D histogram as above, but for dipole vector angles
  """

  #Get positions of oxygen atoms along refVec, then create
  #this array with each entry repeated
  refVec = refVec / np.linalg.norm(refVec)
  zOpos = np.dot(Opos, refVec)
  zOposforH = np.array([[z,z] for z in zOpos]).flatten()

  #Compute all of the angles with respect to the reference
  #Note that the dipole vector of each water molecule will be taken as
  #the sum of the OH bond vectors
  angDip, angPlane = wl.watorient(Opos, Hpos, refVec, boxDim)

  #Set refBins if not set
  if refBins is None:
    refBins = np.arange(np.min(zOpos), np.max(zOpos), 0.2)

  #Same for angBins
  if angBins is None:
    angBins = np.arange(0.0, 180.001, 180.0/500.0)

  #And do 2D histogramming
  plane2Dhist, angEdges, refEdges = np.histogram2d(angPlane, zOposforH, bins=[angBins, refBins], normed=False)
  dip2Dhist, angEdges, refEdges = np.histogram2d(angDip, zOpos, bins=[angBins, refBins], normed=False)

  return plane2Dhist, dip2Dhist

def waterOrientation(Opos, Hpos, boxDim, refVec=[0.0, 0.0, 1.0]):
  """This is a wrapper for the waterlib function watorient.
     Inputs:
     Opos - all water oxygen positions
     Hpos - all water hydrogen positions
     boxDim - box dimensions for imaging
     refVec - the reference vector for water orientation, default
              is the z-direction [1, 0, 0]
     Outputs:
     dipAngs - all angles of dipole vectors with reference vector for all waters
     planeAngs - all angles of plane-normal vector to reference vector for all waters
  """

  #Call watorient to get all angles
  dipAngs, planeAngs = wl.watorient(Opos, Hpos, refVec, boxDim)

  return dipAngs, planeAngs

def binnedVolumePofN(Opos, volBins, numBins, binMask=None):
  """Inputs:
     Opos - array of oxygen 3D coordinates
     volBins - volume (x,y,z coordinate) bin edges tiling the space to place waters into
               Form should be tuple of x, y, and z bin edge arrays
               Bins should be uniform (makes no sense to structure analysis this way otherwise)
     numBins - bin edges for histogramming number of waters in each volume of volBins
     binMask - boolean array of same dimension as number of bins in x, y, z
               Use to exclude some bins by changing certain coordinates to False
     Outputs:
     numWatHist - counts for number of waters in sub-volume of size edgeLxedgeLxedgeL
  """

  #Create mask if necessary
  if binMask is None:
    binMask = np.ones((len(volBins[0])-1, len(volBins[1])-1, len(volBins[2])-1), dtype=bool)
  else:
    if binMask.shape == (len(volBins[0])-1, len(volBins[1])-1, len(volBins[2])-1):
      binMask = binMask
    else:
      print "Dimensions of mask for spatial bins does not match dimensions of spatial bins. Quitting."
      sys.exit(2)

  #Want to use sphere rather than cube for statistics
  #So first find which bin each oxygen belongs to, then find distance
  #to center of bin and see if should exclude or not
  #Written in Fortran for speed
  hist = wl.binongrid(Opos, volBins[0], volBins[1], volBins[2])

  #Use numpy histogramming to count how many oxygens in each cube volume (doesn't use interior spheres)
  #hist, edges = np.histogramdd(Opos, bins=volBins, normed=False)

  #Now histogram number of waters in each subvolume, which will be P(N)
  numWatHist, watedges = np.histogram(hist[binMask].flatten(), bins=numBins, normed=False)

  return numWatHist

#Should also define some function "pointsInVol" that creates binMask based on given set of points or some geometry that should not be included when finding waters (i.e. like a hard sphere or protein)

def HBondsGeneral(accPos, donPos, donHPos, boxL, accInds, donInds, donHInds, distCut=3.5, angCut=150.0):
  """Wraps generalHbonds in the waterlib library to define H-bonds, and also returns their locations.
     Inputs:
     accPos - 3-dimensional vectors of acceptor heavy-atom positions
     donPos - 3D vectors of donor heavy-atom positions (if have multiple hydrogens, must list multiple times)
     donHPos - 3D vector of donor hydrogen positions (should be same length as donPos, which may have duplicates)
     accInds - indices of acceptor atoms
     donInds - indices of donor heavy-atoms
     donHInds - indices of donor hydrogen atoms
     boxL - box dimensions
     distCut - (default 3.5) heavy-atom to heavy-atom distance below which an H-bond may be defined
     angCut - (default 150.0) O-H---O angle cut-off, in degrees, above which an H-bond may be defined
     Outputs:
     NumHB - number of hydrogen bonds for acceptor/donor set provided
     HBlist - NumHB x 2 array with acceptor index in the 1st column and donor index in the 2nd
     HBloc - NumHB x 3 array of h-bond locations, which is halfway between the acceptor and donor H
  """

  #First get H-bond matrix and locations
  HBboolMat = wl.generalhbonds(accPos, donPos, donHPos, boxL, distCut, angCut)
  HBboolMat = np.array(HBboolMat, dtype=bool)

  #Now parse through matrix, counting H-bonds and creating list of index pairs
  NumHB = np.sum(HBboolMat)
  HBlist = (-1)*np.ones((NumHB, 2))
  HBloc = np.zeros((NumHB, 3))
  HBlistCount = 0
  for i, abool in enumerate(HBboolMat):
    theseDonors = donInds[abool]
    if len(theseDonors) > 0:
      theseDonHPos = donHPos[abool]
      #Image donor H location around acceptor
      theseDonHPos = wl.reimage(theseDonHPos, accPos[i], boxL)
      for j, aDon in enumerate(theseDonors):
        HBlist[HBlistCount,:] = [accInds[i], aDon]
        HBloc[HBlistCount] = 0.5*(theseDonHPos[j] + accPos[i])
        HBlistCount += 1

  return NumHB, HBlist, HBloc

def computeSphericalFourierCoeffs(subPos, Pos, BoxDims, lowCut=0.0, highCut=3.413, minDegree=0, maxDegree=12):
  """Computes the vectors of Fourier coefficients for each degree of a spherical harmonic expansion
     as described by Keys, Iacovella, and Glotzer, 2011. subPos is treated as the central atoms,
     while Pos should include the atoms that may potentially be neighbors.
     Inputs:
     subPos - positions of atoms to treat as the central atoms
     Pos - positions of all other atoms, which will be considered for neighbor-searching; can be same as subPos
     BoxDims - box dimensions so that minimum images may be used
     lowCut - (default 0.0) the lower cutoff for the radial shell
     highCut - (default 3.413) the upper cutoff for the radial shell
     minDegree - (default 0) the minimum spherical harmonic degree (l)
     maxDegree - (default 12) the maximum spherical harmonic degree (l)
     Outputs:
     coeffVecs - a len(subPos) x (1 + maxDegree - minDegree) x (2*maxDegree + 1) matrix
                 For each central atom in subPos, a matrix of the complex-valued vectors (as rows)
                 is provided. This still allows magnitudes to be easily evaluated, since real and
                 imaginary parts of zero will contribute nothing to the magnitude
     numNeighbs - number of neighbors for each water molecule (necessary to compute global order parameters
                  or coefficients by multiplying by this and the dividing by the total of the waters
                  to "average" over)
  """

  #Set up the output matrix now since know size
  coeffVecs = np.zeros((len(subPos), 1+maxDegree-minDegree, 2*maxDegree+1), dtype=complex)

  #And array to return number of neighbors for each water
  numNeighbs = np.zeros(len(subPos), dtype='float16')

  #Would be nice to combine neighbor searching with 3-body angle computation or H-bonding code
  #But then harder to efficiently implement different cutoffs...
  #So that might be too ambitious
  #Find neighbors within cutoff for ALL atoms in subPos
  #But make sure using efficient algorithm...
  #If subPos is same as Pos, use allnearneighbors instead
  if np.array_equal(subPos, Pos):
    nearNeighbs = wl.allnearneighbors(Pos, BoxDims, lowCut, highCut).astype(bool)
  else:
    nearNeighbs = wl.nearneighbors(subPos, Pos, BoxDims, lowCut, highCut).astype(bool)

  #Loop over each position in subPos and find neighbor positions in spherical coordinates
  for (i, apos) in enumerate(subPos):
    #Make sure have nearest neighbors...
    if len(Pos[nearNeighbs[i]]) > 0:
      tempPos = wl.reimage(Pos[nearNeighbs[i]], apos, BoxDims) - apos
      numNeighbs[i] = len(tempPos)
      #Compute radial distances... unfortunate that have to do this again, but maybe improve later
      rdists = np.linalg.norm(tempPos, axis=1)
      #And get polar and azimuthal angles
      polarang = np.arccos(tempPos[:,2]/rdists)
      azimang = np.arctan2(tempPos[:,1], tempPos[:,0]) #Using special arctan2 function to get quadrant right
      #Now compute Fourier coefficient vectors (i.e. have complex-valued component of coefficient vector
      #associated with each m value, where m = -l, -l+1, ... , l)
      #Loop over the desired number of coefficients to compute
      for l in range(minDegree, maxDegree + 1):
        thisvec = np.zeros(2*l + 1, dtype=complex)
        #Also note that have one vector for each neighbor, so must loop over neighbors
        for j in range(len(tempPos)):
          thisvec += sph_harm(np.arange(-l, l+1), l, azimang[j], polarang[j])
        thisvec /= len(tempPos)
        #And compute the magnitude of this vector of complex numbers
        coeffVecs[i,l-minDegree,:(2*l+1)] = thisvec

  return coeffVecs, numNeighbs

def get1BodyDOFs(coordO, coordH1, coordH2):
  """Given O, H, and H 3D coordinates, identifies the 6 degrees of freedom for a single water
Note that this is assuming an inhomogeneous system
Vector returned is oxygen x, y, z, followed by the spherical coordinate angles for the
dipole vector relative to the oxygen, and the angle of rotation around the dipole vector
COORDINATES SHOULD ALREADY BE IMAGED.
  """
  dofVec = np.zeros(6)
  dofVec[:3] = coordO[:]

  rOD = 0.5*(coordH1 + coordH2) - coordO
  rOD /= np.linalg.norm(rOD) #Could hard-code the rOD length for speed... maybe later

  rH1H2 = coordH2 - coordH1
  rH1H2 /= np.linalg.norm(rH1H2) #And could also hard-code this, too...

  #rOH1 = coordH1 - coordO
  #rOH1 /= np.linalg.norm(rOH1)

  #rOH2 = coordH2 - coordO
  #rOH2 /= np.linalg.norm(rOH2)

  unitX = np.array([0.0, 0.0, 1.0]) #Arbitrarily pick x axis to define reference plane for rotation about dipole

  #cross1 = np.cross(rOH1, rOH2)
  #cross1 /= np.linalg.norm(cross1)

  crossX = np.cross(rOD, unitX)
  crossX /= np.linalg.norm(crossX)

  dofVec[3] = np.arctan2(rOD[1], rOD[0]) #Making sure to use arctan2 to cover range [-pi, pi]
  dofVec[4] = np.arccos(rOD[2]) #Taking last element is same as dotting with unit Z vector
  dofVec[5] = np.arccos(np.dot(rH1H2, crossX))
  #dofVec[5] = np.arccos(np.dot(cross1, crossX))

  return dofVec

def get2BodyDOFs(coordO1, coordH11, coordH12, coordO2, coordH21, coordH22):
  """Given 3D coordinates for all atoms in two water molecules, computes specifically 2-body degrees of freedom
Note that returns only 6 degrees of freedom, so excludes the DOFs for the first water
ONLY gives those relevant to relative distance and orientation of two waters
Order in returned vector is rO1O2, theta1, theta2, phi, chi1, chi2 (see Lazaridis and Karplus for definitions)
COORDINATES SHOULD ALREADY BE IMAGED
  """
  dofVec = np.zeros(6)

  rO1O2 = coordO2 - coordO1
  dofVec[0] = np.linalg.norm(rO1O2)
  rO1O2 /= dofVec[0]
  rO2O1 = -rO1O2

  rO1D1 = 0.5*(coordH11 + coordH12) - coordO1
  rO1D1 /= np.linalg.norm(rO1D1) #Could hard-code to speed up... may do later

  rO2D2 = 0.5*(coordH21 + coordH22) - coordO2
  rO2D2 /= np.linalg.norm(rO2D2)

  #Need to figure out which H is closer to other oxygen to define rH11H12 according to Lazaridis and Karplus, 1996
  if np.linalg.norm(coordH11 - coordO2) <= np.linalg.norm(coordH12 - coordO2):
    rH11H12 = coordH12 - coordH11
  else:
    rH11H12 = coordH11 - coordH12
  rH11H12 /= np.linalg.norm(rH11H12) #Again, could hard code if wanted

  if np.linalg.norm(coordH21 - coordO1) <= np.linalg.norm(coordH22 - coordO1):
    rH21H22 = coordH22 - coordH21
  else:
    rH21H22 = coordH21 - coordH22
  rH21H22 /= np.linalg.norm(rH21H22)

  cross1 = np.cross(rO1O2, rO1D1)
  cross1 /= np.linalg.norm(cross1)

  cross2 = np.cross(rO2D2, rO2O1)
  cross2 /= np.linalg.norm(cross2)

  dofVec[1] = np.arccos(np.dot(rO1D1, rO1O2))
  dofVec[2] = np.arccos(np.dot(rO2D2, rO2O1))
  dofVec[3] = np.arccos(np.dot(cross1, cross2))
  dofVec[4] = np.arccos(np.dot(rH11H12, cross1))
  dofVec[5] = np.arccos(np.dot(rH21H22, cross2))

  return dofVec

def get3BodyDOFs(coordO1, coordH11, coordH12, coordO2, coordH21, coordH22, coordO3, coordH31, coordH32):
  """Like above function, but gives 6 DOFs pertaining to just the 3-body degrees of freedom
Order in returned vector  is rO1O3 (distance), theta3b (three-body angle),
omega (rotation of 3rd water around O1-O3 vector), then theta3, phi3, and chi3
(last three defined as for the second water in the 2-body DOFs, but for just the third water)
COORDINATES SHOULD ALREADY BE IMAGED
  """
  dofVec = np.zeros(6)

  rO1O2 = coordO2 - coordO1
  rO1O2 /= np.linalg.norm(rO1O2)
  rO2O1 = -rO1O2

  rO1O3 = coordO3 - coordO1
  dofVec[0] = np.linalg.norm(rO1O3)
  rO1O3 /= dofVec[0]
  rO3O1 = -rO1O3

  rO1D1 = 0.5*(coordH11 + coordH12) - coordO1
  rO1D1 /= np.linalg.norm(rO1D1)

  rO3D3 = 0.5*(coordH31 + coordH32) - coordO3
  rO3D3 /= np.linalg.norm(rO3D3)

  if np.linalg.norm(coordH31 - coordO1) <= np.linalg.norm(coordH32 - coordO1):
    rH31H32 = coordH32 - coordH31
  else:
    rH31H32 = coordH31 - coordH32
  rH31H32 /= np.linalg.norm(rH31H32)

  cross12 = np.cross(rO1O2, rO1D1)
  cross12 /= np.linalg.norm(cross12)

  cross13 = np.cross(rO1O3, rO1D1)
  cross13 /= np.linalg.norm(cross13)

  cross31 = np.cross(rO3D3, rO3O1)
  cross31 /= np.linalg.norm(cross31)

  rperp = rO1O3 - np.dot(rO1O2, rO1O3)*rO1O2
  rperp /= np.linalg.norm(rperp)

  dofVec[1] = np.arccos(np.dot(rO1O2, rO1O3))
  dofVec[2] = np.arccos(np.dot(rperp, cross12))
  dofVec[3] = np.arccos(np.dot(rO3D3, rO3O1))
  dofVec[4] = np.arccos(np.dot(cross13, cross31))
  dofVec[5] = np.arccos(np.dot(rH31H32, cross31))

  return dofVec

def distanceMetric1B(vec1, vec2, Rsq=(0.09572**2), sintw=(np.sin(104.52*np.pi/180.0)**2)):
  """Computes distance metric appropriate to 1-body DOFs.
     A direct Euclidean distance is not appropriate since using curvilinear coordinates,
     so this defines a distance utilizing local curvature that is exact for very small
     differences. It comes from Taylor-expanding the formula for Euclidean distance in
     spherical coordinates with respect to both angles to second order.
  """
  diffs = (vec2 - vec1)**2
  dist = np.sqrt(diffs[0] + diffs[1] + diffs[2] + Rsq*diffs[3]
                 + Rsq*np.sin(vec2[3])*np.sin(vec1[3])*diffs[4]
                 + Rsq*sintw*diffs[5])
  return dist

def distanceMetric2B(vec1, vec2, Rsq=(0.09572**2), sintw=(np.sin(104.52*np.pi/180.0)**2)):
  """Computes distance metric appropriate to 2-body DOFs.
     A direct Euclidean distance is not appropriate since using curvilinear coordinates,
     so this defines a distance utilizing local curvature that is exact for very small
     differences. It comes from Taylor-expanding the formula for Euclidean distance in
     spherical coordinates with respect to both angles to second order.
     Note that this includes 1-body degrees of freedom, so expects 12-dimensional vectors.
  """
  diffs = (vec2 - vec1)**2
  dist = np.sqrt(diffs[0] + diffs[1] + diffs[2] + Rsq*diffs[3]
                 + Rsq*np.sin(vec2[3])*np.sin(vec1[3])*diffs[4]
                 + Rsq*sintw*diffs[5]
                 + diffs[6] + Rsq*diffs[7] + Rsq*diffs[8]
                 + Rsq*np.sin(vec2[8])*np.sin(vec1[8])*diffs[9]
                 + Rsq*sintw*diffs[10] + Rsq*sintw*diffs[11])
  return dist

def distanceMetric3B(vec1, vec2, Rsq=(0.09572**2), sintw=(np.sin(104.52*np.pi/180.0)**2)):
  """Computes distance metric appropriate to 3-body DOFs.
     A direct Euclidean distance is not appropriate since using curvilinear coordinates,
     so this defines a distance utilizing local curvature that is exact for very small
     differences. It comes from Taylor-expanding the formula for Euclidean distance in
     spherical coordinates with respect to both angles to second order.
     Note that this includes 1- and 2-body degrees of freedom, so expects 18-dimensional vectors.
  """
  diffs = (vec2 - vec1)**2
  dist = np.sqrt(diffs[0] + diffs[1] + diffs[2] + Rsq*diffs[3]
                 + Rsq*np.sin(vec2[3])*np.sin(vec1[3])*diffs[4]
                 + Rsq*sintw*diffs[5]
                 + diffs[6] + Rsq*diffs[7] + Rsq*diffs[8]
                 + Rsq*np.sin(vec2[8])*np.sin(vec1[8])*diffs[9]
                 + Rsq*sintw*diffs[10] + Rsq*sintw*diffs[11]
                 + diffs[12] + vec2[12]*vec1[12]*diffs[13]
                 + vec2[12]*vec1[12]*np.sin(vec2[13])*np.sin(vec1[13])*diffs[14]
                 + Rsq*diffs[15]
                 + Rsq*np.sin(vec2[15])*np.sin(vec1[15])*diffs[16]
                 + Rsq*sintw*diffs[17])
  return dist


import networkx as nx
from networkx.algorithms import isomorphism
import math
import re

ORIG_PPI = [["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 234) ('type_1' 'type_1') ('type_1' 234) (197 234) (234 197)]", 844],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 478) ('type_1' 'type_1') ('type_1' 478) (197 234) (197 478) (234 197) (478 197)]", 73],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 234) ('type_1' 'type_1') ('type_1' 234) (197 234) (234 197) (234 338) (338 234)]", 53],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 234) ('type_1' 'type_1') ('type_1' 234) (197 234) (234 197) (234 738) (234 800) (738 234) (800 234)]", 25],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 328) ('type_1' 'type_1') ('type_1' 328) (197 478) (478 197) (478 479) (328 479) (479 328) (479 478)]", 4],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 314) ('type_1' 'type_1') ('type_1' 314) (114 314) (314 114) (314 331) (314 1120) (294 331) (331 294) (331 314)]", 11],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 314) ('type_1' 'type_1') ('type_1' 314) (114 314) (314 114) (314 1120)]", 14],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 275) ('type_1' 'type_1') ('type_1' 275) (275 334) (275 375) (275 589) (275 1069) (334 275) (375 275) (589 275) (1069 275)]", 21],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 314) ('type_1' 'type_1') ('type_1' 314) (314 1120)]", 60],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 558) ('type_1' 'type_1') ('type_1' 558) (558 559) (558 571) (558 654) (558 1328) (558 1329) (559 558) (571 558) (654 558) (1328 558) (1329 558)]", 36],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 239) ('type_1' 'type_1') ('type_1' 239) (239 1056) (239 1107) (239 1331) (1056 239)]", 12],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 43) ('type_1' 'type_1') ('type_1' 43) (42 43)]", 7],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 1059) ('type_0' 925) ('type_1' 'type_1') ('type_1' 1059) ('type_1' 925) (925 1185) (1185 925) (1185 1059) (1059 1185)]", 12],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 1440) ('type_0' 749) ('type_0' 1551) ('type_1' 'type_1') ('type_1' 1440) ('type_1' 749) ('type_1' 1551) (749 750) (750 749) (750 1440) (750 1551) (1440 750) (1551 750)]", 4],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 1) ('type_0' 5) ('type_1' 'type_1') ('type_1' 1) ('type_1' 5) (1 12) (1 58) (1 188) (12 1) (12 5) (58 1) (58 5) (58 609) (188 1) (188 5) (5 12) (5 58) (5 188) (5 609) (609 5) (609 58)]", 1]]

CL_PPI =   [["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 215) ('type_1' 'type_1') ('type_1' 215) (215 807)]", 682],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 215) ('type_1' 'type_1') ('type_1' 215) (215 437) (215 807)]", 15],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 602) ('type_1' 'type_1') ('type_1' 602) (215 807) (602 215)]", 13],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 215) ('type_1' 'type_1') ('type_1' 215) (1064 215)]", 488],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 770) ('type_1' 'type_1') ('type_1' 76) (76 770)]", 392],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 64) ('type_1' 'type_1') ('type_1' 64) (5 64) (64 5)]", 56],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 30) ('type_1' 'type_1') ('type_1' 30) (30 503) (503 30) (777 30)]", 1],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 49) ('type_1' 'type_1') ('type_1' 9) (9 49) (49 9)]", 12],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 10) ('type_1' 'type_1') ('type_1' 10) ('type_1' 6) (6 10) (10 6)]", 2],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 9) ('type_0' 10) ('type_1' 'type_1') ('type_1' 9) ('type_1' 10) (9 10)]", 2],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 10) ('type_0' 50) ('type_1' 'type_1') ('type_1' 50) (10 50) (50 10)]", 4],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 20) ('type_0' 244) ('type_1' 'type_1') ('type_1' 33) ('type_1' 6) (3 6) (3 8) (3 33) (3 244) (6 8) (6 20) (8 20) (8 33) (8 244) (33 3) (33 6) (33 20) (244 3) (244 8) (244 20) (244 33) (20 8)]", 1]]

ER_PPI =   [["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 1480) ('type_1' 'type_1') ('type_1' 1281) (1281 1480)]", 1303],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 1643) ('type_1' 'type_1') ('type_1' 1643) (1643 1281)]", 127],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 1437) ('type_1' 'type_1') ('type_1' 1437) (1004 1437)]", 155],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 581) ('type_1' 'type_1') ('type_1' 736) (581 736)]", 102],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 650) ('type_1' 'type_1') ('type_1' 537) (537 650) (650 537)]", 2],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 650) ('type_1' 'type_1') ('type_1' 650) (537 650) (650 537)]", 6],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 32) ('type_1' 'type_1') ('type_1' 170) (4 32) (4 252) (170 4) (170 32) (252 4) (252 170)]", 1],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 17) ('type_0' 74) ('type_1' 'type_1') ('type_1' 8) ('type_1' 13) (1 8) (1 13) (1 17) (3 1) (3 74) (8 3) (8 17) (13 1) (13 3) (17 74) (74 3)]", 1]]

ORIG_BLOGS = [["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 896) ('type_1' 'type_1') ('type_1' 446) (896 446)]", 323],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 896) ('type_1' 'type_1') ('type_1' 446) (896 446) (913 446)]", 6],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 515) ('type_1' 'type_1') ('type_1' 515) (515 446)]", 127],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 772) ('type_1' 'type_1') ('type_1' 772) (446 772)]", 226],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 652) ('type_1' 'type_1') ('type_1' 446) (446 652) (652 446)]", 88],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 1068) ('type_1' 'type_1') ('type_1' 163) (92 163) (1066 163) (1068 163)]", 3],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 887) ('type_1' 'type_1') ('type_1' 887) (885 887) (887 885)]", 155],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 987) ('type_1' 'type_1') ('type_1' 871) (871 987)]", 99],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 193) ('type_0' 265) ('type_1' 'type_1') ('type_1' 265) (193 265) (265 193)]", 39],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 193) ('type_1' 'type_1') ('type_1' 193) (193 383) (193 679) (193 680)]", 8],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 209) ('type_1' 'type_1') ('type_1' 209) ('type_1' 210) (209 210)]", 14],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 4) ('type_0' 293) ('type_1' 'type_1') ('type_1' 4) (293 4)]", 21],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 4) ('type_1' 'type_1') ('type_1' 4) (405 4) (657 4)]", 9],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 135) ('type_1' 'type_1') ('type_1' 317) ('type_1' 135) (135 317) (317 135)]", 13],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 106) ('type_0' 44) ('type_1' 'type_1') ('type_1' 106) ('type_1' 44) (44 106) (106 44)]", 11],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 107) ('type_0' 44) ('type_1' 'type_1') ('type_1' 107) (107 44)]", 16],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 512) ('type_1' 'type_1') ('type_1' 512) ('type_1' 511) (511 512)]", 12],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 426) ('type_0' 427) ('type_1' 'type_1') ('type_1' 426) ('type_1' 427) (426 427)]", 9],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 64) ('type_0' 42) ('type_0' 4) ('type_1' 'type_1') ('type_1' 64) ('type_1' 2) (2 4) (2 42) (4 42) (4 64) (42 2) (42 4) (42 64) (64 2) (64 4) (64 42)]", 1],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 47) ('type_0' 7) ('type_1' 'type_1') ('type_1' 3) ('type_1' 47) (1 3) (1 7) (1 47) (3 1) (3 7) (3 47) (7 1) (7 3) (47 7)]", 1]]

CL_BLOGS = [["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 208) ('type_1' 'type_1') ('type_1' 208) (208 257)]", 191],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 208) ('type_1' 'type_1') ('type_1' 208) (208 257) (776 208)]", 7],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 208) ('type_1' 'type_1') ('type_1' 514) (514 208)]", 237],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 520) ('type_1' 'type_1') ('type_1' 208) (520 208)]", 321],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 208) ('type_1' 'type_1') ('type_1' 208) (776 208)]", 345],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 114) ('type_1' 'type_1') ('type_1' 114) ('type_1' 1155) (114 1155)]", 5],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 128) ('type_1' 'type_1') ('type_1' 70) (70 128) (128 70)]", 50],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 274) ('type_1' 'type_1') ('type_1' 274) (45 274) (274 45)]", 26],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 14) ('type_1' 'type_1') ('type_1' 68) ('type_1' 14) (14 68) (68 14)]", 7],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 29) ('type_0' 14) ('type_1' 'type_1') ('type_1' 14) (14 29) (29 14)]", 10],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 29) ('type_0' 46) ('type_1' 'type_1') ('type_1' 29) (29 46)]", 6],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 29) ('type_0' 14) ('type_1' 'type_1') ('type_1' 29) ('type_1' 14) (14 29) (29 14)]", 4],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 17) ('type_0' 19) ('type_1' 'type_1') ('type_1' 17) ('type_1' 19) (19 17)]", 4],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 1) ('type_0' 12) ('type_0' 5) ('type_1' 'type_1') ('type_1' 1) ('type_1' 5) (1 12) (1 26) (12 1) (12 5) (12 26) (26 1) (26 12) (5 12)]", 1]]

ER_BLOGS = [["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 877) ('type_1' 'type_1') ('type_1' 1121) (877 1121)]", 54],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 1075) ('type_1' 'type_1') ('type_1' 877) (877 1075)]", 1032],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 1109) ('type_1' 'type_1') ('type_1' 1109) (877 1109)]", 48],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 1152) ('type_1' 'type_1') ('type_1' 1152) (1152 1117)]", 18],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 141) ('type_1' 'type_1') ('type_1' 141) (141 1080) (1080 141)]", 4],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 406) ('type_1' 'type_1') ('type_1' 540) (406 540) (540 406)]", 57],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_0' 73) ('type_0' 6) ('type_1' 'type_1') ('type_1' 73) (73 6)]", 5],\
            ["[('type_0' 'type_0') ('type_0' 'type_1') ('type_1' 'type_1') (1 7) (1 23) (5 1) (5 7) (7 23) (10 1) (10 8) (23 1) (23 7) (23 8) (23 10) (8 5) (8 10)]", 1]]

def graph_from_string(the_string):
    pieces = re.split('\[\(| |\) \(|\)\]', the_string)
    pieces = pieces[1:(len(pieces)-1)]
    nodes = set()
    for piece in pieces:
        nodes.add(piece)
    G = nx.DiGraph()
    for node in nodes:
        G.add_node(node)
    idx = 0
    while idx < len(pieces):
        G.add_edge(pieces[idx], pieces[idx + 1])
        idx += 2
    return G

def make_graph_list(the_list):
    return [[graph_from_string(x[0]), x[1]] for x in the_list]

def merge_graph_lists(list_a, list_b, list_c):
    main_list = [[x[0], x[1], 0, 0] for x in list_a]
    for graph_and_count_b in list_b:
        found_match = False
        for i in range(0, len(main_list)):
            gm = isomorphism.DiGraphMatcher(graph_and_count_b[0], main_list[i][0])
            if gm.is_isomorphic():
                main_list[i][2] = graph_and_count_b[1]
                found_match = True
                break
        if not found_match:
            main_list.append([graph_and_count_b[0], 0, graph_and_count_b[1], 0])

    for graph_and_count_c in list_c:
        found_match = False
        for main_gc in main_list:
            gm = isomorphism.DiGraphMatcher(graph_and_count_c[0], main_gc[0])
            if gm.is_isomorphic():
                main_gc[3] = graph_and_count_c[1]
                found_match = True
                break
        if not found_match:
            main_list.append([graph_and_count_c[0], 0, 0, graph_and_count_c[1]])
    return main_list

def merged_probabilities(merged_lists):
    sum_a = 0.0
    sum_b = 0.0
    sum_c = 0.0
    for list_element in merged_lists:
        sum_a += list_element[1] + 1.0
        sum_b += list_element[2] + 1.0
        sum_c += list_element[3] + 1.0
    return [[x[0], (x[1] + 1.0) / sum_a, (x[2] + 1.0) / sum_b, (x[3] + 1.0) / sum_c] for x in merged_lists]

def ratios(prob_list, idx_1, idx_2):
    ratio_list = [[x[0], x[idx_1] / x[idx_2] if x[idx_1] > x[idx_2] else -1.0 * x[idx_2] / x[idx_1]] for x in prob_list]
    ratio_list.sort(key=(lambda x: -1.0 * abs(x[1])))
    return ratio_list

def kl_contributions(prob_list, idx_1, idx_2):
    contributions = [[x[0], -1.0 * x[idx_1] * math.log(x[idx_2] / x[idx_1])] for x in prob_list]
    contributions.sort(key=(lambda x: -1.0 * x[1]))
    return contributions

def KL_divergence(kl_contributions):
    kl = 0.0
    for item in kl_contributions:
        kl += item[1]
    return kl

def display_kl_contributions(some_kl, title):
    print(title)
    for i in range(0, min(len(some_kl), 3)):
        edges = list(some_kl[i][0].edges())
        edges.sort()
        print("%s %s" % (edges, some_kl[i][1]))

orig_ppi_list = make_graph_list(ORIG_PPI)
cl_ppi_list = make_graph_list(CL_PPI)
er_ppi_list = make_graph_list(ER_PPI)
merged_list = merge_graph_lists(orig_ppi_list, cl_ppi_list, er_ppi_list)
probs_list = merged_probabilities(merged_list)
orig_cl_kl_contributions = kl_contributions(probs_list, 1, 2)
orig_er_kl_contributions = kl_contributions(probs_list, 1, 3)
display_kl_contributions(orig_cl_kl_contributions, "PPI: Original vs CL")
display_kl_contributions(orig_er_kl_contributions, "PPI: Original vs ER")
orig_cl_kl = KL_divergence(orig_cl_kl_contributions)
print("orig_cl_kl: %s" % orig_cl_kl)
orig_er_kl = KL_divergence(orig_er_kl_contributions)
print("orig_er_kl: %s" % orig_er_kl)

orig_blogs_list = make_graph_list(ORIG_BLOGS)
cl_blogs_list = make_graph_list(CL_BLOGS)
er_blogs_list = make_graph_list(ER_BLOGS)
merged_list = merge_graph_lists(orig_blogs_list, cl_blogs_list, er_blogs_list)
probs_list = merged_probabilities(merged_list)
orig_cl_kl_contributions = kl_contributions(probs_list, 1, 2)
orig_er_kl_contributions = kl_contributions(probs_list, 1, 3)
display_kl_contributions(orig_cl_kl_contributions, "Blogs: Original vs CL")
display_kl_contributions(orig_er_kl_contributions, "Blogs: Original vs ER")
orig_cl_kl = KL_divergence(orig_cl_kl_contributions)
print("orig_cl_kl: %s" % orig_cl_kl)
orig_er_kl = KL_divergence(orig_er_kl_contributions)
print("orig_er_kl: %s" % orig_er_kl)


from os.path import join, basename
from os import environ
from sys import stderr

import numpy as np
from SetCoverPy import setcover, mathutils
from matplotlib import pyplot as plt
import seaborn as sns
from astropy.io import fits

from redmonster.datamgr.io2 import read_ndArch
from _linelist import __linelist__

# Dict of plates and corresponding MJDS
plates = {8123:56931}

# version of redmonster reductions to use
rmver = 'v1_1_0'

# chi2 threshold to use in set cover distance matrix
mindist = 0.48

rmdir = join( environ['REDMONSTER_SPECTRO_REDUX'], environ['RUN2D'],
             '%s' % rmver)

# loop over all plates
for plate in plates.keys():
    # list of fibers classified as galaxy with zwarning=0
    fibers = []
    zs = []
    # open redmonster file
    hdu = fits.open( join( rmdir, '%s' % plate, 'redmonster-%s-%s.fits'
                          % (plate, plates[plate]) ) )
    # open data file
    idlpath = join( environ['BOSS_SPECTRO_REDUX'], environ['RUN2D'],
                   '%s' % plate, 'spPlate-%s-%s.fits' % (plate, plates[plate]) )
    hduidl = fits.open(idlpath)

    # create wavelength solution
    wave = 10**(hduidl[0].header['COEFF0'] + hduidl[0].header['COEFF1'] *
                np.arange(hduidl[0].header['NAXIS1']))

    # loop over fibers and keep those classed as gal with zwarn=0
    for i,fiber in enumerate(hdu[1].data.FIBERID):
        if fiber != 516:
            if hdu[1].data.ZWARNING[i] == 0:
                if hdu[1].data.CLASS1[i] == 'ssp_galaxy_glob':
                    fibers.append(fiber)
                    zs.append(hdu[1].data.Z1[i])

    # nfibers x nfibers matrix to hold distances
    distmat = np.zeros( (len(fibers), len(fibers)) )

    # binary distance matrix
    #binmat = np.zeros( distmat.shape )

    # calculate distance from each fiber to every other fiber
    for i in range(len(fibers)):
        # set data and such for this fiber
        fiber1 = fibers[i]
        data1 = hduidl[0].data[fiber1]
        sigma1 = hduidl[1].data[fiber1]
        wave1 = wave / (1+zs[i])
        # normalize slow spectrum to mean value between 4500 and 4700 A, and
        # scale errors accordingly
        high1 = np.abs(wave1-4700).argmin()
        low1 = np.abs(wave1-4500).argmin()
        mean1 = np.mean(data1[low1:high1+1])
        data1 /= mean1
        sigma1 *= mean1**2
        # loop over upper half of matrix and fill bottom half with symmetry
        for j in range(i, len(fibers)):
            fiber2 = fibers[j]
            stderr.write("\rTotal: %s i: %s j:%s" % (len(fibers), (i+1), (j+1)) )
            # convert to rest-frame wavelength
            wave2 = wave / (1+zs[j])
            # crop spectra to only overlapping region
            if wave1[0] < wave2[0]:
                low_bound = np.abs(wave1 - wave2[0]).argmin()
                data1_2 = data1[low_bound:]
                sigma1_2 = sigma1[low_bound:]
                wave1_2 = wave1[low_bound:]
                high_bound = np.abs(wave2 - wave1[-1]).argmin()
                data2 = hduidl[0].data[fiber2][:high_bound+1]
                sigma2 = hduidl[1].data[fiber2][:high_bound+1]
                wave2 = wave2[:high_bound+1]
            elif wave1[0] > wave2[0]:
                low_bound = np.abs(wave2 - wave1[0]).argmin()
                data2 = hduidl[0].data[fiber2][low_bound:]
                sigma2 = hduidl[1].data[fiber2][low_bound:]
                wave2 = wave2[low_bound:]
                high_bound = np.abs(wave1 - wave2[-1]).argmin()
                data1_2 = data1[:high_bound+1]
                sigma1_2 = sigma1[:high_bound+1]
                wave1_2 = wave1[:high_bound+1]
            else: # same redshift (most likely same spectrum)
                data1_2 = data1
                sigma1_2 = sigma1
                data2 = hduidl[0].data[fiber2]
                sigma2 = hduidl[1].data[fiber2]
            # normalize fast spectrum to mean value between 4500 and 4700 A, and
            # scale errors accordingly
            low2 = np.abs(wave2-4500).argmin()
            high2 = np.abs(wave2-4700).argmin()
            mean2 = np.mean(data2[low2:high2+1])
            data2 /= mean2
            sigma2 *= mean2**2
            # convert inverse variance to sigma squared and add in quadrature
            variance = (1/sigma1_2) + (1/sigma2)
            # calculate reduced chi2 as distance
            distmat[i][j] = distmat[j][i] = np.sum( (data1_2 - data2)**2 /\
                                                   variance ) / data1.shape[0]
            #if distmat[i][j] < mindist:
                #binmat[i][j] = 1
    print " "
    binmat = distmat < mindist
    cost = np.ones(binmat.shape[0])
    #cost = 1 / hdu

    g = setcover.SetCover(binmat, cost)
    #g.greedy()
    g.SolveSCP()

    # Get the archetype indices
    iarchetype = np.nonzero(g.s)[0]
    print "Number of archetypes: %s\n" % iarchetype.shape

    # How many are covered by each archetype?
    n_rep = np.sum(binmat[:, iarchetype], axis=0)
    for i,arch in enumerate(iarchetype):
        if n_rep[i] == 1:
            print "Archetype #%s represents %s spectra." % (arch, n_rep[i])
        elif n_rep[i] > 1:
            print "Archetype #%s represents %i spectra." % (arch, n_rep[i])

    # Get instances represented by each archetype
    stacks = [] # place to store stacks
    all_counts = []
    restwave = 10**(2.6990 + np.arange(20000) * 0.0001)
    redshifts = []
    fibnums = []
    archfibs = []
    j = -1
    for i in xrange(iarchetype.size):
        if n_rep[i] > 1:
            j += 1
            redshifts.append([])
            fibnums.append([])
            this_stack = np.zeros(20000)
            counts = np.zeros(20000)
            # binary vector of which spectra are represented by this archetype
            itmp = binmat[:, iarchetype[i]]
            # list of fiber numbers represented
            fibrep = np.asarray(fibers)[itmp]
            for fiber in fibrep:
                fiberind = np.where((np.asarray(fibers)-fiber) == 0)[0][0]
                fiberlen = hduidl[0].data[fiber-1].shape[-1]
                this_wave = wave / (1 + zs[fiberind])
                wave0 = np.abs(restwave - this_wave[0]).argmin()
                this_stack[wave0:wave0+fiberlen] += hduidl[0].data[fiber-1]
                counts[wave0:wave0+fiberlen] += 1
                redshifts[j].append(zs[fiberind])
                fibnums[j].append(fiber)
            this_stack /= counts
            stacks.append(this_stack)
            all_counts.append(counts)
            archfibs.append( fibers[iarchetype[i]] )

    # Crop stacks to regions with at least XXX% of max counts
    mastercounts = []
    masterstacks = []
    masterwave = []
    for i,counts in enumerate(all_counts):
        w = (counts >= (0.3*np.max(counts)))
        mastercounts.append(counts[w])
        masterstacks.append(stacks[i][w])
        masterwave.append(restwave[w])


# -*- coding: utf-8 -*-
"""
Created on Wed Aug 21 12:33:36 2019

Combined new simulation using missing data and finite mutation model

@author: marianne aspbury
"""

import msprime
import tsinfer
import numpy as np
from IPython.display import display
from IPython.display import SVG

import pickle
import random # to pick from list of poss roots of subtree
from IndividualClass import Individual # might say unused, but not true, needed to read pickled file
from truncate_ts_samples import truncate_ts_samples as truncate_ts_samples

### import the trees ###
file_loc = 'C:\\Users\\mazmysta\OneDrive - Nexus365\\BDI_proj\\HIV_Transmission_Networks_WillProbert\\19-08-08-first_example_network\\'

with open(file_loc + 'pickled_data_all.pickle', 'rb') as f:
    total_tree = pickle.load(f)

#proof it works
print('Pickle loaded, 1st root num kids: {}'.format(total_tree['-1_-1'].total_children_num))

# Saving list of people with desired num kids beneath
samples_poss = []
desired_overall_children = 5 ## Seemingly my max is 30 with 10 samples, 100 pop size and 1000 length genome.

for key in total_tree:
    if total_tree[key].total_children_num == desired_overall_children:
        samples_poss.append(key)

print(len(samples_poss))
# 8 for kids = 100

#pick only one random sample
random.seed(4) # repeatable choice(s) - sequence same
start = random.choice(samples_poss)
#start='58582_0'
print(start)

print(total_tree[start].infected_others)

def preorder_traversal(u):
    all_nodes = []
    stack = [u]
    while len(stack) > 0:
        v = stack.pop()
        if total_tree[v].direct_children_num > 0: #Returns True if the specified node is not a leaf
            stack.extend(total_tree[v].infected_others)
        all_nodes.append(v)
    return all_nodes

list_of_root_and_kids = preorder_traversal(start)

#print(preorder_traversal(start), len(preorder_traversal(start)))

#print(list(set(list_of_root_and_kids) - set(start)))

# all info required for mini-transmission chain
#for item in list_of_root_and_kids:
#    print(item, total_tree[item].infected_others,
#          total_tree[item].time_infected_others,
#          total_tree[item].age_birth,
#          total_tree[item].age_death)

### msprime assimilation
# Find total num of generations (days) to model over
times_loop = []

for item in list_of_root_and_kids:
    #float to take as number not string, and access list element with [0]
    times_loop.append(float(total_tree[item].time_infected_by[0]))

#since floats, can do maths
total_time = (max(times_loop)-min(times_loop))*365

#Furthest back in time is the smallest date (lowest year), do everything in diffs
gens_list = []
for item in list_of_root_and_kids:
    #float to take as number not string, and access list element with [0]
    diff_calc = 365*(max(times_loop)-float(total_tree[item].time_infected_by[0]))
    gens_list.append(diff_calc + 100) # add 100 so most recent infection is '100 gens (days) in past'

#    print(float(total_tree[item].time_infected_by[0]))

# info on dates so far
for i in range(len(gens_list)):
    print((list_of_root_and_kids)[i],
          #times_loop[i],
          gens_list[i])

######## msprime variable populations setup #########

## source population
PopSource = msprime.PopulationConfiguration(initial_size = 1e8, growth_rate = 0)

## number of populations in present time of model, not including source pop...
final_num_pops = len(list_of_root_and_kids)

## sample_sizes for each population and effective pop sizes
sample_size = 20
infection_size = 1
stable_pop_size = 100
# ## subpops based on infected people, all subpops that want to exist at end of sim (present time) need stated here
pop_list = [PopSource]

#Setting up the end, so all pops exist and at stable pop size, no death in simulation time
for pop in range(final_num_pops):
#    print(pop)
    pop_list.append(
        msprime.PopulationConfiguration(sample_size = sample_size, initial_size = stable_pop_size, growth_rate = 0)
                  )

# no migration between sources accross time, only infection events,
    # so migration matrix is zeros
M = np.zeros((final_num_pops+1,final_num_pops+1))

# Now get transmission events from the data. Use index as population number, but +1 since have fake source pop at index 0.
#for i in list_of_root_and_kids:
#    print(list_of_root_and_kids.index(i) + 1)

####--- new version with sub-pops ---####

## a simple model where independent sub-pop is infection derived from source pop
# if infected by true pop, need to state when diverged from past pop if that's the case
# Oddly source is the destination, i.e. direction of migration is dest -> source if forwards in time view.
# backwards in time means that destination is destination (but it's where the migration has come from)

transfers_list = []
for entry in range(len(pop_list)):
#    print(entry)
    if entry == 0: # ignore 0 since this is the source pop
        pass

    elif entry == 1: # 1 is root so needs own bit
        entry_ID = list_of_root_and_kids[entry-1]
#        print(entry, entry_ID)
        dest_index = 0 #infected from source
        transfer_time = gens_list[entry-1] # time infected still stored.
#        print(transfer_time)
        transfers_list.append(msprime.MassMigration(time = transfer_time, source = entry, dest = dest_index, proportion = 1))

    elif entry > 1: # 1 is root so needs own bit
        # get the index of the infected_by population
        # index of current population is its index in pop_list (index in list_of_root... + 1)
        entry_ID = list_of_root_and_kids[entry-1]
#        print(entry, entry_ID)
        dest_ID = total_tree[entry_ID].infected_by[0]
        dest_index = list_of_root_and_kids.index(dest_ID) + 1
#        print(dest_index, dest_ID)
        transfer_time = gens_list[entry-1]
#        print(transfer_time)
        transfers_list.append(msprime.MassMigration(time = transfer_time, source = entry, dest = dest_index, proportion = 1))

#check as expected
print(transfers_list)

# compare
#for i in range(len(gens_list)):
#    print((list_of_root_and_kids)[i],
#          #times_loop[i],
#          gens_list[i])

## now have set of populations, the transfers for pops (infection events)
    # still need the bottlenecks (pop growth & stabilisation)
    # then can order(sort) the complete demography list by time and simulate

## Bottlenecks: add population growth in so infections from source pop are only ~ 1-5 virions, which then balloons to e.g. ~1000

#### Bottleneck list initiation and creation
Pop_bottleneck_ie_growth_list = []
for entry in range(len(pop_list)):
    if entry > 0: # ignore 0 since this is the source pop. Only need infection time for this so root doesn't need own case

        transfer_time = gens_list[entry-1]

        #infection size setting
        pop_entry_bottleneck_start = msprime.PopulationParametersChange(
                time = transfer_time, initial_size=infection_size, growth_rate=0, population = entry) #i.e. for epochs inf-> 100*entry, this is growth rate

        #growth after infection setting - trial that 20 gens in future (less time back) gives appropriate growth for these params of pop_size = 100, rate = 0.25
        pop_entry_bottleneck_end = msprime.PopulationParametersChange(
                time = transfer_time-20, growth_rate = 0.23, initial_size=stable_pop_size, population = entry) #i.e. for epochs inf-> 100*entry, this is growth rate

        #save to list for manip outside loop
        Pop_bottleneck_ie_growth_list.extend((pop_entry_bottleneck_start, pop_entry_bottleneck_end))

# put all events together then sort them
events = Pop_bottleneck_ie_growth_list + transfers_list
events_sorted = sorted(events, key=lambda x: x.time, reverse=False)

#check
#for event in events_sorted:
#    print(event.__dict__) # just for easier digestion of output

my_history = msprime.DemographyDebugger(
    population_configurations=pop_list, migration_matrix = M,
    demographic_events = events_sorted)

#my_history.print_history()

### plot how pop changes ##
#
#time_steps= range(1,int(np.max(gens_list))+100,2)
## print('pop0:', my_history.population_size_trajectory(time_steps)[:,0])
## print('pop1:', my_history.population_size_trajectory(time_steps)[:,1])
## print('time:', np.array(time_steps))
## plot the populations, matplotlib understands array of y's as multiple y's so don't need to call individually
#
#import matplotlib.pyplot as plt
#plt.figure(1)
##plt.plot(time_steps, my_history.population_size_trajectory(time_steps)[:,0])
##plt.plot(time_steps, my_history.population_size_trajectory(time_steps)[:,1])
##plt.rc('axes', prop_cycle=(cycler(color=['r', 'g', 'b', 'y'])))
#fig, ax = plt.subplots(figsize=(15, 6), dpi=80)
#ax.set_prop_cycle(color=["green", "blue", "red", "orange", "grey", "cyan", "black"][1:])
#plt.plot(time_steps, my_history.population_size_trajectory(time_steps)[:,1:], '--', alpha=0.5) # this will plot each y (pop size var.) separately
#plt.xlim(np.max(time_steps),0) # switch the order of time so present (0) is RHS and past is LHS (max time step)
##plt.xlim(np.max(time_steps),1) # switch the order of time so present (0) is RHS and past is LHS (max time step)
##plt.ylim(np.log(0.5), np.log(150))
##plt.axvline(x=100, color='k', linestyle='-', alpha=0.5) # add a vertical line for migration step
##plt.legend(('1','2','3','4','5','6'),loc='best')
##box = ax.get_position()
##ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
## Put a legend to the right of the current axis
##plt.yscale("log")
##plt.xscale("log")
#ax.legend(('1: ' + list_of_root_and_kids[0],
#           '2: ' + list_of_root_and_kids[1],
#           '3: ' + list_of_root_and_kids[2],
#           '4: ' + list_of_root_and_kids[3],
#           '5: ' + list_of_root_and_kids[4],
#           '6: ' + list_of_root_and_kids[5]),
# loc='center left', bbox_to_anchor=(1.02, 0.5))
#plt.show()
## time = 0 is present, larger time is past

############ Simulation time ############

#file_loc = 'C:\\Users\\mazmysta\\OneDrive - Nexus365\\BDI_proj\\scripts\\PopulationsHIV\\'
#
##how to save
#with open(file_loc + 'pickled_pop_list.pickle', 'wb') as f:
#    pickle.dump(pop_list, f)
#
#with open(file_loc + 'pickled_events_sorted.pickle', 'wb') as f:
#    pickle.dump(events_sorted, f)
#
#with open(file_loc + 'pickled_M.pickle', 'wb') as f:
#    pickle.dump(M, f)
#
## loading example
#with open(file_loc + 'pickled_pop_list.pickle', 'rb') as f:
#    pop_list = pickle.load(f)
#
## simulate this extended simple model
ts2 = msprime.simulate(population_configurations=pop_list, migration_matrix = M,
                       demographic_events = events_sorted,
                       length = 1000,
                       random_seed = 17, recombination_rate = 0.7e-4,
                       mutation_rate=2e-5, end_time=60000)

## simulate this extended simple model
#ts2 = msprime.simulate(population_configurations=pop_list, migration_matrix = M,
#                       demographic_events = events_sorted,
#                       length = 1000,
#                       random_seed = 17, recombination_rate = 0.7e-4)

# 5 pops
colour_map = {0:"grey", 1:"blue", 2:"red", 3:"orange", 4:"green", 5:"cyan", 6:"black"}
node_colours = {u.id: colour_map[u.population] for u in ts2.nodes()}

X = 1
i = 0
for tree in ts2.trees():
    if i < X:
        display(SVG(tree.draw(node_colours = node_colours,
                              height=800, width = 1000,
                              format='SVG',
                              tree_height_scale='log_time')))
        print("Tree {} covers [{:.2f}, {:.2f}); TMRCA = {:.4f}".format(
            tree.index, *tree.interval, tree.time(tree.roots[0])))
        print(tree.branch_length(46))
    else:
        break
    i+=1

#### truncation part

# Having run
# ReadingTransmissionNetworkCSV_OverallChildrenCalcs_save.py
# have a ts2 = simulated data

# ts2 from subsetting to msprime, length 100, picking just one population below
# simplification and subsetting

#ts3 = ts.simplify(samples=[3,8,0,4,7,2,6,5,1,9]
#)

truncated_ts3 = truncate_ts_samples(ts2, average_span=200, random_seed=123)

#print(truncated_ts3)
#<tskit.trees.TreeSequence object at 0x000001ED52232710>

#truncated_ts3.tables
#Out[34]: <tskit.tables.TableCollection at 0x1ed4f6c1860>

#truncated_ts3.tables.nodes
#Out[35]: <tskit.tables.NodeTable at 0x1ed52232908>

#print(truncated_ts3.tables)
#prints the tables - long
X = 10
i = 0
for tree in truncated_ts3.trees():
    if i < X:
        display(SVG(tree.draw(height=800, width = 1000,
                              format='SVG',
                              tree_height_scale='log_time')))
        print("Tree {} covers [{:.2f}, {:.2f}); TMRCA = {:.4f}".format(
            tree.index, *tree.interval, tree.time(tree.roots[0])))
        print(tree.branch_length(46))
    else:
        break
    i+=1

truncated_ts3.dump(file_loc + 'truncated_simulation_tree.trees')
# ts infer

sd = tsinfer.SampleData.from_tree_sequence(truncated_ts3, use_times=False)

ts_inferred = tsinfer.infer(sd, simplify=False)

ts_inferred = ts_inferred.simplify(filter_sites=False, keep_unary=True)

ts_inferred
ts_inferred.dump(file_loc + 'inferred_tree.trees')
#Out[43]: <tskit.trees.TreeSequence at 0x1ed55ca9710>

X = 10
i = 0
for tree in ts_inferred.trees():
    if i < X:
        display(SVG(tree.draw(height=400, width=600,
                              tree_height_scale='rank')))
        print("Tree {} covers [{:.2f}, {:.2f}); TMRCA = {:.4f}".format(
            tree.index, *tree.interval, tree.time(tree.roots[0])))
#        print(tree.branch_length(46))
    else:
        break
    i+=1

print(ts_inferred.genotype_matrix())
print(truncated_ts3.genotype_matrix())

haps = []
for i in ts2.haplotypes():
    haps.append(i)

sequence_IDs = []
for i in range(len(haps)):
    sequence_IDs.append(f'sample_{ts2.samples()[i]}_pop_{ts2.node(i).population}')

#fasta printing
#for i in range(len(haps)):
#    print(f'>{sequence_IDs[i]}\n{haps[i]}')

## write and save fasta file
with open(file_loc + 'test_fasta_file.txt', 'w') as f:
    for i in range(len(haps)):
        f.write(f'>{sequence_IDs[i]}\n{haps[i]}\n')


"""
invariants:
- distinct prototype prefixes have distinct node in the prefix graph, regardless of wildcard matches
- if a node has two or more distinct child nodes, all the node's links are tame

implies that any node with wild links has exactly one child node,
and all rules descending from node match without wildcards through the child prefix

- any node with two or more distinct child nodes must be tame (no wild links)
- all descendent rules of a tamed node must have wildcard disabled at that position
- except for root, wild nodes never have None links, and all links point to same child node
- any two distinct prototype prefixes must have distinct node paths in the prefix tree

if you're adding a new rule, nobody matched prototype (early tame None leaf)
if you're disabling a wildcard, someone matched it (wild node with non-None links)

should never take some wild links and then fail unmatched if a different non-wild link would have matched
all prototypes should have persistent trails with no wildcard links
that way, disabling wildcards for one rule does not break another rule

"""
import numpy as np

class PrefixTreeNode:
    def __init__(self, bound, value=None, added=0):
        self.bound = bound # maximum number of outgoing links
        self.value = value # value of incoming link
        self.added = added # integer timestamp when node is added
        self.tamed = np.iinfo(int).max # integer timestamp when node is tamed
        self.rule = None # for leaf nodes, the associated rule
        self.links = {} # links[v] = child node for value v

    def is_wild(self):
        child_ids = set(map(id, self.links.values()))
        return len(self.links) == self.bound and len(child_ids) == 1

    def tame(self, tamed=0):
        self.links = {value:node for value, node in self.links.items() if node.value == value}
        self.tamed = tamed

    def __str__(self, prefix=""):
        tamed = "" if self.tamed == np.iinfo(int).max else str(self.tamed)
        if self.rule != None: return "%sr%d [%d~>%s]\n" % (prefix, self.rule, self.added, tamed)
        if self.is_wild():
            result = prefix + "* [%d~>%s]\n" % (self.added, tamed)
            result += self.links[0].__str__(prefix+" ")
        else:
            result = ""
            p = " " if len(self.links) == 1 else "|"
            for v in range(self.bound):
                if v in self.links:
                    result += prefix + "%d [%d~>%s]\n" % (v, self.added, tamed)
                    result += self.links[v].__str__(prefix+p)
        return result

    def rules(self):
        if self.rule != None: return [self.rule]
        if self.is_wild(): return self.links[0].rules()
        result = []
        for v in range(self.bound):
            if v in self.links:
                result += self.links[v].rules()
        return result

    def copy(self):
        node = PrefixTreeNode(self.bound, self.value, self.added)
        node.tamed = self.tamed
        node.rule = self.rule
        if self.is_wild():
            child_node = self.links[0].copy()
            node.links = {v: child_node for v in range(self.bound)}
        else:
            node.links = {v: child_node.copy() for v, child_node in self.links.items()}
        return node

    def rewind(self, i):

        # Rewind current node and links
        self.links = {v: node for v, node in self.links.items() if node.added <= i}
        if self.tamed > i and len(self.links) > 0: # should imply len == 1
            _, node = self.links.popitem()
            for v in range(self.bound): self.links[v] = node
            self.tamed = np.iinfo(int).max

        # Recurse on children
        if self.is_wild():
            self.links[0].rewind(i)
        else:
            for child_node in self.links.values(): child_node.rewind(i)

class MacroDatabase:
    def __init__(self, domain, max_rules):

        self.domain = domain
        self.max_rules = max_rules
        self.num_rules = 0
        self.bounds = (7,) * domain.state_size()
        self.root = PrefixTreeNode(self.bounds[0])

        self.prototypes = np.empty((max_rules, domain.state_size()), dtype=int)
        self.wildcards = np.empty((max_rules, domain.state_size()), dtype=bool)
        self.costs = np.empty(max_rules, dtype=int)
        self.macros = [None] * max_rules
        self.permutations = np.empty((max_rules, domain.state_size()), dtype=int)

        self.added = np.ones(max_rules, dtype=int) * np.iinfo(int).max
        self.tamed = np.ones((max_rules, domain.state_size()), dtype=int) * np.iinfo(int).max

    def query(self, state):
        node = self.root
        for k,v in enumerate(state):
            if v not in node.links: return None
            node = node.links[v]
        return node.rule

    def tame(self, node, w, tamed=0):
        node.tame(tamed)
        for r in node.rules():
            self.wildcards[r,w] = False
            self.tamed[r,w] = tamed

    def add_rule(self, prototype, macro, cost, added=0):
        r = self.num_rules
        self.prototypes[r] = prototype
        self.costs[r] = cost
        self.macros[r] = macro
        self.permutations[r] = self.domain.execute(macro, np.arange(self.domain.state_size()))
        self.added[r] = added

        node = self.root
        for k, value in enumerate(prototype):
            child_bound = self.bounds[k+1] if k+1 < len(self.bounds) else 0
            if len(node.links) == 0:
                child_node = PrefixTreeNode(child_bound, value, added)
                node.links = {v: child_node for v in range(node.bound)}
            elif value not in node.links:
                node.links[value] = PrefixTreeNode(child_bound, value, added)
            elif node.links[value].value != value:
                self.tame(node, k, tamed=added)
                node.links[value] = PrefixTreeNode(child_bound, value, added)
            self.wildcards[r,k] = node.is_wild()
            node = node.links[value]
        node.rule = r

        self.num_rules += 1

    def disable(self, r, w, tamed=0):
        node = self.root
        for k in range(w): node = node.links[self.prototypes[r,k]]
        if node.is_wild(): self.tame(node, w, tamed)

    def apply_rule(self, r, state):
        return state[self.permutations[r]].copy()

    def copy(self):
        db = MacroDatabase(self.domain, self.max_rules)
        db.num_rules = self.num_rules
        db.root = self.root.copy()

        db.prototypes = self.prototypes.copy()
        db.wildcards = self.wildcards.copy()
        db.costs = self.costs.copy()
        db.macros = list(self.macros)
        db.permutations = self.permutations.copy()

        db.added = self.added.copy()
        db.tamed = self.tamed.copy()

        return db

    def rewind(self, i):
        # keep all edits up to and including i
        self.root.rewind(i)
        if (self.added > i).any():
            self.num_rules = np.argmax(self.added > i)
        self.added[self.num_rules:] = np.iinfo(int).max
        self.wildcards = self.tamed > i
        self.tamed[self.wildcards] = np.iinfo(int).max
        return self

    def shrink_wrap(self):
        # reduce max_rules to num_rules and shrink arrays in-place for lighter footprint
        self.max_rules = self.num_rules

        self.prototypes = self.prototypes[:self.num_rules].copy()
        self.wildcards = self.wildcards[:self.num_rules].copy()
        self.costs = self.costs[:self.num_rules].copy()
        self.macros = tuple(self.macros[:self.num_rules])
        self.permutations = self.permutations[:self.num_rules].copy()

        self.added = self.added[:self.num_rules].copy()
        self.tamed = self.tamed[:self.num_rules].copy()

        return self

if __name__ == "__main__":

    from cube import CubeDomain
    domain = CubeDomain(2)
    solved = domain.solved_state()

    md = MacroDatabase(domain, max_rules=10)

    result = md.query(solved)
    assert result == None
    print(md.root)

    md.add_rule(solved, (), 0)
    assert md.query(solved) == 0
    assert (md.wildcards[:md.num_rules] == True).all()
    print("-"*24)
    print(md.root)

    # disable all wildcards for solved
    for w in range(len(solved)): md.disable(0, w)

    assert md.query(solved) == 0
    assert (md.wildcards[:md.num_rules] == False).all()
    print("-"*24)
    print(md.root)

    state = domain.perform((0,1,1), solved)
    md.add_rule(state, ((0,1,3),), 1)

    print("-"*24)
    print(md.root)
    assert md.query(solved) == 0
    assert md.query(state) == 1
    assert not (md.wildcards[:md.num_rules] == True).all()
    assert not (md.wildcards[:md.num_rules] == False).all()

    state2 = domain.perform((0,1,2), solved)
    md.add_rule(state2, ((0,1,2),), 1)
    print("-"*24)
    print(md.root)

    state3 = domain.perform((1,1,1), solved)
    md.add_rule(state3, ((1,1,3),), 1)
    print("-"*24)
    print(md.root)

    # simulate adding rules and check queries match after
    md = MacroDatabase(domain, max_rules=10)
    md.add_rule(solved, (), 0)
    for w in range(len(solved)): md.disable(0, w)

    rng = np.random.default_rng()
    for r in range(1, 10):
        state = domain.random_state(20, rng)
        md.add_rule(state, (), 0)

    for r in range(md.num_rules):
        state = md.prototypes[r]
        result = md.query(state)
        brutes = np.flatnonzero(((md.prototypes == state) | md.wildcards).all(axis=1))
        print(r, result, brutes)
        if result not in brutes:
            print("predisable")
            print("-"*24)
            print(md.root)
        assert result in brutes

    for r in range(md.num_rules):
        for w in range(len(state)):
            if rng.uniform() < 0.1:
                md.disable(r, w)

    for r in range(md.num_rules):
        state = md.prototypes[r]
        result = md.query(state)
        brutes = np.flatnonzero(((md.prototypes == state) | md.wildcards).all(axis=1))
        print(r, result, brutes)
        if result not in brutes:
            print("postdisable")
            print("-"*24)
            print(md.root)
        assert result in brutes

    print("-"*24)
    print(md.root)

    # test rewinding
    md = MacroDatabase(domain, max_rules=10)
    md.add_rule(prototype=solved, macro=(), cost=0, added=0)
    for w in range(len(solved)): md.disable(0, w, tamed=0)

    rng = np.random.default_rng()
    for r in range(1, 5):
        state = domain.random_state(20, rng)
        md.add_rule(prototype=state, macro=(), cost=0, added=r)
        md.disable(rng.integers(r, endpoint=True), rng.integers(len(state)), tamed=r)

    for r in range(md.num_rules):
        state = md.prototypes[r]
        result = md.query(state)
        brutes = np.flatnonzero(((md.prototypes == state) | md.wildcards).all(axis=1))
        if result not in brutes:
            print("prerewind")
            print("-"*24)
            print(md.root)
        assert result in brutes

    print("prerewind")
    print("-"*24)
    print(md.root)

    md.rewind(3)
    md.rewind(2)
    for r in range(md.num_rules):
        state = md.prototypes[r]
        result = md.query(state)
        brutes = np.flatnonzero(((md.prototypes == state) | md.wildcards).all(axis=1))
        if result not in brutes:
            print("postrewind")
            print("-"*24)
            print(md.root)
        assert result in brutes

    print("postrewind")
    print("-"*24)
    print(md.root)

    # test copy
    md = MacroDatabase(domain, max_rules=10)
    md.add_rule(prototype=solved, macro=(), cost=0, added=0)
    for w in range(len(solved)): md.disable(0, w, tamed=0)

    rng = np.random.default_rng()
    num_rules = 4
    states = []
    for r in range(1, num_rules-1):
        state = domain.random_state(20, rng)
        md.add_rule(prototype=state, macro=(), cost=0, added=r)
        md.disable(rng.integers(r, endpoint=True), rng.integers(len(state)), tamed=r)
        states.append(state)

    md2 = md.copy()
    state = domain.random_state(20, rng)
    md2.add_rule(prototype=state, macro=(), cost=0, added=num_rules-1)
    for r in range(md.num_rules):
        for w in range(len(state)):
            md.disable(r, w, tamed=num_rules)

    print("orig")
    print("-"*24)
    print(md.root)
    print("copy")
    print("-"*24)
    print(md2.root)

    assert md.num_rules + 1 == md2.num_rules
    assert md.wildcards[:md.num_rules].sum() == 0
    assert md2.wildcards[:md.num_rules].sum() > 0
    assert md.query(state) == None
    assert md2.query(state) != None
    for state in states:
        assert md.query(state) != None
        assert md2.query(state) != None

    # test macro permutations
    md = MacroDatabase(domain, max_rules=1)
    state = domain.perform((0, 1, 1), solved)
    actions = ((1, 1, 1), (2, 1, 1))
    md.add_rule(prototype=solved, macro=actions, cost=0, added=0)
    for w in range(len(solved)):
        if state[w] == solved[w]: md.disable(0, w, tamed=0)

    print("one rule perm")
    print("-"*24)
    print(md.root)

    r = md.query(state)
    assert r == 0
    new_state = md.apply_rule(r, state)
    for action in actions: state = domain.perform(action, state)
    assert (state == new_state).all()

    # compare timing of prefix and brute queries, and array-based prefix
    rule_count = 5000

    md = MacroDatabase(domain, max_rules=rule_count)
    md.add_rule(solved, (), 0)
    for w in range(len(solved)): md.disable(0, w)

    rng = np.random.default_rng()
    for r in range(1, rule_count):
        state = domain.random_state(20, rng)
        md.add_rule(state, (), 0)

    print("counting nodes...")
    def count_nodes(node):
        if node.is_wild():
            return 1 + count_nodes(node.links[0])
        else:
            count = 0
            for v, child in node.links.items():
                count += count_nodes(child)
            return 1 + count
    print(f"{count_nodes(md.root)} total nodes")

    # make prefix tree array query
    class ArrayPrefixTree:
        def __init__(self, root):
            self.links = -np.ones((rule_count * (domain.state_size()+1), 7), dtype=int)
            self.rules = -np.ones(rule_count * (domain.state_size()+1), dtype=int)
            self.num_nodes = 0
            self.link(root)
            # very marginal benefit after tuplifying, otherwise marginally worse
            self.links = tuple(tuple(links) for links in self.links)
            self.rules = tuple(self.rules)
        def add(self, node):
            n = self.num_nodes
            if node.rule != None: self.rules[n] = node.rule
            self.num_nodes += 1
            return n
        def link(self, node):
            n = self.add(node)
            if node.is_wild():
                c = self.link(node.links[0])
                self.links[n,:] = c
            else:
                for v, child in node.links.items():
                    c = self.link(node.links[v])
                    self.links[n,v] = c
            return n
        def query(self, state):
            n = 0
            for k,v in enumerate(state):
                n = self.links[n][v]
                if n == -1: return None
            return self.rules[n]

    apt = ArrayPrefixTree(md.root)

    from time import perf_counter
    prefix_time = 0
    brute_time = 0
    array_time = 0

    for state in md.prototypes:

        assert md.query(state) == apt.query(state)

        start = perf_counter()
        result = md.query(state)
        prefix_time += perf_counter() - start

        start = perf_counter()
        result = np.flatnonzero(((state == md.prototypes) | md.wildcards).all(axis=1))
        brute_time += perf_counter() - start

        start = perf_counter()
        result = apt.query(state)
        array_time += perf_counter() - start

    print("prototype queries:")
    print("prefix time", prefix_time)
    print("brute time", brute_time)
    print("array time", array_time)

    prefix_time = 0
    brute_time = 0
    array_time = 0

    for s in range(rule_count):
        state = domain.random_state(20, rng)

        assert md.query(state) == apt.query(state)

        start = perf_counter()
        result = md.query(state)
        prefix_time += perf_counter() - start

        start = perf_counter()
        result = np.flatnonzero(((state == md.prototypes) | md.wildcards).all(axis=1))
        brute_time += perf_counter() - start

        start = perf_counter()
        result = apt.query(state)
        array_time += perf_counter() - start

    print("random state queries:")
    print("prefix time", prefix_time)
    print("brute time", brute_time)
    print("array time", array_time)


#!/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__ = ('IonExchangeNH3',)

#path to csv with all the inputs
#data_path = '/Users/lewisrowles/opt/anaconda3/lib/python3.8/site-packages/exposan/biogenic_refinery/_ion_exchange_NH3.csv'
#data_path = os.path.abspath(os.path.dirname('_ion_exchange_NH3.csv'))
data_path += 'sanunit_data/_ion_exchange_NH3.tsv'
# data_path += 'sanunit_data/_carbonizer_base.csv'

###
class IonExchangeNH3(SanUnit):
    '''
    Ion Exchange for N recovery from liquid stream. Concentrated NH3 is recovered.

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

    Parameters
    ----------
    ins : WasteStream (liquid), FreshResin, H2SO4

    outs : WasteStream, SpentResin, ConcentratedNH3

    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 = 3
    _N_outs = 3

# in _run: define influent and effluent streams and treatment processes
    def _run(self):
        waste, resin_in, H2SO4 = self.ins
        treated, resin_out, conc_NH3 = self.outs
        treated.copy_like(self.ins[0])
        resin_in.phase = 'l'
        resin_out.phase = 'l'
        conc_NH3.phase = 'l'

        #!!! During storage most N as urea goes to NH3, should that
        # conversion be added or just use total N here?
        N_recovered = waste.imass['NH3'] * (self.N_rec_2) # kg N / hr
        treated.imass['NH3'] =  waste.imass['NH3'] - N_recovered # kg N / hr
        conc_NH3.imass['NH3'] = N_recovered # kg N / hr

        # following Terpeh et al. 2018 estimates for regenerating resin
        # assume volume of eluent = volume of urine
        # !!! cost associated with this influent stream neg compared to cost of resin and acid?
        #conc_NH3.F_vol = waste.F_vol
        # 0.1 M H2SO4 (0.65%) used to regenerate resin

        # !!! need to add SO4 as a component?
        # conc_NH3.imass['SO4'] = 0.1 * 98 * conc_NH3.F_vol # kg SO4 / hr

        # !!! add resin and H2SO4 as influent streams and spent resin as effluent
        resin_demand_influent = (waste.TN / self.resin_lifetime / self.ad_density / 14) # kg resin / m3 treated
        resin_demand_time = resin_demand_influent * waste.F_vol # kg resin / hr
        resin_cost_day = resin_demand_time * 24 * self.cost_resin # $ resin / d
        resin_in.imass['Polystyrene'] = resin_demand_time
        resin_out.imass['Polystyrene'] = resin_demand_time

        acid_demand_influent = waste.TN * self.vol_H2SO4 / self.ad_density / 14 # L acid / L treated
        acid_demand_time = acid_demand_influent * waste.F_vol * 1000 * 1.83 # kg acid / hr
        acid_cost_day = acid_demand_time * 24 / 1000 * self.cost_H2SO4 # $ acid / d
        H2SO4.imass['H2SO4'] = acid_demand_time

        # set values needed for _design and _cost as attributes
        self.volume_treated = waste.F_vol * 1000 * 24 # L liq / d

    #_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

        self.quantity_columns = np.ceil(self.volume_treated /
                                        self.column_daily_loading_rate) # number of 0.4 m columns
        design['PVC'] = PVC_quant = (self.column_length * self.quantity_columns
                                     * self.pvc_mass) # kg PVC
        Tubing_quant = (self.tubing_length * self.quantity_columns
                                           * self.tubing_mass) # kg PE
        Tank_quant = (self.quantity_columns * self.tank_mass / 3) # number of tanks with one tank for three columns
        design['PE'] = PE_quant = Tubing_quant + Tank_quant

        self.construction = (
            Construction(item='PVC', quantity = PVC_quant, quantity_unit = 'kg'),
            Construction(item='PE', quantity = PE_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['PVC'] = (self.cost_PVC_column * self.column_length
                                      * self.quantity_columns)
        self.purchase_costs['Tubing'] = (self.cost_tubing * self.tubing_length
                                         * self.quantity_columns)
        self.purchase_costs['Tank'] = (self.quantity_columns * self.tank_cost / 3) # one tank for three columns
        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
        ix_replacement_parts_annual_cost = 0 # USD/yr only accounts for time running

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

        self.add_OPEX =  (ix_replacement_parts_annual_cost + ix_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