import re
import traceback
from typing import List, Dict, Any, Tuple
import numpy as np
import pandas as pd
import gradio as gr
from fpdf import FPDF

EPS = 1e-9

def parse_coeffs(text: str) -> List[float]:
    if not text or not text.strip():
        return []
    s = text.replace(',', ' ')
    parts = [p for p in s.split() if p.strip()]
    coeffs = []
    for p in parts:
        try:
            coeffs.append(float(eval(p)))
        except Exception:
            raise ValueError(f"Coeficiente inválido: '{p}'")
    return coeffs

def parse_constraints(text: str, nvars: int) -> Tuple[List[Dict[str,Any]], List[int]]:
    lines = [ln.strip() for ln in text.strip().splitlines() if ln.strip()]
    skip_words = ["tal que", "sujeito a", "subject to", "s.t.", "st:"]
    lines = [ln for ln in lines if not any(word in ln.lower() for word in skip_words)]

    free_vars = []
    cons = []
    pattern_free = re.compile(r'x([0-9]+)\s*(livre|free)', flags=re.I)

    # CORREÇÃO: regex robusto para coeficientes negativos ou omitidos
    term_pattern = r'([+-]?\d*(?:\.\d+)?)(x\d+)'

    for ln in lines[:]:
        m = pattern_free.search(ln)
        if m:
            idx = int(m.group(1)) - 1
            if idx < 0 or idx >= nvars:
                raise ValueError(f"Variável livre fora do intervalo: x{idx+1}")
            free_vars.append(idx)
            lines.remove(ln)

    for ln in lines:
        s = ln.replace(" ", "")
        if "<=" in s or "=<" in s:
            s = s.replace("=<", "<=")
            left, right = s.split("<=")
            sense = "<="
        elif ">=" in s or "=>" in s:
            s = s.replace("=>", ">=")
            left, right = s.split(">=")
            sense = ">="
        elif "=" in s:
            left, right = s.split("=")
            sense = "="
        else:
            raise ValueError(f"Faltando <=, >= ou =: '{ln}'")

        try:
            rhs = float(eval(right))
        except Exception:
            raise ValueError(f"RHS inválido em: '{ln}'")

        coeffs = [0.0] * nvars

        # agora os termos negativos funcionam corretamente
        terms = re.findall(term_pattern, left)

        for coef_str, var_str in terms:
            idx = int(var_str[1:]) - 1

            # Trata coeficientes omitidos ou sinal puro
            if coef_str in ["", "+", None]:
                v = 1.0
            elif coef_str == "-":
                v = -1.0
            else:
                v = float(eval(coef_str))

            coeffs[idx] += v

        cons.append({'coeffs': coeffs, 'sense': sense, 'rhs': rhs})

    return cons, sorted(list(set(free_vars)))

def expand_free_variables(nvars: int, c: List[float], constraints: List[Dict[str,Any]], free_vars: List[int]):
    new_c = []
    mapping = {}
    for i in range(nvars):
        if i in free_vars:
            new_c.append(c[i]); mapping[len(new_c)-1] = (i, +1)
            new_c.append(-c[i]); mapping[len(new_c)-1] = (i, -1)
        else:
            new_c.append(c[i]); mapping[len(new_c)-1] = (i, +1)
    new_constraints = []
    for row in constraints:
        coeffs = row['coeffs']
        new_coeffs = []
        for i in range(nvars):
            if i in free_vars:
                new_coeffs.append(coeffs[i])
                new_coeffs.append(-coeffs[i])
            else:
                new_coeffs.append(coeffs[i])
        new_constraints.append({'coeffs': new_coeffs, 'sense': row['sense'], 'rhs': row['rhs']})
    return len(new_c), new_c, new_constraints, mapping

# ---------------- Tableau helpers ----------------
def snapshot_html(tableau: np.ndarray, basis: List[int]) -> str:
    cols = tableau.shape[1]
    html = '<table border="1" style="border-collapse:collapse;font-family:Arial; font-size:12px;">'
    for i in range(tableau.shape[0]):
        html += '<tr>'
        for j in range(cols):
            val = tableau[i, j]
            html += f'<td style="padding:4px;">{val:.6g}</td>'
        html += '</tr>'
    html += '</table>'
    return html

