"""Modelo MILP para escolher a combinacao final de escalas."""

from __future__ import annotations

import numpy as np
import pandas as pd
from scipy.optimize import Bounds, LinearConstraint, milp


def resolver_milp(colunas: pd.DataFrame, rotas: pd.DataFrame, aeronaves: pd.DataFrame) -> pd.DataFrame:
    if colunas.empty:
        return colunas

    n = len(colunas)
    integrality = np.ones(n)
    bounds = Bounds(np.zeros(n), np.ones(n))
    restricoes = []

    # R1 - Cada tripulante aparece no máximo uma vez.
    for trip in sorted(set(colunas["trip1"]).union(set(colunas["trip2"]))):
        linha = np.zeros(n)
        linha[(colunas["trip1"] == trip) | (colunas["trip2"] == trip)] = 1
        restricoes.append(LinearConstraint(linha, -np.inf, 1))

    livres = aeronaves[(aeronaves["disponivel_sede"]) & (~aeronaves["em_pane"]) & (~aeronaves["em_missao_rota"])]["aeronave"].unique()
    for av in livres:
        idx_rota = (colunas["aeronave"] == av) & (colunas["tipo_escala"] == "ROTA_ACIONADA")
        idx_local = (colunas["aeronave"] == av) & (colunas["tipo_escala"] == "MISSAO_LOCAL")
        idx_sbv = (colunas["aeronave"] == av) & (colunas["tipo_escala"] == "SBV")
        if not (idx_rota | idx_sbv).any():
            raise RuntimeError(f"Modelo inviável: não há rota nem SBV candidato para a aeronave {av}.")
        linha = np.zeros(n)
        linha[idx_rota | idx_sbv] = 1
        restricoes.append(LinearConstraint(linha, 1, 1))
        if idx_rota.any():
            linha = np.zeros(n)
            linha[idx_rota] = 1
            restricoes.append(LinearConstraint(linha, -np.inf, 1))
        if idx_local.any():
            linha = np.zeros(n)
            linha[idx_local] = 1
            restricoes.append(LinearConstraint(linha, -np.inf, 1))
        if idx_rota.any() and idx_local.any():
            linha = np.zeros(n)
            linha[idx_rota | idx_local] = 1
            restricoes.append(LinearConstraint(linha, -np.inf, 1))

    for rota_id in range(1, len(rotas) + 1):
        idx = (colunas["tipo_escala"] == "ROTA_ACIONADA") & (colunas["rota_id"] == rota_id)
        if not idx.any():
            raise RuntimeError(f"Modelo inviável: não há coluna candidata para a rota acionada {rota_id}.")
        linha = np.zeros(n)
        linha[idx] = 1
        restricoes.append(LinearConstraint(linha, 1, 1))

    resultado = milp(c=-colunas["score_milp"].to_numpy(dtype=float), integrality=integrality, bounds=bounds, constraints=restricoes)
    if not resultado.success:
        raise RuntimeError(f"MILP não encontrou solução viável: {resultado.message}")
    selecionadas = resultado.x > 0.5
    return colunas.loc[selecionadas].reset_index(drop=True)