#!/usr/bin/env python

import sys, os
import copy
import time
import numpy as np
from skimage import measure as skmeasure
import waterlib as wl
import water_properties as wp
from netCDF4 import Dataset
import parmed as pmd
import pytraj as pt
from pymbar import mbar

Usage="""solute_water_structure topFile inBulk doReweight
  Computes various properties of water in the hydration shell of a solute. This
  is intended to be used with expanded-ensemble simulation inserting a molecule
  into either bulk solution or at an interface. This requires that all configurations
  be reweighted to the appropriate ensemble, which includes removing restraints
  that keep the solute in a specific part of the interface. Two functions related
  to that which computes the solvation free energy, etc. are provided to help with
  this. For simulations in bulk, inBulk should be True to use the appropriate
  reweighting calculation, while if inBulk is False it reweights for a solute at
  an interface - if this is the case, you must have directories named Quad*
  that restrain the solute to different quadrants on the surface. Must be run
  from the directory containing either the four directories with simulations
  at each quadrant restraint or the bulk simulation directory with a single
  trajectory.
  Inputs:
    topFile - topology file associated with trajectories
    inBulk - (default False) flag to say if simulation in bulk or at an interface
    doReweight - (default True) flag for whether or not need to reweight (i.e. in
                                expanded ensemble or not)
  Outputs (all files, no returns):

"""

def getConfigWeightsSurf(kB=0.008314459848, T=298.15):
  """Computes and returns the configuration weights for simulations with a solute
   at an interface.
   Mostly replicates calcdGsolv in genetic_lib, but returns config weights in both
   the fully coupled and decoupled states (also includes pV term - won't matter very
   much for free energy differences, but maybe matters for weighting configurations,
   even though also probably not too much).
  """
  #First define directory structure, spring constants, etc.
  simDirs = ['Quad_0.25X_0.25Y', 'Quad_0.25X_0.75Y', 'Quad_0.75X_0.25Y', 'Quad_0.75X_0.75Y']
  kXY = [10.0, 10.0, 10.0, 10.0] #spring constant in kJ/mol*A^2
  refX = [7.4550, 7.4550, 22.3650, 22.3650]
  refY = [8.6083, 25.8249, 8.6083, 25.8249]
  distRefX = [7.4550, 7.4550, 7.4550, 7.4550]
  distRefY = [8.6083, 8.6083, 8.6083, 8.6083]
  numStates = 19

  #And some constants
  kBT = kB*T
  beta = 1.0 / kBT

  #First make sure all the input arrays have the same dimensions
  numSims = len(simDirs)
  allLens = np.array([len(a) for a in [kXY, refX, refY, distRefX, distRefY]])

  #Want to loop over all trajectories provided, storing solute position information to calculate restraints
  xyPos = None #X and Y coordinates of first heavy atom for all solutes - get shape later
  nSamps = np.zeros((len(simDirs), numStates), dtype=int) #Have as many x-y restraints as sims and same number of lambda states for each
  allPots = np.array([[]]*numStates).T #Potential energies, EXCLUDING RESTRAINT, for each simulation frame and lambda state
                                       #Will also include pV term because may matter for configurations
  xyBox = np.zeros(2)

  for i, adir in enumerate(simDirs):

    topFile = "%s/../sol_surf.top"%adir
    trajFile = "%s/prod.nc"%adir
    alchemicalFile = "%s/alchemical_output.txt"%adir

    #First load in topology and get atom indices
    top = pmd.load_file(topFile)

    #Get solute heavy atoms for each solute
    #Also get indices of surface atoms to use as references later
    #Only taking last united atoms of first SAM molecule we find
    heavyIndices = []
    for res in top.residues:
      if res.name not in ['OTM', 'CTM', 'STM', 'NTM', 'SOL']: #Assumes working with SAM surface...
        thisheavyinds = []
        for atom in res.atoms:
          if not atom.name[0] == 'H':
            thisheavyinds.append(atom.idx)
        heavyIndices.append(thisheavyinds)

    #Make into arrays for easier referencing
    heavyIndices = np.array(heavyIndices)

    #Load in the potential energies, INCLUDING RESTRAINT, at all states for this simulation to figure out frames to skip
    alcDat = np.loadtxt(alchemicalFile)
    startTime = alcDat[0, 1]
    startFrame = int(startTime) - 1 #Be careful here... need write frequency in alchemical file to match exactly with positions
                                    #AND assuming that have written in 1 ps increments...
                                    #Also, first frame in trajectory is NOT at time zero, so subtract 1
    thisPot = alcDat[:, 3:-1]
    thispV = alcDat[:, -1]

    #Next load in the trajectory and get all solute coordinates that matter
    top.rb_torsions = pmd.TrackedList([])
    top = pt.load_parmed(top, traj=False)
    traj = pt.iterload(trajFile, top, frame_slice=(startFrame, -1))
    nFrames = len(traj)
    xyBox = np.array(traj[0].box.values)[:2] #A little lazy, but all boxes should be same and fixed in X and Y dimensions

    thisxyPos = np.zeros((nFrames, len(heavyIndices), 2))
    thisnSamps = np.zeros(numStates, dtype=int)

    #Reference x and y coordinates for this restraint
    thisRefXY = np.array([refX[i], refY[i]])

    for j, frame in enumerate(traj):

      thisPos = np.array(frame.xyz)
      thisXY = thisPos[heavyIndices[:,0]][:, :2] #Takes XY coords for first heavy atom from each solute
      thisxyPos[j,:] = thisXY
      thisnSamps[int(alcDat[j, 2])] += 1 #Lambda states must be indexed starting at 0

      #Also get wrapped positions relative to each reference face
      #AND calculate xy restraint energy to remove by adding this for each solute
      xyEnergy = 0.0
      for k in range(len(heavyIndices)):
        xy = thisXY[k]
        #Then separately reimage around the restraint reference positions to calculate energy
        xy = wl.reimage([xy], thisRefXY, xyBox)[0] - thisRefXY
        xyEnergy += (  0.5*kXY[i]*(0.5*(np.sign(xy[0] - distRefX[i]) + 1))*((xy[0] - distRefX[i])**2)
                     + 0.5*kXY[i]*(0.5*(np.sign(xy[1] - distRefY[i]) + 1))*((xy[1] - distRefY[i])**2) )

      #Remove the restraint energy (only for x-y restraint... z is the same in all simulations)
      thisPot[j,:] -= (xyEnergy / kBT)

      #And also add in pV contribution
      thisPot[j,:] += thispV[j]

    #Add to other things we're keeping track of
    if xyPos is None:
      xyPos = copy.deepcopy(thisxyPos)
    else:
      xyPos = np.vstack((xyPos, thisxyPos))
    nSamps[i,:] = thisnSamps
    allPots = np.vstack((allPots, thisPot))

  #Now should have all the information we need
  #Next, put it into the format that MBAR wants, adding energies as needed
  Ukn = np.zeros((len(simDirs)*numStates, int(np.sum(nSamps))))

  for i in range(len(simDirs)):

    #First get energy of ith type of x-y restraint for all x-y positions
    thisRefXY = np.array([refX[i], refY[i]])
    #Must do by looping over each solute
    xyEnergy = np.zeros(xyPos.shape[0])
    for k in range(len(heavyIndices)):
      xy = wl.reimage(xyPos[:,k,:], thisRefXY, xyBox) - thisRefXY
      xyEnergy += (  0.5*kXY[i]*(0.5*(np.sign(xy[:,0] - distRefX[i]) + 1))*((xy[:,0] - distRefX[i])**2)
                   + 0.5*kXY[i]*(0.5*(np.sign(xy[:,1] - distRefY[i]) + 1))*((xy[:,1] - distRefY[i])**2) )

    #Loop over alchemical states with this restraint and add energy
    for j in range(numStates):

      Ukn[i*numStates+j, :] = allPots[:,j] + (xyEnergy / kBT)

  #Now should be set to run MBAR
  mbarObj = mbar.MBAR(Ukn, nSamps.flatten())

  #Following computePMF in MBAR to get configuration weights with desired potential of interest
  logwCoupled = mbarObj._computeUnnormalizedLogWeights(allPots[:,0])
  logwDecoupled = mbarObj._computeUnnormalizedLogWeights(allPots[:,-1])

  #Also report average solute-system LJ and coulombic potential energies in the fully coupled ensemble
  #(with restraints removed)
  #Just printing these values
  avgQ, stdQ = mbarObj.computeExpectations(allPots[:,0] - allPots[:,4], allPots[:,0])
  avgLJ, stdLJ = mbarObj.computeExpectations(allPots[:,4] - allPots[:,-1], allPots[:,0])
  print("\nAverage solute-system electrostatic potential energy: %f +/- %f"%(avgQ, stdQ))
  print("Average solute-system LJ potential energy: %f +/- %f\n"%(avgLJ, stdLJ))

  #Also print information that can be used to break free energy into components
  #Start by just printing all of the free energies between states
  alldGs, alldGerr = mbarObj.computePerturbedFreeEnergies(allPots.T)
  print("\nAll free energies relative to first (coupled) state:")
  print(alldGs.tolist())
  print(alldGerr.tolist())
  #And the free energy changes associated with just turning on LJ and elctrostatics separately
  dGq = alldGs[0][0] - alldGs[0][4]
  dGqErr = np.sqrt((alldGerr[0][0]**2) + (alldGerr[0][4])**2)
  print("\nElectrostatic dG (with LJ on): %f +/- %f"%(dGq, dGqErr))
  dGlj = alldGs[4][4] - alldGs[4][-1]
  dGljErr = np.sqrt((alldGerr[4][4]**2) + (alldGerr[4][-1])**2)
  print("\nLJ dG (no charges): %f +/- %f"%(dGlj, dGljErr))
  #Now calculate average potential energy differences needed for computing relative entropies
  dUq, dUqErr = mbarObj.computeExpectations(allPots[:,0] - allPots[:,4], allPots[:,0])
  print("\nAverage electrostatic potential energy in fully coupled state: %f +/- %f"%(dUq, dUqErr))
  dUlj, dUljErr = mbarObj.computeExpectations(allPots[:,4] - allPots[:,-1], allPots[:,4])
  print("\nAverage LJ potential energy (no charges) in uncharged state: %f +/- %f"%(dUlj, dUljErr))

  #And return weights after exponentiating log weights and normalizing
  wCoupled = np.exp(logwCoupled)
  wCoupled /= np.sum(wCoupled)
  wDecoupled = np.exp(logwDecoupled)
  wDecoupled /= np.sum(wDecoupled)

  return wCoupled, wDecoupled

