"""Conservatively remove redundant keys from Blender-imported BVH actions.

Usage:
  1. Import a BVH into Blender.
  2. Select its armature.
  3. Open this file in Blender's Scripting workspace and press Run Script.

The original Action is copied to an Action named ``*_before_key_cleanup`` and
given a fake user before any keys are changed. Position and rotation channels
are simplified as synchronized groups, never as unrelated XYZ curves.

Tested with Blender 5.2's layered Action API and compatible with the legacy
Action.fcurves API used by Blender 3.x/early 4.x.
"""

import math

import bpy
from mathutils import Vector


# Conservative defaults for BVHs exported by the SL converter at Scale 1.
# One Blender unit is one metre with that import profile.
POSITION_TOLERANCE = 0.00005          # 0.05 mm, local translation
ROTATION_TOLERANCE_DEGREES = 0.02    # local quaternion angle
SCALE_TOLERANCE = 0.00001
OTHER_CHANNEL_TOLERANCE = 0.000001
CREATE_BACKUP_ACTION = True
FORCE_LINEAR_INTERPOLATION = True


def action_fcurves(obj):
    """Return the F-curves belonging to obj's active Action slot."""
    animation_data = obj.animation_data
    action = animation_data.action if animation_data else None
    if action is None:
        return action, []

    # Blender 3.x and the legacy Action API.
    legacy_fcurves = getattr(action, "fcurves", None)
    if legacy_fcurves is not None:
        return action, list(legacy_fcurves)

    # Blender 4.4+ layered Actions. Limit cleanup to this object's slot.
    slot = getattr(animation_data, "action_slot", None)
    if slot is None:
        return action, []
    fcurves = []
    for layer in action.layers:
        for strip in layer.strips:
            if getattr(strip, "type", None) != "KEYFRAME":
                continue
            try:
                channelbag = strip.channelbag(slot)
            except RuntimeError:
                channelbag = None
            if channelbag is not None:
                fcurves.extend(channelbag.fcurves)
    return action, fcurves


def animated_armatures():
    selected = [obj for obj in bpy.context.selected_objects if obj.type == "ARMATURE" and obj.animation_data]
    if selected:
        return selected
    active = bpy.context.view_layer.objects.active
    if active and active.type == "ARMATURE" and active.animation_data:
        return [active]
    return [obj for obj in bpy.data.objects if obj.type == "ARMATURE" and obj.animation_data and obj.animation_data.action]


def property_owner(obj, data_path):
    owner_path, _, _property_name = data_path.rpartition(".")
    if not owner_path:
        return obj


def quaternion_multiply(left, right):
    lw, lx, ly, lz = left
    rw, rx, ry, rz = right
    return (
        lw * rw - lx * rx - ly * ry - lz * rz,
        lw * rx + lx * rw + ly * rz - lz * ry,
        lw * ry - lx * rz + ly * rw + lz * rx,
        lw * rz + lx * ry - ly * rx + lz * rw,
    )


def normalized_quaternion(values):
    length = math.sqrt(sum(value * value for value in values))
    return tuple(value / length for value in values) if length > 1e-15 else (1.0, 0.0, 0.0, 0.0)


def euler_quaternion(values, order):
    angles = {"X": values[0], "Y": values[1], "Z": values[2]}
    rotation = (1.0, 0.0, 0.0, 0.0)
    for axis in order:
        half = angles[axis] * 0.5
        sine = math.sin(half)
        cosine = math.cos(half)
        axis_rotation = {
            "X": (cosine, sine, 0.0, 0.0),
            "Y": (cosine, 0.0, sine, 0.0),
            "Z": (cosine, 0.0, 0.0, sine),
        }[axis]
        rotation = quaternion_multiply(axis_rotation, rotation)
    return normalized_quaternion(rotation)


def quaternion_angle_error(actual, predicted):
    left = normalized_quaternion(actual)
    right = normalized_quaternion(predicted)
    dot = abs(sum(a * b for a, b in zip(left, right)))
    return 2.0 * math.acos(max(-1.0, min(1.0, dot)))
    try:
        return obj.path_resolve(owner_path)
    except (ValueError, TypeError):
        return obj


def metric_for_group(obj, data_path):
    property_name = data_path.rsplit(".", 1)[-1]
    if property_name == "rotation_euler":
        owner = property_owner(obj, data_path)
        order = getattr(owner, "rotation_mode", "XYZ")
        if order not in {"XYZ", "XZY", "YXZ", "YZX", "ZXY", "ZYX"}:
            order = "XYZ"

        def error(actual, predicted):
            return quaternion_angle_error(euler_quaternion(actual[:3], order), euler_quaternion(predicted[:3], order))

        return "rotation", math.radians(ROTATION_TOLERANCE_DEGREES), error

    if property_name == "rotation_quaternion":
        def error(actual, predicted):
            return quaternion_angle_error(actual[:4], predicted[:4])

        return "rotation", math.radians(ROTATION_TOLERANCE_DEGREES), error

    if property_name == "location":
        return "position", POSITION_TOLERANCE, lambda actual, predicted: (Vector(actual[:3]) - Vector(predicted[:3])).length
    if property_name == "scale":
        return "scale", SCALE_TOLERANCE, lambda actual, predicted: max(abs(a - b) for a, b in zip(actual, predicted))
    return "other", OTHER_CHANNEL_TOLERANCE, lambda actual, predicted: max(abs(a - b) for a, b in zip(actual, predicted))


