from __future__ import annotations
import numpy as np


def mae(pred, obs, axis=None):
    return np.nanmean(np.abs(pred - obs), axis=axis)


def rmse(pred, obs, axis=None):
    return np.sqrt(np.nanmean((pred - obs) ** 2, axis=axis))


def rmae(pred, obs, axis=None, eps=1.0):
    den = np.nanmean(np.abs(obs), axis=axis)
    return 100.0 * mae(pred, obs, axis=axis) / np.maximum(den, eps)


def corr(pred, obs, axis=0):
    p = np.asarray(pred, dtype=float)
    o = np.asarray(obs, dtype=float)
    ok = np.isfinite(p) & np.isfinite(o)
    p = np.where(ok, p, np.nan)
    o = np.where(ok, o, np.nan)
    pm = np.nanmean(p, axis=axis, keepdims=True)
    om = np.nanmean(o, axis=axis, keepdims=True)
    pc = p - pm
    oc = o - om
    num = np.nansum(pc * oc, axis=axis)
    den = np.sqrt(np.nansum(pc**2, axis=axis) * np.nansum(oc**2, axis=axis))
    with np.errstate(invalid='ignore', divide='ignore'):
        r = num / den
    return r


def categorical_scores(forecast_bool, obs_bool, axis=0):
    f = np.asarray(forecast_bool, dtype=bool)
    o = np.asarray(obs_bool, dtype=bool)
    H = np.sum(f & o, axis=axis).astype(float)
    M = np.sum((~f) & o, axis=axis).astype(float)
    F = np.sum(f & (~o), axis=axis).astype(float)
    C = np.sum((~f) & (~o), axis=axis).astype(float)
    with np.errstate(invalid='ignore', divide='ignore'):
        pod = H / (H + M)
        pofd = F / (F + C)
        far = F / (H + F)
        fbi = (H + F) / (H + M)
        csi = H / (H + M + F)
    return {'POD': pod, 'POFD': pofd, 'FAR': far, 'FBI': fbi, 'CSI': csi}