import os


import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from numpy.random import normal
import ray
from ray import tune
from ray.tune.schedulers import ASHAScheduler
import random
from torchvision.transforms import transforms

random.seed(100)

transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5,), (0.5,)),
                                ])

xy_trainPT = torchvision.datasets.MNIST(
    root="./",
    train=True,
    download=True
)

trainset = torchvision.datasets.MNIST(root="./",
                                      train=True,
                                      download=True,
                                      transform=transform)

originalSet = torchvision.datasets.MNIST(root="./",
                                         train=True,
                                         download=True,
                                         transform=transform)

noisyArr = []
originalArr = []

for index, shape in enumerate(trainset):
    noisyArr.append(shape[0].squeeze(dim=0).numpy())
    originalArr.append(originalSet[0][0].squeeze(dim=0).numpy())
    if index == 30000:
        break

noisyArr = np.array(noisyArr)
originalArr = np.array(originalArr)
print('done loading data')

original = originalArr / 255

X_2 = noisyArr / 255

for i in range(len(X_2)):
    norm = abs(np.random.normal(0, 0.3, size=(28, 28)))
    X_2[i] = X_2[i] + norm

pixels = int(784)


class autoencoder(nn.Module):
    def __init__(self, config):
        super(autoencoder, self).__init__()

        size = 28
        kernel = config['convK']
        # print(f"kernal: {kernel}")
        stride = config['convS']
        # print(f"stride: {stride}")
        padding = config['convP']
        #  print(f"padding: {padding}")
        poolK = config['poolK']
        poolS = config['poolS']
        finalOutput = config['actMap']
        self.conv1 = torch.nn.Conv2d(1, finalOutput, kernel_size=kernel, stride=stride, padding=padding)
        self.bn1 = torch.nn.BatchNorm2d(finalOutput)
        self.pool1 = torch.nn.MaxPool2d(stride=poolS, kernel_size=poolK)

        def poolAdjust(originalSize, kernel=poolK, stride=poolS, dilation=1):
            return ((originalSize - (dilation * (kernel - 1)) - 1) // stride) + 1

        def conv2d_size_out(size, kernel_size=kernel, stride=stride, padding=padding):
            return ((size + (padding * 2) - (kernel_size - 1) - 1) // stride) + 1

        convw = poolAdjust(conv2d_size_out(size))
        convh = poolAdjust(conv2d_size_out(size))

        self.linear_input_size = convw * convh * finalOutput

        self.head = torch.nn.Linear(self.linear_input_size, pixels)
        self.flatten = torch.nn.Linear(self.linear_input_size, self.linear_input_size)
        self.func = torch.nn.Hardtanh()
        self.softMax2d = torch.nn.Softmax2d()

    def forward(self, x):
        x = self.bn1(self.conv1(x))
        x = self.pool1(x)
        x = torch.nn.functional.relu(x)
        return self.head(x.view(x.size(0), -1))


def train(config, checkpoint_dir=None, data=None):

    # data = (X_2, original)
    loss_fn = torch.nn.MSELoss()
    model = autoencoder(config)
    optimizer = torch.optim.SGD(model.parameters(), lr=config['lr'], momentum=0.9)
    maxIter = 50000
    batchAmount = config['batchSize']

    if checkpoint_dir:
        checkpoint = os.path.join(checkpoint_dir, "checkpoint")
        model_state, optimizer_state = torch.load(checkpoint)
        model.load_state_dict(model_state)
        optimizer.load_state_dict(optimizer_state)

    for t in range(maxIter):
        epoch_loss = 0

        optimizer.zero_grad()
        idx = np.random.randint(data[0].shape[0], size=batchAmount)  # bootstrapping a subset of the total samples

        X_scaled = torch.unsqueeze(torch.from_numpy(data[0][idx, :]).float(), dim=1)  # creating tensor for convultion

        testValues = torch.from_numpy(
            np.reshape(data[1][idx, :],
                       (batchAmount, -1))
        ).float()  # creating a flattened array for testing

        y_pred = model(X_scaled)  # predict on the subset
        # print(y_pred.size())
        loss = loss_fn(testValues, y_pred)  # get loss on subset

        if t % (maxIter / 10) == 0:
            # print(t, loss.item())
            tune.report(score=loss.item())
            with tune.checkpoint_dir(step=t) as checkpoint_dir:
                path = os.path.join(checkpoint_dir, "checkpoint")
                torch.save(
                    (model.state_dict(), optimizer.state_dict()), path)

        loss.backward()  # get gradient stuff
        optimizer.step()  # optimize
        epoch_loss += loss.item()

def tunerTrain():
    ray.init(_memory=8000000000, object_store_memory=4000000000, _redis_max_memory=8000000000, num_cpus=5),
    resources_per_trial = {"cpu": 4, "extra_cpu": 1},

    searchSpace = {
        'lr': tune.loguniform(1e-4, 9e-1),
        'actMap': tune.grid_search([1, 2]),
        'convK': tune.choice([3, 5, 7, 9]),
        'convS': tune.grid_search([1, 2]),
        'convP': tune.choice([0, 1, 2, 3]),
        'poolK': tune.choice([3, 5, 7, 9]),
        'poolS': tune.grid_search([1, 2]),
        'batchSize': tune.choice([2, 4, 8, 16, 32, 64, 128, 256]),
    }

    analysis = tune.run(tune.with_parameters(train, data=[X_2, original]), num_samples=10, metric='score', mode='min',
                        scheduler=ASHAScheduler(),
                        config=searchSpace)
    dfs = analysis.trial_dataframes
    print(f"Best Config: {analysis.get_best_config('score', mode='min')}")
    df = analysis.results_df
    logdir = analysis.get_best_logdir("score", mode="min")
    print(f"dir of best: {logdir}")
    print(analysis.best_result)
    print(f"Best trial final score: {analysis.get_best_trial('score', mode='min')}")

tunerTrain()