import argparse
from pathlib import Path
import os

import json
from HPOBenchExperimentUtils.utils.runner_utils import get_benchmark_names


def main(args):
    dir_ = Path(args.out_run) / args.exp / "mumbo"
    for rundir in dir_.rglob("run-*"):
        hporh_fl = rundir / "hpobench_runhistory.txt"
        hpotr_fl = rundir / "hpobench_trajectory_v1.txt"
        mumbotr_fl = rundir / "mumbo_trajectory.json"
        newtr_fl = rundir / "hpobench_trajectory_v3.txt"

        for fl in hporh_fl, hpotr_fl, mumbotr_fl:
            assert fl.is_file(), fl

        # 1) Read in all files. We need the runhistory and both trajectories
        hporh, hpotr, mumbotr = ([], [], [])
        for f, a in ((hporh_fl, hporh), (hpotr_fl, hpotr), (mumbotr_fl, mumbotr)):
            with open(f, "r") as fh:
                a.extend([json.loads(line) for line in fh.readlines()])

        # 2) The first 10*D evaluations are the initial design. We will use the configurations from
        # the autogenerated trajectory for these iterations.
        max_call = len(hporh) - 1 - len(mumbotr)
        assert len(mumbotr) == len(hporh[1 + max_call:]), f"Something went wrong in {rundir}"

        final_traj = [hporh[0], ]
        max_bud = {"val": -1}
        for i in hpotr[1:]:
            # All entries with a lower function_call value stem from the initial design
            if i["function_call"] <= max_call:
                final_traj.append(i)
            if max_bud is None or list(max_bud.values())[0] < list(i["fidelity"].values())[0]:
                max_bud = i["fidelity"]

        # 3) Now, finally, add all configurations from the trajectory generated by mumbo. However,
        # we'll need to compute time stamps
        if len(mumbotr) > 0:
            for m, h in zip(mumbotr, hporh[1 + max_call:]):
                mfin = m["finish_time"]
                hfin = h["finish_time"]
                # Add time it took to compute incumbent
                total_time_used = h["total_time_used"] + (mfin - hfin)

                dc = {"start_time": None,
                      "finish_time": hfin,
                      "function_value": None,
                      "fidelity": max_bud,
                      "cost": None,
                      "configuration": m["configuration"],
                      "function_call": h["function_call"],
                      "total_time_used": total_time_used,
                      "total_objective_costs": h["total_objective_costs"],
                      "info": {
                          "fidelity": max_bud,
                      },
                      }
                final_traj.append(dc)

        # Now write the new trajectory
        if args.dry:
          print(final_traj)
          continue
        assert not newtr_fl.is_file()
        with open(newtr_fl, "w") as fh:
            for entry in final_traj:
                json.dump(entry, fh)
                fh.write(os.linesep)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument('--exp', required=True, choices=get_benchmark_names())
    parser.add_argument('--out-run', default="./exp_outputs", type=str)
    parser.add_argument('--nrep', type=int, default=32)
    parser.add_argument('--dry', action="store_true", default=False)

    args, unknown = parser.parse_known_args()
    main(args)
