############################################
# ENANDES PROJECT
# November/2024
# Dividir dados filtrados pkl em arquivos anuais
############################################

from functions_produtividade import *

# filtered
input_pkl_path = '/home/jovyan/Desktop/enandes_v2/tile012010/resultfiltered_data_p2.pkl'
output_directory = '/home/jovyan/Desktop/enandes_v2/modelo_produtividade/dados_ndvi_pkl/p2/'
var_name = 'filtered'

start_year = 2001
end_year = 2024
items_per_year = 23


# Carregar os dados do arquivo PKL original
try:
    print(f"Carregando dados de '{input_pkl_path}'...")
    with open(input_pkl_path, 'rb') as f_in:
        all_data = pickle.load(f_in)
    print("Dados carregados com sucesso.")

    # Verificação básica do tipo de dado (útil para depuração)
    print(f"Tipo de dado carregado: {type(all_data)}")
    if isinstance(all_data, (list, tuple)):
         print(f"Número total de itens carregados: {len(all_data)}")
         # Validação opcional do número de itens
         expected_items = (end_year - start_year + 1) * items_per_year
         if len(all_data) != expected_items:
             print(f"*** Atenção: O número de itens ({len(all_data)}) não corresponde exatamente ao esperado ({expected_items}).")
             print(f"*** Verifique 'items_per_year' ou a estrutura do arquivo. O script continuará, mas os últimos arquivos podem ter menos itens.")
    elif hasattr(all_data, 'shape'): # Para arrays Numpy
        print(f"Shape do array carregado: {all_data.shape}")
        # Validação opcional do número de itens (assumindo que a primeira dimensão é o tempo)
        expected_items = (end_year - start_year + 1) * items_per_year
        if all_data.shape[0] != expected_items:
             print(f"*** Atenção: O número de itens na primeira dimensão ({all_data.shape[0]}) não corresponde exatamente ao esperado ({expected_items}).")
             print(f"*** Verifique 'items_per_year' ou a estrutura do arquivo. O script continuará.")
    else:
        print("Tipo de dado não é lista/tupla ou array Numpy. A lógica de fatiamento pode precisar de ajustes.")


except FileNotFoundError:
    print(f"Erro: Arquivo de entrada '{input_pkl_path}' não encontrado.")
    exit()
except MemoryError:
    print(f"Erro: Memória insuficiente para carregar o arquivo '{input_pkl_path}'.")
    print("Considere abordagens alternativas se o arquivo for muito grande.")
    exit()

# Iterar pelos anos e salvar os arquivos separados
print("\nIniciando a separação e salvamento por ano...")
total_items_processed = 0
for year in range(start_year, end_year + 1):
    # Calcular os índices de início e fim para os dados deste ano
    start_index = (year - start_year) * items_per_year
    end_index = start_index + items_per_year

    # Garantir que não tentemos acessar índices fora dos limites dos dados carregados
    # Isso é importante se o número total de itens não for um múltiplo exato de items_per_year * anos
    actual_end_index = min(end_index, len(all_data) if isinstance(all_data, (list, tuple)) else all_data.shape[0])

    if start_index >= actual_end_index:
        print(f"Aviso: Não há mais dados disponíveis para processar o ano {year}. Interrompendo.")
        break # Para de processar se não houver mais dados

    # Extrair os dados para o ano atual
    try:
        year_data = all_data[start_index:actual_end_index]
        num_items_this_year = len(year_data) if isinstance(year_data, (list, tuple)) else year_data.shape[0]

        if num_items_this_year == 0:
            print(f"Aviso: Nenhum dado encontrado para o ano {year} no intervalo [{start_index}:{actual_end_index}]. Pulando.")
            continue

        # Construir o nome do arquivo de saída
        output_filename = os.path.join(output_directory, f'{var_name}_{year}.pkl')

        # Salvar os dados do ano em um novo arquivo PKL
        print(f"Salvando dados de {year} ({num_items_this_year} itens) em '{output_filename}'...")
        with open(output_filename, 'wb') as f_out:
            pickle.dump(year_data, f_out)
        total_items_processed += num_items_this_year

    except IndexError:
         print(f"Erro de índice ao tentar extrair dados para o ano {year} (intervalo {start_index}:{actual_end_index}). Verifique a estrutura dos dados.")
         continue # Pula para o próximo ano
    except Exception as e:
        print(f"Erro ao processar ou salvar dados para o ano {year}: {e}")
        continue # Pula para o próximo ano

print(f"\nProcessamento concluído!")
print(f"Total de itens processados e distribuídos nos arquivos anuais: {total_items_processed}")
print(f"Verifique o diretório '{output_directory}' para os arquivos gerados.")