from haven import haven_chk as hc
from haven import haven_results as hr
from haven import haven_utils as hu
import torch
import torchvision
import tqdm
import pandas as pd
import pprint
import itertools
import os
import pylab as plt
import exp_configs
import time
import numpy as np

from src import models
from src import datasets
from src import utils as ut

from pprint import pformat


import argparse

from torch.utils.data import sampler
from torch.utils.data.sampler import RandomSampler
from torch.backends import cudnn
from torch.nn import functional as F
from torch.utils.data import DataLoader

cudnn.benchmark = True

import logging

def setupLogging():
    """Setup the logger for this module
    """
    
    # Create the Logger
    root_logger = logging.getLogger()
    root_logger.setLevel(logging.DEBUG)

    handler = logging.StreamHandler()
    logger_formatter = logging.Formatter('%(name)s - %(levelname)s - %(message)s')
    handler.setFormatter(logger_formatter)
    root_logger.addHandler(handler)


def trainval(exp_dict, savedir_base, datadir, reset=False, num_workers=0):
    # bookkeepting stuff
    # ==================
    pprint.pprint(exp_dict)
    exp_id = hu.hash_dict(exp_dict)
    savedir = os.path.join(savedir_base, exp_id)
    if reset:
        hc.delete_and_backup_experiment(savedir)

    os.makedirs(savedir, exist_ok=True)
    hu.save_json(os.path.join(savedir, "exp_dict.json"), exp_dict)
    print("Experiment saved in %s" % savedir)

    logger.info(f'start trainval with experiment dict {pformat(exp_dict)}')

    input('press enter')

    # set seed
    # ==================
    seed = 42
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    # Dataset
    # ==================
    # train set

    train_set = datasets.get_dataset(dataset_dict=exp_dict["dataset"],
                                     split="train",
                                     datadir=datadir,
                                     exp_dict=exp_dict,
                                     dataset_size=exp_dict['dataset_size'])
    # val set
    val_set = datasets.get_dataset(dataset_dict=exp_dict["dataset"],
                                   split="val",
                                   datadir=datadir,
                                   exp_dict=exp_dict,
                                   dataset_size=exp_dict['dataset_size'])

    # test set
    test_set = datasets.get_dataset(dataset_dict=exp_dict["dataset"],
                                   split="test",
                                   datadir=datadir,
                                   exp_dict=exp_dict,
                                   dataset_size=exp_dict['dataset_size'])

    # val_sampler = torch.utils.data.SequentialSampler(val_set)
    val_loader = DataLoader(val_set,
                            # sampler=val_sampler,
                            batch_size=1,
                            collate_fn=ut.collate_fn,
                            num_workers=num_workers)
    test_loader = DataLoader(test_set,
                            # sampler=val_sampler,
                            batch_size=1,
                            collate_fn=ut.collate_fn,
                            num_workers=num_workers)

    # Model
    # ==================
    print('get model')
    model = models.get_model(model_dict=exp_dict['model'],
                             exp_dict=exp_dict,
                             train_set=train_set).cuda()

    # model.opt = optimizers.get_optim(exp_dict['opt'], model)
    model_path = os.path.join(savedir, "model.pth")
    score_list_path = os.path.join(savedir, "score_list.pkl")

    print(model)

    if os.path.exists(score_list_path):
        # resume experiment
        model.load_state_dict(hu.torch_load(model_path))
        score_list = hu.load_pkl(score_list_path)
        s_epoch = score_list[-1]['epoch'] + 1
    else:
        # restart experiment
        score_list = []
        s_epoch = 0

    # Train & Val
    # ==================
    print("Starting experiment at epoch %d" % (s_epoch))
    model.waiting = 0
    model.val_score_best = -np.inf
    
    train_sampler = torch.utils.data.RandomSampler(
                                train_set, replacement=True, 
                                num_samples=2*len(test_set))

    train_loader = DataLoader(train_set,
                            sampler=train_sampler,
                            collate_fn=ut.collate_fn,
                            batch_size=exp_dict["batch_size"], 
                            drop_last=True, 
                            num_workers=num_workers)
    
    for e in range(s_epoch, exp_dict['max_epoch']):
        # Validate only at the start of each cycle
        score_dict = {}
        test_dict = model.val_on_loader(test_loader,
                                savedir_images=os.path.join(savedir, "images"),
                                n_images=3) 
        # Train the model
        train_dict = model.train_on_loader(train_loader)

        # Validate the model
        val_dict = model.val_on_loader(val_loader)
        score_dict["val_score"] = val_dict["val_score"]

        # Get new score_dict
        score_dict.update(train_dict)
        score_dict["epoch"] = e
        score_dict["waiting"] = model.waiting

        model.waiting += 1

        # Add to score_list and save checkpoint
        score_list += [score_dict]

        # Save Best Checkpoint
        score_df = pd.DataFrame(score_list)
        if score_dict["val_score"] >= model.val_score_best:
            test_dict = model.val_on_loader(test_loader,
                                savedir_images=os.path.join(savedir, "images"),
                                n_images=3)  
            score_dict.update(test_dict)
            hu.save_pkl(os.path.join(savedir, "score_list_best.pkl"), score_list)
            # score_df.to_csv(os.path.join(savedir, "score_best_df.csv"))
            hu.torch_save(os.path.join(savedir, "model_best.pth"),
                        model.get_state_dict())
            model.waiting = 0
            model.val_score_best = score_dict["val_score"]
            print("Saved Best: %s" % savedir)

        # Report & Save
        score_df = pd.DataFrame(score_list)
        # score_df.to_csv(os.path.join(savedir, "score_df.csv"))
        print("\n", score_df.tail(), "\n")
        hu.torch_save(model_path, model.get_state_dict())
        hu.save_pkl(score_list_path, score_list)
        print("Checkpoint Saved: %s" % savedir)

        if model.waiting > 100:
            break

    print('Experiment completed et epoch %d' % e)


