import datetime
import glob
import ntpath
import os
import shutil
import traceback

import GPyOpt
from config_bayes import Configuration

from utils import call_subprocess, get_values_from_conff_matrix, load_previous_bayes_experiments


dir_path = os.path.dirname(os.path.realpath(__file__))
PROJECT_ROOT = os.path.join(dir_path, '..')

bayes_config_yaml = os.path.join(dir_path, 'config_bayes.yml')
config = Configuration(bayes_config_yaml)

date_string = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
bayes_params_file = open(os.path.join(PROJECT_ROOT, config.experiments.dir, f"{date_string}_bayes_params.txt"), 'a+')

def call_training_script(gaussian_hyps):
    cmd = 'python -u train.py'
    cmd += ' --epochs ' + gaussian_hyps['epochs'].__str__()
    cmd += ' --batch-size ' + gaussian_hyps['batch-size'].__str__()
    cmd += ' --cfg ' + config.train.cfg.__str__()
    cmd += ' --data ' + config.train.data.__str__()
    cmd += ' --multi-scale ' if gaussian_hyps['multi-scale'] else ""
    cmd += ' --img-size ' + gaussian_hyps['img-size']
    cmd += ' --rect ' if gaussian_hyps['rect'] else ""
    cmd += ' --weights ' + config.train.weights.__str__()
    cmd += ' --device ' + config.train.device.__str__()
    cmd += ' --adam ' if gaussian_hyps['adam'] else ""
    cmd += ' --freeze-layers ' if getattr(config.train, "freeze-layers") else ""
    # cmd += ' --snapshot-every ' if getattr(config.train, "snapshot-every") else ""
    cmd += ' --experiment-dir ' + config.experiments.dir.__str__()

    train_hyps = dict(
        (key, gaussian_hyps[key]) for idx, (key, _) in enumerate(gaussian_hyps.items()) if idx in range(6, 24))
    cmd += f' --hyp \"{train_hyps}\"'

    print("_______ CALLING TRAINING SCRIPT _______")
    print(cmd)

    call_subprocess(cmd)
    return cmd


def move_training_results_to_experiments_dir():
    training_results_dir_path = os.path.join(config.experiments.dir, datetime.datetime.now().strftime(
        '%Y-%m-%d_%H-%M-%S'))  # creating directory accordint to pattern eg: 2020-06-30_17-52-19

    print("_______ CALLING MOVING RESULTS _______")
    print(f"MOVING RESUTLS TO {training_results_dir_path}")

    os.mkdir(training_results_dir_path)

    weights_path = os.path.join(training_results_dir_path, 'best.pt')
    shutil.move('./weights/best.pt', weights_path)  # move best weights

    names_path = open(config.train.data).readlines()[3].split('=')[-1].rstrip()  # read names path from file
    names_file_name = ntpath.basename(names_path)
    experiment_names_path = os.path.join(training_results_dir_path, names_file_name)
    shutil.copy(names_path,
                experiment_names_path)  # copy names file from *.data file to created experiment dir with training results

    tensorboard_dir = './runs'
    tensorboard_events_files = glob.glob(os.path.join(tensorboard_dir, '*'))
    last_modified_events_file = max(tensorboard_events_files, key=os.path.getmtime)
    shutil.move(last_modified_events_file,
                os.path.join(training_results_dir_path))  # saving related tensorboard dir

    shutil.copy2(config.config_path, training_results_dir_path)  # copying configuration yaml

    return weights_path, experiment_names_path, training_results_dir_path


def call_detection_script(gaussian_hyps, weights_path, names_path, dir):
    detect_output_dir = os.path.join(dir, 'output')
    cmd = f"""python -u ./detect.py
          --cfg {config.train.cfg}
          --source {config.detect.source}
          --output {detect_output_dir}
          --names {names_path}
          --weights {weights_path}
          --test-img-size {getattr(config.detect, 'test-img-size')}
          --conf-thres {getattr(config.detect, 'conf-thres')}
          --iou-thres {getattr(config.detect, 'iou-thres')}
          --save-txt"""
    cmd += f" --device {config.train.device}" if config.train.device else ""
    cmd = " ".join(cmd.split())

    print("_______ CALLING DETECTION SCRIPT _______")
    print(cmd)

    call_subprocess(cmd)
    return detect_output_dir


