"""
Le but de projet est de simuler l'évolution d'une épidémie.
Pour cela, on propose d'utiliser un modèle compartimental (SIR)
"""
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from typing import Any
import exercice_initial as init
import argparse
import os
import shutil


def remove_chache_folders(current_repo: str = ""):
    """
    when executing a python program, most os will create several __pycache__
    folders. These folders are not needed after each run and can take some
    memory. This function removes these folders from the disk.
    """
    if current_repo == "":
        new_refs = [elem for elem in os.listdir()]
    else:
        new_refs = [current_repo + "/" + elem for elem in os.listdir(current_repo)]
    for elem in new_refs:
        if os.path.isdir(elem):
            if "__pycache__" in elem:
                shutil.rmtree(elem)
            else:
                remove_chache_folders(current_repo=elem)


class Epidemic:
    def __init__(self, state_names: list) -> None:
        self.states = {}
        for key in state_names:
            self.states[key] = [0.0]
        self.states["sain"] = [0.99]
        self.states["infecté"] = [0.01]
        self.is_initialized = False
        self.delta = 1.0e-2

    def add_relations(self, relations: list):
        assert len(relations) == len(
            self.states
        ), "[ERROR] relations should be defiend for every states"
        self.relations = relations
        self.is_initialized = True

    def compute_an_update(self):
        for cpt, key in enumerate(self.states):
            ### To Complete ###
            pass

    def __call__(self, *args: Any, **kwds: Any) -> Any:
        if self.is_initialized:
            self.compute_an_update()
        pass


if __name__ == "__main__":
    possible_models = {
        "SIR": init.get_SIR,
        ### add more models ###
    }
    parser = argparse.ArgumentParser(
        usage=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        "--model",
        nargs="?",
        type=str,
        default="SIR",
        choices=list(possible_models.keys()),
        help="model to use",
    )
    args = parser.parse_args()
    total_iterations = int(1e4)
    state_names, relations = possible_models[args.model]()
    covid19 = Epidemic(state_names=state_names)
    covid19.add_relations(relations=relations)

    fig = plt.figure()
    plt.xlim([0, 1])
    plt.ylim([0, 1])
    lines = [plt.plot([], [], label=state_name)[0] for state_name in state_names]

    def init():
        for line in lines:
            line.set_data([], [])
        return lines

    x_data = np.linspace(0, 1, total_iterations)
    start_variant = True

    def animate(i):
        covid19()
        current_states = covid19.states
        total = 0
        for cpt, (_, value) in enumerate(current_states.items()):
            if len(value) > len(x_data):
                return lines
            lines[cpt].set_data(x_data[: len(value)], value)
            total += value[-1]
        global start_variant
        if i / total_iterations > 0.5 and args.model == "SIRVDV" and start_variant:
            covid19.states["infecté (variant)"][-1] = 0.01
            start_variant = False
        return lines

    _ = animation.FuncAnimation(
        fig,
        animate,
        init_func=init,
        repeat=False,
        frames=total_iterations,
        blit=True,
        interval=20,
        cache_frame_data=False,
    )

    plt.xlim([0, 1])
    plt.ylim([0, 1])
    plt.legend()
    plt.show()
    fig.savefig(f"{args.model}.png")
    plt.close()
    remove_chache_folders()