"""
grafici_pl.py
=============

Funzioni (gia' pronte) per l'interpretazione geometrica della dualita'.
Non serve modificarle: basta chiamarle su un oggetto Problem_PL.

    disegna_primale(P)            regione ammissibile nel piano (x1, x2)
                                  -> serve un problema con 2 variabili
    disegna_duale(P)              regione ammissibile del duale nel piano (y1, y2):
                                  ogni colonna a_j di A diventa la retta
                                  a_j^T y = c_j, ortogonale al vettore a_j
                                  -> serve un problema con 2 vincoli
    disegna_primale_e_duale(P)    le due figure affiancate

Opzioni comuni:
    mostra_vertici=True    segna le intersezioni fra due rette (soluzioni di base):
                           punti pieni = ammissibili, punti vuoti = non ammissibili
    mostra_normali=True    (solo duale) disegna il vettore a_j su ciascuna retta
    salva="nome.png"       salva la figura su file invece di mostrarla

I vincoli di segno delle variabili (x >= 0, x <= 0, libera) vengono letti
direttamente dal modello Gurobi costruito dalla classe (attributi LB e UB delle
variabili): il disegno e' quindi sempre coerente con il problema risolto.

Requisiti: numpy, matplotlib e la classe Problem_PL con il metodo duale().
"""

import itertools
import numpy as np
import matplotlib.pyplot as plt

TOL = 1e-9
COLORI = ["tab:blue", "tab:orange", "tab:green", "tab:purple", "tab:brown",
          "tab:pink", "tab:olive", "tab:cyan", "tab:gray"]
SIMBOLI = {"<=": r"\leq", ">=": r"\geq", "==": "=", "=": "="}


# ---------------------------------------------------------------------------
# Funzioni di supporto
# ---------------------------------------------------------------------------
def _numero(v):
    """Scrive un numero in modo compatto: 2.0 -> '2', 0.5 -> '0.5'."""
    return f"{v:g}"


def _testo_lineare(coefficienti, nomi):
    """Scrive in LaTeX l'espressione  a1 nome1 + a2 nome2  (es. 'x_1 - 2x_2')."""
    pezzi = []
    for a, nome in zip(coefficienti, nomi):
        if abs(a) < TOL:
            continue
        segno = "-" if a < 0 else "+"
        valore = abs(a)
        coeff = "" if abs(valore - 1) < TOL else _numero(valore)
        pezzi.append((segno, coeff + nome))
    if not pezzi:
        return "0"
    testo = ("-" if pezzi[0][0] == "-" else "") + pezzi[0][1]
    for segno, termine in pezzi[1:]:
        testo += " " + segno + " " + termine
    return testo


def _limiti(P):
    """Limiti inferiori e superiori delle variabili, letti direttamente dal modello
    Gurobi: cosi' il disegno rispetta sempre i vincoli di segno, comunque siano
    stati scritti nella classe."""
    if P.modello is None:
        P.costruisci_modello()
    P.modello.update()
    lb = np.array(P.x.LB, dtype=float)
    ub = np.array(P.x.UB, dtype=float)
    return lb, ub


def _finito(v):
    return abs(v) < 1e20


def _semipiani(P, etichette, nomi):
    """Elenco dei vincoli del problema (righe di A) e dei limiti delle variabili."""
    elenco = []
    for i in range(P.m):
        elenco.append({"a": P.A[i].copy(), "b": P.b[i], "tipo": P.tipo_vincoli[i],
                       "etichetta": etichette[i], "segno": False})
    lb, ub = _limiti(P)
    for j in range(P.n):
        e = np.zeros(P.n)
        e[j] = 1.0
        for valore, tipo, simbolo, finito in ((lb[j], ">=", r"\geq", _finito(lb[j])),
                                              (ub[j], "<=", r"\leq", _finito(ub[j]))):
            if not finito:
                continue
            if abs(valore) < TOL:
                # vincolo di segno: coincide con un asse, che e' gia' disegnato
                elenco.append({"a": e, "b": 0.0, "tipo": tipo, "etichetta": None, "segno": True})
            else:
                elenco.append({"a": e, "b": valore, "tipo": tipo, "segno": False,
                               "etichetta": "$" + nomi[j] + " " + simbolo + " " + _numero(valore) + "$"})
    return elenco


