import tkinter as tk
from tkinter import filedialog
import numpy as np

pixel_pts = None
phys_pts = None
angles = None
M = None
M_inv = None


def compute_affine(pixel_pts, phys_pts):
    num = pixel_pts.shape[0]
    A = np.zeros((2 * num, 6))
    B = np.zeros((2 * num, 1))

    for i in range(num):
        x, y = phys_pts[i]
        u, v = pixel_pts[i]
        A[2*i] = [x, y, 1, 0, 0, 0]
        B[2*i] = u
        A[2*i+1] = [0, 0, 0, x, y, 1]
        B[2*i+1] = v

    X, _, _, _ = np.linalg.lstsq(A, B, rcond=None)
    return np.array([
        [X[0,0], X[1,0], X[2,0]],
        [X[3,0], X[4,0], X[5,0]],
        [0, 0, 1]
    ])


def pixel_to_physical(uv, M_inv):
    p = M_inv @ np.array([uv[0], uv[1], 1.0])
    return p[0]/p[2], p[1]/p[2]


# ---------- GUI ----------
root = tk.Tk()
root.title("仿射标定（含逆矩阵）")
root.geometry("720x580")

def load_txt():
    global pixel_pts, phys_pts, angles, M, M_inv
    path = filedialog.askopenfilename(filetypes=[("TXT", "*.txt")])
    if not path:
        return

    data = np.loadtxt(path)
    pixel_pts = data[:, 0:2]
    phys_pts  = data[:, 2:4]
    angles    = data[:, 4]

    M = compute_affine(pixel_pts, phys_pts)
    M_inv = np.linalg.inv(M)

    txt.delete("1.0", tk.END)

    txt.insert("end", "✅ 标定完成\n\n")

    txt.insert("end", "【仿射矩阵 M（物理 → 像素）】\n")
    txt.insert("end", f"{M}\n\n")

    txt.insert("end", "【逆矩阵 M⁻¹（像素 → 物理）】\n")
    txt.insert("end", f"{M_inv}\n\n")

    txt.insert("end", f"旋转角示例（°）: {angles[:3]}\n")


def calc():
    try:
        u = float(e_u.get())
        v = float(e_v.get())
        x, y = pixel_to_physical([u, v], M_inv)
        lbl.config(text=f"物理坐标: x = {x:.4f}, y = {y:.4f}")
    except Exception as e:
        lbl.config(text=str(e))


# UI
tk.Button(root, text="加载 TXT 文件", command=load_txt).pack(pady=5)

f = tk.Frame(root)
f.pack()
tk.Label(f, text="u").grid(row=0, column=0)
e_u = tk.Entry(f, width=8); e_u.grid(row=0, column=1)
tk.Label(f, text="v").grid(row=0, column=2)
e_v = tk.Entry(f, width=8); e_v.grid(row=0, column=3)
tk.Button(f, text="计算物理坐标", command=calc).grid(row=0, column=4, padx=10)

lbl = tk.Label(root, text="物理坐标: -- , --")
lbl.pack(pady=5)

txt = tk.Text(root)
txt.pack(fill="both", expand=True, padx=5, pady=5)

root.mainloop()