import pandas as pd
import numpy as np
import json
from pathlib import Path
from scipy.signal import savgol_filter
from scipy.stats import pearsonr
import matplotlib.pyplot as plt

# --- CONFIGURAÇÕES DE CAMINHOS ---
BASE_DIR = Path(__file__).parent
PASTA_AMOSTRAS = BASE_DIR / "Gateados 2025" / "Amostra"
ARQUIVO_EXCEL = BASE_DIR / "Gateados 2025" / "BAGUETAS GATEADOS  densidade.xlsx"

def extrair_dados_sg(caminho_arquivo):
    """Lê o .rgp e aplica o Filtro Savitzky-Golay (SG)."""
    try:
        with open(caminho_arquivo, 'r', encoding='utf-8') as f:
            dados = json.load(f)

        drill = np.array(dados['profile']['drill'])
        # Chave 'header' em minúsculo conforme a estrutura real do seu arquivo
        depth = dados['header']['depthMsmt']

        # --- FILTRO SAVITZKY-GOLAY (SG) ---
        # Janela de 31 pontos e polinômio de ordem 3 para preservar picos de densidade
        drill_filtrado = savgol_filter(drill, window_length=31, polyorder=3)

        return {
            "rp_media": np.mean(drill_filtrado),
            "rp_max": np.max(drill_filtrado),
            "profundidade_total": depth
        }
    except Exception as e:
        print(f"Erro no arquivo {caminho_arquivo.name}: {e}")
        return None

# 1. Consolidar Planilhas de Laboratório (Todas as abas/BLOCOS)
print("Lendo planilhas de laboratório (Método Savitzky-Golay)...")
abas = pd.read_excel(ARQUIVO_EXCEL, sheet_name=None)
df_lab = pd.concat(abas.values(), ignore_index=True)

# Garante que os dados essenciais sejam numéricos para evitar erros no merge
df_lab['ME.Bas'] = pd.to_numeric(df_lab['ME.Bas'], errors='coerce')
df_lab['Registro'] = pd.to_numeric(df_lab['Registro'], errors='coerce')
df_lab = df_lab.dropna(subset=['ME.Bas', 'Registro'])

# 2. Processar RGP com Filtro SG
resultados = []
print("Processando arquivos .rgp com Filtro Savitzky-Golay...")

for index, row in df_lab.iterrows():
    # Formatação do nome: Measurement016.rgp
    num_registro = str(int(row['Registro'])).zfill(3)
    nome_arq = f"Measurement{num_registro}.rgp"
    caminho = PASTA_AMOSTRAS / nome_arq

    if caminho.exists():
        dados = extrair_dados_sg(caminho)
        if dados:
            dados['Registro'] = row['Registro']
            resultados.append(dados)

# 3. Cruzamento de Dados e Estatística Final
if resultados:
    df_resisto = pd.DataFrame(resultados)
    df_final = pd.merge(df_lab, df_resisto, on='Registro')

    x = df_final['rp_media']
    y = df_final['ME.Bas']
    r, p = pearsonr(x, y)

    # --- NOMES DE OUTPUT ATUALIZADOS PARA SG ---
    nome_excel_saida = BASE_DIR / "Gateados 2025" / "RESULTADOS_SG.xlsx"
    nome_grafico_saida = BASE_DIR / "Gateados 2025" / "grafico_RESULTADOS_SG.png"

    print(f"\n--- ANÁLISE CONCLUÍDA (SAVITZKY-GOLAY) ---")
    print(f"Correlação (r): {r:.4f}")
    print(f"R² final: {r**2:.4f}")

    # Salvar Planilha Consolidada SG
    df_final.to_excel(nome_excel_saida, index=False)
    print(f"Planilha salva: {nome_excel_saida}")

    # Gerar Gráfico de Dispersão SG
    plt.figure(figsize=(10, 6))
    plt.scatter(x, y, color='teal', alpha=0.6, label='Amostras (Savitzky-Golay)')

    # Linha de tendência linear
    m, b = np.polyfit(x, y, 1)
    plt.plot(x, m*x + b, color='red', linestyle='--', label=f'R² SG = {r**2:.4f}')

    plt.title('Validação Resistógrafo: Filtro Savitzky-Golay (SG)', fontsize=12)
    plt.xlabel('Resistência Média (Processada com SG)', fontsize=10)
    plt.ylabel('Densidade Básica (ME.Bas Laboratório)', fontsize=10)
    plt.legend()
    plt.grid(True, alpha=0.3)

    plt.savefig(nome_grafico_saida)
    print(f"Gráfico salvo: {nome_grafico_saida}")
    plt.show()
else:
    print("Nenhum arquivo .rgp pôde ser processado. Verifique os caminhos no Bazzite.")