#!python
import numpy

from ase.io.vasp import read_vasp, write_vasp
import quantities as pq

def main():
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('-n','--samples', type=int, default=100,
            help='number of samples to average together')
    parser.add_argument('-q','--qho', action='store_true', default=False,
            help='treat as collection of quantum harmonic oscillators')
    parser.add_argument('-T','--temperature', type=float, default=298.0)
    parser.add_argument('POSCAR')
    parser.add_argument('EIGENVECTORS')
    parser.add_argument('FORCECONSTANTS')
    parser.add_argument('EFFECTIVEMASSES')

    args = parser.parse_args()

    atoms = read_vasp(args.POSCAR)
    ews = numpy.loadtxt(args.FORCECONSTANTS, dtype=float)*pq.eV*pq.angstrom**-2
    evs = numpy.loadtxt(args.EIGENVECTORS, dtype=float)
    masses = numpy.loadtxt(args.EFFECTIVEMASSES, dtype=float)*pq.amu
    temperature = args.temperature * pq.K
    samples = args.samples
    qho = args.qho

    ensemble = harmonic_sampler(atoms, ews, evs, masses, temperature, samples, qho)

    f = open('XDATCAR', 'w')
    for i,snapshot in enumerate(ensemble):
        if i == 0:
            write_vasp(f, snapshot, direct=True, sort=False, vasp5=True)
        else:
            f.write(' \n')
            numpy.savetxt(f, snapshot.get_scaled_positions(), fmt=' %19.16f')
            f.flush()
    f.close()

def harmonic_sampler(atoms, ews, evs, masses, temperature, iters, qho=False, rng=None):
    if rng is None:
        rng = numpy.random.RandomState()

    atoms = atoms.copy()
    r0 = atoms.get_positions()

    hbar = pq.constants.hbar
    kB = pq.constants.Boltzmann_constant

    for iter in range(iters):
        r = r0.copy()
        for j in range(len(ews)):
            # Ignore negative and small modes.
            if ews[j] < 1e-2: continue

            beta = 1.0/(kB*temperature)
            mass = masses[j]
            # Vibrational frequency.
            omega = numpy.sqrt(ews[j]/mass)

            # Standard deviation of a classical harmonic oscillator.
            stddev = numpy.sqrt(1/(ews[j]*beta))

            # The standard deviation cannot drop below the ground state
            # quantum harmonic oscillator.
            if qho:
                min_stddev = numpy.sqrt(hbar/(2*mass*omega))
                stddev = max(min_stddev, stddev)

            # Convert to angstrom.
            stddev = stddev.rescale(pq.angstrom).magnitude

            # Displace by a Gaussian random number along the mode.
            g = rng.normal(loc=0.0, scale=stddev)
            v = g*evs[:,j]
            r += v.reshape(len(atoms), 3)

        atoms.set_positions(r)
        yield atoms.copy()


if __name__ == '__main__':
    main()
