"""
The 'plots' module contains all plotting functions for visualizing models.
These functions require the installation of matplotlib for graphics.

## Structure:

- plot_comp_nums
- compare
"""

from . import plt
from . import random
from . import colors


def plot_comp_nums(Model, timesteps, starting_state=None, seed=100):
    """
    This is meant for plotting the number of people in each compartment over a period of time

    - Model: an instance of the `Model` class
    - timesteps: timesteps to plot as range(beg_time, end_time, day_length)
    - starting_state: initial conditions vector: [comp_1_start, comp_2_start, ...]
    - seed: =100, for generating new random colors
    - return: matplotlib plot
    """

    f, ax = plt.subplots(1, 1, figsize=(10, 4))
    integral = Model.integrate(timesteps, starting_state=starting_state)
    compartments = []

    for _ in integral[0]:
        compartments.append([])

    for timestep in integral:
        for compartment in range(len(timestep)):
            compartments[compartment].append(timestep[compartment])

    random.seed(a=seed)
    for layer in range(len(Model.layers)):
        color = (random.random(), random.random(), random.random())
        ax.plot(timesteps, compartments[layer], colors.to_hex(color), linewidth=2,
                label=Model.layer_names[layer])

    legend = ax.legend()
    legend.get_frame().set_alpha(0.5)

    ax.set_xlabel('Time (days)')
    ax.set_ylabel('People')
    ax.yaxis.set_tick_params(length=0)
    ax.xaxis.set_tick_params(length=0)

    for spine in ('top', 'right', 'bottom', 'left'):
        ax.spines[spine].set_visible(False)

    plt.show()


def compare(ranges, title="", subtitle="", markers=None, seed=200):
    """
    This function is used to compare multiple predictions which can range across time ranges.
    Often, this is used to compare real data with model predictions or to show the predictions after
    the real data in one window.

    - ranges: A list of items to plot. Follow the idiom below:
        ```python
        [timesteps (use range(beg, end)), predictions (corresponding to each element in `timesteps`,
        label (in str() format), ... ]
        ```
    - title: ="", (str) the title of the plot
    - subtitle: ="", (str) the subtitle of the plot
    - markers: =[], a list of sequenced markers in the format:\
         ['marker name', [...marker parameters]]\
         'marker name' can be any of:
            - line, param: [y, x1, x2, label]
            - highlighted-box, param: [x1, x2, y1 (axis units), y2 (axis units)]
            - point, param: [label, x, y]
            - arrow, param: [x, y, dx, dy]
    - seed: =200, for generating new random colors
    - return: matplotlib plot (log scale)
    """

    f, ax = plt.subplots(1, 1, figsize=(10, 4))
    plt.yscale('log')

    random.seed(a=seed)
    for r in ranges:
        color = (random.random(), random.random(), random.random())
        ax.plot(r[0], r[1], colors.to_hex(color), linewidth=2,
                label=r[2])

    if markers is not None:
        for marker in markers:

            if marker[0] == 'line':
                plt.hlines(marker[1][0], marker[1][1], marker[1][2], label=marker[1][3])

            if marker[0] == 'highlighted-box':
                plt.axvspan(marker[1][0], marker[1][1], ymin=marker[1][2], ymax=marker[1][3], facecolor='y', alpha=0.25)

            if marker[0] == 'point':
                plt.plot(marker[1][1], marker[1][2], 'ko')
                plt.annotate(marker[1][0], (marker[1][1] + 1, marker[1][2] + 1))

            if marker[0] == 'arrow':
                plt.arrow(marker[1][0], marker[1][1], marker[1][2], marker[1][3])

    legend = ax.legend()
    legend.get_frame().set_alpha(0.5)

    ax.set_xlabel('Time (days)')
    ax.set_ylabel('People')
    ax.yaxis.set_tick_params(length=0)
    ax.xaxis.set_tick_params(length=0)

    for spine in ('top', 'right', 'bottom', 'left'):
        ax.spines[spine].set_visible(False)

    plt.title(title+': '+subtitle)
    plt.show()
