import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
import json
import os
from pathlib import Path
from scipy.signal import savgol_filter

# --- CONFIGURAÇÃO DE CAMINHOS (Respeitando sua estrutura) ---
BASE_DIR = Path(__file__).parent / "Gateados 2025"
PASTA_AMOSTRA = BASE_DIR / "Amostra"
PASTA_OUTPUT_GIFS = BASE_DIR / "Gifs_SG"
PASTA_OUTPUT_GIFS.mkdir(parents=True, exist_ok=True)

# Rodaremos em apenas uma amostra conforme solicitado
ARQUIVO_RGP = PASTA_AMOSTRA / "Measurement016.rgp"

if not ARQUIVO_RGP.exists():
    print(f"❌ Erro: Arquivo {ARQUIVO_RGP.name} não encontrado.")
    exit()

# 1. Carregar Dados Brutos (Sinal Completo)
with open(ARQUIVO_RGP, 'r') as f:
    dados = json.load(f)
y_bruto = np.array(dados['profile']['drill'])
# Limitamos a 800 pontos para o GIF não ficar gigante
y_bruto = y_bruto[:800]
x_bruto = np.arange(len(y_bruto)) * 0.1 # mm

# 2. Configurações dos Processamentos
window_len_sg = 31 # Janela do filtro
poly_order_sg = 3   # Ordem do polinômio

# --- ALGORITMO NOVO (Adaptativo: Detecção de Borda) ---
# Usamos uma janela de estabilidade de 20 pontos
window_stabilidade = 20
# Threshold dinâmico: 15% da força máxima dessa peça específica (não hard-coded)
max_res = np.percentile(y_bruto, 98)
threshold_adaptativo = max_res * 0.15

idx_inicio = 0
for i in range(len(y_bruto) - window_stabilidade):
    if np.mean(y_bruto[i:i+window_stabilidade]) > threshold_adaptativo:
        idx_inicio = i
        break

y_adaptativo = y_bruto[idx_inicio:]
x_adaptativo = (np.arange(len(y_adaptativo)) + idx_inicio) * 0.1 # mm

# Aplicar SG apenas no Lenho Limpo (Rastro Novo)
y_sg_adaptativo = savgol_filter(y_adaptativo, window_len_sg, poly_order_sg)

# --- ALGORITMO ANTIGO (Estático/Hard-coded: Pula 5 pontos) ---
y_estatico = y_bruto[5:]
x_estatico = (np.arange(len(y_estatico)) + 5) * 0.1 # mm
# Aplicar SG no Sinal com possível Casca (Rastro Antigo)
y_sg_estatico = savgol_filter(y_estatico, window_len_sg, poly_order_sg)

# 3. Preparar Figura para Animação
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8), sharex=True)
plt.subplots_adjust(hspace=0.3)

# Configurar Gráfico Superior (Antigo)
ax1.set_xlim(0, x_bruto[-1])
ax1.set_ylim(min(y_bruto)-5, max(y_bruto)+5)
ax1.set_title(f'Processamento Antigo (Estático - SG no Sinal Completo)\nAmostra: {ARQUIVO_RGP.name}', fontsize=12)
ax1.set_ylabel('Resistência')
line_raw_est, = ax1.plot(x_bruto, y_bruto, color='silver', alpha=0.3, label='Bruto')
line_sg_est, = ax1.plot([], [], color='red', lw=2.5, label='Rastro SG (Com Casca)')
# Janela Móvel do SG (fixa em 5 pontos)
rect_est = plt.Rectangle((5*0.1, 0), window_len_sg*0.1, 100, color='red', alpha=0.1)
ax1.add_patch(rect_est)
ax1.legend(loc='upper right')

# Configurar Gráfico Inferior (Novo)
ax2.set_xlim(0, x_bruto[-1])
ax2.set_ylim(min(y_bruto)-5, max(y_bruto)+5)
ax2.set_title(f'Processamento Novo (Adaptativo - Detecção de Borda & Trimming da Casca)\nThreshold Dinâmico (15% Max): {threshold_adaptativo:.1f}', fontsize=12)
ax2.set_xlabel('Profundidade (mm)')
ax2.set_ylabel('Resistência')
line_raw_ad, = ax2.plot(x_adaptativo, y_adaptativo, color='silver', alpha=0.3, label='Lenho Limpo')
line_sg_ad, = ax2.plot([], [], color='purple', lw=2.5, label='Rastro SG (Sem Casca)')
# Janela Móvel do SG (começa no idx_inicio adaptativo)
rect_ad = plt.Rectangle((idx_inicio*0.1, 0), window_len_sg*0.1, 100, color='purple', alpha=0.1)
ax2.add_patch(rect_ad)
# Linha vertical indicando onde o algoritmo dinâmico cortou a casca
line_corte = ax2.axvline(idx_inicio*0.1, color='green', linestyle='--', label=f'Corte Casca ({idx_inicio*0.1:.1f}mm)')
ax2.legend(loc='upper right')

# Preparar listas para armazenar os rastros das animações
x_hist_est, y_hist_est = [], []
x_hist_ad, y_hist_ad = [], []

def update(i):
    # O filtro SG precisa de meia janela de cada lado para calcular o centro
    half = window_len_sg // 2

    # --- Atualização Antiga (Cima - Vermelho) ---
    # Começa no ponto 5 e anda frame a frame
    i_est = i + 5
    if half <= i < len(y_sg_estatico) - half:
        x_hist_est.append(x_estatico[i])
        y_hist_est.append(y_sg_estatico[i])
        line_sg_est.set_data(x_hist_est, y_hist_est)
        # Move o retângulo móvel
        rect_est.set_xy((x_estatico[i - half], ax1.get_ylim()[0]))

    # --- Atualização Nova (Baixo - Roxo) ---
    # Começa no idx_inicio dinâmico e anda frame a frame
    i_ad = i
    if half <= i_ad < len(y_sg_adaptativo) - half:
        x_hist_ad.append(x_adaptativo[i_ad])
        y_hist_ad.append(y_sg_adaptativo[i_ad])
        line_sg_ad.set_data(x_hist_ad, y_hist_ad)
        # Move o retângulo móvel
        rect_ad.set_xy((x_adaptativo[i_ad - half], ax2.get_ylim()[0]))

    return line_sg_est, rect_est, line_sg_ad, rect_ad

# 4. Gerar e Salvar o GIF Comparativo
frames_totais = max(len(y_sg_estatico), len(y_sg_adaptativo))
# Pulamos de 3 em 3 pontos para a animação ser mais rápida
ani = FuncAnimation(fig, update, frames=range(0, frames_totais, 3), blit=True)

nome_saida = PASTA_OUTPUT_GIFS / f"comparativo_antes_depois_SG_{ARQUIVO_RGP.stem}.gif"
print(f"🎬 Iniciando renderização da animação comparativa... Isso pode demorar.")
# fps=20 para visualização fluida
writer = PillowWriter(fps=20)
ani.save(nome_saida, writer=writer)
plt.close()

print(f"✅ SUCESSO: GIF comparativo salvo em {nome_saida}")