# Recopie les clips d'un pack d'animations sur Bridget et Michael (squelettes Genesis 8 du jeu).
# Les chaines portent le meme nom des deux cotes pour que le retarget les aligne.
# Les maillages bridget_gp et michael ne sont pas reimportes.
# Lot AP7 : script generique (ex-RetargetQuaternius.py), pilote par une methode de transformation
# (Tools/anim/methodes/<id>.json ; "ue4" = ancienne methode Quaternius) et par Saved/anim_retarget.json.
#
# Lance par Tools/BuildMod.ps1 -Retarget (UnrealEditor.exe : la copie a besoin de Slate).
# Saved/anim_retarget.json (absent : configuration Quaternius par defaut), champs et defauts :
#   mode "bake"|"rigs" (defaut "bake" ; a defaut de "mode", Saved/anim_to_genesis_step.txt est lu comme avant),
#   methode "ue4", source_mesh /Game/Quaternius/UAL/UAL1_Standard.UAL1_Standard, clips_dir /Game/Quaternius/UAL,
#   out_dir = clips_dir, dest /Game/Quaternius/Retarget (ue4) ou /Game/Retarget_<methode>,
#   corps [{nom,maillage,prefixe}] Bridget (B_) et Michael (M_), exclure [], pose_dir Saved/Measurement,
#   jeton (ecrit par BuildMod, ajoute en fin des lignes de resultat et d'echec : "... jeton <jeton>").
# Methode factice ou copie, ou profil du maillage source different de profil_source : ligne
# "RETARGET echec : <cause>" puis arret.
#
# Phase C (CHG-204) : la pose de retarget de chaque corps est la pose ajustee sur le mannequin
# par la sandbox Mesure (LPPoseMeasure::WriteRetargetPose) : Saved/Measurement/retarget_pose_<Corps>.txt.
# Colonne et cou ont plus d'os chez Genesis : chaines Interpolated. Les os Twist restent hors chaine.
# AP1 : les chaines et la racine viennent des profils Tools/anim/profiles (identification par les noms d'os).

import json
import math
import os

import unreal

DEFAULT_CONFIG = {
    "mode": None,
    "methode": "ue4",
    "source_mesh": "/Game/Quaternius/UAL/UAL1_Standard.UAL1_Standard",
    "clips_dir": "/Game/Quaternius/UAL",
    "out_dir": None,
    "dest": None,
    "corps": [
        {"nom": "Bridget", "maillage": "/Game/Heavy/Daz/Bridget/bridget_gp.bridget_gp", "prefixe": "B_"},
        {"nom": "Michael", "maillage": "/Game/Heavy/Daz/Michael/michael.michael", "prefixe": "M_"},
    ],
    "exclure": [],
    "pose_dir": None,
    "jeton": None,
}
NONE = unreal.Name()
CFG = {}
JETON = None

# Chaines derivees des profils de squelettes (lot AP1) : Tools/anim/profiles/*.json, lus aussi par la
# sandbox (LPSkeletonProfile). Une chaine ne traverse jamais un os *Twist (hors chaine, il suit l'os Bend).
HERE = os.path.dirname(os.path.abspath(__file__))
PROFILE_DIR = os.path.join(HERE, "profiles")
METHOD_DIR = os.path.join(HERE, "methodes")
SIDES = (("l", "Left"), ("r", "Right"))
FINGERS = (("thumb", "Thumb"), ("index", "Index"), ("middle", "Mid"), ("ring", "Ring"), ("pinky", "Pinky"))


class Refusal(Exception):
    pass


def load_config(saved_dir):
    # Configuration par defaut = comportement Quaternius ; Saved/anim_retarget.json la complete champ par champ.
    config = json.loads(json.dumps(DEFAULT_CONFIG))
    path = os.path.join(saved_dir, "anim_retarget.json")
    if os.path.exists(path):
        with open(path, "r", encoding="utf-8-sig") as handle:
            given = json.load(handle)
        global JETON
        JETON = str(given["jeton"]) if given.get("jeton") else None
        for key, value in given.items():
            if key not in DEFAULT_CONFIG:
                raise Refusal("champ inconnu %s dans %s" % (key, path))
            config[key] = value
    if config["mode"] is None:
        config["mode"] = "bake"
        legacy = os.path.join(saved_dir, "anim_to_genesis_step.txt")
        if os.path.exists(legacy):
            with open(legacy, "r", encoding="utf-8") as handle:
                config["mode"] = handle.read().strip().lower() or "bake"
    config["mode"] = str(config["mode"]).strip().lower()
    if config["mode"] not in ("bake", "rigs"):
        raise Refusal("etape inconnue %s" % config["mode"])
    config["clips_dir"] = config["clips_dir"].rstrip("/")
    config["out_dir"] = (config["out_dir"] or config["clips_dir"]).rstrip("/")
    if not config["dest"]:
        config["dest"] = "/Game/Quaternius/Retarget" if config["methode"] == "ue4" else "/Game/Retarget_%s" % config["methode"]
    config["dest"] = config["dest"].rstrip("/")
    if not config["pose_dir"]:
        config["pose_dir"] = os.path.join(saved_dir, "Measurement")
    if not config["corps"]:
        raise Refusal("aucun corps dans la configuration")
    return config


