import unittest
import sys
import os
import numpy as np

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

from spectral_features import (
    calculate_loudness_seq,
    calculate_spectral_centroid_seq,
    calculate_spectral_complexity_seq,
    warp_sequence,
)
from text_utils import (
    generate_spectral_centroid_tags,
    generate_spectral_complexity_tags,
    generate_loudness_contour_tags,
    get_control_tags,
)


class TestSpectralFeatureCalculation(unittest.TestCase):
    """Test spectral feature calculation functions."""

    def test_loudness_seq_basic(self):
        """Test loudness calculation returns correct shape and range."""
        # Create simple audio with known RMS
        sample_rate = 24000
        duration_s = 1.0
        audio = np.ones(int(sample_rate * duration_s)) * 0.5  # RMS should be 0.5

        loudness = calculate_loudness_seq(audio, sample_rate=sample_rate, target_rate=25)

        # Should have ~25 frames for 1 second at 25Hz
        self.assertEqual(len(loudness), 25)
        # Values should be positive
        self.assertTrue(np.all(loudness >= 0))

    def test_loudness_seq_normalized(self):
        """Test loudness normalization to [0,1]."""
        sample_rate = 24000
        duration_s = 1.0
        # Create audio with varying amplitude
        audio = np.sin(2 * np.pi * 440 * np.arange(int(sample_rate * duration_s)) / sample_rate)

        loudness = calculate_loudness_seq(audio, sample_rate=sample_rate, target_rate=25, normalize=True)

        # Should be in [0, 1] range
        self.assertTrue(np.all(loudness >= 0))
        self.assertTrue(np.all(loudness <= 1))

    def test_spectral_centroid_seq_unnormalized(self):
        """Test spectral centroid calculation without normalization returns Hz values."""
        sample_rate = 24000
        duration_s = 20.0
        # Create pure tone at 440 Hz (A4 note)
        audio = np.sin(2 * np.pi * 440 * np.arange(int(sample_rate * duration_s)) / sample_rate)

        centroid = calculate_spectral_centroid_seq(
            audio,
            sample_rate=sample_rate,
            target_rate=1,
            normalize=False,
        )

        # Take only middle 50% of frames to avoid edge effects from padding/smoothing
        total_frames = len(centroid)
        start_idx = total_frames // 4
        end_idx = total_frames - (total_frames // 4)
        middle_centroid = centroid[start_idx:end_idx]

        # Without normalization, should return Hz values close to 440 Hz
        mean_centroid = np.mean(middle_centroid)

        # For a pure 440 Hz tone, centroid should be very close to 440 Hz
        # Middle frames should be accurate with long duration
        # Allow ±5 Hz tolerance
        self.assertGreater(mean_centroid, 440 - 5)
        self.assertLess(mean_centroid, 440 + 5)

        # Verify all values are positive Hz values (not normalized 0-1 range)
        self.assertTrue(np.all(centroid > 1))  # Should be well above normalized range

    def test_spectral_centroid_seq_normalized(self):
        """Test spectral centroid calculation with normalization."""
        sample_rate = 24000
        duration_s = 1.0
        # Create pure tone at 1000 Hz
        audio = np.sin(2 * np.pi * 1000 * np.arange(int(sample_rate * duration_s)) / sample_rate)

        # Calculate loudness for silence detection
        loudness = calculate_loudness_seq(audio, sample_rate=sample_rate, target_rate=1)

        centroid = calculate_spectral_centroid_seq(
            audio,
            sample_rate=sample_rate,
            target_rate=1,
            normalize=True,
            loudness_seq=loudness,
        )

        # Should be in [0, 1] range
        self.assertTrue(np.all(centroid >= 0))
        self.assertTrue(np.all(centroid <= 1))
        # Length depends on padding, should be close to 1 frame per second at 1Hz
        self.assertGreater(len(centroid), 0)

    def test_spectral_complexity_seq_unnormalized(self):
        """Test spectral complexity calculation without normalization."""
        sample_rate = 24000
        duration_s = 20.0
        # Create white noise (should have high complexity, close to 1.0)
        np.random.seed(42)
        audio = np.random.randn(int(sample_rate * duration_s)) * 0.5

        complexity = calculate_spectral_complexity_seq(
            audio,
            sample_rate=sample_rate,
            target_rate=1,
            normalize=False,
        )

        # Take middle 50% to avoid edge effects
        total_frames = len(complexity)
        start_idx = total_frames // 4
        end_idx = total_frames - (total_frames // 4)
        middle_complexity = complexity[start_idx:end_idx]

        # White noise should have very high complexity (close to 1.0)
        mean_complexity = np.mean(middle_complexity)
        # Spectral entropy for white noise should be high (>0.9)
        self.assertGreater(mean_complexity, 0.9)

    def test_spectral_complexity_seq_normalized(self):
        """Test spectral complexity calculation with normalization."""
        sample_rate = 24000
        duration_s = 20.0
        # Create white noise (should have high complexity)
        np.random.seed(42)
        audio = np.random.randn(int(sample_rate * duration_s)) * 0.5

        # Calculate loudness for silence detection
        loudness = calculate_loudness_seq(audio, sample_rate=sample_rate, target_rate=1)

        complexity = calculate_spectral_complexity_seq(
            audio,
            sample_rate=sample_rate,
            target_rate=1,
            normalize=True,
            loudness_seq=loudness,
        )

        # Should be in [0, 1] range
        self.assertTrue(np.all(complexity >= 0))
        self.assertTrue(np.all(complexity <= 1))
        # Should produce valid output
        self.assertGreater(len(complexity), 0)

    def test_spectral_features_silence(self):
        """Test spectral features with silent audio."""
        sample_rate = 24000
        duration_s = 1.0
        audio = np.zeros(int(sample_rate * duration_s))

        loudness = calculate_loudness_seq(audio, sample_rate=sample_rate, target_rate=1)
        centroid = calculate_spectral_centroid_seq(
            audio,
            sample_rate=sample_rate,
            target_rate=1,
            normalize=True,
            loudness_seq=loudness,
        )
        complexity = calculate_spectral_complexity_seq(
            audio,
            sample_rate=sample_rate,
            target_rate=1,
            normalize=True,
            loudness_seq=loudness,
        )

        # Silence should have zero loudness
        self.assertTrue(np.all(loudness == 0))
        # Centroid and complexity should handle silence gracefully
        self.assertTrue(np.all(centroid >= 0))
        self.assertTrue(np.all(complexity >= 0))


class TestSpectralTagGeneration(unittest.TestCase):
    """Test spectral feature tag generation functions."""

    def test_centroid_tags_format(self):
        """Test centroid tags are in correct 0-100 format."""
        # Create normalized test data
        centroid_normalized = np.array([0.5, 0.6, 0.7, 0.8])

        tags = generate_spectral_centroid_tags(centroid_normalized)

        # Should return single tag
        self.assertEqual(len(tags), 1)
        # Tag should start with "spectral_centroid_contour:["
        self.assertTrue(tags[0].startswith("spectral_centroid_contour:["))
        # Tag should end with "]"
        self.assertTrue(tags[0].endswith("]"))
        # Values should be integers in 0-100 range
        values_str = tags[0][len("spectral_centroid_contour:[") : -1]
        values = [int(v) for v in values_str.split(",")]
        self.assertEqual(values, [50, 60, 70, 80])

    def test_complexity_tags_format(self):
        """Test complexity tags are in correct 0-100 format."""
        # Create normalized test data
        complexity_normalized = np.array([0.3, 0.4, 0.5, 0.6])

        tags = generate_spectral_complexity_tags(complexity_normalized)

        # Should return single tag
        self.assertEqual(len(tags), 1)
        # Tag should start with "spectral_complexity_contour:["
        self.assertTrue(tags[0].startswith("spectral_complexity_contour:["))
        # Values should be integers in 0-100 range
        values_str = tags[0][len("spectral_complexity_contour:[") : -1]
        values = [int(v) for v in values_str.split(",")]
        self.assertEqual(values, [30, 40, 50, 60])

    def test_empty_input(self):
        """Test tag generation with empty input."""
        self.assertEqual(generate_spectral_centroid_tags(None), [])
        self.assertEqual(generate_spectral_centroid_tags(np.array([])), [])
        self.assertEqual(generate_spectral_complexity_tags(None), [])
        self.assertEqual(generate_spectral_complexity_tags(np.array([])), [])

    def test_edge_values(self):
        """Test tag generation with edge values (0 and 1)."""
        centroid_edge = np.array([0.0, 0.5, 1.0])
        tags = generate_spectral_centroid_tags(centroid_edge)

        values_str = tags[0][len("spectral_centroid_contour:[") : -1]
        values = [int(v) for v in values_str.split(",")]
        self.assertEqual(values, [0, 50, 100])


class TestSpectralFeaturesIntegration(unittest.TestCase):
    """Test integration of spectral features with control tags."""

    def test_get_control_tags_with_spectral_features(self):
        """Test that spectral features are included in control tags."""
        # Create test data
        loudness_25hz = np.ones(75) * 0.5  # 3 seconds at 25Hz
        loudness_seq = np.array([0.6, 0.7, 0.8])  # 3 seconds at variable rate
        centroid_seq = np.array([0.6, 0.7, 0.8])
        complexity_seq = np.array([0.4, 0.5, 0.6])

        # Get control tags
        tags = get_control_tags(
            sample_duration_s=3.0,
            sample_duration_toks=75,
            do_augment=False,
            loudness_25hz=loudness_25hz,
            loudness_seq=loudness_seq,
            spectral_centroid_seq=centroid_seq,
            spectral_complexity_seq=complexity_seq,
            contour_rate_hz=1.0,
        )

        # Should contain all contour tags with new names
        self.assertIsNotNone(tags)
        self.assertIn("loudness_contour:[", tags)
        self.assertIn("spectral_centroid_contour:[", tags)
        self.assertIn("spectral_complexity_contour:[", tags)

    def test_spectral_tags_in_control_tags_format(self):
        """Test that spectral tags use 0-100 integer format in control tags."""
        centroid_seq = np.array([0.5])
        complexity_seq = np.array([0.7])

        tags = get_control_tags(
            sample_duration_s=1.0,
            sample_duration_toks=25,
            do_augment=False,
            spectral_centroid_seq=centroid_seq,
            spectral_complexity_seq=complexity_seq,
        )

        # Extract spectral tags from control tags string
        self.assertIn("spectral_centroid_contour:[50]", tags)
        self.assertIn("spectral_complexity_contour:[70]", tags)


class TestTimeWarping(unittest.TestCase):
    """Test time-warping augmentation function."""

    def test_warp_no_effect_when_ratio_zero(self):
        """Test that warp_ratio=0 returns unchanged sequence."""
        sequence = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
        warped = warp_sequence(sequence, warp_ratio=0.0, num_anchor_points=3)

        np.testing.assert_array_almost_equal(warped, sequence)

    def test_warp_preserves_length(self):
        """Test that warping preserves sequence length."""
        sequence = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0])
        warped = warp_sequence(sequence, warp_ratio=0.5, num_anchor_points=5)

        self.assertEqual(len(warped), len(sequence))

    def test_warp_preserves_value_range(self):
        """Test that warping preserves value range (no new values created)."""
        np.random.seed(42)
        sequence = np.array([0.1, 0.5, 0.9, 0.3, 0.7, 0.2, 0.8, 0.4, 0.6])
        warped = warp_sequence(sequence, warp_ratio=1.0, num_anchor_points=4)

        # Warped values should be within original range
        self.assertGreaterEqual(np.min(warped), np.min(sequence))
        self.assertLessEqual(np.max(warped), np.max(sequence))

    def test_warp_boundary_preservation(self):
        """Test that first and last values are preserved."""
        sequence = np.array([10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0])
        warped = warp_sequence(sequence, warp_ratio=0.8, num_anchor_points=5)

        # First and last should be close (anchors include boundaries)
        self.assertAlmostEqual(warped[0], sequence[0], places=1)
        self.assertAlmostEqual(warped[-1], sequence[-1], places=1)

    def test_warp_with_feature_calculation(self):
        """Test that warping integrates with feature calculation pipeline."""
        sample_rate = 24000
        duration_s = 10.0
        # Create varying audio
        t = np.arange(int(sample_rate * duration_s)) / sample_rate
        audio = np.sin(2 * np.pi * 100 * t) + 0.5 * np.sin(2 * np.pi * 200 * t)

        # Calculate with warping enabled
        loudness_warped = calculate_loudness_seq(
            audio,
            sample_rate=sample_rate,
            target_rate=1,
            normalize=True,
            apply_warp=True,
            warp_ratio=1.0,
        )

        # Calculate without warping
        loudness_orig = calculate_loudness_seq(
            audio, sample_rate=sample_rate, target_rate=1, normalize=True, apply_warp=False
        )

        # Should have same length and valid values
        self.assertEqual(len(loudness_warped), len(loudness_orig))
        self.assertTrue(np.all(loudness_warped >= 0))
        self.assertTrue(np.all(loudness_warped <= 1))


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