tests/test_gaze_pipeline.py

"""
Diagnostic test for the full gaze pipeline:
  calibration → predict → fusion → focus decision

Tests that looking at screen center reads as focused,
and looking away reads as not focused.
"""
import math
import sys
import os
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from models.gaze_calibration import GazeCalibration, DEFAULT_TARGETS
from models.gaze_eye_fusion import GazeEyeFusion


def _make_landmarks_with_ear(ear_value=0.28):
    """Create a minimal 478-landmark array with given EAR.
    Only the EAR indices (6 per eye) and iris indices need real values."""
    lm = np.full((478, 3), 0.5, dtype=np.float32)
    # Left eye EAR landmarks [33, 160, 158, 133, 153, 145]
    # p1=33, p2=160, p3=158, p4=133, p5=153, p6=145
    # EAR = (|p2-p6| + |p3-p5|) / (2 * |p1-p4|)
    # Set them so EAR ≈ ear_value with horizontal dist = 0.1
    h_dist = 0.1
    v_dist = ear_value * h_dist  # EAR = v_dist / h_dist when both verticals equal
    lm[33] = [0.4, 0.5, 0]   # p1 outer
    lm[133] = [0.5, 0.5, 0]  # p4 inner
    lm[160] = [0.45, 0.5 - v_dist/2, 0]  # p2 top
    lm[145] = [0.45, 0.5 + v_dist/2, 0]  # p6 bottom
    lm[158] = [0.45, 0.5 - v_dist/2, 0]  # p3 top
    lm[153] = [0.45, 0.5 + v_dist/2, 0]  # p5 bottom
    # Right eye — mirror
    lm[362] = [0.6, 0.5, 0]
    lm[263] = [0.5, 0.5, 0]
    lm[385] = [0.55, 0.5 - v_dist/2, 0]
    lm[380] = [0.55, 0.5 + v_dist/2, 0]
    lm[387] = [0.55, 0.5 - v_dist/2, 0]
    lm[373] = [0.55, 0.5 + v_dist/2, 0]
    return lm


def simulate_calibration(noise_std=0.01):
    """Simulate a 9-point calibration where the user looks at each target.

    For each target (screen_x, screen_y), we generate synthetic gaze angles:
      yaw ≈ (screen_x - 0.5) * 0.7 radians  (maps 0..1 to roughly ±20°)
      pitch ≈ (screen_y - 0.5) * 0.5 radians (maps 0..1 to roughly ±15°)
    Plus some noise to simulate real jitter.
    """
    cal = GazeCalibration()

    # Simulate gaze angle for a given screen target
    def target_to_gaze(tx, ty):
        yaw = (tx - 0.5) * 0.7   # ~±20° across screen width
        pitch = (ty - 0.5) * 0.5  # ~±14° across screen height
        return yaw, pitch

    for i, (tx, ty) in enumerate(DEFAULT_TARGETS):
        base_yaw, base_pitch = target_to_gaze(tx, ty)
        n_samples = 45 if i == 0 else 30  # center gets more
        for _ in range(n_samples):
            yaw = base_yaw + np.random.normal(0, noise_std)
            pitch = base_pitch + np.random.normal(0, noise_std)
            cal.collect_sample(yaw, pitch)
        cal.advance()

    ok = cal.fit()
    return cal, ok


def test_calibration_accuracy():
    """Test that calibration maps screen positions correctly."""
    print("\n" + "="*60)
    print("TEST 1: Calibration accuracy")
    print("="*60)

    np.random.seed(42)
    cal, ok = simulate_calibration(noise_std=0.008)
    assert ok, "Calibration fit failed!"
    print(f"  Calibration fitted: {ok}")

    # Test prediction at each target
    max_error = 0
    for tx, ty in DEFAULT_TARGETS:
        yaw = (tx - 0.5) * 0.7
        pitch = (ty - 0.5) * 0.5
        px, py = cal.predict(yaw, pitch)
        err = math.sqrt((px - tx)**2 + (py - ty)**2)
        max_error = max(max_error, err)
        status = "OK" if err < 0.1 else "BAD"
        print(f"  Target ({tx:.2f},{ty:.2f}) → Predicted ({px:.3f},{py:.3f}) "
              f"error={err:.4f} [{status}]")

    print(f"\n  Max error: {max_error:.4f}")
    assert max_error < 0.15, f"Calibration error too high: {max_error:.4f}"
    print("  PASSED")


