import os
import os.path as osp
import random
import datetime
import itertools
import glob
import subprocess
import pyhocon
import glob
import re
import numpy as np
import glog as log
import json
import torch

import sys
sys.path.append('../')

from models import vdgr
from dataloader.dataloader_visdial import VisdialDataset

from dataloader.dataloader_visdial_dense import VisdialDenseDataset


def load_runner(config):
    if config['train_on_dense']:
        return vdgr.DenseRunner(config)
    else:
        return vdgr.SparseRunner(config)
    
def load_dataset(config):
    dataset_eval = None

    if config['train_on_dense']:
        dataset = VisdialDenseDataset(config)
        if config['skip_mrr_eval']:
            temp = config['num_options_dense']
            config['num_options_dense'] = config['num_options']
            dataset_eval = VisdialDenseDataset(config)
            config['num_options_dense'] = temp
        else:
            dataset_eval = VisdialDataset(config)
    else:
        dataset = VisdialDataset(config)
        if config['skip_mrr_eval']:
            dataset_eval = VisdialDenseDataset(config)

    if config['use_trainval']:
        dataset.split = 'trainval'
    else:
        dataset.split = 'train'

    if dataset_eval is not None:
        dataset_eval.split = 'val'
        
    return dataset, dataset_eval


def initialize_from_env(model, mode, eval_dir, model_type, tag=''):
    if "GPU" in os.environ:
        os.environ["CUDA_VISIBLE_DEVICES"] = os.environ['GPU']
    if mode in ['train', 'debug']:
        config = pyhocon.ConfigFactory.parse_file(f"config/{model_type}.conf")[model]
    else:
        path_config = osp.join(eval_dir, 'code', f"config/{model_type}.conf")
        config = pyhocon.ConfigFactory.parse_file(path_config)[model]
        config['log_dir'] = eval_dir
        config['model_config'] = osp.join(eval_dir, 'code/config/bert_base_6layer_6conect.json')
        if config['dp_type'] == 'apex':
            config['dp_type'] = 'ddp'

    if config['dp_type'] == 'dp':
        config['stack_gr_data'] = True

    config['model_type'] = model_type
    if "CUDA_VISIBLE_DEVICES" in os.environ:
        config['num_gpus'] = len(os.environ["CUDA_VISIBLE_DEVICES"].split(','))
        # multi-gpu setting
        if config['num_gpus'] > 1:
            os.environ['MASTER_ADDR'] = '127.0.0.1'
            os.environ['MASTER_PORT'] = '5678'

    if mode == 'debug':
        model += '_debug'

    if tag:
        model += '-' + tag
    if mode in ['train', 'debug']:
        config['log_dir'] = os.path.join(config["log_dir"], model)
        if not os.path.exists(config["log_dir"]):
            os.makedirs(config["log_dir"])
    config['visdial_output_dir'] = osp.join(config['log_dir'], config['visdial_output_dir'])

    config['timestamp'] = datetime.datetime.now().strftime('%m%d-%H%M%S')

    # add the bert config
    config['bert_config'] = json.load(open(config['model_config'], 'r'))
    if mode in ['predict', 'eval']:
        if (not config['loads_start_path']) and (not config['loads_best_ckpt']):
            config['loads_best_ckpt'] = True
            print(f'Setting loads_best_ckpt=True under predict or eval mode')
        if config['num_options_dense'] < 100:
            config['num_options_dense'] = 100
            print('Setting num_options_dense=100 under predict or eval mode')
    if config['visdial_version'] == 0.9:
        config['skip_ndcg_eval'] = True
    
    return config


def set_log_file(fname, file_only=False):
    # if fname already exists, find all log file under log dir,
    # and name the current log file with a new number
    if osp.exists(fname):
        prefix, suffix = osp.splitext(fname)
        log_files = glob.glob(prefix + '*' + suffix)
        count = 0
        for log_file in log_files:
            num = re.search(r'(\d+)', log_file)
            if num is not None:
                num = int(num.group(0))
                count = max(num, count)
        fname = fname.replace(suffix, str(count + 1) + suffix)
    # set log file
    # simple tricks for duplicating logging destination in the logging module such as:
    # logging.getLogger().addHandler(logging.FileHandler(filename))
    # does NOT work well here, because python Traceback message (not via logging module) is not sent to the file,
    # the following solution (copied from : https://stackoverflow.com/questions/616645) is a little bit
    # complicated but simulates exactly the "tee" command in linux shell, and it redirects everything
    if file_only:
        # we only output messages to file, and stdout/stderr receives nothing.
        # this feature is designed for executing the script via ssh:
        # since ssh has a windowing kind of flow control, i.e., if the controller does not read data from a
        # ssh channel and its buffer fills up, the execution machine will not be able to write anything into the
        # channel and the process will be set to sleeping (S) status until someone reads all data from the channel.
        # this is not desired since we do not want to read stdout/stderr from the controller machine.
        # so, here we use a simple solution: disable output to stdout/stderr and only output messages to log file.
        log.logger.handlers[0].stream = log.handler.stream = sys.stdout = sys.stderr = f = open(fname, 'w', buffering=1)
    else:
        # we output messages to both file and stdout/stderr
        tee = subprocess.Popen(['tee', fname], stdin=subprocess.PIPE)
        os.dup2(tee.stdin.fileno(), sys.stdout.fileno())
        os.dup2(tee.stdin.fileno(), sys.stderr.fileno())


def copy_file_to_log(log_dir):
    dirs_to_cp = ['.', 'config', 'dataloader', 'models', 'utils']
    files_to_cp = ['*.py', '*.json', '*.sh', '*.conf']
    for dir_name in dirs_to_cp:
        dir_name = osp.join(log_dir, 'code', dir_name)
        if not osp.exists(dir_name):
            os.makedirs(dir_name)
    for dir_name, file_name in itertools.product(dirs_to_cp, files_to_cp):
        filename = osp.join(dir_name, file_name)
        if len(glob.glob(filename)) > 0:
            os.system(f'cp {filename} {osp.join(log_dir, "code", dir_name)}')
    log.info(f'Files copied to {osp.join(log_dir, "code")}')


def set_random_seed(random_seed):
    torch.manual_seed(random_seed)
    torch.cuda.manual_seed(random_seed)
    random.seed(random_seed)
    np.random.seed(random_seed)


def set_training_steps(config, num_samples):
    if config['parallel'] and config['dp_type'] == 'dp':
        config['num_iter_per_epoch'] = int(np.ceil(num_samples / config['batch_size']))
    else:
        config['num_iter_per_epoch'] = int(np.ceil(num_samples / (config['batch_size'] * config['num_gpus'])))
    if 'train_steps' not in config:
        config['train_steps'] = config['num_iter_per_epoch'] * config['num_epochs']
    if 'warmup_steps' not in config:
        config['warmup_steps'] = int(config['train_steps'] * config['warmup_ratio'])
    return config