import pandas as pd
import numpy as np
import json
from pathlib import Path
from scipy.signal import savgol_filter

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

def extrair_rp_media_sg(caminho_arquivo):
    """Extrai a resistência média processada via Savitzky-Golay."""
    try:
        with open(caminho_arquivo, 'r', encoding='utf-8') as f:
            dados = json.load(f)

        drill = np.array(dados['profile']['drill'])
        # Aplica SG para limpar ruído antes de tirar a média
        drill_filtrado = savgol_filter(drill, window_length=31, polyorder=3)
        return np.mean(drill_filtrado)
    except Exception as e:
        return None

# 1. Carregar e Limpar Dados de Laboratório
print("Lendo dados de laboratório...")
abas = pd.read_excel(ARQUIVO_EXCEL, sheet_name=None)
df_lab = pd.concat(abas.values(), ignore_index=True)

# Limpeza de colunas fantasmagóricas e valores nulos
df_lab = df_lab.loc[:, ~df_lab.columns.str.contains('^Unnamed')]
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 arquivos .rgp para obter a Resistência (RP)
print("Processando arquivos .rgp e extraindo Resistência Média...")
lista_rp = []
for reg in df_lab['Registro']:
    nome_arq = f"Measurement{str(int(reg)).zfill(3)}.rgp"
    caminho = PASTA_AMOSTRAS / nome_arq
    lista_rp.append(extrair_rp_media_sg(caminho) if caminho.exists() else None)

df_lab['RP_Media_SG'] = lista_rp
df_final = df_lab.dropna(subset=['RP_Media_SG'])

# 3. MATEMÁTICA DA ESTIMATIVA (Regressão Linear OLS)
# Encontra a equação: DB = m * RP + c
x = df_final['RP_Media_SG'].values
y = df_final['ME.Bas'].values

m, c = np.polyfit(x, y, 1)

# 4. Aplicar a fórmula para gerar a Densidade Estimada
df_final['DB_Estimada'] = (m * df_final['RP_Media_SG']) + c

# Calcular o Erro Residual (Diferença entre Lab e Estimativa)
df_final['Erro_Absoluto'] = (df_final['ME.Bas'] - df_final['DB_Estimada']).abs()

# 5. Salvar Resultados
df_final.to_excel(OUTPUT_EXCEL, index=False)

print("\n" + "="*30)
print("ESTIMATIVA CONCLUÍDA")
print(f"Equação gerada: DB = ({m:.6f} * RP) + ({c:.6f})")
print(f"Erro Médio da Estimativa: {df_final['Erro_Absoluto'].mean():.4f} g/cm³")
print(f"Arquivo salvo em: {OUTPUT_EXCEL}")
print("="*30)