if __name__ == "__main__":
    setupLogging()
    parser = argparse.ArgumentParser()
    logger = logging.getLogger(__name__)

    parser.add_argument('-e', '--exp_group_list', nargs="+")
    parser.add_argument('-sb', '--savedir_base', required=True)
    parser.add_argument('-d', '--datadir', default=None)
    parser.add_argument("-r", "--reset",  default=0, type=int)
    parser.add_argument("-ei", "--exp_id", default=None)
    parser.add_argument("-j", "--run_jobs", default=0, type=int)
    parser.add_argument("-nw", "--num_workers", type=int, default=0)

    args = parser.parse_args()

    # Collect experiments
    # ===================
    if args.exp_id is not None:
        # select one experiment
        savedir = os.path.join(args.savedir_base, args.exp_id)
        exp_dict = hu.load_json(os.path.join(savedir, "exp_dict.json"))

        exp_list = [exp_dict]

    else:
        # select exp group
        exp_list = []
        for exp_group_name in args.exp_group_list:
            exp_list += exp_configs.EXP_GROUPS[exp_group_name]

    for exp_dict in exp_list:
        # do trainval
        trainval(exp_dict=exp_dict,
                savedir_base=args.savedir_base,
                datadir=args.datadir,
                reset=args.reset,
                num_workers=args.num_workers)

