###########################################
# This is a script for constructing MBTR feature
# from a cif file.
#
#  Prasenjit Sen, September 23, 2025.
#
############################################

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

import csv
import numpy as np

import sys

from dscribe.descriptors import MBTR

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

atoms = read("EntryWithCollCode29929.cif") 

size =10

elements = []
Z_numbers = atoms.numbers

if (len(Z_numbers) <= size):
    for x in Z_numbers:
        if x not in elements:
            elements.append(x)

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

rows = []
rows.append(['Material'] + [f'MBTR{i+1}' for i in range(size)] + ['PBE_hof'])

        #print('System=',name)

if (len(Z_numbers) <= size):

    samples = [atoms]
            #print('Samples=',samples)
    current = samples[0]


# The following line is for the k=1 part of MBTR.
    #mbtr_desc1 = MBTR(species=species, geometry={"function":"atomic_number"}, grid={"min":8.0, "max":21.0, "sigma":0.1, "n":50}, \
#weighting={"function": "unity"}, periodic=True, normalization="l2")

# This line for k=2 part of MBTR
    mbtr_desc2 = MBTR(species=species, geometry={"function":"distance"}, grid={"min":2.0, "max":6.0, "sigma":0.2, "n":100}, \
weighting={"function": "exp", "scale": 0.5, "threshold": 1e-3}, periodic=True, normalization="l2")

#The following lines for the k=3 part of MBTR.

# Now create the k=1 MBTR
    #mbtr1 = mbtr_desc1.create(current)

# Then k=2 MBTR
    mbtr2 = mbtr_desc2.create(current)

# And k=3 MBTR
    #mbtr3 = mbtr_desc3.create(current)

    print(len(mbtr2))
    print(mbtr2)
    sys.exit()

            #rows.append([name] + [f"{x: .5f}" for x in eigenvalues] + [f"{row.pbe_hof: .5f}"])
            #with open('ABSe3_hformMBTR_dataset.csv', 'w', newline='') as f:
            #    writer = csv.writer(f)
            #    writer.writerows(rows)

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