modules/classificador.py

"""
Classificador de Ocupacao - NT-01/2025 CBMGO
Classifica edificacoes em grupos/divisoes e lista exigencias.
Com LLM (cbmgo/llama3-8b-ocupacao) ou regras locais (fallback).
"""
import json
import re

try:
    from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
    import torch
    HAS_TRANSFORMERS = True
except ImportError:
    HAS_TRANSFORMERS = False


# ============================================================
# TABELA DE CLASSIFICACAO NT-01/2025 (Tabela 1)
# ============================================================
TABELA_OCUPACOES = {
    "A-1": {
        "descricao": "Residencia unifamiliar",
        "exemplos": ["casa", "sobrado", "chacara", "sitio"],
        "exigencias_base": ["extintores_portateis", "sinalizacao_emergencia"],
        "area_extintor_m2": 500,
        "capacidade_extintor": "2-A:10-B:C"
    },
    "A-2": {
        "descricao": "Residencia multifamiliar",
        "exemplos": ["apartamento", "condominio", "edificio residencial"],
        "exigencias_base": ["extintores_portateis", "sinalizacao_emergencia", "hidrantes"],
        "area_extintor_m2": 500,
        "capacidade_extintor": "2-A:10-B:C"
    },
    "B-1": {
        "descricao": "Hotel, pousada, hostel",
        "exemplos": ["hotel", "pousada", "hostel", "motel", "hospedagem"],
        "exigencias_base": ["extintores_portateis", "hidrantes", "sinalizacao_emergencia",
                           "iluminacao_emergencia", "deteccao_automatica", "alarme_incendio"],
        "area_extintor_m2": 250,
        "capacidade_extintor": "2-A:20-B:C"
    },
    "C-1": {
        "descricao": "Teatro, cinema, auditorio",
        "exemplos": ["teatro", "cinema", "auditorio", "show", "espetaculo"],
        "exigencias_base": ["extintores_portateis", "hidrantes", "sinalizacao_emergencia",
                           "iluminacao_emergencia", "alarme_incendio", "saidas_emergencia_multiplas"],
        "area_extintor_m2": 250,
        "capacidade_extintor": "2-A:10-B:C"
    },
    "C-2": {
        "descricao": "Igreja, templo, local de culto",
        "exemplos": ["igreja", "templo", "culto", "missa", "evangelico"],
        "exigencias_base": ["extintores_portateis", "hidrantes", "sinalizacao_emergencia",
                           "iluminacao_emergencia"],
        "area_extintor_m2": 300,
        "capacidade_extintor": "2-A:10-B:C"
    },
    "D-1": {
        "descricao": "Comercio em geral, loja, mercado",
        "exemplos": ["loja", "mercado", "supermercado", "comercio", "varejo", "boutique"],
        "exigencias_base": ["extintores_portateis", "hidrantes", "sinalizacao_emergencia"],
        "area_extintor_m2": 400,
        "capacidade_extintor": "2-A:20-B:C"
    },
    "D-2": {
        "descricao": "Escritorio, servico profissional",
        "exemplos": ["escritorio", "consultorio", "clinica", "advocacia", "contabilidade", "banco"],
        "exigencias_base": ["extintores_portateis", "sinalizacao_emergencia"],
        "area_extintor_m2": 500,
        "capacidade_extintor": "2-A:10-B:C"
    },
    "E-1": {
        "descricao": "Escola, creche, universidade",
        "exemplos": ["escola", "colegio", "creche", "universidade", "faculdade", "educacional"],
        "exigencias_base": ["extintores_portateis", "hidrantes", "sinalizacao_emergencia",
                           "iluminacao_emergencia", "saidas_emergencia"],
        "area_extintor_m2": 300,
        "capacidade_extintor": "2-A:10-B:C"
    },
    "F-1": {
        "descricao": "Hospital, clinica com internacao",
        "exemplos": ["hospital", "uti", "pronto-socorro", "internacao"],
        "exigencias_base": ["extintores_portateis", "hidrantes", "sinalizacao_emergencia",
                           "iluminacao_emergencia", "deteccao_automatica", "alarme_incendio",
                           "spda", "chuveiros_automaticos", "plano_emergencia"],
        "area_extintor_m2": 200,
        "capacidade_extintor": "2-A:20-B:C"
    },
    "G-1": {
        "descricao": "Garagem, estacionamento coberto",
        "exemplos": ["garagem", "estacionamento", "parking", "auto"],
        "exigencias_base": ["extintores_portateis", "hidrantes", "sinalizacao_emergencia",
                           "ventilacao_mecanica"],
        "area_extintor_m2": 250,
        "capacidade_extintor": "2-A:40-B:C"
    },
    "H-1": {
        "descricao": "Deposito, armazem, galpao",
        "exemplos": ["deposito", "armazem", "galpao", "estoque", "almoxarifado"],
        "exigencias_base": ["extintores_portateis", "hidrantes", "sinalizacao_emergencia"],
        "area_extintor_m2": 500,
        "capacidade_extintor": "2-A:10-B:C"
    },
    "I-1": {
        "descricao": "Industrial - baixo risco",
        "exemplos": ["fabrica", "industria", "manufatura", "producao", "montagem"],
        "exigencias_base": ["extintores_portateis", "hidrantes", "sinalizacao_emergencia",
                           "iluminacao_emergencia"],
        "area_extintor_m2": 300,
        "capacidade_extintor": "4-A:20-B:C"
    },
    "I-2": {
        "descricao": "Industrial - medio/alto risco",
        "exemplos": ["quimica", "petroleo", "gas", "combustivel", "tinta", "solvente"],
        "exigencias_base": ["extintores_portateis", "hidrantes", "sinalizacao_emergencia",
                           "iluminacao_emergencia", "chuveiros_automaticos", "spda",
                           "deteccao_automatica", "alarme_incendio"],
        "area_extintor_m2": 150,
        "capacidade_extintor": "4-A:40-B:C"
    },
}