def test_fusion_focused_at_center():
    """Test that looking at screen center = focused."""
    print("\n" + "="*60)
    print("TEST 2: Looking at screen center → FOCUSED")
    print("="*60)

    np.random.seed(42)
    cal, ok = simulate_calibration()
    assert ok

    fusion = GazeEyeFusion(cal)
    lm = _make_landmarks_with_ear(0.28)  # eyes open

    # Looking at center: yaw≈0, pitch≈0
    center_yaw = (0.5 - 0.5) * 0.7  # = 0
    center_pitch = (0.5 - 0.5) * 0.5  # = 0

    # Run a few frames to let EMA settle
    for i in range(10):
        result = fusion.update(center_yaw, center_pitch, lm)

    print(f"  gaze_x={result['gaze_x']:.3f} gaze_y={result['gaze_y']:.3f}")
    print(f"  on_screen={result['on_screen']}")
    print(f"  focus_score={result['focus_score']:.3f} (threshold=0.42)")
    print(f"  focused={result['focused']}")
    print(f"  ear={result['ear']:.4f}")

    assert result["on_screen"], "Should be on screen!"
    assert result["focused"], f"Should be focused! score={result['focus_score']}"
    assert 0.35 < result["gaze_x"] < 0.65, f"gaze_x should be near 0.5, got {result['gaze_x']}"
    assert 0.35 < result["gaze_y"] < 0.65, f"gaze_y should be near 0.5, got {result['gaze_y']}"
    print("  PASSED")


def test_fusion_focused_at_edges():
    """Test that looking at screen edges still = focused."""
    print("\n" + "="*60)
    print("TEST 3: Looking at screen edges → FOCUSED")
    print("="*60)

    np.random.seed(42)
    cal, ok = simulate_calibration()
    assert ok
    lm = _make_landmarks_with_ear(0.28)

    edge_targets = [
        (0.15, 0.15, "top-left"),
        (0.85, 0.15, "top-right"),
        (0.15, 0.85, "bottom-left"),
        (0.85, 0.85, "bottom-right"),
        (0.5, 0.15, "top-center"),
        (0.5, 0.85, "bottom-center"),
    ]

    all_pass = True
    for tx, ty, label in edge_targets:
        fusion = GazeEyeFusion(cal)
        yaw = (tx - 0.5) * 0.7
        pitch = (ty - 0.5) * 0.5
        for _ in range(10):
            result = fusion.update(yaw, pitch, lm)

        status = "PASS" if result["focused"] else "FAIL"
        if not result["focused"]:
            all_pass = False
        print(f"  {label:15s} → gaze=({result['gaze_x']:.3f},{result['gaze_y']:.3f}) "
              f"on_screen={result['on_screen']} score={result['focus_score']:.3f} "
              f"[{status}]")

    assert all_pass, "Some edge positions reported unfocused!"
    print("  PASSED")


def test_fusion_unfocused_off_screen():
    """Test that looking far away = not focused."""
    print("\n" + "="*60)
    print("TEST 4: Looking far off screen → NOT FOCUSED")
    print("="*60)

    np.random.seed(42)
    cal, ok = simulate_calibration()
    assert ok
    lm = _make_landmarks_with_ear(0.28)

    off_screen_targets = [
        (2.0, 0.5, "far right"),
        (-1.0, 0.5, "far left"),
        (0.5, 2.0, "far down"),
        (0.5, -1.0, "far up"),
    ]

    all_pass = True
    for tx, ty, label in off_screen_targets:
        fusion = GazeEyeFusion(cal)
        yaw = (tx - 0.5) * 0.7
        pitch = (ty - 0.5) * 0.5
        for _ in range(20):
            result = fusion.update(yaw, pitch, lm)

        status = "PASS" if not result["focused"] else "FAIL"
        if result["focused"]:
            all_pass = False
        print(f"  {label:15s} → gaze=({result['gaze_x']:.3f},{result['gaze_y']:.3f}) "
              f"on_screen={result['on_screen']} score={result['focus_score']:.3f} "
              f"[{status}]")

    assert all_pass, "Some off-screen positions reported focused!"
    print("  PASSED")


def test_fusion_with_closed_eyes():
    """Test that sustained closed eyes = not focused, but brief blinks are OK."""
    print("\n" + "="*60)
    print("TEST 5: Sustained closed eyes → NOT FOCUSED, brief blink → still FOCUSED")
    print("="*60)

    np.random.seed(42)
    cal, ok = simulate_calibration()
    assert ok

    lm_closed = _make_landmarks_with_ear(0.10)  # eyes almost closed
    lm_open = _make_landmarks_with_ear(0.28)

    # 5a: Brief blink (2 frames closed) should NOT trigger unfocused
    fusion = GazeEyeFusion(cal)
    for _ in range(8):
        fusion.update(0, 0, lm_open)
    for _ in range(2):  # 2-frame blink
        result = fusion.update(0, 0, lm_closed)
    print(f"  Brief blink (2 frames): focused={result['focused']} score={result['focus_score']:.3f}")
    assert result["focused"], "Brief blink should NOT trigger unfocused!"

    # 5b: Sustained closure (6+ frames) SHOULD trigger unfocused
    fusion2 = GazeEyeFusion(cal)
    for _ in range(10):
        result2 = fusion2.update(0, 0, lm_closed)

    print(f"  Sustained closure (10 frames): focused={result2['focused']} score={result2['focus_score']:.3f}")
    assert not result2["focused"], f"Sustained closed eyes should be unfocused! score={result2['focus_score']}"
    print("  PASSED")