def getConfigWeightsBulk(alchfile='alchemical_output.txt', kB=0.008314459848, T=298.15):
  """Given an alchemical output file, computes and returns configuration weights in
   both the fully coupled and decoupled ensembles of the solute.
  """
  rawdat = np.loadtxt(alchfile)
  lstates = rawdat[:,2]
  Ukn = rawdat[:,3:-1]
  pV = rawdat[:,-1] #pV term is in last column
  #pV term doesn't matter for free energy differences
  #But does matter for configuration weights (even though it's a small contribution)

  Nsamps = np.zeros(Ukn.shape[1], dtype=int)

  for i in range(Ukn.shape[1]):
    Nsamps[i] = int(np.sum((lstates==i)))
    Ukn[:,i] += pV

  #neworder = np.argsort(lstates)
  #Ukn = Ukn[neworder]

  Ukn /= (kB*T)

  mbarObj = mbar.MBAR(Ukn.T, Nsamps)

  #Following computePMF in MBAR to get configuration weights with desired potential of interest
  logwCoupled = mbarObj._computeUnnormalizedLogWeights(Ukn[:,0])
  logwDecoupled = mbarObj._computeUnnormalizedLogWeights(Ukn[:,-1])

  #Also report average solute-system LJ and coulombic potential energies in the fully coupled ensemble
  #(with restraints removed)
  #Just printing these values
  avgQ, stdQ = mbarObj.computeExpectations(Ukn[:,0] - Ukn[:,4], Ukn[:,0])
  avgLJ, stdLJ = mbarObj.computeExpectations(Ukn[:,4] - Ukn[:,-1], Ukn[:,0])
  print("\nAverage solute-water electrostatic potential energy: %f +/- %f"%(avgQ, stdQ))
  print("Average solute-water LJ potential energy: %f +/- %f\n"%(avgLJ, stdLJ))

  #Also print information that can be used to break free energy into components
  #Start by just printing all of the free energies between states
  alldGs, alldGerr = mbarObj.computePerturbedFreeEnergies(Ukn.T)
  print("\nAll free energies relative to first (coupled) state:")
  print(alldGs.tolist())
  print(alldGerr.tolist())
  #And the free energy changes associated with just turning on LJ and elctrostatics separately
  dGq = alldGs[0][0] - alldGs[0][4]
  dGqErr = np.sqrt((alldGerr[0][0]**2) + (alldGerr[0][4])**2)
  print("\nElectrostatic dG (with LJ on): %f +/- %f"%(dGq, dGqErr))
  dGlj = alldGs[4][4] - alldGs[4][-1]
  dGljErr = np.sqrt((alldGerr[4][4]**2) + (alldGerr[4][-1])**2)
  print("\nLJ dG (no charges): %f +/- %f"%(dGlj, dGljErr))
  #Now calculate average potential energy differences needed for computing relative entropies
  dUq, dUqErr = mbarObj.computeExpectations(Ukn[:,0] - Ukn[:,4], Ukn[:,0])
  print("\nAverage electrostatic potential energy in fully coupled state: %f +/- %f"%(dUq, dUqErr))
  dUlj, dUljErr = mbarObj.computeExpectations(Ukn[:,4] - Ukn[:,-1], Ukn[:,4])
  print("\nAverage LJ potential energy (no charges) in uncharged state: %f +/- %f"%(dUlj, dUljErr))

  #And return weights after exponentiating log weights
  wCoupled = np.exp(logwCoupled)
  wCoupled /= np.sum(wCoupled)
  wDecoupled = np.exp(logwDecoupled)
  wDecoupled /= np.sum(wDecoupled)

  return wCoupled, wDecoupled

def main(args):

  print time.ctime(time.time())

  #Get topology file we're working with
  topFile = args[0]

  #And figure out if we're dealing with solute at surface or in bulk
  if (args[1] == 'True'):
    inBulk = True
  else:
    inBulk = False

  if (args[2] == 'True'):
    doReweight = True
  else:
    doReweight = False

  #Read in topology file now to get information on solute atoms
  top = pmd.load_file(topFile)
  soluteInds = []
  for res in top.residues:
    if res.name not in ['OTM', 'CTM', 'STM', 'NTM', 'SOL']:
      for atom in res.atoms:
        soluteInds.append(atom.idx)

  #Now define how we compute three-body angles with bins and cut-off
  #Shell cut-off
  shellCut = 3.32 #1st minimum distance for TIP4P-Ew water at 298.15 K and 1 bar
  #Number of angle bins
  nAngBins = 100 #500
  #Define bin centers (should be nBins equally spaced between 0 and 180)
  angBinCents = 0.5 * (np.arange(0.0, 180.001, 180.0/nAngBins)[:-1] + np.arange(0.0, 180.001, 180.0/nAngBins)[1:])

  #And distance bins for local oxygen-oxygen RDF calculation
  #(really distance histograms from central oxygens - can normalize however we want, really)
  distBinWidth = 0.05
  nDistBins = int(shellCut / distBinWidth)
  distBins = np.arange(0.0, nDistBins*distBinWidth+0.00001, distBinWidth)
  distBinCents = 0.5 * (distBins[:-1] + distBins[1:])

  #Define the size of the probes used for assessing density fluctuations near solute
  probeRadius = 3.3 # radius in Angstroms; the DIAMETER of a methane so assumes other atoms methane-sized

  #And bins for numbers of waters in probes
  probeBins = np.arange(0.0, 21.00001, 1.0)
  nProbeBins = len(probeBins) - 1 #Will use np.histogram, which includes left edge in bin (so if want up to 20, go to 21)

  #And will record number waters in each solvation shell (histograms of)
  shellBins = np.arange(0.0, 251.00001, 1.0) #probably way too many bins, but don't want to run out
  nShellBins = len(shellBins) - 1

  #Should we do 2D histogram for angles and distances?
  #Or also do 2D histogram for number of waters in probe and three-body angle?
  #Interesting if make probe radius same size as three-body angle cutoff?

  #Finally, also define the bins for computing RDFs of waters near all solute atoms
  #Will use to define 1st and 2nd solvation shells
  rdfBinWidth = 0.2
  rdfMax = 12.00
  rdfBins = np.arange(0.0, rdfMax+0.000001, rdfBinWidth)
  rdfBinCents = 0.5 * (rdfBins[:-1] + rdfBins[1:])
  nRDFBins = len(rdfBinCents)
  rdfBinVols = (4.0*np.pi/3.0)*(rdfBins[1:]**3 - rdfBins[:-1]**3)
  bulkDens = 0.0332 #Roughly right for TIP4P-EW at 298.15 K and 1 bar in inverse Angstroms cubed

  #Need to create a variety of arrays to hold the data we're interested in
  #Will record distributions in both 1st and 2nd solvation shells
  shellCountsCoupled = np.zeros((nShellBins, 2)) #Histograms for numbers of waters in hydration shells of solutes
  probeHistsCoupled = np.zeros((nProbeBins, 2)) #Histograms for numbers waters in probes in 1st and 2nd hydration shells
  angHistsCoupled = np.zeros((nAngBins, 2)) #Histograms of three-body angles for water oxygens within solvation shells
  distHistsCoupled = np.zeros((nDistBins, 2)) #Histograms of distances to water oxygens from central oxygens
  solRDFsCoupled = np.zeros((nRDFBins, len(soluteInds))) #RDFs between each solute atom and water oxygens
  shellCountsDecoupled = np.zeros((nShellBins, 2)) #Same as above, but decoupled state, not coupled state
  probeHistsDecoupled = np.zeros((nProbeBins, 2))
  angHistsDecoupled = np.zeros((nAngBins, 2))
  distHistsDecoupled = np.zeros((nDistBins, 2))
  solRDFsDecoupled = np.zeros((nRDFBins, len(soluteInds)))

  #First need configuration weights to use in computing average quantities
  #But only do if we're using a simuation with an expanded ensemble
  if doReweight:
    if inBulk:
      weightsCoupled, weightsDecoupled = getConfigWeightsBulk(kB=1.0, T=1.0)
      #Using 1 for kB and T because alchemical_output.txt should already have potential energies in kBT
      simDirs = ['.']
    else:
      weightsCoupled, weightsDecoupled = getConfigWeightsSurf()
      simDirs = ['Quad_0.25X_0.25Y', 'Quad_0.25X_0.75Y', 'Quad_0.75X_0.25Y', 'Quad_0.75X_0.75Y']
  else:
    weightsCoupled = np.array([])
    weightsDecoupled = np.array([])
    simDirs = ['.']

  #To correctly match weights up to configurations, need to count frames from all trajectories
  countFrames = 0

  #Next, want to loop over all trajectories and compute RDFs from solute atoms to water oxygens
  #Will use this to define solvation shells for finding other properties
  #Actually, having looked at RDFs, just use 5.5 for first shell and 8.5 for second shell...
  #AND use all atoms, including hydrogens, which have LJ interactions in GAFF2, to define shells
  #Actually now only using heavy atoms... but when look at RDFs, examine all atoms
  for adir in simDirs:

    if doReweight:
      #Before loading trajectory, figure out how many frames to exclude due to weight equilibration
      alcDat = np.loadtxt(adir+'/alchemical_output.txt')
      startTime = alcDat[0, 1]
      startFrame = int(startTime) - 1
    else:
      startFrame = 0

    top = pmd.load_file(topFile)
    top.rb_torsions = pmd.TrackedList([]) #This is just for SAM systems so that it doesn't break pytraj
    top = pt.load_parmed(top, traj=False)
    traj = pt.iterload(adir+'/prod.nc', top, frame_slice=(startFrame, -1))

    if not doReweight:
      weightsCoupled = np.hstack((weightsCoupled, np.ones(len(traj))))
      weightsDecoupled = np.hstack((weightsDecoupled, np.ones(len(traj))))

    print("\nTopology and trajectory loaded from directory %s" % adir)

    owInds = top.select('@OW')
    soluteInds = top.select('!(:OTM,CTM,STM,NTM,SOL)')

    print("\n\tFound %i water oxygens" % len(owInds))
    print("\tFound %i solute atoms" % len(soluteInds))

    for i, frame in enumerate(traj):

      if i%1000 == 0:
        print "On frame %i" % i

      boxDims = np.array(frame.box.values[:3])

      currCoords = np.array(frame.xyz)

      #Wrap based on soluate atom center of geometry and get coordinates of interest
      wrapCOM = np.average(currCoords[soluteInds], axis=0)
      currCoords = wl.reimage(currCoords, wrapCOM, boxDims) - wrapCOM
      owCoords = currCoords[owInds]
      solCoords = currCoords[soluteInds]

      #Loop over solute atoms and find pair-distance histograms with water oxygens
      for j, acoord in enumerate(solCoords):
        solRDFsCoupled[:,j] += (weightsCoupled[countFrames+i]
                                * wl.pairdistancehistogram(np.array([acoord]), owCoords, rdfBinWidth, nRDFBins, boxDims))
        solRDFsDecoupled[:,j] += (weightsDecoupled[countFrames+i]
                                  * wl.pairdistancehistogram(np.array([acoord]), owCoords, rdfBinWidth, nRDFBins, boxDims))
        #Note that pairdistancehistogram is right-edge inclusive, NOT left-edge inclusive
        #In practice, not a big difference

    countFrames += len(traj)

  #Finish by normalizing RDFs properly
  for j in range(len(soluteInds)):
    solRDFsCoupled[:,j] /= rdfBinVols #bulkDens*rdfBinVols
    solRDFsDecoupled[:,j] /= rdfBinVols #bulkDens*rdfBinVols
  if not doReweight:
    solRDFsCoupled /= float(countFrames)
    solRDFsDecoupled /= float(countFrames)

  #And save to file
  np.savetxt('solute-OW_RDFs_coupled.txt', np.hstack((np.array([rdfBinCents]).T, solRDFsCoupled)),
             header='RDF bins (A)    solute atom-OW RDF for solute atom indices %s'%(str(soluteInds)))
  np.savetxt('solute-OW_RDFs_decoupled.txt', np.hstack((np.array([rdfBinCents]).T, solRDFsDecoupled)),
             header='RDF bins (A)    solute atom-OW RDF for solute atom indices %s'%(str(soluteInds)))

  print("\tFound RDFs for water oxygens from solute indices.")

  solShell1Cut = 5.5 #Angstroms from all solute atoms (including hydrogens)
  solShell2Cut = 8.5

  #And now that we know how many frames, we can assign real weights if not reweighting
  if not doReweight:
    weightsCoupled /= float(countFrames)
    weightsDecoupled /= float(countFrames)

  #Reset countFrames so get weights right
  countFrames = 0

  #Repeat looping over trajectories to calculate water properties in solute solvation shell
  for adir in simDirs:

    if doReweight:
      #Before loading trajectory, figure out how many frames to exclude due to weight equilibration
      alcDat = np.loadtxt(adir+'/alchemical_output.txt')
      startTime = alcDat[0, 1]
      startFrame = int(startTime) - 1
    else:
      startFrame = 0

    top = pmd.load_file(topFile)
    top.rb_torsions = pmd.TrackedList([]) #This is just for SAM systems so that it doesn't break pytraj
    top = pt.load_parmed(top, traj=False)
    traj = pt.iterload(adir+'/prod.nc', top, frame_slice=(startFrame, -1))

    print("\nTopology and trajectory loaded from directory %s" % adir)

    owInds = top.select('@OW')
    soluteInds = top.select('!(:OTM,CTM,STM,NTM,SOL)&!(@H=)')
    surfInds = top.select('(:OTM,CTM,STM,NTM)&!(@H=)') #For probe insertions, also include solute and surface heavy atoms

    print("\n\tFound %i water oxygens" % len(owInds))
    print("\tFound %i solute heavy atoms" % len(soluteInds))
    print("\tFound %i non-hydrogen surface atoms" % len(surfInds))

    if len(surfInds) == 0:
      surfInds.dtype=int

    for i, frame in enumerate(traj):

      #if i%10 == 0:
      #  print "On frame %i" % i

      boxDims = np.array(frame.box.values[:3])

      currCoords = np.array(frame.xyz)

      #Wrap based on soluate atom center of geometry and get coordinates of interest
      wrapCOM = np.average(currCoords[soluteInds], axis=0)
      currCoords = wl.reimage(currCoords, wrapCOM, boxDims) - wrapCOM
      owCoords = currCoords[owInds]
      solCoords = currCoords[soluteInds]
      surfCoords = currCoords[surfInds]

      #Now get solvent shells around solute
      shell1BoolMat = wl.nearneighbors(solCoords, owCoords, boxDims, 0.0, solShell1Cut)
      shell1Bool = np.array(np.sum(shell1BoolMat, axis=0), dtype=bool)

      shell2BoolMat = wl.nearneighbors(solCoords, owCoords, boxDims, solShell1Cut, solShell2Cut)
      shell2Bool = np.array(np.sum(shell2BoolMat, axis=0), dtype=bool)

      #And add weight to histogram for numbers of waters in shells
      thisCount1 = int(np.sum(shell1Bool))
      shellCountsCoupled[thisCount1, 0] += weightsCoupled[countFrames+i]
      shellCountsDecoupled[thisCount1, 0] += weightsDecoupled[countFrames+i]

      thisCount2 = int(np.sum(shell2Bool))
      shellCountsCoupled[thisCount2, 1] += weightsCoupled[countFrames+i]
      shellCountsDecoupled[thisCount2, 1] += weightsDecoupled[countFrames+i]

      #And compute water properties of solvent shells, first 3-body angles
      thisAngs1, thisNumAngs1 = wp.getCosAngs(owCoords[shell1Bool], owCoords, boxDims, highCut=shellCut)
      thisAngHist1, thisAngBins1 = np.histogram(thisAngs1, bins=nAngBins, range=[0.0, 180.0], density=False)
      angHistsCoupled[:,0] += weightsCoupled[countFrames+i] * thisAngHist1
      angHistsDecoupled[:,0] += weightsDecoupled[countFrames+i] * thisAngHist1

      thisAngs2, thisNumAngs2 = wp.getCosAngs(owCoords[shell2Bool], owCoords, boxDims, highCut=shellCut)
      thisAngHist2, thisAngBins2 = np.histogram(thisAngs2, bins=nAngBins, range=[0.0, 180.0], density=False)
      angHistsCoupled[:,1] += weightsCoupled[countFrames+i] * thisAngHist2
      angHistsDecoupled[:,1] += weightsDecoupled[countFrames+i] * thisAngHist2

      #And ow-ow pair distance histograms in both shells as well
      thisDistHist1 = wl.pairdistancehistogram(owCoords[shell1Bool], owCoords, distBinWidth, nDistBins, boxDims)
      distHistsCoupled[:,0] += weightsCoupled[countFrames+i] * thisDistHist1
      distHistsDecoupled[:,0] += weightsDecoupled[countFrames+i] * thisDistHist1

      thisDistHist2 = wl.pairdistancehistogram(owCoords[shell2Bool], owCoords, distBinWidth, nDistBins, boxDims)
      distHistsCoupled[:,1] += weightsCoupled[countFrames+i] * thisDistHist2
      distHistsDecoupled[:,1] += weightsDecoupled[countFrames+i] * thisDistHist2

      #Next compute distributions of numbers of waters in probes centered within each shell
      #To do this, create random grid of points in SQUARE that encompasses both shells
      #Then only keep points within each shell based on distance
      #Square will be based on shell cutoffs and min and max coordinates in each dimension of solute
      minSolX = np.min(solCoords[:,0]) - solShell2Cut
      maxSolX = np.max(solCoords[:,0]) + solShell2Cut
      minSolY = np.min(solCoords[:,1]) - solShell2Cut
      maxSolY = np.max(solCoords[:,1]) + solShell2Cut
      minSolZ = np.min(solCoords[:,2]) - solShell2Cut
      maxSolZ = np.max(solCoords[:,2]) + solShell2Cut
      thisGridX = minSolX + np.random.random(500)*(maxSolX - minSolX)
      thisGridY = minSolY + np.random.random(500)*(maxSolY - minSolY)
      thisGridZ = minSolZ + np.random.random(500)*(maxSolZ - minSolZ)
      thisGrid = np.vstack((thisGridX, thisGridY, thisGridZ)).T

      gridBoolMat1 = wl.nearneighbors(solCoords, thisGrid, boxDims, 0.0, solShell1Cut)
      gridBool1 = np.array(np.sum(gridBoolMat1, axis=0), dtype=bool)
      thisNum1 = wl.probegrid(np.vstack((owCoords, surfCoords, solCoords)), thisGrid[gridBool1], probeRadius, boxDims)
      thisProbeHist1, thisProbeBins1 = np.histogram(thisNum1, bins=probeBins, density=False)
      probeHistsCoupled[:,0] += weightsCoupled[countFrames+i] * thisProbeHist1
      probeHistsDecoupled[:,0] += weightsDecoupled[countFrames+i] * thisProbeHist1

      gridBoolMat2 = wl.nearneighbors(solCoords, thisGrid, boxDims, solShell1Cut, solShell2Cut)
      gridBool2 = np.array(np.sum(gridBoolMat2, axis=0), dtype=bool)
      thisNum2 = wl.probegrid(np.vstack((owCoords, surfCoords, solCoords)), thisGrid[gridBool2], probeRadius, boxDims)
      thisProbeHist2, thisProbeBins2 = np.histogram(thisNum2, bins=probeBins, density=False)
      probeHistsCoupled[:,1] += weightsCoupled[countFrames+i] * thisProbeHist2
      probeHistsDecoupled[:,1] += weightsDecoupled[countFrames+i] * thisProbeHist2

    countFrames += len(traj)

  #Should have everything we need, so save to text files
  np.savetxt('solute_shell_hists.txt',
             np.hstack((np.array([shellBins[:-1]]).T, shellCountsCoupled, shellCountsDecoupled)),
             header='Histograms of numbers of waters in first and second solute solvation shells with solvent in coupled (columns 2, 3) and decoupled (columns 4, 5) states')
  np.savetxt('solute_probe_hists.txt',
             np.hstack((np.array([probeBins[:-1]]).T, probeHistsCoupled, probeHistsDecoupled)),
             header='Number waters in probe histograms in first and second solute solvation shells with solvent in coupled (columns 2, 3) and decoupled (columns 4, 5) states')
  np.savetxt('solute_ang_hists.txt',
             np.hstack((np.array([angBinCents]).T, angHistsCoupled, angHistsDecoupled)),
             header='3-body angle histograms in first and second solute solvation shells with solvent in coupled (columns 2, 3) and decoupled (columns 4, 5) states')
  np.savetxt('solute_pair_hists.txt',
             np.hstack((np.array([distBinCents]).T, distHistsCoupled, distHistsDecoupled)),
             header='O-O pair-distance histograms in first and second solute solvation shells with solvent in coupled (columns 2, 3) and decoupled (columns 4, 5) states')

  print time.ctime(time.time())

