import json
import math

# Load robot.json
with open(r'robot.json') as f:
    robot = json.load(f)

# Extract markers by link
markers_by_link = {}
for link_name, link_data in robot['links'].items():
    markers = link_data.get('markers', [])
    if markers:
        markers_by_link[link_name] = markers

print("=" * 70)
print("MARKER DISTRIBUTION BY LINK")
print("=" * 70)
for link, markers in markers_by_link.items():
    print(f"\n{link}: {len(markers)} markers")
    for m in markers:
        pos = m['position']
        print(f"  ID {m['id']:3d}: pos={pos}")

# Check rigid body constraints: distances within each link
print("\n" + "=" * 70)
print("RIGID BODY CONSTRAINT CHECK (Arm1, Ellbow, Arm2)")
print("=" * 70)

for link_name in ['Arm1', 'Ellbow', 'Arm2']:
    if link_name not in markers_by_link:
        continue
    
    markers = markers_by_link[link_name]
    positions = [(m['position'][0], m['position'][1], m['position'][2]) for m in markers]
    
    print(f"\n{link_name}:")
    if len(positions) > 1:
        for i in range(len(positions)):
            for j in range(i+1, len(positions)):
                dx = positions[i][0] - positions[j][0]
                dy = positions[i][1] - positions[j][1]
                dz = positions[i][2] - positions[j][2]
                dist = math.sqrt(dx*dx + dy*dy + dz*dz)
                print(f"  Marker {markers[i]['id']:3d} <-> {markers[j]['id']:3d}: {dist:7.2f} mm")
    else:
        print(f"  Only {len(positions)} marker(s)")

# Check X-distances between links
print("\n" + "=" * 70)
print("X-POSITION CONSTRAINT CHECK (Link relationships)")
print("=" * 70)

# Get marker positions per link
link_x_positions = {}
for link_name in ['Board', 'Base', 'Arm1', 'Ellbow', 'Arm2']:
    if link_name in markers_by_link:
        x_vals = [m['position'][0] for m in markers_by_link[link_name]]
        link_x_positions[link_name] = {
            'min_x': min(x_vals),
            'max_x': max(x_vals),
            'count': len(x_vals)
        }

print("\nX-position ranges by link (local coords):")
for link, data in link_x_positions.items():
    print(f"  {link:10s}: x in [{data['min_x']:7.2f}, {data['max_x']:7.2f}] mm  ({data['count']} markers)")

# Check X-distance between Arm1 and Ellbow
print("\nKey constraint analysis:")
arm1_x = [m['position'][0] for m in markers_by_link['Arm1']]
ellbow_x = [m['position'][0] for m in markers_by_link['Ellbow']]
print(f"  Arm1 local X: {arm1_x}")
print(f"  Ellbow local X: {ellbow_x}")
print(f"  => BOTH rotate around X-axis, so X-spread within each link is FIXED")
print(f"  => X-distance between links is FIXED in world (different local X values)")

# Check Arm2 structure
print("\n" + "=" * 70)
print("ARM2 KINEMATIC CHECK (sin(a) dependency)")
print("=" * 70)

arm2_markers = markers_by_link['Arm2']
print(f"\nArm2 has {len(arm2_markers)} markers")
print("Arm2 marker positions (local coords, before 'a' rotation around Y):")
for m in arm2_markers:
    pos = m['position']
    print(f"  ID {m['id']:3d}: x={pos[0]:8.2f}, y={pos[1]:8.2f}, z={pos[2]:8.2f}")

print("\nWhen Arm2 rotates around Y-axis (variable 'a'):")
print("  Rotation matrix Ry(a) acts on [x_local, y_local, z_local]:")
print("  X_world = 90 + x_local * cos(a) - z_local * sin(a)")
print("  Y_world = y_local (unchanged)")
print("  Z_world = x_local * sin(a) + z_local * cos(a)")
print("\n=> X_world DEPENDS on sin(a) and z_local values!")

# Verify: X differences between Arm2 markers
arm2_x_vals = [m['position'][0] for m in arm2_markers]
arm2_z_vals = [m['position'][2] for m in arm2_markers]
print(f"\nArm2 local X values: {sorted(set(arm2_x_vals))}")
print(f"Arm2 local Z values: {sorted(set(arm2_z_vals))}")
if len(set(arm2_z_vals)) > 1:
    print(f"=> Multiple Z values present: X_world will differ by sin(a) contributions")

print("\n" + "=" * 70)
print("SUMMARY: WHICH CONSTRAINTS MAKE SENSE?")
print("=" * 70)
print("""
1. RIGID BODY DISTANCES within Arm1/Ellbow/Arm2:
   ✓ VALID - Fixed distances between markers on same rigid link
   ✓ HELPS with: Preventing deformation, reducing degrees of freedom
   
2. X-DISTANCES between links (Arm1, Ellbow, Base):
   ✓ VALID - Both rotate around X-axis, so relative X is fixed
   ✓ HELPS with: Preventing sliding along links
   
3. X-POSITION of Arm2 dependent on sin(a):
   ⚠ PARTIALLY VALID - Different markers have different dependencies
   ✓ HELPS with: Constraining the a variable from Z-spread observation
   
IMPLEMENTATION RECOMMENDATION:
- Start with constraints 1 & 2 (rigid body + inter-link X)
- Use constraint 3 as a sanity check on 'a' estimation
""")