class gConcatenate(nn.Module):
    """ g(x, eps) = g([x, eps]), where eps ~ N(0, 1) or Unif(0, 1) """

    def __init__(self, model, noise_dim, noise_type = 'normal', scale = 1.0):
        super(gConcatenate, self).__init__()
        self.model = model
        self.noise_dim = noise_dim
        self.noise_f = torch.randn if noise_type == 'normal' else torch.rand
        self.scale = scale

    def forward(self, x):
        assert len(x.shape) >= 2,'x must have at least 2 dims, where batch is the first'
        eps = self.noise_f(x.shape[0], self.noise_dim).to(x.device)
        eps = self.scale * eps
        return self.model(torch.cat([x, eps], dim = 1))
    
# example
model = MLP(1, [64, 64], 1)
g = gConcatenate(model, noise_dim = 100, noise_type = 'uniform')

def engression_loss(y, preds, return_terms = False):
    """
    Args:
        y (torch.Tensor): True target values (batch_size, output_dim).
        preds (torch.Tensor): Predicted target with independently sampled noise (batch_size, m_samples, output_dim).
        return_terms (bool): If True, return the individual terms of the loss.
    """
    assert y.shape[0] == preds.shape[0] and y.shape[-1] == preds.shape[-1]
    b, d = y.shape
    b, m, d = preds.shape

    # Term 1: the absolute error between the predicted and true values
    term1 = torch.linalg.vector_norm(preds - y[:, None, :], ord = 2, dim = 2).mean()

    # Term 2: pairwise absolute differences between the predicted values
    term2 = torch.tensor(0.0, device = preds.device, dtype = preds.dtype)

    if m > 1:
        # cdist is convinient. The result shape before mean is (n, m, m).
        mean_pairwise_l2_dists = torch.cdist(preds, preds, p = 2).mean()
        term2 = - 0.5 * mean_pairwise_l2_dists * m / (m - 1)
        
    return (term1 + term2, term1, term2) if return_terms else term1 + term2

def train(g, dl, m, lr = 0.001, epochs = 100, verbose = True):

    optimizer = torch.optim.Adam(g.parameters(), lr = lr)

    losses = []
    for _ in range(epochs):
        for x, y in dl:

            g.zero_grad()

            # Generate m samples from the model
            # shape: (batch_size, m, output_dim)
            preds = torch.stack([g(x) for _ in range(m)], dim = 1)

            loss = engression_loss(y, preds)

            loss.backward()
            optimizer.step()

            losses.append(loss.item())
            if verbose: print(loss.item())
    
    return losses

g_stars = {
    'softplus': F.softplus,
    'square': lambda x: (torch.max(x, torch.tensor(0.0)) ** 2) / 2,
    'cubic': lambda x: (x ** 3) / 3,
    'log': lambda x: (x/3 + math.log(3) - 2/3)*(x <= 2) + (torch.log(1 + x*(x > 2)))*(x > 2) 
}

x_train_lims = {
    'softplus': (-2, 2),
    'square': (0, 2),
    'cubic': (-2, 2),
    'log': (0, 2)
}

x_test_lims = {
    'softplus': (-2, 8),
    'square': (0, 6),
    'cubic': (-2, 6),
    'log': (0, 10)
}

def train_data(n, input_dim, type = 'softplus', pre_additive_noise = True):

    g_star = g_stars[type]

    a, b = x_train_lims[type]
    X = torch.rand((n, input_dim)) * (b - a) + a # Unif(a, b)

    sd = 1 if type != 'cubic' else 1.1
    eta = torch.randn((n, input_dim)) * sd
    
    if pre_additive_noise:
        Y = g_star(X + eta)
    else:
        Y = g_star(X) + eta

    return X, Y

# example
X, Y = train_data(1000, 1, type = 'softplus', pre_additive_noise = True)

input_dim, hidden_dim, output_dim = 1, 128, 1
noise_dim, m_train, m_pred = 100, 2, 512

batch_size = 1024
epochs = 20

ds_size = lambda name: 50_000 if name in 'softplus/square' else 100_000 # 50k, 100k
depth = lambda name: 2 if name in 'softplus/square' else 3
hidden_dim = 100
lr = lambda name: 5e-2 #1e-3 #if name in 'softplus/square' else 1e-4

n_runs = 20
pre_additive_noise = True

saved_g_preds = {}
saved_l1_preds = {}

extra = '_post' if not pre_additive_noise else ''
# saved_g_preds = torch.load(f'logs/saved_g_preds{extra}.pt')
# saved_l1_preds = torch.load(f'logs/saved_l1_preds{extra}.pt')


f, axs = plt.subplots(2, 2, figsize = (7, 7))