if __name__ == "__main__":
  main(sys.argv[1:])


# base.py
# Author: Jacob Schreiber <jmschreiber91@gmail.com>

"""
This file contains code that implements the core of the submodular selection
algorithms.
"""

import numpy
from tqdm import tqdm

from ..optimizers import BaseOptimizer
from ..optimizers import NaiveGreedy
from ..optimizers import LazyGreedy
from ..optimizers import ApproximateLazyGreedy
from ..optimizers import TwoStageGreedy
from ..optimizers import StochasticGreedy
from ..optimizers import BidirectionalGreedy
from ..optimizers import GreeDi
from ..optimizers import SieveGreedy
from ..optimizers import OPTIMIZERS

from ..utils import PriorityQueue
from ..utils import check_random_state
from ..utils import _calculate_pairwise_distances

from scipy.sparse import csr_matrix

class BaseSelection(object):
	"""The base selection object.

	This object defines the structures that all submodular selection algorithms
	should follow. All algorithms will have the same public methods and the
	same attributes.

	Parameters
	----------
	n_samples : int
		The number of samples to return.

	initial_subset : list, numpy.ndarray or None, optional
		If provided, this should be a list of indices into the data matrix
		to use as the initial subset, or a group of examples that may not be
		in the provided data should beused as the initial subset. If indices,
		the provided array should be one-dimensional. If a group of examples,
		the data should be 2 dimensional.

	optimizer : string or optimizers.BaseOptimizer, optional
		The optimization approach to use for the selection. Default is
		'two-stage', which makes selections using the naive greedy algorithm
		initially and then switches to the lazy greedy algorithm. Must be
		one of

			'naive' : the naive greedy algorithm
			'lazy' : the lazy (or accelerated) greedy algorithm
			'approximate-lazy' : the approximate lazy greedy algorithm
			'two-stage' : starts with naive and switches to lazy
			'stochastic' : the stochastic greedy algorithm
			'greedi' : the GreeDi distributed algorithm
			'bidirectional' : the bidirectional greedy algorithm

		Default is 'naive'.

	optimizer_kwds : dict or None
		A dictionary of arguments to pass into the optimizer object. The keys
		of this dictionary should be the names of the parameters in the optimizer
		and the values in the dictionary should be the values that these
		parameters take. Default is None.

	reservoir : numpy.ndarray or None
		The reservoir to use when calculating gains in the sieve greedy
		streaming optimization algorithm in the `partial_fit` method.
		Currently only used for graph-based functions. If a numpy array
		is passed in, it will be used as the reservoir. If None is passed in,
		will use reservoir sampling to collect a reservoir. Default is None.

	max_reservoir_size : int
		The maximum size that the reservoir can take. If a reservoir is passed
		in, this value is set to the size of that array. Default is 1000.

	n_jobs : int
		The number of threads to use when performing computation in parallel.
		Currently, this parameter is exposed but does not actually do anything.
		This will be fixed soon.

	random_state : int or RandomState or None, optional
		The random seed to use for the random selection process. Only used
		for stochastic greedy.

	verbose : bool
		Whether to print output during the selection process.

	Attributes
	----------
	n_samples : int
		The number of samples to select.

	ranking : numpy.array int
		The selected samples in the order of their gain with the first number in
		the ranking corresponding to the index of the first sample that was
		selected by the greedy procedure.

	gains : numpy.array float
		The gain of each sample in the returned set when it was added to the
		growing subset. The first number corresponds to the gain of the first
		added sample, the second corresponds to the gain of the second added
		sample, and so forth.
	"""

	def __init__(self, n_samples, initial_subset=None, optimizer='lazy',
		optimizer_kwds={}, reservoir=None, max_reservoir_size=1000,
		n_jobs=1, random_state=None, verbose=False):
		if n_samples <= 0:
			raise ValueError("n_samples must be a positive value.")

		if not isinstance(initial_subset, (list, numpy.ndarray)) and initial_subset is not None:
			raise ValueError("initial_subset must be a list, numpy array, or None")
		if isinstance(initial_subset, (list, numpy.ndarray)):
			initial_subset = numpy.array(initial_subset)

		if not isinstance(optimizer, BaseOptimizer):
			if optimizer not in OPTIMIZERS.keys():
				raise ValueError("Optimizer must be an optimizer object or " \
					"a str in {}.".format(str(OPTIMIZERS.keys())))

		if isinstance(optimizer, BaseOptimizer):
			optimizer.function = self

		if verbose not in (True, False):
			raise ValueError("verbosity must be True or False")

		self.n_samples = n_samples
		self.metric = 'ignore'
		self.random_state = check_random_state(random_state)
		self.optimizer = optimizer
		self.optimizer_kwds = optimizer_kwds
		self.n_jobs = n_jobs
		self.verbose = verbose
		self.initial_subset = initial_subset

		self.ranking = None
		self.idxs = None
		self.gains = None
		self.subset = None
		self.sparse = None
		self._X = None

		self.sieve_current_values_ = None
		self.n_seen_ = 0
		self.reservoir_size = 0 if reservoir is None else reservoir.shape[0]
		self.reservoir = reservoir
		self.max_reservoir_size = max_reservoir_size if reservoir is None else reservoir.shape[0]
		self.update_reservoir_ = reservoir is None

	def fit(self, X, y=None, sample_weight=None, sample_cost=None):
		"""Run submodular optimization to select a subset of examples.

		This method is a wrapper for the full submodular optimization process.
		It takes in some data set (and optionally labels that are ignored
		during this process) and selects `n_samples` from it in the greedy
		manner specified by the optimizer.

		This method will return the selector object itself, not the transformed
		data set. The `transform` method will then transform a data set to the
		selected points, or alternatively one can use the ranking stored in
		the `self.ranking` attribute. The `fit_transform` method will perform
		both optimization and selection and return the selected items.

		Parameters
		----------
		X : list or numpy.ndarray, shape=(n, d)
			The data set to transform. Must be numeric.

		y : list or numpy.ndarray or None, shape=(n,), optional
			The labels to transform. If passed in this function will return
			both the data and th corresponding labels for the rows that have
			been selected.

		sample_weight : list or numpy.ndarray or None, shape=(n,), optional
			The weight of each example. Currently ignored in apricot but
			included to maintain compatibility with sklearn pipelines.

		sample_cost : list or numpy.ndarray or None, shape=(n,), optional
			The cost of each item. If set, indicates that optimization should
			be performed with respect to a knapsack constraint.

		Returns
		-------
		self : BaseGraphSelection
			The fit step returns this selector object.
		"""

		allowed_dtypes = list, numpy.ndarray, csr_matrix

		if not isinstance(X, allowed_dtypes):
			raise ValueError("X must be either a list of lists, a 2D numpy " \
				"array, or a scipy.sparse.csr_matrix.")
		if isinstance(X, numpy.ndarray) and len(X.shape) != 2:
			raise ValueError("X must have exactly two dimensions.")
		if numpy.min(X) < 0.0 and numpy.max(X) > 0.:
			raise ValueError("X cannot contain negative values or must be entirely "\
				"negative values.")
		if self.n_samples > X.shape[0]:
			raise ValueError("Cannot select more examples than the number in" \
				" the data set.")

		if not self.sparse:
			if X.dtype != 'float64':
				X = X.astype('float64')

		if isinstance(self.optimizer, str):
			optimizer = OPTIMIZERS[self.optimizer](function=self,
				verbose=self.verbose, random_state=self.random_state,
				**self.optimizer_kwds)
		else:
			optimizer = self.optimizer

		self._X = X if self._X is None else self._X
		self._initialize(X)

		if self.verbose:
			self.pbar = tqdm(total=self.n_samples, unit_scale=True)

		optimizer.select(X, self.n_samples, sample_cost=sample_cost)

		if self.verbose == True:
			self.pbar.close()

		self.ranking = numpy.array(self.ranking)
		self.gains = numpy.array(self.gains)
		return self

	def partial_fit(self, X, y=None, sample_weight=None, sample_cost=None):
		allowed_dtypes = list, numpy.ndarray, csr_matrix

		if not isinstance(X, allowed_dtypes):
			raise ValueError("X must be either a list of lists, a 2D numpy " \
				"array, or a scipy.sparse.csr_matrix.")
		if isinstance(X, numpy.ndarray) and len(X.shape) != 2:
			raise ValueError("X must have exactly two dimensions.")

		if not self.sparse:
			if X.dtype != 'float64':
				X = X.astype('float64')

		if not isinstance(self.optimizer, SieveGreedy):
			self.optimizer = OPTIMIZERS['sieve'](function=self,
				verbose=self.verbose, random_state=self.random_state,
				**self.optimizer_kwds)

		self._X = X if self._X is None else self._X
		self._initialize(X)

		if self.verbose:
			self.pbar = tqdm(total=self.n_samples, unit_scale=True)

		self.optimizer.select(X, self.n_samples, sample_cost=sample_cost)

		if self.verbose == True:
			self.pbar.close()

		self.ranking = numpy.array(self.ranking)
		self.gains = numpy.array(self.gains)
		self._X = None
		return self

	def transform(self, X, y=None, sample_weight=None):
		"""Transform a data set to include only the selected examples.

		This method will return a selection of X and optionally selections
		of y and sample_weight. The default setting is to select items based
		on the ranking determined in the `fit` step with examples in the same
		order as that ranking. Optionally, the whole data set can be returned,
		with the weights corresponding to samples that were not selected set
		to 0. This setting can be controlled by setting `pipeline=True`.

		Parameters
		----------
		X : list or numpy.ndarray, shape=(n, d)
			The data set to transform. Must be numeric.

		y : list or numpy.ndarray or None, shape=(n,), optional
			The labels to transform. If passed in this function will return
			both the data and the corresponding labels for the rows that have
			been selected. Default is None.

		sample_weight : list or numpy.ndarray or None, shape=(n,), optional
			The sample weights to transform. If passed in this function will
			return the selected labels (y) and the selected samples, even
			if no labels were passed in. Default is None.

		Returns
		-------
		X_subset : numpy.ndarray, shape=(n_samples, d)
			A subset of the data such that n_samples < n and n_samples is the
			integer provided at initialization.

		y_subset : numpy.ndarray, shape=(n_samples,), optional
			The labels that match with the indices of the samples if y is
			passed in. Only returned if passed in.

		sample_weight_subset : numpy.ndarray, shape=(n_samples,), optional
			The weight of each example.
		"""

		r = self.ranking

		if sample_weight is not None:
			if y is None:
				return X[r], None, sample_weight[r]
			else:
				return X[r], y[r], sample_weight[r]

		else:
			if y is None:
				return X[r]
			else:
				return X[r], y[r]

	def fit_transform(self, X, y=None, sample_weight=None, sample_cost=None):
		"""Run optimization and select a subset of examples.

		This method will first perform the `fit` step and then perform the
		`transform` step, returning a transformed data set.

		Parameters
		----------
		X : list or numpy.ndarray, shape=(n, d)
			The data set to transform. Must be numeric.

		y : list or numpy.ndarray or None, shape=(n,), optional
			The labels to transform. If passed in this function will return
			both the data and the corresponding labels for the rows that have
			been selected. Default is None.

		sample_weight : list or numpy.ndarray or None, shape=(n,), optional
			The sample weights to transform. If passed in this function will
			return the selected labels (y) and the selected samples, even
			if no labels were passed in. Default is None.

		sample_cost : list or numpy.ndarray or None, shape=(n,), optional
			The cost of each item. If set, indicates that optimization should
			be performed with respect to a knapsack constraint.

		Returns
		-------
		X_subset : numpy.ndarray, shape=(n_samples, d)
			A subset of the data such that n_samples < n and n_samples is the
			integer provided at initialization.

		y_subset : numpy.ndarray, shape=(n_samples,), optional
			The labels that match with the indices of the samples if y is
			passed in. Only returned if passed in.

		sample_weight_subset : numpy.ndarray, shape=(n_samples,), optional
			The weight of each example.
		"""

		return self.fit(X, y=y, sample_weight=sample_weight,
			sample_cost=sample_cost).transform(X, y=y,
			sample_weight=sample_weight)

	def _initialize(self, X, idxs=None):
		n, d = X.shape
		self._X = X if self._X is None else self._X

		self.sparse = isinstance(X, csr_matrix)
		self.ranking = []
		self.gains = []
		self.subset = numpy.zeros((0, self._X.shape[1]), dtype='float64')

		self.current_values = numpy.zeros(d, dtype='float64')
		self.current_concave_values = numpy.zeros(d, dtype='float64')
		self.mask = numpy.zeros(n, dtype='int8')

		if self.initial_subset is not None:
			if self.initial_subset.ndim == 1:
				if self.initial_subset.dtype == bool:
					self.initial_subset = numpy.where(self.initial_subset == 1)[0]

				if len(self.initial_subset) + self.n_samples > X.shape[0]:
					raise ValueError("When using a mask for the initial subset" \
						" must selected fewer than the size of the subset minus" \
						" the initial subset size, i.e., n_samples < X.shape[0] -"\
						" initial_subset.shape[0].")

				if self.initial_subset.max() > X.shape[0]:
					raise ValueError("When passing in an integer mask for the initial subset"\
						" the maximum value cannot exceed the size of the data set.")
				elif self.initial_subset.min() < 0:
					raise ValueError("When passing in an integer mask for the initial subset"\
						" the minimum value cannot be negative.")

				self.mask[self.initial_subset] = 1

		self.idxs = numpy.where(self.mask == 0)[0]

	def _calculate_gains(self, X, idxs=None):
		raise NotImplementedError

	def _calculate_sieve_gains(self, X, thresholds, idxs):
		n = X.shape[0]
		d = X.shape[1] if self.reservoir is None else self.max_reservoir_size
		l = len(thresholds)

		if self.sieve_current_values_ is None:
			self.sieve_current_values_ = numpy.zeros((l, d),
				dtype='float64')
			self.sieve_selections_ = numpy.zeros((l, self.n_samples),
				dtype='int64') - 1
			self.sieve_gains_ = numpy.zeros((l, self.n_samples),
				dtype='float64') - 1
			self.sieve_n_selected_ = numpy.zeros(l,
				dtype='int64')
			self.sieve_total_gains_ = numpy.zeros(l,
				dtype='float64')
			self.sieve_subsets_ = numpy.zeros((l, self.n_samples,
				self._X.shape[1]), dtype='float64')
		else:
			j = l - self.sieve_current_values_.shape[0]
			if j > 0:
				self.sieve_current_values_ = numpy.vstack([
					self.sieve_current_values_, numpy.zeros((j, d),
						dtype='float64')])
				self.sieve_selections_ = numpy.vstack([
					self.sieve_selections_, numpy.zeros((j, self.n_samples),
						dtype='int64') - 1])
				self.sieve_gains_ = numpy.vstack([self.sieve_gains_,
					numpy.zeros((j, self.n_samples), dtype='float64')])
				self.sieve_n_selected_ = numpy.concatenate([
					self.sieve_n_selected_, numpy.zeros(j, dtype='int64')])
				self.sieve_total_gains_ = numpy.concatenate([
					self.sieve_total_gains_, numpy.zeros(j, dtype='float64')])
				self.sieve_subsets_ = numpy.concatenate([self.sieve_subsets_,
					numpy.zeros((j, self.n_samples, self._X.shape[1]),
						dtype='float64')])

	def _select_next(self, X, gain, idx):
		self.ranking.append(idx)
		self.gains.append(gain)
		self.mask[idx] = True
		self.idxs = numpy.where(self.mask == 0)[0]

		if self.sparse:
			X = self._X[idx:idx+1].toarray()
		else:
			X = self._X[idx:idx+1]

		if self.metric != 'precomputed':
			self.subset = numpy.concatenate([self.subset, X])

