"""generate_input.py

    Provide simple standalone code for generating input file for
    h2mixer.

    Mark A. Caprio
    University of Notre Dame

    + 11/7/16 (mac): Created.
    + 11/20/16 (mac): Use paths from environment.
    + 12/28/16 (mac): Fix reference state list.

"""

import math
import os

# physical constants
k_mN_csqr = 938.92  # (m_N c^2)~938.92 MeV
k_hbar_c = 197.327  # (hbar c)~197.327 Mev fm

# environment
data_dir_h2 = os.environ.get("NCCI_DATA_DIR_H2")
print("NCCI_DATA_DIR_H2",data_dir_h2)
install_dir = os.environ.get("SHELL_INSTALL_DIR")
print("SHELL_INSTALL_DIR",install_dir)

# h2mixer parameter template
params = {
    # stored radial
    "Nmax_orb" : 2,  # for target space
    "Nmax_orb_int" : 20,  # for overlaps from interaction tbmes
    "orbitals_filename" : "orbitals.dat",  # for target space
    "orbitals_int_filename" : "orbitals-int.dat",  # for overlaps from interaction tbmes
    "radial_me_filename_pattern" : "radial-me-{}{}.dat",  # "{}{}" will be replaced by {"r1","r2","k1","k2"}
    "radial_olap_filename" : "radial-olap.dat",
    # stored input TBMEs
    "VNN_filename" : os.path.join(data_dir_h2,"run0164-JISP16-ob-9/JISP16-ob-9-20.bin"),
    "VC_filename" : os.path.join(data_dir_h2,"run0164-JISP16-ob-9/VC-ob-9-20.bin"),
    # scaling/transformation/basis parameters
    "hw" : 20,
    "target_truncation" : ("tb",2),
    "hw_int" : 20,
    "hw_c" : 20,
    "xform_enabled" : False,
    "xform_truncation" : ("ob",9),
    # Hamiltonian parameters
    "A" : 4,   # atomic mass number of nucleus (for kinematic operators)
    "a" : 20,  # coefficient on Ncm in Lawson term
    "hw_cm" : None,  # hw for Ncm in Lawson term; if None, defaults to "hw"
    "use_coulomb" : True,
    # output file format
    "h2_format" : 0  # h2 formats: 0, 15099
}

# MFdn parameter template
mfdn_params = {
    "ndiag" : 1,
    "nuclide" : (2,2),
    "Nshell" : params["Nmax_orb"]+1,
    "Nmin" : 0,
    "Nmax" : 2,
    "Nstep" : 2,
    "Mj" : 0,
    "eigenvectors" : 1,
    "lanczos" : 50,
    "initial_vector" : -2,
    "tolerance" : 1e-6,
    "hw_for_trans" : 20,
    "obs_basename_list" : ["tbme-rrel2","tbme-Ncm"],
    "obdme_multipolarity" : 2,
    "obdme_reference_state_list" : [(0,0,1)]
}



def weight_max_string(truncation):
    """ Convert (rank,cutoff) to "wp wn wpp wnn wpn" string.

    >>> weight_max_string(("ob",4))
        "4 4 8 8 8"
    >>> weight_max_string(("tb",4))
        "4 4 4 4 4"
    """

    (code, N) = truncation
    if (code == "ob"):
        cutoffs = (N,2*N)
    elif (code == "tb"):
        cutoffs = (N,N)

    return "{0[0]} {0[0]} {0[1]} {0[1]} {0[1]}".format(cutoffs)

def oscillator_length(hw):
    """ Calculate oscillator length for given oscillator frequency.

    b(hw) = (hbar c)/[(m_N c^2) (hbar omega)]^(1/2)

    Arguments:
        hw (numeric): hbar omega in MeV

    Returns:
        (float): b in fm
    """

    return k_hbar_c/math.sqrt(k_mN_csqr*hw)