for (name, g_star), ax in zip(g_stars.items(), axs.flatten()):

    t = torch.linspace(x_test_lims[name][0], x_test_lims[name][1], 25)[:, None]

    # Cache to fix plot, etc
    if name in saved_g_preds:
        g_preds = saved_g_preds[name]

    else:
        g_preds = []
        for seed in range(n_runs):

            set_seed(seed)
            # model = MLP(input_dim + noise_dim, [hidden_dim] * depth(name), output_dim)
            model = ResMLP(input_dim + noise_dim, hidden_dim, depth(name), output_dim)
            g = gConcatenate(model, noise_dim = noise_dim, noise_type = 'normal')

            X, Y = train_data(ds_size(name), input_dim, type = name, pre_additive_noise = pre_additive_noise)
            dl = torch.utils.data.DataLoader(list(zip(X, Y)), batch_size = batch_size if batch_size else len(Y), shuffle = True, drop_last = True)

            losses = train(g, dl, m_train, epochs = epochs, verbose = False, lr = lr(name))

            # Predict the median using g: sample m_pred per point and take the median
            with torch.no_grad():
                g.eval()
                g_pred = torch.quantile(
                    torch.stack([g(t) for _ in range(m_pred)], dim = 1), q = 0.5, dim = 1
                )
                g_preds.append(g_pred)
        
        g_preds = torch.stack(g_preds)
        saved_g_preds[name] = g_preds

    # L1 baseline
    if name in saved_l1_preds:
        l1_preds = saved_l1_preds[name]
    else:
        l1_preds = []
        for seed in range(n_runs):
            set_seed(seed)
            model = ResMLP(input_dim, hidden_dim, depth(name), output_dim)
            X, Y = train_data(ds_size(name), input_dim, type = name, pre_additive_noise = True)
            dl = torch.utils.data.DataLoader(list(zip(X, Y)), batch_size = batch_size if batch_size else len(Y), shuffle = True, drop_last = True)
            l1_losses = train_l1(model, dl, epochs = epochs, verbose = False, lr = lr(name))

            with torch.no_grad():
                model.eval()
                l1_preds.append(model(t))

        l1_preds = torch.stack(l1_preds)
        saved_l1_preds[name] = l1_preds


    X, Y = train_data(ds_size(name), 1, type = name, pre_additive_noise = True)
    ax.scatter(X[:5000], Y[:5000], color = 'gray',  alpha = 0.1, s = 1)
    ax.plot(t, g_star(t), label = f'True g', color = 'red')

    ax.plot(t, l1_preds.mean(dim = 0), color = 'tab:blue', label = 'L1')
    ax.fill_between(t.flatten(),
        l1_preds.quantile(0.10, dim = 0).flatten(),
        l1_preds.quantile(0.90, dim = 0).flatten(),
        alpha = 0.2, color = 'tab:blue'
    )

    ax.plot(t, g_preds.mean(dim = 0), color = 'tab:orange', label = 'engression')
    ax.fill_between(t.flatten(),
        g_preds.quantile(0.10, dim = 0).flatten(),
        g_preds.quantile(0.90, dim = 0).flatten(),
        alpha = 0.2, color = 'tab:orange'
    )
    ax.set_title(name)


axs[0, 0].legend()
axs[1, 1].set_ylim((-2, 5))
axs[1, 0].set_ylim((-10, 80))

f.tight_layout()
f.savefig(f'images/figure4{extra}.png', dpi = 300, bbox_inches = 'tight')


torch.save(saved_g_preds, f'logs/saved_g_preds{extra}.pt')
torch.save(saved_l1_preds, f'logs/saved_l1_preds{extra}.pt')

n = 5_000
X = torch.rand(n) * 2 - 1               # X ~ Unif(-1, 1)
Y = X + (torch.randn_like(X) * 1)       # Y|X ~ N(X, 1)
X, Y = X[:, None], Y[:, None]

t = torch.linspace(-1, 1, 25)[:, None]

# True Y|X quantiles if we assume Y|X ~ N(X, 1)
qs = [0.1, 0.5, 0.9]
true_quantiles = {q: t + norm.ppf(q) for q in qs}

# Estimated Y|X quantiles
# [m_train, quantile level] = list of [est quantile], one per batch
quantiles = {}

input_dim, hidden_dim, output_dim = 1, 64, 1
noise_dim = 50
mini_batch_size = 32
m_pred = 512

# noise_scale = 1

stats = defaultdict(list)

m_trains = [2, 4, 8, 16]
noise_scales = [0.1, 1, 2, 4]

for m_train, noise_scale in product(m_trains, noise_scales):

    set_seed(69)
    dl = torch.utils.data.DataLoader(list(zip(X, Y)), batch_size = mini_batch_size, shuffle = False)
    model = ResMLP(in_dim = input_dim + noise_dim, hidden_dim = hidden_dim, n_blocks = 2, out_dim = output_dim)
    g = gConcatenate(model, noise_dim = noise_dim, scale = noise_scale)

    optimizer = torch.optim.Adam(g.parameters(), lr = 1e-4)


    for _ in range(3):
        for x, y in dl:

            g.zero_grad()

            # Generate m samples from the model
            preds = torch.stack([g(x) for _ in range(m_train)], dim = 1)
            loss, loss_1, loss_2 = engression_loss(y, preds, return_terms = True)
            loss.backward()
            optimizer.step()

            with torch.no_grad():

                q_losses = []
                preds = torch.stack([g(t) for _ in range(m_pred)], dim = 1)
                for q in qs:
                    pred_quantile = torch.quantile(preds, q = q, dim = 1)
                    quantiles[m_train, noise_scale, q] = quantiles.get((m_train, noise_scale, q), []) + [pred_quantile]
                    
                    if q in true_quantiles:
                        stats[m_train, noise_scale, q].append(
                            (true_quantiles[q] - pred_quantile.mean(dim = 0)).flatten().abs().mean().item()
                        )

                stats[m_train, noise_scale, 'loss'].append(loss.item())
                stats[m_train, noise_scale, 'loss_1'].append(loss_1.item())
                stats[m_train, noise_scale, 'loss_2'].append(loss_2.item())

torch.save((quantiles, stats), 'logs/m_noise_scale.pt')