class BaseGraphSelection(BaseSelection):
	"""The base graph selection object.

	This object defines the structures that all submodular selection algorithms
	should follow if they operate on a graph, such as pairwise similarity
	measurements. All algorithms will have the same public methods and the same
	attributes.

	NOTE: All ~pairwise~ values in your data must be positive for these
	selection methods to work.

	This implementation allows users to pass in either their own symmetric
	square matrix of similarity values, or a data matrix as normal and a function
	that calculates these pairwise values.

	Parameters
	----------
	n_samples : int
		The number of samples to return.

	metric : str
		The method for converting a data matrix into a square symmetric matrix
		of pairwise similarities. If a string, can be any of the metrics
		implemented in sklearn (see https://scikit-learn.org/stable/modules/
		generated/sklearn.metrics.pairwise_distances.html), including
		"precomputed" if one has already generated a similarity matrix. Note
		that sklearn calculates distance matrices whereas apricot operates on
		similarity matrices, and so a distances.max() - distances transformation
		is performed on the resulting distances. For backcompatibility,
		'corr' will be read as 'correlation'.

	initial_subset : list, numpy.ndarray or None
		If provided, this should be a list of indices into the data matrix
		to use as the initial subset, or a group of examples that may not be
		in the provided data should beused as the initial subset. If indices,
		the provided array should be one-dimensional. If a group of examples,
		the data should be 2 dimensional.

	optimizer : string or optimizers.BaseOptimizer, optional
		The optimization approach to use for the selection. Default is
		'two-stage', which makes selections using the naive greedy algorithm
		initially and then switches to the lazy greedy algorithm. Must be
		one of

			'naive' : the naive greedy algorithm
			'lazy' : the lazy (or accelerated) greedy algorithm
			'approximate-lazy' : the approximate lazy greedy algorithm
			'two-stage' : starts with naive and switches to lazy
			'stochastic' : the stochastic greedy algorithm
			'greedi' : the GreeDi distributed algorithm
			'bidirectional' : the bidirectional greedy algorithm

		Default is 'naive'.

	optimizer_kwds : dict or None
		A dictionary of arguments to pass into the optimizer object. The keys
		of this dictionary should be the names of the parameters in the optimizer
		and the values in the dictionary should be the values that these
		parameters take. Default is None.

	n_neighbors : int or None
		When constructing a similarity matrix, the number of nearest neighbors
		whose similarity values will be kept. The result is a sparse similarity
		matrix which can significantly speed up computation at the cost of
		accuracy. Default is None.

	reservoir : numpy.ndarray or None
		The reservoir to use when calculating gains in the sieve greedy
		streaming optimization algorithm in the `partial_fit` method.
		Currently only used for graph-based functions. If a numpy array
		is passed in, it will be used as the reservoir. If None is passed in,
		will use reservoir sampling to collect a reservoir. Default is None.

	max_reservoir_size : int
		The maximum size that the reservoir can take. If a reservoir is passed
		in, this value is set to the size of that array. Default is 1000.

	n_jobs : int
		The number of threads to use when performing computation in parallel.
		Currently, this parameter is exposed but does not actually do anything.
		This will be fixed soon.

	random_state : int or RandomState or None, optional
		The random seed to use for the random selection process. Only used
		for stochastic greedy.

	verbose : bool
		Whether to print output during the selection process.

	Attributes
	----------
	n_samples : int
		The number of samples to select.

	metric : callable
		A function that takes in a data matrix and converts it to a square
		symmetric matrix.

	ranking : numpy.array int
		The selected samples in the order of their gain.

	gains : numpy.array float
		The gain of each sample in the returned set when it was added to the
		growing subset. The first number corresponds to the gain of the first
		added sample, the second corresponds to the gain of the second added
		sample, and so forth.
	"""

	def __init__(self, n_samples, metric='euclidean',
		initial_subset=None, optimizer='two-stage', optimizer_kwds={},
		n_neighbors=None, reservoir=None, max_reservoir_size=1000,
		n_jobs=1, random_state=None, verbose=False):

		super(BaseGraphSelection, self).__init__(n_samples=n_samples,
			initial_subset=initial_subset, optimizer=optimizer,
			optimizer_kwds=optimizer_kwds, reservoir=reservoir,
			max_reservoir_size=max_reservoir_size, n_jobs=n_jobs,
			random_state=random_state, verbose=verbose)

		self.metric = metric.replace("corr", "correlation")
		self.n_neighbors = n_neighbors

	def fit(self, X, y=None, sample_weight=None, sample_cost=None):
		"""Run submodular optimization to select a subset of examples.

		This method is a wrapper for the full submodular optimization process.
		It takes in some data set (and optionally labels that are ignored
		during this process) and selects `n_samples` from it in the greedy
		manner specified by the optimizer.

		This method will return the selector object itself, not the transformed
		data set. The `transform` method will then transform a data set to the
		selected points, or alternatively one can use the ranking stored in
		the `self.ranking` attribute. The `fit_transform` method will perform
		both optimization and selection and return the selected items.

		Parameters
		----------
		X : list or numpy.ndarray, shape=(n, d)
			The data set to transform. Must be numeric.

		y : list or numpy.ndarray or None, shape=(n,), optional
			The labels to transform. If passed in this function will return
			both the data and th corresponding labels for the rows that have
			been selected.

		sample_weight : list or numpy.ndarray or None, shape=(n,), optional
			The weight of each example. Currently ignored in apricot but
			included to maintain compatibility with sklearn pipelines.

		sample_cost : list or numpy.ndarray or None, shape=(n,), optional
			The cost of each item. If set, indicates that optimization should
			be performed with respect to a knapsack constraint.

		Returns
		-------
		self : BaseGraphSelection
			The fit step returns this selector object.
		"""

		if isinstance(X, csr_matrix) and self.metric not in ("precomputed", "ignore"):
			raise ValueError("Must passed in a precomputed sparse " \
				"similarity matrix or a dense feature matrix.")
		if self.metric == 'precomputed' and X.shape[0] != X.shape[1]:
			raise ValueError("Precomputed similarity matrices " \
				"must be square and symmetric.")

		X_pairwise = _calculate_pairwise_distances(X, metric=self.metric,
			n_neighbors=self.n_neighbors)

		self._X = X
		return super(BaseGraphSelection, self).fit(X_pairwise, y=y,
			sample_weight=sample_weight, sample_cost=sample_cost)

	def partial_fit(self, X, y=None, sample_weight=None, sample_cost=None):
		if self.reservoir is None:
			self.reservoir = numpy.empty((self.max_reservoir_size, X.shape[1]))

		if self.update_reservoir_:
			for i in range(X.shape[0]):
				if self.reservoir_size < self.max_reservoir_size:
					self.reservoir[self.reservoir_size] = X[i]
					self.reservoir_size += 1
				else:
					r = self.random_state.choice(self.n_seen_ + i)
					if r < self.max_reservoir_size:
						self.reservoir[r] = X[i]
						#self.current_values_[:, r] = 0.

		X_pairwise = _calculate_pairwise_distances(X,
			Y=self.reservoir[:self.reservoir_size], metric=self.metric)

		self._X = X
		super(BaseGraphSelection, self).partial_fit(X_pairwise, y=y,
			sample_weight=sample_weight, sample_cost=sample_cost)

		self.current_values = numpy.zeros(self.reservoir_size,
			dtype='float64')
		self.n_seen_ += X.shape[0]

	def _initialize(self, X_pairwise, idxs=None):
		super(BaseGraphSelection, self)._initialize(X_pairwise, idxs=idxs)

	def _calculate_gains(self, X_pairwise):
		super(BaseGraphSelection, self)._calculate_gains(X_pairwise)

	def _calculate_sieve_gains(self, X, thresholds, idxs):
		super(BaseGraphSelection, self)._calculate_sieve_gains(X, thresholds,
			idxs)

	def _select_next(self, X_pairwise, gain, idx):
		super(BaseGraphSelection, self)._select_next(X_pairwise, gain, idx)


