#!/usr/bin/env python3

from argparse import ArgumentParser
from time import time
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import pandas as pd
from os import makedirs
from copy import deepcopy
from linear_module import ActivLinear, calculate_parameters
from convolutional_module import ActivConvolutional
from utils import experiment_slug, one_hot_encode, timethis


class Net(nn.Module):
    def __init__(self, width, n_layers, ds_features, ds_classes, activation, args, image_size):
        super().__init__()
        self.width = width
        self.n_layers = n_layers
        self.ds_features = ds_features
        self.image_size = image_size # Tuple(width, height)
        self.ds_classes = ds_classes
        self.network_type = args.network
        self.args = args
        if args.network == "FCN":
            # Activation is applied before a fully connected layer
            self.layer0 = ActivLinear(ds_features, width, renormalize=args.renormalize) # identity activation
            for i in range(1, self.n_layers):
                setattr(self, f"layer{i}", ActivLinear(width, width, renormalize=args.renormalize, activation=activation))
            self.latest = ActivLinear(width, ds_classes, renormalize=args.renormalize, activation=activation)
        elif args.network == "CNN":
            self.layer0 = ActivConvolutional(ds_features, width, image_size=image_size, renormalize=args.renormalize, kernel_size=args.kernel_size)
            # We are removing kernel_size - 1 pixels from each side
            image_size = (image_size[0] - (args.kernel_size - 1), image_size[1] - (args.kernel_size - 1))
            for i in range(1, self.n_layers):
                setattr(self, f"layer{i}", ActivConvolutional(width, width, image_size=image_size, renormalize=args.renormalize, activation=activation, kernel_size=args.kernel_size))
                image_size = (image_size[0] - (args.kernel_size - 1), image_size[1] - (args.kernel_size - 1))
            in_pixels = image_size[0] * image_size[1]
            self.latest = ActivLinear(width * in_pixels, ds_classes, image_size=image_size, renormalize=args.renormalize, activation=activation)
        else:
            raise Exception(f"Invalid network type in arguments: {args.network}")

    def forward(self, x):
        if self.network_type == "FCN":
            x = torch.flatten(x, start_dim=1, end_dim=3)
        # Normalize inputs dividing by norm such that ||x|| = sqrt(features) for each example
        x = math.sqrt(self.ds_features) * F.normalize(x, p=2, dim=1)
        for i in range(self.n_layers):
            x = getattr(self, f"layer{i}")(x)
        if self.network_type == "CNN":
            x = torch.flatten(x, start_dim=1, end_dim=3)
        x = self.latest(x)
        return x

    def _iterate_layers(self, fn):
        for i in range(self.n_layers):
            yield from fn(getattr(self, f"layer{i}"), i)
        yield from fn(self.latest, "latest")

    def init_weights_(self):
        for i in range(self.n_layers):
            getattr(self, f"layer{i}").init_weights_()
        self.latest.init_weights_()

    def cache_initial_inputs(self):
        for i in range(self.n_layers):
            getattr(self, f"layer{i}").cache_initial_inputs()
        self.latest.cache_initial_inputs()
        
    def measure_input_matrix(self):
        def immeas(layer, i):
            mindiag, maxoff, maxsumoff, meansumoff, maxeig, mineig = layer.measure_input_matrix()
            yield dict(type="input_matrix", layer=i, min_ondiagonal=mindiag, max_offdiagonal=maxoff, max_sumoffdiagonal=maxsumoff, mean_sumoffdiagonal=meansumoff, max_eigenvalue=maxeig, min_eigenvalue=mineig)
        if "input_matrix" not in self.args.no_measure:
            yield from self._iterate_layers(immeas)
        
    def measure_distances(self):
        def dmeas(layer, i):
            distances = layer.measure_distances()
            distances.update(dict(type="distance", layer=i))
            yield distances
        if "distance" not in self.args.no_measure:
            yield from self._iterate_layers(dmeas)

    def measure_input_conditioning(self):
        def icmeas(layer, i):
            low, high = layer.measure_approximate_input_conditioning()
            yield dict(type="iconditioning", layer=i, lowest=low, highest=high)
        if "iconditioning" not in self.args.no_measure:
            yield from self._iterate_layers(icmeas)
        
    def measure_norms(self):
        def nmeas(layer, i):
            matrix2, matrixinf, vector2, vectorinf = layer.measure_norms()
            yield dict(type="norms", layer=i, matrix_norm_2=matrix2, matrix_norm_inf=matrixinf, vector_norm_2=vector2, vector_norm_inf=vectorinf)
        if "norms" not in self.args.no_measure:
            yield from self._iterate_layers(nmeas)