# Mapa simples: grupo NT-01 -> divisao padrao
GRUPO_MAP = {
    "A": "A-2", "B": "B-1", "C": "C-1", "D": "D-1",
    "E": "E-1", "F": "F-1", "G": "G-1", "H": "H-1", "I": "I-1"
}


class ClassificadorOcupacao:
    """
    Classifica edificacoes conforme NT-01/2025 Tabela 1.
    
    Uso:
        clf = ClassificadorOcupacao()
        resultado = clf.classificar("Loja de roupas com 800m2 no centro comercial")
    """
    
    def __init__(self, model_id: str = "cbmgo/llama3-8b-ocupacao", use_llm: bool = False):
        self.model_id = model_id
        self.use_llm = use_llm and HAS_TRANSFORMERS
        self._pipe = None
    
    def _load_llm(self):
        if self._pipe is None and self.use_llm:
            print(f"Carregando modelo: {self.model_id}")
            self._pipe = pipeline(
                "text-generation",
                model=self.model_id,
                device_map="auto",
                max_new_tokens=512,
                temperature=0.1,
            )
    
    def classificar_llm(self, descricao: str) -> dict:
        """Classificacao via LLM fine-tunado."""
        self._load_llm()
        prompt = f"""Classifique a edificacao conforme NT-01/2025 CBMGO Tabela 1.
Responda APENAS em JSON valido.

Edificacao: {descricao}

JSON de resposta:
{{
  "grupo": "letra do grupo (A-I)",
  "divisao": "ex: D-1",
  "descricao_ocupacao": "descricao do tipo",
  "exigencias": ["lista", "de", "sistemas", "obrigatorios"],
  "area_extintor_m2": numero,
  "capacidade_extintor": "ex: 2-A:10-B:C",
  "justificativa": "breve explicacao"
}}"""
        
        saida = self._pipe(prompt)[0]["generated_text"]
        # Extrair JSON da resposta
        match = re.search(r'\{[^{}]*"grupo"[^{}]*\}', saida, re.DOTALL)
        if match:
            try:
                return json.loads(match.group())
            except json.JSONDecodeError:
                pass
        
        # Fallback se LLM falhar
        return self.classificar_regras(descricao)
    
    def classificar_regras(self, descricao: str) -> dict:
        """Classificacao via regras locais (sem GPU)."""
        desc_lower = descricao.lower()
        
        melhor_match = None
        melhor_score = 0
        
        for divisao, info in TABELA_OCUPACOES.items():
            score = 0
            for exemplo in info["exemplos"]:
                if exemplo in desc_lower:
                    score += 2
            # Bonus por palavras parciais
            palavras = desc_lower.split()
            for palavra in palavras:
                for exemplo in info["exemplos"]:
                    if len(palavra) > 3 and (palavra in exemplo or exemplo in palavra):
                        score += 1
            if score > melhor_score:
                melhor_score = score
                melhor_match = divisao
        
        # Fallback para D-1 (comercio geral)
        if melhor_match is None or melhor_score == 0:
            melhor_match = "D-1"
        
        info = TABELA_OCUPACOES[melhor_match]
        grupo = melhor_match[0]
        
        # Calcular exigencias adicionais por area/altura (se mencionado)
        exigencias = list(info["exigencias_base"])
        
        # Detectar area mencionada
        area_match = re.search(r'(\d+)\s*m2', desc_lower)
        area = float(area_match.group(1)) if area_match else 0
        
        altura_match = re.search(r'(\d+)\s*(andar|pavimento|piso|m de altura)', desc_lower)
        altura = float(altura_match.group(1)) if altura_match else 0
        # Converter pavimentos para altura aproximada
        if 'andar' in desc_lower or 'pavimento' in desc_lower:
            altura = altura * 3.0
        
        if area > 1000 and "spda" not in exigencias:
            exigencias.append("spda")
        if altura > 12 and "iluminacao_emergencia" not in exigencias:
            exigencias.append("iluminacao_emergencia")
        if altura > 12 and "saidas_emergencia" not in exigencias:
            exigencias.append("saidas_emergencia")
        
        return {
            "grupo": grupo,
            "divisao": melhor_match,
            "descricao_ocupacao": info["descricao"],
            "exigencias": list(set(exigencias)),
            "area_extintor_m2": info["area_extintor_m2"],
            "capacidade_extintor": info["capacidade_extintor"],
            "justificativa": f"Classificado como {melhor_match} - {info['descricao']} (score={melhor_score})",
            "referencia": "NT-01/2025 CBMGO - Tabela 1"
        }
    
    def classificar(self, descricao: str) -> dict:
        """Interface principal: usa LLM se disponivel, regras locais caso contrario."""
        if self.use_llm and self._pipe is not None:
            return self.classificar_llm(descricao)
        return self.classificar_regras(descricao)
    
    def formatar_resultado(self, resultado: dict) -> str:
        """Formata resultado para exibicao."""
        exig = resultado.get("exigencias", [])
        exig_txt = "\n".join([f"  - {e.replace('_',' ').title()}" for e in exig])
        
        return (
            "=== CLASSIFICACAO NT-01/2025 CBMGO ===\n"
            f"Grupo/Divisao: {resultado.get('divisao','?')} - {resultado.get('descricao_ocupacao','')}\n"
            f"Grupo geral: {resultado.get('grupo','?')}\n\n"
            f"SISTEMAS EXIGIDOS ({len(exig)}):\n{exig_txt}\n\n"
            f"EXTINTORES:\n"
            f"  Area por extintor: {resultado.get('area_extintor_m2','?')} m2\n"
            f"  Capacidade minima: {resultado.get('capacidade_extintor','?')}\n\n"
            f"Justificativa: {resultado.get('justificativa','')}\n"
            f"Ref: {resultado.get('referencia','NT-01/2025 CBMGO')}\n"
        )


# Instancia global
_clf_instance = None

def get_classificador(use_llm: bool = False) -> ClassificadorOcupacao:
    global _clf_instance
    if _clf_instance is None:
        _clf_instance = ClassificadorOcupacao(use_llm=use_llm)
    return _clf_instance


if __name__ == "__main__":
    clf = ClassificadorOcupacao()
    testes = [
        "Loja de roupas com 800m2",
        "Hospital universitario com UTI",
        "Escola municipal de ensino fundamental",
        "Fabrica de produtos quimicos",
        "Edificio residencial de 15 andares",
    ]
    for teste in testes:
        res = clf.classificar(teste)
        print(clf.formatar_resultado(res))
        print("=" * 50)