#!/usr/bin/env python

import yaml
import sys
import numpy
import math
import click
import os

from numerov import solve_numerov, constants, __version__


def read_data(input_file):
    with open(input_file, 'r') as stream:
        try:
            input_data = yaml.load(stream)
            return input_data
        except yaml.YAMLError as e:
            sys.stderr.write(e)
            sys.stderr.write("\n")
            sys.exit(-1)


@click.command()
@click.option('--general', help='File containing general input.')
@click.option('--reduced-mass', help='File containing reduced mass data.')
@click.option('--force-field', help='File containing force-field data.')
@click.option('--exp-values', help='File containing expectation values along displacement.')
def main(general, reduced_mass, force_field, exp_values):

    data = read_data(reduced_mass)
    reduced_mass_amu = data['reduced_mass_amu']

    data = read_data(force_field)
    print(yaml.dump(data))
    x = numpy.array(data['displacements_au'])
    y = numpy.array(data['energies_hartree'])
    # shift potential such that minimum is at zero
    y -= min(y)
    pot_energy_coefs = numpy.polyfit(x, y, data['degree_taylor'])
    fourth_lowest_energy = sorted(y)[3]
    h_x = []
    h_y = []
    for (displacement, pot_energy) in zip(x, y):
        if pot_energy < fourth_lowest_energy:
            h_x.append(displacement)
            h_y.append(pot_energy)

    data = read_data(exp_values)
    print(yaml.dump(data))
    x = numpy.array(data['displacements_au'])
    y = numpy.array(data['exp_values_hartree'])
    exp_value_coefs = numpy.polyfit(x, y, data['degree_taylor'])

    general_data = read_data(general)
    print('energy_precision_hartree: {0}'.format(general_data['energy_precision_hartree']))
    print('num_solutions: {0}'.format(general_data['num_solutions']))
    print('num_steps: {0}'.format(general_data['num_steps']))

    transition_frequency_previous = sys.float_info.max
    displacement_range = (-0.5, 0.5)
    while True:
        qs, psi_squared, energies_hartree, averaged_exp_values_au = solve_numerov(pot_energy_coefs,
                                                                                  exp_value_coefs,
                                                                                  displacement_range,
                                                                                  general_data['num_steps'],
                                                                                  general_data['num_solutions'],
                                                                                  general_data['energy_precision_hartree'],
                                                                                  reduced_mass_amu * constants['amu_to_au'])
        transition_frequency = (energies_hartree[general_data['num_solutions'] - 1] - energies_hartree[0]) * constants['hartree_to_cm1']
        if abs(transition_frequency - transition_frequency_previous) < 1.0e-1:
            qs = numpy.around(qs, decimals=3).tolist()
            psi_squared = numpy.around(psi_squared, decimals=5).tolist()
            break
        displacement_range = (displacement_range[0] - 0.1, displacement_range[1] + 0.1)
        transition_frequency_previous = transition_frequency

    # get harmonic frequency from numerov
    # this can be used as a sanity check to verify whether this matches
    # the provided input frequency
    pot_energy_coefs_harmonic = numpy.polyfit(h_x, h_y, 2)
    _, _, energies_hartree_harmonic, _ = solve_numerov(pot_energy_coefs_harmonic,
                                                       exp_value_coefs,
                                                       displacement_range,
                                                       general_data['num_steps'],
                                                       general_data['num_solutions'],
                                                       general_data['energy_precision_hartree'],
                                                       reduced_mass_amu * constants['amu_to_au'])
    for i in range(len(exp_value_coefs)):
        exp_value_coefs[-i - 1] *= math.factorial(i)
    for i in range(len(pot_energy_coefs)):
        pot_energy_coefs[-i - 1] *= math.factorial(i)

    print('version: {0}'.format(__version__))
    print('displacement_range:')
    print('  min: {0}'.format(displacement_range[0]))
    print('  max: {0}'.format(displacement_range[1]))
    print('exp_value_coefs: {0}'.format(list(exp_value_coefs)))
    print('pot_energy_coefs: {0}'.format(list(pot_energy_coefs)))
    print('pot_energy_coefs_harmonic: {0}'.format(list(pot_energy_coefs_harmonic)))
    print('energies_hartree: {0}'.format(list(energies_hartree)))
    print('energies_hartree_harmonic: {0}'.format(list(energies_hartree_harmonic)))
    print('averaged_exp_values_au: {0}'.format(list(averaged_exp_values_au)))
    print('qs: {0}'.format(qs))
    print('psi_squared: {0}'.format(psi_squared))


if __name__ == '__main__':
    main()
