import math
import matplotlib.pyplot as plt
from matplotlib.backend_bases import MouseButton, PickEvent, MouseEvent, KeyEvent
from geomdl import BSpline, utilities, operations
from geomdl.knotvector import generate
import numpy as np
import matplotlib.axes
import matplotlib.lines
import multiprocessing
from scipy.optimize import differential_evolution


class InteractiveCurve:
    def __init__(self,
                 degree = 4,
                 size = 10,
                 enus: list[tuple[np.float64, np.float64, np.float64]] = None,
                 ):
        plt.rcParams['keymap.back'] = []
        plt.rcParams['keymap.forward'] = []
        plt.rcParams['keymap.home'] = []
        plt.rcParams['keymap.pan'] = []
        plt.rcParams['keymap.xscale'] = []
        plt.rcParams['keymap.yscale'] = []

        self.resize = False
        self.clamped = False
        self.loop = False
        self.k = False
        self.pick = None
        self.near = False
        self.near_value = None

        self.enus = enus
        self.e_bound = None
        self.n_bound = None
        if self.enus is not None and len(self.enus) > 0:
            self.e_bound = min(enu[0] for enu in self.enus), max(enu[0] for enu in self.enus)
            self.n_bound = min(enu[1] for enu in self.enus), max(enu[1] for enu in self.enus)

        self.sample_size = 100

        self.ctrl_pts = np.empty((0, 2), dtype=float)
        self.ctrl_pts_list = self.ctrl_pts.tolist()
        self.spline_curve = BSpline.Curve()
        self.spline_curve.degree = degree

        self.fig, (self.ax_spline, self.ax_curv) = plt.subplots(2, 1, gridspec_kw={'height_ratios': [3, 1]})
        self.update_window_title()

        self.ax_spline.margins(0.2)
        self.ax_spline.set_xlim(-size, size)
        self.ax_spline.set_ylim(-size, size)
        self.ax_spline.set_aspect('equal')

        self.plot_bag_nav, = self.ax_spline.plot([], [], 'm--', lw=1)
        if self.enus is not None and len(self.enus) > 0:
            self.plot_bag_nav.set_data([enu[0] for enu in self.enus], [enu[1] for enu in self.enus])
        self.plot_spline, = self.ax_spline.plot([], [], 'b-', lw=2)
        self.plot_ctrl_ploy, = self.ax_spline.plot([], [], 'k--', lw=1)
        self.plot_ctrl_pt_colors = ['green'] * 0
        self.plot_ctrl_pts = self.ax_spline.scatter([], [], s=80, c='green', picker=True, pickradius=5)
        self.plot_spline_point, = self.ax_spline.plot([], [], 'ro', lw=2)
        self.plot_spline_curv_min, = self.ax_spline.plot([], [], 'yv', lw=1)
        self.plot_spline_curv_max, = self.ax_spline.plot([], [], 'y^', lw=1)

        self.fig.canvas.mpl_connect('pick_event', self.on_pick)
        self.fig.canvas.mpl_connect('button_press_event', self.on_mouse_press)
        self.fig.canvas.mpl_connect('button_release_event', self.on_mouse_release)
        self.fig.canvas.mpl_connect('motion_notify_event', self.on_mouse_move)
        self.fig.canvas.mpl_connect('key_press_event', self.on_key_press)
        self.fig.canvas.mpl_connect('key_release_event', self.on_key_release)

        self.plot_curv, = self.ax_curv.plot([], [], 'b-', lw=2)
        self.plot_curv_point, = self.ax_curv.plot([], [], 'ro', lw=2)
        self.plot_curv_min, = self.ax_curv.plot([], [], 'yv', lw=1)
        self.plot_curv_max, = self.ax_curv.plot([], [], 'y^', lw=1)

        self.update_curve()
        plt.show()

    def update_window_title(self):
        length = 0
        if len(self.ctrl_pts) > self.spline_curve.degree:
            length = operations.length_curve(self.spline_curve)
        title: str = (f"resize/r: {self.resize}  "
                      f"clamped/c: {self.clamped}  "
                      f"loop/l: {self.loop}  "
                      f"k: {self.k}  "
                      f"near/n: {self.near_value}  "
                      f"pick: {self.pick}  "
                      f"ss: {str(self.spline_curve.sample_size)}  "
                      f"len: {length}")
        self.fig.canvas.manager.set_window_title(title)

    def get_ctrl_poly_length(self):
        return np.sum(np.linalg.norm(np.diff(np.array(self.spline_curve.ctrlpts), axis=0), axis=1))

    def calculate_curvature(self, u):
        if len(self.ctrl_pts_list) <= self.spline_curve.degree:
            return None
        der = self.spline_curve.derivatives(u, 2)
        n = der[1][0] * der[2][1] - der[1][1] * der[2][0]
        m = (der[1][0] ** 2 + der[1][1] ** 2) ** 1.5
        if m < 1e-9:
            return 0
        return n / m

    def calculate_greville(self, index):
        if len(self.ctrl_pts_list) <= self.spline_curve.degree:
            return None
        s = 0
        for i in range(self.spline_curve.degree - 1):
            s += self.spline_curve.knotvector[i + 1 + index]
        return s / self.spline_curve.degree

    def update_curve_ctrl_pts(self):
        self.plot_ctrl_pts.set_offsets(self.ctrl_pts)
        self.plot_ctrl_pt_colors = ['green'] * len(self.ctrl_pts)
        if self.pick is not None:
            self.plot_ctrl_pt_colors[self.pick] = "red"
        self.plot_ctrl_pts.set_color(self.plot_ctrl_pt_colors)
        self.plot_ctrl_ploy.set_data(self.ctrl_pts[:, 0], self.ctrl_pts[:, 1])

    def update_ctrl_pts_list(self):
        self.ctrl_pts_list = self.ctrl_pts.tolist()
        if not self.loop:
            return
        if len(self.ctrl_pts) == 0:
            return
        add_size = self.spline_curve.degree # + self.spline_curve.degree + 1
        add_pts = self.ctrl_pts_list
        while len(add_pts) < add_size:
            add_pts += self.ctrl_pts_list
        self.ctrl_pts_list += add_pts[:add_size]

    def update_curve_spline(self):
        if len(self.ctrl_pts_list) <= self.spline_curve.degree:
            return
        self.spline_curve.ctrlpts = self.ctrl_pts_list
        self.spline_curve.knotvector = generate(self.spline_curve.degree, self.spline_curve.ctrlpts_size, clamped=self.clamped)
        self.spline_curve.delta = 1. / self.sample_size
        eval_pts = np.array(self.spline_curve.evalpts)
        self.plot_spline.set_data(eval_pts[:, 0], eval_pts[:, 1])

    def update_curve_curvature(self):
        if not self.k:
            self.plot_spline_curv_min.set_data([], [])
            self.plot_spline_curv_max.set_data([], [])
            self.plot_curv_min.set_data([], [])
            self.plot_curv_max.set_data([], [])
            return
        xu = np.linspace(self.spline_curve.domain[0], self.spline_curve.domain[1], self.spline_curve.sample_size)
        curvatures = [self.calculate_curvature(u) for u in xu]
        self.plot_curv.set_data(xu, curvatures)
        self.ax_curv.relim()
        self.ax_curv.autoscale_view()
        min_index = np.argmin(curvatures)
        max_index = np.argmax(curvatures)
        min_point = self.spline_curve.evaluate_single(xu[min_index].item())
        max_point = self.spline_curve.evaluate_single(xu[max_index].item())
        self.plot_spline_curv_min.set_data([min_point[0]], [min_point[1]])
        self.plot_spline_curv_max.set_data([max_point[0]], [max_point[1]])
        self.plot_curv_min.set_data([xu[min_index]], [curvatures[min_index]])
        self.plot_curv_max.set_data([xu[max_index]], [curvatures[max_index]])

        self.ax_curv.set_title(f"{curvatures[min_index]:>.5f} {curvatures[max_index]:>.5f}", y=0)

    def update_curve_near(self):
        if self.near_value:
            point = self.spline_curve.evaluate_single(self.near_value)
            self.plot_spline_point.set_data([point[0]], [point[1]])
            curv = self.calculate_curvature(self.near_value)
            self.plot_curv_point.set_data([self.near_value], [curv])
        else:
            self.plot_spline_point.set_data([], [])
            self.plot_curv_point.set_data([], [])

    def update_curve_resize(self):
        if self.resize:
            if len(self.ctrl_pts) == 0 and (self.e_bound is None or self.n_bound is None):
                return
            cp_x_min, cp_y_min = math.inf, math.inf
            cp_x_max, cp_y_max = -math.inf, -math.inf
            if len(self.ctrl_pts) > 0:
                cp_x_min, cp_y_min = self.ctrl_pts.min(axis=0)
                cp_x_max, cp_y_max = self.ctrl_pts.max(axis=0)
            e_min, n_min = math.inf, math.inf
            e_max, n_max = -math.inf, -math.inf
            if self.e_bound is not None:
                e_min, e_max = self.e_bound
            if self.n_bound is not None:
                n_min, n_max = self.n_bound
            x_min = min(cp_x_min, e_min)
            x_max = max(cp_x_max, e_max)
            y_min = min(cp_y_min, n_min)
            y_max = max(cp_y_max, n_max)
            width = max(x_max - x_min, 1)
            height = max(y_max - y_min, 1)
            margin_ratio = 0.2
            self.ax_spline.set_xlim(x_min - margin_ratio * width, x_max + margin_ratio * width)
            self.ax_spline.set_ylim(y_min - margin_ratio * height, y_max + margin_ratio * height)

    def update_curve(self):
        self.update_curve_ctrl_pts()
        self.update_ctrl_pts_list()
        self.update_curve_spline()
        self.update_curve_curvature()
        self.update_curve_near()
        self.update_curve_resize()
        self.update_window_title()
        self.fig.canvas.draw_idle()

    def on_pick(self, event: PickEvent):
        if event.artist == self.plot_ctrl_pts:
            if len(event.ind) > 0:
                self.pick = event.ind[0]
        self.update_curve()

    def on_mouse_press(self, event: MouseEvent):
        if event.button == MouseButton.RIGHT:
            if event.xdata is not None and event.ydata is not None:
                self.ctrl_pts = np.vstack([self.ctrl_pts, [event.xdata, event.ydata]])
        self.update_curve()

    def on_mouse_release(self, event: MouseEvent):
        if event.button == MouseButton.LEFT:
            self.pick = None
        if event.button == MouseButton.RIGHT:
            self.pick = None
        self.update_curve()

    def on_mouse_move(self, event: MouseEvent):
        if event.xdata is None or event.ydata is None:
            return
        if self.pick is not None:
            self.ctrl_pts[self.pick] = event.xdata, event.ydata
        self.near_value = None
        if self.near and len(self.ctrl_pts_list) > self.spline_curve.degree:
            closet = differential_evolution(
                func=lambda u: np.sum((np.array(self.spline_curve.evaluate_single(u[0])) - np.array([event.xdata, event.ydata])) ** 2),
                bounds=[(self.spline_curve.domain[0], self.spline_curve.domain[1])]
            )
            if closet.success:
                self.near_value = closet.x[0]
        self.update_curve()

    def on_key_press(self, event: KeyEvent):
        if event.key == 'n':
            self.near = not self.near
        if event.key == 'c':
            self.clamped = not self.clamped
        if event.key == 'r':
            self.resize = not self.resize
        if event.key == 'k':
            self.k = not self.k
        if event.key == 'l':
            self.loop = not self.loop
        if event.key == 'escape':
            self.pick = None
        if self.pick is not None and event.key in ('delete', 'backspace'):
            self.ctrl_pts = np.delete(self.ctrl_pts, [self.pick], 0)
            self.pick = None
        if event.key == 'a' and event.xdata is not None and event.ydata is not None:
            mouse = np.array([event.xdata, event.ydata])
            p = self.ctrl_pts[:-1]
            q = self.ctrl_pts[1:]
            pq = q - p
            pm = mouse - p
            t = np.sum(pq * pm, axis=1) / np.sum(pq * pq, axis=1)
            o = p + pq * t[:, np.newaxis]
            d = np.sum((o - mouse) ** 2, axis=1)
            condition = (t >= 0) & (t <= 1)
            indexes = np.where(condition)[0]
            if len(indexes) > 0:
                distances = d[condition]
                closed_index = indexes[np.argmin(distances)]
                self.ctrl_pts = np.insert(self.ctrl_pts, closed_index + 1, mouse, axis=0)
        if event.key == 'up':
            self.sample_size += 10
            self.sample_size = min(self.sample_size, 2000)
        if event.key == 'down':
            self.sample_size -= 10
            self.sample_size = max(self.sample_size, 10)
        self.update_curve()

    def on_key_release(self, event: KeyEvent):
        pass

class Command:
    def __init__(self):
        self.queue = multiprocessing.Queue()
