from typing import Tuple


def get_SIR() -> Tuple[list, list]:
    beta = 0.2
    gamma = 0.09
    state_names = ["sain", "infecté", "rétabli"]
    relations = (
        lambda x: -beta * x["sain"][-1] * x["infecté"][-1],
        lambda x: beta * x["sain"][-1] * x["infecté"][-1] - gamma * x["infecté"][-1],
        lambda x: gamma * x["infecté"][-1],
    )
    return (state_names, relations)


def get_complex_SIRVD() -> Tuple[list, list]:
    beta = 0.2
    gamma = 0.09
    alpha = 0.6
    sigma = 0.1
    a = 0.01
    b = 0.005
    c = 0.15
    state_names = ["sain", "infecté", "rétabli", "vacciné", "vacciné infecté", "mort"]
    relations = (
        lambda x: -beta * x["sain"][-1] * x["infecté"][-1]
        - beta * x["sain"][-1] * x["vacciné infecté"][-1]
        - alpha * x["sain"][-1] * x["mort"][-1],
        lambda x: beta * x["sain"][-1] * x["infecté"][-1]
        + beta * x["sain"][-1] * x["vacciné infecté"][-1]
        - gamma * x["infecté"][-1]
        - a * x["infecté"][-1],
        lambda x: gamma * x["infecté"][-1] + c * x["vacciné infecté"][-1],
        lambda x: alpha * x["sain"][-1] * x["mort"][-1]
        - sigma * x["infecté"][-1] * x["vacciné"][-1]
        - sigma * x["vacciné infecté"][-1] * x["vacciné"][-1],
        lambda x: sigma * x["infecté"][-1] * x["vacciné"][-1]
        + sigma * x["vacciné infecté"][-1] * x["vacciné"][-1]
        - c * x["vacciné infecté"][-1]
        - b * x["vacciné infecté"][-1],
        lambda x: a * x["infecté"][-1] + b * x["vacciné infecté"][-1],
    )
    return (state_names, relations)


def get_complex_SIRVDV() -> Tuple[list, list]:
    beta = 0.4
    gamma = 0.05
    alpha = 0.6
    sigma = 0.1
    a = 0.01
    b = 0.05
    c = 0.04
    o = 2
    p = 0.01
    q = 0.01
    state_names = [
        "sain",
        "infecté",
        "rétabli (première vague)",
        "vacciné",
        "vacciné infecté",
        "mort",
        "infecté (variant)",
        "rétabli (deuxième vague)",
    ]
    relations = (
        lambda x: -beta * x["sain"][-1] * x["infecté"][-1]
        - beta * x["sain"][-1] * x["vacciné infecté"][-1]
        - alpha * x["sain"][-1] * x["mort"][-1]
        - o * x["sain"][-1] * x["infecté (variant)"][-1],
        lambda x: beta * x["sain"][-1] * x["infecté"][-1]
        + beta * x["sain"][-1] * x["vacciné infecté"][-1]
        - gamma * x["infecté"][-1]
        - a * x["infecté"][-1],
        lambda x: gamma * x["infecté"][-1]
        + c * x["vacciné infecté"][-1]
        - o * x["rétabli (première vague)"][-1] * x["infecté (variant)"][-1],
        lambda x: alpha * x["sain"][-1] * x["mort"][-1]
        - sigma * x["infecté"][-1] * x["vacciné"][-1]
        - sigma * x["vacciné infecté"][-1] * x["vacciné"][-1]
        - o * x["vacciné"][-1] * x["infecté (variant)"][-1],
        lambda x: sigma * x["infecté"][-1] * x["vacciné"][-1]
        + sigma * x["vacciné infecté"][-1] * x["vacciné"][-1]
        - c * x["vacciné infecté"][-1]
        - b * x["vacciné infecté"][-1],
        lambda x: a * x["infecté"][-1]
        + b * x["vacciné infecté"][-1]
        + q * x["infecté (variant)"][-1],
        lambda x: o * x["rétabli (première vague)"][-1] * x["infecté (variant)"][-1]
        + o * x["vacciné"][-1] * x["infecté (variant)"][-1]
        + o * x["sain"][-1] * x["infecté (variant)"][-1]
        - p * x["infecté (variant)"][-1]
        - q * x["infecté (variant)"][-1],
        lambda x: p * x["infecté (variant)"][-1],
    )
    return (state_names, relations)
