import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import math

# ------------------------------------------------------------
# FluxTensor implementation (based on "The Algebra of Flux")
# ------------------------------------------------------------
class FluxTensor:
    """
    A tensor that carries value (v), entropy (s) and flux (t).
    Implements flux addition, multiplication, collapse and evolution.
    """
    def __init__(self, v, s=None, t=None):
        self.v = v                          # stationary component
        if s is None:
            s = torch.zeros_like(v)
        if t is None:
            t = torch.zeros_like(v)
        self.s = s                          # entropy (uncertainty)
        self.t = t                          # flux (velocity)

    def __add__(self, other):
        """Flux addition ⊕ : values add, entropies root-sum-square, fluxes add."""
        if isinstance(other, FluxTensor):
            new_v = self.v + other.v
            new_s = torch.sqrt(self.s**2 + other.s**2 + 1e-8)
            new_t = self.t + other.t
            return FluxTensor(new_v, new_s, new_t)
        else:
            # scalar/tensor addition: treat as zero entropy and zero flux
            other_v = torch.as_tensor(other, device=self.v.device)
            return FluxTensor(self.v + other_v, self.s.clone(), self.t.clone())

    def __mul__(self, other):
        """Flux multiplication ⊗ : product rule for v, error propagation for s."""
        if isinstance(other, FluxTensor):
            new_v = self.v * other.v
            # sqrt( (v1 σ2)^2 + (v2 σ1)^2 )
            new_s = torch.sqrt((self.v * other.s)**2 + (other.v * self.s)**2 + 1e-8)
            new_t = self.v * other.t + other.v * self.t
            return FluxTensor(new_v, new_s, new_t)
        else:
            # multiply by scalar
            scalar = torch.as_tensor(other, device=self.v.device)
            new_v = self.v * scalar
            new_s = self.s * torch.abs(scalar)
            new_t = self.t * scalar
            return FluxTensor(new_v, new_s, new_t)

    def __matmul__(self, other):
        """Matrix multiplication with flux propagation."""
        # For matrix multiplication: (v1 @ v2) and entropy propagation
        if isinstance(other, FluxTensor):
            # Value: standard matmul
            new_v = self.v @ other.v
            # Entropy: propagate through linear transformation
            # σ_out^2 = (|self.v| @ σ_other^2) + (σ_self^2 @ |other.v|^2)  (approximate)
            s_self_sq = self.s**2
            s_other_sq = other.s**2
            # Compute contributions: self.v @ (other.s^2) and (self.s^2) @ (other.v^2)
            term1 = self.v.abs() @ s_other_sq
            term2 = s_self_sq @ (other.v.abs()**2)
            new_s = torch.sqrt(term1 + term2 + 1e-8)
            # Flux: derivative of matrix product
            new_t = self.t @ other.v + self.v @ other.t
            return FluxTensor(new_v, new_s, new_t)
        else:
            # matmul with plain tensor: treat as zero entropy/flux
            other_t = torch.as_tensor(other, device=self.v.device)
            new_v = self.v @ other_t
            new_s = self.s @ (other_t.abs())
            new_t = self.t @ other_t
            return FluxTensor(new_v, new_s, new_t)

    def collapse(self, work):
        """
        Collapse operator 𝒞 : reduce entropy by paying work.
        work can be a scalar or tensor broadcastable to s.
        """
        new_s = torch.clamp(self.s - work, min=0.0)
        return FluxTensor(self.v, new_s, self.t)

    def evolve(self, dt=1.0, entropy_decay=0.0):
        """
        Time evolution (Euler step) for the ODE: dv/dt = t, ds/dt = -decay*s
        """
        new_v = self.v + self.t * dt
        new_s = self.s * (1.0 - entropy_decay * dt)
        new_s = torch.clamp(new_s, min=0.0)
        # flux remains same or could be updated by external dynamics
        return FluxTensor(new_v, new_s, self.t)

    def to(self, device):
        return FluxTensor(self.v.to(device), self.s.to(device), self.t.to(device))

    def detach(self):
        return FluxTensor(self.v.detach(), self.s.detach(), self.t.detach())

    def clone(self):
        return FluxTensor(self.v.clone(), self.s.clone(), self.t.clone())

    def total_entropy(self):
        return self.s.sum().item()


# ------------------------------------------------------------
# Flux Linear Layer: y = W @ x + b, all as FluxTensors
# ------------------------------------------------------------
class FluxLinear(nn.Module):
    def __init__(self, in_features, out_features, init_entropy=0.1):
        super().__init__()
        # Initialize weights as FluxTensor with small entropy and zero flux
        v = torch.empty(out_features, in_features)
        nn.init.kaiming_uniform_(v, a=math.sqrt(5))
        s = torch.full_like(v, init_entropy)
        t = torch.zeros_like(v)
        self.weight = FluxTensor(v, s, t)

        # Bias
        v_bias = torch.zeros(out_features)
        s_bias = torch.full_like(v_bias, init_entropy)
        t_bias = torch.zeros_like(v_bias)
        self.bias = FluxTensor(v_bias, s_bias, t_bias)

    def forward(self, x):
        # x must be FluxTensor
        out = self.weight @ x
        out = out + self.bias
        return out