def _soddisfa(punto, s, tol=1e-7):
    valore = s["a"] @ punto
    if s["tipo"] == "<=":
        return valore <= s["b"] + tol
    if s["tipo"] == ">=":
        return valore >= s["b"] - tol
    return abs(valore - s["b"]) <= tol


def _vertici(semipiani):
    """Intersezioni fra coppie di rette di bordo, con flag di ammissibilita'."""
    punti = []
    for s1, s2 in itertools.combinations(semipiani, 2):
        M = np.array([s1["a"], s2["a"]])
        if abs(np.linalg.det(M)) < TOL:
            continue
        p = np.linalg.solve(M, np.array([s1["b"], s2["b"]]))
        ammissibile = all(_soddisfa(p, s) for s in semipiani)
        if not any(np.allclose(p, q) for q, _ in punti):
            punti.append((p, ammissibile))
    return punti


def _riquadro(punti, extra, limiti):
    if limiti is not None:
        return limiti
    xs = [0.0] + [p[0] for p in punti] + [p[0] for p in extra]
    ys = [0.0] + [p[1] for p in punti] + [p[1] for p in extra]
    cx = (min(xs) + max(xs)) / 2
    cy = (min(ys) + max(ys)) / 2
    meta = 0.8 * max(max(xs) - min(xs), max(ys) - min(ys), 2.0)
    return (cx - meta, cx + meta, cy - meta, cy + meta)


def _taglia(poligono, a, b):
    """Interseca un poligono convesso con il semipiano a^T z <= b."""
    nuovo = []
    n = len(poligono)
    for k in range(n):
        p, q = poligono[k], poligono[(k + 1) % n]
        fp, fq = a @ p - b, a @ q - b
        if fp <= TOL:
            nuovo.append(p)
        if fp * fq < -TOL * TOL:
            t = fp / (fp - fq)
            nuovo.append(p + t * (q - p))
    return nuovo


def _regione(semipiani, riquadro):
    x0, x1, y0, y1 = riquadro
    poligono = [np.array([x0, y0]), np.array([x1, y0]), np.array([x1, y1]), np.array([x0, y1])]
    for s in semipiani:
        if len(poligono) == 0:
            break
        if s["tipo"] in ("<=", "==", "="):
            poligono = _taglia(poligono, s["a"], s["b"])
        if s["tipo"] in (">=", "==", "=") and len(poligono) > 0:
            poligono = _taglia(poligono, -s["a"], -s["b"])
    return poligono


def _segmento(a, b, riquadro):
    """Estremi del tratto della retta a^T z = b che cade nel riquadro."""
    x0, x1, y0, y1 = riquadro
    punti = []
    if abs(a[1]) > TOL:
        for x in (x0, x1):
            y = (b - a[0] * x) / a[1]
            if y0 - TOL <= y <= y1 + TOL:
                punti.append(np.array([x, y]))
    if abs(a[0]) > TOL:
        for y in (y0, y1):
            x = (b - a[1] * y) / a[0]
            if x0 - TOL <= x <= x1 + TOL:
                punti.append(np.array([x, y]))
    if len(punti) < 2:
        return None
    # i due punti piu' lontani fra loro
    migliore = max(itertools.combinations(punti, 2), key=lambda pq: np.linalg.norm(pq[0] - pq[1]))
    return migliore


