"""
DeepLabCut Toolbox (deeplabcut.org)
© A. & M. Mathis Labs

Licensed under GNU Lesser General Public License v3.0
"""


import numpy as np


def extract_cnn_output(outputs, cfg):
    """
    Extract location refinement and score map from DeepLabCut network

    Parameters
    -----------
    outputs : list
        List of outputs from DeepLabCut network.
        Requires 2 entries:
            index 0 is output from Sigmoid
            index 1 is output from pose/locref_pred/block4/BiasAdd

    cfg : dict
        Dictionary read from the pose_cfg.yaml file for the network.

    Returns
    --------
    scmap : ?
        score map

    locref : ?
        location refinement
    """

    scmap = outputs[0]
    scmap = np.squeeze(scmap)
    locref = None
    if cfg["location_refinement"]:
        locref = np.squeeze(outputs[1])
        shape = locref.shape
        locref = np.reshape(locref, (shape[0], shape[1], -1, 2))
        locref *= cfg["locref_stdev"]
    if len(scmap.shape) == 2:  # for single body part!
        scmap = np.expand_dims(scmap, axis=2)
    return scmap, locref


def argmax_pose_predict(scmap, offmat, stride):
    """
    Combines score map and offsets to the final pose

    Parameters
    -----------
    scmap : ?
        score map

    offmat : ?
        offsets

    stride : ?
        ?

    Returns
    --------
    pose :class:`numpy.ndarray`
        pose as a numpy array
    """

    num_joints = scmap.shape[2]
    pose = []
    for joint_idx in range(num_joints):
        maxloc = np.unravel_index(
            np.argmax(scmap[:, :, joint_idx]), scmap[:, :, joint_idx].shape
        )
        offset = np.array(offmat[maxloc][joint_idx])[::-1]
        pos_f8 = np.array(maxloc).astype("float") * stride + 0.5 * stride + offset
        pose.append(np.hstack((pos_f8[::-1], [scmap[maxloc][joint_idx]])))
    return np.array(pose)


def get_top_values(scmap, n_top=5):
    batchsize, ny, nx, num_joints = scmap.shape
    scmap_flat = scmap.reshape(batchsize, nx * ny, num_joints)
    if n_top == 1:
        scmap_top = np.argmax(scmap_flat, axis=1)[None]
    else:
        scmap_top = np.argpartition(scmap_flat, -n_top, axis=1)[:, -n_top:]
        for ix in range(batchsize):
            vals = scmap_flat[ix, scmap_top[ix], np.arange(num_joints)]
            arg = np.argsort(-vals, axis=0)
            scmap_top[ix] = scmap_top[ix, arg, np.arange(num_joints)]
        scmap_top = scmap_top.swapaxes(0, 1)

    Y, X = np.unravel_index(scmap_top, (ny, nx))
    return Y, X


def multi_pose_predict(scmap, locref, stride, num_outputs):
    Y, X = get_top_values(scmap[None], num_outputs)
    Y, X = Y[:, 0], X[:, 0]
    num_joints = scmap.shape[2]
    DZ = np.zeros((num_outputs, num_joints, 3))
    for m in range(num_outputs):
        for k in range(num_joints):
            x = X[m, k]
            y = Y[m, k]
            DZ[m, k, :2] = locref[y, x, k, :]
            DZ[m, k, 2] = scmap[y, x, k]

    X = X.astype("float32") * stride + 0.5 * stride + DZ[:, :, 0]
    Y = Y.astype("float32") * stride + 0.5 * stride + DZ[:, :, 1]
    P = DZ[:, :, 2]

    pose = np.empty((num_joints, num_outputs * 3), dtype="float32")
    pose[:, 0::3] = X.T
    pose[:, 1::3] = Y.T
    pose[:, 2::3] = P.T

    return pose