#Internal imports
import dataloader
import model
import test
import my_train
import create_submission
import utils

#External imports
import yaml
import torch
import wandb
import logging
import torch.optim
import torch.nn as nn
import os
import argparse

def optimizer(cfg, network):
    result = {"Adam" : torch.optim.Adam(network.parameters())}
    return result[cfg["Optimizer"]]

def train(args, cfg):
    rootDir = args.rootDir if args.rootDir != None else cfg["LogDir"]

    logging.basicConfig(filename= rootDir + 'main_unit_test.log', level=logging.INFO)
    

    use_cuda = torch.cuda.is_available()
    trainpath           = args.PATHTOTRAININGSET if args.PATHTOTRAININGSET != None else cfg["Dataset"]["_DEFAULT_TRAIN_FILEPATH"]
    num_days            = int(cfg["Dataset"]["num_days"])
    batch_size          = int(cfg["Dataset"]["batch_size"])
    num_workers         = int(cfg["Dataset"]["num_workers"])
    valid_ratio         = float(cfg["Dataset"]["valid_ratio"])
    max_num_samples     = eval(cfg["Dataset"]["max_num_samples"])

    epochs          = int(cfg["Training"]["Epochs"])
    log_freq        = int(cfg["Wandb"]["log_freq"])
    log_interval    = int(cfg["Wandb"]["log_interval"])

    dataset_transform = cfg["Dataset"]["Transform"]

    if not args.no_wandb:
        wandb.init(entity = "wherephytoplankton", project = "Kaggle phytoplancton", config = {"batch_size": batch_size, "epochs": epochs})

    # Re-compute the statistics or use the stored ones
    approx_stats = cfg["ApproximativeStats"]


    if approx_stats:
        MEAN    = eval(cfg["ApproximativeMean"])
        STD     = eval(cfg["ApproximativeSTD"])
        MAX    = eval(cfg["ApproximativeMaxi"])
        MIN    = eval(cfg["ApproximativeMini"])
    else :
        MEAN, STD, MAX, MIN = dataloader.get_stats_train_dataset(trainpath,
        num_days,
        batch_size,
        num_workers,
        use_cuda,
        valid_ratio,
        overwrite_index=True,
        max_num_samples=max_num_samples,
        train_transform=None,
        valid_transform=None
    )

    train_loader, valid_loader = dataloader.get_dataloaders(
        trainpath,
        num_days,
        batch_size,
        num_workers,
        use_cuda,
        valid_ratio,
        overwrite_index =  True,
        max_num_samples=max_num_samples,
        train_transform= eval(dataset_transform),
        valid_transform=eval(dataset_transform)
    )

    if use_cuda :
        device = torch.device('cuda')
    else :
        device = toch.device('cpu')

    network = model.build_model(cfg, 14)

    network = network.to(device)

    model.initialize_model(cfg, network)

    f_loss = model.RMSLELoss()

    optimizer = optimizer(cfg, network)

    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, 
        'min', 
        patience = 5, 
        threshold = 0.2,
        factor = 0.5
    )

    logdir, raw_run_name = utils.create_unique_logpath(rootDir, cfg["Model"]["Name"])
    network_checkpoint = model.ModelCheckpoint(logdir + "/best_model.pt", network)

    if not args.no_wandb:
        wandb.run.name = raw_run_name
        wandb.watch(network, log_freq = log_freq)

    if args.detect_anomaly:
        torch.autograd.set_detect_anomaly(True)

    best_val_loss = None
    for t in range(cfg["Training"]["Epochs"]):
        print("Epoch {}".format(t))
        my_train.train(args, network, train_loader, f_loss, optimizer, device, log_interval)

        val_loss = test.test(network, valid_loader, f_loss, device)

        if best_val_loss != None:
            if val_loss < best_val_loss :
                network_checkpoint.update(val_loss)

        scheduler.step(val_loss)


        print(" Validation : Loss : {:.4f}".format(val_loss))
        if not args.no_wandb:
            wandb.log({"val_loss": val_loss})

    utils.write_summary(logdir, network, optimizer, val_loss)

    logging.info(f"Best model saved in folder {logdir}")


def test(args):

    dataset_transform = cfg["Dataset"]["Transform"]
    rootDir = args.rootDir if args.rootDir != None else cfg["LogDir"]

    if use_cuda :
        device = torch.device('cuda')
    else :
        device = toch.device('cpu')

    logdir, raw_run_name = utils.create_unique_logpath(rootDir, cfg["Model"]["Name"])

    model_path = args.PATHTOCHECKPOINT

    network = model.build_model(cfg, 14)

    network = model.to(device)

    network.load_state_dict(torch.load(model_path))

    create_submission.create_submission(args, network, eval(dataset_transform), device, rootDir, logdir)

    logging.info(f"The submission csv file has been created in the folder : {logdir}")

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument(
    "--no_wandb",
    action="store_true",
    help="If specified, no log will be sent to wandb. Especially useful when running batch jobs.",
    )

    parser.add_argument(
    "--detect_anomaly",
    action="store_true",
    help="If specified, torch.autograd.set_detect_anomaly(True) will be activated",
    )

    parser.add_argument(
    "--rootDir",
    default=None,
    help="Directory in which the log files will be stored"
    )

    parser.add_argument(
        "--PATHTOTESTSET",
    default=None,
    help="Path of the file on which the model will be tested on"
    )

    parser.add_argument(
        "--PATHTOTRAININGSET",
        default=None,
        help="Path of the file on which the model will be trained on"
    )

    parser.add_argument(
        "--PATHTOCHECKPOINT",
        default="./logs/BestBidirectionalLSTM/best_model.pt",
        help="Path of the model to load"
    )

    parser.add_argument(
        "command", 
        choices=["train", "test"]
    )

    args = parser.parse_args()
    
    config_file = open("config.yml")
    cfg = yaml.load(config_file, Loader=yaml.FullLoader)
    
    eval(f"{args.command}(args)")