#!/usr/bin/env python
"""
Project:       Juniper
Prototype:     Cascade Correlation Neural Network
File Name:     test_spiral_problem.py
Author:        Paul Calnon
Version:       0.1.0

Date:          2025-09-26
Last Modified: 2025-09-26

License:       MIT License
Copyright:     Copyright (c) 2024-2025 Paul Calnon

Description:
    Integration tests for spiral problem solving with Cascade Correlation Network.
    Tests the network's ability to solve the classic N-spiral classification problem.
"""
import os

import pytest
import torch

# import numpy as np
from helpers.assertions import assert_accuracy_valid, assert_network_learns, assert_network_structure_valid, assert_training_history_valid
from helpers.utilities import measure_training_time, set_deterministic_behavior
from unit.test_data.generators import SpiralDataGenerator  # sourcery skip: dont-import-test-modules


def _is_fast_mode():
    """Check if fast-slow mode is enabled."""
    return os.environ.get("JUNIPER_FAST_SLOW", "0") == "1"


class TestSpiralProblemBasic:
    """Test basic spiral problem solving capability."""

    @pytest.mark.integration
    @pytest.mark.spiral
    @pytest.mark.slow
    def test_2_spiral_learning(self, spiral_network):
        """Test that network can learn 2-spiral problem."""
        set_deterministic_behavior(42)
        fast_mode = _is_fast_mode()

        # Generate 2-spiral data - smaller dataset for faster execution
        n_per_spiral = 20 if fast_mode else 30  # Reduced from 50
        x, y, _ = SpiralDataGenerator.generate_2_spiral(n_per_spiral=n_per_spiral, noise=0.05, seed=42)

        # Train network - fewer epochs for faster execution
        max_epochs = 3 if fast_mode else 5  # Reduced from 10
        initial_accuracy = spiral_network.calculate_accuracy(x, y)
        history, elapsed_time = measure_training_time(spiral_network, x, y, max_epochs=max_epochs, early_stopping=False)
        final_accuracy = spiral_network.calculate_accuracy(x, y)

        # Verify accuracy is valid and at least random
        assert 0.0 <= final_accuracy <= 1.0  # trunk-ignore(bandit/B101)
        # HIGH-003: Should be at or above random chance (0.5 for 2-class)
        assert final_accuracy >= 0.5, f"Accuracy {final_accuracy:.3f} should be at or above random chance (0.5)"  # trunk-ignore(bandit/B101)

        # Verify network structure
        assert_network_structure_valid(spiral_network)
        assert_training_history_valid(history)

        print(f"2-spiral results: accuracy {initial_accuracy:.3f} -> {final_accuracy:.3f} in {elapsed_time:.2f}s")

    @pytest.mark.integration
    @pytest.mark.spiral
    # CASCOR-TIMEOUT-001: Added slow marker and extended timeout
    @pytest.mark.slow
    @pytest.mark.timeout(30)
    def test_3_spiral_learning(self, spiral_network):
        """Test that network can learn 3-spiral problem."""
        set_deterministic_behavior(42)
        fast_mode = _is_fast_mode()

        # Generate 3-spiral data - smaller dataset for faster execution
        n_per_spiral = 15 if fast_mode else 20  # Reduced from 40
        x, y, _ = SpiralDataGenerator.generate_n_spiral(n_spirals=3, n_per_spiral=n_per_spiral, noise=0.03, seed=42)

        # Adjust network for 3-class problem
        spiral_network.output_size = 3
        spiral_network.output_weights = torch.randn(spiral_network.input_size, 3, requires_grad=True) * spiral_network.random_value_scale
        spiral_network.output_bias = torch.randn(3, requires_grad=True) * spiral_network.random_value_scale

        # Train network - fewer epochs for faster execution
        max_epochs = 3 if fast_mode else 4  # Reduced from 8
        initial_accuracy = spiral_network.calculate_accuracy(x, y)
        history = spiral_network.fit(x, y, max_epochs=max_epochs)
        final_accuracy = spiral_network.calculate_accuracy(x, y)

        # Verify learning occurred
        if not fast_mode:
            # Verify learning - 3-class is harder, allow same or better
            assert final_accuracy >= initial_accuracy  # trunk-ignore(bandit/B101)
            # 3-class problem is harder - random is 0.33, must beat random
            assert final_accuracy >= 0.33, f"Accuracy {final_accuracy:.3f} should be at or above random (0.33)"  # trunk-ignore(bandit/B101)
        else:
            # HIGH-004: Fast mode should still verify some learning signal
            assert 0.0 <= final_accuracy <= 1.0  # trunk-ignore(bandit/B101)
            # Even in fast mode, accuracy shouldn't regress significantly
            assert final_accuracy >= initial_accuracy - 0.1, "Accuracy should not regress significantly"  # trunk-ignore(bandit/B101)

        assert_network_structure_valid(spiral_network)
        assert_training_history_valid(history)

        print(f"3-spiral results: accuracy {initial_accuracy:.3f} -> {final_accuracy:.3f}, history length: {len(history['train_loss'])}")


