import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from collections import Counter

# ============================================================
# THE META-PROBLEM: Encode sorting "algorithm" in weights
# ============================================================

class SortingMatrixLearner(nn.Module):
    """
    A network that learns to sort by encoding comparison patterns
    in its parameter matrices. The "algorithm" emerges from training.
    """
    def __init__(self, input_len=4, hidden_dim=16):
        super().__init__()
        self.input_len = input_len
        
        # These matrices will "learn" the sorting algorithm
        # Matrix A: learns comparison logic
        self.W_comp = nn.Parameter(torch.randn(input_len, hidden_dim) * 0.5)
        self.b_comp = nn.Parameter(torch.zeros(hidden_dim))
        
        # Matrix B: learns ranking/scoring
        self.W_rank = nn.Parameter(torch.randn(hidden_dim, input_len) * 0.3)
        self.b_rank = nn.Parameter(torch.zeros(input_len))
        
        # Matrix C: learns ordering reconstruction
        self.W_order = nn.Parameter(torch.randn(input_len, input_len) * 0.2)
        self.b_order = nn.Parameter(torch.zeros(input_len))
        
        # Temperature for softmax (learnable)
        self.temp = nn.Parameter(torch.tensor(1.0))
        
    def forward(self, x):
        # x shape: (batch, input_len)
        batch_size = x.shape[0]

        # Learned comparison signals
        h = torch.tanh(x @ self.W_comp + self.b_comp)

        # Learned ranking scores for each input element.
        # W_rank maps hidden features back to one score per input position.
        scores = h @ self.W_rank + self.b_rank

        # Convert scores into a soft rank for each element by comparing all
        # element pairs. Smaller scores should map to earlier output positions.
        temperature = torch.abs(self.temp) + 1e-3
        pairwise = (scores.unsqueeze(2) - scores.unsqueeze(1)) / temperature
        soft_ranks = torch.sigmoid(pairwise).sum(dim=-1)

        # Build a soft permutation matrix that assigns each element to an
        # output position. W_order acts as a learned bias over assignments.
        positions = torch.arange(self.input_len, device=x.device, dtype=x.dtype)
        order_bias = self.W_order + self.b_order.unsqueeze(0)
        perm_logits = -(soft_ranks.unsqueeze(2) - positions.view(1, 1, -1)) ** 2
        perm_logits = perm_logits + order_bias.unsqueeze(0)
        perm = torch.softmax(perm_logits, dim=1)

        # Reconstruct the sorted vector as a weighted sum over input elements
        # for each target output position.
        sorted_vals = torch.matmul(perm.transpose(1, 2), x.unsqueeze(-1)).squeeze(-1)
        return sorted_vals

# ============================================================
# DIFFERENTIABLE SORTING LAYER - Learn comparisons
# ============================================================

class LearnableSortLayer(nn.Module):
    """
    A layer that learns comparison patterns.
    The weights encode which elements to compare and how to merge.
    """
    def __init__(self, width):
        super().__init__()
        # These weights will "learn" a sorting network pattern
        self.merge_weights = nn.Parameter(torch.randn(width, width) * 0.1)
        self.compare_mask = nn.Parameter(torch.eye(width))  # Initially identity
        
    def forward(self, x):
        # Apply learned comparison patterns
        # Elements learn to compare with each other through these weights
        diff = x.unsqueeze(2) - x.unsqueeze(1)  # All pairs comparison
        learned_compare = torch.sigmoid(diff @ self.compare_mask)
        
        # Merge based on learned patterns
        output = x + (learned_compare @ self.merge_weights).mean(dim=2)
        return output

# ============================================================
# COMPETITION: Discover which sorting "algorithm" emerges
# ============================================================

def generate_sorting_pairs(n_elements, n_samples=2000):
    """Generate input-output pairs for sorting training."""
    X, y = [], []
    for _ in range(n_samples):
        # Random sequence
        seq = torch.rand(n_elements)
        # Sorted sequence (ground truth)
        sorted_seq = torch.sort(seq)[0]
        X.append(seq)
        y.append(sorted_seq)
    return torch.stack(X), torch.stack(y)

