np.random.seed(42)
a = np.random.randn(1)
b = np.random.randn(1)
print(a, b)

# Computes our model's predicted output
yhat = a + b * x_train

# How wrong is our model? That's the error!
error = (y_train - yhat)

# It is a regression, so it computes mean squared error (MSE)
loss = (error ** 2).mean()
print(f"loss: {loss}")

# Computes gradients for both "a" and "b" parameters
a_grad = -2 * error.mean()
b_grad = -2 * (x_train * error).mean()
print(a_grad, b_grad)

# Sets learning rate
lr = 1e-1

# Updates parameters using gradients and the learning rate
print(a, b)
a = a - lr * a_grad
b = b - lr * b_grad

print(a, b)

# Define number of epochs
n_epochs = 1000

# Step 0
np.random.seed(42)
a = np.random.randn(1)
b = np.random.randn(1)

for epoch in range(n_epochs):
    # Step 1:
    # Compute our model's predicted output
    yhat = a + b * x_train

    # Step 2:
    # How wrong is our model? That's the error!
    error = (y_train - yhat)
    # It is a regression, so it computes mean squared error (MSE)
    loss = (error ** 2).mean()

    # Step 3:
    # Compute gradients for both "a" and "b" parameters
    a_grad = -2 * error.mean()
    b_grad = -2 * (x_train * error).mean()

    # Step 4:
    # Update parameters using gradients and the learning rate
    a -= lr * a_grad
    b -= lr * b_grad
print(a, b)

# Sanity Check: do we get the same results as our gradient descent?
from sklearn.linear_model import LinearRegression
linr = LinearRegression()
linr.fit(x_train, y_train)
print(linr.intercept_, linr.coef_[0])

# Step 0
torch.manual_seed(42)

a = torch.randn(1, requires_grad=True, dtype=torch.float, device=device)
b = torch.randn(1, requires_grad=True, dtype=torch.float, device=device)

# Step 1
# Compute our model's predicted output
yhat = a + b * x_train_tensor

# Step 2
# How wrong is our model? That's the error!
error = (y_train_tensor - yhat)
# It is a regression, so it computes mean squared error (MSE)
loss = (error ** 2).mean()

# Step 3
# No more manual computation of gradients!
loss.backward()

# Computes gradients for both "a" and "b" parameters
# a_grad = -2 * error.mean()
# b_grad = -2 * (x_train_tensor * error).mean()

print(f"Pytorch gradient: {a.grad}, {b.grad}")

a.grad.zero_(), b.grad.zero_()

# Step 0
torch.manual_seed(42)
a = torch.randn(1, requires_grad=True, dtype=torch.float, device=device)
b = torch.randn(1, requires_grad=True, dtype=torch.float, device=device)

# Step 1
# Compute our model's predicted output
yhat = a + b * x_train_tensor

# Step 2
# How wrong is our model? That's the error!
error = (y_train_tensor - yhat)
# It is a regression, so it computes mean squared error (MSE)
loss = (error ** 2).mean()

# Step 3
# No more manual computation of gradients!
loss.backward()
# Computes gradients for both "a" and "b" parameters
# a_grad = -2 * error.mean()
# b_grad = -2 * (x_train_tensor * error).mean()
print(a.grad, b.grad)

# Step 4
# Update parameters using gradients and the learning rate
with torch.no_grad(): # what is that?!
    a -= lr * a.grad
    b -= lr * b.grad

# PyTorch is "clingy" to its computed gradients, we need to tell it to let it go...
a.grad.zero_()
b.grad.zero_()

print(a.grad, b.grad)

lr = 1e-1
n_epochs = 1000

# Step 0
torch.manual_seed(42)
a = torch.randn(1, requires_grad=True, dtype=torch.float, device=device)
b = torch.randn(1, requires_grad=True, dtype=torch.float, device=device)

