import argparse
import os
import re
import json
import numpy as np
import itertools

import pandas as pd
import pyscipopt as pyopt
import ecole as ec


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-d",
        "--dataset_path",
        type=str,
        required=True,
        help="Dataset with statistics, i.e. dataset of MILP instances that have been run.",
    )
    parser.add_argument(
        "-i",
        "--instances_path",
        type=str,
        required=True,
        help="Dataset with instances (either train or valid dir).",
    )
    parser.add_argument("-o", "--output_file", type=str, required=True)
    args = parser.parse_args()

    dataset_as_df = sweep_directories_to_dataframe(
        dataset_path=args.dataset_path,
        instances_path=args.instances_path,
    )

    dataset_as_df.to_csv(args.output_file)


def sweep_directories_to_dataframe(dataset_path: str, instances_path: str):
    """Walks file structure generated by Max's data gen script and builds csv file."""

    dataset_path = dataset_path

    instance_file_dataset = []
    config_vector_dataset = []
    time_limit_dataset = []
    initial_primal_bound_dataset = []
    initial_dual_bound_dataset = []
    primal_dual_integral_dataset = []

    config_encodings = _build_config_encodings()

    path_tree = os.walk(dataset_path, topdown=False)  # only 'leaf' directories
    path_regex = re.compile(
        ".*item_placement_(\d+)/config-(\d+)/seed-(\d+)/nsweep-(\d+).*"
    )

    for dir_path, _, files in path_tree:
        if match := path_regex.fullmatch(dir_path):
            instance_name, config_id, seed, nsweep = match.groups()

            with open(os.path.join(dir_path, "args.json")) as json_file:
                args = json.load(json_file)
                time_limit = args["timelimit"]

            with open(os.path.join(dir_path, "stats.json")) as json_file:
                stats = json.load(json_file)
                primal_dual_integral = stats["PrimalDualIntegral"]

            with open(
                os.path.join(instances_path, f"item_placement_{instance_name}.json")
            ) as json_file:
                initial_primal_dual = json.load(json_file)
                initial_primal = initial_primal_dual["primal_bound"]
                initial_dual = initial_primal_dual["dual_bound"]

            instance_file_dataset.append(f"item_placement_{instance_name}.mps.gz")
            initial_primal_bound_dataset.append(initial_primal)
            initial_dual_bound_dataset.append(initial_dual)
            time_limit_dataset.append(time_limit)
            config_vector_dataset.append(config_encodings[int(config_id)])
            primal_dual_integral_dataset.append(primal_dual_integral)

    dataset_info_pd = pd.DataFrame(
        {
            "instance_file": instance_file_dataset,
            "presolve_config_encoding": [c[0] for c in config_vector_dataset],
            "heuristic_config_encoding": [c[1] for c in config_vector_dataset],
            "separating_config_encoding": [c[2] for c in config_vector_dataset],
            "initial_primal_bound": initial_primal_bound_dataset,
            "initial_dual_bound": initial_dual_bound_dataset,
            "time_limit": time_limit_dataset,
            "time_limit_primal_dual_integral": primal_dual_integral_dataset,
        }
    )

    return dataset_info_pd


def _build_config_encodings():
    param_to_one_hot = {
        "OFF": 0,
        "DEFAULT": 1,
        "FAST": 2,
        "AGGRESSIVE": 3,
    }
    parameter_combinations = itertools.product(
        param_to_one_hot.keys(), repeat=3
    )  # e.g. ["DEFAULT", "AGGRESSIVE", "OFF"] for presolve, heuristic, separating

    config_encodings = [
        [param_to_one_hot[p] for p in comb] for comb in parameter_combinations
    ]
    assert len(config_encodings) == 4 * 4 * 4
    assert len(config_encodings[0]) == 3

    return config_encodings


if __name__ == "__main__":
    main()