import os
import errno
import json
import numpy as np
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt
from Stele.processing.processing_hsg import helper_functions as procHSGHelp
from .pmt import PMT
from .high_sideband_pmt import HighSidebandPMT

np.set_printoptions(linewidth=500)

class HighSidebandPMTOld(PMT):
    """
    Old version: Replaced March 01, 2017

    Class initialized by loading in data set.

    Multiple copies of the same sideband were stacked as raw data and combined,
    effectively causing (2) 10-pt scans to be treated the same as (1) 20pt
    scan.  This works well until you have photon counted pulses.
    """
    def __init__(self, file_path, verbose=False):
        """
        Initializes a SPEX spectrum.  It'll open a single file, then read
        the data from that file using .add_sideband().  The super's init will
        handle the parameters and the description.

        attributes:
            self.parameters - dictionary of important experimental parameters,
                              created in PMT
            self.sb_dict - keys are sideband order, values are PMT data arrays
            self.sb_list - sorted list of included sidebands

        :param file_path: path to the current file
        :type file_path: str
        :param verbose: Flag to see the nitty gritty details
        :type verbose: bool
        :return:
        """
        # Creates the json parameters dictionary
        super(HighSidebandPMT, self).__init__(file_path)
        self.fname = file_path
        self.parameters["files included"] = [file_path]
        with open(file_path, 'r') as f:
            sb_num = int(f.readline()[1:])
        raw_temp = np.genfromtxt(file_path, comments='#', delimiter=',')[3:, :]
        self.initial_sb = sb_num
        self.initial_data = np.array(raw_temp)
        self.sb_dict = {sb_num: np.array(raw_temp)}
        self.sb_list = [sb_num]

    def add_sideband(self, other):
        """
        This bad boy will add another PMT sideband object to the sideband
        spectrum of this object.  It handles when you measure the same sideband
        twice.  It assumes both are equally "good"

        It currently doesn't do any sort of job combining dictionaries or
        anything, but it definitely could, if you have two incomplete
        dictionaries

        :param other: the new sideband data to add to the larger spectrum.
                      Add means append, no additino is performed
        :type other: HighSidebandPMT
        :return:
        """
        """
        This bad boy will add another PMT sideband object to the sideband
        spectrum of this object

        It currently doesn't do any sort of job combining dictionaries or
        anything, but it definitely could
        """
        self.parameters["files included"].append(other.fname)

        if other.initial_sb in self.sb_list:
            self.sb_list.append(other.initial_sb)

        # Make things comma delimited?
        try:
            self.sb_dict[other.initial_sb].vstack((other.initial_data))
        except Exception:
            self.sb_dict[other.initial_sb] = np.array(other.initial_data)

    def process_sidebands(self, verbose=False):
        """
        This bad boy will clean up the garbled mess that is the object before
        hand, including clearing out misfired shots and doing the averaging.

        Affects:
            self.sb_dict = Averages over sidebands

        Creates:
            self.sb_list = The sideband orders included in this object.

        :param verbose: Flag to see the nitty gritty details.
        :type verbose: bool
        :return: None
        """

        for sb_num, sb in list(self.sb_dict.items()):
            if sb_num == 0:
                # This way the FEL doesn't need to be on during laser
                # line measurement
                fire_condition = -np.inf
            else:
                # Say FEL fired if the cavity dump signal is
                # more than half the mean of the cavity dump signal
                fire_condition = np.mean(sb[:, 2]) / 2
            frequencies = sorted(list(set(sb[:, 0])))

            temp = None
            for freq in frequencies:
                data_temp = np.array([])
                for point in sb:
                    if point[0] == freq and point[2] > fire_condition:
                        data_temp = np.hstack((data_temp, point[3]))
                try:
                    temp = np.vstack((temp, np.array([
                        freq, np.mean(data_temp),
                        np.std(data_temp) / np.sqrt(len(data_temp))
                        ])))
                except Exception:
                    temp = np.array([
                        freq, np.mean(data_temp),
                        np.std(data_temp) / np.sqrt(len(data_temp))
                        ])
            # turn NIR freq into eV
            temp[:, 0] = temp[:, 0] / 8065.6
            temp = temp[temp[:, 0].argsort()]
            self.sb_dict[sb_num] = np.array(temp)
        self.sb_list = sorted(self.sb_dict.keys())
        if verbose:
            print("Sidebands included", self.sb_list)

    def integrate_sidebands(self, verbose=False):
        """
        This method will integrate the sidebands to find their strengths, and
        then use a magic number to define the width, since they are currently
        so utterly undersampled for fitting.

        It is currently the preferred method for calculating sideband
        strengths.  self.fit_sidebands is probably better with better-sampled
        lines.

        Creates:
        self.sb_results = full list of integrated data. Column order is:
                          [sb order, Freq (eV), "error" (eV), Integrate area
                          (arb.), area error, "Linewidth" (eV),
                          "Linewidth error" (eV)
        self.full_dict = Dictionary where the SB order column is removed and
                         turned into the keys.  The values are the rest of that
                         sideband's results.

        :param verbose: Flag to see the nitty gritty details
        :type verbose: bool
        :return: None
        """
        self.full_dict = {}
        for sideband in list(self.sb_dict.items()):
            index = np.argmax(sideband[1][:, 1])
            nir_frequency = sideband[1][index, 0]
            area = np.trapz(np.nan_to_num(
                sideband[1][:, 1]), sideband[1][:, 0]
                )
            # Divide by the step size?
            error = np.sqrt(np.sum(np.nan_to_num(
                sideband[1][:, 2]) ** 2)) / 8065.6
            if verbose:
                print("order", sideband[0])
                print("area", area)
                print("error", error)
                print("ratio", area / error)
            details = np.array([
                sideband[0], nir_frequency, 1 / 8065.6, area, error,
                2 / 8065.6, 1 / 8065.6
                ])
            if area < 0:
                if verbose:
                    print("area less than 0", sideband[0])
                continue
            # Two seems like a good cutoff?
            elif area < 1.5 * error:
                if verbose:
                    print("I did not keep sideband ", sideband[0])
                continue
            try:
                self.sb_results = np.vstack((self.sb_results, details))
            except Exception:
                self.sb_results = np.array(details)
            self.full_dict[sideband[0]] = details[1:]
        try:
            self.sb_results = self.sb_results[self.sb_results[:, 0].argsort()]

        except (IndexError, AttributeError):
            # IndexError where there's only one sideband
            # AttributeError when there aren't any (one sb which wasn't fit)
            pass

    def fit_sidebands(self, plot=False, verbose=False):
        """
        This method will fit a gaussian to each of the sidebands provided in
        the self.sb_dict and make a list just like in the EMCCD version.  It
        will also use the standard error of the integral of the PMT peak as the
        error of the gaussian area instead of that element from the covariance
        matrix.  Seems more legit.

        attributes:
        self.sb_results: the numpy array that contains all of the fit info just
                         like it does in the CCD class.
        self.full_dict = A dictionary version of self.sb_results

        :param plot: Flag to see the results plotted
        :type plot: bool
        :param verbose: Flag to see the nitty gritty details
        :type verbose: bool
        :return: None
        """
        sb_fits = {}
        for sideband in list(self.sb_dict.items()):
            if verbose:
                print("Sideband number", sideband[0])
                print("Sideband data:\n", sideband[1])
            index = np.argmax(sideband[1][:, 1])
            nir_frequency = sideband[1][index, 0]
            peak = sideband[1][index, 1]
            width_guess = 0.0001  # Yep, another magic number
            p0 = [nir_frequency, peak * width_guess, width_guess, 0.00001]

            if verbose:
                x_vals = np.linspace(np.amin(sideband[1][:, 0]),
                                     np.amax(sideband[1][:, 0]), num=50)
                plt.plot(x_vals, procHSGHelp.gauss(x_vals, *p0),
                         label="fit :{}".format(sideband[1]))
                print("p0:", p0)
            try:
                coeff, var_list = curve_fit(
                    gauss, sideband[1][:, 0], sideband[1][:, 1],
                    sigma=sideband[1][:, 2], p0=p0
                    )
                coeff[1] = abs(coeff[1])
                coeff[2] = abs(coeff[2])
                if verbose:
                    print("coeffs:", coeff)
                    print("stdevs:", np.sqrt(np.diag(var_list)))
                    print("integral", np.trapz(
                        sideband[1][:, 1], sideband[1][:, 0]
                        ))
                # The error on where the sideband is should be small
                if np.sqrt(np.diag(var_list))[0] / coeff[0] < 0.5:
                    sb_fits[sideband[0]] = np.concatenate((
                        np.array([sideband[0]]), coeff,
                        np.sqrt(np.diag(var_list))
                        ))
                    # print "error then:", sb_fits[sideband[0]][6]
                    relative_error = (
                        np.sqrt(sum([
                            x ** 2 for x in sideband[1][index - 1:index + 2, 2]
                            ]))
                        / np.sum(sideband[1][index - 1:index + 2, 1]))
                    if verbose:
                        print("relative error:", relative_error)
                    sb_fits[sideband[0]][6] = coeff[1] * relative_error
                    # print "error now:", sb_fits[sideband[0]][6]
                    if plot:
                        x_vals = np.linspace(np.amin(sideband[1][:, 0]),
                                             np.amax(sideband[1][:, 0]), num=50
                                             )
                        plt.plot(x_vals, procHSGHelp.gauss(x_vals, *coeff))
                        # plt.plot(x_vals, gauss(x_vals, *p0))
                else:
                    print("what happened?")
            except Exception:
                print("God damn it, Leroy.\nYou couldn't fit this.")
                sb_fits[sideband[0]] = None

        for result in sorted(sb_fits.keys()):
            try:
                self.sb_results = np.vstack((self.sb_results, sb_fits[result]))
            except Exception:
                self.sb_results = np.array(sb_fits[result])

        self.sb_results = self.sb_results[:, [0, 1, 5, 2, 6, 3, 7, 4, 8]]
        self.sb_results = self.sb_results[:, :7]
        if verbose:
            print("And the results, please:\n", self.sb_results)

        self.full_dict = {}
        for sb in self.sb_results:
            self.full_dict[sb[0]] = np.asarray(sb[1:])

    def laser_line(self, verbose=False):
        """
        This method is designed to scale everything in the PMT to the
        conversion efficiency based on our measurement of the laser line with a
        fixed attenuation.

        Creates:
            self.parameters['normalized?'] = Flag to specify if the laser has
            been accounted for.

        :return: None
        """

        if 0 not in self.sb_list:
            self.parameters['normalized?'] = False
            return
        else:
            laser_index = np.where(self.sb_results[:, 0] == 0)[0][0]
            if verbose:
                print("sb_results", self.sb_results[laser_index, :])
                print("laser_index", laser_index)

            laser_strength = np.array(self.sb_results[laser_index, 3:5])

            if verbose:
                print("Laser_strength", laser_strength)

            for sb in self.sb_results:
                sb[4] = (sb[3] / laser_strength[0]) * np.sqrt(
                    (sb[4] / sb[3]) ** 2 + (laser_strength[1] / laser_strength[0]) ** 2)
                sb[3] = sb[3] / laser_strength[0]
            for sb in list(self.full_dict.values()):
                sb[3] = (sb[2] / laser_strength[0]) * np.sqrt(
                    (sb[3] / sb[2]) ** 2 + (laser_strength[1] / laser_strength[0]) ** 2)
                sb[2] = sb[2] / laser_strength[0]
            self.parameters['normalized?'] = True

    def save_processing(self, file_name, folder_str, marker='', index=''):
        """
        This will save all of the self.proc_data and the results from the
        fitting of this individual file.

        Format:
        spectra_fname = file_name + '_' + marker + '_' + str(index) + '.txt'
        fit_fname = file_name + '_' + marker + '_' + str(index) + '_fits.txt'

        Inputs:
        file_name = the beginning of the file name to be saved
        folder_str = the location of the folder where the file will be saved,
                     will create the folder, if necessary.
        marker = I...I don't know what this was originally for
        index = used to keep these files from overwriting themselves when in a
                list

        Outputs:
        Two files:
            self.proc_data = the continuous spectrum
            self.sb_results = the individual sideband details

        :param file_name: The base name for the saved file
        :type file_name: str
        :param folder_str: The full name for the folder hte file is saved it.
                           Folder can be created
        :type folder_str: str
        :param marker: Marker for the file, appended to file_name, often the
                       self.parameters['series']
        :type marker: str
        :param index: used to keep these files from overwriting themselves when
                      marker is the same
        :type index: str or int
        :return: None
        """
        try:
            os.mkdir(folder_str)
        except OSError as e:
            if e.errno == errno.EEXIST:
                pass
            else:
                raise

        spectra_fname = file_name + '_' + marker + '_' + str(index) + '.txt'
        fit_fname = file_name + '_' + marker + '_' + str(index) + '_fits.txt'
        self.save_name = spectra_fname
        # self.parameters["files included"] = list(self.files)
        try:
            parameter_str = json.dumps(
                self.parameters, sort_keys=True, indent=4,
                separators=(',', ': ')
                )
        except Exception:
            print("Source: PMT.save_images\nJSON FAILED")
            print("Here is the dictionary that broke JSON:\n", self.parameters)
            return
        parameter_str = parameter_str.replace('\n', '\n#')

        # Make the number of lines constant so importing is easier
        # for num in range(99 - num_lines): parameter_str += '\n#'
        num_lines = parameter_str.count('#')
        parameter_str += '\n#' * (99 - num_lines)

        origin_import_spec = '\nNIR frequency,Signal,Standard error\neV,arb. u.,arb. u.\n,{:.3f},'.format(
            self.parameters["fieldStrength"]["mean"])
        spec_header = '#' + parameter_str + origin_import_spec

        origin_import_fits = '\nCenter energy,error,Amplitude,error,Linewidth,error\neV,,arb. u.,,eV,,\n,,'  # + marker
        fits_header = '#' + parameter_str + origin_import_fits

        for sideband in sorted(self.sb_dict.keys()):
            # TODO: REPLACE ALL INSTANCES OF ERRORS AS CONTROL STATEMENTS WITH
            #       CORRECT IF/ELSE OR SIMILAR LOGIC
            try:
                complete = np.vstack((complete, self.sb_dict[sideband]))
            except Exception:
                complete = np.array(self.sb_dict[sideband])

        np.savetxt(
            os.path.join(folder_str, spectra_fname), complete, delimiter=',',
            header=spec_header, comments='', fmt='%0.6e')

        try:
            np.savetxt(os.path.join(folder_str, fit_fname), self.sb_results,
                       delimiter=',',
                       header=fits_header, comments='', fmt='%0.6e')
        except AttributeError:
            # Catch the error that happens if you save something without files
            print("warning, couldn't save fit file (no sidebands found?)")

        print("Saved PMT spectrum.\nDirectory: {}".format(
            os.path.join(folder_str, spectra_fname)))


