"""
Unit tests for ML prediction models
"""

import pytest
import numpy as np
import sys
import os

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

from models import (
    TaskCompletionPredictor,
    DifficultyClassifier,
    StressEstimator,
    DurationPredictor,
    ProductivityForecaster,
    EnsemblePredictor,
)


class TestTaskCompletionPredictor:
    """Tests for TaskCompletionPredictor class"""

    @pytest.fixture
    def predictor(self):
        return TaskCompletionPredictor()

    def test_predict_probability_range(self, predictor):
        """Test that prediction is within valid range"""
        features = {
            "complexity_normalized": 0.6,
            "pri_attention_demand": 0.5,
            "trait_conscientiousness": 0.7,
            "completion_rate": 0.7,
        }
        
        probability = predictor.predict(features)
        
        assert 0 <= probability <= 1

    def test_high_conscientiousness_improves_probability(self, predictor):
        """Test that high conscientiousness increases completion probability"""
        base_features = {
            "complexity_normalized": 0.6,
            "pri_attention_demand": 0.5,
            "completion_rate": 0.7,
        }
        
        low_consc = {**base_features, "trait_conscientiousness": 0.3}
        high_consc = {**base_features, "trait_conscientiousness": 0.9}
        
        prob_low = predictor.predict(low_consc)
        prob_high = predictor.predict(high_consc)
        
        assert prob_high > prob_low

    def test_high_complexity_reduces_probability(self, predictor):
        """Test that high complexity reduces completion probability"""
        base_features = {
            "pri_attention_demand": 0.5,
            "trait_conscientiousness": 0.7,
            "completion_rate": 0.7,
        }
        
        low_complex = {**base_features, "complexity_normalized": 0.2}
        high_complex = {**base_features, "complexity_normalized": 1.0}
        
        prob_low = predictor.predict(low_complex)
        prob_high = predictor.predict(high_complex)
        
        assert prob_low > prob_high


class TestDifficultyClassifier:
    """Tests for DifficultyClassifier class"""

    @pytest.fixture
    def classifier(self):
        return DifficultyClassifier()

    def test_classify_returns_valid_level(self, classifier):
        """Test that classification returns valid difficulty level"""
        features = {"complexity_normalized": 0.6, "duration_normalized": 0.25}
        
        level, confidence = classifier.classify(features)
        
        assert level in ["EASY", "MODERATE", "HARD"]
        assert 0 <= confidence <= 1

    def test_low_complexity_is_easy(self, classifier):
        """Test that low complexity tasks are classified as easy"""
        features = {"complexity_normalized": 0.1, "duration_normalized": 0.1}
        
        level, _ = classifier.classify(features)
        
        assert level == "EASY"

    def test_high_complexity_is_hard(self, classifier):
        """Test that high complexity tasks are classified as hard"""
        features = {"complexity_normalized": 1.0, "duration_normalized": 1.0, "time_pressure": 1.0}
        
        level, _ = classifier.classify(features)
        
        assert level == "HARD"


class TestStressEstimator:
    """Tests for StressEstimator class"""

    @pytest.fixture
    def estimator(self):
        return StressEstimator()

    def test_estimate_returns_valid_range(self, estimator):
        """Test that stress estimate is within valid range"""
        features = {
            "complexity_normalized": 0.6,
            "pri_urgency_score": 0.5,
            "time_pressure": 0.14,
        }
        
        stress = estimator.estimate(features)
        
        assert 1 <= stress <= 10

    def test_urgent_task_high_stress(self, estimator):
        """Test that urgent tasks have higher stress"""
        normal = {
            "complexity_normalized": 0.6,
            "pri_urgency_score": 0.5,
            "time_pressure": 0.1,
        }
        urgent = {
            "complexity_normalized": 0.6,
            "pri_urgency_score": 0.9,
            "time_pressure": 1.0,
        }
        
        stress_normal = estimator.estimate(normal)
        stress_urgent = estimator.estimate(urgent)
        
        assert stress_urgent > stress_normal

    def test_high_neuroticism_increases_stress(self, estimator):
        """Test that high neuroticism increases stress via personality"""
        base_features = {
            "complexity_normalized": 0.6,
            "pri_urgency_score": 0.5,
            "time_pressure": 0.3,
        }
        
        low_neuro_personality = {"neuroticism": 20}
        high_neuro_personality = {"neuroticism": 80}
        
        stress_low = estimator.estimate(base_features, low_neuro_personality)
        stress_high = estimator.estimate(base_features, high_neuro_personality)
        
        assert stress_high > stress_low


class TestDurationPredictor:
    """Tests for DurationPredictor class"""

    @pytest.fixture
    def predictor(self):
        return DurationPredictor()

    def test_predict_returns_positive_duration(self, predictor):
        """Test that predicted duration is positive"""
        estimated_duration = 60
        difficulty = "MODERATE"
        category = "WORK"
        
        duration = predictor.predict(estimated_duration, difficulty, category)
        
        assert duration > 0

    def test_complex_tasks_take_longer(self, predictor):
        """Test that harder tasks have longer predicted duration"""
        estimated_duration = 30
        category = "WORK"
        
        duration_easy = predictor.predict(estimated_duration, "EASY", category)
        duration_hard = predictor.predict(estimated_duration, "HARD", category)
        
        assert duration_hard > duration_easy


class TestProductivityForecaster:
    """Tests for ProductivityForecaster class"""

    @pytest.fixture
    def forecaster(self):
        return ProductivityForecaster()

    def test_forecast_returns_correct_days(self, forecaster):
        """Test that forecast returns correct number of days"""
        historical_data = {
            "prod_monday": 70,
            "prod_tuesday": 75,
        }
        
        forecast = forecaster.forecast(historical_data, days=7)
        
        assert len(forecast) == 7

    def test_forecast_valid_productivity_range(self, forecaster):
        """Test that forecasted productivity is valid"""
        historical_data = {}
        
        forecast = forecaster.forecast(historical_data, days=5)
        
        for day in forecast:
            assert 0 <= day["predicted_productivity"] <= 100
            assert day["recommended_tasks"] >= 0


class TestEnsemblePredictor:
    """Tests for EnsemblePredictor class"""

    @pytest.fixture
    def ensemble(self):
        return EnsemblePredictor()

    def test_predict_returns_all_metrics(self, ensemble):
        """Test that ensemble returns all required metrics"""
        features = {
            "complexity_normalized": 0.6,
            "pri_attention_demand": 0.5,
            "trait_conscientiousness": 0.7,
        }
        task_data = {
            "estimated_duration": 60,
            "category": "WORK",
        }
        
        result = ensemble.predict(features, task_data)
        
        assert "completion_probability" in result
        assert "difficulty_level" in result
        assert "stress_level" in result
        assert "predicted_duration" in result

    def test_predictions_are_consistent(self, ensemble):
        """Test that predictions are consistent for same input"""
        features = {
            "complexity_normalized": 0.6,
            "pri_attention_demand": 0.5,
            "trait_conscientiousness": 0.7,
        }
        task_data = {
            "estimated_duration": 60,
            "category": "WORK",
        }
        
        result1 = ensemble.predict(features, task_data)
        result2 = ensemble.predict(features, task_data)
        
        assert result1["completion_probability"] == result2["completion_probability"]
        assert result1["difficulty_level"] == result2["difficulty_level"]


if __name__ == "__main__":
    pytest.main([__file__, "-v"])