"""Offline skin/motion comparison. Never modifies the source FBXs or Unity assets.

Mixamo clips are skeleton-only. Their evaluated world bone matrices are applied
to the same-name bones of the original skinned model, preserving its bind pose.
This explicitly handles the differing bind/rest poses; it is not Unity Humanoid.
"""
import hashlib
import json
import math
import sys
from pathlib import Path

import bpy
import numpy as np
from mathutils import Matrix, Vector

REFERENCE_POINTS = None
GEOMETRY_CHECKS = {}


def coords(obj, evaluated=False):
    item = obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) if evaluated else obj
    data = item.to_mesh() if evaluated else item.data
    result = np.array([item.matrix_world @ v.co for v in data.vertices], dtype=np.float64)
    if evaluated:
        item.to_mesh_clear()
    return result


def imported(path):
    before = set(bpy.data.objects)
    if path.suffix.lower() == '.blend':
        bpy.ops.wm.open_mainfile(filepath=str(path))
    elif path.suffix.lower() == '.fbx':
        bpy.ops.import_scene.fbx(filepath=str(path))
    else:
        bpy.context.scene.render.fps = 30
        bpy.ops.import_scene.gltf(filepath=str(path))
    return [o for o in bpy.data.objects if o not in before]


def configure_render(height, center):
    scene = bpy.context.scene
    scene.render.engine = 'CYCLES'
    scene.cycles.samples = 16
    scene.render.resolution_x, scene.render.resolution_y = 512, 640
    scene.render.resolution_percentage = 100
    scene.view_settings.view_transform = 'AgX'
    world = bpy.data.worlds.new('AuditWorld')
    world.use_nodes = True
    world.node_tree.nodes['Background'].inputs[0].default_value = (.35, .38, .42, 1)
    world.node_tree.nodes['Background'].inputs[1].default_value = .7
    scene.world = world
    camera = bpy.data.objects.new('AuditCamera', bpy.data.cameras.new('AuditCamera'))
    scene.collection.objects.link(camera)
    scene.camera = camera
    camera.data.type = 'ORTHO'
    camera.data.ortho_scale = height*1.3
    for direction, power in [((2, -3, 4), 180), ((-3, 1, 2), 120)]:
        lamp = bpy.data.objects.new('AuditLight', bpy.data.lights.new('AuditLight', 'AREA'))
        scene.collection.objects.link(lamp)
        lamp.location = center + Vector(direction)*height
        lamp.rotation_euler = (center-lamp.location).to_track_quat('-Z', 'Y').to_euler()
        lamp.data.energy, lamp.data.size = power*height**2, height*3
    return camera


