"""ob.py -- one-body operator definitions and functions

This file defines one-body operator functions and documents the syntax
for one-body matrix element sources.

A one-body source is defined by the dictionary mapping:

    id: parameter_dict

Here parameters_dict must contain exactly on of the following items, which
defines the one-body source type:

    "filename": input filename str

    "builtin": "kinematic"|"am"|"isospin"|"solid-harmonic"

    "linear-combination": CoefficientDict({"a": c_a, "b": c_b, ...})

    "tensor-product": [factor_a_id, factor_b_id]

Supporting information is provided by the supplementary items:

    "qn": tuple of (J0,g0,Tz0)

    "coefficient": float coefficient -- for "tensor-product" only

    "orbital_filename": orbital filename str (for generation in a larger space)
    -- for "builtin" only

A one-body operator (a.k.a target) is defined by the tuple:

    (name, qn, operator_id)

where operator_id is the id of an obme source.

Predefined sets of one-body observables can be generated by adding
the following to the list ob_observable_sets:

   "E0": E0p, E0n
   "intrinsic-E1": E1p, E1n
   "M1": M1, DLp, DLn, DSp, DSn
   "E<n>": E<n>p, E<n>n
   "M<n>": M<n>, M<n>lp, M<n>ln, M<n>sp, M<n>sn
   "F": F
   "GT: GT
     => "beta": F, GT

Patrick J. Fasano
University of Notre Dame

- 09/09/20 (pjf): Created, with some content from operators.py
- 09/12/20 (pjf): Overhaul logic for generating one-body targets and sources.
- 11/13/20 (pjf): Use constants module.
- 07/09/21 (pjf):
    + Fix EL scaling.
    + Add definitions for rtz, iktz, ltz, stz.
    + Fix solid_harmonic_source for ik solid harmonics.
    + Use pattern matching for setting up solid harmonic sources.
    + Print warning when overriding with user-defined sources.
- 07/14/21 (pjf):
    + Add definitions for r.rtz and ik.iktz.
    + Add E1 intrinsic correction.
- 02/11/22 (mac): Define observable set name "intrinsic-E1" for
    intrinsic E1 to avoid overloading lab frame E1.
- 02/15/22 (pjf): Fix tz vs. tau0 scaling in intrinsic-E1 operator.
- 06/28/22 (mac): Fix E0n operator in E0 observable set.
- 08/03/22 (pjf): Generate both Tz-raising and lowering operators for beta decay.
- 08/06/22 (pjf): Fix phases on Tz-lowering beta decay operators so that they
    are related by (-1)^Tz0.
- 01/30/23 (pjf): Rename j0->J0 and tz0->Tz0.

"""
import collections
import math
import re

import mcscript.exception
import mcscript.utils

from .. import (
    constants,
    environ,
    utils,
    modes,
)
from . import tb

################################################################
# kinematic operators
################################################################

k_kinematic_operators = {
    "identity": {"builtin": "kinematic", "qn": (0,0,0)},
    "r":        {"builtin": "kinematic", "qn": (1,1,0)},
    "r.r":      {"builtin": "kinematic", "qn": (0,0,0)},
    "ik":       {"builtin": "kinematic", "qn": (1,1,0)},
    "ik.ik":    {"builtin": "kinematic", "qn": (0,0,0)},
    "rtz":      {"tensor-product": ["r","tz"], "qn": (1,1,0)},
    "r.rtz":    {"tensor-product": ["r.r","tz"], "qn": (0,0,0)},
    "iktz":     {"tensor-product": ["ik","tz"], "qn": (1,1,0)},
    "ik.iktz":  {"tensor-product": ["ik.ik","tz"], "qn": (0,0,0)},
}

################################################################
# angular momentum operators
################################################################

k_am_operators = {
    "j":  {"builtin": "am", "qn": (1,0,0)},
    "j2": {"builtin": "am", "qn": (1,0,0)},
    "l":  {"builtin": "am", "qn": (1,0,0)},
    "l2": {"builtin": "am", "qn": (1,0,0)},
    "s":  {"builtin": "am", "qn": (1,0,0)},
    "s2": {"builtin": "am", "qn": (1,0,0)},
    "ltz": {"tensor-product": ["l", "tz"], "qn": (1,0,0)},
    "stz": {"tensor-product": ["s", "tz"], "qn": (1,0,0)},
    "sp": {"tensor-product": ["delta_p", "s"], "qn": (1,0,0)},
    "sn": {"tensor-product": ["delta_n", "s"], "qn": (1,0,0)},
    "sp2": {"tensor-product": ["delta_p", "s2"], "qn": (0,0,0)},
    "sn2": {"tensor-product": ["delta_n", "s2"], "qn": (0,0,0)},
}

