"""Fixed gradient histograms and optional coarse appearance; no learned extractor.

HOG-like: 9 unsigned soft orientation bins, 4x4-pixel cells,
2x2-cell overlapping blocks, L2-Hys normalization. Exactly 1764 values.
"""
import numpy as np


def descriptor(rgb, hybrid=False, sparse=False):
    rgb = np.asarray(rgb).reshape(-1, 3, 32, 32).astype(np.float32) / 255
    # Square-root compression reduces dependence on illumination.
    image = np.sqrt(rgb)
    gx, gy = np.zeros_like(image), np.zeros_like(image)
    gx[:, :, :, 1:-1] = image[:, :, :, 2:] - image[:, :, :, :-2]
    gy[:, :, 1:-1, :] = image[:, :, 2:, :] - image[:, :, :-2, :]
    magnitude = np.hypot(gx, gy)
    channel = magnitude.argmax(axis=1)[:, None]
    mag = np.take_along_axis(magnitude, channel, axis=1)[:, 0]
    dx = np.take_along_axis(gx, channel, axis=1)[:, 0]
    dy = np.take_along_axis(gy, channel, axis=1)[:, 0]
    angle = (np.arctan2(dy, dx) % np.pi) * (9 / np.pi)
    lower = np.floor(angle).astype(int) % 9
    fraction = angle - np.floor(angle)
    hist = np.empty((len(rgb), 8, 8, 9), dtype=np.float32)
    for b in range(9):
        weighted = mag * ((lower == b) * (1 - fraction) + ((lower + 1) % 9 == b) * fraction)
        hist[:, :, :, b] = weighted.reshape(-1, 8, 4, 8, 4).sum(axis=(2, 4))
    blocks = np.stack([hist[:, y:y+2, x:x+2].reshape(len(rgb), -1)
                       for y in range(7) for x in range(7)], axis=1)
    blocks /= np.sqrt((blocks * blocks).sum(axis=2, keepdims=True) + 1e-10)
    blocks = np.minimum(blocks, 0.2)
    blocks /= np.sqrt((blocks * blocks).sum(axis=2, keepdims=True) + 1e-10)
    values = blocks.reshape(len(rgb), -1)
    if sparse:
        # Keep the 256 strongest local orientation features per image.
        keep = np.argpartition(values, -256, axis=1)[:, -256:]
        sparse_values = np.zeros_like(values)
        np.put_along_axis(sparse_values, keep, np.take_along_axis(values, keep, axis=1), axis=1)
        values = sparse_values
    if hybrid:
        color = rgb.reshape(-1, 3, 4, 8, 4, 8).mean(axis=(3, 5)).reshape(len(rgb), -1)
        gray = (rgb[:, 0] * .299 + rgb[:, 1] * .587 + rgb[:, 2] * .114)
        gray = gray.reshape(-1, 8, 4, 8, 4).mean(axis=(2, 4)).reshape(len(rgb), -1)
        values = np.concatenate([values, color * .25, gray * .25, (1 - gray) * .25], axis=1)
    encoded = np.zeros((len(rgb), 2048), dtype=np.uint8)
    encoded[:, :values.shape[1]] = np.clip(np.rint(values * 3 * 255), 0, 255).astype(np.uint8)
    return encoded