import datetime
import glob
import io
import ntpath
import os
import shutil
import subprocess

from config import Configuration


def call_training_script(config):
    cmd = '/home/tomekb/miniconda3/envs/conda3.7/bin/python -u /home/tomekb/yolov3/train.py '
    cmd += f"--experiment-dir {config.experiments.dir}"
    cmd += config.train.get_args_string()  # getting rest of train arguments

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

    os.chdir('..')  # change to project root directory
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
    for line in io.TextIOWrapper(process.stdout, encoding="utf-8"):  # print output of training process to console
        print(line)
    return cmd


def move_training_results_to_experiments_dir(config):
    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('/home/tomekb/yolov3/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'
    last_modified_tensorboard_dir = max(glob.glob(os.path.join(tensorboard_dir, '*/')), key=os.path.getmtime)
    shutil.move(last_modified_tensorboard_dir, os.path.join(training_results_dir_path))  # saving related tensorboard dir


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

    # for test purposes only
    # shutil.copy2('/home/tomekb/yolov3/experiments/1/best.pt', training_results_dir_path)

    return weights_path, experiment_names_path, training_results_dir_path


def call_detection_script(config, weights_path, names_path, dir):
    detect_output_dir = os.path.join(dir, 'output')
    cmd = f"""/home/tomekb/miniconda3/envs/conda3.7/bin/python -u /home/tomekb/yolov3/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 += " --agnostic-nms" if getattr(config.detect, 'agnostic-nms') else ""
    cmd += " --agument" if getattr(config.detect, 'augment') else ""
    cmd += f" --device {config.train.device}" if config.train.device else ""

    cmd = " ".join(cmd.split())

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

    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
    for line in io.TextIOWrapper(process.stdout, encoding="utf-8"):  # print output of  process to console
        print(line)

    return detect_output_dir


def call_generate_confussion_matrix(detect_output_dir, config, names_path, train_results_dir):
    labels_dir = getattr(config.confussion_matrix, 'labels-dir')

    cmd = f"node ./our_scripts/generate-confusion-matrix.js {detect_output_dir} {labels_dir} {names_path} > {train_results_dir}/confussion-matrix.tsv"
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
    print("_______ CALLING CONFUSSION MATRIX SCRIPT _______")
    print(cmd)
    for line in io.TextIOWrapper(process.stdout, encoding="utf-8"):  # print output of  process to console
        print(line)


if __name__ == '__main__':
    config = Configuration()

    train_cmd = call_training_script(config)
    weights_path, names_path, train_results_dir = move_training_results_to_experiments_dir(config)
    detect_output_dir = call_detection_script(config, weights_path, names_path, train_results_dir)
    call_generate_confussion_matrix(detect_output_dir, config, names_path, train_results_dir)