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

CWD = Path.cwd()
PASTA_AMOSTRA = CWD / "Gateados 2025" / "Amostra"
PASTA_OUTPUT_GIFS = CWD / "Gateados 2025" / "Gifs_SMA"
PASTA_OUTPUT_GIFS.mkdir(parents=True, exist_ok=True)

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

        y_bruto = np.array(dados['profile']['drill'])[:600]
        x = np.arange(len(y_bruto)) * 0.1

        window_len = 31
        half = window_len // 2

        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_sma, = ax.plot([], [], color='orange', lw=2.5, label='SMA (Média Móvel)')
        line_mean, = ax.plot([], [], color='black', lw=2, label='Média da Janela')
        rect = plt.Rectangle((0, 0), window_len*0.1, 100, color='orange', alpha=0.1)
        ax.add_patch(rect)

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

        y_sma_list, x_sma_list = [], []

        def update(i):
            if i < half or i >= len(y_bruto) - half:
                return line_sma, line_mean, rect

            # Fatiamento seguro garantindo tamanhos iguais
            idx_start = i - half
            idx_end = i + half + 1 # +1 para garantir que pegue o número correto de pontos

            win_x = x[idx_start:idx_end]
            win_y = y_bruto[idx_start:idx_end]

            # Verifica se os tamanhos batem antes de plotar
            if len(win_x) != len(win_y):
                return line_sma, line_mean, rect

            current_mean = np.mean(win_y)
            y_sma_list.append(current_mean)
            x_sma_list.append(x[i])

            line_mean.set_data(win_x, [current_mean] * len(win_x))
            line_sma.set_data(x_sma_list, y_sma_list)
            rect.set_xy((win_x[0], ax.get_ylim()[0]))
            return line_sma, line_mean, rect

        ani = FuncAnimation(fig, update, frames=range(0, len(y_bruto), 3), blit=True)
        nome_gif = PASTA_OUTPUT_GIFS / f"animacao_SMA_{caminho_rgp.stem}.gif"
        ani.save(nome_gif, writer=PillowWriter(fps=20))
        plt.close(fig)
        print(f"✓ Gerado SMA: {nome_gif.name}")

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

arquivos_rgp = sorted(list(PASTA_AMOSTRA.glob("*.rgp")))
for arq in arquivos_rgp:
    gerar_gif_sma(arq)