#!/usr/bin/env python3
import os
import re
import sys

import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd
import matplotlib as mpl


def obj_timeline(file_contents: bytes):
    objectives = []
    times = []
    for line in file_contents:
        if line.strip() == "==========":
            break
        match = re.match(r"objective\s=\s(\d+)", line)
        if match:
            objectives.append(int(match.group(1)))
            continue
        match = re.match(r"%\stime elapsed:\s(\d+\.\d+)\ss", line)
        if match:
            times.append(float(match.group(1)))
            continue

    assert len(objectives) > 0
    assert len(objectives) == len(times)

    return [(times[i], objectives[i]) for i in range(len(objectives))]


TAGMAP = {
    "original": "Base",
    "restart": "Restart Based LNS",
    "replay": "LNS Replay",
}

if __name__ == "__main__":
    folder = sys.argv[1]
    CONFIG = ["original", "restart"]
    solver = sys.argv[2]
    if solver == "Gecode":
        CONFIG.append("replay")
    statistics = dict()
    instances = set()
    # Read all the files
    for config in CONFIG:
        statistics[config] = dict()
        for root, dirs, files in os.walk(folder + "/" + config):
            for name in files:
                if not name.endswith(".sol"):
                    continue
                components = name[:-(4)].split(".")
                data = components[0]
                instances.add(data)
                seed = 1
                if len(components) > 1:
                    assert len(components) == 2
                    seed = components[1]

                if data not in statistics[config]:
                    statistics[config][data] = []

                with open(os.path.join(root, name)) as f:
                    contents = f.readlines()

                timeline = obj_timeline(contents)
                statistics[config][data].append(timeline)

    baseline = 0
    for data in instances:
        baseline += statistics["original"][data][0][0][1]

    times = []
    cumulative = []
    tag = []

    def emit(time, obj, conf):
        times.append(time)
        cumulative.append(obj)
        tag.append(TAGMAP[conf])

    for config in CONFIG:
        events = dict()
        for data in instances:
            length = len(statistics[config][data])
            for i in range(length):
                timeline = statistics[config][data][i]
                for j in range(1, len(timeline)):
                    if timeline[j][0] not in events:
                        events[timeline[j][0]] = 0
                    events[timeline[j][0]] += (
                        timeline[j][1] - timeline[j - 1][1]
                    ) / length

        sorted_events = sorted(events)
        incumbent = baseline
        emit(0, incumbent, config)
        for i in sorted_events:
            incumbent += events[i]
            emit(i, incumbent, config)
        emit(120, incumbent, config)

    df = pd.DataFrame(
        data={
            "Time (s)": times,
            "Cumulative Objective": cumulative,
            "Solver Version": tag,
        }
    )

    plot = sns.lineplot(
        data=df, x="Time (s)", y="Cumulative Objective", hue="Solver Version"
    )
    plot.figure.savefig("output.pdf")