from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import sys
__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.abspath(os.path.join(__dir__, '../')))

import subprocess
import numpy as np

from ppcls.utils import config


def get_result(log_dir):
    log_file = "{}/train.log".format(log_dir)
    with open(log_file, "r") as f:
        raw = f.read()
    res = float(raw.split("best metric: ")[-1].split("]")[0])
    return res


def search_train(search_list, base_program, base_output_dir, search_key,
                 config_replace_value, model_name, search_times=1):
    best_res = 0.
    best = search_list[0]
    all_result = {}
    for search_i in search_list:
        program = base_program.copy()
        for v in config_replace_value:
            program += ["-o", "{}={}".format(v, search_i)]
            if v == "Arch.name":
                model_name = search_i
        res_list = []
        for j in range(search_times):
            output_dir = "{}/{}_{}_{}".format(base_output_dir, search_key, search_i, j).replace(".", "_")
            program += ["-o", "Global.output_dir={}".format(output_dir)]
            process = subprocess.Popen(program)
            process.communicate()
            res = get_result("{}/{}".format(output_dir, model_name))
            res_list.append(res)
        all_result[str(search_i)] = res_list

        if np.mean(res_list) > best_res:
            best = search_i
            best_res = np.mean(res_list)
    all_result["best"] = best
    return all_result


def search_strategy():
    args = config.parse_args()
    configs = config.get_config(args.config, overrides=args.override, show=False)
    base_config_file = configs["base_config_file"]
    distill_config_file = configs["distill_config_file"]
    model_name = config.get_config(base_config_file)["Arch"]["name"]
    gpus = configs["gpus"]
    gpus = ",".join([str(i) for i in gpus])
    base_program = ["python3.7", "-m", "paddle.distributed.launch", "--gpus={}".format(gpus),
                    "tools/train.py", "-c", base_config_file]
    base_output_dir = configs["output_dir"]
    search_times = configs["search_times"]
    search_dict = configs.get("search_dict")
    all_results = {}
    for search_i in search_dict:
        search_key = search_i["search_key"]
        search_values = search_i["search_values"]
        replace_config = search_i["replace_config"]
        res = search_train(search_values, base_program, base_output_dir,
                           search_key, replace_config, model_name, search_times)
        all_results[search_key] = res
        best = res.get("best")
        for v in replace_config:
            base_program += ["-o", "{}={}".format(v, best)]

    teacher_configs = configs.get("teacher", None)
    if teacher_configs is not None:
        teacher_program = base_program.copy()
        # remove incompatible keys
        teacher_rm_keys = teacher_configs["rm_keys"]
        rm_indices = []
        for rm_k in teacher_rm_keys:
            for ind, ki in enumerate(base_program):
              if rm_k in ki:
                rm_indices.append(ind)
        for rm_index in rm_indices[::-1]:
            teacher_program.pop(rm_index)
            teacher_program.pop(rm_index-1)
        replace_config = ["Arch.name"]
        teacher_list = teacher_configs["search_values"]
        res = search_train(teacher_list, teacher_program, base_output_dir, "teacher", replace_config, model_name)
        all_results["teacher"] = res
        best = res.get("best")
        t_pretrained = "{}/{}_{}_0/{}/best_model".format(base_output_dir, "teacher", best, best)
        base_program += ["-o", "Arch.models.0.Teacher.name={}".format(best),
                         "-o", "Arch.models.0.Teacher.pretrained={}".format(t_pretrained)]
    output_dir = "{}/search_res".format(base_output_dir)
    base_program += ["-o", "Global.output_dir={}".format(output_dir)]
    final_replace = configs.get('final_replace')
    for i in range(len(base_program)):
      base_program[i] = base_program[i].replace(base_config_file, distill_config_file)
      for k in final_replace:
        v = final_replace[k]
        base_program[i] = base_program[i].replace(k, v)

    process = subprocess.Popen(base_program)
    process.communicate()
    print(all_results, base_program)


if __name__ == '__main__':
    search_strategy()