"""
# -*- coding: utf-8 -*-
#
# Copyright 2021 Michael Büsch <m@bues.ch>
#
# Licensed under the Apache License version 2.0
# or the MIT license, at your option.
# SPDX-License-Identifier: Apache-2.0 OR MIT
#
"""

__all__ = [
    "minibatches",
    "proportion_equal",
    "array_to_boolint",
    "collapse_bool_nodes",
]

from typing import Generator, Iterator, Tuple
import numpy as np

def cut_batch(batch: np.ndarray,
              minibatch_size: int) -> Generator:
    size = batch.shape[0]
    return (
        batch[i:i+minibatch_size]
        for i in range(0, size, minibatch_size)
    )

def minibatches(x_batch: np.ndarray,
                y_batch: np.ndarray,
                minibatch_size: int = 128)\
                -> Iterator[Tuple[Generator, Generator]]:
    """Split a batch into mini batches.
    """
    assert x_batch.shape[0] == y_batch.shape[0]
    return zip(cut_batch(x_batch, minibatch_size),
               cut_batch(y_batch, minibatch_size))

def proportion_equal(a: np.ndarray,
                     b: np.ndarray,
                     rtol: float = 1e-05,
                     atol: float = 1e-08) -> float:
    """Returns the proportion of equal elements in a and b
    as a number between 0.0 and 1.0.
    """
    return np.mean(np.isclose(a, b, rtol, atol))

def array_to_boolint(y, thres=0.75):
    """Convert an array to a new array with {0, 1} values.
    Values below the threshold are translated to 0.
    Other values are translated to 1.
    """
    return (y >= thres).astype(np.int)

def collapse_bool_nodes(y, thres=0.75):
    """Collapse the multiple boolean nodes into one integer node.
    """
    assert y.ndim == 2
    nr_nodes = y.shape[1]
    y = array_to_boolint(y, thres)
    # Scale nodes.
    y *= np.fromiter((1 << i for i in range(nr_nodes)),
                      dtype=y.dtype)
    # Collapse nodes.
    y = np.sum(y, axis=1)
    return y

# vim: ts=4 sw=4 expandtab