class Model(nn.Module):
    '''
    Used in MNIST experiments.
    A two-layer linear ReLU network with dropout and max norm regularization.
    '''
    def __init__(self, hidden_size, max_norm = 2.0, drop_rate = 0.5):
        super(Model, self).__init__()
        self.max_norm = max_norm
        self.layers = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28 * 28, hidden_size),
            nn.ReLU(),
            nn.Dropout(drop_rate),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(drop_rate),
            nn.Linear(hidden_size, 10)
        )

    def forward(self, x):
        # Clip the weights to the maximum allowed norm
        if self.max_norm is not None:
            with torch.no_grad():
                for layer in self.modules():
                    if isinstance(layer, nn.Linear):
                        norm = layer.weight.data.norm(2, dim=1, keepdim=True)
                        desired = torch.clamp(norm, max=self.max_norm)
                        layer.weight.data *= (desired / norm)
                    
        return self.layers(x)

# Regular cross-entropy loss
def hard_loss(outputs, labels, criterion, *args):
    return criterion(outputs, labels)

# Distillation loss
def soft_loss(outputs, labels, criterion, examples, big_model, T, a):
    with torch.no_grad():
        big_model.eval()
        soft_labels = F.softmax(big_model(examples) / T, dim = 1)

    return a * (T ** 2) * criterion(outputs / T, soft_labels) + (1 - a) * criterion(outputs, labels)

# Hyperparameters
num_epochs = 100
batch_size = 128
lr = 0.01
patience = 7
big_model_size = 1200
small_model_size = 800

loader = lambda ds, shuffle = False: DataLoader(ds, batch_size = batch_size, shuffle = shuffle)
val_loader   = loader(val_dataset)
test_loader  = loader(test_dataset)

# Train the big model
train_dataset.dataset.transform = aug_transform
train_loader = loader(train_dataset, shuffle = True)

big_model = Model(1200).to(device)
big_train_history = train(big_model, hard_loss, train_loader, val_loader, num_epochs, lr, patience)
test_loss, test_accuracy = evaluate_model(big_model, test_loader)

save_results(big_model, big_train_history, 'big_model')
test_loss, test_accuracy

def train_small_model(train_dataset, val_dataset, seed, loss, model_size = small_model_size):
    set_seed(seed)
    # no augmentation
    train_dataset.dataset.transform = reg_transform
    train_loader = loader(train_dataset, shuffle = True)
    val_loader = loader(val_dataset)

    # or regularization
    small_model = Model(model_size, max_norm = None, drop_rate = 0.0).to(device)
    small_train_history = train(small_model, loss, train_loader, val_loader, num_epochs, lr, patience)
    return small_model, small_train_history

# small_model, small_train_history = train_small_model(train_dataset, val_dataset, seed = 42, loss = hard_loss, T = 1.0)
# save_results(small_model, small_train_history, 'small_model')
# evaluate_model(small_model, test_loader)

temperature, a = 4.0, 0.5

distilled_model, distilled_train_history = train_small_model(
    train_dataset = train_dataset, val_dataset = val_dataset, seed = 42,
    loss = functools.partial(soft_loss, big_model = big_model, T = temperature, a = a)
)
save_results(distilled_model, distilled_train_history, 'distilled_model')
evaluate_model(distilled_model, test_loader)

# Remove all 3s from the dataset
train_dataset = datasets.MNIST(root = DATA_DIR, train = True,  download = True, transform = aug_transform)
train_dataset = torch.utils.data.Subset(train_dataset, np.where(train_dataset.targets != 3)[0])

