#!/usr/bin/env python3
"""
Script para gerar plots espaciais com Grads (AGCM/CSFR)
Autor: Thaisa Lopes
Descrição: Gera scripts .gs e imagens para cada domínio, métrica e modelo
"""

import os
import re
from pathlib import Path

# ============================================================================
# CONFIGURAÇÕES GLOBAIS
# ============================================================================

TIPO = "PRESS"
IDENTIFICADOR_MODELO = "AGCM"  # Mude para "CSFR" quando necessário
DIRETORIO_BASE = f"/mnt/beegfs/thaisa.lopes/SCANTEC/dataout/SMNA_x_NCEP.bkp/{IDENTIFICADOR_MODELO}_{TIPO}/"

TEMPOS = [0, 24, 48, 72, 96]


# ============================================================================
# FUNÇÕES AUXILIARES
# ============================================================================

def extrair_variaveis_do_ctl(caminho_ctl):
    """
    Extrai a lista de variáveis do arquivo .ctl
    """
    variaveis = []
    
    try:
        with open(caminho_ctl, 'r') as f:
            linhas = f.readlines()
        
        dentro_vars = False
        
        for linha in linhas:
            linha = linha.strip()
            
            if not linha:
                continue
            
            if re.search(r'^vars\s+\d+', linha, re.IGNORECASE):
                dentro_vars = True
                continue
            
            if dentro_vars and re.search(r'^endvars', linha, re.IGNORECASE):
                break
            
            if dentro_vars and not linha.startswith('#'):
                partes = linha.split()
                if partes:
                    nome_var = re.sub(r'[^a-zA-Z0-9_]', '', partes[0])
                    if nome_var and nome_var not in variaveis:
                        variaveis.append(nome_var)
        
        if variaveis:
            print(f"     📋 Variáveis encontradas: {len(variaveis)}")
            return variaveis
            
    except Exception as e:
        print(f"     ⚠️ Erro ao ler .ctl: {e}")
    
    return []


def identificar_arquivos_ctl(caminho_dominio):
    """
    Identifica arquivos .ctl no domínio
    """
    arquivos_encontrados = []
    for arquivo in os.listdir(caminho_dominio):
        if arquivo.endswith('.ctl'):
            match = re.search(r'(MEAN|RMSE|VIES)(BAM_NCEP|BAM_SMNA)', arquivo)
            if match:
                caminho_completo = os.path.join(caminho_dominio, arquivo)
                variaveis = extrair_variaveis_do_ctl(caminho_completo)
                
                if variaveis:
                    arquivos_encontrados.append({
                        'arquivo': arquivo,
                        'metrica': match.group(1),
                        'modelo': match.group(2),
                        'caminho_completo': caminho_completo,
                        'variaveis': variaveis
                    })
            else:
                match = re.search(r'(MEAN|RMSE|VIES)(BAM_[A-Z]+)', arquivo)
                if match:
                    caminho_completo = os.path.join(caminho_dominio, arquivo)
                    variaveis = extrair_variaveis_do_ctl(caminho_completo)
                    
                    if variaveis:
                        arquivos_encontrados.append({
                            'arquivo': arquivo,
                            'metrica': match.group(1),
                            'modelo': match.group(2),
                            'caminho_completo': caminho_completo,
                            'variaveis': variaveis
                        })
    return arquivos_encontrados


def gerar_script_gs_completo(dominio, metrica, modelo, ctl_path, output_dir, variaveis):
    """
    Gera script .gs completo para uma combinação de métrica e modelo
    """
    
    nome_ctl = os.path.basename(ctl_path)
    
    # Criar subpasta específica para a métrica
    pasta_metrica = os.path.join(output_dir, metrica)
    Path(pasta_metrica).mkdir(parents=True, exist_ok=True)
    
    linhas = []
    linhas.append(f"* Script Grads para {dominio} - {metrica} - {modelo} ({IDENTIFICADOR_MODELO})")
    linhas.append(f"* Gerado automaticamente")
    linhas.append("")
    linhas.append("'reinit'")
    linhas.append(f"'open {nome_ctl}'")
    linhas.append("")
    
    # Loop sobre tempos
    for idx, tempo in enumerate(TEMPOS):
        t_index = idx + 1
        linhas.append(f"* ========== Tempo: {tempo}h ==========")
        linhas.append(f"'set t {t_index}'")
        linhas.append("")
        
        for var in variaveis:
            # Caminho relativo para a subpasta da métrica
            nome_imagem = f"{IDENTIFICADOR_MODELO}_{dominio}_{metrica}_{modelo}_{var}_{tempo:03d}h.png"
            caminho_imagem = f"{metrica}/{nome_imagem}"
            
            linhas.append(f"* Variavel: {var}")
            linhas.append("'set gxout shaded'")
            linhas.append(f"'d {var}'")
            linhas.append("'set gxout contour'")
            linhas.append(f"'d {var}'")
            linhas.append("'cbarn'")
            linhas.append(f"'draw title {IDENTIFICADOR_MODELO} | {dominio} | {metrica} | {modelo} | {var} | {tempo}h'")
            linhas.append(f"'printim {output_dir}/{caminho_imagem} x1000 y800'")
            linhas.append("'c'")
            linhas.append("")
    
    # 🔥 CRÍTICO: Sair do Grads
    linhas.append("'quit'")
    
    return '\n'.join(linhas)


