#!/usr/bin/env python3
"""
test_terrain_classifier.py — Unit tests for terrain classification logic.

Tests the core algorithms without requiring ROS2 or hardware.
"""

import unittest
import numpy as np
from unittest.mock import Mock, patch

import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from saltybot_bringup.terrain_classifier_node import (
    IMUVibrationAnalyzer,
    CNNTextureAnalyzer,
    TerrainClassifier,
    TerrainDecision,
)


class TestIMUVibrationAnalyzer(unittest.TestCase):
    """Test IMU vibration analysis and FFT roughness computation."""

    def setUp(self):
        self.analyzer = IMUVibrationAnalyzer(window_size=100, sample_rate=100.0)

    def test_empty_buffer_returns_zero(self):
        roughness, accel_mag = self.analyzer.compute_roughness()
        self.assertEqual(roughness, 0.0)
        self.assertEqual(accel_mag, 0.0)

    def test_smooth_surface_low_roughness(self):
        for _ in range(100):
            self.analyzer.add_sample(0.0, 0.0, 0.0)

        roughness, accel_mag = self.analyzer.compute_roughness()
        self.assertLess(roughness, 0.2)

    def test_rough_surface_high_roughness(self):
        for i in range(100):
            vibration = 2.0 * np.sin(2 * np.pi * 20 * i / 100.0)
            self.analyzer.add_sample(vibration, vibration, vibration)

        roughness, accel_mag = self.analyzer.compute_roughness()
        self.assertGreater(roughness, 0.3)

    def test_buffer_size_respected(self):
        self.assertEqual(self.analyzer.accel_buffer.maxlen, 100)
        for i in range(150):
            self.analyzer.add_sample(i, i, i)

        self.assertEqual(len(self.analyzer.accel_buffer), 100)

    def test_reset_clears_buffer(self):
        for i in range(50):
            self.analyzer.add_sample(i, i, i)

        self.assertGreater(len(self.analyzer.accel_buffer), 0)
        self.analyzer.reset()
        self.assertEqual(len(self.analyzer.accel_buffer), 0)

    def test_roughness_is_normalized(self):
        for i in range(100):
            vibration = 5.0 * np.sin(2 * np.pi * 30 * i / 100.0)
            self.analyzer.add_sample(vibration, vibration, vibration)

        roughness, _ = self.analyzer.compute_roughness()
        self.assertGreaterEqual(roughness, 0.0)
        self.assertLessEqual(roughness, 1.0)


class TestTerrainClassifier(unittest.TestCase):
    """Test terrain classification decision logic."""

    def setUp(self):
        self.classifier = TerrainClassifier()

    def test_pavement_classification(self):
        decision = self.classifier.classify(
            imu_roughness=0.1,
            accel_magnitude=9.85,
            cnn_features=None
        )

        self.assertEqual(decision.terrain_type, self.classifier.TERRAIN_PAVEMENT)
        self.assertGreater(decision.confidence, 0.7)
        self.assertEqual(decision.recommended_speed_ratio, 1.0)

    def test_grass_classification(self):
        decision = self.classifier.classify(
            imu_roughness=0.25,
            accel_magnitude=9.9,
            cnn_features=None
        )

        self.assertEqual(decision.terrain_type, self.classifier.TERRAIN_GRASS)
        self.assertGreater(decision.confidence, 0.5)
        self.assertEqual(decision.recommended_speed_ratio, 0.8)

    def test_gravel_classification(self):
        decision = self.classifier.classify(
            imu_roughness=0.55,
            accel_magnitude=10.0,
            cnn_features=None
        )

        self.assertEqual(decision.terrain_type, self.classifier.TERRAIN_GRAVEL)
        self.assertGreater(decision.confidence, 0.5)
        self.assertEqual(decision.recommended_speed_ratio, 0.5)

    def test_sand_classification(self):
        decision = self.classifier.classify(
            imu_roughness=0.6,
            accel_magnitude=10.3,
            cnn_features=None
        )

        self.assertEqual(decision.terrain_type, self.classifier.TERRAIN_SAND)
        self.assertGreater(decision.confidence, 0.5)
        self.assertEqual(decision.recommended_speed_ratio, 0.4)

    def test_all_terrains_have_speed_recommendations(self):
        for terrain_type in range(6):
            self.assertIn(
                terrain_type,
                self.classifier.SPEED_RECOMMENDATIONS
            )

            speed_ratio = self.classifier.SPEED_RECOMMENDATIONS[terrain_type]
            self.assertGreater(speed_ratio, 0.0)
            self.assertLessEqual(speed_ratio, 1.0)

    def test_decision_fields_are_valid(self):
        decision = self.classifier.classify(0.3, 9.9)

        self.assertGreaterEqual(decision.terrain_type, 0)
        self.assertLessEqual(decision.terrain_type, 5)

        self.assertGreaterEqual(decision.roughness, 0.0)
        self.assertLessEqual(decision.roughness, 1.0)

        self.assertGreaterEqual(decision.confidence, 0.0)
        self.assertLessEqual(decision.confidence, 1.0)

        self.assertGreater(decision.recommended_speed_ratio, 0.0)
        self.assertLessEqual(decision.recommended_speed_ratio, 1.0)


class TestTerrainDecisionDataclass(unittest.TestCase):
    """Test TerrainDecision dataclass."""

    def test_terrain_decision_creation(self):
        decision = TerrainDecision(
            terrain_type=1,
            roughness=0.3,
            confidence=0.8,
            recommended_speed_ratio=0.9
        )

        self.assertEqual(decision.terrain_type, 1)
        self.assertEqual(decision.roughness, 0.3)
        self.assertEqual(decision.confidence, 0.8)
        self.assertEqual(decision.recommended_speed_ratio, 0.9)


if __name__ == '__main__':
    unittest.main()