def load_method(method_id, method_dir=None):
    path = os.path.join(method_dir or METHOD_DIR, "%s.json" % method_id)
    if not os.path.exists(path):
        raise Refusal("methode inconnue %s" % method_id)
    with open(path, "r", encoding="utf-8") as handle:
        method = json.load(handle)
    if method.get("factice") or method.get("copie"):
        raise Refusal("methode %s non implementee (factice ou copie)" % method_id)
    return method


def suffix(config):
    # Les assets de la methode ue4 gardent leurs noms historiques ; les autres sont suffixes par la methode.
    return "" if config["methode"] == "ue4" else "_" + config["methode"]


def source_rig_name(config):
    return "IK_Quaternius" if config["methode"] == "ue4" else "IK_Source" + suffix(config)


def load_profiles():
    profiles = []
    for name in sorted(os.listdir(PROFILE_DIR)):
        if name.endswith(".json"):
            with open(os.path.join(PROFILE_DIR, name), "r", encoding="utf-8") as handle:
                profiles.append(json.load(handle))
    if not profiles:
        raise RuntimeError("aucun profil de squelette dans %s" % PROFILE_DIR)
    return profiles


def role(profile, name, index=0):
    value = profile["roles"].get(name)
    if value is None:
        return None
    if isinstance(value, list):
        return value[index] if -len(value) <= index < len(value) else None
    return value if index in (0, -1) else None


def identify_names(names, label="?"):
    # Meme regle que LPSkeletonProfile::Identify : au moins 95 % des os du profil presents ; a egalite de
    # score, le profil qui reconnait le plus d'os (le mannequin UE5 contient tous les noms d'os UE4).
    names = set(names)
    best, best_score, best_hits = None, 0.0, 0
    for profile in load_profiles():
        bones = []
        for value in profile["roles"].values():
            bones.extend(value if isinstance(value, list) else [value])
        hits = sum(1 for bone in bones if bone in names)
        score = round(hits / float(len(bones)), 6)
        if score > best_score or (score == best_score and hits > best_hits):
            best, best_score, best_hits = profile, score, hits
    if best is None or best_score < 0.95:
        raise RuntimeError("squelette inconnu %s (meilleur profil %.0f%%)" % (label, best_score * 100.0))
    return best, best_score


def identify(mesh):
    best, best_score = identify_names(bone_names(mesh), mesh.get_name())
    unreal.log_warning("RETARGET profil %s -> %s (%.0f%%)" % (mesh.get_name(), best["id"], best_score * 100.0))
    return best


def chains(profile):
    # (nom de chaine, premier os, dernier os) : memes noms des deux cotes pour auto_map_chains.
    made = [("Spine", role(profile, "spine", 0), role(profile, "spine", -1)),
            ("Neck", role(profile, "neck", 0), role(profile, "head"))]
    for side, word in SIDES:
        made += [
            (word + "Collar", role(profile, "clavicle_" + side), role(profile, "upperarm_" + side)),
            (word + "Forearm", role(profile, "lowerarm_" + side), role(profile, "lowerarm_" + side)),
            (word + "Hand", role(profile, "hand_" + side), role(profile, "hand_" + side)),
            (word + "Thigh", role(profile, "thigh_" + side), role(profile, "thigh_" + side)),
            (word + "Calf", role(profile, "calf_" + side), role(profile, "calf_" + side)),
            (word + "Foot", role(profile, "foot_" + side), role(profile, "foot_" + side)),
            (word + "Toe", role(profile, "toe_" + side), role(profile, "toe_" + side)),
        ]
        for finger, daz in FINGERS:
            made.append((word + daz, role(profile, finger + "_" + side, 0), role(profile, finger + "_" + side, -1)))
    missing = [name for name, start, end in made if not start or not end]
    if missing:
        raise RuntimeError("profil %s : chaines sans os %s" % (profile["id"], ", ".join(missing)))
    return made