################################################################
# isospin operators
################################################################
k_isospin_operators = {
    "tz": {"builtin": "isospin", "qn": (0,0,0)},
    "t+": {"builtin": "isospin", "qn": (0,0,+1)},
    "t-": {"builtin": "isospin", "qn": (0,0,-1)},
    "delta_p": {"linear-combination": mcscript.utils.CoefficientDict({"identity": 0.5, "tz": +1.0}), "qn": (0,0,0)},
    "delta_n": {"linear-combination": mcscript.utils.CoefficientDict({"identity": 0.5, "tz": -1.0}), "qn": (0,0,0)},
}

################################################################
# ladder operators
################################################################
k_ladder_operators_native = {
    "c+": {"builtin": "ladder", "qn": (1,1,0)},
    "c-": {"builtin": "ladder", "qn": (1,1,0)},
}

def ladder_operators_generic(hw):
    """Generate ladder operator source info for generic basis.

    Arguments:
        hw (float): hw of basis

    Returns:
        (dict): operator sources
    """

    qn = (1,1,0)
    b = utils.oscillator_length(hw)
    sources = {}
    sources["c+"] = {
        "linear-combination": mcscript.utils.CoefficientDict({
            "r": 1/b, "ik": -b
        }),
        "qn": qn
    }
    sources["c"] = {
        "linear-combination": mcscript.utils.CoefficientDict({
            "r": 1/b, "ik": b
        }),
        "qn": qn
    }

    return sources

################################################################
# solid harmonic operators
################################################################
def solid_harmonic_source(coordinate, order, J0=None):
    """Generate solid harmonic source info.

    Arguments:
        coordinate (str): "r" or "ik"
        order (int): power of coordinate (i.e. r^n)
        J0 (int, optional): rank of spherical harmonic (defaults to order)

    Returns:
        (id (str), source (dict)): id and source dictionary
    """
    if J0 is None:
        J0 = order
    if coordinate not in {"r", "ik"}:
        raise mcscript.exception.ScriptError("unknown coordinate {}".format(coordinate))

    qn = (J0,J0%2,0)  # (J0,g0,Tz0)
    identifier = "{:s}{:d}Y{:d}".format(coordinate, order, J0)
    source_dict = {
        "builtin": "solid-harmonic", "coordinate": coordinate, "order": order, "qn": qn
    }
    return (identifier, source_dict)

################################################################
# predefined sets
################################################################

k_h2mixer_builtin = {
    "identity",
    "j",  # this is safe, as our orbitals are always eigenstates of j
    "l2", "s2", "j2",
    "tz",
}

