import pylab

filename = 'Fe-BCC.hdf5'

# Load the HeisenbergExchange object
heisenberg_exchange = nlread(filename, HeisenbergExchange)[0]

exchange_stiffness = heisenberg_exchange.exchangeStiffness()

# Calculate the extrapolated exchange stiffness using default values.
extrapolated_exchange_stiffness, exchange_stiffness_vec, damping_factor_range, p_fit \
    = heisenberg_exchange.extrapolatedExchangeStiffness()


# Calculate the extrapolated exchange stiffness using a 4th order polynomial fit.
extrapolated_exchange_stiffness_4, exchange_stiffness_vec_4, damping_factor_range_4, p_fit_4 \
    = heisenberg_exchange.extrapolatedExchangeStiffness(
        polynomial_order=4)

# Calculate the extrapolated exchange stiffness using a custom damping factor range,
extrapolated_exchange_stiffness_custom, exchange_stiffness_vec_custom, damping_factor_range_custom, p_fit_custom \
    = heisenberg_exchange.extrapolatedExchangeStiffness(
        damping_factor_range=numpy.linspace(0.1, 0.4, 10)*Ang**-1,
        )

# Print the results
nlprint('-------------------------------------------------------')
nlprint('Exchange stiffness:')
nlprint('-------------------------------------------------------')
nlprint('A_ex                               : {: .1f}'.format(exchange_stiffness))
nlprint('A_ex (extrapolated, default)       : {: .1f}'.format(extrapolated_exchange_stiffness))
nlprint('A_ex (extrapolated, 4th order fit  : {: .1f}'.format(extrapolated_exchange_stiffness_4))
nlprint('A_ex (extrapolated, custom range)  : {: .1f}'.format(extrapolated_exchange_stiffness_custom))
nlprint('-------------------------------------------------------')

# Plot the data and fits
damping_factor_range_fit = numpy.linspace(0, damping_factor_range[-1].inUnitsOf(Ang**-1), 100)

pylab.figure()
pylab.plot(0, exchange_stiffness.inUnitsOf(meV*Ang**2), 'm*', label='No extrapolation')
pylab.plot(0, extrapolated_exchange_stiffness.inUnitsOf(meV*Ang**2), 'rs', label='Extrapolated result, default')
pylab.plot(0, extrapolated_exchange_stiffness_4.inUnitsOf(meV*Ang**2), 'gs', label='Extrapolated result, 4th order fit')
pylab.plot(0, extrapolated_exchange_stiffness_custom.inUnitsOf(meV*Ang**2), 'bs', label='Extrapolated result, custom range')
pylab.plot(damping_factor_range.inUnitsOf(Ang**-1), exchange_stiffness_vec.inUnitsOf(meV*Ang**2), 'ko', label='Data')
pylab.plot(damping_factor_range_fit, numpy.polyval(p_fit, damping_factor_range_fit), 'k-', label='Fit, default')
pylab.plot(damping_factor_range_fit, numpy.polyval(p_fit_4, damping_factor_range_fit), 'g--', label='Fit, 4th order fit')
pylab.plot(damping_factor_range_fit, numpy.polyval(p_fit_custom, damping_factor_range_fit), 'b--', label='Fit, custom range')
pylab.xlabel('Damping factor  ($1/\\AA$)')
pylab.ylabel('Exchange stiffness (mev $\\AA^2$)')
pylab.legend(loc=0)
pylab.show()