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
from tqdm import tqdm
import numpy as np

# ----------------------------------------------------------------------
# Custom layer: Operator Series Linear
# Implements: y = combine_{k=1..K} ( f_k(x) )   where combine is one of the
# non-linear operators (geometric mean, multiplication, softmax, phase-locked).
# Each f_k(x) = sum_i W_ki * x_i with a sinusoidal basis expansion.
# ----------------------------------------------------------------------
class OperatorSeriesLinear(nn.Module):
    def __init__(self, in_features, out_features, K=4, operator='geometric_mean'):
        """
        Args:
            in_features: input dimension
            out_features: output dimension
            K: number of series terms
            operator: 'geometric_mean', 'multiplication', 'softmax', 'phase_locked'
        """
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.K = K
        self.operator = operator

        # Learnable basis functions: each term has its own frequency and phase
        # We use a sinusoidal expansion: f_k(x) = sum_i a_{k,i} * sin(omega_k * x_i + phi_k)
        # For simplicity, we learn a separate linear projection per term.
        self.W = nn.Parameter(torch.randn(K, out_features, in_features) * 0.1)
        self.bias = nn.Parameter(torch.zeros(K, out_features))

        # Additional learnable parameters for phase-locked operator
        if operator == 'phase_locked':
            self.phase_bias = nn.Parameter(torch.randn(K, out_features, in_features) * 0.1)

    def forward(self, x):
        # x shape: (batch, in_features)
        batch = x.shape[0]
        # Compute each term's linear response
        terms = []
        for k in range(self.K):
            # Linear part (could be replaced by any differentiable basis)
            term_k = F.linear(x, self.W[k], self.bias[k])  # (batch, out_features)
            # Apply non-linear activation to each term (optional)
            term_k = torch.sin(term_k)  # bounded non-linearity
            terms.append(term_k)

        # Stack terms: (K, batch, out_features)
        terms = torch.stack(terms, dim=0)

        # Apply the chosen operator to combine the K terms
        if self.operator == 'geometric_mean':
            # Stable geometric mean with proper gradient flow
            # Shift terms to be positive for log, then restore with bias
            terms_shifted = terms + 1.0  # shift to [0, 2] range
            # Use softplus to ensure positivity and smooth gradients
            terms_pos = F.softplus(terms_shifted)
            # Geometric mean: exp(mean(log(terms)))
            log_terms = torch.log(terms_pos + 1e-8)
            combined = torch.exp(torch.mean(log_terms, dim=0))
        elif self.operator == 'multiplication':
            combined = torch.prod(terms, dim=0)
        elif self.operator == 'softmax':
            # Weighted combination using softmax attention weights
            # Compute attention weights over K terms
            # Use mean of terms for stable softmax computation
            attention_weights = F.softmax(torch.mean(terms, dim=2, keepdim=True), dim=0)
            combined = torch.sum(attention_weights * terms, dim=0)
        elif self.operator == 'phase_locked':
            # phase-locked coupling: sum_{i<j} sin(theta_i - theta_j)
            # Here each term_k is treated as a phase angle
            # We compute pairwise sine differences across the K terms
            combined = torch.zeros(batch, self.out_features, device=x.device)
            for i in range(self.K):
                for j in range(i+1, self.K):
                    combined = combined + torch.sin(terms[i] - terms[j])
        else:
            raise ValueError(f"Unknown operator: {self.operator}")

        return combined

# ----------------------------------------------------------------------
# CNN feature extractor + Operator Series classifier
# ----------------------------------------------------------------------
class OperatorSeriesNet(nn.Module):
    def __init__(self, operator='geometric_mean', K=4):
        super().__init__()
        # Feature extractor (small CNN for CIFAR-10)
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.dropout = nn.Dropout(0.25)

        # After convs: 128 * 4 * 4 = 2048 features (CIFAR-10 32x32 -> 4x4 after 3 pools)
        self.fc_in = 128 * 4 * 4
        # Operator series layer as classifier
        self.op_series = OperatorSeriesLinear(self.fc_in, 10, K=K, operator=operator)

    def forward(self, x):
        # CNN front-end
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = self.dropout(x)
        x = x.view(x.size(0), -1)   # flatten
        # Operator series classification
        x = self.op_series(x)
        return x

# ----------------------------------------------------------------------
# Training and testing utilities
# ----------------------------------------------------------------------
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    train_loss = 0
    correct = 0
    for data, target in tqdm(train_loader, desc=f'Epoch {epoch}'):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
        pred = output.argmax(dim=1, keepdim=True)
        correct += pred.eq(target.view_as(pred)).sum().item()
    train_loss /= len(train_loader.dataset)
    accuracy = 100. * correct / len(train_loader.dataset)
    print(f'Train set: Average loss: {train_loss:.4f}, Accuracy: {correct}/{len(train_loader.dataset)} ({accuracy:.2f}%)')
    return accuracy

def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in tqdm(test_loader, desc='Testing'):
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.cross_entropy(output, target, reduction='sum').item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
    test_loss /= len(test_loader.dataset)
    accuracy = 100. * correct / len(test_loader.dataset)
    print(f'Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)')
    return accuracy

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

    # CIFAR-10 data
    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])
    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])

    trainset = torchvision.datasets.CIFAR10(root='../data', train=True, download=True, transform=transform_train)
    testset = torchvision.datasets.CIFAR10(root='../data', train=False, download=True, transform=transform_test)
    train_loader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
    test_loader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)

    # Choose an operator: 'geometric_mean', 'multiplication', 'softmax', 'phase_locked'
    operator = 'geometric_mean'   # try others!
    K = 4                         # series order
    model = OperatorSeriesNet(operator=operator, K=K).to(device)
    print(f"Model with operator={operator}, K={K} has {sum(p.numel() for p in model.parameters()):,} parameters")

    optimizer = optim.Adam(model.parameters(), lr=0.003, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)

    best_acc = 0
    for epoch in range(1, 21):
        train_acc = train(model, device, train_loader, optimizer, epoch)
        test_acc = test(model, device, test_loader)
        scheduler.step()
        if test_acc > best_acc:
            best_acc = test_acc
            torch.save(model.state_dict(), f'best_model_{operator}_K{K}.pth')
        print(f"Best test accuracy so far: {best_acc:.2f}%\n")

    print(f"Final best test accuracy for operator {operator}: {best_acc:.2f}%")

if __name__ == '__main__':
    main()