def train_and_extract_algorithm(n_elements=4, epochs=500):
    """Train network and analyze what algorithm emerged."""
    
    print("=" * 70)
    print("PHASE 1: Training network to encode sorting in weight matrices")
    print("=" * 70)
    
    X_train, y_train = generate_sorting_pairs(n_elements, n_samples=3000)
    
    model = SortingMatrixLearner(input_len=n_elements, hidden_dim=32)
    optimizer = optim.Adam(model.parameters(), lr=0.02)
    criterion = nn.MSELoss()
    
    # Training loop - optimizer writes sorting "algorithm" into weights
    losses = []
    for epoch in range(epochs):
        optimizer.zero_grad()
        output = model(X_train)
        loss = criterion(output, y_train)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
        
        if epoch % 100 == 0:
            print(f"  Epoch {epoch:4d} | Loss: {loss.item():.6f}")
    
    print(f"\n✓ Final loss: {losses[-1]:.6f}")
    return model, losses

# ============================================================
# ANALYSIS: Reverse-engineer the discovered algorithm
# ============================================================

def analyze_learned_algorithm(model, X_test):
    """Extract what the weight matrices learned."""
    
    print("\n" + "=" * 70)
    print("PHASE 2: Analyzing weight matrices to discover algorithm")
    print("=" * 70)
    
    print("\n📊 WEIGHT MATRIX ANALYSIS:")
    print("-" * 50)
    
    # Analyze W_rank - what priorities did it learn?
    rank_weights = model.W_rank.data
    print(f"\nRank Weight Matrix (hidden→score):")
    print(f"  Shape: {rank_weights.shape}")
    print(f"  Mean: {rank_weights.mean():.4f}, Std: {rank_weights.std():.4f}")
    
    # Find which input positions have highest influence on ranking
    influence = rank_weights.abs().mean(dim=0)
    print(f"\n  Input position influence scores: {influence.tolist()}")
    most_important = influence.argmax().item()
    print(f"  → Network focuses most on position {most_important}")
    
    # Analyze W_comp - what comparisons did it learn?
    comp_weights = model.W_comp.data
    print(f"\nComparison Weight Matrix (input→hidden):")
    print(f"  Shape: {comp_weights.shape}")
    
    # What patterns emerge in comparison weights?
    comp_patterns = comp_weights @ comp_weights.T
    print(f"  Comparison patterns (approx):")
    
    # Analyze W_order - what ordering did it learn?
    order_weights = model.W_order.data
    print(f"\nOrder Reconstruction Matrix:")
    print(f"  Dominant pattern (diagonal vs off-diagonal):")
    diag = order_weights.diagonal().mean()
    off_diag = (order_weights.sum() - order_weights.trace()) / (order_weights.numel() - order_weights.shape[0])
    print(f"    Diagonal avg: {diag:.4f}, Off-diagonal avg: {off_diag:.4f}")
    
    return {
        'rank_weights': rank_weights,
        'comp_weights': comp_weights,
        'order_weights': order_weights
    }

def test_sorting_accuracy(model, n_tests=100):
    """Test how well the learned weights sort."""
    
    print("\n" + "=" * 70)
    print("PHASE 3: Testing sorting accuracy")
    print("=" * 70)
    
    correct = 0
    total = 0
    errors = []
    
    for _ in range(n_tests):
        X_test = torch.rand(4)
        with torch.no_grad():
            sorted_pred = model(X_test.unsqueeze(0)).squeeze()
        sorted_true = torch.sort(X_test)[0]
        
        # Check if order is correct (allow small numerical errors)
        if torch.allclose(sorted_pred, sorted_true, atol=0.1):
            correct += 1
        else:
            error = (sorted_pred - sorted_true).abs().mean()
            errors.append(error.item())
        total += 1
    
    accuracy = correct / total * 100
    avg_error = np.mean(errors) if errors else 0
    
    print(f"\n  Accuracy: {accuracy:.1f}%")
    print(f"  Average error: {avg_error:.4f}")
    
    # Show some examples
    print("\n  Sample predictions:")
    for i in range(3):
        X = torch.rand(4)
        with torch.no_grad():
            pred = model(X.unsqueeze(0)).squeeze()
        true = torch.sort(X)[0]
        print(f"    Input: {torch.round(X * 100) / 100}")
        print(f"    Pred:  {torch.round(pred * 100) / 100}")
        print(f"    True:  {torch.round(true * 100) / 100}")
        print()