# Split training data into train and validation sets
train_size = int(0.9 * len(train_dataset))
val_size = len(train_dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(
    train_dataset, [train_size, val_size]
)

# Train small without distillation
small_no_3, small_no_3_history = train_small_model(
    train_dataset = train_dataset, val_dataset = val_dataset,
    seed = 42, loss = hard_loss, model_size = 800
)

# Train small with distillation
distilled_no_3, distilled_no_3_history = train_small_model(
    train_dataset = train_dataset, val_dataset = val_dataset,
    seed = 42, loss = functools.partial(soft_loss, big_model = big_model, T = 4.0, a = 0.5),
    model_size = 800
)

class RandX(Dataset):
    """Injects example noise into dataset by replacing x% of inputs with random gaussian N(0, 1) noise"""
    def __init__(self, dataset, x = 1.0):
        self.dataset = dataset
        self.x = x

        self.modified = {}
        for idx, (img, _) in enumerate(self.dataset):
            if np.random.rand() <= x:
                self.modified[idx] = torch.randn_like(img)
        torch.save(self.modified, os.path.join(dataset.root, 'randX_modified'))

    def __len__(self): return len(self.dataset)

    def __getitem__(self, idx):
        X, y = self.dataset[idx]
        return self.modified.get(idx, X), y 

class RandY(Dataset):
    """Injects example noise into dataset by replacing y% of labels with random labels"""
    def __init__(self, dataset, y = 1.0):
        self.dataset = dataset
        self.y = y

        self.modified = {}
        for idx in range(len(self.dataset)):
            if np.random.rand() <= y:
                self.modified[idx] = np.random.randint(0, len(self.dataset.classes))
        torch.save(self.modified, os.path.join(dataset.root, 'randY_modified'))

    def __len__(self): return len(self.dataset)

    def __getitem__(self, idx):
        X, y = self.dataset[idx]
        return X, self.modified.get(idx, y)

def train(model, train, val, optimizer, criterion = nn.CrossEntropyLoss(), epochs = 100, batch_size = 256, save_path = 'models/tmp'):
    model.to(device)
    train = DataLoader(train, batch_size = batch_size, shuffle = True)
    val = DataLoader(val, batch_size = batch_size, shuffle = False)
    best_loss = np.inf
    for epoch in range(epochs):
        model.train()
        for images, labels in train:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            loss = criterion(model(images), labels)
            loss.backward()
            optimizer.step()
        val_loss, val_acc = eval_model(model, val)
        if val_loss < best_loss:
            best_loss = val_loss
            torch.save(model.state_dict(), save_path)
        print(f'Epoch {epoch + 1}/{epochs}, Loss: {val_loss:.4f}, Accuracy: {val_acc:.4f}')

def _init_weights(m):
    if 'Linear' in str(type(m)):
        nn.init.xavier_uniform_(m.weight)
        nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
    elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
        nn.init.constant_(m.weight, 1)
        nn.init.constant_(m.bias, 0)

def initialize_model(model, data_loader, in_device = False):
    with torch.no_grad():
        imgs, _ = next(iter(data_loader))
        if in_device:model.to(device); model(imgs.to(device))
        _ = model(imgs)
        model.apply(_init_weights)

class MLP(nn.Module):
    def __init__(self, n_classes = 10):
        super().__init__()
        self.model = nn.Sequential(
            nn.Flatten(),
            nn.LazyLinear(4096),
            nn.ReLU(),
            nn.LazyLinear(4096),
            nn.ReLU(),
            nn.LazyLinear(n_classes)
        )

    def forward(self, x):
        return self.model(x)

def add_missclassified(missclassified, model, test_set, batch_size = 256):
    model = model.to(device); model.eval()
    test = DataLoader(test_set, batch_size = batch_size, shuffle = False)
    i = 0
    with torch.no_grad():
        for images, labels in test:
            images, labels = images.to(device), labels.to(device)
            _, pred = torch.max(model(images), 1)
            missclassified[i:i + test.batch_size] += (pred != labels).float()
            i += test.batch_size

def gen_fig_1(epochs = 1, n_inits = 100):
    training_sets = [train_set, RandX(train_set, x = 1.0), RandY(train_set, y = 1.0)]
    for training_set in training_sets:
        missclassified = torch.zeros(len(test_set)).to(device)
        for _ in range(n_inits):
            m = MLP()
            initialize_model(m, DataLoader(train_set, batch_size = 256))
            train(m, training_set, test_set, optim.SGD(m.parameters(), lr = 0.01), epochs = epochs)
            add_missclassified(missclassified, m, test_set)
        missclassified /= n_inits
        torch.save(missclassified, f'logs/missclassified_epochs={epochs}_' + training_set.__class__.__name__)

gen_fig_1(epochs = 1, n_inits = 100)
gen_fig_1(epochs = 10, n_inits = 100)

f, (ax1, ax2) = plt.subplots(1, 2, figsize = (9, 4), sharey = True)

for ax, epochs in zip([ax1, ax2], [1, 10]):

    for fname in sorted([f for f in os.listdir('logs') if 'missclassified' in f and f'epochs={epochs}_' in f]):
        missclassified = torch.load(os.path.join('logs', fname))
        p = (1 - missclassified).sort().values.to('cpu').numpy()
        ax.plot(p, label = fname.split('_')[-1])
    
    # Plot binomially sampled points
    randX_mean = 1 - torch.load(os.path.join('logs', 'missclassified_RandX')).mean().item()
    bin_data = np.random.binomial(n = 100, p = randX_mean, size = 10000) / 100
    bin_data.sort()
    ax.plot(bin_data, label = 'Binomial_X')
    ax.set_title(f'Epoch = {epochs}')

ax1.legend()
f.supylabel('P(correct)')
f.supxlabel('Example(sorted by P(correct))')
f.tight_layout()

n = 10
f, ax = plt.subplots(2, n, figsize = (n - 2, 2))
for i, idx in enumerate(torch.sort(p).indices[:n]):
    img, label = test_set[idx]
    ax[0][i].imshow(img.permute(1, 2, 0).numpy())
    ax[0][i].axis('off')
    ax[0][i].set_title(test_set.classes[label].replace('mobile', ''))

for i, idx in enumerate(torch.sort(p).indices[-n:]):
    img, label = test_set[idx]
    ax[1][i].imshow(img.permute(1, 2, 0).numpy())
    ax[1][i].axis('off')
    ax[1][i].set_title(test_set.classes[label].replace('mobile', ''))

f.suptitle('Hardest (top) and easiest (bottom) examples')
f.tight_layout()

class ConvNet(nn.Module):
    def __init__(self, num_classes):
        super(ConvNet, self).__init__()

        self.model = nn.Sequential(
            nn.LazyConv2d(out_channels=200, kernel_size=5),  # LazyConv2d to infer input channels
            nn.BatchNorm2d(200),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=3),

            nn.LazyConv2d(out_channels=200, kernel_size=5),  # Another LazyConv2d
            nn.BatchNorm2d(200),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=3),

            nn.Flatten(),  # Flatten for the fully connected layer
            nn.LazyLinear(out_features=384),  # LazyLinear to infer input features
            nn.BatchNorm1d(384),
            nn.ReLU(),

            nn.LazyLinear(out_features=192),  # Another LazyLinear
            nn.BatchNorm1d(192),
            nn.ReLU(),

            nn.LazyLinear(out_features=num_classes)
        )

    def forward(self, x):
        return self.model(x)
    
