from __future__ import annotations

from functools import partial
from typing import Callable, Sequence, Tuple

import flax.linen as nn
import jax
import jax.numpy as jnp

from ..ops import distance_similarity, non_max_suppression
from ..typing import Array, ArrayLike
from .common import DefaultUnpicklerMixin


def generate_predictions(module, outputs):
    """
    Produce a list of proposal locations based on predication map, remove redundency with non_max_suppression
    """
    locations = outputs["pred_locs"]
    scores = jax.nn.sigmoid(outputs["logits"])

    distance_threshold = module.nms_threshold
    output_size = module.max_output
    topk = module.pre_nms_topk
    score_threshold = module.min_score

    # sort and nms
    if topk <= 0 or topk > scores.size:
        topk = scores.size

    scores, indices = jax.lax.top_k(scores, topk)
    locations = locations[indices]

    threshold = 1 / distance_threshold / distance_threshold
    sel = non_max_suppression(
        scores,
        locations,
        output_size,
        threshold,
        score_threshold,
        return_selection=True,
    )

    idx_of_selected = jnp.argwhere(sel, size=output_size, fill_value=-1).squeeze(-1)

    scores = jnp.where(idx_of_selected >= 0, scores[idx_of_selected], -1.0)
    locations = jnp.where(
        idx_of_selected[:, None] >= 0, locations[idx_of_selected], -1.0
    )

    return dict(
        scores=scores,
        locations=locations,
    )


class LPN(nn.Module, DefaultUnpicklerMixin):
    """Location detection head

    Attributes:
        nms_threshold: non-max-supression threshold, if performing nms on detected locations.
        pre_nms_topk: max number of detections to be processed regardless of nms, ignored if negative
        max_output: number of detection outputs
    """

    # network hyperparams
    feature_scale: int = 4
    dtype: Any = None

    # detection hyperparams
    nms_threshold: float = 8.0
    pre_nms_topk: int = -1
    max_output: int = 512
    min_score: float = 0.2

    @nn.compact
    def _block(self, feature: ArrayLike) -> dict:
        x = feature
        height, width = x.shape[:-1]

        logits = nn.Conv(1, (1, 1), dtype=self.dtype)(x)
        regressions = nn.Conv(2, (1, 1), dtype=self.dtype)(x)

        ref_locs = (
            jnp.moveaxis(jnp.mgrid[:height, :width] + 0.5, 0, -1) * self.feature_scale
        )
        locs = ref_locs + regressions * self.feature_scale

        return dict(
            regressions=regressions.reshape(-1, 2),
            logits=logits.reshape(-1),
            ref_locs=ref_locs.reshape(-1, 2),
            pred_locs=locs.reshape(-1, 2),
        )

    def __call__(self, feature: ArrayLike, mask: ArrayLike | None = None) -> dict:
        network_outputs = self._block(feature)
        predictions = generate_predictions(self, network_outputs)

        if mask is not None:
            pred_locs = tuple(
                jnp.floor(predictions["locations"]).transpose().astype(int)
            )
            masked = mask[pred_locs]
            predictions["scores"] = jnp.where(masked, predictions["scores"], 0)
            predictions["locations"] = jnp.where(
                masked[:, None], predictions["locations"], -1
            )

        return dict(
            detector=network_outputs,
            predictions=predictions,
        )