for epoch in range(n_epochs):
    # Step 1
    # Compute our model's predicted output
    yhat = a + b * x_train_tensor

    # Step 2
    # How wrong is our model? That's the error!
    error = (y_train_tensor - yhat)
    # It is a regression, so it computes mean squared error (MSE)
    loss = (error ** 2).mean()

    # Step 3
    # No more manual computation of gradients!
    loss.backward()

    # Step 4
    # Update parameters using gradients and the learning rate
    with torch.no_grad():
        a -= lr * a.grad
        b -= lr * b.grad

    # PyTorch is "clingy" to its computed gradients, we need to tell it to let it go...
    a.grad.zero_()
    b.grad.zero_()

print(a, b)

print("loss", loss)

loss.detach().cpu().numpy()

print(loss.item(), loss.tolist())

# Our parameters
torch.manual_seed(42)
a = torch.randn(1, requires_grad=True, dtype=torch.float, device=device)
b = torch.randn(1, requires_grad=True, dtype=torch.float, device=device)

# Learning rate
lr = 1e-1

# Defines a SGD optimizer to update the parameters
optimizer = optim.SGD([a, b], lr=lr)

n_epochs = 1000

for epoch in range(n_epochs):
    # Step 1
    yhat = a + b * x_train_tensor

    # Step 2
    error = y_train_tensor - yhat
    loss = (error ** 2).mean()

    # Step 3, compute gradients
    loss.backward()

    # Step 4, apply gradients to update parameters
    # No more manual update!
    # with torch.no_grad():
    #     a -= lr * a.grad
    #     b -= lr * b.grad
    optimizer.step()

    # No more telling PyTorch to let gradients go!
    # a.grad.zero_()
    # b.grad.zero_()
    optimizer.zero_grad()

print(a, b)

class ManualLinearRegression(nn.Module):
  def __init__(self):
      super().__init__()
      # parameter tensors need their gradient
      a = torch.randn(1, requires_grad=True, dtype=torch.float)
      b = torch.randn(1, requires_grad=True, dtype=torch.float)

      # To make "a" and "b" real parameters of the model, we need to
      # wrap them with nn.Parameter
      self.a = nn.Parameter(a)
      self.b = nn.Parameter(b)

  def forward(self, x):
      # Computes the outputs / predictions
      return self.a + self.b * x

dummy = ManualLinearRegression()
list(dummy.parameters())
# Returns: [Parameter containing:
# tensor([2.6584], requires_grad=True),
# Parameter containing:
# tensor([1.2004], requires_grad=True)]

dummy.state_dict()
# Returns: OrderedDict([('a', tensor([2.6584])), ('b', tensor([1.2004]))])

optimizer.state_dict()

torch.manual_seed(42)
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# Now we can create a model and send it at once to the device
model = ManualLinearRegression().to(device)

# We can also inspect its parameters using its state_dict
print(model.state_dict())

yhat = model(x_train_tensor)

lr = 1e-1
n_epochs = 1000

loss_fn = nn.MSELoss(reduction='mean')

# Now the optimizers uses the parameters from the model
optimizer = optim.SGD(model.parameters(), lr=lr)

for epoch in range(n_epochs):
    # Sets model to training mode
    model.train()

    # Step 1
    # No more manual prediction!
    # yhat = a + b * x_tensor
    yhat = model(x_train_tensor)

    # Step 2, compute loss, sum of errors
    loss = loss_fn(yhat, y_train_tensor)
    # Step 3, compute gradients
    loss.backward()
    # Step 4, update parameters, and zero out gradient
    optimizer.step()
    optimizer.zero_grad()

print(model.state_dict())

class LayerLinearRegression(nn.Module):
    def __init__(self):
        super().__init__()
        # Instead of our custom parameters, we use a Linear layer with
        # single input and single output
        self.linear = nn.Linear(1, 1)

    def forward(self, x):
        # Now it only takes a call to the layer to make predictions
        return self.linear(x)

dummy = LayerLinearRegression()
list(dummy.parameters())
# -> OrderedDict([('linear.weight', tensor([[0.4591]])),
#                 ('linear.bias', tensor([-0.7359]))])
dummy.state_dict()
# -> LayerLinearRegression( (linear): Linear(in_features=1, out_features=1, bias=True))