def train_conv(model, train, val, optimizer, criterion = nn.CrossEntropyLoss(), epochs = 100, batch_size = 256):
    model.to(device)
    train = DataLoader(train, batch_size = batch_size, shuffle = True)
    val = DataLoader(val, batch_size = batch_size, shuffle = False)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size = 15, gamma = 0.5)
    for epoch in range(epochs):
        model.train()
        for images, labels in train:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            loss = criterion(model(images), labels)
            loss.backward()
            optimizer.step()
        scheduler.step()
        val_loss, val_acc = eval_model(model, val)
        print(f'Epoch {epoch + 1}/{epochs}, Loss: {val_loss:.4f}, Accuracy: {val_acc:.4f}')

cifar10_model = ConvNet(num_classes = n_classes)
randY_model = ConvNet(num_classes = n_classes)

for model, dataset in zip([cifar10_model, randY_model], [train_set, RandY(train_set)]):
    initialize_model(model, DataLoader(dataset, batch_size = 256))
    train_conv(model, dataset, test_set, optim.SGD(model.parameters(), lr = 0.01, momentum = 0.9), epochs = 10)
    torch.save(model.state_dict(), f'models/fig2_{model.__class__.__name__}')

f, (ax1, ax2) = plt.subplots(1, 2, figsize = (9, 4), sharey = True)

for name, model, ax in zip(['Cifar 10', 'RandY'], [cifar10_model, randY_model], [ax1, ax2]):

    model.eval()
    kernel_weights = model.model[0].weight[:7 * 12].detach().cpu().clone()
    kernel_weights = (kernel_weights - kernel_weights.min()) / (kernel_weights.max() - kernel_weights.min())
    filter_img = utils.make_grid(kernel_weights, nrow = 12, padding = 1)
    ax.imshow(filter_img.permute(1, 2, 0))
    ax.set_title(name)
    ax.axis('off')

def standard_normal(shape):
    r = torch.randn(shape)
    r = r.to(device)
    return r

def lass(model, x, alpha = 0.25 / 255, beta = 0.2 / 255, r = 0.3 / 255, eta = standard_normal, max_iter = 10):
    """
    Langevin Adversarial Sample Search (LASS).
    Finds a perturbation of x that changes the model's prediction.
    
        labels: Tensor of true labels corresponding to the input x.
        alpha: Step size for the gradient sign method.
        beta: Scaling factor for the noise.
        r: Clipping radius for adversarial perturbations.
        eta: Noise process.
    """
    # Orignal prediction
    with torch.no_grad():
        pred_on_x = model(x).argmax(dim=1)
    

    x_adv = x.clone().detach().requires_grad_(True)
    converged = False
    iter_count = 0

    while not converged and iter_count < max_iter:
        iter_count += 1

        # Forward pass to get model output
        x_adv.requires_grad_(True)
        output = model(x_adv)

        # Compute gradient of the output with respect to input
        loss = F.cross_entropy(output, pred_on_x)  # Use actual labels
        loss.backward()

        # Compute the perturbation
        gradient_sign = x_adv.grad.sign()
        delta = alpha * gradient_sign + beta * eta(x_adv.shape)

        with torch.no_grad():
            
            x_adv += delta

            # Apply the clipping to each dimension so that each pixel is in the range [x - r, x + r]
            x_adv = torch.clamp(x_adv, x - r, x + r)

            # Check if the adversarial example has changed the model's prediction
            new_output = model(x_adv)
            if not torch.equal(output.argmax(dim=1), new_output.argmax(dim=1)):
                converged = True
                x_hat = x_adv.clone().detach()

        # Zero the gradients for the next iteration
        model.zero_grad()
        if x_adv.grad is not None: x_adv.grad.zero_()

    return converged, x_hat if converged else None

def compute_csr(model, test_set, n_examples = None, shuffle = False, **lass_kwargs):
    if n_examples is None: n_examples = len(test_set)
    model = model.to(device)
    csr = 0
    for i, (images, labels) in enumerate(DataLoader(test_set, batch_size = 1, shuffle = shuffle)):
        if i == n_examples: break
        images, labels = images.to(device), labels.to(device)
        converged, _ = lass(model, images, **lass_kwargs)
        if converged: csr += 1
    return csr / n_examples

model = cifar10_model
model.eval(); model.to(device)
set_seed(5)

f, (ax1, ax2) = plt.subplots(1, 2, figsize = (9, 4), sharey = True)
for i, (x, y) in enumerate(DataLoader(test_set, batch_size = 1, shuffle = True)):
    x = x.to(device)
    y = y.to(device)
    with torch.no_grad():
        if model(x).argmax() != y: continue
    converged, adv = lass(model, x)
    if converged:
        ax1.imshow(x.squeeze().cpu().permute(1, 2, 0))
        ax2.imshow(adv.squeeze().cpu().permute(1, 2, 0))
        ax1.axis('off'); ax2.axis('off')
        with torch.no_grad():
            ax1.set_title(f'Original. Predicted: {test_set.classes[model(x).argmax().item()]}')
            ax2.set_title(f'Adversarial. Predicted: {test_set.classes[model(adv).argmax().item()]}')
        break

def train_fig9(model, train, val, optimizer, criterion = nn.CrossEntropyLoss(), epochs = 100, batch_size = 256):
    val_accs, csrs = [], []
    model.to(device)
    train = DataLoader(train, batch_size = batch_size, shuffle = True)
    val = DataLoader(val, batch_size = batch_size, shuffle = False)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size = 15, gamma = 0.5)
    for epoch in range(epochs):
        model.train()
        for images, labels in train:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            loss = criterion(model(images), labels)
            loss.backward()
            optimizer.step()
        scheduler.step()
        val_loss, val_acc = eval_model(model, val)
        val_accs.append(val_acc)
        print(f'Epoch {epoch + 1}/{epochs}, Loss: {val_loss:.4f}, Accuracy: {val_acc:.4f}')

        if epoch % 10 == 0:
            csr = compute_csr(model, val.dataset, n_examples = 500, r = 40 / 255)
            csrs.append(csr)
            print(f'CSR: {csr:.4f}')
    
    return val_accs, csrs