INTERPOLATED = ("Spine", "Neck")


def make_factory(factory_class):
    factory = factory_class()
    factory.set_editor_property("edit_after_new", False)
    return factory


def vsub(a, b):
    return unreal.MathLibrary.subtract_vector_vector(a, b)


def vnorm(v):
    return unreal.MathLibrary.normal(v)


def vdot(a, b):
    return a.x * b.x + a.y * b.y + a.z * b.z


def vcross(a, b):
    return unreal.MathLibrary.cross_vector_vector(a, b)


def bone_names(mesh):
    pose = mesh.skeleton.get_reference_pose()
    return [str(name) for name in unreal.AnimPoseExtensions.get_bone_names(pose)]


def pose_file(body):
    return os.path.join(CFG["pose_dir"], "retarget_pose_%s.txt" % body)


def read_pose(body):
    path = pose_file(body)
    if not os.path.exists(path):
        raise RuntimeError("pose de retarget absente %s : ouvrir la sandbox Mesure une fois" % path)
    offsets = []
    with open(path, "r", encoding="utf-8-sig") as handle:
        for line in handle:
            cells = line.split()
            if len(cells) != 5 or cells[0].startswith("#"):
                continue
            offsets.append((cells[0], unreal.Quat(float(cells[1]), float(cells[2]), float(cells[3]), float(cells[4]))))
    return offsets


def apply_pose(ctrl, mesh, body):
    side = unreal.RetargetSourceOrTarget.TARGET
    names = set(bone_names(mesh))
    applied = 0
    for bone, quat in read_pose(body):
        if bone not in names:
            unreal.log_warning("RETARGET pose %s os absent %s" % (body, bone))
            continue
        ctrl.set_rotation_offset_for_retarget_pose_bone(bone, quat, side)
        applied += 1
    unreal.log_warning("RETARGET pose %s %d os (pose %s ajustee)" % (body, applied, CFG["pose"]))


def make_rig(name, mesh_path):
    asset_path = CFG["dest"] + "/" + name
    if unreal.EditorAssetLibrary.does_asset_exist(asset_path):
        unreal.EditorAssetLibrary.delete_asset(asset_path)
    mesh = unreal.load_asset(mesh_path)
    profile = identify(mesh)
    root = role(profile, "hips")
    rig_chains = chains(profile)
    rig = unreal.AssetToolsHelpers.get_asset_tools().create_asset(
        name, CFG["dest"], unreal.IKRigDefinition, make_factory(unreal.IKRigDefinitionFactory))
    ctrl = unreal.IKRigController.get_controller(rig)
    ctrl.set_skeletal_mesh(mesh)
    ctrl.set_retarget_root(root)
    for chain, start, end in rig_chains:
        made = str(ctrl.add_retarget_chain(chain, start, end, NONE))
        unreal.log_warning("RETARGET chaine %s %s %s -> %s" % (name, chain, start, made))
        if made == "None":
            raise RuntimeError("chaine impossible %s %s %s" % (name, chain, start))
    unreal.EditorAssetLibrary.save_loaded_asset(rig)
    return rig, mesh, rig_chains