model = nn.Sequential(nn.Linear(1, 1)).to(device)

def make_train_step(model, loss_fn, optimizer):
    # Builds function that performs a step in the train loop
    def train_step(x, y):
        # Sets model to TRAIN mode
        model.train()
        # Step 1: Make predictions
        yhat = model(x)
        # Step 2: Compute loss
        loss = loss_fn(yhat, y)
        # Step 3: Compute gradients
        loss.backward()
        # Step 4: Update parameters and zeroes gradients
        optimizer.step()
        optimizer.zero_grad()
        # Returns the loss
        return loss.item()

    # Returns the function that will be called inside the train loop
    return train_step

lr = 1e-1

# Create a MODEL, a LOSS FUNCTION and an OPTIMIZER
model = nn.Sequential(nn.Linear(1, 1)).to(device)
loss_fn = nn.MSELoss(reduction='mean')
optimizer = optim.SGD(model.parameters(), lr=lr)

# Create the train_step function for our model, loss function and optimizer
train_step = make_train_step(model, loss_fn, optimizer)

n_epochs = 1000

losses = []
# For each epoch...
for epoch in range(n_epochs):
    # Performs one train step and returns the corresponding loss
    loss = train_step(x_train_tensor, y_train_tensor)
    losses.append(loss)

# Checks model's parameters
print(model.state_dict())

plt.plot(losses[:200])
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.yscale('log')
plt.savefig('/tmp/trainingloss.png', bbox_inches="tight")

from torch.utils.data import TensorDataset
train_data = TensorDataset(x_train_tensor, y_train_tensor)
print(train_data[0])

from torch.utils.data.dataset import random_split

# build tensors from numpy arrays BEFORE split
x_tensor = torch.from_numpy(x).float()
y_tensor = torch.from_numpy(y).float()

# build dataset containing ALL data points
dataset = TensorDataset(x_tensor, y_tensor)

# perform the split (could do, e.g. a [60,20,20] split as well)
train_dataset, val_dataset = random_split(dataset, [80, 20])

# build a loader of each set
train_loader = DataLoader(dataset=train_dataset, batch_size=16)
val_loader = DataLoader(dataset=val_dataset, batch_size=20)

checkpoint = {'epoch': n_epochs,
              'model_state_dict': model.state_dict(),
              'optimizer_state_dict': optimizer.state_dict(),
              'loss': losses,
              'val_loss': val_losses}

torch.save(checkpoint, 'model_checkpoint.pth')

checkpoint = torch.load('model_checkpoint.pth')

model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

epoch = checkpoint['epoch']
losses = checkpoint['loss']
val_losses = checkpoint['val_loss']

from torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau, MultiStepLR

optimizer = optim.SGD(model.parameters(), lr=lr)
scheduler = ReduceLROnPlateau(optimizer, 'min')

#scheduler = StepLR(optimizer, step_size=30, gamma=0.5)
#scheduler = MultiStepLR(optimizer, milestones=[30,80], gamma=0.1)

def train_loop_with_scheduler(model, loss_fn, optimizer, scheduler, n_epochs, train_loader, val_loader=None):
    # Device from where the model parameters (hence, the model) are
    device = next(model.parameters()).device.type
    # Create the train_step function for our model, loss function and optimizer
    train_step = make_train_step(model, loss_fn, optimizer)

    losses = []
    val_losses = []
    learning_rates = []

    for epoch in range(n_epochs):
        # TRAINING
        batch_losses = []
        for x_batch, y_batch in train_loader:
            x_batch = x_batch.to(device)
            y_batch = y_batch.to(device)

            loss = train_step(x_batch, y_batch)
            batch_losses.append(loss)

        losses.append(np.mean(batch_losses))

        # VALIDATION
        if val_loader is not None:
            val_loss = validation(model, loss_fn, val_loader)
            val_losses.append(val_loss)

        print(f"Epoch {epoch} complete...")

        # SCHEDULER
        if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
            scheduler.step(val_loss)
        else:
            scheduler.step()

        learning_rates.append(optimizer.state_dict()['param_groups'][0]['lr'])

    return losses, val_losses, learning_rates

