#!/usr/bin/env python3
"""
Quaternion GUI workbench + interactive 3D axes visualiser (NO normalisation).


- QA * QB is exactly ONE Hamilton product (no hidden extra maths)
- Plotting does NOT normalise the quaternion and does NOT do q*p*q̄ (no "two multiplies" for plotting)
- Still plots axes for "impure" cases, but overlays a warning on the graph
- Your rule: IMPURE if R >= 1e-6 (note: not abs(), exactly as stated)
- Shows a RED vector arrow for the vector part (i, j, k)

Dependencies:
    python -m pip install numpy matplotlib
"""

import math
import tkinter as tk
from tkinter import ttk, messagebox

import numpy as np

import matplotlib
matplotlib.use("TkAgg")
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk
from matplotlib.figure import Figure


EPS_IMPURE_R = 1e-6
EPS_NORM = 1e-12


class Quaternion:
    __slots__ = ("w", "x", "y", "z")

    def __init__(self, w: float, x: float, y: float, z: float):
        self.w = float(w)
        self.x = float(x)
        self.y = float(y)
        self.z = float(z)

    def as_tuple(self):
        return (self.w, self.x, self.y, self.z)

    def norm2(self) -> float:
        return self.w * self.w + self.x * self.x + self.y * self.y + self.z * self.z

    def norm(self) -> float:
        return math.sqrt(self.norm2())

    def conjugate(self):
        return Quaternion(self.w, -self.x, -self.y, -self.z)

    def inverse(self):
        n2 = self.norm2()
        if n2 < EPS_NORM:
            raise ZeroDivisionError("Quaternion norm is ~0; cannot invert.")
        c = self.conjugate()
        return Quaternion(c.w / n2, c.x / n2, c.y / n2, c.z / n2)

    #  "impure" if |R|>= 1e-6  
    def is_impure(self, eps_r=EPS_IMPURE_R) -> bool:
        return self.w >= abs(eps_r)

    def __add__(self, other):
        return Quaternion(self.w + other.w, self.x + other.x, self.y + other.y, self.z + other.z)

    def __sub__(self, other):
        return Quaternion(self.w - other.w, self.x - other.x, self.y - other.y, self.z - other.z)

    def __mul__(self, other):
        # Hamilton product
        w1, x1, y1, z1 = self.as_tuple()
        w2, x2, y2, z2 = other.as_tuple()
        return Quaternion(
            w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
            w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
            w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
            w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2,
        )

    def __truediv__(self, other):
        return self * other.inverse()


def quat_to_matrix_no_normalise(q: Quaternion) -> np.ndarray:
    """
    Standard quaternion->matrix formula, but DOES NOT normalise q.

    If q is not unit, the resulting matrix will generally not be orthonormal.
    That's useful here for visualising intermediates without 'helpful' coercion.
    """
    w, x, y, z = q.w, q.x, q.y, q.z

    r00 = 1.0 - 2.0 * (y * y + z * z)
    r01 = 2.0 * (x * y - w * z)
    r02 = 2.0 * (x * z + w * y)

    r10 = 2.0 * (x * y + w * z)
    r11 = 1.0 - 2.0 * (x * x + z * z)
    r12 = 2.0 * (y * z - w * x)

    r20 = 2.0 * (x * z - w * y)
    r21 = 2.0 * (y * z + w * x)
    r22 = 1.0 - 2.0 * (x * x + y * y)

    return np.array([[r00, r01, r02],
                     [r10, r11, r12],
                     [r20, r21, r22]], dtype=float)


