"""."""

import numpy as np

from kinematic_tracker.association.association_metric import AssociationMetric

from association_studio.common import FVT, OVT

from .tracker_factory import get_low_order_nu_scenes_tracker
from .valid_metrics import ValidMetrics


class ScoreDriver:
    def __init__(self) -> None:
        self.valid_metrics = ValidMetrics()
        self.probe_1z = np.zeros((1, 10))
        tracker = get_low_order_nu_scenes_tracker()
        self.ref_1z = np.array([(0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 4.0, 2.0, 1.0)])
        tracker.advance(0, self.ref_1z)
        self._ref_kfs = [tracker.tracks_c[0][0].kkf.kalman_filter]
        self.metric_names = np.array(self.valid_metrics.get_valid_metrics(), dtype=str)
        num_metrics = len(self.metric_names)
        self.metric_drivers: OVT = 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.MAHALANOBIS_SIZE_MODULATED.value)

    def _init_metric_drivers(self) -> None:
        for metric_num, metric_name in enumerate(self.metric_names):
            self.metric_drivers[metric_num] = self.valid_metrics.get_driver(metric_name)

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

    def set_ref_z(self, ref_z: FVT) -> None:
        self._ref_kfs[0].statePre[:, 0] = ref_z

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

    def compute_scores(self, vec_z: FVT) -> None:
        self.probe_1z[0] = vec_z
        for idx, driver in enumerate(self.metric_drivers):
            self.scores[idx] = driver.compute_metric(self.probe_1z, self._ref_kfs)[0, 0]
