# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/metrics.ipynb.

# %% auto 0
__all__ = ['rounder', 'common_format_dict', 'denorm_mae', 'common', 'sizeof_fmt']

# %% ../nbs/metrics.ipynb 1
from typing import Callable

import numpy as np
from sklearn.metrics import mean_absolute_error
import pandas as pd

import enilm
import enilm.norm

# %% ../nbs/metrics.ipynb 3
def denorm_mae(
        pred_fn: Callable[[np.ndarray], np.ndarray], 
        x_chunks: np.ndarray, 
        y_chunks: np.ndarray, 
        y_mean: float,
        y_std: float,
    ) -> float:
    "Denormalized MAE"
    n = len(x_chunks)
    n_splits = 10
    while not (n / n_splits).is_integer():
        n_splits -= 1
        if n_splits < 0:
            raise ValueError

    x_splits = np.split(x_chunks, n_splits)
    y_splits = np.split(y_chunks, n_splits)
    model_mae = []
    for x, y in zip(x_splits, y_splits):
        model_mae.append(mean_absolute_error(
            # enilm.norm.denormalize(y.squeeze(), y_mean, y_std),
            enilm.norm.denormalize(y, y_mean, y_std),
            # enilm.norm.denormalize(model(np.expand_dims(x, axis=2)).numpy(), y_mean, y_std),
            # enilm.norm.denormalize(model(x).numpy(), y_mean, y_std),
            enilm.norm.denormalize(pred_fn(x), y_mean, y_std),
        ))

    return np.average(model_mae)


# %% ../nbs/metrics.ipynb 7
def common(series: pd.Series):
    return {
        "count": series.count(),
        "nunique": series.nunique(),
        "mean": series.mean(),
        "std": series.std(),
        "sum": series.sum(),
        "min": series.min(),
        "max": series.max(),
        "memory_usage": series.memory_usage(),
        "quantile (.25)": series.quantile(q=0.25),
        "quantile (.5)": series.quantile(q=0.5),  # median
        "quantile (.75)": series.quantile(q=0.75),
    }

# %% ../nbs/metrics.ipynb 10
rounder = lambda v: f"{round(v):,}"

def sizeof_fmt(num, suffix="B"):
    for unit in ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"]:
        if abs(num) < 1024.0:
            return f"{num:3.1f}{unit}{suffix}"
        num /= 1024.0
    return f"{num:.1f}Yi{suffix}"

common_format_dict = {
    "count": "{:,}",
    "nunique": "{:,}",
    "mean": rounder,
    "std": rounder,
    "sum": rounder,
    "min": rounder,
    "max": rounder,
    "quantile (.25)": rounder,
    "quantile (.5)": rounder,
    "quantile (.75)": rounder,
    "memory_usage": sizeof_fmt,
}