################################################################
# obme sources
################################################################
def generate_ob_observable_sets(task):
    """
    """
    ob_observables = []
    obme_sources = {}
    ob_observable_sets = task.get("ob_observable_sets", [])
    for name in ob_observable_sets:
        # special case for E0
        if name == "E0":
            qn = (0,0,0)
            ob_observables += [
                ("E0p", qn, "E0p"),
                ("E0n", qn, "E0n"),
            ]
            obme_sources["r.r"] = k_kinematic_operators["r.r"]
            obme_sources["E0p"] = {
                "tensor-product": ["delta_p", "r.r"],
                "qn": qn
            }
            obme_sources["E0n"] = {
                "tensor-product": ["delta_n", "r.r"],
                "qn": qn
            }
            continue

        # special case for (intrinsic) E1
        #
        # see equation (30) of "intrinsic" [Caprio 2020, doi:10.1088/1361-6471/ab9d38]
        #
        # note different factor of 1/2 from eq. (30), arising from $t_z$ vs $tau_0$
        if name == "intrinsic-E1":
            nuclide = task["nuclide"]
            (Z,N) = nuclide
            A = sum(nuclide)
            qn = (1,1,0)
            ob_observables += [
                ("E1p", qn, "E1p"),
                ("E1n", qn, "E1n"),
            ]
            obme_sources["E1p"] = {
                "linear-combination": {"r1Y1tz": 1., "r1Y1": -1.*(Z-N)/A},
                "qn": (1,1,0),
            }
            obme_sources["E1n"] = {
                "linear-combination": {"r1Y1tz": -1., "r1Y1": 1.*(Z-N)/A},
                "qn": (1,1,0),
            }
            continue

        # electric transitions (general)
        match = re.fullmatch(r"E([0-9]+)", name)
        if match:
            order = int(match.group(1))
            qn = (order,order%2,0)
            (J0, _, _) = qn
            ob_observables += [
                ("E{}p".format(order), qn, "E{}p".format(order)),
                ("E{}n".format(order), qn, "E{}n".format(order)),
            ]
            solid_harmonic_id = "r{:d}Y{:d}".format(order, J0)
            obme_sources["E{}p".format(order)] = {
                "tensor-product": ["delta_p", solid_harmonic_id],
                "qn": qn
            }
            obme_sources["E{}n".format(order)] = {
                "tensor-product": ["delta_n", solid_harmonic_id],
                "qn": qn
            }

        # special case for M1
        if name == "M1":
            qn = (1,0,0)
            coefficient = math.sqrt(3/(4*math.pi))
            ob_observables += [
                ("Dlp", qn, "Dlp"),
                ("Dln", qn, "Dln"),
                ("Dsp", qn, "Dsp"),
                ("Dsn", qn, "Dsn"),
                ("M1", qn, "M1"),
            ]
            obme_sources["l"] = k_am_operators["l"]
            obme_sources["s"] = k_am_operators["s"]
            obme_sources["Dlp"] = {
                "tensor-product": ["delta_p", "l"], "coefficient": coefficient, "qn": qn
            }
            obme_sources["Dln"] = {
                "tensor-product": ["delta_n", "l"], "coefficient": coefficient, "qn": qn
            }
            obme_sources["Dsp"] = {
                "tensor-product": ["delta_p", "s"], "coefficient": coefficient, "qn": qn
            }
            obme_sources["Dsn"] = {
                "tensor-product": ["delta_n", "s"], "coefficient": coefficient, "qn": qn
            }
            obme_sources["M1"] = {
                "linear-combination": {
                    "Dlp": 1.0,
                    # "Dln": 0.0,
                    "Dsp": constants.k_gp,
                    "Dsn": constants.k_gn,
                },
                "qn": qn,
            }
            continue

        # magnetic transitions (general)
        match = re.fullmatch(r"M([0-9]+)", name)
        if match:
            order = int(match.group(1))
            J0 = order
            if order == 0:
                raise mcscript.exception.ScriptError("you must construct additional magnetic monopoles")
            qn = (order,(order-1)%2,0)
            l_coefficient = math.sqrt((2*order+1)*order) * 2/(order+1)
            s_coefficient = math.sqrt((2*order+1)*order)
            ob_observables += [
                ("M{}lp".format(order), qn, "M{}lp".format(order)),
                ("M{}ln".format(order), qn, "M{}ln".format(order)),
                ("M{}sp".format(order), qn, "M{}sp".format(order)),
                ("M{}sn".format(order), qn, "M{}sn".format(order)),
                ("M{}".format(order), qn, "M{}".format(order)),
            ]
            obme_sources["l"] = k_am_operators["l"]
            obme_sources["s"] = k_am_operators["s"]
            solid_harmonic_id = "r{:d}Y{:d}".format(order-1, J0-1)
            obme_sources["l"+solid_harmonic_id] = {
                "tensor-product": ["l", solid_harmonic_id], "qn": qn
            }
            obme_sources["s"+solid_harmonic_id] = {
                "tensor-product": ["s", solid_harmonic_id], "qn": qn
            }
            obme_sources["M{}lp".format(order)] = {
                "tensor-product": ["delta_p", "l"+solid_harmonic_id],
                "coefficient": l_coefficient,
                "qn": qn
            }
            obme_sources["M{}ln".format(order)] = {
                "tensor-product": ["delta_n", "l"+solid_harmonic_id],
                "coefficient": l_coefficient,
                "qn": qn
            }
            obme_sources["M{}sp".format(order)] = {
                "tensor-product": ["delta_p", "s"+solid_harmonic_id],
                "coefficient": s_coefficient,
                "qn": qn
            }
            obme_sources["M{}sn".format(order)] = {
                "tensor-product": ["delta_n", "s"+solid_harmonic_id],
                "coefficient": s_coefficient,
                "qn": qn
            }
            obme_sources["M1"] = {
                "linear-combination": {
                    "M{}lp".format(order): 1.0,
                    # "M{}ln".format(order): 0.0,
                    "M{}sp".format(order): constants.k_gp,
                    "M{}sn".format(order): constants.k_gn,
                },
                "qn": qn,
            }

        if name in {"F", "beta"}:
            # note: we can use the same observable name twice since only one
            # will ever be applicable for a given pair of nuclei
            ob_observables += [("F", (0,0,+1), "t+")]
            ob_observables += [("F", (0,0,-1), "F-")]
            obme_sources["t+"] = k_isospin_operators["t+"]
            obme_sources["t-"] = k_isospin_operators["t-"]
            # we define F- with a negative phase so that F- = (-1)^Tz0 F+
            obme_sources["F-"] = {"linear-combination": {"t-": -1}, "qn": (0,0,-1)}

        if name in {"GT", "beta"}:
            ob_observables += [("GT", (1,0,+1), "GT+")]
            ob_observables += [("GT", (1,0,-1), "GT-")]
            obme_sources["t+"] = k_isospin_operators["t+"]
            obme_sources["t-"] = k_isospin_operators["t-"]
            obme_sources["s"] = k_am_operators["s"]
            # note that the Gamow-Teller operator is defined as sigma*tau/2,
            # which simplifies to 2*s*t -- we do NOT introduce the additional
            # factor of 1/sqrt(3) present in Suhonen 7.21, since this operator
            # is used with the standard equation for one-body operators (Suhonen
            # 4.25) rather than the (oddly beta-decay-specific) conventional
            # equation Suhonen 7.19
            obme_sources["GT+"] = {"tensor-product": ["s", "t+"], "coefficient": 2.0, "qn": (1,0,+1)}
            # we define GT- with a negative phase so that GT- = (-1)^Tz0 GT+
            obme_sources["GT-"] = {"tensor-product": ["s", "t-"], "coefficient": -2.0, "qn": (1,0,-1)}

    return (ob_observables, obme_sources)