def run_script(params):
    """ Generate shell run script.

    Arguments:
        params (dict): parameter dictionary

    Returns:
        (list of str): input file lines
    """

    lines = []

    # orbital generation
    lines.append("{install_dir}/bin/orbital-gen --oscillator {Nmax_orb} {orbitals_filename}".format(install_dir=install_dir,**params))
    lines.append("{install_dir}/bin/orbital-gen --oscillator {Nmax_orb_int} {orbitals_int_filename}".format(install_dir=install_dir,**params))

    # radial operator generation
    for operator_type in ["r","k"]:
        for power in [1,2]:
            radial_me_filename = params["radial_me_filename_pattern"].format(operator_type,power)
            lines.append("{install_dir}/bin/radial-gen --kinematic {} {} oscillator orbitals.dat {}".format(operator_type,power,radial_me_filename,install_dir=install_dir))

    # radial overlap generation
    if (params["xform_enabled"]):
        hw = params["hw"]
        hw_int = params["hw_int"]
        b_ratio = math.sqrt(hw_int/hw)
        lines.append("{install_dir}/bin/radial-gen --overlaps {b_ratio} oscillator {orbitals_int_filename} {radial_olap_filename}".format(install_dir=install_dir,b_ratio=b_ratio,**params))

    # invoke h2mixer
    lines.append("{install_dir}/bin/h2mixer < h2mixer.in".format(install_dir=install_dir))

    # ensure terminal line
    lines.append("")

    return "\n".join(lines)

def h2mixer_input(params):
    """ Generate input for h2mixer.

    Arguments:
        params (dict): parameter dictionary

    Returns:
        (list of str): input file lines
    """

    lines = []

    # initial comment
    lines.append("# h2mixer input generated by h2_input.py")
    lines.append("# given parameters:")
    lines.append("#   {}".format(params))
    lines.append("")

    # global mode definitions
    target_weight_max = weight_max_string(params["target_truncation"])
    lines.append("set-target-indexing {orbitals_filename} {target_weight_max}".format(target_weight_max=target_weight_max,**params))
    ## lines.append("set-target-indexing-oscillator {target_truncation[0]} {target_truncation[1]}".format(**params))
    lines.append("set-target-multipolarity 0 0 0")
    lines.append("set-output-format {h2_format}".format(**params))
    lines.append("set-mass {A}".format(**params))
    lines.append("")

    # radial operator inputs
    for operator_type in ["r","k"]:
        for power in [1,2]:
            radial_me_filename = params["radial_me_filename_pattern"].format(operator_type,power)
            lines.append("define-radial-operator {} {} {}".format(operator_type,power,radial_me_filename))
    lines.append("")

    # sources: identity and kinematic operators
    lines.append("define-source operator identity")
    lines.append("define-source operator Ursqr")
    lines.append("define-source operator Vr1r2")
    lines.append("define-source operator Uksqr")
    lines.append("define-source operator Vk1k2")
    lines.append("")

    # sources: VNN
    if (params["xform_enabled"]):
        xform_weight_max = weight_max_string(params["xform_truncation"])
        lines.append("define-source xform VNN {VNN_filename} {xform_weight_max} {radial_olap_filename}".format(xform_weight_max=xform_weight_max,**params))
    else:
        lines.append("define-source input VNN {VNN_filename}".format(**params))

    # sources: Coulomb
    if (params["use_coulomb"]):
        lines.append("define-source input VC {VC_filename}".format(**params))
    lines.append("")

    # extract Hamiltonian parameters for convenience
    A = params["A"]
    a = params["a"]
    hw = params["hw"]
    hw_c = params["hw_c"]
    hw_cm = params["hw_cm"]
    if (hw_cm is None):
        hw_cm = hw

    # target: Hamiltonian
    coef_Uksqr = (A-1)/(2*A)*hw+1/(2*A)*a*(hw/hw_cm)
    coef_Vk1k2 = -1/A*hw+1/A*a*(hw/hw_cm)
    coef_Ursqr = 1/(2*A)*a*(hw_cm/hw)
    coef_Vr1r2 = 1/A*a*(hw_cm/hw)
    coef_identity = -3/2*a
    lines.append("define-target tbme-H.bin")
    lines.append("  add-source Ursqr {:e}".format(coef_Ursqr))
    lines.append("  add-source Vr1r2 {:e}".format(coef_Vr1r2))
    lines.append("  add-source Uksqr {:e}".format(coef_Uksqr))
    lines.append("  add-source Vk1k2 {:e}".format(coef_Vk1k2))
    lines.append("  add-source identity {:e}".format(coef_identity))
    lines.append("  add-source VNN {}".format(1.))
    if (params["use_coulomb"]):
        coef_VC = math.sqrt(hw/hw_c)
        lines.append("  add-source VC {:e}".format(coef_VC))

    # target: radius squared
    coef_Ursqr = (A-1)*(oscillator_length(hw)/A)**2
    coef_Vr1r2 = -2*(oscillator_length(hw)/A)**2
    lines.append("define-target tbme-rrel2.bin")
    lines.append("  add-source Ursqr {:e}".format(coef_Ursqr))
    lines.append("  add-source Vr1r2 {:e}".format(coef_Vr1r2))

    # target: NCM
    #
    # These are the "a" terms from the Hamiltonian, but without the "a".
    coef_Uksqr = 1/(2*A)*(hw/hw_cm)
    coef_Vk1k2 = 1/A*(hw/hw_cm)
    coef_Ursqr = 1/(2*A)*(hw_cm/hw)
    coef_Vr1r2 = 1/A*(hw_cm/hw)
    coef_identity = -3/2
    lines.append("define-target tbme-Ncm.bin")
    lines.append("  add-source Ursqr {:e}".format(coef_Ursqr))
    lines.append("  add-source Vr1r2 {:e}".format(coef_Vr1r2))
    lines.append("  add-source Uksqr {:e}".format(coef_Uksqr))
    lines.append("  add-source Vk1k2 {:e}".format(coef_Vk1k2))
    lines.append("  add-source identity {:e}".format(coef_identity))

    # target: Trel (diagnostic)
    #
    # These are the "non-a" terms from the Hamiltonian.
    coef_Uksqr = (A-1)/(2*A)*hw
    coef_Vk1k2 = -1/A*hw
    lines.append("define-target tbme-Trel.bin")
    lines.append("  add-source Uksqr {:e}".format(coef_Uksqr))
    lines.append("  add-source Vk1k2 {:e}".format(coef_Vk1k2))

    # target: VNN (diagnostic)
    lines.append("define-target tbme-VNN.bin")
    lines.append("  add-source VNN {}".format(1.))

    # target: VC (diagnostic)
    lines.append("define-target tbme-VC.bin")
    coef_VC = math.sqrt(hw/hw_c)
    lines.append("  add-source VC {:e}".format(coef_VC))

    # ensure terminal line
    lines.append("")

    return "\n".join(lines)