torch.manual_seed(42)

# builds tensors from numpy arrays BEFORE split
x_tensor = torch.from_numpy(x).float()
y_tensor = torch.from_numpy(y).float()

# builds dataset containing ALL data points
dataset = TensorDataset(x_tensor, y_tensor)

# performs the split
train_dataset, val_dataset = random_split(dataset, [80, 20])

# builds a loader of each set
train_loader = DataLoader(dataset=train_dataset, batch_size=16)
val_loader = DataLoader(dataset=val_dataset, batch_size=20)

# defines learning rate
lr = 1e-1

# Create a MODEL, a LOSS FUNCTION and an OPTIMIZER (and SCHEDULER)
model = nn.Sequential(nn.Linear(1, 1)).to(device)
loss_fn = nn.MSELoss(reduction='mean')
optimizer = optim.SGD(model.parameters(), lr=lr)

scheduler = ReduceLROnPlateau(optimizer, 'min')
#scheduler = StepLR(optimizer, step_size=30, gamma=0.5)
#scheduler = MultiStepLR(optimizer, milestones=[30,80], gamma=0.1)

n_epochs = 1000

losses, val_losses, l_rates = train_loop_with_scheduler(model, loss_fn, optimizer, scheduler, n_epochs, train_loader, val_loader)

print(model.state_dict())

fig, ax1 = plt.subplots()
ax2_col = "green"
ax1.set_xlim(0, 100)
ax2 = ax1.twinx()

ax1.plot(losses, label='Training Loss', linestyle="-", linewidth=2)
ax1.plot(val_losses, label='Validation Loss', linestyle=":", linewidth=2)
ax1.set_xlabel('Epochs')
ax1.set_ylabel('Loss')
ax1.set_yscale("log")

ax2.set_yscale("log")
ax2.plot(l_rates, label='Learning rate', linestyle="--", linewidth=1, color=ax2_col)
ax2.set_ylabel('Learning Rate', color=ax2_col)
ax2.tick_params(axis='y', labelcolor=ax2_col)

ax1.legend()
fig.tight_layout()
fig.savefig("/tmp/lr_scheduler.png", bbox_inches="tight")
plt.show()

model = nn.DataParallel(model)

from ax import optimize
best_parameters, best_values, _, _ = optimize(
    parameters=[
        {"name": "x1",
         "type": "range",
         "bounds": [-10.0, 10.0],},
        {"name": "x2",
         "type": "range",
         "bounds": [-10.0, 10.0],},],
    evaluation_function=booth,
    minimize=True,)print(best_parameters)

import numpy as np
import torch
import torch.optim as optim
import torch.nn as nn
import torchvision as tv

import torch.nn.functional as F
from torch.utils.data import DataLoader

from torchviz import make_dot
from torch.optim.lr_scheduler import ReduceLROnPlateau

epochs = 20
batch_size = 32
path = "/tmp"

# Compression of factor 24.5, assuming the input is 784 float
embedding_dim = 32

class autoencoder(nn.Module):
    def __init__(self, dim=32, **kwargs):
        super().__init__()
        assert(dim <= 64)

        # Note: We define the nn.<model> in __init__ because they have
        # learnable parameters. Most handy to use the torch.nn module
        # for that.

        # Encoder, 4 fully connected / "dense" layers
        self.fc1 = nn.Linear(kwargs["input_shape"], 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, dim)

        # Decoder, the encoder in reverse
        self.fc6 = nn.Linear(dim, 64)
        self.fc7 = nn.Linear(64, 128)
        self.fc8 = nn.Linear(128, kwargs["input_shape"])

    def forward(self, x):
        # Pooling and ReLU don't have learnable parameters, so usually
        # goes here. More convinient to use torch.nn.functional for them.

        # Encoder
        x = self.fc1(x)
        x = F.relu(x)         # (ReLU(x) = max(0,x))

        x = self.fc2(x)
        x = F.relu(x)

        x = self.fc3(x)
        x = F.relu(x)

        # Decoder
        x = self.fc6(x)
        x = F.relu(x)

        x = self.fc7(x)
        x = F.relu(x)

        x = self.fc8(x)
        # Output layer, for scaling between 0 to 1
        # x = torch.sigmoid(x)
        x = torch.tanh(x)
        return x