def make_retargeter(name, source_rig, target_rig, source_mesh, target_mesh, body, source_chains):
    asset_path = CFG["dest"] + "/" + name
    if unreal.EditorAssetLibrary.does_asset_exist(asset_path):
        unreal.EditorAssetLibrary.delete_asset(asset_path)
    retargeter = unreal.AssetToolsHelpers.get_asset_tools().create_asset(
        name, CFG["dest"], unreal.IKRetargeter, make_factory(unreal.IKRetargetFactory))
    ctrl = unreal.IKRetargeterController.get_controller(retargeter)
    ctrl.set_ik_rig(unreal.RetargetSourceOrTarget.SOURCE, source_rig)
    ctrl.set_ik_rig(unreal.RetargetSourceOrTarget.TARGET, target_rig)
    ctrl.set_preview_mesh(unreal.RetargetSourceOrTarget.SOURCE, source_mesh)
    ctrl.set_preview_mesh(unreal.RetargetSourceOrTarget.TARGET, target_mesh)
    ctrl.auto_map_chains(unreal.AutoMapChainType.EXACT, True)
    # Pose de retarget : la pose ajustee de la sandbox, et non plus l'alignement par le maillage,
    # les reglages Daz et les demi-tours d'epaule (pole / roll180), qu'elle remplace.
    apply_pose(ctrl, target_mesh, body)
    for chain_name, _, _ in source_chains:
        settings = ctrl.get_retarget_chain_settings(chain_name)
        fk = settings.get_editor_property("fk")
        mode = unreal.RetargetRotationMode.INTERPOLATED if chain_name in INTERPOLATED else unreal.RetargetRotationMode.ONE_TO_ONE
        fk.set_editor_property("rotation_mode", mode)
        fk.set_editor_property("translation_mode", unreal.RetargetTranslationMode.NONE)
        fk.set_editor_property("pole_vector_matching", 0.0)
        settings.set_editor_property("fk", fk)
        ctrl.set_retarget_chain_settings(chain_name, settings)
    unreal.EditorAssetLibrary.save_loaded_asset(retargeter)
    return retargeter


def source_clips():
    prefixes = tuple(body["prefixe"] for body in CFG["corps"])
    clips = []
    for path in unreal.EditorAssetLibrary.list_assets(CFG["clips_dir"], False, False):
        asset = unreal.load_asset(path)
        if not isinstance(asset, unreal.AnimSequence):
            continue
        name = asset.get_name()
        if name.startswith(prefixes) or name in CFG["exclure"]:
            continue
        clips.append(unreal.EditorAssetLibrary.find_asset_data(path))
    return clips


def delete_copies(names):
    # Seulement les copies des clips source. Pas tout asset dont le nom commence par un prefixe de corps.
    for name in names:
        for body in CFG["corps"]:
            for folder in (CFG["out_dir"] + "/", "/Game/"):
                path = folder + body["prefixe"] + name
                if unreal.EditorAssetLibrary.does_asset_exist(path):
                    unreal.EditorAssetLibrary.delete_asset(path)


def tag(line):
    # AP7 : BuildMod ne retient que les lignes de resultat portant son jeton (cle "jeton" de anim_retarget.json).
    return "%s jeton %s" % (line, JETON) if JETON else line


def finish(line):
    unreal.EditorLoadingAndSavingUtils.save_dirty_packages(True, True)
    unreal.log_warning(tag(line))
    unreal.SystemLibrary.quit_editor()


def anim_pose(sequence, time):
    try:
        options = unreal.AnimPoseEvaluationOptions()
        for prop in ("should_retarget", "b_should_retarget"):
            try:
                options.set_editor_property(prop, False)
            except Exception:
                pass
        return unreal.AnimPoseExtensions.get_anim_pose_at_time(sequence, time, options)
    except Exception as exc:
        unreal.log_warning("RETARGET pose api %s" % exc)
        return unreal.AnimPoseExtensions.get_anim_pose_at_time(sequence, time)


def hand_ahead(sequence, chest_name, hand_name, left_name, right_name, time, pelvis_name="pelvis"):
    pose = anim_pose(sequence, time)
    space = unreal.AnimPoseSpaces.WORLD

    def loc(bone):
        return unreal.AnimPoseExtensions.get_bone_pose(pose, bone, space).translation

    chest = loc(chest_name)
    pelvis = loc(pelvis_name)
    up = vnorm(vsub(chest, pelvis))
    side = vnorm(vsub(loc(left_name), loc(right_name)))
    forward = vnorm(vcross(side, up))
    return vdot(vsub(loc(hand_name), chest), forward)