class App(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("Quaternion Workbench + Axes Visualiser (raw quaternion)")
        self.geometry("1180x700")

        # QA, QB default to identity
        self.qa_vars = [
            tk.StringVar(value="1.0"),
            tk.StringVar(value="0.0"),
            tk.StringVar(value="0.0"),
            tk.StringVar(value="0.0"),
        ]
        self.qb_vars = [
            tk.StringVar(value="1.0"),
            tk.StringVar(value="0.0"),
            tk.StringVar(value="0.0"),
            tk.StringVar(value="0.0"),
        ]
        self.qr_vars = [
            tk.StringVar(value=""),
            tk.StringVar(value=""),
            tk.StringVar(value=""),
            tk.StringVar(value=""),
        ]

        self.status_var = tk.StringVar(value="Ready.")

        self._build_ui()
        self._build_plot()
        self._update_plot_from_result(None)

    # ---------------- UI ----------------

    def _build_ui(self):
        outer = ttk.Frame(self, padding=10)
        outer.pack(fill="both", expand=True)

        left = ttk.Frame(outer)
        left.pack(side="left", fill="y", padx=(0, 12))

        right = ttk.Frame(outer)
        right.pack(side="right", fill="both", expand=True)

        qa_frame = ttk.LabelFrame(left, text="Quaternion QA (R, i, j, k)", padding=10)
        qb_frame = ttk.LabelFrame(left, text="Quaternion QB (R, i, j, k)", padding=10)
        btn_frame = ttk.LabelFrame(left, text="Operations", padding=10)
        qr_frame = ttk.LabelFrame(left, text="Result QR (R, i, j, k)", padding=10)

        qa_frame.pack(fill="x", pady=(0, 10))
        qb_frame.pack(fill="x", pady=(0, 10))
        btn_frame.pack(fill="x", pady=(0, 10))
        qr_frame.pack(fill="x", pady=(0, 10))

        self._make_quat_row(qa_frame, self.qa_vars)
        ttk.Button(qa_frame, text="Conjugate QA", command=self.on_conjugate_qa).grid(
            row=2, column=0, columnspan=4, sticky="ew", pady=(8, 0)
        )

        self._make_quat_row(qb_frame, self.qb_vars)
        self._make_quat_row(qr_frame, self.qr_vars, readonly=True)

        ttk.Button(btn_frame, text="QA + QB", command=lambda: self.on_op("+")).grid(row=0, column=0, sticky="ew", padx=2, pady=2)
        ttk.Button(btn_frame, text="QA - QB", command=lambda: self.on_op("-")).grid(row=0, column=1, sticky="ew", padx=2, pady=2)
        ttk.Button(btn_frame, text="QA * QB", command=lambda: self.on_op("*")).grid(row=1, column=0, sticky="ew", padx=2, pady=2)
        ttk.Button(btn_frame, text="QA / QB", command=lambda: self.on_op("/")).grid(row=1, column=1, sticky="ew", padx=2, pady=2)

        ttk.Separator(btn_frame, orient="horizontal").grid(row=2, column=0, columnspan=2, sticky="ew", pady=(10, 10))

        ttk.Button(btn_frame, text="Copy QR → QA", command=self.copy_qr_to_qa).grid(row=3, column=0, sticky="ew", padx=2, pady=2)
        ttk.Button(btn_frame, text="Copy QR → QB", command=self.copy_qr_to_qb).grid(row=3, column=1, sticky="ew", padx=2, pady=2)

        ttk.Button(btn_frame, text="Clear Result", command=self.clear_result).grid(row=4, column=0, columnspan=2, sticky="ew", padx=2, pady=(8, 2))

        for c in range(2):
            btn_frame.columnconfigure(c, weight=1)

        status = ttk.Label(left, textvariable=self.status_var, wraplength=360, foreground="#444")
        status.pack(fill="x", pady=(6, 0))

        plot_frame = ttk.LabelFrame(
            right,
            text="3D View (interactive: drag to rotate, scroll to zoom)",
            padding=8,
        )
        plot_frame.pack(fill="both", expand=True)
        self.plot_frame = plot_frame

    def _make_quat_row(self, parent, vars4, readonly=False):
        labels = ["R", "i", "j", "k"]
        for idx, lab in enumerate(labels):
            ttk.Label(parent, text=lab).grid(row=0, column=idx, sticky="w")

        for idx, var in enumerate(vars4):
            ent = ttk.Entry(parent, textvariable=var, width=10)
            if readonly:
                ent.state(["readonly"])
            ent.grid(row=1, column=idx, sticky="ew", padx=2)
            parent.columnconfigure(idx, weight=1)

    # ---------------- Plot ----------------

    def _build_plot(self):
        fig = Figure(figsize=(6.4, 5.4), dpi=100)
        ax = fig.add_subplot(111, projection="3d")
        self.fig = fig
        self.ax = ax

        canvas = FigureCanvasTkAgg(fig, master=self.plot_frame)
        canvas.draw()
        canvas.get_tk_widget().pack(fill="both", expand=True)

        toolbar = NavigationToolbar2Tk(canvas, self.plot_frame)
        toolbar.update()

        self.canvas = canvas

    def _draw_axes(self, ax, origin, basis, label_prefix="A", alpha=1.0):
        o = origin
        for idx, name in enumerate(["x", "y", "z"]):
            d = basis[:, idx]
            ax.quiver(o[0], o[1], o[2], d[0], d[1], d[2], length=1.0, normalize=True, alpha=alpha)
            tip = o + d
            ax.text(tip[0], tip[1], tip[2], f"{label_prefix}{name}", alpha=alpha)

    def _finalise_axes(self, ax):
        ax.set_xlim(-1.2, 1.2)
        ax.set_ylim(-1.2, 1.2)
        ax.set_zlim(-1.2, 1.2)
        ax.set_xlabel("X")
        ax.set_ylabel("Y")
        ax.set_zlabel("Z")
        ax.view_init(elev=20, azim=35)

    def _update_plot_from_result(self, q: Quaternion | None):
        """
        Plotting rules:
        - DO NOT normalise q for plotting
        - DO NOT do q*p*q̄ (no quaternion multiplies needed for the plot)
        - Always draw R axes, even in 'impure' cases
        - Overlay warning when (rule) R >= 1e-6
        """
        ax = self.ax
        ax.clear()

        # World axes
        self._draw_axes(ax, origin=np.zeros(3), basis=np.eye(3), label_prefix="W", alpha=0.25)

        if q is None:
            ax.set_title("No result quaternion yet")
            self._finalise_axes(ax)
            self.canvas.draw()
            return

        # RED vector arrow for vector part (i, j, k)
        v = np.array([q.x, q.y, q.z], dtype=float)
        if float(np.linalg.norm(v)) > 0:
            ax.quiver(
                0, 0, 0,
                v[0], v[1], v[2],
                length=1.0,
                #normalize=True,
                normalize=False,
                alpha=0.95,
                color="red",
            )
            ax.text(v[0], v[1], v[2], "v=(i,j,k)", alpha=0.95)
        else:
            ax.text(0, 0, 0, "v=(0,0,0)", alpha=0.95)

        # Degenerate quaternion (can't form meaningful axes)
        if q.norm() < EPS_NORM:
            ax.text2D(0.03, 0.95, "⚠ Quaternion norm ~0 — cannot show axes", transform=ax.transAxes)
            ax.set_title("Degenerate quaternion")
            self._finalise_axes(ax)
            self.canvas.draw()
            return

        # Result axes from raw quaternion -> matrix (no normalisation)
        Rm = quat_to_matrix_no_normalise(q)
        basis = np.column_stack([Rm[:, 0], Rm[:, 1], Rm[:, 2]])
        self._draw_axes(ax, origin=np.zeros(3), basis=basis, label_prefix="R", alpha=0.9)

        # Warning overlay per your rule: impure if R >= 1e-6
        if q.is_impure(EPS_IMPURE_R):
            ax.text2D(
                0.03, 0.95,
                "⚠ IMPURE (rule: |R| ≥ 1e-6)\nAxes drawn from RAW quaternion (no normalisation)",
                transform=ax.transAxes
            )
            ax.set_title("Result axes + IMPURE warning")
        else:
            ax.set_title("Result axes (raw quaternion; no normalisation)")

        self._finalise_axes(ax)
        self.canvas.draw()

    # ---------------- Actions ----------------

    def parse_quat(self, vars4, name: str) -> Quaternion:
        try:
            vals = [float(v.get().strip()) for v in vars4]
            return Quaternion(*vals)
        except ValueError:
            raise ValueError(f"{name}: please enter numeric values for R, i, j, k.")

    def set_result(self, q: Quaternion):
        for var, val in zip(self.qr_vars, q.as_tuple()):
            var.set(f"{val:.12g}")
        self._update_plot_from_result(q)

    def clear_result(self):
        for var in self.qr_vars:
            var.set("")
        self.status_var.set("Ready.")
        self._update_plot_from_result(None)

    def copy_qr_to_qa(self):
        if not self.qr_vars[0].get().strip():
            return
        for src, dst in zip(self.qr_vars, self.qa_vars):
            dst.set(src.get())
        self.status_var.set("Copied QR → QA")

    def copy_qr_to_qb(self):
        if not self.qr_vars[0].get().strip():
            return
        for src, dst in zip(self.qr_vars, self.qb_vars):
            dst.set(src.get())
        self.status_var.set("Copied QR → QB")

    def on_conjugate_qa(self):
        try:
            qa = self.parse_quat(self.qa_vars, "QA")
            qc = qa.conjugate()
            for var, val in zip(self.qa_vars, qc.as_tuple()):
                var.set(f"{val:.12g}")
            self.status_var.set("QA conjugated.")
        except Exception as e:
            messagebox.showerror("Error", str(e))

    def on_op(self, op: str):
        try:
            qa = self.parse_quat(self.qa_vars, "QA")
            qb = self.parse_quat(self.qb_vars, "QB")

            if op == "+":
                qr = qa + qb
            elif op == "-":
                qr = qa - qb
            elif op == "*":
                qr = qa * qb
            elif op == "/":
                qr = qa / qb
            else:
                raise ValueError("Unknown operation")

            self.set_result(qr)
            self.status_var.set(f"Computed QR = QA {op} QB")
        except Exception as e:
            messagebox.showerror("Error", str(e))


if __name__ == "__main__":
    App().mainloop()