def trainstep(model, device, optimizer, test_loader, data=None, target=None, args=None):
    """
    Does a step of training and measuring together.
    Measures characteristics of the current model, and emits them.

    We measure:
    - Test loss and percentage of correct predictions
    - Conditioning of each layer (largest and lowest eigenvalue of kernel matrix)
    - Conditioning with respect to the x variable! Kernel on gradients
    - Measure Lipschitzianity of the network for the use in generalization bounds
    - Distance of each layer from its initialization in L2 and L\infty norm
    """
    lfs = dict(
        CE=F.cross_entropy,
        MSE=F.mse_loss,
        PMSE=F.mse_loss,
        KL=F.kl_div,
    )
    loss_function = lfs[args.loss_function]
    
    loss = None
    taken_time = None
    if target is not None:
        model.train()
        data, target = data.to(device), target.to(device)
        onehot_target = one_hot_encode(target, args.classes)
        # Measure how long it takes for a forward-backward pass
        tstart = time()
        optimizer.zero_grad()
        output = None
        output = model(data)
        if args.loss_function == "CE":
            # Here the logsoftmax is inside the loss function already
            loss = loss_function(output, target, reduction="mean")
        elif args.loss_function == "KL":
            # Don't calculate it in this way since it is exactly the same as doing CE.
            # output should be interpreted as log-likelihoods
            logoutput = F.log_softmax(output, dim=1)
            loss = loss_function(logoutput, onehot_target, reduction="batchmean")
        elif args.loss_function == "MSE":
            # output should be interpreted as log-likelihoods
            logoutput = F.softmax(output, dim=1)
            loss = loss_function(logoutput, onehot_target, reduction="mean")
        elif args.loss_function == "PMSE":
            # Use plain MSE-loss
            loss = loss_function(output, onehot_target - 1.0 / args.classes, reduction="mean")
        else:
            raise ValueError(f"Loss function {args.loss_function} not found!")
        loss.backward()
        loss = loss.item()
        tend = time()
        taken_time = tend - tstart
    else:
        # Perform a single forward pass to measure things at the start
        data = data.to(device)
        model(data)
        if "distance" in args.no_measure:
            print("Skip input caching")
        else:
            model.cache_initial_inputs()
    yield from model.measure_input_matrix()
    # Input conditioning requires untouched weights
    yield from model.measure_input_conditioning()
    # --- Important: Evaluation of the model must be performed after other
    #                measures because of caching in forward pass
    # We now change the weights in the model to allow for correct test loss testing
    if target is not None:
        optimizer.step()
    # And then we measure things that require updated weights
    yield from model.measure_norms()
    yield from model.measure_distances()
    # Measure amount of correct classification and test loss
    model.eval()
    test_loss = 0.0
    correct = 0.0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            # Convert target to one-hot-encoded tensor
            onehot_target = one_hot_encode(target, args.classes)
            output = model(data)
            if args.loss_function == "CE":
                # Here the logsoftmax is inside the loss function already
                test_loss += loss_function(output, target, reduction="sum").item()
            elif args.loss_function == "KL":
                # output should be interpreted as log-likelihoods
                logoutput = F.log_softmax(output, dim=1)
                test_loss += loss_function(logoutput, onehot_target, reduction="mean").item()
            elif args.loss_function == "MSE":
                # output should be interpreted as log-likelihoods
                logoutput = F.softmax(output, dim=1)
                test_loss += data.size()[0] * loss_function(logoutput, onehot_target, reduction="mean").item()
            elif args.loss_function == "PMSE":
                test_loss += data.size()[0] * loss_function(output, onehot_target - 1.0 / args.classes, reduction="mean").item()
            else:
                raise ValueError(f"Unknown loss function {args.loss_function}")
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    correct /= len(test_loader.dataset)
    yield dict(train_loss=loss, test_loss=test_loss, correct=correct, taken_time=taken_time)