def audit(label, mesh_path, clip_path, out):
    global REFERENCE_POINTS
    bpy.ops.wm.read_factory_settings(use_empty=True)
    objects = imported(mesh_path)
    arm = next(o for o in objects if o.type == 'ARMATURE')
    mesh = max((o for o in objects if o.type == 'MESH'), key=lambda o: len(o.data.vertices))
    rest = coords(mesh)
    up = int(np.argmax(np.ptp(rest, axis=0)))
    normalized = rest[:, [0, 2, 1]].copy() if up == 1 else rest.copy()
    if up == 1:
        normalized[:, 1] *= -1
    lo, hi = normalized.min(0), normalized.max(0)
    height = float(hi[2]-lo[2])
    if REFERENCE_POINTS is None:
        REFERENCE_POINTS = normalized.copy()
    if str(mesh_path) not in GEOMETRY_CHECKS:
        matches=[]
        for sign in (1,-1):
            candidate=normalized.copy()
            candidate[:,:2] *= sign
            distances=[]
            for start in range(0,len(candidate),128):
                distance2=((candidate[start:start+128,None,:]-REFERENCE_POINTS[None,:,:])**2).sum(2)
                distances.extend(np.sqrt(distance2.min(1)).tolist())
            matches.append({'yaw_degrees': 0 if sign==1 else 180, 'max_nearest_vertex_distance_m': max(distances),
                            'mean_nearest_vertex_distance_m': float(np.mean(distances))})
        GEOMETRY_CHECKS[str(mesh_path)] = min(matches,key=lambda m:m['mean_nearest_vertex_distance_m'])
    groups = [g.name for g in mesh.vertex_groups]
    weights = np.zeros((len(rest), len(groups)))
    for v in mesh.data.vertices:
        for g in v.groups:
            weights[v.index, g.group] = g.weight
    edges = np.array([e.vertices[:] for e in mesh.data.edges])
    lengths = np.linalg.norm(rest[edges[:, 1]]-rest[edges[:, 0]], axis=1)
    valid = lengths > height*.002
    z = (normalized[:, 2]-lo[2])/height
    hand = (np.abs(normalized[:, 0]) > height*.255) & (z > .55) & (z < .67)
    neck = (z > .67) & (z < .78) & (np.abs(normalized[:, 0]) < height*.12)
    head = z > .78
    regions = {'all': np.ones(len(rest), dtype=bool), 'hands': hand, 'neck': neck, 'head': head}
    region_edges = {n: valid & mask[edges[:, 0]] & mask[edges[:, 1]] for n, mask in regions.items()}
    foot_groups = [i for i,n in enumerate(groups) if any(k in n for k in ('Foot', 'Toe', 'Shin', 'UpLeg')) or n.endswith('Leg')]
    head_groups = [i for i,n in enumerate(groups) if any(k in n for k in ('Head', 'Neck'))]
    def region_summary(mask):
        w = weights[mask]
        return {'vertices': int(mask.sum()), 'average_weights': {n: float(w[:,i].mean()) for i,n in enumerate(groups) if w[:,i].max() > .01},
                'max_foot_influence': float(w[:,foot_groups].sum(1).max()),
                'max_head_neck_influence': float(w[:,head_groups].sum(1).max())}
    def norm_point(p):
        return [p.x, -p.z, p.y] if up == 1 else list(p)
    bones = {b.name: {'head': norm_point(arm.matrix_world @ b.head_local), 'tail': norm_point(arm.matrix_world @ b.tail_local),
                       'parent': b.parent.name if b.parent else None} for b in arm.data.bones}
    result = {'label': label, 'mesh_source': str(mesh_path), 'clip_source': str(clip_path) if clip_path else None,
        'geometry_match_to_first_mixamo': GEOMETRY_CHECKS[str(mesh_path)],
        'mesh': {'vertices': len(rest), 'triangles': sum(len(p.vertices)-2 for p in mesh.data.polygons), 'height_m': height,
                 'bounds_normalized_z_up': [lo.tolist(), hi.tolist()], 'hands': region_summary(hand), 'head': region_summary(head),
                 'max_influences_over_0_00001': int((weights > .00001).sum(1).max()),
                 'max_influences_over_0_01': int((weights > .01).sum(1).max()),
                 'max_weight_discarded_if_top4': float(np.sort(weights,axis=1)[:,:-4].sum(1).max()),
                 'weight_sum_range': [float(weights.sum(1).min()), float(weights.sum(1).max())],
                 'rest_edge_lengths_sha256': hashlib.sha256(np.round(np.sort(lengths), 5).tobytes()).hexdigest()},
        'bones': bones}
    # Motion-independent LBS stress test: head yaw 30 degrees and both upper arms
    # down 70 degrees, in normalized world coordinates. Same angles for both rigs.
    # Descendants inherit their parent's deformation; joint pivots remain those
    # of each delivered rig, so this tests the combined bind placement + weights.
    canonical = Matrix.Rotation(math.pi/2,4,'X') if up == 1 else Matrix.Identity(4)
    bind = {b.name: canonical @ arm.matrix_world @ b.matrix_local for b in arm.data.bones}
    def stress(arm_down, head_yaw):
        deform = {}
        for b in arm.data.bones:
            inherited = deform[b.parent.name] if b.parent else Matrix.Identity(4)
            posed = inherited @ bind[b.name]
            suffix = b.name.split(':')[-1]
            angle, axis = 0, 'Z'
            if suffix == 'Head':
                angle = head_yaw
            elif suffix in ('LeftArm','RightArm'):
                angle, axis = arm_down*(1 if bind[b.name].translation.x > 0 else -1), 'Y'
            if angle:
                pivot = posed.translation
                posed = Matrix.Translation(pivot) @ Matrix.Rotation(math.radians(angle),4,axis) @ Matrix.Translation(-pivot) @ posed
            deform[b.name] = posed @ bind[b.name].inverted()
        points4 = np.column_stack((normalized,np.ones(len(normalized))))
        deformed = np.zeros((len(rest),3))
        for i,name in enumerate(groups):
            deformed += weights[:,i,None]*(points4 @ np.array(deform[name]).T)[:,:3]
        ratios = np.linalg.norm(deformed[edges[:,1]]-deformed[edges[:,0]],axis=1)/np.maximum(lengths,1e-12)
        return {n: {'max_edge_stretch': float(ratios[mask].max()), 'edges_over_2x': int((ratios[mask] > 2).sum())}
                for n,mask in region_edges.items()}
    result['controlled_lbs_stress'] = {'neutral': stress(0,0), 'arms_down_70_head_yaw_30': stress(70,30)}
    source = None
    if clip_path:
        objects = imported(clip_path)
        source = next(o for o in objects if o.type == 'ARMATURE')
        assert set(arm.pose.bones.keys()) == set(source.pose.bones.keys()), 'skeleton names differ'
    action_arm = source or arm
    action = action_arm.animation_data.action if action_arm.animation_data else None
    if not action:
        return result
    first, last = map(int, action.frame_range)
    fps = bpy.context.scene.render.fps/bpy.context.scene.render.fps_base
    samples, endposes, worldposes = [], [], []
    named = lambda suffix: next((b.name for b in action_arm.pose.bones if b.name == suffix or b.name.endswith(':'+suffix)), None)
    hips_name = named('Hips')
    tracks = {s: named(s) for s in ['Hips','Head','LeftFoot','RightFoot','LeftToeBase','RightToeBase']}
    for frame in range(first, last+1):
        bpy.context.scene.frame_set(frame)
        bpy.context.view_layer.update()
        if source:
            matrices = {b.name: arm.matrix_world.inverted() @ source.matrix_world @ b.matrix for b in source.pose.bones}
            for b in arm.pose.bones:
                b.matrix = matrices[b.name]
                bpy.context.view_layer.update()
        deformed = coords(mesh, True)
        ratios = np.linalg.norm(deformed[edges[:, 1]]-deformed[edges[:, 0]], axis=1)/np.maximum(lengths, 1e-12)
        worst = int(np.argmax(np.where(valid, ratios, 0)))
        values = {n: {'max': float(ratios[mask].max()) if mask.any() else None,
                      'p99': float(np.percentile(ratios[mask],99)) if mask.any() else None,
                      'over_2x': int((ratios[mask] > 2).sum())} for n,mask in region_edges.items()}
        feet = {s: [float(deformed[weights[:,i] > .25,2].min()) for i,n in enumerate(groups) if s+'Foot' in n and (weights[:,i] > .25).any()]
                for s in ['Left','Right']}
        samples.append({'frame': frame, 'regions': values,
                        'floor_min_z': float(deformed[:,2].min()), 'feet_min_z': feet,
                        'hips_world': list(action_arm.matrix_world @ action_arm.pose.bones[hips_name].head),
                        'worst_edge': {'ids': edges[worst].tolist(), 'rest': normalized[edges[worst]].tolist(),
                                       'ratio': float(ratios[worst]), 'weights': [{n: float(weights[v,i]) for i,n in enumerate(groups) if weights[v,i]>.00001} for v in edges[worst]]}})
        poses = {b.name: b.matrix_basis.to_quaternion().copy() for b in action_arm.pose.bones}
        world = {n: list(action_arm.matrix_world @ action_arm.pose.bones[b].head) for n,b in tracks.items() if b}
        worldposes.append(world)
        if frame in (first, last):
            endposes.append(poses)
    seam = {n: min(float(endposes[0][n].rotation_difference(endposes[-1][n]).angle),2*math.pi-float(endposes[0][n].rotation_difference(endposes[-1][n]).angle))*180/math.pi for n in endposes[0]}
    result['animation'] = {'range': [first,last], 'fps': fps, 'duration_between_endpoints_s': (last-first)/fps,
        'samples': samples, 'joint_tracks': worldposes, 'loop_local_rotation_gap_degrees': seam,
        'summary': {n: {'max_edge_stretch': max(s['regions'][n]['max'] for s in samples if s['regions'][n]['max'] is not None),
                        'max_p99_edge_stretch': max(s['regions'][n]['p99'] for s in samples if s['regions'][n]['p99'] is not None),
                        'max_edges_over_2x': max(s['regions'][n]['over_2x'] for s in samples)} for n in regions}}
    print('METRICS', label, result['animation']['summary'], 'SEAM', sorted(seam.items(),key=lambda x:x[1],reverse=True)[:4])
    if label.startswith(('Mixamo','TokenRig41','TokenRigLocal')) and '--metrics-only' not in sys.argv:
        mid = first+(last-first)//3
        # Re-evaluate one representative source frame before rendering.
        bpy.context.scene.frame_set(mid)
        bpy.context.view_layer.update()
        matrices = {b.name: arm.matrix_world.inverted() @ source.matrix_world @ b.matrix for b in source.pose.bones}
        for b in arm.pose.bones:
            b.matrix = matrices[b.name]
            bpy.context.view_layer.update()
        points = coords(mesh, True)
        center = Vector(((points[:,0].min()+points[:,0].max())/2, (points[:,1].min()+points[:,1].max())/2, (points[:,2].min()+points[:,2].max())/2))
        camera = configure_render(height, center)
        for view, direction in [('front',(0,-3,.12)), ('side',(3,0,.12)), ('back',(0,3,.12))]:
            camera.location = center+Vector(direction)*height
            camera.rotation_euler = (center-camera.location).to_track_quat('-Z','Y').to_euler()
            bpy.context.scene.render.filepath = str(out/(label.replace(' ','_')+'_'+view+'.png'))
            bpy.ops.render.render(write_still=True)
    return result


def main():
    root, out = map(Path, sys.argv[sys.argv.index('--')+1:][:2])
    out.mkdir(parents=True, exist_ok=True)
    base = root/'AIClient/Assets/Bundles/Characters/Human/test/Boy_mixamo'
    soma = root/'Tools/output/campguide_studio/rig_soma'
    jobs = [('Mixamo '+name, base/'CampGuideStudio.fbx', base/('CampGuideStudio@'+name+'.fbx')) for name in ['Walking','Breathing Idle','Happy Idle']]
    jobs += [('SOMA '+name, soma/'kimodo'/name/'baked.glb',None) for name in ['Idle','Walk']]
    results=[]
    for label,mesh,clip in jobs:
        result=audit(label,mesh,clip,out)
        results.append(result)
        (out/'comparison.json').write_text(json.dumps(results,indent=2),encoding='utf-8')
    print('AUDIT_DONE',out/'comparison.json')


if __name__ == '__main__':
    main()