"""."""

import copy

import matplotlib.pyplot as plt
import numpy as np

from matplotlib.backend_bases import MouseEvent, PickEvent
from matplotlib.container import BarContainer
from matplotlib.patches import Circle, Polygon
from pyquaternion import Quaternion
from scipy.spatial.transform import Rotation

from association_studio.boxes import from_vec_z, rotate_by_quat
from association_studio.score_driver.metric_labels import get_labels
from association_studio.score_driver.score_driver import ScoreDriver

from .metric_mode import MetricMode
from .show_coo_mode import ShowCooMode


class AssociationStudioBackend:
    def __init__(self) -> None:
        self.score_driver = ScoreDriver()
        self.show_coo_mode = ShowCooMode.TRANSLATION_AND_ROTATION
        self.metric_mode = MetricMode()
        self.box_det = from_vec_z(self.score_driver.det_1z[0])
        self.box_trk = from_vec_z(np.array((0.3, 0.7, 0.0, 1.0, 0.0, 0.0, -0.4, 2.0, 4.0, 1.0)))
        self.box_trk_start_pose_change = copy.deepcopy(self.box_trk)
        bcc = self.box_det.bottom_corners()[:2, :].T
        self.det_poly = Polygon(bcc, edgecolor='k', facecolor='yellow', closed=True)
        self.det_circle = Circle(bcc[2:].mean(axis=0), 0.35, edgecolor='k', facecolor='yellow')

        bcc = self.box_trk.bottom_corners()[:2, :].T
        self.trk_poly = Polygon(bcc, edgecolor='b', facecolor='none', closed=True, picker=True)
        trk_xy = bcc[2:].mean(axis=0)
        self.trk_circle = Circle(trk_xy, 0.35, edgecolor='b', facecolor='w', picker=True)
        self.score_driver.compute_score(self.box_trk)

    def next_show_coo_mode(self) -> None:
        self.show_coo_mode = ShowCooMode(self.show_coo_mode.next())

    def next_metric(self) -> None:
        self.metric_mode.next()
        self.score_driver.set_association_metric(self.metric_mode.metric.value)
        self.score_driver.compute_score(self.box_trk)

    def start_pose_change(self) -> None:
        self.box_trk_start_pose_change = copy.deepcopy(self.box_trk)

    def change_trk_pose(self, pick: PickEvent, event: MouseEvent) -> None:
        me = pick.mouseevent

        if pick.artist == self.trk_poly:
            translation = np.array((event.xdata - me.xdata, event.ydata - me.ydata, 0.0))
            self.box_trk = copy.deepcopy(self.box_trk_start_pose_change)
            self.box_trk.translate(translation)
        elif pick.artist == self.trk_circle:
            self.box_trk = copy.deepcopy(self.box_trk_start_pose_change)
            c = self.box_trk.center
            a = np.array((event.xdata - c[0], event.ydata - c[1], 0.0))
            b = np.array((me.xdata - c[0], me.ydata - c[1], 0.0))
            q_ab = Quaternion(Rotation.align_vectors(a, b)[0].as_quat(scalar_first=True))
            rotate_by_quat(q_ab, self.box_trk.orientation)

        bcc = self.box_trk.bottom_corners()[:2, :].T
        self.trk_poly.set_xy(bcc)
        self.trk_circle.set_center(bcc[2:].mean(axis=0))
        self.score_driver.compute_score(self.box_trk)

    def get_title(self) -> str:
        return f'{self.metric_mode.metric.value} {self.score_driver.score:5.4f}'

    def get_text(self) -> str:
        return self.show_coo_mode.format(self.box_trk)

    def get_horizontal_bar_container(self, ax_metrics: plt.Axes) -> BarContainer:
        scores = self.score_driver.scores
        names = get_labels(self.score_driver.metric_names)
        return ax_metrics.barh(range(scores.shape[0]), scores, tick_label=names)

    def update_horizontal_bar_container(self, bar: BarContainer) -> None:
        for bar, new_h in zip(bar, self.score_driver.scores):
            bar.set_width(new_h)
