"""."""

import numpy as np

from kinematic_tracker.association.association_metric import AssociationMetric
from nuscenes.utils.data_classes import Box

from association_studio.boxes import fill_vec_z
from association_studio.score_driver.tracker_factory import get_low_order_nu_scenes_tracker


class ScoreDriver:
    def __init__(self):
        self.det_1z = np.array([(0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 2.0, 4.0, 1.0)])
        self.buf_1z = np.zeros((1, 10))
        tracker = get_low_order_nu_scenes_tracker()
        tracker.advance(0, self.det_1z)
        self.ref_kalman_filters = [tracker.tracks_c[0][0].kkf.kalman_filter]
        valid_metrics = [
            mt.value for mt in AssociationMetric if mt != AssociationMetric.UNKNOWN_METRIC
        ]
        self.metric_names = np.array(valid_metrics, dtype=str)
        num_metrics = len(self.metric_names)
        self.metric_drivers: np.ndarray[tuple[int], np.dtype[np.object_]] = np.zeros(
            num_metrics, object
        )
        self.scores = np.zeros(num_metrics)
        self._init_metric_drivers()
        self.metric_mask = np.zeros(num_metrics, bool)
        self.metric_index = -1
        self.score = 0.0
        self.set_association_metric(AssociationMetric.SIZE_MODULATED_MAHALANOBIS.value)

    def _init_metric_drivers(self) -> None:
        tracker = get_low_order_nu_scenes_tracker()
        for metric_num, metric_name in enumerate(self.metric_names):
            tracker.set_association_metric(metric_name)
            self.metric_drivers[metric_num] = tracker.metric_driver

    def set_association_metric(self, metric: str) -> None:
        self.metric_mask[:] = metric == self.metric_names
        self.metric_index = np.argmax(self.metric_mask)

    def compute_score(self, box: Box) -> None:
        self.compute_scores(box)
        assert self.metric_index >= 0
        self.score = self.scores[self.metric_index]

    def compute_scores(self, box: Box) -> None:
        fill_vec_z(box, self.buf_1z[0])
        for idx, driver in enumerate(self.metric_drivers):
            self.scores[idx] = driver.compute_metric(self.buf_1z, self.ref_kalman_filters)[0, 0]