def mfdn_input(mfdn_params):
    """ Generate input for mfdn v14 beta06.

    Arguments:
        mfdn_params (dict): MFDn parameter dictionary

    Returns:
        (list of str): input file lines
    """

    lines = []

    # base parameters
    twice_Mj = int(2*mfdn_params["Mj"])
    lines.append("{:d}  # IFLAGMBSI".format(0))
    lines.append("{ndiag:d}  # ndiag (0: no spares, automatic ndiag)".format(**mfdn_params))
    lines.append("{:d}  # number of classes".format(2))
    lines.append("{nuclide[0]:d}  # protons (class 1 particles)".format(**mfdn_params))
    lines.append("{nuclide[1]:d}  # protons (class 2 particles)".format(**mfdn_params))
    lines.append("1 {Nshell:d}  # min, max # S.P. shells for class 1 particles".format(**mfdn_params))
    lines.append("1 {Nshell:d}  # min, max # S.P. shells for class 2 particles".format(**mfdn_params))
    lines.append("{Nmin:d} {Nmax:d} {Nstep:d}  # N_min, N_max, delta_N".format(**mfdn_params))
    lines.append("{:d}   # Total 2 M_j".format(twice_Mj))
    lines.append("{eigenvectors:d} {lanczos:d} {initial_vector:d} {tolerance:e}  # number of eigenvalues/vectors, max number of its, ...)".format(**mfdn_params))
    lines.append("{:d} {:d}  # rank of input Hamiltonian/interaction".format(2,2))
    lines.append("{hw_for_trans:g} {k_mN_csqr:g}  # h-bar*omega, Nucleon mass (MeV) ".format(k_mN_csqr=k_mN_csqr,**mfdn_params))

    # tbo parameters
    lines.append("tbme-H")  # Hamiltonian basename
    num_obs = 2 + len(mfdn_params["obs_basename_list"])
    lines.append("{:d}   # number of observables (J, T, R2, ...)".format(num_obs))

    lines += mfdn_params["obs_basename_list"]


    # obdme parameters
    lines.append("{enable_obd:d} {twice_multipolarity:d} # static one-body density matrix elements (0: no one-body densities), twice multipolarity".format(
        enable_obd=1,twice_multipolarity=2*mfdn_params["obdme_multipolarity"]
    ))
    lines.append("{num_reference_states:d} {max_delta_J:d} #number of reference states for transitions (0: no transitions, -1: all2all), max delta2J (?)".format(
        num_reference_states=len(mfdn_params["obdme_reference_state_list"]),
        max_delta_J=2*mfdn_params["obdme_multipolarity"]
    ))
    for reference_state in mfdn_params["obdme_reference_state_list"]:
        lines.append("{:d} {:d} {:d}".format(2*reference_state[0],reference_state[1],reference_state[2]))

    # ensure terminal line
    lines.append("")

    return "\n".join(lines)