def get_obme_targets_h2mixer(task, tbme_targets):
    """Get OBME target list for h2mixer for a given set of TBME targets.

    These are the sources required internally by h2mixer, for upgrading to
    two-body via U or V.

    Arguments:
        task (dict): as described in module docstring
        tbme_targets (list of target dicts): list of TBME target dictionaries

    Returns:
        (set of str): set of OBME targets for h2mixer
    """
    # accumulate obme targets
    obme_targets = set()

    # extract dependencies from tbme targets
    #   postfix is irrelevant for this purpose
    tbme_sources = tb.get_tbme_sources(task, tbme_targets, postfix="")
    for tbme_source in tbme_sources.values():
        if "operatorU" in tbme_source:
            obme_targets.add(tbme_source["operatorU"])
        if "operatorV" in tbme_source:
            obme_targets.update(tbme_source["operatorV"])

    return obme_targets

def get_obme_targets_observables(task):
    """Get OBME target set for observables.

    These are the sources required for one-body observables, i.e. for input to
    obscalc-ob.

    Arguments:
        task (dict): as described in module docstring

    Returns:
        (set of str): set of observable OBME targets
    """
    # accumulate obme targets
    obme_targets = set()

    for (basename, qn, operator) in generate_ob_observable_sets(task)[0]:
        obme_targets.add(operator)
    for (basename, qn, operator) in task.get("ob_observables", []):
        obme_targets.add(operator)

    return obme_targets

def get_obme_targets_obmixer(task):
    """Get OBME target set for output by obmixer.

    These are the sources for observables plus the sources which cannot be
    generated within h2mixer.

    Arguments:
        task (dict): as described in module docstring

    Returns:
        (set of str): set of OBME targets for obmixer
    """
    # accumulate obme targets
    obme_targets = set()

    # get h2mixer targets and connected sources
    tbme_targets_by_qn = tb.get_tbme_targets(task)
    obme_targets_h2mixer = set()
    for tbme_targets in tbme_targets_by_qn.values():
        obme_targets_h2mixer.update(
            get_obme_targets_h2mixer(task, tbme_targets)  # postfix doesn't matter
        )
    obme_sources_h2mixer = get_obme_sources(task, obme_targets_h2mixer)

    # iterate over sources and accumulate those which cannot
    # be generated by h2mixer as targets for obmixer
    for (identifier, source) in obme_sources_h2mixer.items():
        if ("builtin" in source) and (identifier not in k_h2mixer_builtin):
            obme_targets.add(identifier)

    # add observable targets
    obme_targets.update(get_obme_targets_observables(task))

    return obme_targets

