import numpy as np
import pytest

from zeromodel.core import ZeroModel
from zeromodel.memory import ZeroMemory
from zeromodel.tools.spatial_optimizer import SpatialOptimizer


class TestSpatialOptimizer:
    """Test suite for the SpatialOptimizer class (ZeroModel's Spatial Calculus)."""
    
    @pytest.fixture(autouse=True)
    def setup_optimizer(self):
        """Setup optimizer for each test."""
        self.optimizer = SpatialOptimizer(Kc=2, Kr=2, alpha=0.9)
    

    def test_basic_functionality(self):
        """Test basic methods of SpatialOptimizer with simple data."""
        # 1. Create simple score matrix (3 sources, 2 metrics)
        X = np.array([
            [0.8, 0.2],  # Source 0: high metric 0, low metric 1
            [0.2, 0.8],  # Source 1: low metric 0, high metric 1
            [0.5, 0.5]   # Source 2: medium on both
        ])
        
        # 2. Test top_left_mass calculation
        # With identity ordering, top-left is [0.8, 0.2] for first row
        optimizer = SpatialOptimizer(Kc=2, Kr=2, alpha=0.9)
        mass = optimizer.top_left_mass(X)
        expected = 0.8 * (0.9 ** 0) + 0.2 * (0.9 ** 1) + 0.2 * (0.9 ** 1) + 0.8 * (0.9 ** 2)
        assert np.isclose(mass, expected), f"Expected {expected}, got {mass}"
        
        # 3. Test column ordering
        u = np.array([0.7, 0.3])  # Metric 0 is more interesting
        cidx, Xc = optimizer.order_columns(X, u)
        assert np.array_equal(cidx, [0, 1]), "Columns should be in original order"
        
        # Reverse interest (metric 1 more interesting)
        u = np.array([0.3, 0.7])
        cidx, Xc = optimizer.order_columns(X, u)
        assert np.array_equal(cidx, [1, 0]), "Columns should be reversed"
        assert np.array_equal(Xc[:, 0], X[:, 1]), "Column 1 should be first"
        
        # 4. Test row ordering
        w = np.array([0.7, 0.3])  # Weight metric 0 more heavily
        ridx, Y = optimizer.order_rows(X, w)
        assert np.array_equal(ridx, [0, 2, 1]), "Source 0 should be highest ranked"
        
        # 5. Test phi_transform
        u = np.array([0.3, 0.7])  # Metric 1 more interesting
        w = np.array([0.3, 0.7])  # Weight metric 1 more heavily
        Y, ridx, cidx = optimizer.phi_transform(X, u, w)
        
        # Verify top-left has highest value
        assert Y[0, 0] >= Y[0, 1], "Top-left should be highest in first row"
        assert Y[0, 0] >= Y[1, 0], "Top-left should be highest in first column"
        
        print("✅ Basic functionality test completed successfully")

    def test_metric_graph_and_layout(self):
        """Test metric graph construction and canonical layout."""
        # 1. Create column order history (3 metrics, 4 time steps)
        col_orders = [
            np.array([0, 1, 2]),  # Time 0: natural order
            np.array([1, 0, 2]),  # Time 1: metrics 0 and 1 swapped
            np.array([1, 2, 0]),  # Time 2: metric 0 moves to end
            np.array([2, 1, 0])   # Time 3: reverse order
        ]
        
        # 2. Test metric graph
        W = self.optimizer.metric_graph(col_orders)
        assert W.shape == (3, 3), "Graph should be 3x3 for 3 metrics"
        assert np.all(W >= 0) and np.all(W <= 1), "Edge weights should be in [0,1]"
        
        # Verify symmetry (should be symmetric graph)
        assert np.allclose(W, W.T), "Graph should be symmetric"
        
        # Metric 0 and 1 appear together more than 0 and 2
        assert W[0, 1] > W[0, 2], "Metrics 0 and 1 should have stronger connection"
        
        # 3. Test canonical layout
        layout = self.optimizer.compute_canonical_layout(W)
        assert len(layout) == 3, "Layout should have 3 metrics"
        assert set(layout) == {0, 1, 2}, "Layout should contain all metrics"
        
        # In this pattern, metrics 1 and 2 should be closer in layout
        # than 0 and 2 (based on co-occurrence patterns)
        pos = {m: i for i, m in enumerate(layout)}
        assert abs(pos[1] - pos[2]) < abs(pos[0] - pos[2]), \
            "Metrics 1 and 2 should be closer in canonical layout"
        
        print("✅ Metric graph and layout test completed successfully")

    def test_learn_weights_simple(self):
        """Test weight learning with a simple, controlled dataset."""
        # Create dataset
        # Replace with a proper array structure
        series = [np.array([[0.9, np.random.random()], 
            [0.1, np.random.random()]]) for _ in range(10)]
        
        # Learn weights
        optimizer = SpatialOptimizer(Kc=2, Kr=2, alpha=0.9)
        w_star = optimizer.learn_weights(series)
        
        # Verify weights are reasonable, not that they always improve TL mass
        assert np.all(w_star >= 0)
        assert np.isclose(np.sum(w_star), 1.0, atol=1e-3)
        
    def test_learn_weights_with_column_mean(self):
        """Test weight learning with u_mode='col_mean'."""
        # 1. Create time series with drifting metric importance
        series = []
        for t in range(15):
            # Early on, metric 0 is more important
            # Later, metric 1 becomes more important
            importance_factor = 0.2 + 0.8 * (t / 14)
            
            X = np.zeros((8, 2))
            for i in range(8):
                # Metric 0: decreasing relevance over time
                X[i, 0] = (0.9 - 0.8 * (t / 14)) * (0.9 - i * 0.1)
                # Metric 1: increasing relevance over time
                X[i, 1] = importance_factor * (0.9 - i * 0.1)
            
            series.append(X)
        
        # 2. Initialize optimizer with column mean mode
        optimizer = SpatialOptimizer(Kc=2, Kr=4, alpha=0.9, u_mode="col_mean")
        
        # 3. Learn weights
        w_star = optimizer.learn_weights(series, iters=100)
        
        # 4. Verify learned weights balance both metrics
        # Since importance shifts over time, we expect weights to be more balanced
        print(f"Learned weights with column mean: {w_star}")
        assert 0.3 < w_star[0] < 0.7, "Weight for metric 0 should be moderate"
        assert 0.3 < w_star[1] < 0.7, "Weight for metric 1 should be moderate"
        assert np.isclose(w_star[0] + w_star[1], 1.0, atol=1e-5), "Weights should sum to ~1"
        
        print("✅ Column mean mode weight learning test completed successfully")

    def test_end_to_end_optimization(self, tmp_path):
        """Test complete end-to-end optimization workflow with proper expectations."""
        # 1. Generate realistic training history
        metric_names = ["loss", "val_loss", "acc", "val_acc", "grad_norm"]
        historical_data = []
        
        # Simulate 30 training epochs
        for epoch in range(30):
            # Create realistic patterns
            if epoch < 15:
                # Initial training phase
                loss = 1.0 - (epoch * 0.04)
                val_loss = 1.0 - (epoch * 0.035)
                acc = 0.4 + (epoch * 0.02)
                val_acc = 0.35 + (epoch * 0.018)
            else:
                # Stronger overfitting phase
                loss = 0.4 - ((epoch-15) * 0.01)
                val_loss = 0.5 + ((epoch-15) * 0.025)
                acc = 0.7 + ((epoch-15) * 0.005)
                val_acc = 0.65 - ((epoch-15) * 0.007)
            
            grad_norm = 0.8 - (epoch * 0.015)
            
            # Create score matrix (5 sources, 5 metrics)
            X = np.array([
                [loss, val_loss, acc, val_acc, grad_norm],
                [loss-0.05, val_loss-0.05, acc+0.05, val_acc+0.05, grad_norm-0.05],
                [loss-0.1, val_loss-0.1, acc+0.1, val_acc+0.1, grad_norm-0.1],
                [loss-0.15, val_loss-0.15, acc+0.15, val_acc+0.15, grad_norm-0.15],
                [loss-0.2, val_loss-0.2, acc+0.2, val_acc+0.2, grad_norm-0.2]
            ])
            
            historical_data.append(X)
        
        # 2. Run end-to-end optimization
        optimizer = SpatialOptimizer(Kc=5, Kr=3)
        optimizer.apply_optimization(historical_data)
        
        # 3. Verify results with proper expectations
        assert optimizer.metric_weights is not None, "Metric weights should be learned"
        assert len(optimizer.metric_weights) == 5, "Should have weights for all metrics"
        assert np.isclose(np.linalg.norm(optimizer.metric_weights), 1.0, atol=5e-4), "Weights should be normalized"
        
        print(f"Learned metric weights: {optimizer.metric_weights}")
        print(f"Metric names: {metric_names}")
        
        # In overfitting detection, validation metrics should have higher weights than training metrics
        # This is the key principle of ZeroModel's spatial calculus - the intelligence is in the structure
        assert optimizer.metric_weights[1] > optimizer.metric_weights[0], \
            "Validation loss should have higher weight than training loss during overfitting"
        
        assert optimizer.metric_weights[3] < optimizer.metric_weights[2], \
            "Validation accuracy should have lower weight than training accuracy during overfitting"
        
        # Check that the weights make sense for overfitting detection
        # Validation loss (index 1) and training accuracy (index 2) should be the top signals
        top_indices = np.argsort(-optimizer.metric_weights)[:2]
        assert 1 in top_indices, "Validation loss should be among top 2 signals"
        assert 2 in top_indices, "Training accuracy should be among top 2 signals"
        
        # Verify the canonical layout prioritizes overfitting signals
        assert optimizer.canonical_layout is not None
        # In canonical layout, lower index = higher priority
        val_loss_idx = np.where(optimizer.canonical_layout == 1)[0][0]
        loss_idx = np.where(optimizer.canonical_layout == 0)[0][0]
        assert val_loss_idx < loss_idx, "Validation loss should be prioritized before training loss"

    def test_edge_cases(self):
        """Test SpatialOptimizer edge cases and error handling."""
        # 1. Test with invalid parameters
        with pytest.raises(ValueError, match="Kc must be positive"):
            SpatialOptimizer(Kc=0)
        
        # Test Kr AFTER Kc (since Kc is validated first)
        with pytest.raises(ValueError, match="Kr must be positive"):
            SpatialOptimizer(Kc=2, Kr=0)  # Provide valid Kc first
               
        print("✅ Edge cases test completed successfully")


    def test_real_world_integration(self, tmp_path):
        """Test integration with ZeroModel in a realistic scenario."""
        # 1. Set up ZeroMemory to collect metrics
        metric_names = ["loss", "val_loss", "acc", "val_acc", "lr", "grad_norm"]
        zeromemory = ZeroMemory(
            metric_names=metric_names,
            buffer_steps=100,
            tile_size=8,
            selection_k=24
        )
        
        # 2. Simulate training process
        for epoch in range(25):
            # Create realistic metrics
            if epoch < 15:
                # Initial training phase
                metrics = {
                    "loss": 1.0 - epoch * 0.04,
                    "val_loss": 1.0 - epoch * 0.035,
                    "acc": 0.4 + epoch * 0.02,
                    "val_acc": 0.35 + epoch * 0.018,
                    "lr": 0.1,
                    "grad_norm": 0.8 - epoch * 0.02
                }
            else:
                # FIXED: Stronger overfitting phase
                metrics = {
                    "loss": 0.4 - (epoch-15) * 0.01,
                    "val_loss": 0.5 + (epoch-15) * 0.025,  # FIXED: Increased slope
                    "acc": 0.7 + (epoch-15) * 0.005,
                    "val_acc": 0.65 - (epoch-15) * 0.007,  # FIXED: Steeper decline
                    "lr": 0.01,
                    "grad_norm": 0.5 - (epoch-15) * 0.01
                }
            
            # Log metrics
            zeromemory.log(step=epoch, metrics=metrics)
        
        # 3. Extract historical data for optimization
        historical_data = []
        for epoch in range(25):
            # Get the score matrix for this epoch
            # For this test, we'll reconstruct it from the metrics
            X = np.zeros((1, len(metric_names)))
            for i, name in enumerate(metric_names):
                # Use the last logged value
                X[0, i] = zeromemory.buffer_values[zeromemory.buffer_head - 1, i]
            historical_data.append(X)
        
        # 4. Optimize spatial organization
        optimizer = SpatialOptimizer(Kc=6, Kr=5)
        optimizer.apply_optimization(historical_data)
        
        # 5. Create ZeroModel with optimized layout
        # FIXED: Using correct API for ZeroModel initialization
        zeromodel = ZeroModel(metric_names=metric_names)
        
        # Process current metrics
        current_metrics = {
            "loss": 0.35,
            "val_loss": 0.65,
            "acc": 0.75,
            "val_acc": 0.60,
            "lr": 0.01,
            "grad_norm": 0.45
        }
        
        # Convert to score matrix
        score_matrix = np.array([list(current_metrics.values())])

        # Prepare with optimized layout
        if optimizer.canonical_layout is not None:
            ordered_metrics = [metric_names[i] for i in optimizer.canonical_layout]
            sql_query = f"SELECT * FROM virtual_index ORDER BY {', '.join(ordered_metrics)} DESC"
            zeromodel.prepare(
                score_matrix=score_matrix,
                sql_query=sql_query
            )

        # 7. Verify decision (metric 0)
        doc_idx, confidence = zeromodel.get_decision_by_metric(0)
        assert doc_idx == 0, "Should select the only document"
        assert 0 <= confidence <= 1.0, "Confidence should be in [0,1]"

        # 8. Extract critical tile via encoder
        from zeromodel.vpm.encoder import VPMEncoder
        tile = VPMEncoder('float32').get_critical_tile(zeromodel.sorted_matrix, tile_size=3)
        assert isinstance(tile, bytes), "Tile should be bytes"
        assert len(tile) > 4, "Tile should have header and data"

        print("✅ Real-world integration test completed successfully")