[((40, 18, 40, 37), 'logging.getLogger', 'logging.getLogger', ({}, {}), '()', False, 'import logging\n'), ((43, 14, 43, 37), 'logging.StreamHandler', 'logging.StreamHandler', ({}, {}), '()', False, 'import logging\n'), ((44, 23, 44, 82), 'logging.Formatter', 'logging.Formatter', ({(44, 41, 44, 81): '"""%(name)s - %(levelname)s - %(message)s"""'}, {}), "('%(name)s - %(levelname)s - %(message)s')", False, 'import logging\n'), ((52, 4, 52, 27), 'pprint.pprint', 'pprint.pprint', ({(52, 18, 52, 26): 'exp_dict'}, {}), '(exp_dict)', False, 'import pprint\n'), ((53, 13, 53, 35), 'haven.haven_utils.hash_dict', 'hu.hash_dict', ({(53, 26, 53, 34): 'exp_dict'}, {}), '(exp_dict)', True, 'from haven import haven_utils as hu\n'), ((54, 14, 54, 48), 'os.path.join', 'os.path.join', ({(54, 27, 54, 39): 'savedir_base', (54, 41, 54, 47): 'exp_id'}, {}), '(savedir_base, exp_id)', False, 'import os\n'), ((58, 4, 58, 39), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((69, 4, 69, 24), 'numpy.random.seed', 'np.random.seed', ({(69, 19, 69, 23): 'seed'}, {}), '(seed)', True, 'import numpy as np\n'), ((70, 4, 70, 27), 'torch.manual_seed', 'torch.manual_seed', ({(70, 22, 70, 26): 'seed'}, {}), '(seed)', False, 'import torch\n'), ((71, 4, 71, 36), 'torch.cuda.manual_seed_all', 'torch.cuda.manual_seed_all', ({(71, 31, 71, 35): 'seed'}, {}), '(seed)', False, 'import torch\n'), ((77, 16, 81, 75), 'src.datasets.get_dataset', 'datasets.get_dataset', (), '', False, 'from src import datasets\n'), ((83, 14, 87, 73), 'src.datasets.get_dataset', 'datasets.get_dataset', (), '', False, 'from src import datasets\n'), ((90, 15, 94, 73), 'src.datasets.get_dataset', 'datasets.get_dataset', (), '', False, 'from src import datasets\n'), ((97, 17, 101, 52), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader\n'), ((102, 18, 106, 52), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader\n'), ((116, 17, 116, 51), 'os.path.join', 'os.path.join', ({(116, 30, 116, 37): 'savedir', (116, 39, 116, 50): '"""model.pth"""'}, {}), "(savedir, 'model.pth')", False, 'import os\n'), ((117, 22, 117, 61), 'os.path.join', 'os.path.join', ({(117, 35, 117, 42): 'savedir', (117, 44, 117, 60): '"""score_list.pkl"""'}, {}), "(savedir, 'score_list.pkl')", False, 'import os\n'), ((121, 7, 121, 38), 'os.path.exists', 'os.path.exists', ({(121, 22, 121, 37): 'score_list_path'}, {}), '(score_list_path)', False, 'import os\n'), ((141, 19, 146, 52), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader\n'), ((202, 13, 202, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((203, 13, 203, 40), 'logging.getLogger', 'logging.getLogger', ({(203, 31, 203, 39): '__name__'}, {}), '(__name__)', False, 'import logging\n'), ((56, 8, 56, 48), 'haven.haven_chk.delete_and_backup_experiment', 'hc.delete_and_backup_experiment', ({(56, 40, 56, 47): 'savedir'}, {}), '(savedir)', True, 'from haven import haven_chk as hc\n'), ((59, 17, 59, 55), 'os.path.join', 'os.path.join', ({(59, 30, 59, 37): 'savedir', (59, 39, 59, 54): '"""exp_dict.json"""'}, {}), "(savedir, 'exp_dict.json')", False, 'import os\n'), ((124, 21, 124, 49), 'haven.haven_utils.load_pkl', 'hu.load_pkl', ({(124, 33, 124, 48): 'score_list_path'}, {}), '(score_list_path)', True, 'from haven import haven_utils as hu\n'), ((172, 19, 172, 43), 'pandas.DataFrame', 'pd.DataFrame', ({(172, 32, 172, 42): 'score_list'}, {}), '(score_list)', True, 'import pandas as pd\n'), ((187, 19, 187, 43), 'pandas.DataFrame', 'pd.DataFrame', ({(187, 32, 187, 42): 'score_list'}, {}), '(score_list)', True, 'import pandas as pd\n'), ((191, 8, 191, 48), 'haven.haven_utils.save_pkl', 'hu.save_pkl', ({(191, 20, 191, 35): 'score_list_path', (191, 37, 191, 47): 'score_list'}, {}), '(score_list_path, score_list)', True, 'from haven import haven_utils as hu\n'), ((219, 18, 219, 62), 'os.path.join', 'os.path.join', ({(219, 31, 219, 48): 'args.savedir_base', (219, 50, 219, 61): 'args.exp_id'}, {}), '(args.savedir_base, args.exp_id)', False, 'import os\n'), ((111, 12, 113, 49), 'src.models.get_model', 'models.get_model', (), '', False, 'from src import models\n'), ((123, 30, 123, 55), 'haven.haven_utils.torch_load', 'hu.torch_load', ({(123, 44, 123, 54): 'model_path'}, {}), '(model_path)', True, 'from haven import haven_utils as hu\n'), ((220, 32, 220, 70), 'os.path.join', 'os.path.join', ({(220, 45, 220, 52): 'savedir', (220, 54, 220, 69): '"""exp_dict.json"""'}, {}), "(savedir, 'exp_dict.json')", False, 'import os\n'), ((62, 55, 62, 72), 'pprint.pformat', 'pformat', ({(62, 63, 62, 71): 'exp_dict'}, {}), '(exp_dict)', False, 'from pprint import pformat\n'), ((152, 47, 152, 78), 'os.path.join', 'os.path.join', ({(152, 60, 152, 67): 'savedir', (152, 69, 152, 77): '"""images"""'}, {}), "(savedir, 'images')", False, 'import os\n'), ((178, 24, 178, 68), 'os.path.join', 'os.path.join', ({(178, 37, 178, 44): 'savedir', (178, 46, 178, 67): '"""score_list_best.pkl"""'}, {}), "(savedir, 'score_list_best.pkl')", False, 'import os\n'), ((180, 26, 180, 65), 'os.path.join', 'os.path.join', ({(180, 39, 180, 46): 'savedir', (180, 48, 180, 64): '"""model_best.pth"""'}, {}), "(savedir, 'model_best.pth')", False, 'import os\n'), ((175, 47, 175, 78), 'os.path.join', 'os.path.join', ({(175, 60, 175, 67): 'savedir', (175, 69, 175, 77): '"""images"""'}, {}), "(savedir, 'images')", False, 'import os\n')]