# ------------------------------------------------------------
# Iterative Flux MLP: y = mlp(x, y)
# ------------------------------------------------------------
class IterativeFluxMLP(nn.Module):
    """
    Refinement network: takes input x and previous state y, outputs new y.
    After several refinement steps, the final y is mapped to logits.
    """
    def __init__(self, input_dim=784, hidden_dim=128, state_dim=64, num_classes=10,
                 num_refinements=5, work_per_step=0.05, entropy_decay=0.02):
        super().__init__()
        self.state_dim = state_dim
        self.num_refinements = num_refinements
        self.work_per_step = work_per_step
        self.entropy_decay = entropy_decay

        # Encoder: x -> embedding
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, state_dim)
        )

        # Refinement MLP: takes [x_embed, y] -> new y
        self.refiner = nn.Sequential(
            nn.Linear(state_dim + state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, state_dim)
        )

        # Classifier: final state -> logits
        self.classifier = nn.Linear(state_dim, num_classes)

        # We will convert plain tensors to FluxTensors before refinement
        # For simplicity, we keep a fixed initial entropy for y0
        self.init_entropy = 0.5

    def forward(self, x):
        # x is a plain tensor [batch, 784]
        batch = x.shape[0]

        # Encode input to embedding (plain tensor)
        x_emb = self.encoder(x)          # [batch, state_dim]

        # Initial state y0 as FluxTensor with zero value but high entropy
        y = FluxTensor(
            v=torch.zeros(batch, self.state_dim, device=x.device),
            s=torch.full((batch, self.state_dim), self.init_entropy, device=x.device),
            t=torch.zeros(batch, self.state_dim, device=x.device)
        )

        # Iterative refinement: y_{t+1} = refiner( concat(x_emb, y_t) )
        for step in range(self.num_refinements):
            # Concatenate x_emb (as FluxTensor with zero entropy) and current y
            x_flux = FluxTensor(x_emb, torch.zeros_like(x_emb), torch.zeros_like(x_emb))
            # Concat along feature dimension
            cat_v = torch.cat([x_flux.v, y.v], dim=-1)
            cat_s = torch.cat([x_flux.s, y.s], dim=-1)
            cat_t = torch.cat([x_flux.t, y.t], dim=-1)
            cat_flux = FluxTensor(cat_v, cat_s, cat_t)

            # Apply refiner (which uses linear layers with flux arithmetic)
            # For simplicity we simulate a two-layer MLP with FluxLinear
            # Because our refiner is a plain nn.Sequential, we must convert it
            # to use FluxLinear layers. We'll re-implement inside the loop using
            # custom flux operations, or we can define a flux version of the refiner.
            # For clarity, we define a small flux MLP inline.
            h1 = self._flux_linear(cat_flux, self.refiner[0].weight, self.refiner[0].bias)
            h1 = self._flux_relu(h1)
            h2 = self._flux_linear(h1, self.refiner[2].weight, self.refiner[2].bias)
            # h2 is the new y (FluxTensor)
            y = h2

            # Apply collapse: pay work to reduce entropy
            y = y.collapse(work=self.work_per_step)

            # Evolve (optional): let entropy decay naturally
            y = y.evolve(dt=1.0, entropy_decay=self.entropy_decay)

        # Final classification: take only the value component of y
        logits = self.classifier(y.v)
        return logits, y

    def _flux_linear(self, x, weight, bias):
        """Manual flux linear transformation using FluxTensor arithmetic."""
        # weight: [out, in], bias: [out]
        # Standard linear: output = x @ W.T + b
        out_v = x.v @ weight.T
        out_s = x.s @ weight.T.abs()
        out_t = x.t @ weight.T
        out = FluxTensor(out_v, out_s, out_t) + bias
        return out

    def _flux_relu(self, x):
        """ReLU on value component; entropy and flux pass through unchanged."""
        v_new = F.relu(x.v)
        return FluxTensor(v_new, x.s, x.t)


# ------------------------------------------------------------
# Training and testing
# ------------------------------------------------------------
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    total_loss = 0
    correct = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        data = data.view(data.size(0), -1)   # flatten
        optimizer.zero_grad()
        logits, _ = model(data)
        loss = F.cross_entropy(logits, target)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        pred = logits.argmax(dim=1)
        correct += pred.eq(target).sum().item()

        if batch_idx % 100 == 0:
            print(f'Train Epoch {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
                  f'({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')
    avg_loss = total_loss / len(train_loader)
    acc = 100. * correct / len(train_loader.dataset)
    print(f'====> Epoch {epoch} Average loss: {avg_loss:.4f}, Accuracy: {acc:.2f}%')
    return avg_loss, acc

def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            data = data.view(data.size(0), -1)
            logits, _ = model(data)
            test_loss += F.cross_entropy(logits, target, reduction='sum').item()
            pred = logits.argmax(dim=1)
            correct += pred.eq(target).sum().item()
    test_loss /= len(test_loader.dataset)
    acc = 100. * correct / len(test_loader.dataset)
    print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {acc:.2f}%\n')
    return test_loss, acc

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # Data loading
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    train_dataset = datasets.MNIST('../data', train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST('../data', train=False, transform=transform)
    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

    # Model
    model = IterativeFluxMLP(
        input_dim=784,
        hidden_dim=128,
        state_dim=64,
        num_classes=10,
        num_refinements=32,
        work_per_step=0.05,
        entropy_decay=0.02
    ).to(device)

    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # Training loop
    for epoch in range(42):
        train_loss, train_acc = train(model, device, train_loader, optimizer, epoch)
        test_loss, test_acc = test(model, device, test_loader)

    print("Training finished.")

if __name__ == "__main__":
    main()