#! /usr/bin/env python

import math

from atsim.potentials import EAMPotential, Potential
from atsim.potentials.eam_tabulation import TABEAM_EAMTabulation


def makeFunc(a, b, r_e, c):
    # Creates functions of the form used for density function.
    # Functional form also forms components of pair potential.
    def func(r):
        return (a * math.exp(-b*(r/r_e - 1)))/(1+(r/r_e - c)**20.0)
    return func


def makePairPotAA(A, gamma, r_e, kappa,
                  # Function factory that returns functions parameterised for homogeneous pair interactions
                  B, omega, lamda):
    f1 = makeFunc(A, gamma, r_e, kappa)
    f2 = makeFunc(B, omega, r_e, lamda)

    def func(r):
        return f1(r) - f2(r)
    return func


def makePairPotAB(dens_a, phi_aa, dens_b, phi_bb):
    # Function factory that returns functions parameterised for heterogeneous pair interactions
    def func(r):
        return 0.5 * ((dens_b(r)/dens_a(r) * phi_aa(r)) + (dens_a(r)/dens_b(r) * phi_bb(r)))
    return func


def makeEmbed(rho_e, rho_s, F_ni, F_i, F_e, eta):
    # Function factory returning parameterised embedding function.
    rho_n = 0.85*rho_e
    rho_0 = 1.15*rho_e

    def e1(rho):
        return sum([F_ni[i] * (rho/rho_n - 1)**float(i) for i in range(4)])

    def e2(rho):
        return sum([F_i[i] * (rho/rho_e - 1)**float(i) for i in range(4)])

    def e3(rho):
        return F_e * (1.0 - eta*math.log(rho/rho_s)) * (rho/rho_s)**eta

    def func(rho):
        if rho < rho_n:
            return e1(rho)
        elif rho_n <= rho < rho_0:
            return e2(rho)
        return e3(rho)
    return func


def makePotentialObjects():
    # Potential parameters
    r_eCu = 2.556162
    f_eCu = 1.554485
    gamma_Cu = 8.127620
    omega_Cu = 4.334731
    A_Cu = 0.396620
    B_Cu = 0.548085
    kappa_Cu = 0.308782
    lambda_Cu = 0.756515

    rho_e_Cu = 21.175871
    rho_s_Cu = 21.175395
    F_ni_Cu = [-2.170269, -0.263788, 1.088878, -0.817603]
    F_i_Cu = [-2.19, 0.0, 0.561830, -2.100595]
    eta_Cu = 0.310490
    F_e_Cu = -2.186568

    r_eAl = 2.863924
    f_eAl = 1.403115
    gamma_Al = 6.613165
    omega_Al = 3.527021
    # A_Al      = 0.134873
    A_Al = 0.314873
    B_Al = 0.365551
    kappa_Al = 0.379846
    lambda_Al = 0.759692

    rho_e_Al = 20.418205
    rho_s_Al = 23.195740
    F_ni_Al = [-2.807602, -0.301435, 1.258562, -1.247604]
    F_i_Al = [-2.83, 0.0, 0.622245, -2.488244]
    eta_Al = 0.785902
    F_e_Al = -2.824528

    # Define the density functions
    dens_Cu = makeFunc(f_eCu, omega_Cu, r_eCu, lambda_Cu)
    dens_Al = makeFunc(f_eAl, omega_Al, r_eAl,  lambda_Al)

    # Finally, define embedding functions for each species
    embed_Cu = makeEmbed(rho_e_Cu, rho_s_Cu, F_ni_Cu, F_i_Cu, F_e_Cu, eta_Cu)
    embed_Al = makeEmbed(rho_e_Al, rho_s_Al, F_ni_Al, F_i_Al, F_e_Al, eta_Al)

    # Wrap them in EAMPotential objects
    eam_potentials = [
        EAMPotential("Al", 13, 26.98, embed_Al, dens_Al),
        EAMPotential("Cu", 29, 63.55, embed_Cu, dens_Cu)]

    # Define pair functions
    pair_CuCu = makePairPotAA(A_Cu, gamma_Cu, r_eCu, kappa_Cu,
                              B_Cu, omega_Cu, lambda_Cu)

    pair_AlAl = makePairPotAA(A_Al, gamma_Al, r_eAl, kappa_Al,
                              B_Al, omega_Al, lambda_Al)

    pair_AlCu = makePairPotAB(dens_Cu, pair_CuCu, dens_Al, pair_AlAl)

    # Wrap them in Potential objects
    pair_potentials = [
        Potential('Al', 'Al', pair_AlAl),
        Potential('Cu', 'Cu', pair_CuCu),
        Potential('Al', 'Cu', pair_AlCu)]

    return eam_potentials, pair_potentials


def main():
    eamPotentials, pairPotentials = makePotentialObjects()

    # Perform tabulation
    # Make tabulation
    cutoff_rho = 100.0
    nrho = 2000

    cutoff = 6
    nr = 2000

    tabulation = TABEAM_EAMTabulation(
        pairPotentials, eamPotentials, cutoff, nr, cutoff_rho, nrho)

    with open("TABEAM", 'w') as outfile:
        tabulation.write(outfile)


if __name__ == '__main__':
    main()
