###########################################
# This is a script for exctracting materials from the abse3.db
# database. 
#
#  Prasenjit Sen, July 9, 2025.
#
############################################

from ase import Atoms
from ase.io import read
from ase.db import connect

import csv
import numpy as np


np.set_printoptions(threshold=np.inf, formatter={'float': lambda x: "{0:9.5f}".format(x)})

db = connect('abse3.db')

atom_max = 30 #Maximum number of atoms in a unit cell in the database.
elements =[]
rows = []
rows.append(['Material'] + ['PBE_hof'])


for row in db.select():
    atoms = row.toatoms()
    Z_numbers = atoms.numbers
 
    for x in Z_numbers:
            if x not in elements:
                elements.append(x)

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

try:
    for row in db.select():
        atoms = row.toatoms()
        Z_numbers = atoms.numbers

        name = row.name + '-' + row.prototype

        if (len(Z_numbers) <= atom_max):

            samples = [atoms]
            #print('Samples=',samples)
            current = samples[0]
            #print('Current=', current)
# The previous lines extract the information about the atoms object
# in 'current' from the list 'samples'. This will later be used to construct
# features.


            rows.append([name] + [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), '>atom_max')
            print('Compound=', name)

except Exception as e:
        #pass
        print(e)