def log_hands(label):
    # Releve seulement : main devant (+) ou derriere (-) le buste sur le premier clip cuit du pack.
    # La mesure fine est la tournee de la sandbox. Sans clip cuit ou sans roles : on saute.
    profile = CFG["source_profile"]
    chest_src = role(profile, "spine", -1)
    hips_src = role(profile, "hips")
    hands_src = (role(profile, "hand_l"), role(profile, "hand_r"))
    collars_src = (role(profile, "clavicle_l"), role(profile, "clavicle_r"))
    if not (chest_src and hips_src and all(hands_src) and all(collars_src)):
        unreal.log_warning("RETARGET mesure sautee : roles du profil source incomplets")
        return
    first = None
    for body in CFG["corps"]:
        for data in source_clips():
            name = str(data.asset_name)
            if unreal.EditorAssetLibrary.does_asset_exist(CFG["out_dir"] + "/" + body["prefixe"] + name):
                first = name
                break
        if first:
            break
    if first is None:
        unreal.log_warning("RETARGET mesure sautee : aucun clip cuit")
        return
    spec = [(CFG["clips_dir"] + "/" + first, chest_src, hips_src, hands_src, collars_src)]
    for body in CFG["corps"]:
        spec.append((CFG["out_dir"] + "/" + body["prefixe"] + first, "chestUpper", "pelvis", ("lHand", "rHand"), ("lCollar", "rCollar")))
    for path, chest, pelvis, hands, collars in spec:
        leaf = path.split("/")[-1]
        sequence = unreal.load_asset(path + "." + leaf)
        if not sequence:
            unreal.log_warning("RETARGET mesure absente %s" % path)
            continue
        for hand in hands:
            ahead = hand_ahead(sequence, chest, hand, collars[0], collars[1], 0.0, pelvis)
            unreal.log_warning("RETARGET main %s %s %s devant %.1f" % (label, leaf, hand, ahead))


def bake_named(clips, source_mesh, bodies):
    # bodies : [(corps, maillage, retargeter)] ; renvoie (nombre de copies par corps, nombre deplaces).
    made = []
    for body, mesh, retargeter in bodies:
        made.append(unreal.IKRetargetBatchOperation.duplicate_and_retarget(
            clips, source_mesh, mesh, retargeter, "", "", body["prefixe"], "", True))
    moved = 0
    for data in [item for group in made for item in group]:
        package = str(data.package_name)
        name = str(data.asset_name)
        dest = CFG["out_dir"] + "/" + name
        if package == dest or unreal.EditorAssetLibrary.rename_asset(package, dest):
            moved += 1
        else:
            unreal.log_error("RETARGET deplacement rate %s" % package)
    unreal.EditorLoadingAndSavingUtils.save_dirty_packages(True, True)
    return [len(group) for group in made], moved


def run():
    global CFG
    CFG = load_config(unreal.Paths.project_saved_dir())
    method = load_method(CFG["methode"])
    CFG["pose"] = method["pose"]
    mode = CFG["mode"]
    unreal.log_warning("RETARGET etape %s" % mode)
    suf = suffix(CFG)
    probe = unreal.load_asset(CFG["source_mesh"])
    if not probe:
        raise Refusal("maillage source introuvable %s" % CFG["source_mesh"])
    source_profile, _ = identify_names(bone_names(probe), probe.get_name())
    if source_profile["id"] != method["profil_source"]:
        raise Refusal("profil du maillage source %s (%s) different de profil_source %s de la methode %s" % (
            probe.get_name(), source_profile["id"], method["profil_source"], CFG["methode"]))
    CFG["source_profile"] = source_profile
    unreal.log_warning("RETARGET methode %s profil %s calques %s" % (CFG["methode"], source_profile["id"], method["pose"]))
    source_rig, source_mesh, source_chains = make_rig(source_rig_name(CFG), CFG["source_mesh"])
    bodies = []
    for body in CFG["corps"]:
        rig, mesh, _ = make_rig("IK_" + body["nom"] + suf, body["maillage"])
        retargeter = make_retargeter("RTG_" + body["nom"] + suf, source_rig, rig, source_mesh, mesh, body["nom"], source_chains)
        bodies.append((body, mesh, retargeter))
    if mode == "rigs":
        finish("RETARGET rigs prets")
        return
    names = [str(data.asset_name) for data in source_clips()]
    delete_copies(names)
    clips = source_clips()
    unreal.log_warning("RETARGET clips source %d" % len(clips))
    counts, moved = bake_named(clips, source_mesh, bodies)
    try:
        log_hands("final")
    except Exception as exc:
        unreal.log_warning("RETARGET mesure finale ratee %s" % exc)
    done = " ".join("%s %d" % (body["nom"].lower(), count) for (body, _, _), count in zip(bodies, counts))
    finish("RETARGET %s deplaces %d pose %s ajustee" % (done, moved, method["pose"]))


def main():
    try:
        run()
    except Refusal as exc:
        unreal.log_warning(tag("RETARGET echec : %s" % exc))
        unreal.SystemLibrary.quit_editor()
    except Exception as exc:
        unreal.log_warning(tag("RETARGET echec : %s" % exc))
        unreal.SystemLibrary.quit_editor()
        raise


if not os.environ.get("LP_RETARGET_NOMAIN"):
    main()