import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import numpy as np
from collections import deque

# -------------------------------
# MLP Model for CIFAR-10
# -------------------------------
class MLP(nn.Module):
    def __init__(self, input_size=3072, hidden_sizes=[512, 256, 128], num_classes=10, dropout=0.2):
        super().__init__()
        layers = []
        prev = input_size
        for h in hidden_sizes:
            layers.append(nn.Linear(prev, h))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(dropout))
            prev = h
        layers.append(nn.Linear(prev, num_classes))
        self.net = nn.Sequential(*layers)
        
        # For storing activations and weights per layer (used in loss)
        self.activations = []
        self.weights = []
        self.biases = []
        self._register_hooks()
        
    def _register_hooks(self):
        def hook_fn(module, input, output):
            if isinstance(module, nn.Linear):
                self.activations.append(output.detach())
                self.weights.append(module.weight)
                self.biases.append(module.bias)
        for m in self.net:
            if isinstance(m, nn.Linear):
                m.register_forward_hook(hook_fn)
                
    def forward(self, x):
        self.activations.clear()
        self.weights.clear()
        self.biases.clear()
        x = x.view(x.size(0), -1)
        return self.net(x)

# -------------------------------
# CCT-32 Loss Implementation
# -------------------------------
class CCT32Loss(nn.Module):
    def __init__(self, model, lambda_standard=1.0, device='cuda'):
        super().__init__()
        self.model = model
        self.lambda_standard = lambda_standard
        self.device = device
        
        # History buffers for terms that need previous states
        self.prev_weights = deque(maxlen=2)   # for Q08, Q10
        self.prev_hidden = deque(maxlen=5)    # for Q09, Q10
        self.prev_loss = None                 # for Q09
        
        # Learnable layer weights (simplified: fixed scalar per term)
        # In practice these could be learnable, but we fix for stability
        self.alpha = torch.ones(8, device=device)   # Layer 1
        self.beta  = torch.ones(8, device=device)   # Layer 2
        self.gamma = torch.ones(8, device=device)   # Layer 3
        self.delta = torch.ones(8, device=device)   # Layer 4
        
        # Additional hyperparameters
        self.chaos_threshold = 1.0
        self.bias_bound = 1.0
        self.phase_bound = 10.0
        self.efficiency_threshold = 0.5
        self.curvature_threshold = 5.0
        self.compression_target = 0.5
        
    def forward(self, outputs, targets, step=None):
        """
        outputs: logits from model (batch, num_classes)
        targets: true labels (batch,)
        step: optional training step index
        """
        batch_size = outputs.size(0)
        probs = F.softmax(outputs, dim=1)
        log_probs = F.log_softmax(outputs, dim=1)
        
        # Standard cross-entropy
        ce_loss = F.cross_entropy(outputs, targets)
        
        # Get model internals
        activations = self.model.activations      # list of tensors [batch, dim]
        weights = self.model.weights              # list of weight matrices
        biases = self.model.biases                # list of bias vectors
        
        # Default zero for missing terms
        L_S = torch.tensor(0.0, device=self.device)
        L_P = torch.tensor(0.0, device=self.device)
        L_Q = torch.tensor(0.0, device=self.device)
        L_T = torch.tensor(0.0, device=self.device)
        
        # -------------------- Layer 1: Stationary Entropy --------------------
        # Q01: law consistency (hidden state consistent with stationary laws)
        #   subtract when consistent (low norm deviation)
        if len(activations) > 0:
            h_norm = torch.norm(activations[-1], dim=1).mean()
            law_dev = torch.abs(h_norm - 1.0)  # assume unit norm stationary law
            L_S = L_S - self.alpha[0] * law_dev   # subtract
        
        # Q02: gradient conservation violation (penalize large gradient change)
        #   additive when violated: compute gradient of loss w.r.t. weights
        if step is not None and step > 0 and self.prev_loss is not None:
            grad_norm = 0.0
            for p in self.model.parameters():
                if p.grad is not None:
                    grad_norm += p.grad.norm().item()
            grad_change = abs(grad_norm - self.prev_loss)  # approximate
            L_S = L_S + self.alpha[1] * grad_change
        
        # Q03: weight symmetry violation (subtract when symmetric)
        #   approximate: for each layer, penalty for non-symmetry (not typical for MLP)
        sym_loss = 0.0
        for w in weights:
            if w.size(0) == w.size(1):  # only square matrices
                sym_loss += torch.norm(w - w.T) / w.numel()
        L_S = L_S - self.alpha[2] * sym_loss   # subtract when symmetric
        
        # Q04: phase space bound violation (penalize activations out of bounds)
        #   additive
        bound_penalty = 0.0
        for a in activations:
            bound_penalty += torch.mean(F.relu(torch.abs(a) - self.phase_bound))
        L_S = L_S + self.alpha[3] * bound_penalty
        
        # Q05: bias boundedness violation (additive)
        bias_penalty = 0.0
        for b in biases:
            bias_penalty += torch.mean(F.relu(torch.abs(b) - self.bias_bound))
        L_S = L_S + self.alpha[4] * bias_penalty
        
        # Q06: attractor manifold distance (subtract when on manifold)
        #   attractor = one-hot encoded targets
        target_onehot = F.one_hot(targets, num_classes=outputs.size(1)).float()
        dist_to_attractor = torch.norm(probs - target_onehot, dim=1).mean()
        L_S = L_S - self.alpha[5] * dist_to_attractor
        
        # Q07: invariant preservation (subtract when invariant holds)
        #   use sum of activations as invariant
        if len(activations) > 0:
            inv = torch.mean(activations[-1], dim=1).sum()
            if hasattr(self, 'prev_inv'):
                inv_change = torch.abs(inv - self.prev_inv)
                L_S = L_S - self.alpha[6] * inv_change
            self.prev_inv = inv.detach()
        
        # Q08: fixed point stability (subtract when stable: small weight change)
        if len(self.prev_weights) > 0:
            w_prev = self.prev_weights[-1]
            w_change = 0.0
            for w_curr, w_prev_layer in zip(weights, w_prev):
                w_change += torch.norm(w_curr - w_prev_layer)
            L_S = L_S - self.alpha[7] * w_change
        # Store current weights for next step
        self.prev_weights.append([w.clone().detach() for w in weights])
        
        # -------------------- Layer 2: Probability Entropy --------------------
        # Q09: prediction uncertainty growth (subtract when narrowing)
        entropy = -torch.sum(probs * log_probs, dim=1).mean()
        if len(self.prev_hidden) > 0:
            entropy_change = entropy - self.prev_entropy
            L_P = L_P - self.beta[0] * torch.abs(entropy_change)  # subtract if decreasing
        self.prev_entropy = entropy.detach()
        
        # Q10: limit cycle detection (subtract when periodic)
        if len(self.prev_hidden) >= 2:
            h_curr = activations[-1] if activations else outputs
            h_prev = self.prev_hidden[-1]
            periodicity = torch.norm(h_curr - h_prev, dim=1).mean()
            L_P = L_P - self.beta[1] * periodicity   # subtract if low (periodic)
        self.prev_hidden.append(activations[-1].detach() if activations else outputs.detach())
        
        # Q11: output distribution entropy (subtract when low)
        L_P = L_P - self.beta[2] * entropy   # subtract entropy (encourage peaked)
        
        # Q12: chaos threshold (additive if Lyapunov approx > threshold)
        #   use gradient norm of output w.r.t. input as proxy
        grad_out = torch.autograd.grad(outputs.sum(), outputs, create_graph=True)[0]
        lyap = torch.norm(grad_out, p='fro')
        chaos_penalty = F.relu(lyap - self.chaos_threshold)
        L_P = L_P + self.beta[3] * chaos_penalty
        
        # Q13: ODE trajectory deviation (additive)
        #   assume simple ODE: dh/dt = -h, compare consecutive hidden states
        if len(self.prev_hidden) >= 2:
            h_curr = activations[-1] if activations else outputs
            h_prev = self.prev_hidden[-2]
            ode_pred = h_prev - h_prev  # trivial: zero change? Better: dt=1, dh = -h_prev
            ode_dev = torch.norm(h_curr - (h_prev - h_prev), dim=1).mean()
            L_P = L_P + self.beta[4] * ode_dev
        
        # Q14: variance explosion (additive)
        var_out = torch.var(outputs, dim=0).mean()
        var_penalty = F.relu(var_out - 10.0)   # threshold 10
        L_P = L_P + self.beta[5] * var_penalty
        
        # Q15: KL divergence from target (additive)
        kl_div = F.kl_div(log_probs, target_onehot, reduction='batchmean')
        L_P = L_P + self.beta[6] * kl_div
        
        # Q16: conditional entropy minimization (subtract when minimized)
        #   H(Y|X) = -E[log p(y|x)]
        cond_entropy = -torch.mean(torch.gather(log_probs, 1, targets.unsqueeze(1)))
        L_P = L_P - self.beta[7] * cond_entropy
        
        # -------------------- Layer 3: Question Collapse Entropy --------------------
        # Q17: mutual information gain (subtract when high)
        H_Y = entropy
        H_YgX = cond_entropy
        mi = H_Y - H_YgX
        L_Q = L_Q - self.gamma[0] * mi
        
        # Q18: collapse potential of neuron (subtract when positive)
        #   Δ_i = H(T) - H(T|h_i) approx using last hidden layer
        H_T = -torch.mean(torch.sum(target_onehot * torch.log(target_onehot+1e-8), dim=1))
        if len(activations) > 0:
            # compute H(T | h_last) by binning activations (simplified)
            h_last = activations[-1]
            with torch.no_grad():
                _, idx = torch.sort(h_last, dim=0)
                # crude conditional entropy: treat top/bottom halves as conditions
                mid = batch_size // 2
                cond_ent_h = 0.0
                for split in [idx[:mid], idx[mid:]]:
                    sub_targets = targets[split]
                    p = torch.bincount(sub_targets, minlength=num_classes).float() / len(sub_targets)
                    p = p + 1e-8
                    cond_ent_h -= torch.sum(p * torch.log(p))
            collapse_pot = H_T - cond_ent_h
            L_Q = L_Q - self.gamma[1] * F.relu(collapse_pot)
        
        # Q19: information efficiency (additive if below threshold)
        #   efficiency = Δ_i / (parameter count per neuron)
        if len(activations) > 0:
            param_per_neuron = weights[-1].numel() / weights[-1].size(0) if weights else 1.0
            eff = collapse_pot / param_per_neuron
            L_Q = L_Q + self.gamma[2] * F.relu(self.efficiency_threshold - eff)
        
        # Q20: hypothesis space pruning (subtract when large pruning)
        #   use margin as proxy: larger margin -> more pruning
        probs_sorted, _ = torch.sort(probs, dim=1, descending=True)
        margin = probs_sorted[:, 0] - probs_sorted[:, 1]
        L_Q = L_Q - self.gamma[3] * margin.mean()
        
        # Q21: path uniqueness (subtract when low entropy over paths)
        #   treat last layer logits as "paths"
        path_entropy = -torch.sum(probs * log_probs, dim=1).mean()
        L_Q = L_Q - self.gamma[4] * path_entropy
        
        # Q22: conditional question dependency (tracking only, not used in loss)
        #   could regularize, but we omit to keep simple
        
        # Q23: insufficient work budget (additive if remaining entropy > target)
        target_entropy = 0.5   # threshold
        remaining = H_YgX - target_entropy
        L_Q = L_Q + self.gamma[5] * F.relu(remaining)
        
        # Q24: net information gain (subtract when positive)
        net_gain = collapse_pot - param_per_neuron if 'collapse_pot' in locals() else 0.0
        L_Q = L_Q - self.gamma[6] * F.relu(net_gain)
        
        # -------------------- Layer 4: Taylor-Token Expansion --------------------
        # Q25: residual at order n (additive)
        #   treat each layer as an expansion order, residual = cross-entropy after that layer
        residual = ce_loss   # final residual
        L_T = L_T + self.delta[0] * residual
        
        # Q26: manifold curvature (additive if high)
        #   approx by hessian trace of output w.r.t. input
        if activations:
            hessian = torch.autograd.functional.hessian(lambda x: outputs.sum(), 
                                                        activations[-1].detach().requires_grad_(True))
            curvature = torch.norm(hessian, p='fro')
            L_T = L_T + self.delta[1] * F.relu(curvature - self.curvature_threshold)
        
        # Q27: Taylor order comparison (additive if next order more informative)
        #   compare loss reduction between layers
        if len(weights) >= 2:
            # using intermediate activations as pseudo-orders
            mid_out = self.model.net[:-2](outputs)  # not exact, but illustrative
            mid_ce = F.cross_entropy(mid_out, targets) if mid_out.size(-1)==num_classes else ce_loss
            improvement = mid_ce - ce_loss
            L_T = L_T + self.delta[2] * F.relu(improvement)
        
        # Q28: token distribution diffuseness (subtract when concentrated)
        token_entropy = -torch.sum(probs * log_probs, dim=1).mean()
        L_T = L_T - self.delta[3] * token_entropy
        
        # Q29: Taylor series convergence (additive until converged)
        #   difference between cumulative sum of token contributions and H(T)
        cum_contrib = mi   # mutual info as cum contribution
        conv_gap = torch.abs(cum_contrib - H_T)
        L_T = L_T + self.delta[4] * conv_gap
        
        # Q30: semantic compression ratio (additive if undercompressed)
        input_bits = 3072 * 8   # each pixel 8 bits
        comp_bits = sum(p.numel() for p in self.model.parameters()) * 32   # 32-bit floats
        ratio = comp_bits / input_bits
        L_T = L_T + self.delta[5] * F.relu(self.compression_target - ratio)
        
        # Q31: semantic convergence to target (subtract when converged)
        #   difference between expansion sum and H(T)
        L_T = L_T - self.delta[6] * conv_gap
        
        # Q32: resolution threshold matching (subtract when matched)
        #   use number of hidden layers as resolution
        used_res = len(weights)
        optimal_res = 3   # assume 3 hidden layers optimal
        res_mismatch = abs(used_res - optimal_res)
        L_T = L_T - self.delta[7] * res_mismatch
        
        # Combine all layers with standard loss
        total = self.lambda_standard * ce_loss + L_S + L_P + L_Q + L_T
        # Store previous loss for Q02
        self.prev_loss = total.detach().item()
        
        return total

