import numpy as np
import json
import matplotlib.pyplot as plt
from pathlib import Path
from scipy.signal import savgol_filter, find_peaks

# --- CONFIGURAÇÕES DE CAMINHOS ---
BASE_DIR = Path(__file__).parent / "Gateados 2025"
PASTA_AMOSTRAS = BASE_DIR / "Amostra"
# Nova pasta onde as imagens serão salvas
PASTA_GRAFICOS = BASE_DIR / "Graficos_Savgol_Processados"

def processar_todas_amostras():
    """Varre a pasta de amostras e processa todos os arquivos .rgp encontrados."""

    # Cria a pasta de destino se ela não existir
    PASTA_GRAFICOS.mkdir(parents=True, exist_ok=True)

    # Lista todos os arquivos com extensão .rgp
    arquivos_rgp = list(PASTA_AMOSTRAS.glob("*.rgp"))

    if not arquivos_rgp:
        print(f"Nenhum arquivo .rgp foi encontrado na pasta: {PASTA_AMOSTRAS}")
        return

    print(f"Iniciando o processamento em lote de {len(arquivos_rgp)} amostras...")
    print("-" * 40)

    for caminho_arquivo in arquivos_rgp:
        gerar_e_salvar_grafico(caminho_arquivo)

    print("-" * 40)
    print(f"Processamento finalizado! Gráficos salvos em: {PASTA_GRAFICOS}")

def gerar_e_salvar_grafico(caminho):
    """Lê um arquivo .rgp, aplica o filtro SG e salva a imagem do gráfico."""
    nome_arq = caminho.name
    # Extrai apenas o nome sem a extensão para usar no título e no salvamento
    nome_base = caminho.stem

    try:
        with open(caminho, 'r', encoding='utf-8') as f:
            dados = json.load(f)

        drill = np.array(dados['profile']['drill'])
        x = np.arange(len(drill))

        # Parâmetros SG
        w_suave, w_deriv = 51, 101
        drill_sg = savgol_filter(drill, window_length=w_suave, polyorder=3)
        delta_rp = savgol_filter(drill, window_length=w_deriv, polyorder=3, deriv=1)

        # Margem de Segurança
        margem = max(w_suave, w_deriv) // 2
        x_l = x[margem:-margem]
        bruto_l = drill[margem:-margem]
        suave_l = drill_sg[margem:-margem]
        delta_l = delta_rp[margem:-margem]

        # --- DETECÇÃO DE PICOS (X) ---
        indices_picos, _ = find_peaks(delta_l, distance=50, prominence=0.05)
        total_x = len(indices_picos)

        # Zonas Adaptativas
        r_min, r_max = suave_l.min(), suave_l.max()
        intervalo = r_max - r_min
        z1, z2 = r_min + (intervalo * 0.33), r_min + (intervalo * 0.66)

        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(16, 11), sharex=True, constrained_layout=True)

        # --- GRÁFICO 1: PERFIL DE RESISTÊNCIA ---
        ax1.axhspan(z2, r_max, color='green', alpha=0.09, label='Zona Alta Densidade')
        ax1.axhspan(z1, z2, color='yellow', alpha=0.09, label='Zona Média')
        ax1.axhspan(r_min, z1, color='red', alpha=0.09, label='Zona Baixa Densidade')

        ax1.plot(x_l, bruto_l, color='gray', alpha=0.6, linewidth=0.5, label='Leitura Bruta')
        ax1.plot(x_l, suave_l, color='#1a252f', linewidth=2.5, label='Perfil Suavizado (SG)')

        ax1.set_title(f'Perfil de Resistência - {nome_base}', fontsize=15, fontweight='bold', pad=30)
        ax1.set_ylabel('Resistência à Perfuração (RP)', fontsize=11)

        ax1.set_xlabel('Profundidade da Perfuração (Pontos)', fontsize=11)
        ax1.tick_params(labelbottom=True)

        ax1.grid(True, which='both', linestyle='--', alpha=0.15)
        ax1.legend(loc='lower left', bbox_to_anchor=(0, 1.02), ncol=5, frameon=False, fontsize=10)

        # --- GRÁFICO 2: VARIAÇÃO DE RESISTÊNCIA (ANÉIS DE CRESCIMENTO) ---
        ax2.fill_between(x_l, 0, delta_l, where=(delta_l >= 0), color='#27ae60', alpha=0.35, label='Início LT')
        ax2.fill_between(x_l, 0, delta_l, where=(delta_l < 0), color='#c0392b', alpha=0.35, label='Início LI')

        ax2.scatter(x_l[indices_picos], delta_l[indices_picos], color='darkred', marker='x', s=50, linewidths=1.5,
                    label=f'Picos de Transição ({total_x})')

        ax2.plot(x_l, delta_l, color='black', linewidth=1, alpha=0.6)

        ax2.set_title('Variação de Resistência (Anéis de Crescimento)', fontsize=13, fontweight='bold', pad=30)
        ax2.set_xlabel('Profundidade da Perfuração (Pontos)', fontsize=11)
        ax2.set_ylabel('ΔRP', fontsize=11)

        ax2.axhline(0, color='black', linewidth=1.5)
        ax2.grid(True, alpha=0.15)
        ax2.legend(loc='lower left', bbox_to_anchor=(0, 1.02), ncol=4, frameon=False, fontsize=10)

        # --- SALVAMENTO ---
        caminho_salvar = PASTA_GRAFICOS / f"{nome_base}_analise.png"

        # dpi=150 gera uma imagem de excelente qualidade para relatórios (~2400x1650 pixels)
        plt.savefig(caminho_salvar, dpi=150, bbox_inches='tight')

        # O plt.close(fig) é obrigatório para não estourar a memória RAM rodando dezenas de arquivos
        plt.close(fig)

        print(f"[OK] {nome_arq} -> {total_x} picos detectados.")

    except Exception as e:
        # Bloco try/except garante que se um arquivo .rgp estiver corrompido, o loop não para.
        print(f"[ERRO] Falha ao processar {nome_arq}: {e}")

if __name__ == "__main__":
    processar_todas_amostras()