"""Train the model"""

import argparse

import logging

import os, shutil

import numpy as np

import pandas as pd

from sklearn.utils.class_weight import compute_class_weight

import torch

import torch.optim as optim

import torchvision.models as models

from torch.autograd import Variable

from torch.utils.tensorboard import SummaryWriter

from tqdm import tqdm

# from torchsummary import summary

import utils

import json

import model.net as net

import model.data_loader as data_loader

from evaluate import evaluate

parser = argparse.ArgumentParser()

parser.add_argument('--data-dir', default='data', help="Directory containing the dataset")

parser.add_argument('--model-dir', default='experiments', help="Directory containing params.json")

parser.add_argument('--setting-dir', default='settings', help="Directory with different settings")

parser.add_argument('--setting', default='collider-prognosticfactor', help="Directory contain setting.json, experimental setting, data-generation, regression model etc")

parser.add_argument('--fase', default='xybn', help='fase of training model, see manuscript for details. x, y, xy, bn, or feature')

parser.add_argument('--experiment', default='', help="Manual name for experiment for logging, will be subdir of setting")

parser.add_argument('--restore-file', default=None,

                    help="Optional, name of the file in --model_dir containing weights to reload before \

                    training")  # 'best' or 'train'

parser.add_argument('--restore-last', action='store_true', help="continue a last run")

parser.add_argument('--restore-warm', action='store_true', help="continue on the run called 'warm-start.pth'")

parser.add_argument('--use-last', action="store_true", help="use last state dict instead of 'best' (use for early stopping manually)")

parser.add_argument('--cold-start', action='store_true', help="ignore previous state dicts (weights), even if they exist")

parser.add_argument('--warm-start', dest='cold_start', action='store_false', help="start from previous state dict")

parser.add_argument('--disable-cuda', action='store_true', help="Disable Cuda")

parser.add_argument('--no-parallel', action="store_false", help="no multiple GPU", dest="parallel")

parser.add_argument('--parallel', action="store_true", help="multiple GPU", dest="parallel")

parser.add_argument('--gpu', default=0, type=int, help='if not running in parallel (=all gpus), only use this gpu')

parser.add_argument('--intercept', action="store_true", help="dummy run for getting intercept baseline results")

# parser.add_argument('--visdom', action='store_true', help='generate plots with visdom')

# parser.add_argument('--novisdom', dest='visdom', action='store_false', help='dont plot with visdom')

parser.add_argument('--monitor-grads', action='store_true', help='keep track of mean norm of gradients')

parser.set_defaults(parallel=False, cold_start=True, use_last=False, intercept=False, restore_last=False, save_preds=False,

                    monitor_grads=False, restore_warm=False

                    # visdom=False

                    )

def train(model, optimizer, loss_fn, dataloader, metrics, params, setting, writer=None, epoch=None):

    """Train the model on `num_steps` batches

    Args:

        model: (torch.nn.Module) the neural network

        optimizer: (torch.optim) optimizer for parameters of model

        loss_fn: a function that takes batch_output and batch_labels and computes the loss for the batch

        dataloader: (DataLoader) a torch.utils.data.DataLoader object that fetches training data

        metrics: (dict) a dictionary of functions that compute a metric using the output and labels of each batch

        params: (Params) hyperparameters

        num_steps: (int) number of batches to train on, each of size params.batch_size

    """

    global train_tensor_keys, logdir

    # set model to training mode

    model.train()

    # summary for current training loop and a running average object for loss

    summ = []

    loss_avg = utils.RunningAverage()

    # create storate for tensors for OLS after minibatches

    ts = []

    Xs = []

    Xtrues = []

    Ys = []

    Xhats = []

    Yhats = []

    Zhats = []

    # Use tqdm for progress bar

    with tqdm(total=len(dataloader)) as progress_bar:

        for i, batch in enumerate(dataloader):

            summary_batch = {}

            # put batch on cuda

            batch = {k: v.to(params.device) for k, v in batch.items()}

            if not (setting.covar_mode and epoch > params.suppress_t_epochs):

                batch["t"] = torch.zeros_like(batch['t'])

            Xs.append(batch['x'].detach().cpu())

            Xtrues.append(batch['x_true'].detach().cpu())

            # compute model output and loss

            output_batch = model(batch['image'], batch['t'].view(-1,1), epoch)

            Yhats.append(output_batch['y'].detach().cpu())

            # calculate loss

            if args.fase == "feature":

                # calculate loss for z directly, to get clear how well this can be measured

                loss_fn_z = torch.nn.MSELoss()

                loss_z = loss_fn_z(output_batch["y"].squeeze(), batch["z"])

                loss   = loss_z

                summary_batch["loss_z"] = loss_z.item()

            else:

                loss_fn_y = torch.nn.MSELoss()

                loss_y = loss_fn_y(output_batch["y"].squeeze(), batch["y"])

                loss   = loss_y

                summary_batch["loss_y"] = loss_y.item()

            # calculate loss for colllider x

            loss_fn_x = torch.nn.MSELoss()

            loss_x = loss_fn_x(output_batch["bnx"].squeeze(), batch["x"])

            summary_batch["loss_x"] = loss_x.item()

            if not params.alpha == 1:

                # possibly weigh down contribution of estimating x

                loss_x *= params.alpha

                summary_batch["loss_x_weighted"] = loss_x.item()

            # add x loss to total loss

            loss += loss_x            

            # add least squares regression on final layer

            if params.do_least_squares:

                X    = batch["x"].view(-1,1)

                t    = batch["t"].view(-1,1)

                Z    = output_batch["bnz"]

                if Z.ndimension() == 1:

                    Z.unsqueeze_(1)

                Xhat = output_batch["bnx"]

                # add intercept

                Zi = torch.cat([torch.ones_like(t), Z], 1)

                # add treatment info

                Zt = torch.cat([Zi, t], 1)

                Y  = batch["y"].view(-1,1)

                # regress y on final layer, without x

                betas_y = net.cholesky_least_squares(Zt, Y, intercept=False)

                y_hat   = Zt.matmul(betas_y).view(-1,1)

                mse_y  = ((Y - y_hat)**2).mean()

                summary_batch["regr_b_t"] = betas_y[-1].item()

                summary_batch["regr_loss_y"] = mse_y.item()

                # regress x on final layer without x

                betas_x = net.cholesky_least_squares(Zi, Xhat, intercept=False)

                x_hat   = Zi.matmul(betas_x).view(-1,1)

                mse_x  = ((Xhat - x_hat)**2).mean()

                # store all tensors for single pass after epoch

                Xhats.append(Xhat.detach().cpu())

                Zhats.append(Z.detach().cpu())

                ts.append(t.detach().cpu())

                Ys.append(Y.detach().cpu())

                summary_batch["regr_loss_x"] = mse_x.item()

            # add loss_bn only after n epochs

            if params.bottleneck_loss and epoch > params.bn_loss_lag_epochs:

                # only add to loss when bigger than margin

                if params.bn_loss_margin_type == "dynamic-mean":

                    # for each batch, calculate loss of just using mean for predicting x

                    mse_x_mean = ((X - X.mean())**2).mean()

                    loss_bn = torch.max(torch.zeros_like(mse_x), mse_x_mean - mse_x)

                elif params.bn_loss_margin_type == "fixed":

                    mse_diff = params.bn_loss_margin - mse_x

                    loss_bn = torch.max(torch.zeros_like(mse_x), mse_diff)

                else:

                    raise NotImplementedError(f'bottleneck loss margin type not implemented: {params.bn_loss_margin_type}')

                # possibly reweigh bottleneck loss and add to total loss

                summary_batch["loss_bn"] = loss_bn.item()

                # note is this double?

                if loss_bn > params.bn_loss_margin:

                    loss_bn *= params.bottleneck_loss_wt

                    loss    += loss_bn

            # perform parameter update

            optimizer.zero_grad()

            loss.backward()

            optimizer.step()

            summary_batch['loss'] = loss.item()

            summ.append(summary_batch)

            # if necessary, write out tensors

            if params.monitor_train_tensors and (epoch % params.save_summary_steps == 0):

                tensors = {}

                for tensor_key in train_tensor_keys:

                    if tensor_key in batch.keys():

                        tensors[tensor_key] = batch[tensor_key].squeeze().numpy()

                    elif tensor_key.endswith("hat"):

                        tensor_key = tensor_key.split("_")[0]

                        if tensor_key in output_batch.keys():

                            tensors[tensor_key+"_hat"] = output_batch[tensor_key].detach().cpu().squeeze().numpy()

                    else:

                        assert False, f"key not found: {tensor_key}"

                # print(tensors)

                df = pd.DataFrame.from_dict(tensors, orient='columns')

                df["epoch"] = epoch

                with open(os.path.join(logdir, 'train-tensors.csv'), 'a') as f:

                    df[["epoch"]+train_tensor_keys].to_csv(f, header=False)

            # update the average loss

            loss_avg.update(loss.item())

            progress_bar.set_postfix(loss='{:05.3f}'.format(loss_avg()))

            progress_bar.update()

    # visualize gradients

    if epoch % params.save_summary_steps == 0 and args.monitor_grads:

        abs_gradients = {}

        for name, param in model.named_parameters():

            try: # patch here, there were names / params that were 'none'

                abs_gradients[name] = np.abs(param.grad.cpu().numpy()).mean()

                writer.add_histogram("grad-"+name, param.grad, epoch)

                writer.add_scalars("mean-abs-gradients", abs_gradients, epoch)

            except:

                pass

    # compute mean of all metrics in summary

    metrics_mean = {metric:np.nanmean([x[metric] for x in summ]) for metric in summ[0]}

    # collect tensors

    Xhat = torch.cat(Xhats,0).view(-1,1)

    Yhat = torch.cat(Yhats,0).view(-1,1)

    Zhat = torch.cat(Zhats,0)

    t    = torch.cat(ts,0)

    X    = torch.cat(Xs,0)

    Xtrue= torch.cat(Xtrues,0)

    Y    = torch.cat(Ys,0)

    if params.do_least_squares:

        # after the minibatches, do a single OLS on the whole data

        Zi = torch.cat([torch.ones_like(t), Zhat], 1)

        # add treatment info

        Zt = torch.cat([Zi, t], 1)

        # add x for biased version

        XZt = torch.cat([torch.ones_like(t), Xhat, Zhat, t], 1)

        betas_y_bias       = net.cholesky_least_squares(XZt, Y, intercept=False)

        betas_y_causal     = net.cholesky_least_squares(Zt, Y, intercept=False)

        model.betas_bias   = betas_y_bias

        model.betas_causal = betas_y_causal

        metrics_mean["regr_bias_coef_t"]   = betas_y_bias.squeeze()[-1]

        metrics_mean["regr_bias_coef_z"]   = betas_y_bias.squeeze()[-2]

        metrics_mean["regr_causal_coef_t"] = betas_y_causal.squeeze()[-1]

        metrics_mean["regr_causal_coef_z"] = betas_y_causal.squeeze()[-2]

    # create some plots

    xx_scatter    = net.make_scatter_plot(X.numpy(), Xhat.numpy(), xlabel='x', ylabel='xhat') 

    xtruex_scatter= net.make_scatter_plot(Xtrue.numpy(), Xhat.numpy(), xlabel='xtrue', ylabel='xhat') 

    xyhat_scatter = net.make_scatter_plot(X.numpy(), Yhat.numpy(), c=t.numpy(), xlabel='x', ylabel='yhat')

    yy_scatter    = net.make_scatter_plot(Y.numpy(), Yhat.numpy(), c=t.numpy(), xlabel='y', ylabel='yhat') 

    writer.add_figure('x-xhat/train', xx_scatter, epoch+1)

    writer.add_figure('xtrue-xhat/train', xtruex_scatter, epoch+1)

    writer.add_figure('x-yhat/train', xyhat_scatter, epoch+1)

    writer.add_figure('y-yhat/train', yy_scatter, epoch+1)

    metrics_string = " ; ".join("{}: {:05.3f}".format(k, v) for k, v in metrics_mean.items())

    logging.info("- Train metrics: " + metrics_string)

    return metrics_mean

def train_and_evaluate(model, train_dataloader, val_dataloader, optimizer, loss_fn, metrics, params, setting, args,

                       writer=None, logdir=None, restore_file=None):

    """Train the model and evaluate every epoch.

    Args:

        model: (torch.nn.Module) the neural network

        train_dataloader: (DataLoader) a torch.utils.data.DataLoader object that fetches training data

        val_dataloader: (DataLoader) a torch.utils.data.DataLoader object that fetches validation data

        optimizer: (torch.optim) optimizer for parameters of model

        loss_fn: a function that takes batch_output and batch_labels and computes the loss for the batch

        metrics: (dict) a dictionary of functions that compute a metric using mnisthe output and labels of each batch

        params: (Params) hyperparameters

        model_dir: (string) directory containing config, weights and log

        restore_file: (string) optional- name of file to restore from (withoutmnistits extension .pth.tar)

        covar_mode: (bool) does the data-loader give back covariates / additional data

    """

    # setup directories for data

    setting_home = setting.home

    if not args.fase == "feature":

        data_dir = os.path.join(setting_home, "data")

    else:

        if setting.mode3d:

            data_dir = "data"

        else:

            data_dir = "slices"

    covar_mode = setting.covar_mode

    x_frozen = False

    best_val_metric = 0.0

    if "loss" in setting.metrics[0]:

        best_val_metric = 1.0e6

    val_preds = np.zeros((len(val_dataloader.dataset), params.num_epochs))

    for epoch in range(params.num_epochs):

        # Run one epoch

        logging.info(f"Epoch {epoch+1}/{params.num_epochs}; setting: {args.setting}, fase {args.fase}, experiment: {args.experiment}")

        # compute number of batches in one epoch (one full pass over the training set)

        train_metrics = train(model, optimizer, loss_fn, train_dataloader, metrics, params, setting, writer, epoch)

        print(train_metrics)

        for metric_name in train_metrics.keys():

            metric_vals = {'train': train_metrics[metric_name]}

            writer.add_scalars(metric_name, metric_vals, epoch+1)

        # for name, param in model.named_parameters():

        #     writer.add_histogram(name, param.clone().cpu().data.numpy(), epoch+1)

        if epoch % params.save_summary_steps == 0:

            # Evaluate for one epoch on validation set

            valid_metrics, outtensors = evaluate(model, loss_fn, val_dataloader, metrics, params, setting, epoch, writer) 

            valid_metrics["intercept"] = model.regressor.fc.bias.detach().cpu().numpy()

            print(valid_metrics) 

            for name, module in model.regressor.named_children():

                if name == "t":

                    valid_metrics["b_t"] = module.weight.detach().cpu().numpy()

                elif name == "zt":

                    weights = module.weight.detach().cpu().squeeze().numpy().reshape(-1)

                    for i, weight in enumerate(weights):

                        valid_metrics["b_zt"+str(i)] = weight

                else:

                    pass

            for metric_name in valid_metrics.keys():

                metric_vals = {'valid': valid_metrics[metric_name]}

                writer.add_scalars(metric_name, metric_vals, epoch+1)

            # create plots

            val_df = val_dataloader.dataset.df

            xx_scatter    = net.make_scatter_plot(val_df.x.values, outtensors['xhat'], xlabel='x', ylabel='xhat') 

            xtruex_scatter= net.make_scatter_plot(val_df.x_true.values, outtensors['xhat'], xlabel='x', ylabel='xhat') 

            xyhat_scatter = net.make_scatter_plot(val_df.x.values, outtensors['predictions'], c=val_df.t, xlabel='x', ylabel='yhat')

            zyhat_scatter = net.make_scatter_plot(val_df.z.values, outtensors['predictions'], c=val_df.t, xlabel='z', ylabel='yhat')

            yy_scatter    = net.make_scatter_plot(val_df.y.values, outtensors['predictions'], c=val_df.t, xlabel='yhat', ylabel='y') 

            writer.add_figure('x-xhat/valid', xx_scatter, epoch+1)

            writer.add_figure('xtrue-xhat/valid', xtruex_scatter, epoch+1)

            writer.add_figure('x-yhat/valid', xyhat_scatter, epoch+1)

            writer.add_figure('z-yhat/valid', zyhat_scatter, epoch+1)

            writer.add_figure('y-yhat/valid', yy_scatter, epoch+1)

            if params.save_preds:

                # writer.add_histogram("predictions", preds)

                if setting.num_classes == 1:

                    val_preds[:, epoch] = np.squeeze(outtensors['predictions'])

                    # write preds to file

                    pred_fname = os.path.join(setting.home, setting.fase+"-fase", "preds_val.csv")

                    with open(pred_fname, 'ab') as f:

                        np.savetxt(f, preds.T, newline="")

                np.save(os.path.join(setting.home, setting.fase+"-fase", "preds.npy"), preds)

            else:

                val_metric = valid_metrics[setting.metrics[0]]

            if "loss" in str(setting.metrics[0]):

                is_best = val_metric<=best_val_metric

            else:

                is_best = val_metric>=best_val_metric

            # Save weights

            state_dict = model.state_dict()

            optim_dict = optimizer.state_dict()

            state = {

                'epoch': epoch+1,

                'state_dict': state_dict,

                'optim_dict': optim_dict

            }

            utils.save_checkpoint(state,

                                is_best=is_best,

                                checkpoint=logdir)

            # If best_eval, best_save_path

            valid_metrics["epoch"] = epoch

            if is_best:

                logging.info("- Found new best {}: {:.3f}".format(setting.metrics[0], val_metric))

                best_val_metric = val_metric

                # Save best val metrics in a json file in the model directory

                best_json_path = os.path.join(logdir, "metrics_val_best_weights.json")

                utils.save_dict_to_json(valid_metrics, best_json_path)

            # Save latest val metrics in a json file in the model directory

            last_json_path = os.path.join(logdir, "metrics_val_last_weights.json")

            utils.save_dict_to_json(valid_metrics, last_json_path)

    # final evaluation

    writer.export_scalars_to_json(os.path.join(logdir, "all_scalars.json"))

    if args.save_preds:

        np.save(os.path.join(setting.home, setting.fase + "-fase", "val_preds.npy"), val_preds)

if __name__ == '__main__':

    # Load the parameters from json file

    args = parser.parse_args()

    # Load information from last setting if none provided:

    last_defaults = utils.Params("last-defaults.json")

    if args.setting == "":

        print("using last default setting")

        args.setting = last_defaults.dict["setting"]

        for param, value in last_defaults.dict.items():

            print("{}: {}".format(param, value))

    else:

        with open("last-defaults.json", "r+") as jsonFile:

            defaults = json.load(jsonFile)

            tmp = defaults["setting"]

            defaults["setting"] = args.setting

            jsonFile.seek(0)  # rewind

            json.dump(defaults, jsonFile)

            jsonFile.truncate()

    # setup visdom environment

    # if args.visdom:

        # from visdom import Visdom

        # viz = Visdom(env=f"lidcr_{args.setting}_{args.fase}_{args.experiment}")

    # load setting (data generation, regression model etc)

    setting_home = os.path.join(args.setting_dir, args.setting)

    setting = utils.Params(os.path.join(setting_home, "setting.json"))

    setting.home = setting_home

    # when not specified in call, grab model specification from setting file

    if setting.cnn_model == "":

        json_path = os.path.join(args.model_dir, "t-suppression", args.experiment+".json")

    else:

        json_path = os.path.join(args.model_dir, setting.cnn_model, 'params.json')

    assert os.path.isfile(json_path), "No json configuration file found at {}".format(json_path)

    if not os.path.exists(os.path.join(setting.home, args.fase + "-fase")):

        os.makedirs(os.path.join(setting.home, args.fase + "-fase"))

    shutil.copy(json_path, os.path.join(setting_home, args.fase + "-fase", "params.json"))

    params = utils.Params(json_path)

    # covar_mode = setting.covar_mode

    # mode3d = setting.mode3d

    parallel = args.parallel

    params.device = None

    if not args.disable_cuda and torch.cuda.is_available():

        params.device = torch.device('cuda')

        params.cuda = True

        # switch gpus for better use when running multiple experiments

        if not args.parallel:

            torch.cuda.set_device(int(args.gpu))

    else:

        params.device = torch.device('cpu')

    # adapt fase

    setting.fase = args.fase

    setting.metrics = pd.Series(setting.metrics).drop_duplicates().tolist()

    print("metrics {}:".format(setting.metrics))

    # Set the random seed for reproducible experiments

    torch.manual_seed(230)

    if params.cuda: torch.cuda.manual_seed(230)

    # Set the logger

    logdir=os.path.join(setting_home, setting.fase+"-fase", "runs")

    if not args.experiment == '':

        logdir=os.path.join(logdir, args.experiment)

    if not os.path.isdir(logdir):

        os.makedirs(logdir)

    # copy params as backupt to logdir

    shutil.copy(json_path, os.path.join(logdir, "params.json"))

    # utils.set_logger(os.path.join(args.model_dir, 'train.log'))

    utils.set_logger(os.path.join(logdir, 'train.log'))

    # Create the input data pipeline

    logging.info("Loading the datasets...")

    # fetch dataloaders

    dataloaders = data_loader.fetch_dataloader(args, params, setting, ["train", "valid"])

    train_dl = dataloaders['train']

    valid_dl = dataloaders['valid']

    if setting.num_classes > 1 and params.balance_classes:

        train_labels = train_dl.dataset.df[setting.outcome[0]].values

        class_weights = compute_class_weight('balanced', np.unique(train_labels), train_labels)

    # valid_dl = train_dl

    logging.info("- done.")

    if args.intercept:

        assert len(setting.outcome) == 1, "Multiple outcomes not implemented for intercept yet"

        print("running intercept mode")

        mu = valid_dl.dataset.df[setting.outcome].values.mean()

        def new_forward(self, x, data, mu=mu):

            intercept = torch.autograd.Variable(mu * torch.ones((x.shape[0],1)), requires_grad=False).to(params.device, non_blocking=True)

            bn_activations = torch.autograd.Variable(torch.zeros((x.shape[0],)), requires_grad=False).to(params.device, non_blocking=True)

            return {setting.outcome[0]: intercept, "bn": bn_activations}

        net.Net3D.forward = new_forward

        params.num_epochs = 1

        setting.metrics = []

        logdir = os.path.join(logdir, "intercept")

    if setting.mode3d:

        model = net.Net3D(params, setting).to(params.device)

    else:

        model = net.CausalNet(params, setting).to(params.device)

    optimizers = {'sgd': optim.SGD, 'adam': optim.Adam}

    if parallel:

        print("parallel mode")

        model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))

    if params.momentum > 0:

        optimizer = optimizers[params.optimizer](model.parameters(), lr=params.learning_rate, weight_decay=params.wd, momentum=params.momentum)

    else:

        optimizer = optimizers[params.optimizer](model.parameters(), lr=params.learning_rate, weight_decay=params.wd)

    # if params.use_mi:

    #     optimizer.add_param_group({'params': mine.parameters()})

    if setting.covar_mode and params.lr_t_factor != 1:

        optimizer = net.speedup_t(model, params)

    if args.restore_last and (not args.cold_start):

        print("Loading state dict from last running setting")

        utils.load_checkpoint(os.path.join(setting.home, args.fase + "-fase", "last.pth.tar"), model, strict=False)

    elif args.restore_warm:

        utils.load_checkpoint(os.path.join(setting.home, 'warm-start.pth.tar'), model, strict=False)

    else:

        pass

    # fetch loss function and metrics

    if setting.num_classes > 1 and params.balance_classes:

        loss_fn = net.get_loss_fn(setting, weights=class_weights)

    else:

        loss_fn = net.get_loss_fn(setting)

    # metrics = {metric:net.all_metrics[metric] for metric in setting.metrics}

    metrics = None

    if params.monitor_train_tensors:

        print(f"Recording all train tensors")

        import csv   

        train_tensor_keys = ['t','x', 'z', 'y', 'x_hat', 'z_hat', 'y_hat']

        with open(os.path.join(logdir, 'train-tensors.csv'), 'w') as f:

            writer = csv.writer(f)

            writer.writerow(['epoch']+train_tensor_keys)

    # Train the model

    # print(model)

    # print(summary(model, (3, 224, 224), batch_size=1))

    logging.info("Starting training for {} epoch(s)".format(params.num_epochs))

    for split, dl in dataloaders.items():

        logging.info("Number of %s samples: %s" % (split, str(len(dl.dataset))))

        # logging.info("Number of valid examples: {}".format(len(valid.dataset)))

    with SummaryWriter(logdir) as writer:

        # train(model, optimizer, loss_fn, train_dl, metrics, params)

        train_and_evaluate(model, train_dl, valid_dl, optimizer, loss_fn, metrics, params, setting, args,

                           writer, logdir, args.restore_file)