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'])
        depth = dados['header']['depthMsmt']

        # Filtro SG para preservação de 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
print("Lendo planilhas de laboratório (SG)...")
abas = pd.read_excel(ARQUIVO_EXCEL, sheet_name=None)
df_lab = pd.concat(abas.values(), ignore_index=True)

# --- LIMPEZA DE COLUNAS VAZIAS (UNNAMED) ---
# Remove colunas que começam com 'Unnamed' ou que são totalmente Nulas
df_lab = df_lab.loc[:, ~df_lab.columns.str.contains('^Unnamed')]
df_lab = df_lab.dropna(axis=1, how='all')

# Garante tipos numéricos para o 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
resultados = []
print("Processando arquivos .rgp...")

for index, row in df_lab.iterrows():
    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 e Saída
if resultados:
    df_resisto = pd.DataFrame(resultados)
    df_final = pd.merge(df_lab, df_resisto, on='Registro')

    # Remove qualquer coluna fantasma que possa ter surgido no merge
    df_final = df_final.loc[:, ~df_final.columns.str.contains('^Unnamed')]

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

    # Nomes de output
    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 (SG) ---")
    print(f"R² final: {r**2:.4f}")

    # Salvar sem as colunas vazias
    df_final.to_excel(nome_excel_saida, index=False)
    print(f"Planilha limpa salva: {nome_excel_saida}")

    # Gráfico
    plt.figure(figsize=(10, 6))
    plt.scatter(x, y, color='teal', alpha=0.6, label='Amostras (SG)')
    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)')
    plt.xlabel('Resistência Média (SG)')
    plt.ylabel('Densidade Básica (Lab)')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.savefig(nome_grafico_saida)
    plt.show()