def get_obme_sources(task, targets):
    """Get OBME sources for task.

    Arguments:
        task (dict): as described in module docstring
        targets (set): set of targets to generate

    Returns:
        (OrderedDict of dict): id to source mapping, sorted in reverse
            topological order
    """
    # accumulate sources
    obme_sources = {}

    # gather pre-defined sources first
    obme_sources.update(**k_kinematic_operators, **k_am_operators, **k_isospin_operators)
    if task.get("basis_mode") in {modes.BasisMode.kDirect, modes.BasisMode.kDilated}:
        obme_sources.update(**k_ladder_operators_native)
    else:
        obme_sources.update(**ladder_operators_generic(task["hw"]))

    # add sources from observable sets
    obme_sources.update(**generate_ob_observable_sets(task)[1])

    # set up solid harmonics
    solid_harmonic_re = re.compile(r"(r|ik)([0-9]+)Y([0-9]+)")
    iv_solid_harmonic_re = re.compile(r"(r|ik)([0-9]+)Y([0-9]+)tz")
    required_solid_harmonics = set()
    for target_id in targets:
        match = solid_harmonic_re.match(target_id)
        if match:
            required_solid_harmonics.add(match.group(0))
        match = iv_solid_harmonic_re.match(target_id)
        if match:
            required_solid_harmonics.add(match.group(0))
    for source in obme_sources.values():
        for source_id in list(source.get("linear-combination", {}).keys()) + source.get("tensor-product", []):
            match = solid_harmonic_re.match(source_id)
            if match:
                required_solid_harmonics.add(match.group(0))
            match = iv_solid_harmonic_re.match(source_id)
            if match:
                required_solid_harmonics.add(match.group(0))
    for solid_harmonic_id in required_solid_harmonics:
        match = solid_harmonic_re.fullmatch(solid_harmonic_id)
        if match:
            coordinate = match.group(1)
            order = int(match.group(2))
            J0 = int(match.group(3))
            (source_id, source) = solid_harmonic_source(coordinate, order, J0)
            assert source_id == solid_harmonic_id
            obme_sources[source_id] = source

        match = iv_solid_harmonic_re.fullmatch(solid_harmonic_id)
        if match:
            coordinate = match.group(1)
            order = int(match.group(2))
            J0 = int(match.group(3))
            (source_id, source) = solid_harmonic_source(coordinate, order, J0)
            assert source_id+"tz" == solid_harmonic_id
            obme_sources[solid_harmonic_id] = {
                "tensor-product": [source_id, "tz"],
                "qn": source["qn"],
            }

    # override with user-defined sources
    user_obme_sources = task.get("obme_sources", [])
    for (source_id, source) in user_obme_sources:
        if source_id in obme_sources:
            print("WARN: overriding obme source '{:s}' with {}".format(source_id, source))
        obme_sources[source_id] = source

    # construct dependency graph
    obme_dependency_graph = {}
    for (source_id, source) in obme_sources.items():
        if source.get("linear-combination"):
            obme_dependency_graph[source_id] = source["linear-combination"].keys()
        elif source.get("tensor-product"):
            obme_dependency_graph[source_id] = source["tensor-product"]
        else:
            obme_dependency_graph[source_id] = []

    # construct minimal set of sources (in reverse topological order)
    sorted_obme_sources = collections.OrderedDict()
    for id_ in reversed(mcscript.utils.topological_sort(obme_dependency_graph, targets)):
        sorted_obme_sources[id_] = obme_sources[id_]

    return sorted_obme_sources

def get_obme_sources_h2mixer(task, targets, postfix):
    """Get OBME sources for task (for use by h2mixer).

    This modifies the list of OBME sources generated by get_obme_sources()
    by determining which sources are written to file by obmixer, and reading
    them from file instead of regenerating them.

    Arguments:
        task (dict): as described in module docstring
        targets (set): set of targets to generate

    Returns:
        (OrderedDict of dict): id to source mapping, sorted in reverse
            topological order
    """
    # input normalization
    if not isinstance(targets, set):
        targets = set(targets)

    # get obme sources
    obme_sources = get_obme_sources(task, targets)

    # sources which were previously a target for obmixer should now be file inputs
    #   this is achieved by erasing their dependency information; this
    #   turns them into leaf nodes
    for identifier in (get_obme_targets_obmixer(task) & obme_sources.keys()):
        obme_sources[identifier] = {
            "filename": environ.obme_filename(postfix, identifier),
            "qn": obme_sources[identifier]["qn"]
        }

    # re-construct dependency graph
    obme_dependency_graph = {}
    for (source_id, source) in obme_sources.items():
        if source.get("linear-combination"):
            obme_dependency_graph[source_id] = source["linear-combination"].keys()
        elif source.get("tensor-product"):
            obme_dependency_graph[source_id] = source["tensor-product"]
        else:
            obme_dependency_graph[source_id] = []

    # re-construct minimal set of sources (in reverse topological order)
    sorted_obme_sources = collections.OrderedDict()
    for id_ in reversed(mcscript.utils.topological_sort(obme_dependency_graph, targets)):
        sorted_obme_sources[id_] = obme_sources[id_]

    return sorted_obme_sources
