import torch
import torch.nn as nn
import lam

# Ternary path.
def l():
    print(lam.lamda)


def _mean_ternarize(W):
    with torch.no_grad():
        m = W.abs().mean()
        return torch.clamp(torch.round(W / m), min=-1, max=1)


def linear_Ternarize(W, num_steps=100):
    # Fall back to the mean-based threshold when MOMA is disabled.
    if not getattr(lam, 'use_moma', True):
        return _mean_ternarize(W)
    with torch.no_grad():
        N = W.numel()
        phi=lam.c_val
        if phi>0.3:
            num_steps=200
        # Mean absolute weight under the L1 definition.
        mu = W.abs().sum() / N
        # Standard deviation term used by the search objective.
        std = ((W - mu).pow(2).sum() / N).sqrt()

        # Search interval for the threshold.
        gamma_min = (1 - phi) * mu
        gamma_max = (1 + phi) * mu

        # Sweep candidate thresholds and keep the best result.
        min_loss = float('inf')
        best_W_tn = None

        for step in range(num_steps + 1):
            # Current threshold.
            r = gamma_min + (gamma_max - gamma_min) * step / num_steps

            # Negative and positive activation rates.
            p = (W < -r).float().sum() / N
            q = (W > r).float().sum() / N

            # Moments induced by the ternary distribution.
            g1_r = q - p
            g2_r = p + q - g1_r**2

            # Standard deviation implied by g2(r).
            std_r = g2_r.sqrt()

            # Match the ternary moments to the full-precision weights.
            mean_diff = (g1_r - mu).abs()
            std_diff = (std_r - std).abs()
            loss = mean_diff + lam.lamda * std_diff

            # Keep the best ternary projection found so far.
            if loss < min_loss:
                min_loss = loss
                best_W_tn = torch.where(W > r, 1,
                                        torch.where(W < -r, -1, 0))

        return best_W_tn

class TnLinear(nn.Linear):
    def __init__(self, *args, **kwargs):
        super(TnLinear, self).__init__(*args, **kwargs)

    def forward(self, x):
        w = self.weight
        w_tn = w + (linear_Ternarize(w.data) - w).detach()
        output = nn.functional.linear(x, w_tn, self.bias)
        return output

class MLP2(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(MLP2, self).__init__()
        self.fc1 = TnLinear(input_size, hidden_size)
        self.relu1 = nn.ReLU()
        self.fc2 = TnLinear(hidden_size, output_size)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu1(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# Full-precision path.
class MLP3(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(MLP3, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu1(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x


class LeNet5(nn.Module):
    def __init__(self, input_channels=1, num_classes=10):
        super(LeNet5, self).__init__()
        # First convolutional layer: 1 input channel, 6 output channels, 5x5 kernel
        self.conv1 = nn.Conv2d(input_channels, 6, kernel_size=5, stride=1, padding=2)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        # Second convolutional layer: 6 input channels, 16 output channels, 5x5 kernel
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # Fully connected layers
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 16 channels, 5x5 feature map
        self.relu3 = nn.ReLU()
        self.fc2 = nn.Linear(120, num_classes)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        # First convolutional layer with ReLU and 2x2 max pooling
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        # Second convolutional layer with ReLU and 2x2 max pooling
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        # Flatten the tensor
        x = x.view(-1, 16 * 5 * 5)
        # Fully connected layers with ReLU
        x = self.fc1(x)
        x = self.relu3(x)
        # Output layer with softmax activation
        x = self.fc2(x)
        x = self.softmax(x)
        return x




class LeNet5_tn(nn.Module):
    def __init__(self, input_channels=1, num_classes=10):
        super(LeNet5_tn, self).__init__()
        # First convolutional layer: 1 input channel, 6 output channels, 5x5 kernel
        self.conv1 = nn.Conv2d(input_channels, 6, kernel_size=5, stride=1, padding=2)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        # Second convolutional layer: 6 input channels, 16 output channels, 5x5 kernel
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # Fully connected layers
        self.fc1 = TnLinear(16 * 5 * 5, 120)  # 16 channels, 5x5 feature map
        self.relu3 = nn.ReLU()
        self.fc2 = TnLinear(120, num_classes)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        # First convolutional layer with ReLU and 2x2 max pooling
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        # Second convolutional layer with ReLU and 2x2 max pooling
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        # Flatten the tensor
        x = x.view(-1, 16 * 5 * 5)
        # Fully connected layers with ReLU
        x = self.fc1(x)
        x = self.relu3(x)
        # Output layer with softmax activation
        x = self.fc2(x)
        x = self.softmax(x)
        return x