for dataset in [train_set, RandX(train_set, x = 1.0), RandY(train_set, y = 1.0)]:
    set_seed(seed  = 42)
    print(dataset.__class__.__name__)
    model = ConvNet(num_classes = n_classes)
    initialize_model(model, DataLoader(dataset, batch_size = 256))
    val_accs, csrs = train_fig9(model, dataset, test_set, optim.SGD(model.parameters(), lr = 0.01, momentum = 0.9), epochs = 141)
    torch.save({'val_accs': val_accs, 'csrs': csrs}, f'logs/fig9_r=45_{dataset.__class__.__name__}')

plt.plot(np.arange(0, 141, 10), torch.load('logs/fig9_CIFAR10')['csrs'], 'b', label = 'CIFAR10')
plt.plot(np.arange(0, 141, 10), [0] * 15, 'r--', label = 'RandX (?)')
plt.plot(np.arange(0, 141, 10), torch.load('logs/fig9_RandY')['csrs'], 'g', label = 'RandY')

plt.legend()
plt.xlabel('Epoch')
plt.ylabel('Critical Sample Ratio (CSR)')

# Note that we sample from a standard normal distribution
X = torch.randn((5, 2))
rand_planes = torch.randn((2, 2)) - 0.5

def lims(X, ax, eps = 0.1):
    m = X[:, ax].abs().max().item()
    m += eps
    return (-m, m)

plt.scatter(X[:, 0], X[:, 1])
plt.xlim(lims(X, 0)); plt.ylim(lims(X, 1))

# Axes
plt.axline((0,0), (0,1), c = '0', linewidth = 0.8)
plt.axline((0,0), (1,0), c = '0', linewidth = 0.8)

# Random planes
for v in rand_planes:
    x, y = v.tolist()
    plt.axline((0,0), (y, -x), c = 'g', linewidth = 0.8)

plt.scatter(*X[0].tolist())
plt.scatter(*X[2].tolist())
plt.xlim(lims(X, 0)); plt.ylim(lims(X, 1))
plt.axline((0,0), (0,1), c = '0', linewidth = 0.8)
plt.axline((0,0), (1,0), c = '0', linewidth = 0.8)

x, y = rand_planes[0].tolist()
plt.axline((0,0), (y, -x), c = 'g', linewidth = 0.8, label = 'hyperplane')
plt.quiver(*([0], [0]), x, y, color = '0.5', label = 'v')
plt.legend()

import torch
import torch.nn as nn
import torch.optim as optim

from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader, Dataset
from PIL import Image

import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import os, itertools, time

os.makedirs('logs', exist_ok = True)
os.makedirs('models', exist_ok = True)

seed = 42
np.random.seed(seed)
torch.manual_seed(seed)

device = torch.device(
    'cuda' if torch.cuda.is_available() else
    ('mps' if torch.backends.mps.is_available() else
    'cpu')
)

def eval_model(model, test, criterion = nn.CrossEntropyLoss()):
    model.eval()
    correct, loss = 0, 0.0
    with torch.no_grad():
        for images, labels in test:
            images, labels = images.to(device), labels.to(device)
            _, pred = torch.max(model(images), 1)
            correct += (pred == labels).float().sum().item()
            loss += criterion(model(images), labels).item()
    return loss / len(test.dataset), correct / len(test.dataset)

device

train_set = datasets.CIFAR10('./data', download = True, train = True, transform = transforms.ToTensor())
val_set = datasets.CIFAR10('./data', download = True, train = False, transform = transforms.ToTensor())

def channel_means_stds(dataset):
    imgs = torch.stack([img for img, _ in train_set])
    return imgs.mean(dim = [0, 2, 3]), imgs.std(dim = [0, 2, 3])

means, stds = channel_means_stds(train_set)
print(f'Training channel-wise\n\tmeans: {means}\n\tstds: {stds}')

means, stds = channel_means_stds(val_set)
print(f'Validation channel-wise\n\tmeans: {means}\n\tstds: {stds}')

train_transform = transforms.Compose([
    transforms.RandomCrop(32, padding = 4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(means, stds),
])

# We do not perform data augmentation on the validation set
val_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(means, stds),
])

train_set.transform = train_transform
val_set.transform = val_transform

train_loader = DataLoader(train_set, batch_size = 128, shuffle = True)
val_loader = DataLoader(val_set, batch_size = 128, shuffle = False)

def xavier_init(m):
    if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)

def train_epoch(model, train, optimizer, criterion):
    # Trains the model for one epoch
    model.train()
    train_loss, correct = 0.0, 0
    for images, labels in train:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        output = model(images)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
        _, pred = torch.max(output, 1)
        correct += (pred == labels).float().sum().item()
    return train_loss / len(train.dataset), correct / len(train.dataset)