def sampled_group(fcurves):
    frames = sorted({round(point.co.x, 8) for fcurve in fcurves for point in fcurve.keyframe_points})
    width = max(fcurve.array_index for fcurve in fcurves) + 1
    by_index = {fcurve.array_index: fcurve for fcurve in fcurves}
    values = []
    for frame in frames:
        values.append([by_index[index].evaluate(frame) if index in by_index else 0.0 for index in range(width)])
    return frames, values


def interpolated(left, right, amount):
    return [a + (b - a) * amount for a, b in zip(left, right)]


def simplify_indices(frames, values, tolerance, error_metric):
    """Vector RDP: retain the minimum keys needed within the supplied error."""
    if len(frames) <= 2:
        return set(range(len(frames)))
    keep = {0, len(frames) - 1}
    pending = [(0, len(frames) - 1)]
    while pending:
        left, right = pending.pop()
        duration = frames[right] - frames[left]
        if right <= left + 1 or abs(duration) < 1e-12:
            continue
        worst_index = None
        worst_error = tolerance
        for index in range(left + 1, right):
            amount = (frames[index] - frames[left]) / duration
            error = error_metric(values[index], interpolated(values[left], values[right], amount))
            if error > worst_error:
                worst_error = error
                worst_index = index
        if worst_index is not None:
            keep.add(worst_index)
            pending.append((left, worst_index))
            pending.append((worst_index, right))
    return keep


def clean_group(obj, data_path, fcurves):
    frames, original_values = sampled_group(fcurves)
    kind, tolerance, error_metric = metric_for_group(obj, data_path)
    keep_indices = simplify_indices(frames, original_values, tolerance, error_metric)
    keep_frames = {frames[index] for index in keep_indices}
    keys_before = sum(len(fcurve.keyframe_points) for fcurve in fcurves)

    for fcurve in fcurves:
        for point in list(fcurve.keyframe_points):
            if round(point.co.x, 8) not in keep_frames:
                fcurve.keyframe_points.remove(point, fast=True)
        if FORCE_LINEAR_INTERPOLATION:
            for point in fcurve.keyframe_points:
                point.interpolation = "LINEAR"
        fcurve.update()

    # Measure the actual reduced F-curves against every original keyed frame.
    width = max(fcurve.array_index for fcurve in fcurves) + 1
    by_index = {fcurve.array_index: fcurve for fcurve in fcurves}
    maximum_error = 0.0
    for frame, original in zip(frames, original_values):
        reduced = [by_index[index].evaluate(frame) if index in by_index else 0.0 for index in range(width)]
        maximum_error = max(maximum_error, error_metric(original, reduced))

    return {
        "kind": kind,
        "keys_before": keys_before,
        "keys_after": sum(len(fcurve.keyframe_points) for fcurve in fcurves),
        "maximum_error": maximum_error,
        "tolerance": tolerance,
    }


def clean_armature(obj):
    action, fcurves = action_fcurves(obj)
    if action is None or not fcurves:
        raise RuntimeError(f"{obj.name}: no editable F-curves found in the active Action slot")

    if CREATE_BACKUP_ACTION:
        backup = action.copy()
        backup.name = f"{action.name}_before_key_cleanup"
        backup.use_fake_user = True

    groups = {}
    for fcurve in fcurves:
        groups.setdefault(fcurve.data_path, []).append(fcurve)

    result = {
        "object": obj.name,
        "action": action.name,
        "groups": len(groups),
        "keys_before": 0,
        "keys_after": 0,
        "max_position_error": 0.0,
        "max_rotation_error": 0.0,
    }
    for data_path, group in groups.items():
        cleaned = clean_group(obj, data_path, group)
        result["keys_before"] += cleaned["keys_before"]
        result["keys_after"] += cleaned["keys_after"]
        if cleaned["kind"] == "position":
            result["max_position_error"] = max(result["max_position_error"], cleaned["maximum_error"])
        elif cleaned["kind"] == "rotation":
            result["max_rotation_error"] = max(result["max_rotation_error"], cleaned["maximum_error"])
    return result


def main():
    targets = animated_armatures()
    if not targets:
        raise RuntimeError("No animated armature found. Import a BVH and select its armature first.")

    results = [clean_armature(obj) for obj in targets]
    for result in results:
        removed = result["keys_before"] - result["keys_after"]
        percent = removed / max(1, result["keys_before"]) * 100
        print(
            "BVH KEY CLEANUP: "
            f"{result['object']} / {result['action']} — removed {removed:,} of {result['keys_before']:,} keys "
            f"({percent:.1f}%), kept {result['keys_after']:,}; "
            f"max local position error {result['max_position_error'] * 1000:.4f} mm, "
            f"max local rotation error {math.degrees(result['max_rotation_error']):.5f}°"
        )

    if not bpy.app.background:
        total_removed = sum(result["keys_before"] - result["keys_after"] for result in results)

        def draw_message(self, _context):
            self.layout.label(text=f"Removed {total_removed:,} redundant keys from {len(results)} armature(s).")
            self.layout.label(text="Original Action backup(s) are preserved in the Blender file.")

        bpy.context.window_manager.popup_menu(draw_message, title="BVH key cleanup complete", icon="INFO")
    return results


if __name__ == "__main__":
    main()