import matplotlib.pyplot as plt

# Increase the number of MD steps to converge the histograms.
STEPS = 5000

plt.xlabel("Kinetic energy K / Hartree")
plt.ylabel("Probability P(K)")

# Set up configuration
bulk_configuration = BulkConfiguration(
    bravais_lattice=FaceCenteredCubic(5.4306*Angstrom),
    elements=[Silicon, Silicon],
    fractional_coordinates=[[0.0, 0.0, 0.0], [0.25, 0.25, 0.25]]
    )

# Set calculator
calculator = TremoloXCalculator(parameters=Tersoff_Si_1988b())
bulk_configuration.setCalculator(calculator)

# Run MD simulations with each thermostat.
for thermostat_class in NVTBerendsen, NVTBussiDonadioParrinello:
    # Set up MD method
    method = thermostat_class(
        time_step=1*femtoSecond,
        reservoir_temperature=300*Kelvin,
        thermostat_timescale=100*femtoSecond,
        heating_rate=0*Kelvin/picoSecond,
        initial_velocity=ZeroVelocities(),
    )

    # Run long MD simulation
    md_trajectory = MolecularDynamics(
        bulk_configuration,
        constraints=[],
        trajectory_filename='trajectory.hdf5',
        steps=STEPS,
        log_interval=50,
        method=method
    )

    # Choose color and label for thermostat
    color = next(plt.gca()._get_lines.prop_cycler)["color"]
    label = thermostat_class.__name__

    # Kinetic energy distribution from samples in MD trajectory.
    temperatures = numpy.unique(md_trajectory.reservoirTemperatures())
    assert len(temperatures) == 1, 'Reservoir temperature should not change during MD run.'
    temperature = temperatures[0]
    beta = (1.0/(boltzmann_constant * temperature)).inUnitsOf(Hartree**(-1))
    ekin = md_trajectory.kineticEnergies().inUnitsOf(Hartree)

    # Plot histogram of kinetic energies
    plt.hist(
        ekin,
        bins='auto', alpha=0.5,
        color=color, stacked=True, density=True, label=f'Thermostatted MD, {label}'
    )

# Find number of degrees of freedom
n = bulk_configuration.numberOfAtoms()
Nf = 3*n

# To avoid overflows reorganize the expression for P(K),
# so the logs of large quantities cancel before exponentiating.
def logGamma(x):
    """ Expansion of log(Gamma(x)) for large x """
    lG = (x-0.5) * numpy.log(x) - x + 0.5*numpy.log(2.0*numpy.pi)
    return lG

# For K=0.0, log(beta^*K) gives inf.
ekin = ekin[ekin > 0.0]
ekin.sort()

probability_canonical = numpy.exp(
    Nf/2.0 * numpy.log(beta*ekin) - beta*ekin - logGamma(Nf/2.0)
) / ekin

# Plot canonical distribution
plt.plot(
    ekin, probability_canonical,
    color=color, lw=2, label=f'Canonical distribution, T={temperature.inUnitsOf(Kelvin)} K'
)

plt.legend()
plt.show()