import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

# ============================================================
# 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
        
        # FIXED: Weight matrices must allow x @ W to work
        # x shape: (batch, input_len) = (batch, 4)
        # We need W_rank: (input_len, hidden_dim) = (4, 16) → result (batch, 16)
        
        # Matrix A: learns comparison logic (encodes how to score each element)
        self.W_comp = nn.Parameter(torch.randn(input_len, hidden_dim) * 0.5)  # (4, 16)
        self.b_comp = nn.Parameter(torch.zeros(hidden_dim))  # (16,)
        
        # Matrix B: learns ranking/scoring (encodes which value is "smallest")
        self.W_rank = nn.Parameter(torch.randn(input_len, hidden_dim) * 0.3)  # (4, 16)
        self.b_rank = nn.Parameter(torch.zeros(input_len))  # (4,)
        
        # Matrix C: learns ordering reconstruction
        self.W_order = nn.Parameter(torch.randn(hidden_dim, input_len) * 0.2)  # (16, 4)
        self.b_order = nn.Parameter(torch.zeros(input_len))  # (4,)
        
        # Temperature for softmax (learnable)
        self.temp = nn.Parameter(torch.tensor(1.0))
        
    def forward(self, x):
        # x shape: (batch, input_len)
        
        # Learned comparison signals - each input position scores itself
        h = torch.tanh(x @ self.W_comp + self.b_comp)  # (batch, 16)
        
        # Learned ranking scores - direct mapping to element importance
        ranks = x @ self.W_rank + self.b_rank  # (batch, 4) — score per element
        
        # Reconstruct sorted output using learned attention
        # Elements with LOWER learned ranks should come first
        neg_ranks = -ranks  # Negate so smallest rank → highest attention
        attn = torch.softmax(neg_ranks / (torch.abs(self.temp) + 0.1), dim=-1)
        
        # Apply attention-weighted sum to reorder
        # attn sums to 1, used as weights to reconstruct
        sorted_vals = attn * x  # Element-wise: weighted by learned priority
        
        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
        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):
        seq = torch.rand(n_elements)
        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=16)
    optimizer = optim.Adam(model.parameters(), lr=0.02)
    criterion = nn.MSELoss()
    
    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

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)
    
    rank_weights = model.W_rank.data
    print(f"\nRank Weight Matrix (input→score):")
    print(f"  Shape: {rank_weights.shape}")
    print(f"  Mean: {rank_weights.mean():.4f}, Std: {rank_weights.std():.4f}")
    
    # Analyze diagonal dominance (position-based vs value-based)
    diag = rank_weights.diagonal().abs()
    off_diag_mask = ~torch.eye(rank_weights.shape[0], dtype=bool)
    off_diag = rank_weights.abs()[off_diag_mask]
    
    print(f"\n  Diagonal (position) weights: {diag.tolist()}")
    print(f"  Off-diagonal (value) avg: {off_diag.mean():.4f}")
    
    influence = rank_weights.abs().mean(dim=1)
    print(f"  Input position influence: {influence.tolist()}")
    most_important = influence.argmax().item()
    print(f"  → Network focuses most on position {most_important}")
    
    comp_weights = model.W_comp.data
    print(f"\nComparison Weight Matrix (input→hidden):")
    print(f"  Shape: {comp_weights.shape}")
    
    order_weights = model.W_order.data
    print(f"\nOrder Reconstruction Matrix:")
    print(f"  Dominant diagonal: {order_weights.diagonal().mean():.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
        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}")
    
    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: {X.round(2).tolist()}")
        print(f"    Pred:  {pred.round(2).tolist()}")
        print(f"    True:  {true.round(2).tolist()}")
        print()

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

def hypothesize_algorithm(model):
    """Infer what algorithm the network discovered."""
    
    print("\n" + "=" * 70)
    print("PHASE 4: Algorithm hypothesis - What did it discover?")
    print("=" * 70)
    
    rank_weights = model.W_rank.data
    
    diag_dominance = rank_weights.diagonal().abs().mean()
    off_diag_mask = ~torch.eye(rank_weights.shape[0], dtype=bool)
    off_diag_avg = rank_weights.abs()[off_diag_mask].mean()
    
    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 * 1.5:
        print("\n  🏆 DISCOVERED: POSITION-WEIGHTED SORT")
        print("     Network learned that specific positions should rank higher.")
        print("     Strategy: Treat index positions as priority levels.")
    else:
        print("\n  🏆 DISCOVERED: VALUE-BASED RANKING")
        print("     Network learned actual value comparisons.")
        print("     Strategy: Score elements by their magnitude, reorder by score.")
    
    print(f"\n  Learned temperature: {model.temp.item():.4f}")
    if model.temp.item() > 0.5:
        print("     → High temperature: fuzzy/soft ranking")
    else:
        print("     → Low temperature: sharp/rigid ranking")

# ============================================================
# 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)")
    
    # 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: Maps input values to ranking scores            │
    │          (Encodes: "which values should come first?")   │
    │                                                         │
    │  W_comp:  Transforms input for comparison               │
    │          (Encodes: "how to evaluate each element")      │
    │                                                         │
    │  W_order: Reconstructs sorted output from attention     │
    │          (Encodes: "how to blend elements by rank")     │
    └─────────────────────────────────────────────────────────┘
    
    The optimizer wrote the solution into the matrices through
    500 gradient updates over truth table pairs (X_train, y_train).
    """)