def primal_simplex_tableau(T: np.ndarray, basis: List[int], max_iters=1000) -> Tuple[np.ndarray, List[int], List[Dict[str,Any]]]:
    m = T.shape[0] - 1
    ncols = T.shape[1]
    path = []
    path.append({'tableau': T.copy(), 'basis': basis.copy(), 'html': snapshot_html(T, basis)})

    it = 0
    while it < max_iters:
        it += 1
        obj_row = T[-1, :-1]
        entering_candidates = np.where(obj_row < -EPS)[0]
        if entering_candidates.size == 0:
            break
        entering = int(entering_candidates[0])
        ratios = np.full(m, np.inf)
        for i in range(m):
            a = T[i, entering]
            if a > EPS:
                ratios[i] = T[i, -1] / a
        if np.all(np.isinf(ratios)):
            raise Exception('Unbounded LP')
        leaving = int(np.argmin(ratios))
        piv = T[leaving, entering]
        T[leaving, :] = T[leaving, :] / piv
        for i in range(m+1):
            if i == leaving: continue
            T[i, :] = T[i, :] - T[i, entering] * T[leaving, :]
        basis[leaving] = entering
        path.append({'tableau': T.copy(), 'basis': basis.copy(), 'html': snapshot_html(T, basis)})
    return T, basis, path

# ---------------- Two-Phase implementation----------------
def build_tableau_two_phase(c: List[float], constraints: List[Dict[str,Any]], sense: str = 'max'):
    obj_mult = 1.0
    if sense == 'min':
        obj_mult = -1.0
    c_adj = [ci * obj_mult for ci in c]

    n = len(c_adj)
    m = len(constraints)

    slacks = 0
    artificials = 0
    for row in constraints:
        if row['sense'] == '<=':
            slacks += 1
        elif row['sense'] == '>=':
            slacks += 1
            artificials += 1
        else:
            artificials += 1

    total_cols = n + slacks + artificials + 1
    T = np.zeros((m + 1, total_cols))

    slack_idx = n
    artificial_idx = n + slacks

    basis = []
    art_positions = []
    s_counter = 0
    a_counter = 0

    for i, row in enumerate(constraints):
        coeffs = row['coeffs']
        T[i, :n] = coeffs
        if row['sense'] == '<=':
            T[i, slack_idx + s_counter] = 1.0
            basis.append(slack_idx + s_counter)
            s_counter += 1
        elif row['sense'] == '>=':
            T[i, slack_idx + s_counter] = -1.0
            T[i, artificial_idx + a_counter] = 1.0
            basis.append(artificial_idx + a_counter)
            art_positions.append(artificial_idx + a_counter)
            s_counter += 1
            a_counter += 1
        else:  # equality
            T[i, artificial_idx + a_counter] = 1.0
            basis.append(artificial_idx + a_counter)
            art_positions.append(artificial_idx + a_counter)
            a_counter += 1
        T[i, -1] = row['rhs']

    # Phase I objective: minimize sum of artificials.
    # Convert to maximization for our tableau solver: maximize (-sum a_j)
    # So c_phase1 (for maximization) = -1 for each artificial column.
    c_phase1 = np.zeros(total_cols - 1)
    for a in art_positions:
        c_phase1[a] = -1.0

    # In tableau we store -c in last row, so set T[-1, :-1] = -c_phase1
    T[-1, :-1] = -c_phase1

    # But because artificials are in basis, we must adjust objective row:
    # T[-1, :] = -c + sum_{i in basis} c_Bi * row_i, where c_Bi = c_phase1[basis_i]
    for i in range(m):
        bi = basis[i]
        cBi = c_phase1[bi] if bi < len(c_phase1) else 0.0
        if abs(cBi) > EPS:
            T[-1, :] += cBi * T[i, :]

    return T, basis, (n, slacks, artificials), art_positions, c_adj


