#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
session_client.py — 电脑端：一次运行 = 一个会话。

流程全自动，不再手动 scp 图片/坐标：
  1. 启动 -> hello（握手）-> capture（板子抓帧，只回传图像）
  2. 图像显示在窗口里；鼠标点一个像素
  3. 只把 (u, v) 发给板子；板子用本地缓存的同帧点云算
     p_cam / p_lidar / p_map，并把 p_map 记成目标（点云不出板）
  4. 你移动板子，界面每 200 ms 问一次 status，实时显示 Δ 坐标（只给坐标，不给距离）

  需要电脑端算（要看候选点/平面拟合细节、离线复盘）时：
     python session_client.py --board 192.168.43.50 --compute pc
  这时才会把同帧点云拉回电脑，算完把 p_map 发回板子。

用法（电脑）：
  python session_client.py --board 192.168.43.50
  python session_client.py --board 192.168.43.50 --radius 12
  python session_client.py --board 192.168.43.50 --no-gui --u 800 --v 648 --watch

依赖：pynng、numpy、Pillow（Tkinter 随 Python 自带）
"""
import argparse
import csv
import io
import os
import socket
import time
import zlib

import numpy as np
from PIL import Image, ImageTk
import tkinter as tk

import pynng

import target_math
import wire

MAX_W, MAX_H = 1100, 800


# --------------------------------------------------------------------- 网络
def detect_local_ip(board_ip, port):
    """问内核：去板子的路由走哪张网卡，返回本机在这个网段上的 IP。"""
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect((board_ip, port))   # UDP connect 不发包，只选路
        return s.getsockname()[0]
    finally:
        s.close()


class BoardClient:
    def __init__(self, board_ip, port, timeout_ms=8000):
        self.board_ip = board_ip
        self.port = int(port)
        self.client_ip = detect_local_ip(board_ip, self.port)
        self.session = wire.new_session(board_ip, self.client_ip)
        self.sock = pynng.Req0(dial='tcp://%s:%d' % (board_ip, self.port))
        self.sock.recv_max_size = 32 * 1024 * 1024
        self.sock.send_timeout = 5000
        self.sock.recv_timeout = int(timeout_ms)

    def request(self, typ, blob=b'', **fields):
        msg = {'session': self.session, 'type': typ, 'ts_ns': time.time_ns()}
        msg.update(fields)
        self.sock.send(wire.pack(msg, blob))
        hdr, out_blob = wire.unpack(self.sock.recv())
        wire.check_session(hdr, self.session)
        return hdr, out_blob

    def new_session(self, note=''):
        """开始一个新会话（新的 id，仍然带同一个 client_ip / board_ip）。"""
        self.session = wire.new_session(self.board_ip, self.client_ip, note)
        return self.session

    def close(self):
        try:
            self.sock.close()
        except Exception:
            pass


# --------------------------------------------------------------------- 计算
def quat_to_rot(q):
    x, y, z, w = [float(v) for v in q]
    return np.array([
        [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
        [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
        [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)],
    ], dtype=np.float64)


def decode_cloud(hdr, blob):
    if not hdr.get('ok'):
        raise RuntimeError(hdr.get('error', 'cloud 请求失败'))
    if hdr.get('encoding') != 'zlib_f32':
        raise RuntimeError('未知点云编码: %s' % hdr.get('encoding'))
    raw = zlib.decompress(blob)
    n = int(hdr.get('n_points', 0))
    return np.frombuffer(raw, dtype='<f4').reshape(n, 3)


def compute_target(xyz, odom, u, v, radius, fit_plane, index=None):
    """像素 -> (locate 结果, map 系坐标)。u/v 用标定分辨率(1600x1296)。"""
    if index is not None:
        res = index.locate(u, v, radius=radius, fit_plane=fit_plane)
    else:
        res = target_math.locate(xyz, u, v, radius=radius, fit_plane=fit_plane)
    p_map = None
    if res['found'] and odom:
        R = quat_to_rot(odom['orientation_xyzw'])
        t = np.asarray(odom['position'], dtype=np.float64)
        p_map = R @ np.asarray(res['p_lidar'], dtype=np.float64) + t
    return res, p_map


def fmt3(p):
    return '(%.4f, %.4f, %.4f)' % (float(p[0]), float(p[1]), float(p[2]))


def result_text(res, p_map, u, v):
    if not res['found']:
        return ('像素(标定系): (%.1f, %.1f)\n\n'
                '未找到点云（候选 %d）\n'
                '可能：无回波 / 超出视场 / 半径太小\n'
                '试试把半径调大（如 15~20），或点平整、点云密的区域'
                % (u, v, res.get('n_cand', 0)))

    lines = [
        '像素(标定系): (%.1f, %.1f)' % (u, v),
        '命中像素: (%.1f, %.1f)  偏离 %.2f px' % (res['hit_u'], res['hit_v'], res['px_err']),
        '候选数: %d' % res['n_cand'],
        '平面拟合: %s' % ('是' if res['plane_used'] else '否'),
    ]
    if res.get('multi_layer'):
        lines.append('! 候选跨多个深度层，可能擦过边缘')
    if res.get('confidence') is not None:
        lines.append('置信度: %.0f' % res['confidence'])
    lines += [
        '',
        'p_cam（相机系）  : %s' % fmt3(res['p_cam']),
        'p_lidar（雷达系）: %s' % fmt3(res['p_lidar']),
    ]
    if p_map is not None:
        lines.append('★ p_map（地图系）: %s' % fmt3(p_map))
        lines.append('  （map 原点 = 建图初始位置）')
    else:
        lines.append('p_map: 不可用（本帧没有 SLAM 位姿）')
    return '\n'.join(lines)


# --------------------------------------------------------------------- 界面
class App:
    def __init__(self, root, cli, args):
        self.root = root
        self.cli = cli
        self.args = args
        self.image = None
        self.img_w = self.img_h = 0
        self.calib_w = self.calib_h = 0
        self.scale = 1.0
        self.xyz = None
        self.index = None
        self.capture_odom = None
        self.target = None
        self.poll_job = None
        self.compute_mode = args.compute
        try:
            hello, _ = self.cli.request('hello')
            if not hello.get('board_can_compute'):
                if self.compute_mode == 'board':
                    print('板子不能本地计算，自动切到 --compute pc')
                self.compute_mode = 'pc'
        except Exception as exc:
            print('hello 失败: %s' % exc)
        self._build_ui()
        self.capture()

    # ---------------------------------------------------------------- UI
    def _build_ui(self):
        self.root.title('灵枢 · 会话式精度测试')

        left = tk.Frame(self.root)
        left.pack(side=tk.LEFT, padx=8, pady=8)
        right = tk.Frame(self.root)
        right.pack(side=tk.RIGHT, fill=tk.Y, padx=8, pady=8)

        self.canvas = tk.Canvas(left, width=MAX_W, height=MAX_H,
                                cursor='crosshair', bg='#111')
        self.canvas.pack()
        self.canvas.bind('<Button-1>', self.on_click)

        bar = tk.Frame(right)
        bar.pack(anchor='w', fill=tk.X)
        tk.Button(bar, text='新会话', command=self.new_session).pack(side=tk.LEFT)
        tk.Button(bar, text='重新抓帧', command=self.capture).pack(side=tk.LEFT, padx=6)
        tk.Button(bar, text='清除目标', command=self.reset_target).pack(side=tk.LEFT)

        pf = tk.Frame(right)
        pf.pack(anchor='w', pady=6)
        tk.Label(pf, text='候选半径(px):').pack(side=tk.LEFT)
        self.radius_var = tk.StringVar(value=str(self.args.radius))
        tk.Entry(pf, textvariable=self.radius_var, width=6).pack(side=tk.LEFT)
        self.fit_var = tk.BooleanVar(value=not self.args.no_fit)
        tk.Checkbutton(pf, text='平面拟合', variable=self.fit_var).pack(side=tk.LEFT, padx=8)

        self.text = tk.Text(right, width=58, height=18, font=('Consolas', 10))
        self.text.pack(fill=tk.BOTH, expand=True)

        tk.Label(right, text='实时位姿 / 误差（只给坐标）',
                 anchor='w').pack(fill=tk.X, pady=(8, 0))
        self.status_var = tk.StringVar(value='等待抓帧…')
        tk.Label(right, textvariable=self.status_var, justify='left', anchor='w',
                 font=('Consolas', 12), fg='#0a7').pack(fill=tk.X)

        self.info_var = tk.StringVar(value='')
        tk.Label(right, textvariable=self.info_var, justify='left', anchor='w',
                 font=('Consolas', 9), fg='#666').pack(fill=tk.X, pady=(6, 0))

    def set_text(self, s):
        self.text.configure(state=tk.NORMAL)
        self.text.delete('1.0', tk.END)
        self.text.insert(tk.END, s)
        self.text.configure(state=tk.DISABLED)

    def set_status(self, s):
        self.status_var.set(s)

    # ---------------------------------------------------------------- 抓帧
    def capture(self):
        self.cancel_poll()
        self.target = None
        self.set_status('正在抓帧…')
        self.root.update_idletasks()

        try:
            h, img_blob = self.cli.request('capture')
        except Exception as exc:
            self.set_status('抓帧失败: %s' % exc)
            return
        if not h.get('ok'):
            self.set_status('抓帧失败: %s' % h.get('error'))
            return

        try:
            im = Image.open(io.BytesIO(img_blob)).convert('RGB')
        except Exception as exc:
            self.set_status('图像解码失败: %s' % exc)
            return

        self.image = im
        self.img_w, self.img_h = im.size
        self.calib_w = int(h.get('calib_width') or self.img_w)
        self.calib_h = int(h.get('calib_height') or self.img_h)

        # 默认板子端算，点云不出板；只有 --compute pc 才拉同帧点云
        self.xyz = None
        self.index = None
        cloud_err = None
        if self.compute_mode == 'pc':
            try:
                hc, cblob = self.cli.request('cloud')
                self.xyz = decode_cloud(hc, cblob)
            except Exception as exc:
                cloud_err = str(exc)

        self.capture_odom = h.get('odom')

        # 显示
        self.scale = min(MAX_W / self.img_w, MAX_H / self.img_h, 1.0)
        disp = im.resize((max(1, int(self.img_w * self.scale)),
                          max(1, int(self.img_h * self.scale))), Image.LANCZOS)
        self.tkimg = ImageTk.PhotoImage(disp)
        self.canvas.config(width=disp.width, height=disp.height)
        self.canvas.delete('all')
        self.canvas.create_image(0, 0, anchor='nw', image=self.tkimg)

        odom = self.capture_odom
        info = [
            '会话     : %s' % self.cli.session['id'][:8],
            '链路     : %s -> %s:%d' % (self.cli.client_ip, self.cli.board_ip, self.cli.port),
            '图像     : %dx%d（标定 %dx%d）' % (self.img_w, self.img_h,
                                               self.calib_w, self.calib_h),
            '计算     : %s' % ('板子端（点云不出板）' if self.compute_mode == 'board'
                                else '电脑端（已拉回点云）'),
            '点云     : %s' % (('%d 点（已拉回）' % self.xyz.shape[0])
                                if self.xyz is not None
                                else ('留在板子' if self.compute_mode == 'board'
                                      else '获取失败: %s' % cloud_err)),
            '图像/点云: %.2f ms' % float(h.get('stamp_delta_ms') or 0.0),
        ]
        if odom:
            info.append('抓帧位姿 : %s' % fmt3(odom['position']))
        else:
            info.append('抓帧位姿 : 无（板子没给 odom）')
        info.append('')
        info.append('点击图像选择目标点…')
        self.set_text('\n'.join(info))
        self.info_var.set('会话 %s | %s -> %s' % (
            self.cli.session['id'][:8], self.cli.client_ip, self.cli.board_ip))
        self.set_status('已抓帧，请点选目标')

    # ---------------------------------------------------------------- 点选
    def on_click(self, event):
        if self.image is None:
            self.set_status('还没有图像，先抓帧')
            return
        if self.compute_mode == 'pc' and self.xyz is None:
            self.set_status('点云没拉到，无法本地计算；可改用默认的板子端计算')
            return

        u_img = event.x / self.scale
        v_img = event.y / self.scale
        u_cal = u_img * self.calib_w / self.img_w
        v_cal = v_img * self.calib_h / self.img_h

        try:
            radius = float(self.radius_var.get())
        except ValueError:
            radius = self.args.radius
        fit_plane = bool(self.fit_var.get())

        if self.compute_mode == 'board':
            # 点云不出板：只发像素，板子算完把 p_map 记成目标
            try:
                res, _ = self.cli.request('click', u=u_cal, v=v_cal,
                                          radius=radius, fit_plane=fit_plane)
            except Exception as exc:
                self.set_status('点选失败: %s' % exc)
                return
            p_map = res.get('p_map')
        else:
            # 电脑端算：先拉点云，本地算，再把 p_map 回传板子
            try:
                if self.index is None:
                    self.index = target_math.TargetIndex(self.xyz)
                res, p_map = compute_target(self.xyz, self.capture_odom,
                                            u_cal, v_cal, radius, fit_plane,
                                            index=self.index)
            except Exception as exc:
                self.set_status('计算失败: %s' % exc)
                return
            if res['found'] and p_map is not None:
                try:
                    h, _ = self.cli.request(
                        'set_target', u=u_cal, v=v_cal,
                        p_cam=[float(x) for x in res['p_cam']],
                        p_lidar=[float(x) for x in res['p_lidar']],
                        p_map=[float(x) for x in p_map])
                except Exception as exc:
                    self.set_status('回传目标失败: %s' % exc)
                    return
                if not h.get('ok'):
                    self.set_status('板子拒绝目标: %s' % h.get('error'))
                    return

        self.draw_marker(event.x, event.y)
        self.set_text(result_text(res, p_map, u_cal, v_cal))

        if not res.get('found'):
            self.set_status('没找到点，换个位置或调大半径')
            return
        if p_map is None:
            self.set_status('找到了点，但本帧没有位姿，无法跟踪')
            return

        self.target = p_map
        self.set_status('目标已锁定，移动板子…')
        self.log_row(u_cal, v_cal, res, p_map)
        self.schedule_poll()

    def draw_marker(self, x, y):
        self.canvas.delete('marker')
        self.canvas.create_oval(x - 7, y - 7, x + 7, y + 7,
                                outline='#00ff66', width=2, tags='marker')
        self.canvas.create_line(x - 15, y, x + 15, y, fill='#00ff66', tags='marker')
        self.canvas.create_line(x, y - 15, x, y + 15, fill='#00ff66', tags='marker')

    # ---------------------------------------------------------------- 轮询
    def schedule_poll(self):
        self.cancel_poll()
        self.poll_job = self.root.after(self.args.poll_ms, self.poll_status)

    def cancel_poll(self):
        if self.poll_job is not None:
            try:
                self.root.after_cancel(self.poll_job)
            except Exception:
                pass
            self.poll_job = None

    def poll_status(self):
        self.poll_job = None
        if self.target is None:
            return
        try:
            h, _ = self.cli.request('status')
        except Exception as exc:
            self.set_status('status 失败: %s' % exc)
            self.schedule_poll()
            return

        if not h.get('ok') or not h.get('valid'):
            self.set_status('等待里程计…（%s）' % h.get('error', '位姿停更'))
        else:
            pos = h['pos']
            lines = ['pos=(%+.4f, %+.4f, %+.4f)  yaw=%+.1f°'
                     % (pos[0], pos[1], pos[2], h.get('yaw_deg', 0.0))]
            d = h.get('delta')
            if d:
                lines.append('Δ=(%+.4f, %+.4f, %+.4f)' % (d[0], d[1], d[2]))
            self.set_status('\n'.join(lines))
        self.schedule_poll()

    # ---------------------------------------------------------------- 其它
    def reset_target(self):
        self.cancel_poll()
        self.target = None
        self.set_status('已清除目标')
        try:
            self.cli.request('reset')
        except Exception:
            pass

    def new_session(self):
        """开一个新会话：新 id + 重新抓帧（每次测试互不干扰）。"""
        self.cancel_poll()
        self.target = None
        self.cli.new_session()
        self.capture()

    def log_row(self, u, v, res, p_map):
        if self.args.no_log:
            return
        path = self.args.log
        new = not os.path.exists(path)
        with open(path, 'a', newline='', encoding='utf-8-sig') as f:
            w = csv.writer(f)
            if new:
                w.writerow(['session_id', 'client_ip', 'board_ip',
                            'u_cal', 'v_cal',
                            'p_cam_x', 'p_cam_y', 'p_cam_z',
                            'p_lidar_x', 'p_lidar_y', 'p_lidar_z',
                            'p_map_x', 'p_map_y', 'p_map_z',
                            'plane_used', 'n_cand', 'px_err', 'ts_ns'])
            w.writerow([
                self.cli.session['id'], self.cli.client_ip, self.cli.board_ip,
                round(u, 2), round(v, 2)]
                + [round(float(res['p_cam'][i]), 4) for i in range(3)]
                + [round(float(res['p_lidar'][i]), 4) for i in range(3)]
                + [round(float(p_map[i]), 4) for i in range(3)]
                + [int(res['plane_used']), int(res['n_cand']),
                   round(float(res['px_err']), 2), time.time_ns()])


# --------------------------------------------------------------------- 无界面
def print_status(h):
    if not h.get('ok'):
        print('status: %s' % h.get('error'))
        return
    pos = h.get('pos')
    line = ('pos=(%+.4f, %+.4f, %+.4f)  yaw=%+.1f°'
            % (pos[0], pos[1], pos[2], h.get('yaw_deg', 0.0))) if pos else 'pos=无'
    if not h.get('valid'):
        line += '  [位姿无效]'
    d = h.get('delta')
    if d:
        line += '\nΔ=(%+.4f, %+.4f, %+.4f)' % (d[0], d[1], d[2])
    print(line, flush=True)


def run_cli(cli, args):
    hello, _ = cli.request('hello')
    mode = args.compute
    if mode == 'board' and not hello.get('board_can_compute'):
        print('板子不能本地计算，自动切到 --compute pc', flush=True)
        mode = 'pc'
    print('hello: ok=%s node=%s calib=%sx%s compute=%s'
          % (hello.get('ok'), hello.get('node'), hello.get('calib_width'),
             hello.get('calib_height'), mode), flush=True)

    h, img_blob = cli.request('capture')
    if not h.get('ok'):
        raise SystemExit('capture 失败: %s' % h.get('error'))
    print('capture: %d B, %sx%s, 点云 %d 点, 时间差 %.2f ms'
          % (len(img_blob), h.get('image_width'), h.get('image_height'),
             h.get('n_points', 0), h.get('stamp_delta_ms', 0.0)), flush=True)
    if h.get('odom'):
        print('抓帧位姿: %s' % fmt3(h['odom']['position']), flush=True)

    if mode == 'board':
        res, _ = cli.request('click', u=args.u, v=args.v,
                             radius=args.radius, fit_plane=not args.no_fit)
        p_map = res.get('p_map')
        print('\n' + result_text(res, p_map, args.u, args.v) + '\n', flush=True)
        if not (res.get('found') and p_map is not None):
            return
    else:
        hc, cblob = cli.request('cloud')
        xyz = decode_cloud(hc, cblob)
        print('cloud: %d 点 (%.0f KB -> %.0f KB)'
              % (xyz.shape[0], hc.get('raw_bytes', 0) / 1024.0,
                 hc.get('zlib_bytes', 0) / 1024.0), flush=True)

        idx = target_math.TargetIndex(xyz)
        res, p_map = compute_target(xyz, h.get('odom'), args.u, args.v,
                                    args.radius, not args.no_fit, index=idx)
        print('\n' + result_text(res, p_map, args.u, args.v) + '\n', flush=True)
        if not (res['found'] and p_map is not None):
            return

        r, _ = cli.request('set_target', u=args.u, v=args.v,
                           p_cam=[float(x) for x in res['p_cam']],
                           p_lidar=[float(x) for x in res['p_lidar']],
                           p_map=[float(x) for x in p_map])
        print('set_target: ok=%s target=%s' % (r.get('ok'), r.get('target')), flush=True)

    if args.watch:
        n = 0
        while args.watch_count is None or n < args.watch_count:
            st, _ = cli.request('status')
            print_status(st)
            n += 1
            if args.watch_count is None or n < args.watch_count:
                time.sleep(args.poll_ms / 1000.0)


# --------------------------------------------------------------------- 入口
def parse_args():
    ap = argparse.ArgumentParser(description='会话式精度测试客户端')
    ap.add_argument('--board', required=True, help='板子 IP，例如 192.168.43.50')
    ap.add_argument('--port', type=int, default=5555)
    ap.add_argument('--radius', type=float, default=10.0, help='候选像素半径')
    ap.add_argument('--compute', choices=['board', 'pc'], default='board',
                    help='board=点云不出板（默认）；pc=拉点云到电脑算')
    ap.add_argument('--no-fit', action='store_true', help='关闭平面拟合')
    ap.add_argument('--no-gui', action='store_true')
    ap.add_argument('--u', type=float, default=None, help='标定系像素 u（配合 --no-gui）')
    ap.add_argument('--v', type=float, default=None, help='标定系像素 v（配合 --no-gui）')
    ap.add_argument('--watch', action='store_true', help='点选后持续打印 Δ')
    ap.add_argument('--watch-count', type=int, default=None)
    ap.add_argument('--poll-ms', type=int, default=200)
    ap.add_argument('--timeout-ms', type=int, default=8000)
    ap.add_argument('--log', default='session_log.csv')
    ap.add_argument('--no-log', action='store_true')
    return ap.parse_args()


def main():
    args = parse_args()
    cli = BoardClient(args.board, args.port, args.timeout_ms)
    try:
        if args.no_gui:
            if args.u is None or args.v is None:
                raise SystemExit('--no-gui 需要同时给 --u 和 --v（标定系像素）')
            run_cli(cli, args)
        else:
            root = tk.Tk()
            App(root, cli, args)
            try:
                root.mainloop()
            finally:
                cli.close()
    finally:
        cli.close()


if __name__ == '__main__':
    main()