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

# --- CONFIGURAÇÕES DE CAMINHOS ---
CWD = Path.cwd()
# Ajuste conforme sua estrutura no Bazzite
PASTA_AMOSTRA = CWD / "Gateados 2025" / "Amostra"
PASTA_OUTPUT_GIFS = CWD / "Gateados 2025" / "Gifs_SG"
PASTA_OUTPUT_GIFS.mkdir(parents=True, exist_ok=True)

def gerar_gif_sg(caminho_rgp):
    try:
        with open(caminho_rgp, 'r') as f:
            dados = json.load(f)

        # 600 pontos para manter o arquivo leve e didático
        y_bruto = np.array(dados['profile']['drill'])[:600]
        x = np.arange(len(y_bruto)) * 0.1

        window_len = 51  # Janela maior para SG ser mais visível
        poly_order = 3

        fig, ax = plt.subplots(figsize=(10, 6))
        ax.set_xlim(0, x[-1])
        ax.set_ylim(min(y_bruto)-5, max(y_bruto)+5)

        line_raw, = ax.plot(x, y_bruto, color='silver', alpha=0.4, label='Sinal Bruto')
        line_sg, = ax.plot([], [], color='purple', lw=2.5, label='SG (Suavizado)')
        line_poly, = ax.plot([], [], color='red', lw=1.5, label='Polinômio Local')
        rect = plt.Rectangle((0, 0), window_len*0.1, 100, color='red', alpha=0.1)
        ax.add_patch(rect)

        ax.set_title(f'Processamento Savitzky-Golay: {caminho_rgp.name}')
        ax.legend(loc='upper right')

        y_sg_list, x_sg_list = [], []

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

            # Janela de dados
            win_x = x[i - half : i + half]
            win_y = y_bruto[i - half : i + half]

            # Ajuste Polinomial
            coeffs = np.polyfit(win_x, win_y, poly_order)
            poly_y = np.polyval(coeffs, win_x)

            y_sg_list.append(poly_y[half])
            x_sg_list.append(x[i])

            line_poly.set_data(win_x, poly_y)
            line_sg.set_data(x_sg_list, y_sg_list)
            rect.set_xy((win_x[0], ax.get_ylim()[0]))
            return line_sg, line_poly, rect

        # Frames pulando de 2 em 2 para velocidade de renderização
        ani = FuncAnimation(fig, update, frames=range(0, len(y_bruto), 2), blit=True)
        nome_gif = PASTA_OUTPUT_GIFS / f"animacao_SG_{caminho_rgp.stem}.gif"

        ani.save(nome_gif, writer=PillowWriter(fps=20))
        plt.close(fig)
        print(f"✓ Gerado SG: {nome_gif.name}")

    except Exception as e:
        print(f"✗ Erro em {caminho_rgp.name}: {e}")

# --- EXECUÇÃO ---
arquivos_rgp = list(PASTA_AMOSTRA.glob("*.rgp"))
print(f"Iniciando processamento SG de {len(arquivos_rgp)} arquivos...")

for arq in arquivos_rgp:
    gerar_gif_sg(arq)

print(f"\nConcluído! GIFs SG em: {PASTA_OUTPUT_GIFS}")