import ast
import io
import os
import subprocess
from glob import glob

import numpy as np


def call_subprocess(cmd):
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True)
    error = False
    for process_line_output in io.TextIOWrapper(process.stdout,
                                                encoding="utf-8"):  # print output of training process to console
        if 'Traceback' in process_line_output:
            error = True
        print(process_line_output)
    if error:
        raise RuntimeError("An error occured during calling subprocess")


def get_values_from_conff_matrix(path):
    lines = open(path, 'r').readlines()[:6]
    d = {}
    for l in lines:
        key, value, *_ = l.split("\t")
        d.update({key.replace(":", ""): int(value)})
    return d


def get_bayes_params_as_dict(x):
    return {
        '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
    }


def load_previous_bayes_experiments(experiments_dir):
    paths = list(glob(os.path.join(experiments_dir, '*bayes_params.txt')))
    if len(paths) == 0:
        print("No bayes files found")
        return None, None

    y_values = []
    x_values = []

    for p in paths:
        file = open(p, 'r')
        lines = file.readlines()
        for line in lines:
            try:
                dict_str, y_val = line.split('\t')
                bayes_dict = ast.literal_eval(dict_str)
                bayes_values = dict_to_numpy(bayes_dict)
                x_values.append(bayes_values)
                y_values.append(float(y_val))

                print("Loaded values from prevous experiments ", dict_str, y_val)
            except:
                raise Exception(f"Cannot parse line {line} from file {p}")
    if not y_values or not x_values:
        print("No bayes files found")
        return None, None

    return np.array(x_values), np.array(y_values).reshape((len(y_values), 1))


def dict_to_numpy(d):
    x = []
    for key, value in d.items():
        if type(value) == bool:
            x.append(int(value))
        elif key == 'img-size':
            start_img_size, end_img_size = value.split(" ")
            x.append(int(start_img_size))
            x.append(int(end_img_size))
        else:
            x.append(float(value))
    return x


if __name__ == '__main__':
    get_values_from_conff_matrix('/home/tomekb/yolov3/experiments/2020-08-17_02-05-43/confussion-matrix.tsv')