""" Saves a txt file with a subset of data from an ase database.
    When analysing very large ase db files, it may be useful to save selected
    subsets of data in txt files.

    Dependencies
    ------------
    ASE 3.9.0 or newer.

    <fname>.db : ASE SQLite3 file.
        A database file that contains the fields 'wftop', 'energy'.

    Input parameters
    ----------------
    U_SHE : float.
        Make a subset of states with the work relevant function for this
        potential versus the standard hydrogen electrode (SHE).
"""
from sys import argv
import numpy as np
import ase.db
from scipy import constants

# Define parameters.
U_SHE = float(argv[1])
phi_SHE = 4.44
# Define limits for the work function.
DeltaU = 0.1
U_SHEmin = U_SHE - DeltaU
U_SHEmax = U_SHE + DeltaU
phi_min = U_SHEmin + phi_SHE
phi_max = U_SHEmax + phi_SHE

# Connect to the ASE database
fname = 'sample_3k.db'
print('connecting to '+fname)
c = ase.db.connect(fname)
# Select everything.
s = c.select()

# Loop over all atomic structures in the db.
data = []
for d in s:
    # Store work function in a variable.
    phi = float(d.wftop)
    # Check if this atomic state has a work function within the interval.
    if phi <= phi_min or phi >= phi_max:
        # If not, continue the loop ignoring the rest if this iteration.
        continue
    # Get more data about each atomic state.
    dbid = int(d.id)
    numbers = d.numbers
    n = len([i for i in numbers if i == 1]) - 64
    Ekin = float(d.Ekin)
    Epot = float(d.energy)
    Etot = Ekin + Epot
    wfbot = float(d.wfbot)
    # Append the data in a list of lists.
    data.append([Epot, Ekin, phi, n, wfbot, dbid])
# Put the data in a numpy matrix.
a = np.array(data)
# You can save this matrix for later use...
# Dump the matrix in a txt file.
name = 'dat_USHE_'+str(int(round(1000*U_SHE)))+'mV'
np.savetxt(name+'.txt', np.hstack([np.vstack(a[:, 0]), np.vstack(a[:, 1]),
                                   np.vstack(a[:, 2]), np.vstack(a[:, 3]),
                                   np.vstack(a[:, 4]), np.vstack(a[:, 5])]))
print(np.shape(a), 'dataset saved.')

# ... or you can do various analyses on the matrix.
# E.g. Find the most stable (potential energy) state at a given pH, U_RHE.
pH = 7
kB = constants.k / constants.e
T = 300
E_H2 = -7.540484
U_RHE = U_SHE + 2.303 * kB * T * pH
DeltaE = a[:, 0] - a[:, 3] * (E_H2/2 - U_RHE)
j = np.argmin(DeltaE)
print('Most stable atomic structure at pH = ' + str(pH) + ', U = ' +
      str(U_RHE) + 'V vs. RHE.', 'database id = ' + str(int(a[j, 5])),
                                 'n = ' + str(int(a[j, 3])),
                                 'Epot = ' + str(a[j, 0]),
                                 'Ekin = ' + str(a[j, 1]))