#!/usr/bin/env python
import os
import numpy as np
import time
import copy
import sys

import argparse

ang_2_bohr = 1.0/0.52917721067
hart_2_ev = 27.21138602

import atomistic_tools.cp2k_grid_orbitals as cgo
import atomistic_tools.cp2k_stm_sts as css
from atomistic_tools import common
from atomistic_tools.cube import Cube

from mpi4py import MPI

comm = MPI.COMM_WORLD
mpi_rank = comm.Get_rank()
mpi_size = comm.Get_size()

parser = argparse.ArgumentParser(
    description='Puts the CP2K orbitals on grid and calculates STM.')

### ----------------------------------------------------------------------
### Input and output files
parser.add_argument(
    '--cp2k_input_file',
    metavar='FILENAME',
    required=True,
    help='CP2K input of the SCF calculation.')
parser.add_argument(
    '--basis_set_file',
    metavar='FILENAME',
    required=True,
    help='File containing the used basis sets.')
parser.add_argument(
    '--xyz_file',
    metavar='FILENAME',
    required=True,
    help='.xyz file containing the geometry.')
parser.add_argument(
    '--wfn_file',
    metavar='FILENAME',
    required=True,
    help='Restart file containing the final wavefunction.')
parser.add_argument(
    '--hartree_file',
    metavar='FILENAME',
    required=True,
    help='Cube file containing the hartree potential.')
parser.add_argument(
    '--output_file',
    metavar='FILENAME',
    default="./stm.npz",
    help='File, where to save the STM/STS output')
parser.add_argument(
    '--orb_output_file',
    metavar='FILENAME',
    default="./orb.npz",
    help='File, where to save the orbital output')
### ----------------------------------------------------------------------
### Parameters for putting orbitals on grid
parser.add_argument(
    '--eval_region',
    type=str,
    nargs=6,
    metavar='X',
    required=True,
    help=common.eval_region_description
)
parser.add_argument(
    '--dx',
    type=float,
    metavar='DX',
    required=True,
    help='Spatial step for the grid (angstroms).')
parser.add_argument(
    '--eval_cutoff',
    type=float,
    metavar='D',
    default=16.0,
    help=("Size of the region around the atom where each"
          " orbital is evaluated (only used for 'G' region).")
)
parser.add_argument(
    '--extrap_extent',
    type=float,
    metavar='H',
    default=4.0,
    required=True,
    help="The extent of the extrapolation region. (angstrom)")