def _direzione_illimitata(P):
    """Una direzione d (con |d_i| <= 1) lungo cui ci si muove restando ammissibili e
    l'obiettivo migliora: si risolve un piccolo problema di PL sul 'cono di recessione'."""
    import gurobipy as gp
    from gurobipy import GRB
    m = gp.Model()
    m.Params.OutputFlag = 0
    d = m.addMVar(2, lb=-1.0, ub=1.0)
    for i in range(P.m):
        riga = P.A[i] @ d
        t = P.tipo_vincoli[i]
        if t == "<=":
            m.addConstr(riga <= 0)
        elif t == ">=":
            m.addConstr(riga >= 0)
        else:
            m.addConstr(riga == 0)
    lb, ub = _limiti(P)
    for j in range(P.n):
        if _finito(lb[j]):
            m.addConstr(d[j] >= 0)
        if _finito(ub[j]):
            m.addConstr(d[j] <= 0)
    m.setObjective(P.c @ d, GRB.MAXIMIZE if P.senso == "max" else GRB.MINIMIZE)
    m.optimize()
    if m.Status != GRB.OPTIMAL:
        return None
    v = np.array(d.X)
    if np.linalg.norm(v) < TOL:
        return None
    return v / np.linalg.norm(v)


# ---------------------------------------------------------------------------
# Disegno di una regione nel piano
# ---------------------------------------------------------------------------
def _disegna(P, ax, nomi, etichette, titolo, mostra_vertici, mostra_normali, limiti):
    if P.modello is None or P.stato == "non ancora risolto":
        P.risolvi()
    semipiani = _semipiani(P, etichette, nomi)
    vertici = _vertici(semipiani)
    extra = []
    if P.stato == "ottimo":
        extra.append(P.soluzione())
    # il riquadro si adatta ai vertici ammissibili (o a tutti, se non ce ne sono)
    punti = [p for p, amm in vertici if amm]
    if len(punti) == 0:
        punti = [p for p, _ in vertici]
    direzione_ill = None
    ammissibili = [p for p, amm in vertici if amm]
    if P.stato == "illimitato" and ammissibili:
        direzione_ill = _direzione_illimitata(P)
        if direzione_ill is not None:
            xs = [p[0] for p in punti] + [0.0]
            ys = [p[1] for p in punti] + [0.0]
            ampiezza = max(max(xs) - min(xs), max(ys) - min(ys), 2.0)
            extra = extra + [ammissibili[0] + 1.5 * ampiezza * direzione_ill]
    riquadro = _riquadro(punti, extra, limiti)
    x0, x1, y0, y1 = riquadro
    lato = max(x1 - x0, y1 - y0)

    # regione ammissibile
    poligono = _regione(semipiani, riquadro)
    if len(poligono) >= 2:
        xs = [p[0] for p in poligono]
        ys = [p[1] for p in poligono]
        area = 0.5 * abs(sum(xs[k] * ys[(k + 1) % len(xs)] - xs[(k + 1) % len(xs)] * ys[k]
                             for k in range(len(xs))))
        if area > 1e-6 * lato * lato:
            ax.fill(xs, ys, color="tab:green", alpha=0.18, zorder=1)
            ax.plot(xs + [xs[0]], ys + [ys[0]], color="tab:green", lw=1.2, alpha=0.6, zorder=1)
        else:
            # regione "schiacciata" (segmento o semiretta): la evidenziamo con un tratto spesso
            p, q = max(itertools.combinations(poligono, 2),
                       key=lambda pq: np.linalg.norm(pq[0] - pq[1]))
            ax.plot([p[0], q[0]], [p[1], q[1]], color="tab:green", lw=7, alpha=0.35,
                    solid_capstyle="round", zorder=1)

    # assi
    ax.axhline(0, color="black", lw=0.8, zorder=2)
    ax.axvline(0, color="black", lw=0.8, zorder=2)

    # rette dei vincoli
    k = 0
    for s in semipiani:
        if s["segno"]:
            continue
        seg = _segmento(s["a"], s["b"], riquadro)
        colore = COLORI[k % len(COLORI)]
        k += 1
        if seg is None:
            continue
        p, q = seg
        ax.plot([p[0], q[0]], [p[1], q[1]], "-", color=colore, lw=2,
                label=s["etichetta"], zorder=3)
        if mostra_normali:
            m = (p + q) / 2
            v = s["a"] / np.linalg.norm(s["a"]) * 0.12 * lato
            ax.annotate("", xy=m + v, xytext=m,
                        arrowprops=dict(arrowstyle="->", color=colore, lw=1.8), zorder=4)

    # soluzioni di base (intersezioni di due rette)
    if mostra_vertici:
        for p, ammissibile in vertici:
            if not (x0 <= p[0] <= x1 and y0 <= p[1] <= y1):
                continue
            if ammissibile:
                ax.plot(p[0], p[1], "o", color="black", ms=6, zorder=5)
            else:
                ax.plot(p[0], p[1], "o", mfc="white", mec="gray", ms=6, zorder=5)

    # ottimo, curva di livello e direzione di miglioramento
    if P.stato == "ottimo":
        xo = P.soluzione()
        valore = P.Valore()
        seg = _segmento(P.c, valore, riquadro) if np.linalg.norm(P.c) > TOL else None
        if seg is not None:
            p, q = seg
            ax.plot([p[0], q[0]], [p[1], q[1]], "--", color="tab:red", lw=1.3,
                    label="curva di livello ottima", zorder=3)
        if np.linalg.norm(P.c) > TOL:
            direzione = P.c if P.senso == "max" else -P.c
            v = direzione / np.linalg.norm(direzione) * 0.15 * lato
            ax.annotate("", xy=xo + v, xytext=xo,
                        arrowprops=dict(arrowstyle="-|>", color="tab:red", lw=2), zorder=6)
        ax.plot(xo[0], xo[1], "*", color="tab:red", ms=16, mec="darkred", zorder=7,
                label="ottimo (" + _numero(round(xo[0], 4)) + ", " + _numero(round(xo[1], 4)) + ")")
        titolo += "\nottimo, valore = " + _numero(round(valore, 4))
    else:
        titolo += "\n" + P.stato
        # problema illimitato: una semiretta ammissibile lungo cui l'obiettivo migliora
        if direzione_ill is not None:
            p0 = ammissibili[0]
            q = p0 + 3 * lato * direzione_ill
            ax.plot([p0[0], q[0]], [p0[1], q[1]], ":", color="tab:red", lw=2, zorder=5)
            ax.annotate("", xy=p0 + 0.25 * lato * direzione_ill, xytext=p0,
                        arrowprops=dict(arrowstyle="-|>", color="tab:red", lw=2.5), zorder=6)
            ax.plot([], [], ":", color="tab:red", lw=2,
                    label="semiretta ammissibile: l'obiettivo cresce senza limite"
                    if P.senso == "max" else
                    "semiretta ammissibile: l'obiettivo decresce senza limite")
            # per confronto: il gradiente dell'obiettivo (spesso esce dalla regione!)
            gradiente = P.c if P.senso == "max" else -P.c
            if np.linalg.norm(gradiente) > TOL:
                g = gradiente / np.linalg.norm(gradiente) * 0.18 * lato
                ax.annotate("", xy=p0 + g, xytext=p0,
                            arrowprops=dict(arrowstyle="-|>", color="dimgray", lw=1.5,
                                            linestyle="--"), zorder=6)
                ax.plot([], [], "--", color="dimgray", lw=1.5,
                        label="gradiente: direzione di massima crescita")

    ax.set_xlim(x0, x1)
    ax.set_ylim(y0, y1)
    ax.set_aspect("equal", adjustable="box")
    ax.set_xlabel("$" + nomi[0] + "$")
    ax.set_ylabel("$" + nomi[1] + "$")
    ax.set_title(titolo, fontsize=11)
    ax.grid(alpha=0.25)
    ax.legend(fontsize=8, loc="best", framealpha=0.9)


