# Copyright 2025 Thinking Machines Lab
#
# Licensed under the Apache License, Version 2.0
#
# Modifications:
# - Adapted for HPC-AI cloud fine-tuning workflow
# Copyright © 2025 HPC-AI.COM

import pytest

from hpcai.lib.chunked_fwdbwd_helpers import (
    _unique,
    _metrics_reduction,
    REDUCE_MAP,
)
from hpcai.types import ForwardBackwardOutput, LossFnOutput, TensorData


class TestMetricsReductionWithUnique:
    """Test the _metrics_reduction function with unique metrics."""

    def create_forward_backward_output(self, metrics: dict, num_loss_fn_outputs: int = 1) -> ForwardBackwardOutput:
        """Helper to create ForwardBackwardOutput for testing."""
        # LossFnOutput is Dict[str, TensorData], so create sample data
        tensor_data = TensorData(
            data=[0.0],
            dtype="float32",
            shape=[1]
        )
        loss_fn_outputs = [{"loss": tensor_data} for _ in range(num_loss_fn_outputs)]
        return ForwardBackwardOutput(
            loss_fn_output_type="test",
            metrics=metrics,
            loss_fn_outputs=loss_fn_outputs,
        )

    def test_unique_reduction_single_result(self):
        """Test unique reduction with single result."""
        results = [
            self.create_forward_backward_output({"clock_cycle:unique": 12345})
        ]

        reduced = _metrics_reduction(results)

        assert "clock_cycle:unique" in reduced
        assert reduced["clock_cycle:unique"] == 12345

    def test_unique_reduction_multiple_results(self):
        """Test unique reduction with multiple results creates additional keys."""
        results = [
            self.create_forward_backward_output({"clock_cycle:unique": 12345}),
            self.create_forward_backward_output({"clock_cycle:unique": 12346}),
            self.create_forward_backward_output({"clock_cycle:unique": 12347}),
        ]

        reduced = _metrics_reduction(results)

        # First value should be the main key
        assert reduced["clock_cycle:unique"] == 12345
        # Additional values should get suffixed keys
        assert reduced["clock_cycle:unique_1"] == 12346
        assert reduced["clock_cycle:unique_2"] == 12347

    def test_unique_reduction_with_other_metrics(self):
        """Test unique reduction alongside other metric types."""
        results = [
            self.create_forward_backward_output({
                "clock_cycle:unique": 100,
                "loss:mean": 0.5,
                "accuracy:max": 0.8,
            }),
            self.create_forward_backward_output({
                "clock_cycle:unique": 101,
                "loss:mean": 0.6,
                "accuracy:max": 0.9,
            }),
        ]

        reduced = _metrics_reduction(results)

        # Unique metric behavior
        assert reduced["clock_cycle:unique"] == 100
        assert reduced["clock_cycle:unique_1"] == 101

        # Other metrics should work as normal
        assert reduced["loss:mean"] == 0.55
        assert reduced["accuracy:max"] == 0.9

    def test_unique_reduction_empty_results(self):
        """Test unique reduction with empty results."""
        results = []
        reduced = _metrics_reduction(results)
        assert reduced == {}

    def test_unique_reduction_missing_metric(self):
        """Test unique reduction when some results don't have the metric."""
        results = [
            self.create_forward_backward_output({"clock_cycle:unique": 100}),
            self.create_forward_backward_output({"other_metric:mean": 0.5}),  # Missing clock_cycle
        ]

        reduced = _metrics_reduction(results)

        # Neither metric should be present since they're not in all results
        # This matches the actual behavior of _metrics_reduction which requires
        # all results to have a metric for it to be processed
        assert "clock_cycle:unique" not in reduced
        assert "other_metric:mean" not in reduced
        assert len(reduced) == 0   # Should be empty

    def test_unique_reduction_with_float_values(self):
        """Test unique reduction with float values."""
        results = [
            self.create_forward_backward_output({"timestamp:unique": 1234567890.123}),
            self.create_forward_backward_output({"timestamp:unique": 1234567890.456}),
        ]

        reduced = _metrics_reduction(results)

        assert reduced["timestamp:unique"] == 1234567890.123
        assert reduced["timestamp:unique_1"] == 1234567890.456

    def test_invalid_reduction_type_raises_assertion(self):
        """Test that invalid reduction types raise AssertionError."""
        results = [
            self.create_forward_backward_output({"invalid:nonexistent": 100})
        ]

        with pytest.raises(AssertionError, match="Invalid reduction nonexistent"):
            _metrics_reduction(results)