### ----------------------------------------------------------------------
### Gas phase analysis parameters
parser.add_argument(
    '--n_homo',
    type=int,
    metavar='N',
    default=0,
    help="Number of HOMO orbitals to analyse.")
parser.add_argument(
    '--n_lumo',
    type=int,
    metavar='N',
    default=0,
    help="Number of LUMO orbitals to analyse.")
parser.add_argument(
    '--orb_heights',
    nargs='*',
    type=float,
    metavar='H',
    help="List of heights for constant height orbital pictures (wrt topmost atom).")
parser.add_argument(
    '--orb_isovalues',
    nargs='*',
    type=float,
    metavar='C',
    help="List of charge density isovalues for constant current orbital pictures")
parser.add_argument(
    '--orb_fwhms',
    nargs='*',
    type=float,
    default=[0.02],
    help="Full width at half maximum for orbital STS gaussian broadening. (eV)")
### ----------------------------------------------------------------------
### Slab system analysis parameters
###
### Option 1: continuous selection
parser.add_argument(
    '--energy_range',
    nargs=3,
    type=float,
    metavar='E',
    help='Selection of STM/STS energy values based on a range: min, max and differential.')
###
### Option 2: discrete selection
parser.add_argument(
    '--energies',
    nargs='*',
    type=float,
    metavar='E',
    help='Discrete energies where to run the STM/STS.')
### ----------------------------------------------------------------------
### Parameters for STM/STS series
parser.add_argument(
    '--heights',
    nargs='*',
    type=float,
    metavar='H',
    help="List of heights for constant height STM pictures (wrt topmost atom).")
parser.add_argument(
    '--isovalues',
    nargs='*',
    type=float,
    metavar='C',
    help="List of charge density isovalues for constant current STM pictures.")
parser.add_argument(
    '--fwhms',
    nargs='*',
    type=float,
    default=[0.1],
    help="Full width at half maximum for STS gaussian broadening. (eV)")

time0 = time.time()

### ------------------------------------------------------
### Parse args for only one rank to suppress duplicate stdio
### ------------------------------------------------------

args = None
args_success = False
try:
    if mpi_rank == 0:
        args = parser.parse_args()
        args_success = True
finally:
    args_success = comm.bcast(args_success, root=0)

if not args_success:
    print(mpi_rank, "exiting")
    exit(0)

args = comm.bcast(args, root=0)

### ------------------------------------------------------
### Energy values for STM/STS
### ------------------------------------------------------

if args.energies is not None:
    e_arr = np.array(args.energies)
elif args.energy_range is not None:
    emin, emax, de = args.energy_range
    e_arr = np.arange(emin, emax+de/2, de)
else:
    e_arr = None

max_fwhm = np.max(args.fwhms)
if e_arr is not None:
    sel_emin = np.min(e_arr) - 2.0*max_fwhm
    sel_emax = np.max(e_arr) + 2.0*max_fwhm
else:
    sel_emin = None
    sel_emax = None

### ------------------------------------------------------
### Evaluate orbitals on the real-space grid
### ------------------------------------------------------

cp2k_grid_orb = cgo.Cp2kGridOrbitals(mpi_rank, mpi_size, mpi_comm=comm, single_precision=True)
cp2k_grid_orb.read_cp2k_input(args.cp2k_input_file)
cp2k_grid_orb.read_xyz(args.xyz_file)
cp2k_grid_orb.center_atoms_to_cell()
cp2k_grid_orb.read_basis_functions(args.basis_set_file)
cp2k_grid_orb.load_restart_wfn_file(args.wfn_file,
                                    emin=sel_emin, emax=sel_emax,
                                    n_occ=args.n_homo, n_virt=args.n_lumo
)

print("R%d/%d: loaded wfn, %.2fs"%(mpi_rank, mpi_size, (time.time() - time0)))
sys.stdout.flush()
time1 = time.time()

eval_reg = common.parse_eval_region_input(args.eval_region, cp2k_grid_orb.ase_atoms, cp2k_grid_orb.cell)

cp2k_grid_orb.calc_morbs_in_region(args.dx,
                                x_eval_region = eval_reg[0],
                                y_eval_region = eval_reg[1],
                                z_eval_region = eval_reg[2],
                                reserve_extrap = args.extrap_extent,
                                eval_cutoff = args.eval_cutoff)

print("R%d/%d: evaluated wfn, %.2fs"%(mpi_rank, mpi_size, (time.time() - time1)))
sys.stdout.flush()
time1 = time.time()

### ------------------------------------------------------
### Extrapolate orbitals
### ------------------------------------------------------

hart_cube = Cube()
hart_cube.read_cube_file(args.hartree_file)
extrap_plane_z = eval_reg[2][1] / ang_2_bohr - np.max(cp2k_grid_orb.ase_atoms.positions[:, 2])
hart_plane = hart_cube.get_plane_above_topmost_atom(extrap_plane_z) - cp2k_grid_orb.ref_energy/hart_2_ev

cp2k_grid_orb.extrapolate_morbs(hart_plane=hart_plane)

print("R%d/%d: extrapolated wfn, %.2fs"%(mpi_rank, mpi_size, (time.time() - time1)))
sys.stdout.flush()
time1 = time.time()

### ------------------------------------------------------
### Set up STM object
### ------------------------------------------------------

stm = css.STM(mpi_comm = comm, cp2k_grid_orb = cp2k_grid_orb)
stm.gather_global_energies()
stm.divide_by_space()

### ------------------------------------------------------
### Run STM-STS analysis for orbitals
### ------------------------------------------------------

orb_heights = args.orb_heights if args.orb_heights is not None else []
orb_isovalues = args.orb_isovalues if args.orb_isovalues is not None else []
orb_fwhms = args.orb_fwhms if args.orb_fwhms is not None else []

if len(orb_fwhms) != 0 and (len(orb_heights) != 0 or len(orb_isovalues) != 0):

    orbital_list = list(range(-args.n_homo + 1, args.n_lumo + 1))

    stm.create_orbital_images(orbital_list, orb_heights, orb_isovalues)

    orbital_list_wrt_0 = list(np.array(orbital_list) + stm.cgo.i_homo_glob[0])
    orbital_energies = stm.global_morb_energies[0][orbital_list_wrt_0]

    stm.calculate_stm_maps(orb_fwhms, orb_isovalues, orb_heights, orbital_energies)

    stm.collect_and_save_orb_maps(path=args.orb_output_file)

### ------------------------------------------------------
### Run STM-STS analysis for general energies
### ------------------------------------------------------

heights = args.heights if args.heights is not None else []
isovalues = args.isovalues if args.isovalues is not None else []
fwhms = args.fwhms if args.fwhms is not None else []

if e_arr is not None and len(fwhms) != 0 and (len(heights) != 0 or len(isovalues) != 0):

    stm.calculate_stm_maps(fwhms, isovalues, heights, e_arr)

    stm.collect_and_save_stm_maps(path=args.output_file)

print("R%d/%d: finished, total time: %.2fs"%(mpi_rank, mpi_size, (time.time() - time0)))


"""Plots Figures 1, EDF Figures 1 and 7 in Gillett et al."""
import logging
import os
from pprint import pformat
import numpy as np
import csv
import matplotlib
from scipy import stats
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import esmvaltool.diag_scripts.attribute.ncblendmask_esmval as ncbm
import random
import math
import sys

from esmvaltool.diag_scripts.shared import (group_metadata, run_diagnostic,
                                            select_metadata, sorted_metadata)
from esmvaltool.diag_scripts.shared._base import (
    ProvenanceLogger, get_diagnostic_filename, get_plot_filename)
from esmvaltool.diag_scripts.shared.plot import quickplot

logger = logging.getLogger(os.path.basename(__file__))
logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))

def main(cfg):
    matplotlib.use('Agg')
    plt.ioff() #Turn off interactive plotting.
    # Get a description of the preprocessed data that we will use as input, as well as other flags.
    input_data = cfg['input_data'].values()
    auxiliary_data_dir=cfg['auxiliary_data_dir']
    plot_dir=cfg['plot_dir']
    output_file_type=cfg['output_file_type']
    obs=cfg['obs']
    if obs=='had5':
        obs_file=auxiliary_data_dir+'/HadCRUT.5.0.1.0.anomalies.ensemble_mean.nc'
        hadlabel='HadCRUT5'
        ensobs=auxiliary_data_dir+'/HadCRUT.5.0.1.0.anomalies.' #If set apply multi-model analysis using ensemble obs data.
    elif obs=='had4':
        obs_file=auxiliary_data_dir+'/HadCRUT.4.6.0.0.median.nc'  #Updated to end of 2019.
        hadlabel='HadCRUT4'
        ensobs=auxiliary_data_dir+'/HadCRUT.4.6.0.0.anomalies.' #If set apply multi-model analysis using ensemble obs data.
    else:
        exit('Obs not recognised')
    sftlf_file=auxiliary_data_dir+'/CNRM-CM6-1-5x5-sftlf.nc' #Hard-coded path to sftlf file for CNRM-CM6 on a 5x5 grid.

    grouped_input_data = group_metadata(
        input_data, 'dataset', sort='ensemble')
    logger.info(
        "Group input data by model and sort by ensemble:"
        "\n%s", pformat(grouped_input_data))
    type (grouped_input_data)
    nmodel=len(grouped_input_data)

#Initialise variables.
    experiments=['historical-ssp245','hist-GHG','hist-aer','hist-nat','hist-volc','hist-sol','hist-stratO3','hist-CO2','hist-GHG-ssp245-GHG','hist-aer-ssp245-aer','hist-nat-ssp245-nat']
    labels=['Anthropogenic and natural forcings','Greenhouse gases','Aerosols','Natural forcings']
    nexp=len(experiments)-3 #Subtract three to account for repetition of hist-GHG, hist-nat, hist-aer.
    diag_name='gmst02' #Annual mean GMST.
    ldiag=85 #length of diagnostic,hard-coded.
    years=list(range(1850,2020,2)) #Used for plotting.
    anom_max=500 #arbitrary max size for number of anomalies.
    mean_diag=np.zeros((ldiag,nexp,nmodel))
    mean_gmst_comp_warming=np.zeros((ldiag,nexp,nmodel))
    mean_ann_warming=np.zeros((ldiag,nexp,nmodel))
    mm_ann_warming=np.zeros((ldiag,nexp))
    range_ann_warming=np.zeros((ldiag,nexp,2)) # 5-95% range.
    nensmax=50
    msval=1e20
    all_ann_warming=np.full((ldiag,nexp,nmodel,nensmax),msval)
    all_ann_warming_comp=np.full((ldiag,nexp,nmodel,nensmax),msval)
    all_ann_warming_gsat=np.full((ldiag,nexp,nmodel,nensmax),msval)
    ens_sizes=np.zeros((nexp,nmodel))
    model_names=[]
    ensobs_diag=[]
    ensobs_dec_warming=[]

#Set up figure including colours.
    font = {'size'   : 5}
    matplotlib.rc('font', **font)
    mm_conv=0.0394
    mod_cols=np.array([[0,73,73],[255,255,109],[0,146,146],[255,109,182],[255,182,119],[146,0,0],[0,109,219],[182,109,255],[109,182,255],[182,219,255],[73,0,146],[146,73,0],[219,209,0],[36,255,36]])/256.
    cols=np.array([[0,0,0],[196,121,0],[178,178,178],[0,52,102],[0,79,0],[200,0,0],[0,200,0],[0,0,200],[112,160,205]])/256.
    shade_cols=np.array([[128,128,128,128],[204,174,113,128],[191,191,191,128],[67,147,195,128],[223,237,195,128],[255,150,150,128],[150,255,150,128],[150,150,255,128],[91,174,178,128]])/256.

    plt.figure(figsize=[88*mm_conv,113*mm_conv])
    plt.subplot(211)

