from ase.io import read
from kaldo.conductivity import Conductivity
from kaldo.forceconstants import ForceConstants
from kaldo.phonons import Phonons
import numpy as np

### Set up force constant objects via interface to LAMMPS ####

# Replicate the unit cell 'nrep'=1 time
nrep = 1
supercell = np.array([nrep, nrep, nrep])

# Load in computed 2nd, 3rd IFCs from LAMMPS outputs
forceconstants = ForceConstants.from_folder(folder='fc_aSi512',
                                            supercell=supercell,
                                            format='lammps')

# Configure phonon object
# 'is_classic': specify if the system is classic, True for classical and False for quantum
# 'temperature: temperature (Kelvin) at which simulation is performed
# 'folder': name of folder containing phonon property and thermal conductivity calculations
# 'storage': Format to storage phonon properties ('formatted' for ASCII format data, 'numpy'
#            for python numpy array and 'memory' for quick calculations, no data stored)

phonons_config = {
    'is_classic': False,
    'temperature': 300,  #'temperature'=300K
    'folder': 'ALD_aSi512',
    'third_bandwidth': 0.5 / 4.135,  # 0.5 eV is used here.
    'broadening_shape': 'triangle',
### Set up the coordinates of the system and the force constant calculations ####

# Define the system according to ASE style. 'a': lattice parameter (Angstrom)
atoms = bulk('C', 'diamond', a=3.566)

# Replicate the unit cell 'nrep'=3 times
nrep = 3
supercell = np.array([nrep, nrep, nrep])

# Configure force constant calculator
forceconstants_config = {
    'atoms': atoms,
    'supercell': supercell,
    'folder': 'fc_c_diamond'
}
forceconstants = ForceConstants(**forceconstants_config)

# Define input information for the ase LAMMPSlib calculator
# Terosff potential for carbon (C.tersoff) is used for this example.
lammps_inputs = {
    'lmpcmds':
    ['pair_style tersoff', 'pair_coeff * * forcefields/C.tersoff C'],
    'keep_alive': True,
    'log_file': 'lammps-c-diamond.log'
}

# Compute 2nd and 3rd IFCs with the defined calculators
forceconstants.second.calculate(LAMMPSlib(**lammps_inputs))
forceconstants.third.calculate(LAMMPSlib(**lammps_inputs))

### Set up the phonon object and the anharmonic properties calculations ####
# Example: silicon bulk, hiPhive
# Computes: anharmonic properties and thermal conductivity for silicon bulk (2 atoms per cell)
# Uses: hiPhive

from kaldo.conductivity import Conductivity
from kaldo.forceconstants import ForceConstants
from kaldo.phonons import Phonons
import numpy as np

# Config force constants object by loading in the IFCs
# from hiphive calculations
forceconstants = ForceConstants.from_folder('hiPhive_si_bulk',
                                            supercell=[3, 3, 3],
                                            format='hiphive')

### Set up the phonon object and the anharmonic properties calculations ####

# Configure phonon object
# 'k_points': number of k-points
# 'is_classic': specify if the system is classic, True for classical and False for quantum
# 'temperature: temperature (Kelvin) at which simulation is performed
# 'folder': name of folder containing phonon property and thermal conductivity calculations
# 'storage': Format to storage phonon properties ('formatted' for ASCII format data, 'numpy'
#            for python numpy array and 'memory' for quick calculations, no data stored)

# Define the k-point mesh using 'kpts' parameter
k_points = 7  #'k_points'=7 k points in each direction
phonons_config = {
    'kpts': [k_points, k_points, k_points],
    'is_classic': False,
    'temperature': 300,  #'temperature'=300K