def hypothesize_algorithm(model):
    """Infer what algorithm the network discovered."""
    
    print("=" * 70)
    print("PHASE 4: Algorithm hypothesis - What did it discover?")
    print("=" * 70)
    
    rank_weights = model.W_rank.data
    
    # Check if network learned a "position-based" sort
    # (treating indices as priorities - like selection sort)
    diag_dominance = rank_weights.diagonal().abs().mean()
    off_diag_avg = (rank_weights.sum() - rank_weights.trace()) / (rank_weights.numel() - rank_weights.shape[0])
    
    print("\n🔬 Algorithm Detection:")
    print(f"  Diagonal weight dominance: {diag_dominance:.4f}")
    print(f"  Off-diagonal weight avg:   {off_diag_avg:.4f}")
    
    if diag_dominance > off_diag_avg * 2:
        print("\n  🏆 DISCOVERED: POSITION-WEIGHTED SORT")
        print("     Network learned that early positions should rank higher.")
        print("     This is similar to SELECTION SORT strategy:")
        print("     - Find minimum, place at position 0")
        print("     - Find next minimum, place at position 1...")
        print("     - The weight matrix encodes: 'smaller indices = higher priority'")
    else:
        print("\n  🏆 DISCOVERED: VALUE-BASED RANKING")
        print("     Network learned actual value comparisons.")
        print("     This is similar to RANKING SORT strategy:")
        print("     - Score each element by its value")
        print("     - Reconstruct order based on learned scores")
    
    # Check what temperature revealed
    print(f"\n  Learned temperature: {model.temp.item():.4f}")
    if model.temp.item() > 0.5:
        print("     → High temperature: fuzzy/soft ranking (probabilistic)")
    else:
        print("     → Low temperature: sharp/rigid ranking (deterministic)")
    
    return {
        'type': 'position-weighted' if diag_dominance > off_diag_avg * 2 else 'value-based',
        'temp': model.temp.item()
    }

def visualize_weight_matrix(matrix, name):
    """Create ASCII visualization of weight matrix."""
    print(f"\n{name} (ASCII heatmap):")
    print("  ", end="")
    for i in range(matrix.shape[1]):
        print(f"{i:>5}", end="")
    print()
    print("  " + "-" * (matrix.shape[1] * 5 + 1))
    
    for i in range(matrix.shape[0]):
        print(f"{i}|", end=" ")
        for j in range(matrix.shape[1]):
            val = matrix[i, j].item()
            if val > 0.1:
                print("█", end="   ")
            elif val < -0.1:
                print("▓", end="   ")
            elif abs(val) > 0.05:
                print("░", end="   ")
            else:
                print(" ", end="   ")
        print()
    print("  Legend: █ positive, ▓ negative, ░ small, (space) near-zero")

# ============================================================
# MAIN: Run the experiment
# ============================================================

if __name__ == "__main__":
    print("\n" + "=" * 70)
    print("🧠 THOUGHT EXPERIMENT: Can weights discover a sorting algorithm?")
    print("=" * 70)
    print("""
    HYPOTHESIS: Through gradient descent over (X_train, y_train) pairs,
    the optimizer will update weight matrices until they ENCODE a sorting
    strategy - effectively "writing" an algorithm into the parameters.
    """)
    
    # Train
    model, losses = train_and_extract_algorithm(n_elements=4, epochs=500)
    
    # Analyze
    weights = analyze_learned_algorithm(model, X_test=None)
    
    # Visualize matrices
    visualize_weight_matrix(model.W_rank.data, "Rank Weight Matrix (W_rank)")
    visualize_weight_matrix(model.W_comp.data[:4, :8], "Comparison Matrix (W_comp, first rows)")
    
    # Test
    test_sorting_accuracy(model, n_tests=100)
    
    # Hypothesize
    hypothesis = hypothesize_algorithm(model)
    
    print("\n" + "=" * 70)
    print("CONCLUSION")
    print("=" * 70)
    print(f"""
    The network discovered a {hypothesis['type']} sorting strategy!
    
    The "algorithm" is encoded entirely in three weight matrices:
    
    ┌─────────────────────────────────────────────────────────┐
    │  W_rank: Learns which input positions matter most      │
    │          (Encodes: "where should I look for min?")     │
    │                                                         │
    │  W_comp:  Learns comparison patterns                   │
    │          (Encodes: "how do I compare two values?")     │
    │                                                         │
    │  W_order: Learns ordering reconstruction               │
    │          (Encodes: "how do I build the sorted output?")│
    └─────────────────────────────────────────────────────────┘
    
    The optimizer wrote the solution into the matrices through
    {500} gradient updates over truth table pairs (X_train, y_train).
    
    The discovered "algorithm" is not explicit code - it's a 
    distributed representation in continuous parameter space.
    """)