def run_two_phase(c, constraints, sense='max'):    
    #FASE I
    T0, basis0, (n_orig, n_slack, n_art), art_positions, c_adj = build_tableau_two_phase(
        c, constraints, sense
    )

    try:
        T1, basis1, path1 = primal_simplex_tableau(T0.copy(), basis0.copy())
    except Exception as e:
        return {
            'status': 'phase1_failed',
            'error': str(e),
            'trace': traceback.format_exc()
        }

    phase1_obj = float(T1[-1, -1])

    # se sum(a_j) != 0 → inviável
    if abs(phase1_obj) > 1e-6:
        return {
            'status': 'infeasible',
            'phase1_obj': phase1_obj,
            'phase1_path': path1,
            'tableau_phase1': T1
        }

    # ---------- REMOVER ARTIFICIAIS ----------
    art_cols = set(art_positions)
    old_ncols = T1.shape[1] - 1
    keep_cols = [j for j in range(old_ncols) if j not in art_cols]

    # construir tableau da fase II (T2)
    T2 = np.zeros((T1.shape[0], len(keep_cols) + 1))
    for i, col in enumerate(keep_cols):
        T2[:, i] = T1[:, col]
    T2[:, -1] = T1[:, -1]

    # nova base
    basis2 = []
    for bi in basis1:
        if bi in art_cols:
            basis2.append(None)
        else:
            basis2.append(keep_cols.index(bi))

    # corrigir linhas onde a base ficou None
    used = set([b for b in basis2 if b is not None])
    m = T2.shape[0] - 1

    for i in range(m):
        if basis2[i] is None:
            replaced = False
            for j in range(T2.shape[1] - 1):
                if j not in used and abs(T2[i, j]) > EPS:
                    piv = T2[i, j]
                    T2[i, :] = T2[i, :] / piv
                    for r in range(m+1):
                        if r != i:
                            T2[r, :] -= T2[r, j] * T2[i, :]
                    basis2[i] = j
                    used.add(j)
                    replaced = True
                    break
            if not replaced:
                basis2[i] = None

    #FASE II — definir objetivo original
    c_full = []
    for col in keep_cols:
        if col < len(c_adj):
            c_full.append(c_adj[col])
        else:
            c_full.append(0.0)

    c_full = np.array(c_full)

    T2[-1, :-1] = -c_full

    for i in range(m):
        bi = basis2[i]
        if bi is not None and bi < len(c_full):
            coef = c_full[bi]
            if abs(coef) > EPS:
                T2[-1, :] += coef * T2[i, :]

    # preencher bases ausentes
    for i in range(m):
        if basis2[i] is None:
            for j in range(T2.shape[1]-1):
                if j not in used:
                    basis2[i] = j
                    used.add(j)
                    break

    #SIMPLEX FASE II
    try:
        T_final, basis_final, path2 = primal_simplex_tableau(T2.copy(), basis2.copy())
    except Exception as e:
        return {
            'status': 'phase2_failed',
            'error': str(e),
            'phase1_path': path1,
            'trace': traceback.format_exc()
        }

    #EXTRAI X*, REDUCED COSTS E DUAL (GERAL)
    x = [0.0] * n_orig

    for i, bi in enumerate(basis_final):
        if bi is not None:
            oldcol = keep_cols[bi]
            if oldcol < n_orig:
                x[oldcol] = float(T_final[i, -1])

    z = float(T_final[-1, -1])

    # custos reduzidos apenas variáveis originais
    reduced = []
    for j in range(n_orig):
        if j in keep_cols:
            colpos = keep_cols.index(j)
            z_j = -T_final[-1, colpos]
            reduced.append(round(c_adj[j] - z_j, 8))
        else:
            reduced.append(0.0)

    # Reconstruir matriz A (somente colunas originais) e b, c (originais)
    A_orig = np.array([row['coeffs'] for row in constraints], dtype=float)  # m x n_orig
    b_vec = np.array([row['rhs'] for row in constraints], dtype=float)
    cvec = np.array(c[:n_orig], dtype=float)

    # Construir matriz M das colunas que permaneceram no tableau (T_final[:m, :-1])
    M = T_final[:m, :-1].copy()  # m x ncols_keep

    # Basis matrix B (colunas básicas da fase II) — usar basis_final (índices em 0..ncols_keep-1)
    # Garantir que não haja None; se houver, já tentamos preencher antes.
    if any(bi is None for bi in basis_final):
        # Em casos degenerados, preencher com pseudo-solução: y zeros
        y_star = np.zeros(m)
    else:
        B = M[:, basis_final]  # m x m
        # montar c_B (custos das colunas básicas)
        cB = np.zeros(m)
        for i, bi in enumerate(basis_final):
            if bi < len(c_full):
                cB[i] = c_full[bi]
            else:
                cB[i] = 0.0
        # y^T = cB^T * B^{-1}
        try:
            Binv = np.linalg.inv(B)
            y_star = (cB @ Binv)  # shape (m,)
        except np.linalg.LinAlgError:
            # fallback: tentar solução via least squares
            try:
                y_star, *_ = np.linalg.lstsq(B.T, cB, rcond=None)
            except Exception:
                y_star = np.zeros(m)

    # dual objective b^T y
    dual_obj = float(b_vec @ y_star)

    # Definir folgas/violação das desigualdades do dual dependendo do sentido primal
    # Se primal == 'max' -> dual é min b^T y s.t. A^T y >= c  => slack = A^T y - c  (>=0)
    # Se primal == 'min' -> dual is max b^T y s.t. A^T y <= c => slack = c - A^T y  (>=0)
    if sense == 'max':
        dual_slacks = (A_orig.T @ y_star) - cvec
    else:
        dual_slacks = cvec - (A_orig.T @ y_star)

    # Preços-sombra (y) - ajustar sinal/convenção para exibição: mantemos y_star como calculado.
    shadow = []
    for i, row in enumerate(constraints):
        shadow.append(round(float(y_star[i]), 8))

    return {
        'status': 'optimal',
        'x': [round(v, 8) for v in x],
        'obj': round(z, 8),
        'y_dual': [round(float(v), 8) for v in y_star],
        'dual_obj': round(dual_obj, 8),
        'dual_slacks': [round(float(v), 8) for v in dual_slacks],
        'A': A_orig.tolist(),
        'b': b_vec.tolist(),
        'c': cvec.tolist(),
        'path_phase1': path1,
        'path_phase2': path2,
        'tableau_final': T_final,
        'basis_final': basis_final,
        'reduced_costs': reduced,
        'shadow_prices': shadow
    }