# -------------------------------
# Training and Evaluation
# -------------------------------
def train_model(model, device, trainloader, optimizer, criterion, epochs=20):
    model.train()
    for epoch in range(epochs):
        running_loss = 0.0
        correct = 0
        total = 0
        for i, (inputs, targets) in enumerate(trainloader):
            inputs, targets = inputs.to(device), targets.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets, step=epoch*len(trainloader)+i)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
            
            if i % 100 == 99:
                print(f'Epoch {epoch+1}, Batch {i+1}: Loss {running_loss/100:.4f}, Acc {100.*correct/total:.2f}%')
                running_loss = 0.0
        print(f'Epoch {epoch+1} finished. Accuracy: {100.*correct/total:.2f}%')
        
def test_model(model, device, testloader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, targets in testloader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
    print(f'Test Accuracy: {100.*correct/total:.2f}%')

# -------------------------------
# Main
# -------------------------------
def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f'Using device: {device}')
    
    # Prepare CIFAR-10
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
    ])
    trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True, num_workers=2)
    testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
    testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)
    
    model = MLP(input_size=3072, hidden_sizes=[512, 256, 128], num_classes=10).to(device)
    criterion = CCT32Loss(model, lambda_standard=1.0, device=device)
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    print("Starting training with CCT-32 loss...")
    train_model(model, device, trainloader, optimizer, criterion, epochs=20)
    test_model(model, device, testloader)

if __name__ == '__main__':
    main()