import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from pathlib import Path
import re


def parse_trc(trc_path):
   """Parse a .trc file and return marker names, frame data, and metadata."""
   with open(trc_path, 'r') as f:
       lines = f.readlines()


   meta_keys = lines[2].strip().split('\t')
   meta_vals = lines[3].strip().split('\t')
   metadata = dict(zip(meta_keys, meta_vals))


   marker_line = lines[3].strip().split('\t')


   header_lines = 0
   for i, line in enumerate(lines):
       if line.strip() and not line.startswith(('PathFileType', 'DataRate',
                                                 'Frame', '\t')):
           try:
               float(line.strip().split('\t')[0])
               header_lines = i
               break
           except ValueError:
               continue


   raw_markers = lines[3].strip().split('\t')
   markers = [m for m in raw_markers if m and m not in ('Frame#', 'Time')]


   marker_names = []
   for m in markers:
       if m and (not marker_names or m != marker_names[-1]):
           marker_names.append(m)


   data_lines = lines[header_lines:]
   data = []
   for line in data_lines:
       vals = line.strip().split('\t')
       if len(vals) > 2:
           try:
               row = [float(v) if v else np.nan for v in vals]
               data.append(row)
           except ValueError:
               continue


   data = np.array(data)
   return marker_names, data, metadata




trc_files = sorted((work_dir / "pose-3d").glob("*.trc"))
print(f"📊 Found {len(trc_files)} TRC file(s):")
for f in trc_files:
   print(f"   {f.name}")


trc_file = None
for f in trc_files:
   if 'filt' in f.name.lower() and 'augm' not in f.name.lower():
       trc_file = f
       break
if trc_file is None and trc_files:
   trc_file = trc_files[0]


if trc_file:
   print(f"\n📈 Visualizing: {trc_file.name}")
   marker_names, data, metadata = parse_trc(trc_file)
   print(f"   Markers: {len(marker_names)}")
   print(f"   Frames:  {data.shape[0]}")
   print(f"   Marker names: {marker_names[:10]}{'...' if len(marker_names) > 10 else ''}")


   frames = data[:, 0].astype(int) if data.shape[1] > 0 else []
   times = data[:, 1] if data.shape[1] > 1 else []
   coords = data[:, 2:]


   n_markers = len(marker_names)


   mid_frame = len(data) // 2
   fig = plt.figure(figsize=(16, 6))


   ax1 = fig.add_subplot(131, projection='3d')
   xs = coords[mid_frame, 0::3][:n_markers]
   ys = coords[mid_frame, 1::3][:n_markers]
   zs = coords[mid_frame, 2::3][:n_markers]


   ax1.scatter(xs, ys, zs, c="dodgerblue", s=40, alpha=0.8, edgecolors="navy")
   for i, name in enumerate(marker_names[:len(xs)]):
       if i % 3 == 0:
           ax1.text(xs[i], ys[i], zs[i], f' {name}', fontsize=6, alpha=0.7)


   ax1.set_xlabel('X (m)')
   ax1.set_ylabel('Y (m)')
   ax1.set_zlabel('Z (m)')
   ax1.set_title(f'3D Keypoints (Frame {int(frames[mid_frame])})', fontsize=10)


   ax2 = fig.add_subplot(132)
   key_markers = ['RAnkle', 'LAnkle', 'RWrist', 'LWrist', 'Nose']
   colors_map = {'RAnkle': 'red', 'LAnkle': 'blue', 'RWrist': 'orange',
                 'LWrist': 'green', 'Nose': 'purple'}


   for mkr in key_markers:
       if mkr in marker_names:
           idx = marker_names.index(mkr)
           z_col = idx * 3 + 2
           if z_col < coords.shape[1]:
               ax2.plot(times, coords[:, z_col],
                        label=mkr, color=colors_map.get(mkr, 'gray'),
                        linewidth=1.2, alpha=0.8)


   ax2.set_xlabel('Time (s)')
   ax2.set_ylabel('Z position (m)')
   ax2.set_title('Vertical Trajectories', fontsize=10)
   ax2.legend(fontsize=8, loc="best")
   ax2.grid(True, alpha=0.3)


   ax3 = fig.add_subplot(133)
   if len(times) > 1:
       dt = np.diff(times)
       dt[dt == 0] = 1e-6
       for mkr in ['RAnkle', 'RWrist']:
           if mkr in marker_names:
               idx = marker_names.index(mkr)
               x_col, y_col, z_col = idx * 3, idx * 3 + 1, idx * 3 + 2
               if z_col < coords.shape[1]:
                   dx = np.diff(coords[:, x_col])
                   dy = np.diff(coords[:, y_col])
                   dz = np.diff(coords[:, z_col])
                   speed = np.sqrt(dx**2 + dy**2 + dz**2) / dt
                   speed = np.clip(speed, 0, np.nanpercentile(speed, 99))
                   ax3.plot(times[1:], speed, label=mkr,
                            color=colors_map.get(mkr, 'gray'),
                            linewidth=0.8, alpha=0.7)


       ax3.set_xlabel('Time (s)')
       ax3.set_ylabel('Speed (m/s)')
       ax3.set_title('Marker Speeds (Quality Check)', fontsize=10)
       ax3.legend(fontsize=8)
       ax3.grid(True, alpha=0.3)


   plt.tight_layout()
   plt.savefig(work_dir / 'trajectory_analysis.png', dpi=150, bbox_inches="tight")
   plt.show()
   print("✅ Trajectory plots saved to trajectory_analysis.png")


else:
   print("⚠ No TRC file found to visualize.")