def gerar_script_execucao_dominio(dominio):
    """
    Gera script bash para executar todos os scripts de um domínio
    """
    
    linhas = [
        "#!/bin/bash",
        f"# Script para executar plots do domínio {dominio} ({IDENTIFICADOR_MODELO}_{TIPO})",
        "",
        f"cd {DIRETORIO_BASE}/{dominio}",
        "",
        "echo \"==========================================\"",
        f"echo \"Processando domínio: {dominio} ({IDENTIFICADOR_MODELO}_{TIPO})\"",
        "echo \"==========================================\"",
        "echo \"\"",
        "",
        "for gs in grads_scripts/*.gs; do",
        "    if [ -f \"$gs\" ]; then",
        "        echo \"Executando: $(basename $gs)\"",
        "        # 🔥 CRÍTICO: Usar -lbc e redirecionar entrada para evitar travamentos",
        "        grads -lbc \"$gs\" < /dev/null",
        "        echo \"Finalizado: $(basename $gs)\"",
        "        echo \"\"",
        "    fi",
        "done",
        "",
        "echo \"==========================================\"",
        f"echo \"✅ {dominio} ({IDENTIFICADOR_MODELO}_{TIPO}) concluído!\"",
        "echo \"==========================================\""
    ]
    
    return '\n'.join(linhas)


def gerar_script_execucao_todos(dominios):
    """
    Gera script bash para executar todos os domínios
    """
    
    linhas = [
        "#!/bin/bash",
        f"# Script para executar todos os plots espaciais ({IDENTIFICADOR_MODELO}_{TIPO})",
        "",
        f"DIR_BASE=\"{DIRETORIO_BASE}\"",
        "",
        "echo \"==========================================\"",
        f"echo \"INICIANDO PLOTS ESPACIAIS - {IDENTIFICADOR_MODELO}_{TIPO}\"",
        "echo \"==========================================\"",
        "echo \"\"",
        ""
    ]
    
    for dominio in dominios:
        linhas.extend([
            f"echo \"\"",
            f"echo \"📁 Processando domínio: {dominio}\"",
            f"echo \"------------------------------------------\"",
            f"cd \"$DIR_BASE/{dominio}\"",
            "",
            f"for gs in grads_scripts/*.gs; do",
            f"    if [ -f \"$gs\" ]; then",
            f"        echo \"  Executando: $(basename $gs)\"",
            f"        grads -lbc \"$gs\" < /dev/null",
            f"        echo \"  Finalizado: $(basename $gs)\"",
            f"    fi",
            f"done",
            "",
            f"echo \"✅ {dominio} ({IDENTIFICADOR_MODELO}_{TIPO}) concluído!\"",
            ""
        ])
    
    linhas.extend([
        "echo \"\"",
        "echo \"==========================================\"",
        f"echo \"✅ PROCESSAMENTO {IDENTIFICADOR_MODELO}_{TIPO} CONCLUÍDO!\"",
        "echo \"==========================================\""
    ])
    
    return '\n'.join(linhas)


# ============================================================================
# FUNÇÃO PRINCIPAL
# ============================================================================

