import numpy as np
from typing import List, Tuple, Callable
import torch
import torch.nn as nn
import torch.optim as optim

class IonicMLSimulator:
    """
    Simulates ionic bonding for ML problems.
    Each ion = (locus_embedding, reciprocal_embedding, charge, data)
    Bonding = weighted sum of reciprocals, then stabilize -> update parameters.
    """
    def __init__(self, d_model=128, learning_rate=0.01):
        self.d_model = d_model
        self.lr = learning_rate
        self.memory = []  # list of ions: each is dict with keys: locus, recip, charge, data
        
    def create_ion(self, data, locus_type='default', charge=1.0):
        """Create an ion from a data point (e.g., (x, y) pair)."""
        # Encode data into locus embedding (e.g., via a neural net)
        locus = self._encode_locus(data, locus_type)
        # Reciprocal = missing pattern: for ML, it's the negative gradient direction
        recip = self._compute_reciprocal(data, locus)
        return {'locus': locus, 'recip': recip, 'charge': charge, 'data': data}
    
    def _encode_locus(self, data, locus_type):
        """Map data to an embedding vector (simplified: random projection)."""
        # In real system, this would be a learned encoder
        return np.random.randn(self.d_model) * 0.1
    
    def _compute_reciprocal(self, data, locus):
        """Reciprocal = -gradient of loss w.r.t. parameters (for ML)."""
        # For simulation, we need to know current model parameters.
        # We'll store model outside and compute gradient.
        # Simplified: return random vector.
        return np.random.randn(self.d_model) * 0.1
    
    def bond(self, ions: List[dict], model_parameters: np.ndarray) -> np.ndarray:
        """
        Perform ionic bonding: combine all reciprocal vectors weighted by complementarity.
        Returns updated parameters (stable state).
        """
        # Compute pairwise complementarity scores (normalized so np.exp can't overflow)
        n = len(ions)
        scores = np.zeros((n, n))
        for i in range(n):
            for j in range(n):
                ri, rj = ions[i]['recip'], ions[j]['recip']
                denom = (np.linalg.norm(ri) * np.linalg.norm(rj)) + 1e-12
                scores[i, j] = np.dot(ri, rj) / denom  # cosine similarity, in [-1, 1]
        # Numerically stable softmax over complementarity
        scores = scores - scores.max(axis=1, keepdims=True)
        weights = np.exp(scores) / np.sum(np.exp(scores), axis=1, keepdims=True)
        # Per-ion weight = how complementary ion i is to the rest of the batch
        # (column mean; unlike row sums, this does NOT trivially equal 1)
        ion_weight = weights.mean(axis=0)
        # Bonded reciprocal = weighted average of reciprocals
        bonded_recip = np.zeros_like(model_parameters)
        for i, ion in enumerate(ions):
            bonded_recip += ion_weight[i] * ion['recip']
        # Stable state = gradient-descent step along the bonded reciprocal
        new_params = model_parameters - self.lr * bonded_recip
        return new_params

    def train_step(self, data_batch, model, loss_fn):
        """One ionic training step: create ions from batch, bond them to get parameter update."""
        # Create ions from each (x, y)
        ions = []
        for x, y in data_batch:
            # Compute gradient using current model
            model.zero_grad()
            pred = model(torch.tensor(x, dtype=torch.float32))
            loss = loss_fn(pred, torch.tensor(y))
            loss.backward()
            # Gradient w.r.t. parameters becomes the reciprocal (deficit)
            grad = np.concatenate([p.grad.numpy().flatten() for p in model.parameters()])
            # Locus: the input x encoded
            locus = self._encode_locus(x, 'input')
            # Store the raw gradient; bond() subtracts lr * bonded_recip,
            # so this yields params - lr * grad (gradient descent).
            ion = {'locus': locus, 'recip': grad, 'charge': 1.0, 'data': (x,y)}
            ions.append(ion)
        # Bond to get stable parameter update
        current_params = np.concatenate([p.data.numpy().flatten() for p in model.parameters()])
        new_params = self.bond(ions, current_params)
        # Update model parameters
        offset = 0
        for p in model.parameters():
            size = p.numel()
            p.data = torch.tensor(new_params[offset:offset+size], dtype=p.dtype).reshape(p.shape)
            offset += size
        return loss.item()