import numpy as np
import matplotlib.pyplot as plt

from solvers.model_smt import run_Z3
from solvers.model_cp import run_CPLEX, solving_MAPF
from utils.environments import *

"""
Run this file to replicate the experiment described in Extensive Testing subsection of the report.
This test was performed on a two-dimensional grid generated by NetworkX by varying the size 
of the graph and the number of agents.
"""

MIN_SIZE = 2
MAX_SIZE = 5
SEED = 42


def extensive_test(num_agents):
    """
    Test different utils usually increasing difficulty by means of graph size and number of agents.

    :param num_agents: the list containing for each graph the number of agents
    :return:
    """

    if len(num_agents) != MAX_SIZE - MIN_SIZE + 1:
        raise ValueError("sizes and number of agents must be equal.")

    time_CPLEX = []
    time_Z3 = []
    memory_usage_Z3 = []
    memory_usage_CPLEX = []
    number_of_conflicts_Z3 = []
    decisions_Z3 = []

    size = MIN_SIZE

    env_index = 1
    sep = "=" * 50

    while size <= MAX_SIZE:

        upper_bound = 2 * size
        number_of_agents = num_agents[env_index - 1]

        print(sep)
        try:
            agents, edges, graph = environments(nx.grid_2d_graph, number_of_agents, SEED, n=size, m=size)

            min_shortest_path, max_shortest_path = min_max_shortest_path(graph, agents)

        except Exception as e:
            print(e)
            return

        makespan = max_shortest_path
        print("ENVIRONMENT %d (%d agents and %d vertices)" % (env_index, number_of_agents, len(edges)))

        # --------------------------------------------------------------------------------------------------------------
        # Z3
        print(sep)
        print("Z3")
        check, solve_time, memory_usage, number_of_conflicts, decisions, _ = run_Z3(edges, agents, makespan)

        while not check and makespan <= upper_bound:
            makespan += 1
            check, solve_time, memory_usage, number_of_conflicts, decisions, _ = run_Z3(edges, agents, makespan)

        if not check and makespan >= upper_bound:
            print("Unsatisfiable")

        time_Z3.append(solve_time)
        memory_usage_Z3.append(memory_usage)
        number_of_conflicts_Z3.append(number_of_conflicts)
        decisions_Z3.append(decisions)

        # --------------------------------------------------------------------------------------------------------------
        # CPLEX
        print(sep)
        print("CPLEX")
        print("\nStep 1) Searching for optimal number of layers\n")
        check, ret, num_layers, solve_time, memory_usage, _, _ = \
            solving_MAPF(agents, edges, upper_bound, min_shortest_path)

        print("\nStep 2) Solving with %d layers\n" % num_layers)
        if check:
            _, _, solve_time, memory_usage, _, _, _ = \
                run_CPLEX(edges, agents, ret, num_layers)
        else:
            print("CPLEX: unsatisfiable")

        print(sep)

        time_CPLEX.append(solve_time)
        memory_usage_CPLEX.append(memory_usage)

        nx.draw(graph, with_labels=True)
        plt.show()

        size += 1
        env_index += 1

    """
    Statistics:
    
    time: Solving time
    memory_usage: Memory occupied in MB
    number_of_conflicts: A high number of conflicts could indicate that the solver didn't proceed very goal-directed
    decisions: Number of decision_points
    """

    fig, ax = plt.subplots(2, 2)
    size_range = np.linspace(MIN_SIZE, MAX_SIZE, MAX_SIZE - MIN_SIZE + 1)

    ax[0, 0].plot(size_range, number_of_conflicts_Z3, "-b", label='Z3')
    ax[0, 0].set_xlabel("Graph's size")
    ax[0, 0].set_ylabel("Number of conflicts")
    plt.yscale('log')

    ax[0, 1].plot(size_range, decisions_Z3, "-b", label='Z3')
    ax[0, 1].set_xlabel("Graph's size")
    ax[0, 1].set_ylabel("Number of decision points")
    plt.yscale('log')

    fig.suptitle("Z3 vs CPLEX")
    ax[1, 0].plot(size_range, time_Z3, "-b", label='Z3')
    ax[1, 0].plot(size_range, time_CPLEX, "-r", label='CPLEX')
    ax[1, 0].set_xlabel("Graph's size")
    ax[1, 0].set_ylabel("time")
    plt.yscale('log')

    ax[1, 1].plot(size_range, memory_usage_Z3, "-b", label='Z3')
    ax[1, 1].plot(size_range, memory_usage_CPLEX, "-r", label='CPLEX')
    ax[1, 1].set_xlabel("Graph's size")
    ax[1, 1].set_ylabel("Memory's usage (MB)")
    plt.yscale('log')

    plt.legend(loc='best')

    plt.show()


extensive_test([i for i in range(MIN_SIZE, MAX_SIZE + 1)])
