###########################################
# This is a script for exctracting materials from the abse3.db
# database. After exctracting materials information as ase Atoms
# objects, their sine matrix is constructed, and the eigenvalues
# are calculated. Finally, material names, sine matrix eigenvalues
# and the PBE heat of formation are written row-wise as a csv file.
#
#  Prasenjit Sen, July 9, 2025.
#
############################################

# Import necessary things from ASE
# Import Atoms, read, connect from appropriate modules

#Import usual python stuff
# import numpy and csv

from numpy import linalg as LA

# Import what is required from DScribe
# Import sine matrix from DScribe


db = connect('abse3.db')

elements =[]
rows = []
# Following is the header line for the data file.
rows.append(['Material'] + [f'smeig{i+1}' for i in range(30)] + ['PBE_hof'])

# Use toatoms attribute of each line in db to store structure
# as ASE atoms object. Find the elements constituting the materials
# (as their Z number) in the database, and store them in the list elements. 
# This is just to know what elements appear in the data set. You may use this
# to choose only certain elements as the central atom.

for row in db.select():
    #save structure as atoms object
    Z_numbers = atoms.numbers
 
    # store the Z_numbers in elements list

species = sorted(elements)
print('Species=',species)

try:
    for row in db.select():
        # save structure as atoms object
        Z_numbers = atoms.numbers

        #name+protocol create a unique label for each material.
        name = row.name + '-' + row.prototype

        if (len(Z_numbers) <= 30):

            samples = [atoms]
            current = samples[0]
# The previous lines extract the information about the atoms object
# in 'current' from the list 'samples'.

            #sm_desc = [write how you want to set the parameters to create 
            # the sine matrix.] 
            sine_matrix = sm_desc.create(current)

            SM_2d = sine_matrix.reshape #Complete the statement to reshape
            # the sine matrix as a 2D array.

            # Here, eigenvalues are created
            eigenvalues = LA.eigvalsh(SM_2d)

            # You may save the sine matrix elements, or the eigenvalues.
            rows.append([name] + [f"{x: .5f}" for x in eigenvalues] + [f"{row.pbe_hof: .5f}"])


            with open('hform_dataset_ABSe3.csv', 'w', newline='') as f:
                writer = csv.writer(f)
                writer.writerows(rows)

        else:
            print('# atoms= ', len(Z_numbers), '>30')
            print('Compound=', name)

except Exception as e:
        #pass
        print(e)