# ---------------- Helpers & PDF ----------------
def clean_vector(vec):
    try:
        return [float(v) for v in vec]
    except:
        return vec


# ---------------- Gradio handler ----------------
def run_algorithms(nvars_str, objective_str, cons_str, sense, mode):
    try:
        nvars = int(nvars_str)
        if nvars <= 0:
            return 'Erro: nvars deve ser inteiro positivo', '', '', '', ''
        c = parse_coeffs(objective_str)
        if len(c) != nvars:
            return 'Erro: coeficientes do objetivo não correspondem a nvars', '', '', '', ''
        constraints, free_vars = parse_constraints(cons_str, nvars)
        if free_vars:
            nvars, c, constraints, mapping = expand_free_variables(nvars, c, constraints, free_vars)
    except Exception as e:
        return f'Erro ao ler entrada: {e}', '', '', '', ''
        
    res = run_two_phase(c, constraints, sense)
    status = res.get('status')
    # infeasible detected in Phase I
    if status == 'infeasible':
        return f"Problema inviável (Fase I obj = {res.get('phase1_obj')})", '', '', '', ''

    # Phase I failed
    if status == 'phase1_failed':
        return f"Erro na Fase I: {res.get('error','(sem detalhe)')}", '', '', '', ''

    if status == 'optimal':
        x_primal = res['x']
        z_primal = res['obj']
        reduced = res.get('reduced_costs', [])
        shadow = res.get('shadow_prices', [])
        T_final = res.get('tableau_final', None)
        path_primal = res.get('path_phase2', [])
        path_phase1 = res.get('path_phase1', [])
        y_dual = res.get('y_dual', [])
        dual_obj = res.get('dual_obj', None)
        dual_slacks = res.get('dual_slacks', [])
        A = res.get('A', [])
        b = res.get('b', [])
        cvec = res.get('c', [])
    else:
        return f"Erro na resolução: status inesperado '{status}' - {res.get('error','')}", '', '', '', ''

    steps_html_phase2 = ""
    for idx, step in enumerate(path_primal):
        steps_html_phase2 += f"<h4>Fase II — Passo {idx+1} — Base: {step.get('basis','?')}</h4>"
        steps_html_phase2 += snapshot_html(np.array(step['tableau']), step.get('basis', [])) + "<br/>"

    steps_html_phase1 = ""
    for idx, step in enumerate(path_phase1):
        steps_html_phase1 += f"<h4>Fase I — Passo {idx+1} — Base: {step.get('basis','?')}</h4>"
        steps_html_phase1 += snapshot_html(np.array(step['tableau']), step.get('basis', [])) + "<br/>"

    df = pd.DataFrame({'Variável': [f'x{i+1}' for i in range(len(x_primal))], 'Valor': x_primal})
    solution_html = df.to_html(index=False)
    solution_html += f"<p><b>Valor ótimo (estimado) = {z_primal:.6g}</b></p>"

    x_primal = clean_vector(x_primal); reduced = clean_vector(reduced); shadow = clean_vector(shadow)
    z_primal = float(z_primal)

    model_txt = f"Objective ({'min' if sense=='min' else 'max'}): {c}\nConstraints:\n"
    for r in constraints:
        model_txt += f"  {r['coeffs']} {r['sense']} {r['rhs']}\n"

    summary = ""
    summary += f"Solução primal x* = {x_primal}\n"
    summary += f"Z_primal (estimado) = {z_primal:.6g}\n\n"
    summary += f"Solução dual y* = {y_dual}\n"
    summary += f"Valor dual b^T y = {dual_obj}\n"
    summary += f"Folgas/violação dual (A^T y - c) = {dual_slacks}\n\n"
    summary += f"Preços-sombra (dual interpretado) = {shadow}\n"
    summary += f"Custos reduzidos (vars originais) = {reduced}\n"

    return model_txt, solution_html, steps_html_phase1, steps_html_phase2, summary

