import os
import h5py

#import numpy as np
from ase.io import read
from hiphive import ClusterSpace, StructureContainer, ForceConstantPotential
from hiphive.cutoffs import estimate_maximum_cutoff
from hiphive.cutoffs import Cutoffs
from hiphive.utilities import extract_parameters
from hiphive import ForceConstants
from hiphive import enforce_rotational_sum_rules
from hiphive.calculators import ForceConstantCalculator
from trainstation import Optimizer

prim = read('POSCAR-unitcell')
supercell = read('SPOSCAR')

# Define a cluster space using the largest cutoff you can
#max_cutoff = estimate_maximum_cutoff(supercell) - 0.1
#cutoffs = [max_cutoff]  # only second order needed
cutoffs = [4.0,5.0]
cs = ClusterSpace(prim, cutoffs)

# import the phonopy force constants using the correct supercell also provided by phonopy
fcs = ForceConstants.read_phono3py(supercell, 'fc3.hdf5')
# Find the parameters that best fits the force constants given you cluster space
parameters = extract_parameters(fcs, cs)

fcp = ForceConstantPotential(cs, parameters)

# Enforce the rotational sum rules
parameters_rot = enforce_rotational_sum_rules(cs, parameters, ['Huang', 'Born-Huang'])

# use the new parameters to make a fcp and then create the force constants and write to a phonopy file
fcp = ForceConstantPotential(cs, parameters_rot)
fcs = fcp.get_force_constants(supercell)
fcs.write_to_phono3py('force_constants.hdf5', format='hdf5')