def _etichette_primale(P, nomi):
    etichette = []
    for i in range(P.m):
        etichette.append("vincolo " + str(i + 1) + ": $" + _testo_lineare(P.A[i], nomi) + " "
                         + SIMBOLI[P.tipo_vincoli[i]] + " " + _numero(P.b[i]) + "$")
    return etichette


def _etichette_duale(D, nomi):
    # il vincolo j del duale corrisponde alla colonna a_j di A
    etichette = []
    for j in range(D.m):
        etichette.append("colonna $a_" + str(j + 1) + "$: $" + _testo_lineare(D.A[j], nomi) + " "
                         + SIMBOLI[D.tipo_vincoli[j]] + " " + _numero(D.b[j]) + "$")
    return etichette


def _segni(P, nomi):
    lb, ub = _limiti(P)
    testo = []
    for j in range(P.n):
        if _finito(lb[j]) and abs(lb[j]) < TOL and not _finito(ub[j]):
            testo.append("$" + nomi[j] + r" \geq 0$")
        elif _finito(ub[j]) and abs(ub[j]) < TOL and not _finito(lb[j]):
            testo.append("$" + nomi[j] + r" \leq 0$")
        elif not _finito(lb[j]) and not _finito(ub[j]):
            testo.append("$" + nomi[j] + "$ libera")
        else:
            testo.append("$" + _numero(lb[j]) + r" \leq " + nomi[j] + r" \leq " + _numero(ub[j]) + "$")
    return ", ".join(testo)