# ---------------- Gradio UI ----------------

with gr.Blocks() as demo:
    gr.Markdown("# Simplex — Duas Fases (Fase I / Fase II) (Dual Geral)")
    with gr.Row():
        with gr.Column(scale=1):
            nvars = gr.Textbox(label='Número de variáveis (n)', value='2')
            objective = gr.Textbox(label='Coeficientes da função objetivo (ex: \"60 30\")', value='60 30')
            cons = gr.Textbox(label='Restrições (uma por linha). Ex.: 2x1 + 3x2 <= 300', lines=6,
                              value='2x1 + 4x2 >= 40\n3x1 + 2x2 >= 50')
            sense = gr.Radio(['max','min'], value='max', label='Tipo de objetivo')
            run = gr.Button('Executar Simplex (Duas Fases)')
        with gr.Column(scale=2):
            model_out = gr.Textbox(label='Função objetivo e restrições (modelo)', lines=6)
            solution_out = gr.HTML(label='Solução ótima (tabela)')
            steps_phase2_out = gr.HTML(label='Passos do Simplex (Phase II tableaus)')
            steps_phase1_out = gr.HTML(label='Passos do Simplex (Phase I tableaus)')
            summary_out = gr.Textbox(label='Resumo', lines=12)

    run.click(run_algorithms, inputs=[nvars, objective, cons, sense, gr.State(value='primal_and_dual')], outputs=[model_out, solution_out, steps_phase2_out, steps_phase1_out, summary_out])
   
    gr.Examples(examples=[["2","60 30","2x1 + 4x2 >= 40\n3x1 + 2x2 >= 50","max"]], inputs=[nvars, objective, cons, sense])

if __name__ == '__main__':   
    demo.launch(ssr_mode=False)