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

__all__ = ['CACHE_FOLDER', 'log', 'custom_theme', 'console', 'filter_spectrum', 'process_bin', '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 fire import Fire
from .utils import *
from .constants import SPECTRAL_BLOCKS
from .parser import *
CACHE_FOLDER = Path.cwd() / ".cache"

# Cell
logging.basicConfig(
    level="NOTSET",
    format="%(message)s",
    datefmt="[%X]",
    handlers=[RichHandler(rich_tracebacks=True)],
)

log = logging.getLogger("rich")

# Cell
custom_theme = Theme({"info": "dim cyan", "warning": "magenta", "danger": "bold red"})
console = Console(theme=custom_theme)

# Cell
def filter_spectrum(
    df: pd.DataFrame,
    time_start: str = None,
    time_stop: str = None,
    freq_start: str = None,
    freq_stop: str = None,
) -> pd.DataFrame:
    """Recebe o arquivo de espectro df e retorna de acordo com os filtros

    Args:
        df (pd.DataFrame): Arquivo de espectro. Timestamp como linhas e frequências como colunas
        time_start (str): Timestamp de início. Se None filtra desde o início do arquivo
        time_stop (str): Timestamp de fim. Se None filtra até o fim do arquivo
        freq_start (str): Filtro inicial de frequência. Se None retorna desde a menor frequências
        freq_stop (str): Filtro Final de frequência. Se None retorna até a maior frequência.

    Returns:
        pd.DataFrame: DataFrame com Frequência, min, max e mean após os filtros aplicados.
    """
    df = df.copy()
    if time_start is None:
        time_start = "01/01/2000"
    if time_stop is None:
        time_stop = "31/12/2100"
    try:
        time_start = pd.to_datetime(time_start)
        time_stop = pd.to_datetime(time_stop)
    except pd.errors.ParserError:
        log.error(
            f"[bold red blink] Datas inválidas! Verifique as strings de data {freq_start} e {freq_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,
            )
            time_start = 0
            time_stop = df.shape[0]

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

    if freq_start is None:
        freq_start = 0
    if freq_stop is None:
        freq_stop = np.inf

    filtered_cols = df.columns[(float(freq_start) <= cols) & (cols <= float(freq_stop))]
    filtered_rows = df.index[(time_start <= rows) & (rows <= time_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}
    )

# Internal Cell
def read_meta(filename):
    ext = filename.suffix
    if ext == ".csv":
        df = pd.read_csv(filename)
    elif ext == ".xlsx":
        df = pd.read_excel(filename, engine="openpyxl")
    elif ext == ".fth":
        df = pd.read_feather(filename)
        if "wallclock_datetime" in df.columns:
            df.set_index("wallclock_datetime", inplace=True)
    else:
        raise ValueError(f"Extension {ext} not implemented")
    return df

# Cell
def process_bin(
    entrada: str,
    saida: str,
    recursivo: bool = False,
    pastas: Iterable[str] = None,
    levels: bool = False,
    substituir: bool = False,
    dtype: str = "float16",
) -> None:
    """Recebe uma pasta ou arquivo bin, processa e salva os metadados e espectro na saida.

    Args:
        entrada (str): Caminho para a Pasta ou Arquivo .bin
        saida (str): Pasta onde salvar os arquivos processados
        recursivo (bool, optional): Buscar os arquivos de entrada recursivamente. Defaults to False.
        pastas (Iterable[str], optional): Limitar a busca a essas pastas. Defaults to None.
        levels (bool, optional): Extrair e salvar os dados de espectro. Defaults to False.
        substituir (bool, optional): Reprocessar arquivos já processados?. Defaults to False.
        dtype (str, optional): Tipo de dados a salvar o espectro. Defaults to "float16".
    """

    entrada = Path(entrada)
    if entrada.is_file():
        lista_bins = [entrada]
    else:
        lista_bins = get_files(
            entrada, extensions=[".bin"], recurse=recursivo, folders=pastas
        )
    parsed_bins = {}
    meta_path = Path(f"{saida}/meta")
    levels_path = Path(f"{saida}/levels")
    meta_path.mkdir(exist_ok=True, parents=True)
    levels_path.mkdir(exist_ok=True, parents=True)
    log_meta = Path(f"{saida}/log_meta.txt")
    log_levels = Path(f"{saida}/log_levels.txt")
    if substituir:
        done_meta = set()
        done_levels = set()
    else:

        done_meta = (
            set(log_meta.read_text().split("\n")) if log_meta.exists() else set()
        )
        done_levels = (
            set(log_levels.read_text().split("\n")) if log_levels.exists() else set()
        )

    console.rule("Lista de Arquivos a serem processados", style="bold red")
    console.print(
        [f.name for f in lista_bins],
        style="bold white",
        overflow="fold",
        justify="left",
    )
    if not lista_bins:
        console.print(":sleeping: Nenhum arquivo .bin a processar :zzz:")
        return

    if not levels:
        lista_bins = [f for f in lista_bins if f.name not in done_meta]
    else:
        lista_bins = [f for f in lista_bins if f.name not in done_levels]

    if not lista_bins:
        console.print(":sleeping: Nenhum arquivo novo a processar :zzz:")
        console.print(
            ":point_up: use --substituir no terminal ou substituir=True na chamada caso queira reprocessar os bins e sobrepôr os arquivos existentes :wink:"
        )
        return

    try:

        with Progress(transient=True, auto_refresh=False) as progress:
            bins = progress.track(
                lista_bins,
                total=len(lista_bins),
                description="[green]Processando Blocos Binários",
            )

            for file in bins:
                progress.console.print(f"[cyan]Processando Blocos de: [red]{file.name}")
                parsed_bins[file.name] = parse_bin(file)
                progress.refresh()

            lista_meta = [(k, v) for k, v in parsed_bins.items() if k not in done_meta]

            if lista_meta:
                blocks = progress.track(
                    lista_meta,
                    total=len(lista_meta),
                    description="[cyan]Exportando Metadados",
                )
                for file, block in blocks:
                    progress.console.print(f"[cyan]Extraindo Metadados de: [red]{file}")
                    export_meta(file, block, meta_path, ext=".fth")
                    done_meta.add(file)
                    progress.refresh()
            if levels:
                lista_levels = lista_meta = [
                    (k, v) for k, v in parsed_bins.items() if k not in done_levels
                ]
                if lista_levels:
                    bins = progress.track(
                        lista_levels,
                        total=len(lista_levels),
                        description="[grey]Exportando Dados de Espectro",
                    )
                    for file, block_obj in bins:
                        progress.console.print(
                            f"[grey]Extraindo Espectro de: [red]{file}"
                        )
                        meta_index = []
                        blocks = block_obj["blocks"]
                        for (tipo, tid) in blocks.keys():
                            if tipo not in SPECTRAL_BLOCKS:
                                continue
                            meta_file = Path(
                                f"{meta_path}/{file}-B_{tipo}_TId_{tid}.fth"
                            )
                            if not meta_file.exists():
                                export_meta(
                                    file,
                                    block_obj,
                                    meta_path,
                                    ext=".fth",
                                )
                                done_meta.add(file)
                            meta_df = read_meta(meta_file)
                            meta_index.append(meta_df.index.tolist())
                        export_level(
                            file,
                            block_obj,
                            levels_path,
                            ext=".fth",
                            index=meta_index,
                            dtype=dtype,
                        )
                        done_levels.add(file)
                        progress.refresh()
        console.print("kbô :satisfied:")
    finally:
        log_meta.write_text("\n".join(sorted(list(done_meta))))
        log_levels.write_text("\n".join(sorted(list(done_levels))))

# Internal Cell
def appended_mean(row: pd.Series)->float:
    """Recebe uma linha do DataFrame e retorna sua média ponderada pela coluna Count

    Args:
        row (pd.Series): Linha do DataFrame

    Returns:
        float: Média Ponderada da linha pela coluna Count
    """
    return (row["Count"] * row["Mean"]).sum() / row["Count"].sum()

# Cell
def extract_bin_stats(
    filename: str,
    time_start: str = None,
    time_stop: str = None,
    freq_start: str = None,
    freq_stop: str = None,
    cache: str = CACHE_FOLDER,
) -> pd.DataFrame:
    """Recebe o caminho para um arquivo CRFS bin e retorna um dataframe com o resumo estatístico dos dados de espectro

    Args:
        filename (str): Caminho para o arquivo bin
        time_start (str): Timestamp de início. Se None filtra desde o início do arquivo
        time_stop (str): Timestamp de fim. Se None filtra até o fim do arquivo
        freq_start (str): Filtro inicial de frequência. Se None retorna desde a menor frequências
        freq_stop (str): Filtro Final de frequência. Se None retorna até a maior frequência.
        cache (str, optional): Caminho para a pasta de cache. Default é criar uma pasta oculta .cache no diretório atual.

    Returns:
        pd.DataFrame: Dataframe contendo o resumo estatístico do arquivo
    """

    cache = Path(cache)
    cache.mkdir(exist_ok=True, parents=True)
    filename = Path(filename)
    while True:
        cached_files = get_files(cache / "levels")
        # TODO filter based on metadata
        cached_levels = cached_files.filter(lambda name: filename.stem in str(name))
        if not cached_levels:
            process_bin(entrada=filename, saida=cache, levels=True)
        else:
            break
    dfs = cached_levels.map(pd.read_feather)
    spectra = dfs.map(filter_spectrum, time_start=time_start, time_stop=time_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