import torch
import torch.nn as nn
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from torch.utils.data import random_split
from torch.utils.data import Subset
import torch.nn.functional as F
import pandas as pd
import os



# Transformations (tensor + normalisation classique MNIST)
Tr = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,), (0.3081,)) ])

import torch
from torchvision import datasets, transforms
from torch.utils.data import random_split
from torch.utils.data import Subset


class SubDataset(torch.utils.data.Dataset):
    def __init__(self, base_dataset, indices, transform=None):
        self.base_dataset = base_dataset
        self.indices      = indices
        self.transform    = transform

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

    def __getitem__(self, i):
        x, y = self.base_dataset[self.indices[i]]
        if self.transform:
            x = self.transform(x)
        return x, y

def LoadDS(dataset, transform_train, transform_valid, split=0.8) :

    nb = len(dataset)
    n_train = int(split * nb)

    g = torch.Generator().manual_seed(0)
    idx = torch.randperm(nb, generator=g).tolist()
    train_idx = idx[:n_train]
    val_idx   = idx[n_train:]

    train_ds = SubDataset(dataset, train_idx, transform_train)  # 80% des images
    val_ds   = SubDataset(dataset, val_idx,   transform_valid)  # 20% des images

    print("Dataset LOADED")
    return train_ds, val_ds


dataset = datasets.MNIST(".\data", download = True, transform = None)
train_ds, val_ds = LoadDS(dataset, Tr,Tr)




#########################################################



class MNISTCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 8, kernel_size=3, padding=1)   # petit
        self.conv2 = nn.Conv2d(8, 16, kernel_size=3, padding=1)  # petit
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(16 * 7 * 7, 32)
        self.fc2 = nn.Linear(32, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool(x)               # 28x28 -> 14x14
        x = F.relu(self.conv2(x))
        x = self.pool(x)               # 14x14 -> 7x7
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x


################################################################


class Scenario : pass

def createScenario(lr):

    S = Scenario()
    S.log_folder = r"C:\log\test1"
    S.CSVname = os.path.join(S.log_folder, "LR_" + format(lr, ".0e") +".csv")


    # train/val
    S.epochs = 100
    S.batch_size = 32
    S.train_batch = DataLoader(train_ds, batch_size=S.batch_size, shuffle=True)
    S.valid_batch = DataLoader(val_ds,   batch_size=S.batch_size, shuffle=False)


    # model / optimizer
    S.device    = "cpu"  # torch.device("cuda" if torch.cuda.is_available() else "cpu")
    S.model     = MNISTCNN()
    S.loss_fn   = nn.CrossEntropyLoss()
    S.optimizer = torch.optim.SGD(S.model.parameters() , lr= lr )

    return S




#######################################################################

from torch.utils.tensorboard import SummaryWriter
def train(S):

    logger = SummaryWriter(S.log_folder)
    info  = []

    S.model = S.model.to(S.device)
    for epoch in range(S.epochs):

        # TRAIN

        train_GlobLoss = train_accuracy = 0
        S.model.train()
        for x, y in S.train_batch:
            x = x.to(S.device)
            y = y.to(S.device)

            S.optimizer.zero_grad()
            logits = S.model(x)
            loss = S.loss_fn(logits,y)
            loss.backward()
            train_GlobLoss += loss.item() * x.size(0)
            S.optimizer.step()

            # accuracy
            predictions     = logits.argmax(dim=1)
            train_accuracy += (predictions == y).sum().item()

        # VALIDATION

        valid_GlobLoss = valid_accuracy = 0
        S.model.eval()
        with torch.no_grad():
            for x, y in S.valid_batch:
                x = x.to(S.device)
                y = y.to(S.device)
                logits = S.model(x)
                loss = S.loss_fn(logits, y)
                valid_GlobLoss += loss.item() * x.size(0)

                # accuracy
                predictions     = logits.argmax(dim=1)
                valid_accuracy += (predictions == y).sum().item()

        # indicateurs
        nbtrain, nbvalid = len(S.train_batch.dataset), len(S.valid_batch.dataset)
        train_GlobLoss  /= nbtrain
        train_accuracy  /= nbtrain
        valid_GlobLoss  /= nbvalid
        valid_accuracy  /= nbvalid

        print( f"{epoch+1}/{S.epochs} - "
            f"tLoss {train_GlobLoss:.3f}  - vLoss {valid_GlobLoss:.3f} - "
            f"tAcc  {train_accuracy:.3f}  -  vAcc {valid_accuracy:.3f}")

        logger.add_scalars("Loss",      {"train": train_GlobLoss, "valid": valid_GlobLoss}, epoch+1)
        logger.add_scalars("Accuracy",  {"train": train_accuracy, "valid": valid_accuracy}, epoch+1)

        info.append([epoch+1,train_accuracy,valid_accuracy,train_GlobLoss,valid_GlobLoss])

    logger.close()
    # CSV
    columns = ["epoch", "acc/train", "acc/valid", "loss/train", "loss/valid"]
    df = pd.DataFrame(info, columns=columns)
    df.to_csv( S.CSVname , index = False )



for lr in  [1e-3, 3e-4, 1e-4, 3e-5 ] :
    S = createScenario(lr)
    train(S)