def train(model, train, val, init_lr, plateau_patience = 20):

    optimizer = optim.SGD(model.parameters(), lr = init_lr, momentum = 0.9, weight_decay = 5e-4)
    scheduler  = optim.lr_scheduler.MultiStepLR(optimizer, milestones = [82, 123], gamma = 0.1)
    criterion = nn.CrossEntropyLoss()

    model.to(device)

    init_epochs = 165

    epoch = 0
    plateau_count = 0
    best_loss = None

    train_losses, train_accs, val_losses, val_accs = [], [], [], []

    while epoch < init_epochs and plateau_count < plateau_patience:

        # Train the model for an epoch
        loss, acc = train_epoch(model, train, optimizer, criterion)
        train_losses.append(loss)
        train_accs.append(acc)

        # Evaluate the model on the validation set
        val_loss, val_acc = eval_model(model, val, criterion)
        val_losses.append(val_loss)
        val_accs.append(val_acc)

        # Update the learning rate
        scheduler.step(val_loss)

        # Check for a plateau
        if best_loss is None or val_loss < best_loss:
            best_loss = val_loss
            plateau_count = 0
        else:
            plateau_count += 1
        
        epoch += 1

        # "If learning doesn’t plateu for some number of epochs,
        # we roughly double the number of epochs until it does."
        if epoch == init_epochs and plateau_count < plateau_patience:
            init_epochs *= 2

        print(f'Epoch {epoch}/{init_epochs} | Learning Rate: {optimizer.param_groups[0]["lr"]} | '
              f'Training loss: {train_losses[-1]:.4f} | '
              f'Validation loss: {val_losses[-1]:.4f} | '
              f'Validation accuracy: {val_accs[-1]:.4f}')
        
    return train_losses, train_accs, val_losses, val_accs

MODELS_DIR = 'models'
LOGS_DIR = 'logs'

for lr, bn in [(0.0001, False), (0.0001, True), (0.003, True), (0.1, True)]:

    s = f'lr={lr}' + ('_bn' if bn else '')
    print(s)

    model = models.resnet101(num_classes = 10)
    model.apply(xavier_init)

    if not bn: disable_bn(model)

    torch.save(model, f'{MODELS_DIR}/{s}_init.pth')
    data = train(model, train_loader, val_loader, init_lr = lr)
    torch.save(model, f'{MODELS_DIR}/{s}_end.pth')
    torch.save(data, f'{LOGS_DIR}/{s}.pth')

# code-summary: Plot results

get_x = lambda x: 100 * np.arange(1, len(x) + 1) / len(x)

f, (ax1, ax2) = plt.subplots(1, 2, figsize = (9, 4), sharey = True)

for fname in os.listdir(LOGS_DIR):

    train_losses, train_accs, val_losses, val_accs = torch.load(f'{LOGS_DIR}/{fname}')
    ax1.plot(get_x(train_accs), train_accs, label = fname[:-4])
    ax2.plot(get_x(val_accs), val_accs, label = fname[:-4])
    print(f'{fname[:-4]} took {len(train_accs)} epochs')


ax1.legend(); ax2.legend()
ax1.set_ylabel('Training accuracy'); ax2.set_ylabel('Validation accuracy')
ax1.set_xlabel('% of training'); ax2.set_xlabel('% of training')
f.tight_layout()

def init_func(f, also_linear = True):
    def init(m):
        if isinstance(m, nn.Conv2d):
            f(m.weight)
        if also_linear and isinstance(m, nn.Linear):
            f(m.weight)
    return init

criterion = nn.CrossEntropyLoss()
images, labels = next(iter(val_loader))
images, labels = images.to(device), labels.to(device)

with_bn, without_bn = {}, {}

for init_f in [
    nn.init.xavier_uniform_, nn.init.kaiming_uniform_,
    nn.init.kaiming_normal_, lambda x: nn.init.kaiming_normal_(x, mode = 'fan_out', nonlinearity='relu')]:

    init_name = init_f.__name__[:-1]
    if 'lambda' in init_name: init_name = '(default) kaiming_normal fan_out'
    
    for linear in (True, False):

        lin_name = ' w/linear lyrs' if linear else ''
        
        for bn, d in zip((True, False), (with_bn, without_bn)):
                
                model = models.resnet101(num_classes = 10)
                model.to(device)
                if not bn: disable_bn(model)
                model.apply(init_func(init_f, linear))

                model.train()
                output = model(images)
                loss = criterion(output, labels)
                loss.backward()

                model.eval()
                d[init_name + lin_name] = model.layer3[9].conv1.weight.grad.view(-1).cpu().detach().numpy()

for init_name, v in with_bn.items():
    f, (ax1, ax2) = plt.subplots(1, 2, figsize = (8, 3))
    sns.histplot(v, ax = ax1)
    ax1.set_title('with BN')
    sns.histplot(without_bn[init_name], ax = ax2)
    ax2.set_title('without BN')
    ax2.set_ylabel('')
    ax2.set_ylim(0, 1500)
    f.suptitle(init_name)
    f.tight_layout()