model = autoencoder(32, input_shape=28*28)
model.fc3.weight

class autoencoder(nn.Module):
    def __init__(self, dim=32, **kwargs):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(kwargs["input_shape"], 128),
            nn.ReLU(True),
            nn.Linear(128, 64),
            nn.ReLU(True),
            nn.Linear(64, dim))
        self.decoder = nn.Sequential(
            nn.Linear(dim, 64),
            nn.ReLU(True),
            nn.Linear(64, 128),
            nn.ReLU(True),
            nn.Linear(128, kwargs["input_shape"]),
            nn.Tanh())

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

model = autoencoder(32, input_shape=28*28)
model.encoder[4].weight

print(model.encoder)
 Sequential(
   (0): Linear(in_features=784, out_features=128, bias=True)
   (1): ReLU(inplace=True)
   (2): Linear(in_features=128, out_features=64, bias=True)
   (3): ReLU(inplace=True)
   (4): Linear(in_features=64, out_features=32, bias=True))

def to_img(x):
    "De-Normalize MNIST images, for checking training"
    x = 0.5 * (x + 1)
    x = x.clamp(0, 1)
    x = x.view(x.size(0), 1, 28, 28)
    return x

# use GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# create a model from `AE` autoencoder class
# load it to the specified device, either gpu or cpu
# model = ae(embedding_dim, input_shape=28*28).to(device)
model = autoencoder(embedding_dim, input_shape=28*28).to(device)

# create an optimizer object
# Adam optimizer with learning rate 1e-3; can also play with weight_decay=1e-5
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# Mean-Squared Error Loss. Defaults to reduction='mean', such that we
# normalize with number of samples, or else we can not compare loss
# from different sized batches (like train vs test/validation batch).
criterion = nn.MSELoss(reduction='mean')

# Load MNIST dataset
transform = tv.transforms.Compose([tv.transforms.ToTensor()])
train_dataset = tv.datasets.MNIST(
    root=path, train=True, transform=transform, download=True)
test_dataset = tv.datasets.MNIST(
    root=path, train=False, transform=transform, download=True)

# Create dataloader, (use 4 sub-processes for data loading)
train_loader = DataLoader(
    train_dataset, batch_size=128, shuffle=True, num_workers=4, pin_memory=True)