def main():
    parser = ArgumentParser(description="Train and evaluate a given neural network.")
    parser.add_argument("--dataset", type=str, default="CIFAR10", help="Name of the dataset to train on")
    parser.add_argument("--n_layers", type=int, help="Number of layers (excluding last) to create")
    parser.add_argument("--width", type=int, help="Width of each layer to create")
    parser.add_argument("--examples", type=int, help="Number of examples to train on")
    parser.add_argument("--batch_size", type=int, default=None, help="Batch size for training")
    parser.add_argument("--epochs", type=int, help="Number of epochs to train")
    parser.add_argument("--lr", type=float, help="Learning rate")
    parser.add_argument("--network", type=str, default="FCN", help="Type of network to use (FCN vs CNN)")
    parser.add_argument("--kernel_size", type=int, default=3, help="Size of convolutional kernels used")
    parser.add_argument("--renormalize", type=str, default="False", help="Whether to renormalize after activations")
    parser.add_argument("--classes", type=int, default=10, help="Number of classes")
    parser.add_argument("--activation", type=str, default="ReLU", help="Activation function to use")
    parser.add_argument("--loss_function", type=str, help="Loss function to use")
    parser.add_argument("--seed", type=int, default=0, help="Random seed for reproducibility")
    parser.add_argument("--data_folder", type=str, default="./data", help="Folder where datasets are downloaded")
    parser.add_argument("--output_folder", type=str, default="./results", help="Folder where to store results of the experiment")
    parser.add_argument("--disable-gpu-check", default=False, action="store_true", help="Whether to disable to single GPU usage check.")
    parser.add_argument("--no_measure", type=str, nargs="*", default=[], help="Measures to not perform.")
    args = parser.parse_args()
    args.renormalize = True if args.renormalize in ["True", "true", "ok", "yes"] else False
    
    if args.batch_size is None:
        print(f"No batch size passed, will use the number of examples ({args.examples})")
        args.batch_size = args.examples
    run_slug = experiment_slug(args)
    # Manual seeding for reproducibility!
    torch.manual_seed(args.seed)
    
    use_cuda = torch.cuda.is_available()
    print(f"Will use CUDA: {use_cuda}")
    # Check we are seeing (and thus using) only a single device
    print(f"Seeing {torch.cuda.device_count()} GPUs")
    assert torch.cuda.device_count() <= 1, "Code is seeing more than one GPU, do not waste resources!"
    print(f"Using CUDA Device", torch.cuda.current_device(), torch.cuda.get_device_name())
    device = torch.device("cuda" if use_cuda else "cpu")
    
    transform = transforms.Compose([
        transforms.ToTensor(),
    ])
    # Dynamically extract the dataset from torchvision
    ds = getattr(datasets, args.dataset)
    train_ds = ds(args.data_folder, train=True, download=True, transform=transform)
    # Generate a random stratified split of the dataset
    # 1. Store the target indices for all the records
    targets = torch.tensor([target for _, target in train_ds], dtype=torch.int)
    targets_num = 1 + targets.max().item()
    print(f"Number of dataset targets: {targets_num}")
    assert args.examples % targets_num == 0, f"# of examples {args.examples} is not divisible by number of targets {targets_num}"
    # 2. For each target generate a random permutation of its indices in the dataset
    full_idxs = None
    for tgt in range(targets_num):
        tgt_idx = torch.nonzero(targets == tgt, as_tuple=True)[0]
        # Select a random permutation of the indices
        perm_idx = torch.randperm(len(tgt_idx), generator=torch.Generator())
        # Select only some indices in targets
        perm_tgt_idx = tgt_idx[perm_idx[torch.arange(args.examples // targets_num)]]
        if full_idxs is None:
            full_idxs = perm_tgt_idx
        else:
            full_idxs = torch.cat((full_idxs, perm_tgt_idx))
    # 3. Subset only those indices
    train_ds = torch.utils.data.Subset(train_ds, full_idxs)
    # # Generate a random stratified subset of the dataset using sklearn
    # # train_ds_idx, _ = train_test_split(np.arange(len(train_ds)), test_size=args.examples, shuffle=True, stratify=train_ds.targets)
    # permuted_indices = torch.randperm(len(train_ds), generator=torch.Generator())
    # example_indices = permuted_indices[torch.arange(args.examples)]
    # train_ds = torch.utils.data.Subset(train_ds, example_indices)
    #     train_ds = torch.utils.data.Subset(train_ds, train_ds_idx)
    test_ds = ds(args.data_folder, train=False, download=True, transform=transform)
    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=args.batch_size, shuffle=False, num_workers=4)
    test_loader = torch.utils.data.DataLoader(test_ds, batch_size=args.batch_size, num_workers=4)

    ds_features = 0
    image_size = (0, 0)
    ds_classes = args.classes
    for data, target in train_loader:
        _, channels, sizex, sizey = data.size()
        image_size = (sizex, sizey)
        features = dict(FCN=sizex * sizey * channels, CNN=channels)
        ds_features = features[args.network]
        print(f"Dataset features {ds_features} to {ds_classes} outputs")
        break

    model = Net(
        width=args.width,
        n_layers=args.n_layers,
        ds_features=ds_features,
        ds_classes=ds_classes,
        activation=args.activation,
        args=args,
        image_size=image_size,
    ).to(device)
    model.init_weights_()
    optimizer = optim.SGD(model.parameters(), lr=args.lr)

    measures = {}
    def add_to_measures(data):
        "Adds the dict to the right kind of measure"
        dtype = data["type"] if "type" in data else "general"
        if dtype not in measures:
            measures[dtype] = []
        if "type" in data:
            del data["type"]
        measures[dtype].append(data)

    def find_measures(dtype, **kwargs):
        found = []
        for element in measures[dtype]:
            for name, attr in kwargs.items():
                if element[name] != attr:
                    break
            else:
                found.append(element)
        return found
        
    def save_measures():
        "Saves in CSV the created measures"
        makedirs(args.output_folder, exist_ok=True)
        for dtype, meas in measures.items():
            df = pd.DataFrame.from_records(meas)
            for argument in ["dataset", "examples", "network", "loss_function", "n_layers", "width",
                             "batch_size", "lr", "activation", "epochs", "kernel_size", "seed", "renormalize"]:
                if args.network == "FCN" and argument == "kernel_size":
                    df[argument] = None
                    continue
                df[argument] = getattr(args, argument)
            df.to_csv(f"{args.output_folder}/{dtype}@{run_slug}.csv", index=False)
        
    general_stats = dict(
        dataset=args.dataset,
        n_layers=args.n_layers,
        width=args.width,
        examples=args.examples,
        batch_size=args.batch_size,
        epochs=args.epochs,
        lr=args.lr,
        seed=args.seed,
    )
    for epoch in range(args.epochs + 1):
        # Measure the model before training
        for batch_idx, (data, target) in enumerate(train_loader):
            stats = dict(epoch=epoch, batch=batch_idx)
            stats.update(general_stats)
            print(f"Stats: {stats}")
            curtarget = target if epoch > 0 else None
            for pmeas in trainstep(model, device, optimizer, test_loader, data=data, target=curtarget, args=args):
                newstats = deepcopy(stats)
                newstats.update(pmeas)
                add_to_measures(newstats)
                print(f"Measure: {pmeas}")
    save_measures()


if __name__ == "__main__":
    main()
