# Set up lattice
lattice = FaceCenteredCubic(5.4306*Angstrom)

# Define elements
elements = [Silicon, Silicon]

# Define coordinates
fractional_coordinates = [[0.,   0.,   0.  ],
                          [0.25, 0.25, 0.25]]

# Set up configuration
unit_configuration = BulkConfiguration(
    bravais_lattice=lattice,
    elements=elements,
    fractional_coordinates=fractional_coordinates
    )
reps = 3
bulk_configuration = unit_configuration.repeat(reps, reps, reps)
bulk_configuration = bulk_configuration.copyAndDeleteAtoms(5)

# Set up calculator
potentialSet = Tersoff_Si_2005()
calculator = TremoloXCalculator(parameters=potentialSet)
calculator.setVerletListsDelta(1.0*Angstrom)

bulk_configuration.setCalculator(calculator)
bulk_configuration.update()

# Set up density matrix
dynamical_matrix = DynamicalMatrix(
    bulk_configuration,
    filename='Silicon_defect.hdf5',
    object_id='dynamical_matrix',
    repetitions=Automatic,
    atomic_displacement=0.01*Angstrom,
    acoustic_sum_rule=True,
    finite_difference_method=Central,
    force_tolerance=1e-08*Hartree/Bohr**2,
    processes_per_displacement=None,
    log_filename_prefix='forces_displacement_',
    use_wigner_seitz_scheme=False,
    )
dynamical_matrix.update()

# Compute the vibrational inverse participation ratio
vibrational_ipr = VibrationalInverseParticipationRatio(
    configuration=bulk_configuration,
    dynamical_matrix=dynamical_matrix
    )

print('\nNumber of atoms N:', bulk_configuration.numberOfAtoms())
print(f'Expected IPR for delocalized states 1/N: {1 / bulk_configuration.numberOfAtoms():.3f}')

# Obtain results
energies = vibrational_ipr.energies().inUnitsOf(meV)
ipr = vibrational_ipr.inverseParticipationRatio()
idx = numpy.argmax(ipr)
print(f'Mode at energy {energies.ravel()[idx]:.2f} meV',
      f'has the highest IPR: {ipr.ravel()[idx]:.3f}')
