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

# --- CONFIGURAÇÕES ---
BASE_DIR = Path(__file__).parent
ARQUIVO_RGP = BASE_DIR / "Gateados 2025" / "Amostra" / "Measurement016.rgp"

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

# 2. Configurações do Filtro
window_len = 51
poly_order = 3

# 3. Preparar a Figura
fig, ax = plt.subplots(figsize=(10, 6))
ax.set_xlim(0, x[-1])
ax.set_ylim(min(y_bruto)-0.1, max(y_bruto)+0.1)
line_raw, = ax.plot(x, y_bruto, color='silver', alpha=0.4, label='Sinal Bruto (Ruído)')
line_sg, = ax.plot([], [], color='purple', lw=2, label='Valores Suavizados (SG)')
line_poly, = ax.plot([], [], color='red', lw=1.5, label='Polinômio Local')
rect = plt.Rectangle((0, 0), window_len*0.1, max(y_bruto)+1, color='red', alpha=0.1)
ax.add_patch(rect)

ax.set_title(f'Animação Filtro Savitzky-Golay - {ARQUIVO_RGP.name}')
ax.set_xlabel('Profundidade (mm)')
ax.set_ylabel('Resistência')
ax.legend(loc='upper right')

# Função de Animação
y_smooth = []
x_smooth = []

def update(i):
    if i < window_len // 2 or i > len(y_bruto) - window_len // 2:
        return line_sg, line_poly, rect

    # Define a janela atual
    start = i - window_len // 2
    end = i + window_len // 2
    window_x = x[start:end]
    window_y = y_bruto[start:end]

    # Ajusta polinômio local
    poly_coeffs = np.polyfit(window_x, window_y, poly_order)
    poly_y = np.polyval(poly_coeffs, window_x)

    # Valor suavizado é o centro do polinômio
    current_smooth_val = poly_y[window_len // 2]
    y_smooth.append(current_smooth_val)
    x_smooth.append(x[i])

    # Atualiza elementos gráficos
    line_poly.set_data(window_x, poly_y)
    line_sg.set_data(x_smooth, y_smooth)
    rect.set_xy((window_x[0], ax.get_ylim()[0]))

    return line_sg, line_poly, rect

# 4. Gerar e Salvar o GIF
print("Gerando GIF... Isso pode levar um minuto.")
ani = FuncAnimation(fig, update, frames=range(len(y_bruto)), interval=20, blit=True)
writer = PillowWriter(fps=30)
ani.save(BASE_DIR / "animacao_sg_processamento.gif", writer=writer)
plt.show()

print("GIF salvo com sucesso!")