# это скрипт с функциями, используемыми из эксперимента Anylogic
# и для обработки результатов в виде 3D-графика

from typing import Optional
try:
    import os
    import numpy as np
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import matplotlib.ticker as plticker
    from  matplotlib.colors import LinearSegmentedColormap
    from mpl_toolkits.mplot3d import Axes3D
 
    IMPORT_SUCCESS = True
except ModuleNotFoundError:
    IMPORT_SUCCESS = False
    

def get_num_workers(iteration_index: int) -> int:
    ### функция возвращает 4 и затем увеличивает на 1 каждые 3 итерации
    return (iteration_index // 3) + 4

def get_arrival_rate(iteration_index: int) -> float:
    ### функция возвращает 1.5 и затем увеличивает на 0.5 каждые 3 итерации
    return (iteration_index % 3)/2 + 1.5

def reset():
    """
      очиcтка файла журнала и удаление cтарого графика
      функция вызывается в начале эксперимента
    """
    with open("model_outputs_log.csv", "w") as _:
        pass
    
    try:
        os.remove("model_outputs_3d.png")
    except FileNotFoundError:
        pass

def log_outputs(num_workers: int, arrival_rate: float, queue_time: float):
    """
      запись данных ввода/вывода в файл журнала
      функция вызывается в конце каждого прогона
    """
    with open("model_outputs_log.csv", "a") as outfile:
        outfile.write(f"{num_workers},{arrival_rate},{queue_time}\n")

def save_results() -> Optional[bool]:
    """
      чтение файла журнала и сохранение 3D графика
      функция возвращает имя файла или None
    """
    if not IMPORT_SUCCESS:
        return None

    fig = plt.figure()
    fig.set_tight_layout(True)

    ax = fig.add_subplot(projection='3d')
    ax.set_title("Результаты эксперимента")
    ax.set_xlabel("число работников")
    ax.set_ylabel("число пациентов в мин") 
    ax.set_zlabel("время ожидания (мин)")
    
    ex_array = np.loadtxt("model_outputs_log.csv", delimiter=",", ndmin=2)
    if ex_array.size > 0:  # в файле нужны данные
        x = ex_array[:,0]
        y = ex_array[:,1]
        #z = ex_array[:,2]
        z = [2, 3, 4, 5, 8, 5, 7]
        
        # настройка цвета по z значениям
        #cmap = LinearSegmentedColormap.from_list('gr',["green", "yellow", "red"], N=256) 
        #colors = cmap((z-min(z))/(max(z)-min(z)) )

        # можно сделать столбики меньше чем расстояние между точками
        width, depth = 1*0.5, 0.5*0.5
        
        ax.bar3d(x, y, np.zeros_like(z), width, depth, z, shade=True, color='crimson')

        # изменение разметки осей x, y
        ax.xaxis.set_major_locator(plticker.MultipleLocator(base=1.0))
        ax.yaxis.set_major_locator(plticker.MultipleLocator(base=0.5))
        
    filename = "model_outputs_3d.png"
    fig.savefig(filename)
    return filename