def call_generate_confussion_matrix(detect_output_dir, names_path, train_results_dir):
    labels_dir = getattr(config.confussion_matrix, 'labels-dir')
    conff_matrix_path = os.path.join(train_results_dir, 'confussion-matrix.tsv')
    cmd = f"node ./our_scripts/generate-confusion-matrix.js {detect_output_dir} {labels_dir} {names_path} > {conff_matrix_path}"
    print("_______ CALLING CONFUSSION MATRIX SCRIPT _______")
    print(cmd)
    call_subprocess(cmd)
    return conff_matrix_path


def yolov3(x):
    bayes_hyps = {
        'epochs': int(x[:, 0]),
        'batch-size': int(x[:, 1]),
        'multi-scale': bool(x[:, 2]),
        'img-size': f"{int(x[:, 3])} {int(x[:, 4])}",
        'rect': bool(x[:, 5]),
        'adam': bool(x[:, 6]),
        'giou': float(x[:, 7]),  # train hyps start index
        'cls': float(x[:, 8]),
        'cls_pw': float(x[:, 9]),
        'obj': float(x[:, 10]),
        'obj_pw': float(x[:, 11]),
        'iou_t': float(x[:, 12]),
        'lr0': float(x[:, 13]),
        'lrf': float(x[:, 14]),
        'momentum': float(x[:, 15]),
        'weight_decay': float(x[:, 16]),
        'fl_gamma': float(x[:, 17]),
        'hsv_h': float(x[:, 18]),
        'hsv_s': float(x[:, 19]),
        'hsv_v': float(x[:, 20]),
        'degrees': float(x[:, 21]),
        'translate': float(x[:, 22]),
        'scale': float(x[:, 23]),
        'shear': float(x[:, 24]),  # train hyps end index
    }

    try:
        call_training_script(bayes_hyps)
        weights_path, names_path, train_results_dir = move_training_results_to_experiments_dir()
        detect_output_dir = call_detection_script(bayes_hyps, weights_path, names_path, train_results_dir)
        conf_matrix_path = call_generate_confussion_matrix(detect_output_dir, names_path, train_results_dir)

        y_dict = get_values_from_conff_matrix(conf_matrix_path) 

        # tutaj wzór na wyliczanie funkcji
        y_val = (1 - (y_dict['right'] * 10 - y_dict['false positives'] * 3 - y_dict['mistakes']) / y_dict['labeled']) / 30

        # zapisywanie do pliku zadeklarowanego globalnie
        line = "\t".join([bayes_hyps.__str__(), str(y_val)])
        print('###### line ########')
        print(line)
        bayes_params_file.writelines([line, '\n'])
        return y_val
    except:
        tb = traceback.format_exc()
        y_max_val = 1
        print("An error occured during running training-detect-confussion process \n", tb)
        print(f"Returning {y_max_val} from current bayessian iteration")
        line = "\t".join([bayes_hyps.__str__(), str(y_max_val)])
        bayes_params_file.writelines([line, '\n'])
        return y_max_val
# na jakiej rozdzielczości jest puszczana detekcja
if __name__ == '__main__':
    bounds = config.get_bayes_bounds()

    os.chdir(PROJECT_ROOT)  # change to project root directory
    # wczytywanie z poprzednich eksperymentów plik bayes_params
    X, Y =  load_previous_bayes_experiments(config.experiments.dir)

    constraints = [
        {
            'name':'img_size_constraint',
            'constraint': '(x[:,3] - x[:,4])' # img-size-start - img-size-end <= 0
        }
    ]


    bayes_optimizer = GPyOpt.methods.BayesianOptimization(f=yolov3, domain=bounds, X=X, Y=Y, verbosity=True,
                                                          initial_design_numdata=5, constraints=constraints, )
    bayes_optimizer.run_optimization(config.bayes.iterations, verbosity=True)
    bayes_params_file.close()