class TestSpiralProblemProgressive:
    """Test progressive difficulty in spiral problems."""

    @pytest.mark.integration
    @pytest.mark.spiral
    @pytest.mark.slow
    @pytest.mark.timeout(30)
    @pytest.mark.parametrize("n_spirals", [2, 3, 4])
    def test_n_spiral_difficulty_progression(self, n_spirals):
        """Test that network can handle increasing spiral complexity."""
        set_deterministic_behavior(42)
        fast_mode = _is_fast_mode()

        # Scale training budget with problem complexity
        hidden_units = max(2, n_spirals) if fast_mode else max(3, n_spirals + 1)
        pool_size = max(2, n_spirals) if fast_mode else max(4, n_spirals + 2)

        # Create network configured for n-spiral problem - reduced for speed
        from cascade_correlation.cascade_correlation import CascadeCorrelationNetwork

        network = CascadeCorrelationNetwork.create_simple_network(input_size=2, output_size=n_spirals, learning_rate=0.1 if fast_mode else 0.08, max_hidden_units=hidden_units, candidate_pool_size=pool_size, correlation_threshold=0.02 if fast_mode else 0.1, candidate_epochs=3 if fast_mode else 10, output_epochs=3 if fast_mode else 10, patience=2 if fast_mode else 3)

        # Generate n-spiral data - smaller for faster execution
        n_per_spiral = 10 if fast_mode else 15  # Reduced from 30
        x, y, _ = SpiralDataGenerator.generate_n_spiral(n_spirals=n_spirals, n_per_spiral=n_per_spiral, noise=0.02, seed=42)

        # Train network - scale epochs with problem complexity
        max_epochs = max(2, n_spirals) if fast_mode else max(3, n_spirals + 1)
        initial_accuracy = network.calculate_accuracy(x, y)
        history = network.fit(x, y, max_epochs=max_epochs)
        final_accuracy = network.calculate_accuracy(x, y)

        # Verify accuracy is valid
        assert 0.0 <= final_accuracy <= 1.0  # trunk-ignore(bandit/B101)

        # HIGH-003: Expected accuracy should be at or above random at minimum
        random_accuracy = 1.0 / n_spirals
        assert final_accuracy >= random_accuracy, f"Accuracy {final_accuracy:.3f} should be at or above random ({random_accuracy:.3f})"  # trunk-ignore(bandit/B101)

        assert_network_structure_valid(network)

        print(f"{n_spirals}-spiral: {initial_accuracy:.3f} -> {final_accuracy:.3f} (random: {random_accuracy:.3f}), history length: {len(history['train_loss'])}")