#Loop over models, then experiments, then ensemble members.
    for mm, dataset in enumerate(grouped_input_data):
        logger.info("*************** Processing model %s", dataset)
        model_names.append(dataset)
        lbl=dataset
        grouped_model_input_data = group_metadata(
            grouped_input_data[dataset], 'exp', sort='ensemble')
        for exp in grouped_model_input_data:
            logger.info("***** Processing experiment %s", exp)
            exp_string = [experiments.index(i) for i in experiments if exp == i]
            experiment = exp_string[0]
            #Label hist-nat-ssp245-nat as hist-nat, hist-ghg-ssp245-ghg as hist-ghg etc
            #(some models' hist-nat ends in 2014 so is merged with ssp245-nat).
            if experiment > 7: experiment=experiment-7
            print ('*** Experiment',exp,'Index:',experiment)
            grouped_exp_input_data = group_metadata(
              grouped_model_input_data[exp], 'ensemble', sort='variable_group')
            nens=len(grouped_exp_input_data)
            ens_sizes[experiment,mm]=nens
            exp_diags=np.zeros((ldiag,nens))
            exp_ann_warming=np.zeros((ldiag,nens))
            exp_gmst_comp_warming=np.zeros((ldiag,nens))

            for ee, ensemble in enumerate(grouped_exp_input_data):
                logger.info("** Processing ensemble %s", ensemble)
                for attributes in grouped_exp_input_data[ensemble]:
                    logger.info("Processing variable %s", attributes['variable_group'])
                    file=attributes['filename']
                logger.info("*************** Files for blend and mask %s", file)
                #Calculate masked and blended GMST for individual simulation.
                (exp_diags[:,ee],obs_diag, dec_warming, obs_dec_warming, ann_warming, gmst_comp_warming)=ncbm.ncblendmask_esmval(
                    file,obs_file,diag_name,ensobs,ensobs_diag,ensobs_dec_warming, cfg=cfg)
                ensobs='' #Set to empty string so that ensemble obs diagnostics are only calculated on the first iteration.
                #Take anomalies relative to 1850-1900.
                exp_diags[:,ee]=exp_diags[:,ee]-np.mean(exp_diags[0:int(np.round((1901-1850)/2)),ee])
                obs_diag=obs_diag-np.mean(obs_diag[0:int(np.round((1901-1850)/2))])
                exp_ann_warming[:,ee]=ann_warming
                exp_gmst_comp_warming[:,ee]=gmst_comp_warming
                #Plot first ensemble member of historical.
                if exp=="historical-ssp245" and ee==0:
                    alpha_ens=1. if ee==0 else 0.2
                    ls='dashed' if mm > 7 else 'solid'
                    plt.plot(years,exp_diags[:,ee],color=mod_cols[mm],linewidth=0.5,label=lbl,zorder=1-ee,alpha=alpha_ens,linestyle=ls)
                    lbl=""
            mean_diag[:,experiment,mm]=np.mean(exp_diags,axis=1)
            mean_ann_warming[:,experiment,mm]=np.mean(exp_ann_warming,axis=1)
            mean_gmst_comp_warming[:,experiment,mm]=np.mean(exp_gmst_comp_warming,axis=1)
            all_ann_warming[:,experiment,mm,0:nens]=exp_diags #Use GMST.
            all_ann_warming_comp[:,experiment,mm,0:nens]=exp_gmst_comp_warming #Use spatially complete GMST.
            all_ann_warming_gsat[:,experiment,mm,0:nens]=exp_ann_warming #Use GSAT.

    #Write GMST and GSAT timeseries to CSV file for other applications (not needed for Gillett et al. plots).
    with open(plot_dir+'/cmip6_gmst.csv', mode='w') as file:
        data_writer=csv.writer(file,delimiter=',',quotechar='"', quoting=csv.QUOTE_MINIMAL)
        data_writer.writerow(['CMIP6 DAMIP models HadCRUT4-masked blended GMST (Cowtan et al., 2015) and globally-complete GSAT'])
        for experiment in range(nexp):
            data_writer.writerow(['Experiment:',experiments[experiment]])
            for mm, dataset in enumerate(grouped_input_data):
                data_writer.writerow([dataset])
                for ee in range(int(ens_sizes[experiment,mm])):
                    data_writer.writerow(['Ensemble member',ee])
                    data_writer.writerow(['Year, GMST_complete, GMST_HadCRUT4_masked, GSAT'])
                    for yy in range(ldiag):
                        data_writer.writerow([years[yy],all_ann_warming[yy,experiment,mm,ee],all_ann_warming_comp[yy,experiment,mm,ee],all_ann_warming_gsat[yy,experiment,mm,ee]])

#Calculate ratio of GSAT to GMST warming for individual simulations.
    denom=np.mean(all_ann_warming[int((2010-1850)/2):int((2020-1850)/2),0,:,:],axis=0)
    ratio_by_model=np.mean(all_ann_warming_gsat[int((2010-1850)/2):int((2020-1850)/2),0,:,:],axis=0)/denom
    copy_ratio_by_model=np.reshape(ratio_by_model[:,:],nmodel*nensmax)
#Equivalent calculation for spatially-complete GMST.
    denom=np.mean(all_ann_warming_comp[int((2010-1850)/2):int((2020-1850)/2),0,:,:],axis=0)
    ratio_by_model_comp=np.mean(all_ann_warming_gsat[int((2010-1850)/2):int((2020-1850)/2),0,:,:],axis=0)/denom
    copy_ratio_by_model_comp=np.reshape(ratio_by_model_comp[:,:],nmodel*nensmax)

    plt.plot(years,obs_diag,color='black',linewidth=1,label=hadlabel)
    plt.plot(years,np.mean(mean_diag[:,0,:],axis=1),color=cols[1],linewidth=1,label='Model mean GMST')
    plt.plot(years,np.mean(mean_ann_warming[:,0,:],axis=1),color='red',linewidth=1,label='Model mean GSAT')
    plt.plot([1850,2020],[0,0],color='silver',linewidth=0.5,ls='--',zorder=0)
    plt.axis([1850,2020,-2,3])
    plt.xlabel('Year')
    plt.ylabel('Global mean temperature anomaly ($^\circ$C)')
    plt.legend(loc=2, ncol=2, fontsize='x-small', fancybox=False, frameon=False, columnspacing=0.5,handlelength=1)
    for experiment in range(nexp):
        wts=np.zeros((nmodel,nensmax))
        for mm in range(nmodel):
            wts[mm,0:int(ens_sizes[experiment,mm])]=1./ens_sizes[experiment,mm]
        wts=np.reshape(wts,nmodel*nensmax)/np.sum(wts)
        if experiment==0:
#Calculate ratio of GSAT to obs-masked GMST warming 5-95% range.
           sort_ratio=np.sort(copy_ratio_by_model)
           sort_index=np.argsort(copy_ratio_by_model)
           cdf=np.cumsum(wts[sort_index])
           range_ratio=[sort_ratio[cdf>=0.05][0],sort_ratio[cdf>=0.95][0]]
           print ('5-95% range of ratio for GSAT/obs-masked GMST',range_ratio)
#Equivalent calculation for spatially-complete GMST
           sort_ratio=np.sort(copy_ratio_by_model_comp)
           sort_index=np.argsort(copy_ratio_by_model_comp)
           cdf=np.cumsum(wts[sort_index])
           range_ratio=[sort_ratio[cdf>=0.05][0],sort_ratio[cdf>=0.95][0]]
           print ('5-95% range of ratio for GSAT/spatially-complete GMST',range_ratio)
        for yy in range(ldiag):
                year_warming=np.reshape(all_ann_warming[yy,experiment,:,:],nmodel*nensmax)
                sort_warming=np.sort(year_warming)
                sort_index=np.argsort(year_warming)
                cdf=np.cumsum(wts[sort_index])
                range_ann_warming[yy,experiment,:]=[sort_warming[cdf>=0.05][0],sort_warming[cdf>=0.95][0]]
                mm_ann_warming[yy,experiment]=np.sum(year_warming*wts)
    plt.text (1825,2.25,'a',fontsize =7,fontweight='bold', va='center', ha='center')
    plt.subplot(212)
    zzs=[3,1,0,2]
    for experiment in range(4):
        offset=0
        plt.fill_between(years,range_ann_warming[:,experiment,0]+offset,range_ann_warming[:,experiment,1]+offset,color=shade_cols[experiment+1,:],zorder=zzs[experiment])
        plt.plot(years,mm_ann_warming[:,experiment]+offset,color=cols[experiment+1,:],linewidth=1,label=labels[experiment],zorder=zzs[experiment]+4)

    plt.plot(years,obs_diag,color='black',linewidth=1,label=hadlabel,zorder=8)
    plt.axis([1850,2020,-2,3])
    plt.plot([1850,2020],[0,0],color='black',linewidth=0.5,ls='--',zorder=0)
    plt.xlabel('Year')
    plt.ylabel('Global mean surface temperature anomaly ($^\circ$C)')
    plt.legend(loc=2, ncol=2, fontsize='x-small', fancybox=False, frameon=False, columnspacing=0.5,handlelength=1)
    plt.text (1825,2.25,'b',fontsize =7,fontweight='bold', va='center', ha='center')
    plt.savefig(plot_dir+'/Fig1_'+obs+'.'+output_file_type)
    plt.close()

#Plot Extended Data Fig 1 showing all DAMIP GMST timeseries.
    fig=plt.figure(figsize=[88*mm_conv,176*mm_conv])
    ax1=fig.add_subplot(111)
    for experiment in range(nexp):
        offset=experiment*-1.5
        plt.fill_between(years,range_ann_warming[:,experiment,0]+offset,range_ann_warming[:,experiment,1]+offset,color=shade_cols[experiment+1,:])
        plt.plot([1850,2025],[offset,offset],color='black',linewidth=0.5)
        plt.plot(years,mm_ann_warming[:,experiment]+offset,color=cols[experiment+1,:],linewidth=0.5,label=experiments[experiment])
        plt.text(1860,offset+0.4,experiments[experiment])
    plt.plot(years,obs_diag,color='black',linewidth=0.5,label=hadlabel,zorder=8)
    ax1.set_xlim(1850,2020)
    ax1.set_xlabel('Year')
    ax1.set_ylim(-11,2)
    plt.yticks(np.arange(27)*0.5-11,['','','-1.0','-0.5','0.0','0.5','1.0','','-1.0','-0.5','0.0','0.5','1.0','','-1.0','-0.5','0.0','0.5','1.0','','-1.0','-0.5','0.0','0.5','1.0','',''])
    ax1.set_ylabel('Global mean surface temperature change ($^\circ$C)')
    ax2=ax1.twinx()
    ax2.set_ylim(-11,2)
    plt.yticks(np.arange(27)*0.5-11,['-0.5',' 0.0',' 0.5',' 1.0','','-1.0','-0.5',' 0.0',' 0.5',' 1.0','','-1.0','-0.5',' 0.0',' 0.5',' 1.0','','-1.0','-0.5',' 0.0',' 0.5',' 1.0','','','','',''])
    plt.savefig(plot_dir+'/supplement_timeseries'+obs+'.'+output_file_type)
    plt.close()

#Calculate uncertainty in GMST and GSAT warming in obs (as reported in Gillett et al.).
    ensobs_dec_warming=np.sort(ensobs_dec_warming)
    enssize=len(ensobs_dec_warming)
    print ('Obs GMST warming: mean, 5, 95%',np.mean(ensobs_dec_warming),ensobs_dec_warming[math.floor((enssize-1)*0.05)],ensobs_dec_warming[math.ceil((enssize-1)*0.95)])
    nmc=10000
    obs_gsat=np.zeros((nmc))
    for mc_counter in range(nmc):
      ens=random.randint(0,enssize-1)
      mm=random.randint(0,nmodel-1)
      ee=random.randint(0,ens_sizes[0,mm]-1)
      obs_gsat[mc_counter]=ensobs_dec_warming[ens]*np.mean(all_ann_warming_gsat[int((2010-1850)/2):int((2020-1850)/2),0,mm,ee],axis=0)/np.mean(all_ann_warming[int((2010-1850)/2):int((2020-1850)/2),0,mm,ee],axis=0)
    obs_gsat=np.sort(obs_gsat)
    print ('Obs GSAT warming: mean, 5, 95%',np.mean(obs_gsat),obs_gsat[math.floor((nmc-1)*0.05)],obs_gsat[math.ceil((nmc-1)*0.95)])

if __name__ == '__main__':

    with run_diagnostic() as config:
        main(config)