import dpdata
import numpy as np


def get_amber_fp(
    cutoff: float,
    target: str = ":1",
    out: str = None,
) -> dpdata.MultiSystems:
    """(qmwater-qmwater_not)-(mmwater-mmwater_not)
    = (qmwater-mmwater) - (qmwater_not-mmwater_not).
    """
    ms = dpdata.MultiSystems()
    ep = r"@%EP"
    if cutoff > 0.0:
        interactwith = f"({target})<@{cutoff:f}&!{ep}"
    else:
        interactwith = target

    s_qmwater = dpdata.LabeledSystem("qmwater", fmt="amber/md/qmmm", qm_region=target)
    s_qmwater = s_qmwater.remove_atom_names("EP")
    s_mmwater = dpdata.LabeledSystem("mmwater", fmt="amber/md/qmmm", qm_region=target)
    s_mmwater = s_mmwater.remove_atom_names("EP")
    assert np.all(s_qmwater["coords"] == s_mmwater["coords"])
    s_allcorr = s_mmwater.correction(s_qmwater)

    s_qmwaternot = dpdata.LabeledSystem("qmwater_not", fmt="amber/md/qmmm")
    s_qmwaternot = s_qmwaternot.remove_atom_names("EP")
    s_mmwaternot = dpdata.LabeledSystem("mmwater_not", fmt="amber/md/qmmm")
    s_mmwaternot = s_mmwaternot.remove_atom_names("EP")
    assert np.all(s_qmwaternot["coords"] == s_mmwaternot["coords"])
    s_notcorr = s_mmwaternot.correction(s_qmwaternot)

    # manually process
    natom_not = s_notcorr.get_natoms()
    assert np.all(s_allcorr["coords"][:, -natom_not:] == s_notcorr["coords"])
    s_allcorr.data["energies"] -= s_notcorr["energies"]
    s_allcorr.data["forces"][:, -natom_not:] -= s_notcorr["forces"]

    s_allcorr = s_allcorr.pick_by_amber_mask(
        "noepw.parm7", interactwith, pass_coords=True, nopbc=True
    )
    for ss in s_allcorr:
        ms.append(ss)

    if out:
        ms.to_deepmd_npy(out)
    return ms


def run(args):
    get_amber_fp(
        cutoff=args.cutoff,
        target=args.qm_region,
        out=args.out,
    )