test_loader = DataLoader(
    test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

train_losses = []
val_losses = []
for epoch in range(1, epochs+1):
    # TRAINING
    loss = 0
    for data in train_loader:
        # "_" = labels, don't need them
        batch_features, _ = data

        # Flatten / reshape mini-batch data from [N, 28, 28] to [N, 784] matrix
        batch_features = batch_features.view(-1, 784)

        # Load it to the active device, where the model lives
        batch_features = batch_features.to(device)

        # Forward pass, predict/reconstruct output
        outputs = model(batch_features)

        # Compute training reconstruction loss
        train_loss = criterion(outputs, batch_features)

        # Don't accumulate gradients on subsequent backward passes
        optimizer.zero_grad()

        # Backward pass: compute gradient of the loss with respect to
        # model parameters
        train_loss.backward()

        # Perform single parameter update step based on current gradients
        optimizer.step()

        # Add the mini-batch training loss to epoch loss
        loss += train_loss.item()

    # compute the epoch training loss
    loss = loss / len(train_loader)
    train_losses.append(loss)

    # VALIDATION
    val_loss = 0
    # no gradients in validation!
    with torch.no_grad():
        for data in test_loader:
            batch_features, _ = data

            # Flatten / reshape mini-batch data from [N, 28, 28] to [N, 784] matrix
            batch_features = batch_features.view(-1, 784)

            # Load it to the active device, where the model lives
            batch_features = batch_features.to(device)

            # set model in eval mode
            model.eval()

            # make prediction
            outputs = model(batch_features)

            # Compute mini-batch val. loss reconstruction to epoch loss
            val_loss += criterion(outputs, batch_features).item()

    val_loss = val_loss / len(test_loader)
    val_losses.append(val_loss)

    # Display the epoch training loss
    print(f"epoch: {epoch}/{epochs}\t train_loss: {loss:.6f}\t val_loss: {val_loss:.6f}")

    # Print out example image every 10 epoch
    if epoch % 10 == 0:
        pic = to_img(outputs.cpu().data)
        tv.utils.save_image(pic, f'{path}/image_{epoch}.png')

def plot_train_val(train, val, path):
    "Plot train and validation curves, on log and lin plots"
    import matplotlib.pyplot as plt

    plt.style.use('fivethirtyeight')
    fig, ax = plt.subplots(1, 2, figsize=(12, 4))

    ax[0].plot(train, label='Training Loss')
    ax[0].plot(val, label='Validation Loss')
    ax[0].set_yscale('log')
    ax[0].set_xlabel('Epochs')
    ax[0].set_ylabel('Loss')
    ax[0].set_title("Loss on log-lin")
    ax[0].legend()

    ax[1].plot(train, label='Training Loss')
    ax[1].plot(val, label='Validation Loss')
    ax[1].set_xlabel('Epochs')
    ax[1].set_ylabel('Loss')
    ax[1].set_title("Loss on lin-lin")
    ax[1].legend()

    fig.savefig(path + "/val_loss.png", bbox_inches="tight")
    fig.show()


plot_train_val(train_losses, val_losses, path)


def plot_images(loader, model, outpath):
    """
    Plot input and output images of the autoencoder,
    for comparison

    Parameters
    ----------
    loader: torch.utils.data.dataloader.DataLoader
        Pytorch loader for data, of MNIST 28x28 images

    model: pytorch model
        Autoencoder, returns image of same format as input

    outpath: str
        Folder to put data in
    """
    import matplotlib.pyplot as plt

    # obtain one batch of test images
    dataiter = iter(loader)
    images, labels = dataiter.next()
    images_flatten = images.view(images.size(0), -1)

    # Get sample outputs
    output = model(images_flatten.to(device))

    # Output is resized into a batch of images,
    output = output.view(loader.batch_size, 1, 28, 28)

    # Turn off gradient on tensor, move to CPU, such that -> numpy
    output = output.detach().cpu().numpy()

    # Prep images for display
    images = images.numpy()

    # plot the first ten input images and then reconstructed images
    fig, axes = plt.subplots(nrows=2, ncols=10, sharex=True,
                             sharey=True, figsize=(25, 4))

    # input images on top row, reconstructions on bottom
    for images, row in zip([images, output], axes):
        for img, ax in zip(images, row):
            ax.imshow(np.squeeze(img), cmap='gray')
            ax.get_xaxis().set_visible(False)
            ax.get_yaxis().set_visible(False)
    fig.savefig(outpath + "/mnist_ae.png", bbox_inches="tight")
    fig.show()


plot_images(test_loader, model, path)

# # plot net
# dot = make_dot(output)
# dot.format = 'png'
# dot.render(path + "/net2")


def save(model, outpath='model_checkpoint.pth'):
    checkpoint = {'epoch': epochs,
                  'model_state_dict': model.state_dict(),
                  'optimizer_state_dict': optimizer.state_dict(),
                  'loss': train_losses,
                  'val_loss': val_losses}
    torch.save(checkpoint, outpath)


def load(model, outpath='model_checkpoint.pth'):
    checkpoint = torch.load(outpath)

    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

    epochs = checkpoint['epoch']
    train_losses = checkpoint['loss']
    val_losses = checkpoint['val_loss']
    return epochs, train_losses, val_losses, model