def fig3(model, init_lr, log_lrs, log_batches = [1, 4, 7]):

    # batch -> list of relative losses for each lr
    out = {}

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr = init_lr, momentum = 0.9, weight_decay = 5e-4)

    for i, (images, labels) in enumerate(train_loader):

        images, labels = images.to(device), labels.to(device)

        if i in log_batches:
            
            torch.save(model.state_dict(), f'models/model_state.tmp')
            torch.save(optimizer.state_dict(), f'models/optimizer_state.tmp')

            rel_losses = []
            for lr in log_lrs:

                for param_group in optimizer.param_groups: param_group['lr'] = lr

                optimizer.zero_grad()
                output = model(images)
                current_loss = criterion(output, labels)
                current_loss.backward()
                optimizer.step()

                with torch.no_grad():
                    output = model(images)
                    tmp_loss = criterion(output, labels)
                    rel_loss = (tmp_loss / current_loss).item()
                
                    # print learning rate, current loss, tmp loss, relative loss (at 4 decimal places)
                    print(f'{lr:.5f} {current_loss:.5f} {tmp_loss:.5f} {rel_loss:.5f}')
                    rel_losses.append(rel_loss)
            
                model.load_state_dict(torch.load('models/model_state.tmp'))
                optimizer.load_state_dict(torch.load('models/optimizer_state.tmp'))   

                # If loss is nan of int, break. Unlikely to recover.
                if torch.isnan(tmp_loss).item() or torch.isinf(tmp_loss).item():
                    rel_losses.pop()
                    print('breaking')
                    break

            out[i] = rel_losses
        
        if i == max(log_batches): break

        # take the actual step
        optimizer.zero_grad()
        output = model(images)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()
    
    return out


lrs = np.logspace(-5.5, 0.5, 80)

# With BN
model = models.resnet101(num_classes = 10)
model.to(device)
model.apply(init_net)
with_bn = fig3(model, init_lr = 0.1, log_lrs = lrs)

# Without BN
model = models.resnet101(num_classes = 10)
model.to(device)
disable_bn(model)
model.apply(init_net)
without_bn = fig3(model, init_lr = 0.0001, log_lrs = lrs)

f, axs = plt.subplots(2, 3, figsize = (12, 6))

for (batch, rel_losses), ax in zip(with_bn.items(), axs[0]):
    ax.plot(lrs, rel_losses)
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.set_ylim(0, 1.5)
    ax.set_xlim(10**-5.5, 0.4)
    ax.set_title(f'Batch {batch} with BN')
    ax.axhline(1, color = 'grey', linestyle = '--')

for (batch, rel_losses), ax in zip(without_bn.items(), axs[1]):
    ax.plot(lrs[:len(rel_losses)], rel_losses)
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.set_ylim(10**-2, 10**3)
    ax.set_xlim(10**-5.5, 0.4)
    ax.set_title(f'Batch {batch} w/o BN')
    ax.axhline(1, color = 'grey', linestyle = '--')

axs[0, 0].set_ylabel('Relative loss')
axs[1, 0].set_ylabel('Relative loss')
f.tight_layout()

df = []

def log_activation_stats(layer_name, key):
    def hook(module, input, output):
        with torch.no_grad():
            df.append({
                'bn': key,
                'layer': layer_name,
                'mean': output[:, 0, :, :].mean().abs().item(),
                'std': output[:, 0, :, :].std().abs().item()
            })
    return hook


for _ in range(10):

    for bn in [True, False]:

        model = models.resnet101(num_classes = 10)
        model.to(device)
        if not bn: disable_bn(model)
        model.apply(init_net)

        # Layers to log activations from
        log_layers = [] # (n, name, layer)
        n = 1
        for name, layer in model.named_modules():
            if 'conv' in name:
                n += 1
                if n in list(range(5, 101 + 12, 12)):
                    log_layers.append((n, name, layer))

        for n, _, layer in log_layers: layer.register_forward_hook(log_activation_stats(n, str(bn)))
        for images, _ in val_loader: model(images.to(device))

df = pd.DataFrame(df)

f, (ax1, ax2) = plt.subplots(1, 2, figsize = (9, 4))

sns.lineplot(data = df, x = 'layer', y = 'mean', hue = 'bn', ax = ax1)
ax1.set_yscale('log')

sns.lineplot(data = df, x = 'layer', y = 'std', hue = 'bn', ax = ax2)
ax2.set_yscale('log')

f.tight_layout()

f, (ax1, ax2) = plt.subplots(1, 2, figsize = (8, 3), sharey = True, sharex = True)

images, labels = next(iter(val_loader))
images, labels = images.to(device), labels.to(device)

for bn, ax in zip([False, True], [ax1, ax2]):

    model = models.resnet101(num_classes = 10)
    model.to(device)
    if not bn: disable_bn(model)
    model.apply(init_net)

    out = model(images)
    out.retain_grad()
    loss = nn.CrossEntropyLoss()(out, labels)
    loss.backward()

    ax = sns.heatmap(out.grad.cpu().detach().numpy(), cmap = 'viridis', ax = ax)
    ax.set_xticks([]); ax.set_yticks([0, 40, 80, 120]); ax.set_yticklabels([0, 40, 80, 120])
    ax.set_xlabel('Classes')
    ax.set_ylabel('Images in batch' if not bn else '')
    ax.set_title('With BN' if bn else 'Without BN')

f.tight_layout()

def fig6(init_func = init_net):

    f, (ax1, ax2) = plt.subplots(1, 2, figsize = (8, 3), sharex = True, sharey = True)

    for bn, ax in zip([False, True], [ax1, ax2]):

        model = models.resnet101(num_classes = 10)
        model.to(device)
        if not bn: disable_bn(model)
        model.apply(init_func)

        avg_grad = torch.zeros((128, 10), device = device)

        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            out = model(images)
            out.retain_grad()
            loss = nn.CrossEntropyLoss()(out, labels)
            loss.backward()
            if out.grad.shape == avg_grad.shape: avg_grad += out.grad
        
        avg_grad /= len(val_loader)

        ax = sns.heatmap(avg_grad.cpu().detach().numpy(), cmap = 'viridis', ax = ax)
        ax.set_xlabel('Classes')
        ax.set_ylabel('Images in batch')
        ax.set_title('With BN' if bn else 'Without BN')

    f.tight_layout()

fig6()