def main():
    """Função principal do script"""
    print("="*70)
    print(f"GERADOR DE PLOTS ESPACIAIS - GRADS ({IDENTIFICADOR_MODELO}_{TIPO})")
    print("="*70)
    print(f"Diretório base: {DIRETORIO_BASE}")
    print(f"Tempos: {TEMPOS}")
    print(f"Total de tempos: {len(TEMPOS)}")
    print()
    
    if not os.path.exists(DIRETORIO_BASE):
        print(f"❌ ERRO: Diretório base não encontrado: {DIRETORIO_BASE}")
        return
    
    dominios = [d for d in os.listdir(DIRETORIO_BASE) 
                if os.path.isdir(os.path.join(DIRETORIO_BASE, d))
                and d not in ['fig', 'fig_espacial', 'grads_scripts', '__pycache__']]
    
    dominios.sort()
    
    print(f"Domínios encontrados: {len(dominios)}")
    print(f"Lista: {dominios}\n")
    
    total_scripts = 0
    dominios_com_arquivos = []
    
    for dominio in dominios:
        print(f"{'='*70}")
        print(f"📁 Processando domínio: {dominio}")
        print(f"{'='*70}")
        
        caminho_dominio = os.path.join(DIRETORIO_BASE, dominio)
        
        pasta_fig = os.path.join(caminho_dominio, 'fig_espacial')
        pasta_scripts = os.path.join(caminho_dominio, 'grads_scripts')
        
        # Criar estrutura de pastas
        Path(pasta_fig).mkdir(parents=True, exist_ok=True)
        Path(pasta_scripts).mkdir(parents=True, exist_ok=True)
        
        # Criar subpastas para cada métrica (apenas MEAN, RMSE, VIES)
        for metrica in ['MEAN', 'RMSE', 'VIES']:
            Path(os.path.join(pasta_fig, metrica)).mkdir(parents=True, exist_ok=True)
        
        arquivos = identificar_arquivos_ctl(caminho_dominio)
        
        if not arquivos:
            print(f"  ⚠️ Nenhum arquivo .ctl válido encontrado")
            continue
        
        dominios_com_arquivos.append(dominio)
        
        print(f"  📄 Arquivos .ctl encontrados: {len(arquivos)}")
        for item in arquivos:
            print(f"     - {item['arquivo']}")
            print(f"       Métrica: {item['metrica']}, Modelo: {item['modelo']}")
            print(f"       Variáveis: {len(item['variaveis'])}")
        
        print(f"\n  🔧 Gerando scripts Grads...")
        
        for item in arquivos:
            if not item['variaveis']:
                print(f"     ⚠️ Nenhuma variável encontrada para {item['arquivo']}, pulando...")
                continue
            
            print(f"     Gerando: {len(item['variaveis'])} variáveis × {len(TEMPOS)} tempos = {len(item['variaveis']) * len(TEMPOS)} imagens")
            print(f"     📁 Salvando em: fig_espacial/{item['metrica']}/")
            
            script = gerar_script_gs_completo(
                dominio,
                item['metrica'],
                item['modelo'],
                item['caminho_completo'],
                'fig_espacial',
                item['variaveis']
            )
            
            nome_script = f"{IDENTIFICADOR_MODELO}_{dominio}_{item['metrica']}_{item['modelo']}.gs"
            caminho_script = os.path.join(pasta_scripts, nome_script)
            
            with open(caminho_script, 'w') as f:
                f.write(script)
            
            print(f"     ✅ {nome_script}")
            total_scripts += 1
        
        script_exec = gerar_script_execucao_dominio(dominio)
        caminho_exec = os.path.join(caminho_dominio, f"executar_{IDENTIFICADOR_MODELO}_{dominio}.sh")
        
        with open(caminho_exec, 'w') as f:
            f.write(script_exec)
        
        os.chmod(caminho_exec, 0o755)
        
        print(f"\n  🚀 Script de execução: executar_{IDENTIFICADOR_MODELO}_{dominio}.sh")
        print()
    
    if dominios_com_arquivos:
        script_mestre = gerar_script_execucao_todos(dominios_com_arquivos)
        caminho_mestre = os.path.join(DIRETORIO_BASE, f"executar_todos_{IDENTIFICADOR_MODELO}.sh")
        
        with open(caminho_mestre, 'w') as f:
            f.write(script_mestre)
        
        os.chmod(caminho_mestre, 0o755)
        
        print("="*70)
        print("📊 RESUMO FINAL")
        print("="*70)
        print(f"✅ Modelo: {IDENTIFICADOR_MODELO}")
        print(f"✅ Tipo: {TIPO}")
        print(f"✅ Domínios processados: {len(dominios_com_arquivos)}")
        print(f"✅ Scripts .gs gerados: {total_scripts}")
        print()
        print("📁 Estrutura de pastas criada:")
        print("   fig_espacial/")
        print("   ├── MEAN/")
        print("   ├── RMSE/")
        print("   └── VIES/")
        print()
        print("🚀 Para executar TUDO de uma vez:")
        print(f"  cd {DIRETORIO_BASE}")
        print(f"  ./executar_todos_{IDENTIFICADOR_MODELO}.sh")
        print()
        print("🚀 Ou executar um domínio específico:")
        print(f"  cd {DIRETORIO_BASE}/AFRIC")
        print(f"  ./executar_{IDENTIFICADOR_MODELO}_AFRIC.sh")
        print("="*70)
    
    print(f"\n✅ Processamento {IDENTIFICADOR_MODELO}_{TIPO} finalizado!")


if __name__ == "__main__":
    main()