from collections import defaultdict
import logging
import time
from typing import List

from .decoder import Decoder
from . import utils
from ..annotation import AnnotationDet
from .. import headmeta, visualizer

LOG = logging.getLogger(__name__)


class CifDet(Decoder):
    occupancy_visualizer = visualizer.Occupancy()
    max_detections_before_nms = 120

    def __init__(self, head_metas: List[headmeta.CifDet], *, visualizers=None):
        super().__init__()
        self.metas = head_metas

        # prefer decoders with more classes
        self.priority = -1.0  # prefer keypoints over detections
        self.priority += sum(m.n_fields for m in head_metas) / 1000.0

        self.visualizers = visualizers
        if self.visualizers is None:
            self.visualizers = [visualizer.CifDet(meta) for meta in self.metas]

        self.timers = defaultdict(float)

    @classmethod
    def factory(cls, head_metas):
        # TODO: multi-scale
        return [
            CifDet([meta])
            for meta in head_metas
            if isinstance(meta, headmeta.CifDet)
        ]

    def __call__(self, fields):
        start = time.perf_counter()

        if self.visualizers:
            for vis, meta in zip(self.visualizers, self.metas):
                vis.predicted(fields[meta.head_index])

        cifhr = utils.CifDetHr().fill(fields, self.metas)
        seeds = utils.CifDetSeeds(cifhr.accumulated).fill(fields, self.metas)
        occupied = utils.Occupancy(cifhr.accumulated.shape, 2, min_scale=2.0)

        annotations = []
        for v, f, x, y, w, h in seeds.get():
            if occupied.get(f, x, y):
                continue
            ann = AnnotationDet(self.metas[0].categories).set(
                f + 1, v, (x - w / 2.0, y - h / 2.0, w, h))
            annotations.append(ann)
            if len(annotations) >= self.max_detections_before_nms:
                break
            occupied.set(f, x, y, 0.1 * min(w, h))

        self.occupancy_visualizer.predicted(occupied)

        annotations = utils.nms.Detection().annotations(annotations)
        # annotations = sorted(annotations, key=lambda a: -a.score)

        LOG.info('annotations %d, decoder = %.3fs', len(annotations), time.perf_counter() - start)
        return annotations