# ---------------------------------------------------------------------------
# Funzioni da usare
# ---------------------------------------------------------------------------
def disegna_primale(P, ax=None, mostra_vertici=True, limiti=None, salva=None):
    """Regione ammissibile del primale nel piano (x1, x2). Serve n = 2."""
    if P.n != 2:
        print("disegna_primale: servono 2 variabili, il problema ne ha", P.n)
        return
    nuova = ax is None
    if nuova:
        fig, ax = plt.subplots(figsize=(6.5, 6))
    nomi = ["x_1", "x_2"]
    titolo = "Primale (" + P.senso + ") nello spazio delle $x$   [" + _segni(P, nomi) + "]"
    _disegna(P, ax, nomi, _etichette_primale(P, nomi), titolo,
             mostra_vertici, False, limiti)
    if nuova:
        _mostra_o_salva(salva)


def disegna_duale(P, ax=None, mostra_vertici=True, mostra_normali=True, limiti=None, salva=None):
    """Regione ammissibile del duale nel piano (y1, y2). Serve m = 2:
    ogni colonna a_j di A e' la retta a_j^T y = c_j."""
    if P.m != 2:
        print("disegna_duale: servono 2 vincoli nel primale, il problema ne ha", P.m)
        return
    D = P.duale()
    nuova = ax is None
    if nuova:
        fig, ax = plt.subplots(figsize=(6.5, 6))
    nomi = ["y_1", "y_2"]
    titolo = "Duale (" + D.senso + ") nello spazio delle $y$   [" + _segni(D, nomi) + "]"
    _disegna(D, ax, nomi, _etichette_duale(D, nomi), titolo,
             mostra_vertici, mostra_normali, limiti)
    if nuova:
        _mostra_o_salva(salva)


def disegna_primale_e_duale(P, mostra_vertici=True, mostra_normali=True, salva=None):
    """Le due figure affiancate."""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 6))
    if P.n == 2:
        disegna_primale(P, ax=ax1, mostra_vertici=mostra_vertici)
    else:
        ax1.text(0.5, 0.5, "il primale ha " + str(P.n) + " variabili:\nnon si disegna nel piano",
                 ha="center", va="center", transform=ax1.transAxes)
        ax1.set_axis_off()
    if P.m == 2:
        disegna_duale(P, ax=ax2, mostra_vertici=mostra_vertici, mostra_normali=mostra_normali)
    else:
        ax2.text(0.5, 0.5, "il primale ha " + str(P.m) + " vincoli:\nil duale non si disegna nel piano",
                 ha="center", va="center", transform=ax2.transAxes)
        ax2.set_axis_off()
    fig.tight_layout()
    _mostra_o_salva(salva)


def _mostra_o_salva(salva):
    if salva is None:
        plt.show()
    else:
        plt.savefig(salva, dpi=150, bbox_inches="tight")
        plt.close()
        print("figura salvata in", salva)