#!/usr/bin/env python

# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Tests for SARM utility functions.

Tests the implementation of SARM paper formulas:
- Formula (1): compute_temporal_proportions - dataset-level temporal proportions
- Formula (2): compute_tau, compute_cumulative_progress - progress labels
"""

import pytest
import numpy as np
import torch

from lerobot.policies.sarm.sarm_utils import SubtaskAnnotation, Subtask, Timestamp
from lerobot.policies.sarm.sarm_utils import (
    compute_temporal_proportions,
    compute_tau,
    compute_cumulative_progress_batch,
)

def make_annotation(subtasks: list[tuple[str, int, int]]) -> SubtaskAnnotation:
    """Helper to create SubtaskAnnotation from list of (name, start_sec, end_sec)."""
    return SubtaskAnnotation(
        subtasks=[
            Subtask(
                name=name,
                timestamps=Timestamp(
                    start=f"{start // 60:02d}:{start % 60:02d}",
                    end=f"{end // 60:02d}:{end % 60:02d}"
                )
            )
            for name, start, end in subtasks
        ]
    )


class TestComputeTemporalProportions:
    """Tests for compute_temporal_proportions (SARM Paper Formula 1).
    
    Formula: ᾱ_k = (1/M) × Σ_i (L_{i,k} / T_i)
    
    Key insight: This averages the PROPORTION of each subtask within each trajectory,
    giving equal weight to all trajectories regardless of absolute length.
    """
    
    def test_basic_two_trajectories_equal_proportions(self):
        """Test with two trajectories that have equal proportions."""
        # Both trajectories: subtask1 = 50%, subtask2 = 50%
        # Traj 1: T=100s, subtask1=50s, subtask2=50s
        # Traj 2: T=200s, subtask1=100s, subtask2=100s
        annotations = {
            0: make_annotation([('subtask1', 0, 50), ('subtask2', 50, 100)]),
            1: make_annotation([('subtask1', 0, 100), ('subtask2', 100, 200)]),
        }
        
        result = compute_temporal_proportions(annotations)
        
        # Both should be 0.5
        assert abs(result['subtask1'] - 0.5) < 1e-6
        assert abs(result['subtask2'] - 0.5) < 1e-6
    
    def test_paper_example_different_from_avg_durations(self):
        """Test that compute_temporal_proportions differs from naive average duration approach.
        
        This is the key test showing the difference between:
        - Paper formula: average of (L_i,k / T_i)
        - Naive approach: mean(L_i,k) / sum(mean(L_i,j))
        """
        # Episode 1: T=100s, subtask1=80s, subtask2=20s (proportions: 0.8, 0.2)
        # Episode 2: T=200s, subtask1=40s, subtask2=160s (proportions: 0.2, 0.8)
        annotations = {
            0: make_annotation([('subtask1', 0, 80), ('subtask2', 80, 100)]),
            1: make_annotation([('subtask1', 0, 40), ('subtask2', 40, 200)]),
        }
        
        result = compute_temporal_proportions(annotations)
        
        # Paper formula: 
        # ᾱ_1 = (1/2) × (80/100 + 40/200) = (1/2) × (0.8 + 0.2) = 0.5
        # ᾱ_2 = (1/2) × (20/100 + 160/200) = (1/2) × (0.2 + 0.8) = 0.5
        assert abs(result['subtask1'] - 0.5) < 1e-6
        assert abs(result['subtask2'] - 0.5) < 1e-6
    
    def test_single_trajectory(self):
        """Test with a single trajectory."""
        # T=100s, reach=30s, grasp=20s, lift=50s
        annotations = {
            0: make_annotation([('reach', 0, 30), ('grasp', 30, 50), ('lift', 50, 100)]),
        }
        
        result = compute_temporal_proportions(annotations)
        
        assert abs(result['reach'] - 0.3) < 1e-6
        assert abs(result['grasp'] - 0.2) < 1e-6
        assert abs(result['lift'] - 0.5) < 1e-6
    
    def test_sum_to_one(self):
        """Test that proportions always sum to 1."""
        # Three episodes with varying proportions
        annotations = {
            0: make_annotation([('a', 0, 10), ('b', 10, 50), ('c', 50, 100)]),  # 0.1, 0.4, 0.5
            1: make_annotation([('a', 0, 20), ('b', 20, 70), ('c', 70, 100)]),  # 0.2, 0.5, 0.3
            2: make_annotation([('a', 0, 30), ('b', 30, 90), ('c', 90, 100)]),  # 0.3, 0.6, 0.1
        }
        
        result = compute_temporal_proportions(annotations)
        
        total = sum(result.values())
        assert abs(total - 1.0) < 1e-6
    
    def test_empty_annotations_returns_empty(self):
        """Test that empty annotations returns empty dict."""
        result = compute_temporal_proportions({})
        assert result == {}
    
    def test_uniform_proportions(self):
        """Test with uniform proportions across subtasks."""
        # Each subtask takes 25% of each episode
        annotations = {
            0: make_annotation([('a', 0, 25), ('b', 25, 50), ('c', 50, 75), ('d', 75, 100)]),
            1: make_annotation([('a', 0, 50), ('b', 50, 100), ('c', 100, 150), ('d', 150, 200)]),
        }
        
        result = compute_temporal_proportions(annotations)
        
        for name in ['a', 'b', 'c', 'd']:
            assert abs(result[name] - 0.25) < 1e-6


class TestComputeTau:
    """Tests for compute_tau (within-subtask progress).
    
    Formula: τ_t = (t - s_k) / (e_k - s_k) ∈ [0, 1]
    """
    
    def test_at_start(self):
        """τ should be 0 at subtask start."""
        tau = compute_tau(current_frame=10, subtask_start=10, subtask_end=50)
        assert tau == 0.0
    
    def test_at_end(self):
        """τ should be 1 at subtask end."""
        tau = compute_tau(current_frame=50, subtask_start=10, subtask_end=50)
        assert tau == 1.0
    
    def test_at_middle(self):
        """τ should be 0.5 at subtask midpoint."""
        tau = compute_tau(current_frame=30, subtask_start=10, subtask_end=50)
        assert abs(tau - 0.5) < 1e-6
    
    def test_quarter_progress(self):
        """Test τ at 25% through subtask."""
        tau = compute_tau(current_frame=20, subtask_start=0, subtask_end=80)
        assert abs(tau - 0.25) < 1e-6
    
    def test_zero_duration_subtask(self):
        """τ should be 1.0 for zero-duration subtask."""
        tau = compute_tau(current_frame=10, subtask_start=10, subtask_end=10)
        assert tau == 1.0
    
    def test_clamps_below_zero(self):
        """τ should be clamped to 0 if frame is before subtask."""
        tau = compute_tau(current_frame=5, subtask_start=10, subtask_end=50)
        assert tau == 0.0
    
    def test_clamps_above_one(self):
        """τ should be clamped to 1 if frame is after subtask."""
        tau = compute_tau(current_frame=60, subtask_start=10, subtask_end=50)
        assert tau == 1.0
    
    def test_float_inputs(self):
        """Test with float frame indices (from interpolation)."""
        tau = compute_tau(current_frame=25.5, subtask_start=10.0, subtask_end=50.0)
        expected = (25.5 - 10.0) / (50.0 - 10.0)
        assert abs(tau - expected) < 1e-6


class TestComputeCumulativeProgressBatchScalar:
    """Tests for compute_cumulative_progress_batch with scalar inputs (normalized progress y_t).
    
    Formula: y_t = P_{k-1} + ᾱ_k × τ_t ∈ [0, 1]
    """
    
    def test_first_subtask_start(self):
        """y should be 0 at start of first subtask."""
        proportions = [0.3, 0.5, 0.2]
        y = compute_cumulative_progress_batch(tau=0.0, stage_indices=0, alpha=proportions)
        assert y == 0.0
    
    def test_first_subtask_end(self):
        """y should equal ᾱ_1 at end of first subtask."""
        proportions = [0.3, 0.5, 0.2]
        y = compute_cumulative_progress_batch(tau=1.0, stage_indices=0, alpha=proportions)
        assert abs(y - 0.3) < 1e-6
    
    def test_second_subtask_start(self):
        """y should equal P_1 at start of second subtask."""
        proportions = [0.3, 0.5, 0.2]
        y = compute_cumulative_progress_batch(tau=0.0, stage_indices=1, alpha=proportions)
        assert abs(y - 0.3) < 1e-6
    
    def test_second_subtask_end(self):
        """y should equal P_2 at end of second subtask."""
        proportions = [0.3, 0.5, 0.2]
        y = compute_cumulative_progress_batch(tau=1.0, stage_indices=1, alpha=proportions)
        assert abs(y - 0.8) < 1e-6  # 0.3 + 0.5
    
    def test_third_subtask_end(self):
        """y should be 1.0 at end of last subtask."""
        proportions = [0.3, 0.5, 0.2]
        y = compute_cumulative_progress_batch(tau=1.0, stage_indices=2, alpha=proportions)
        assert abs(y - 1.0) < 1e-6
    
    def test_midpoint_of_subtask(self):
        """Test progress at midpoint of a subtask."""
        proportions = [0.4, 0.6]
        # At τ=0.5 in subtask 1: y = P_0 + ᾱ_1 × 0.5 = 0 + 0.4 × 0.5 = 0.2
        y = compute_cumulative_progress_batch(tau=0.5, stage_indices=0, alpha=proportions)
        assert abs(y - 0.2) < 1e-6
        
        # At τ=0.5 in subtask 2: y = P_1 + ᾱ_2 × 0.5 = 0.4 + 0.6 × 0.5 = 0.7
        y = compute_cumulative_progress_batch(tau=0.5, stage_indices=1, alpha=proportions)
        assert abs(y - 0.7) < 1e-6
    
    def test_uniform_proportions(self):
        """Test with uniform proportions."""
        proportions = [0.25, 0.25, 0.25, 0.25]
        
        # At end of each subtask, progress should be 0.25, 0.5, 0.75, 1.0
        for i in range(4):
            y = compute_cumulative_progress_batch(tau=1.0, stage_indices=i, alpha=proportions)
            expected = (i + 1) * 0.25
            assert abs(y - expected) < 1e-6


class TestComputeCumulativeProgressBatchTensor:
    """Tests for compute_cumulative_progress_batch with tensor inputs (GPU batch version)."""
    
    def test_tensor_matches_scalar_version(self):
        """Test that tensor version matches scalar version."""
        proportions = [0.3, 0.5, 0.2]
        alpha = torch.tensor(proportions, dtype=torch.float32)
        cumulative = torch.zeros(len(proportions) + 1, dtype=torch.float32)
        cumulative[1:] = torch.cumsum(alpha, dim=0)
        
        test_cases = [
            (0.0, 0),  # start of subtask 0
            (1.0, 0),  # end of subtask 0
            (0.0, 1),  # start of subtask 1
            (0.5, 1),  # middle of subtask 1
            (1.0, 2),  # end of subtask 2
        ]
        
        for tau_val, stage_idx in test_cases:
            # Scalar version
            expected = compute_cumulative_progress_batch(tau_val, stage_idx, proportions)
            
            # Tensor version (single element)
            tau = torch.tensor([[[tau_val]]])  # (1, 1, 1)
            stages = torch.tensor([[stage_idx]])  # (1, 1)
            result = compute_cumulative_progress_batch(tau, stages, alpha, cumulative)
            
            assert abs(result[0, 0, 0].item() - expected) < 1e-6
    
    def test_batch_processing(self):
        """Test batch processing with multiple samples."""
        proportions = [0.4, 0.6]
        alpha = torch.tensor(proportions, dtype=torch.float32)
        cumulative = torch.zeros(3, dtype=torch.float32)
        cumulative[1:] = torch.cumsum(alpha, dim=0)
        
        # Batch of 2 samples, sequence length 3
        tau = torch.tensor([
            [[0.0], [0.5], [1.0]],  # sample 1
            [[0.0], [0.5], [1.0]],  # sample 2
        ])
        stages = torch.tensor([
            [0, 0, 0],  # sample 1: all in subtask 0
            [1, 1, 1],  # sample 2: all in subtask 1
        ])
        
        result = compute_cumulative_progress_batch(tau, stages, alpha, cumulative)
        
        # Sample 1: subtask 0 with tau 0, 0.5, 1.0 -> y = 0, 0.2, 0.4
        assert abs(result[0, 0, 0].item() - 0.0) < 1e-6
        assert abs(result[0, 1, 0].item() - 0.2) < 1e-6
        assert abs(result[0, 2, 0].item() - 0.4) < 1e-6
        
        # Sample 2: subtask 1 with tau 0, 0.5, 1.0 -> y = 0.4, 0.7, 1.0
        assert abs(result[1, 0, 0].item() - 0.4) < 1e-6
        assert abs(result[1, 1, 0].item() - 0.7) < 1e-6
        assert abs(result[1, 2, 0].item() - 1.0) < 1e-6
    
    def test_auto_compute_cumulative_prior(self):
        """Test that cumulative_prior is auto-computed when not provided."""
        proportions = [0.3, 0.5, 0.2]
        alpha = torch.tensor(proportions, dtype=torch.float32)
        
        tau = torch.tensor([[[0.5]]])
        stages = torch.tensor([[1]])
        
        # Without cumulative_prior (should auto-compute)
        result = compute_cumulative_progress_batch(tau, stages, alpha)
        
        # Expected: P_0 + alpha_1 * 0.5 = 0.3 + 0.5 * 0.5 = 0.55
        assert abs(result[0, 0, 0].item() - 0.55) < 1e-6


class TestEndToEndProgressLabeling:
    """End-to-end tests for progress label computation."""
    
    def test_consistent_semantic_meaning(self):
        """Test that same subtask completion maps to same progress across trajectories.
        
        This is the key semantic property: "end of subtask 1" should always 
        mean the same progress value regardless of trajectory speed.
        """
        proportions = [0.3, 0.5, 0.2]
        
        # Fast trajectory: subtask 1 ends at frame 30 (of 100)
        tau_fast = compute_tau(30, 0, 30)  # = 1.0
        y_fast = compute_cumulative_progress_batch(tau_fast, 0, proportions)
        
        # Slow trajectory: subtask 1 ends at frame 90 (of 300)
        tau_slow = compute_tau(90, 0, 90)  # = 1.0
        y_slow = compute_cumulative_progress_batch(tau_slow, 0, proportions)
        
        # Both should map to same progress (0.3 = end of subtask 1)
        assert abs(y_fast - y_slow) < 1e-6
        assert abs(y_fast - 0.3) < 1e-6
    
    def test_monotonic_within_subtask(self):
        """Test that progress is monotonically increasing within a subtask."""
        proportions = [0.4, 0.6]
        
        prev_y = -1
        for tau in np.linspace(0, 1, 11):
            y = compute_cumulative_progress_batch(tau, 0, proportions)
            assert y > prev_y or (tau == 0 and y == 0)
            prev_y = y
    
    def test_continuous_across_subtasks(self):
        """Test that progress is continuous at subtask boundaries."""
        proportions = [0.3, 0.5, 0.2]
        
        # End of subtask 0 (tau=1.0)
        y_end_0 = compute_cumulative_progress_batch(1.0, 0, proportions)
        
        # Start of subtask 1 (tau=0.0)
        y_start_1 = compute_cumulative_progress_batch(0.0, 1, proportions)
        
        # Should be equal (P_1 = 0.3)
        assert abs(y_end_0 - y_start_1) < 1e-6
        
        # End of subtask 1
        y_end_1 = compute_cumulative_progress_batch(1.0, 1, proportions)
        
        # Start of subtask 2
        y_start_2 = compute_cumulative_progress_batch(0.0, 2, proportions)
        
        # Should be equal (P_2 = 0.8)
        assert abs(y_end_1 - y_start_2) < 1e-6