#!/usr/bin/env python3
"""El modelo de Chiu y Jain: dos flujos, un enlace y una señal común de congestión.

En cada paso ambos reciben la misma señal: si x1 + x2 supera la capacidad, bajan; si no,
suben. Se comparan las cuatro combinaciones de aumento y reducción.

    python3 phase.py aimd --x1 10 --x2 70 --steps 12
    python3 phase.py all --png fases.png
"""
import argparse

RULES = {
    # nombre: (aumento, reducción)
    "aimd": (lambda x: x + 5, lambda x: x * 0.5),
    "aiad": (lambda x: x + 5, lambda x: max(x - 20, 0)),
    "mimd": (lambda x: x * 1.25, lambda x: x * 0.5),
    "miad": (lambda x: x * 1.25, lambda x: max(x - 20, 0)),
}


def walk(rule, x1, x2, capacity=100, steps=200):
    """Devuelve la lista de puntos (x1, x2) por los que pasa el sistema."""
    increase, decrease = RULES[rule]
    points = [(x1, x2)]
    for _ in range(steps):
        change = decrease if x1 + x2 > capacity else increase
        x1, x2 = change(x1), change(x2)
        points.append((x1, x2))
    return points


def jain(x1, x2):
    if x1 == 0 and x2 == 0:
        return 0.0
    return (x1 + x2) ** 2 / (2 * (x1 * x1 + x2 * x2))


def plot(path, x1, x2, capacity):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    figure, axes = plt.subplots(1, 4, figsize=(14, 3.6), sharex=True, sharey=True)
    for ax, rule in zip(axes, RULES):
        points = walk(rule, x1, x2, capacity, steps=60)
        ax.plot([0, capacity], [capacity, 0], color="gray", label="eficiencia")
        ax.plot([0, capacity], [0, capacity], color="gray", linestyle="--", label="equidad")
        ax.plot(*zip(*points), marker=".", linewidth=0.7)
        ax.set_title(rule.upper())
        ax.set_xlim(0, capacity * 1.3)
        ax.set_ylim(0, capacity * 1.3)
        ax.set_xlabel("x1")
    axes[0].set_ylabel("x2")
    axes[0].legend()
    figure.tight_layout()
    figure.savefig(path, dpi=150)


def main():
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("rule", choices=[*RULES, "all"])
    parser.add_argument("--x1", type=float, default=10)
    parser.add_argument("--x2", type=float, default=70)
    parser.add_argument("--capacity", type=float, default=100)
    parser.add_argument("--steps", type=int, default=12)
    parser.add_argument("--png", default=None)
    args = parser.parse_args()
    if args.png:
        plot(args.png, args.x1, args.x2, args.capacity)
        print(f"gráfico en {args.png}")
        return
    if args.rule == "all":
        print(f"{'regla':<6} {'x1':>7} {'x2':>7} {'suma':>7} {'|x1-x2|':>8} {'Jain':>6}")
        for rule in RULES:
            # El último punto: el sistema oscila alrededor de la capacidad, pero x1/x2 se ve bien.
            last = walk(rule, args.x1, args.x2, args.capacity)[-1]
            print(f"{rule:<6} {last[0]:>7.1f} {last[1]:>7.1f} {last[0] + last[1]:>7.1f} "
                  f"{abs(last[0] - last[1]):>8.1f} {jain(*last):>6.3f}")
        return
    print(f"{'paso':>4} {'x1':>7} {'x2':>7} {'suma':>7} {'|x1-x2|':>8} {'Jain':>6}")
    for number, (x1, x2) in enumerate(walk(args.rule, args.x1, args.x2, args.capacity,
                                           args.steps)):
        print(f"{number:>4} {x1:>7.1f} {x2:>7.1f} {x1 + x2:>7.1f} {abs(x1 - x2):>8.1f} "
              f"{jain(x1, x2):>6.3f}")


if __name__ == "__main__":
    main()