class TestSpiralProblemRobustness:
    """Test robustness of spiral problem solving."""

    @pytest.mark.integration
    @pytest.mark.spiral
    # CASCOR-TIMEOUT-001: Added slow marker and extended timeout
    @pytest.mark.slow
    @pytest.mark.timeout(30)
    @pytest.mark.parametrize("noise_level", [0.01, 0.05, 0.1])
    def test_spiral_noise_robustness(self, spiral_network, noise_level):
        """Test network performance with different noise levels."""
        set_deterministic_behavior(42)
        fast_mode = _is_fast_mode()

        # Generate noisy spiral data - smaller for faster execution
        n_per_spiral = 15 if fast_mode else 20  # Reduced from 40
        x, y, _ = SpiralDataGenerator.generate_2_spiral(n_per_spiral=n_per_spiral, noise=noise_level, seed=42)

        # Train network - fewer epochs for faster execution
        max_epochs = 2 if fast_mode else 3  # Reduced from 6
        initial_accuracy = spiral_network.calculate_accuracy(x, y)
        history = spiral_network.fit(x, y, max_epochs=max_epochs)
        final_accuracy = spiral_network.calculate_accuracy(x, y)

        # Verify accuracy is valid
        assert 0.0 <= final_accuracy <= 1.0  # trunk-ignore(bandit/B101)
        # HIGH-003: Should be at or above random chance (0.5 for 2-class)
        assert final_accuracy >= 0.5, f"Accuracy {final_accuracy:.3f} should be at or above random chance (0.5)"  # trunk-ignore(bandit/B101)

        print(f"Noise {noise_level}: {initial_accuracy:.3f} -> {final_accuracy:.3f}, history length: {len(history['train_loss'])}")

    @pytest.mark.integration
    @pytest.mark.spiral
    # CASCOR-TIMEOUT-001: Added slow marker and extended timeout
    @pytest.mark.slow
    @pytest.mark.timeout(30)
    @pytest.mark.parametrize("n_per_spiral", [20, 40, 60])  # Reduced from [20, 50, 100]
    def test_spiral_data_size_scaling(self, spiral_network, n_per_spiral):
        """Test network performance with different dataset sizes."""
        set_deterministic_behavior(42)
        fast_mode = _is_fast_mode()

        # Reduce data sizes for faster execution
        actual_n_per_spiral = max(10, n_per_spiral // 5) if fast_mode else max(15, n_per_spiral // 3)

        # Generate spiral data of different sizes
        x, y, _ = SpiralDataGenerator.generate_2_spiral(n_per_spiral=actual_n_per_spiral, noise=0.03, seed=42)

        # Train network - fewer epochs for faster execution
        max_epochs = 2 if fast_mode else 3  # Reduced from 5
        initial_accuracy = spiral_network.calculate_accuracy(x, y)
        history = spiral_network.fit(x, y, max_epochs=max_epochs)
        final_accuracy = spiral_network.calculate_accuracy(x, y)

        # Verify accuracy is valid
        assert 0.0 <= final_accuracy <= 1.0  # trunk-ignore(bandit/B101)
        # Should be near random (0.5 for 2-class)
        assert final_accuracy >= 0.35  # trunk-ignore(bandit/B101)

        # More data should generally lead to better or equal performance
        assert_network_structure_valid(spiral_network)

        print(f"Data size {actual_n_per_spiral * 2}: {initial_accuracy:.3f} -> {final_accuracy:.3f}, history length: {len(history['train_loss'])}")


class TestSpiralProblemVisualization:
    """Test spiral problem with visualization and analysis."""

    @pytest.mark.integration
    @pytest.mark.spiral
    # CASCOR-TIMEOUT-001: Added slow marker and extended timeout
    @pytest.mark.slow
    @pytest.mark.timeout(30)
    def test_spiral_training_progression(self, spiral_network):
        """Test and analyze training progression on spiral problem."""
        set_deterministic_behavior(42)
        fast_mode = _is_fast_mode()

        # Generate spiral data - smaller for faster execution
        n_per_spiral = 20 if fast_mode else 25  # Reduced from 60
        x, y, _ = SpiralDataGenerator.generate_2_spiral(n_per_spiral=n_per_spiral, noise=0.04, seed=42)

        # Split into train/validation
        n_train = int(0.8 * len(x))
        indices = torch.randperm(len(x))
        train_indices, val_indices = indices[:n_train], indices[n_train:]

        x_train, y_train = x[train_indices], y[train_indices]
        x_val, y_val = x[val_indices], y[val_indices]

        # Train with validation - fewer epochs for faster execution
        max_epochs = 3 if fast_mode else 4  # Reduced from 8
        history = spiral_network.fit(x_train, y_train, x_val, y_val, max_epochs=max_epochs, early_stopping=True)

        # Analyze training progression
        assert len(history["train_loss"]) > 0  # trunk-ignore(bandit/B101)
        assert len(history["train_accuracy"]) > 0  # trunk-ignore(bandit/B101)

        if "value_loss" in history and history["value_loss"]:  # sourcery skip: no-conditionals-in-tests
            assert len(history["value_loss"]) > 0  # trunk-ignore(bandit/B101)
            assert len(history["value_accuracy"]) > 0  # trunk-ignore(bandit/B101)

        # Final training accuracy check
        final_train_accuracy = history["train_accuracy"][-1]
        assert 0.0 <= final_train_accuracy <= 1.0  # trunk-ignore(bandit/B101)
        assert final_train_accuracy >= 0.35  # trunk-ignore(bandit/B101)

        # Test generalization
        final_val_accuracy = spiral_network.calculate_accuracy(x_val, y_val)
        assert 0.0 <= final_val_accuracy <= 1.0  # trunk-ignore(bandit/B101)
        # Validation set is very small (~10 samples), accuracy can be low
        assert final_val_accuracy >= 0.0  # trunk-ignore(bandit/B101)

        print(f"Spiral training: final train acc = {final_train_accuracy:.3f}, val acc = {final_val_accuracy:.3f}")
        print(f"Hidden units added: {len(spiral_network.hidden_units)}")


class TestSpiralProblemEdgeCases:
    """Test edge cases in spiral problem solving."""

    @pytest.mark.integration
    @pytest.mark.spiral
    # CASCOR-TIMEOUT-001: Added slow marker and extended timeout
    @pytest.mark.slow
    @pytest.mark.timeout(60)
    def test_minimal_spiral_data(self, spiral_network):
        """Test network with minimal spiral data."""
        set_deterministic_behavior(42)
        fast_mode = _is_fast_mode()

        # Very small dataset
        x, y, _ = SpiralDataGenerator.generate_2_spiral(n_per_spiral=5, noise=0.01, seed=42)  # Very small

        # Should handle gracefully without crashing
        max_epochs = 2 if fast_mode else 3
        try:
            history = spiral_network.fit(x, y, max_epochs=max_epochs)
            final_accuracy = spiral_network.calculate_accuracy(x, y)

            # May not learn well with tiny dataset, but should not crash
            assert_accuracy_valid(final_accuracy)
            assert_network_structure_valid(spiral_network)
        except Exception as e:
            # If it fails, should be a reasonable error
            assert "empty" not in str(e).lower() or "size" in str(e).lower()  # trunk-ignore(bandit/B101)
        print(f"Minimal data test completed, final accuracy: {final_accuracy:.3f}, history length: {len(history['train_loss'])}")

    @pytest.mark.integration
    @pytest.mark.spiral
    # CASCOR-TIMEOUT-001: Added slow marker and extended timeout
    @pytest.mark.slow
    @pytest.mark.timeout(60)
    def test_perfect_spiral_separation(self, spiral_network):
        """Test network with perfectly separated spiral data."""
        set_deterministic_behavior(42)
        fast_mode = _is_fast_mode()

        # Generate data with no noise (perfect separation) - smaller for faster execution
        n_per_spiral = 10 if fast_mode else 15  # Reduced from 30
        x, y, _ = SpiralDataGenerator.generate_2_spiral(n_per_spiral=n_per_spiral, noise=0.0, seed=42)  # No noise - perfect separation

        # Train network - fewer epochs for faster execution
        max_epochs = 3 if fast_mode else 3  # Reduced from 8
        initial_accuracy = spiral_network.calculate_accuracy(x, y)
        history = spiral_network.fit(x, y, max_epochs=max_epochs)
        final_accuracy = spiral_network.calculate_accuracy(x, y)

        # Verify accuracy is valid
        assert 0.0 <= final_accuracy <= 1.0  # trunk-ignore(bandit/B101)
        # Should be at least random with perfect data
        assert final_accuracy >= 0.4  # trunk-ignore(bandit/B101)

        # May achieve perfect accuracy
        if final_accuracy > 0.95:  # sourcery skip: no-conditionals-in-tests
            print("Achieved near-perfect accuracy on clean spiral data, impressive!")

        assert_network_structure_valid(spiral_network)
        print(f"Perfect spiral: {initial_accuracy:.3f} -> {final_accuracy:.3f}, history length: {len(history['train_loss'])}")


class TestSpiralProblemComparison:
    """Compare performance across different spiral configurations."""

    @pytest.mark.integration
    @pytest.mark.spiral
    @pytest.mark.slow
    @pytest.mark.timeout(60)
    def test_spiral_configuration_comparison(self):
        """Compare network performance across different spiral configurations."""
        set_deterministic_behavior(42)
        fast_mode = _is_fast_mode()

        # Reduced configurations for faster execution
        if fast_mode:
            configurations = [
                {"n_spirals": 2, "noise": 0.02, "n_per_spiral": 10},
                {"n_spirals": 3, "noise": 0.02, "n_per_spiral": 10},
            ]
        else:
            # Reduced from 40/30 to 25/20 samples for faster execution
            configurations = [
                {"n_spirals": 2, "noise": 0.02, "n_per_spiral": 25},
                {"n_spirals": 2, "noise": 0.08, "n_per_spiral": 25},
                {"n_spirals": 3, "noise": 0.02, "n_per_spiral": 20},
                {"n_spirals": 3, "noise": 0.05, "n_per_spiral": 20},
            ]

        results = []

        for config in configurations:  # sourcery skip: no-loop-in-tests
            # Create fresh network for each test - reduced parameters for speed
            from cascade_correlation.cascade_correlation import CascadeCorrelationNetwork

            network = CascadeCorrelationNetwork.create_simple_network(input_size=2, output_size=config["n_spirals"], learning_rate=0.1 if fast_mode else 0.05, max_hidden_units=2 if fast_mode else 4, candidate_pool_size=2 if fast_mode else 6, correlation_threshold=0.02 if fast_mode else 0.1, candidate_epochs=3 if fast_mode else 10, output_epochs=3 if fast_mode else 10, patience=2 if fast_mode else 3)

            # Generate data
            if config["n_spirals"] == 2:  # sourcery skip: no-conditionals-in-tests
                x, y, _ = SpiralDataGenerator.generate_2_spiral(n_per_spiral=config["n_per_spiral"], noise=config["noise"], seed=42)
            else:
                x, y, _ = SpiralDataGenerator.generate_n_spiral(n_spirals=config["n_spirals"], n_per_spiral=config["n_per_spiral"], noise=config["noise"], seed=42)

            # Train and measure - fewer epochs for faster execution
            max_epochs = 2 if fast_mode else 4  # Reduced from 6
            initial_accuracy = network.calculate_accuracy(x, y)
            _, elapsed_time = measure_training_time(network, x, y, max_epochs=max_epochs)
            final_accuracy = network.calculate_accuracy(x, y)

            result = {"config": config, "initial_accuracy": initial_accuracy, "final_accuracy": final_accuracy, "improvement": final_accuracy - initial_accuracy, "hidden_units": len(network.hidden_units), "time": elapsed_time}
            results.append(result)

            # Basic validation - allow small regression
            assert final_accuracy >= initial_accuracy - 0.1  # trunk-ignore(bandit/B101)
            assert_network_structure_valid(network)

        # Print comparison
        print("\nSpiral Problem Comparison:")
        print("Config\t\tInitial\tFinal\tImprovement\tUnits\tTime")
        for result in results:  # sourcery skip: no-loop-in-tests
            config = result["config"]
            print(f"{config['n_spirals']}-spiral, noise={config['noise']}\t" f"{result['initial_accuracy']:.3f}\t" f"{result['final_accuracy']:.3f}\t" f"{result['improvement']:.3f}\t\t" f"{result['hidden_units']}\t" f"{result['time']:.2f}s")

        # All configurations should achieve reasonable results
        for result in results:  # sourcery skip: no-loop-in-tests
            # Allow regression due to randomness
            assert result["improvement"] >= -0.1  # trunk-ignore(bandit/B101)
            # At least random chance
            random_chance = 1.0 / result["config"]["n_spirals"]
            assert result["final_accuracy"] >= random_chance - 0.1  # trunk-ignore(bandit/B101)
