# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/05_filter.ipynb (unless otherwise specified).

__all__ = ['CACHE_FOLDER', 'filter_spectrum', 'appended_mean', 'extract_bin_stats']

# Cell
from datetime import datetime
from typing import *
import os
import logging
from fastcore.xtras import Path
from fastcore.script import call_parse, Param, store_true
import numpy as np
import pandas as pd
from rich.progress import Progress
from rich.console import Console
from rich.theme import Theme
from rich.logging import RichHandler
from rfpy.utils import *
from rfpy.constants import SPECTRAL_BLOCKS
from rfpy.parser import *
from rfpy.main import process_bin
CACHE_FOLDER = Path(r"C:\Users\rsilva\Downloads\saida")
if not CACHE_FOLDER.exists():
    CACHE_FOLDER = Path.cwd() / 'cache'
    CACHE_FOLDER.mkdir(exist_ok=True, parents=True)

# Cell
def filter_spectrum(df, start, stop, freq_start, freq_stop):
    df = df.copy()
    try:
        start = pd.to_datetime(start)
        stop = pd.to_datetime(stop)
    except pd.errors.ParserError:
        log.error(f"[bold red blink] Datas inválidas! Verifique as strings de data {start} e {stop}")

    try:
        df.set_index('index', inplace=True)
        df.index = pd.to_datetime(df.index)
    except pd.errors.KeyError:
        if not isinstance(df.index, pd.DatetimeIndex):
            log.warning(
                f"Não foi passado uma coluna ou índice com datetime a ser filtrado, todas as linhas serão processadas",
                exc_info=True
            )
            start = 0
            stop = df.shape[0]

    cols = df.columns.values.astype('float')
    rows = df.index.values

    filtered_cols = df.columns[(float(freq_start) <= cols) & (cols <= float(freq_stop))]
    filtered_rows = df.index[(start <= rows) & (rows <= stop)]
    if len(filtered_cols) == 0 or len(filtered_rows) == 0:
        return None
    count = filtered_rows.shape[0]
    array = df.loc[filtered_rows, filtered_cols].values
    freq = filtered_cols.values.astype('float32')
    min_ = array.min(axis=0)
    max_ = array.max(axis=0)
    mean = array.mean(axis=0)
    return pd.DataFrame({'Frequency': freq, 'Min': min_, 'Max': max_, 'Mean': mean, 'Count': count})

# Cell
def appended_mean(row):
    return (row['Count'] * row['Mean']).sum() / row['Count'].sum()

@call_parse
def extract_bin_stats(filename: Param("Caminho para o arquivo .bin", str),
                      start: Param("Timestamp do Início", str),
                      stop: Param("Timestamp do Fim", str),
                      freq_start: Param("Frequência Inicial (MHz)", str),
                      freq_stop: Param("Frequência Final (MHz)", str)):

    filename = Path(filename)
    while True:
        cached_files = get_files(CACHE_FOLDER / 'levels')
        #TODO filter based on metadata
        cached_levels = cached_files.filter(lambda name: filename.stem in str(name))
        if not cached_levels:
            process_bin(filename, CACHE_FOLDER, levels=True)
        else:
            break
    dfs = cached_levels.map(pd.read_feather)
    spectra = dfs.map(filter_spectrum, start=start, stop=stop, freq_start=freq_start, freq_stop=freq_stop)
    spectra = [s for s in spectra if s is not None]
    out = pd.DataFrame(columns=['Frequency', 'Min', 'Max', 'Mean'])
    if not spectra:
        log.warning(
                f"Os parâmetros repassados não correspondem a nenhum dado espectral do arquivo",
                exc_info=True
            )
        return out
    spectra = pd.concat(spectra)
    gb  = spectra.groupby('Frequency')
    out['Frequency'] = spectra.Frequency.unique()
    out['Min'] = gb.min()['Min'].values
    out['Max'] = gb.max()['Max'].values
    out['Mean'] = gb.apply(appended_mean).values
    return out