def test_l2cs_cosine_scoring():
    """Test the L2CSPipeline cosine scoring directly."""
    print("\n" + "="*60)
    print("TEST 6: L2CS cosine scoring (no calibration)")
    print("="*60)

    YAW_THRESHOLD = 22.0
    PITCH_THRESHOLD = 20.0

    test_angles = [
        (0, 0, "dead center"),
        (5, 3, "slightly off"),
        (10, 8, "moderate off"),
        (15, 12, "near edge"),
        (20, 18, "at threshold"),
        (25, 22, "beyond threshold"),
        (35, 30, "way off"),
    ]

    for yaw_deg, pitch_deg, label in test_angles:
        yaw_t = min(yaw_deg / YAW_THRESHOLD, 1.0)
        pitch_t = min(pitch_deg / PITCH_THRESHOLD, 1.0)
        yaw_score = 0.5 * (1.0 + math.cos(math.pi * yaw_t))
        pitch_score = 0.5 * (1.0 + math.cos(math.pi * pitch_t))
        gaze_score = 0.55 * yaw_score + 0.45 * pitch_score
        focused = gaze_score >= 0.52

        print(f"  yaw={yaw_deg:3d}° pitch={pitch_deg:3d}° → "
              f"score={gaze_score:.3f} focused={focused} [{label}]")

    print("  (informational — no assertion)")


def test_derotation_consistency():
    """Test that derotation produces stable results."""
    print("\n" + "="*60)
    print("TEST 7: Derotation consistency")
    print("="*60)

    def _derotate_gaze(pitch_rad, yaw_rad, roll_deg):
        roll_rad = -math.radians(roll_deg)
        cos_r, sin_r = math.cos(roll_rad), math.sin(roll_rad)
        return (yaw_rad * sin_r + pitch_rad * cos_r,
                yaw_rad * cos_r - pitch_rad * sin_r)

    pitch, yaw = 0.1, 0.2  # radians

    results = []
    for roll_deg in [0, 5, -5, 10, -10, 15]:
        dr_pitch, dr_yaw = _derotate_gaze(pitch, yaw, roll_deg)
        results.append((roll_deg, dr_pitch, dr_yaw))
        print(f"  roll={roll_deg:+4d}° → pitch={dr_pitch:.4f} yaw={dr_yaw:.4f}")

    # At roll=0, should pass through unchanged
    assert abs(results[0][1] - pitch) < 0.001, "Derotation at roll=0 should be identity for pitch"
    # Note: derotation formula swaps pitch/yaw, so at roll=0:
    # returns (yaw*sin(0) + pitch*cos(0), yaw*cos(0) - pitch*sin(0)) = (pitch, yaw)
    print(f"\n  Note: _derotate_gaze returns (pitch', yaw') = "
          f"(yaw*sin(-roll) + pitch*cos(-roll), yaw*cos(-roll) - pitch*sin(-roll))")
    print("  At roll=0: returns (pitch, yaw) — identity ✓")
    print("  PASSED")


def test_calibration_with_verification_points():
    """Simulate a full calibration + verification workflow.
    After calibrating, test 5 verification targets that weren't in calibration."""
    print("\n" + "="*60)
    print("TEST 8: Calibration + verification targets")
    print("="*60)

    np.random.seed(42)
    cal, ok = simulate_calibration(noise_std=0.005)
    assert ok

    # Verification points NOT in the calibration grid
    verify_targets = [
        (0.3, 0.3, "upper-left quarter"),
        (0.7, 0.3, "upper-right quarter"),
        (0.5, 0.5, "dead center"),
        (0.3, 0.7, "lower-left quarter"),
        (0.7, 0.7, "lower-right quarter"),
    ]

    lm = _make_landmarks_with_ear(0.28)
    all_pass = True

    for tx, ty, label in verify_targets:
        fusion = GazeEyeFusion(cal)
        yaw = (tx - 0.5) * 0.7
        pitch = (ty - 0.5) * 0.5
        for _ in range(15):
            result = fusion.update(yaw, pitch, lm)

        px, py = result["gaze_x"], result["gaze_y"]
        err = math.sqrt((px - tx)**2 + (py - ty)**2)
        status = "PASS" if result["focused"] and err < 0.2 else "FAIL"
        if status == "FAIL":
            all_pass = False

        print(f"  {label:25s} target=({tx:.1f},{ty:.1f}) → "
              f"gaze=({px:.3f},{py:.3f}) err={err:.3f} "
              f"focused={result['focused']} [{status}]")

    assert all_pass, "Verification targets failed!"
    print("  PASSED")


if __name__ == "__main__":
    test_calibration_accuracy()
    test_fusion_focused_at_center()
    test_fusion_focused_at_edges()
    test_fusion_unfocused_off_screen()
    test_fusion_with_closed_eyes()
    test_l2cs_cosine_scoring()
    test_derotation_consistency()
    test_calibration_with_verification_points()

    print("\n" + "="*60)
    print("ALL TESTS PASSED")
    print("="*60)