"""
示例 3：实时曲线与幅值图显示

本示例演示如何在读取传感器结果的同时：
1. 显示形变矢量场
2. 显示流场幅值热力图
3. 使用 PyQtGraph 实时绘制法向力与切向力曲线
4. 按键控制录像开始与结束

SDK version: 0.3.1
"""

import os
import sys
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
import cv2
import orisys
import numpy as np
import time
import argparse
from datetime import datetime
try:
    import pyqtgraph as pg
except:
    print("PyQtGraph is not installed. Please install it using 'pip install pyqt5 pyqtgraph'.")
    exit()

def create_plot():
    p1 = win.addPlot(title="法向力")
    p1.setLabel('bottom', '时间', units='s')
    win.nextRow()
    p2 = win.addPlot(title="切向力 X")
    p2.setLabel('bottom', '时间', units='s')
    win.nextRow()
    p3 = win.addPlot(title="切向力 Y")
    p3.setLabel('bottom', '时间', units='s')
    plot_af = p1.plot(pen='k')
    plot_tfx = p2.plot(pen='k')
    plot_tfy = p3.plot(pen='k')
    return p1, p2, p3, plot_af, plot_tfx, plot_tfy

def update(plot_af, plot_tfx, plot_tfy, data, p1, p2, p3):
    plot_af.setData(data[:,:,0])
    p1.setYRange(min(data[:,1,0]),max(data[:,1,0]))
    plot_tfx.setData(data[:,:,1])
    p2.setYRange(min(data[:,1,1]),max(data[:,1,1]))
    plot_tfy.setData(data[:,:,2])
    p3.setYRange(min(data[:,1,2]),max(data[:,1,2]))

# 深度图
def draw_magnitude_map(vfield, threshold=3, colormap=cv2.COLORMAP_JET) -> np.ndarray:
        """
        绘制流场幅值伪彩色图。
        
        Args:
            vfield: 需要可视化的二维矢量场
            threshold: 最小显示阈值；低于阈值的区域显示为低值颜色
            colormap: OpenCV 颜色映射类型（默认 COLORMAP_JET）
        
        Returns:
            np.ndarray: 400x400 的 BGR 伪彩图像
        """
        flow_field = vfield
            
        # 计算每个像素点的幅值
        magnitude = np.sqrt(flow_field[:, :, 0]**2 + flow_field[:, :, 1]**2)
        
        # 先降采样再插值回原尺寸，使显示更平滑
        downsample_factor = 10
        h, w = magnitude.shape
        magnitude_downsampled = magnitude[::downsample_factor, ::downsample_factor]
        magnitude_smoothed = cv2.resize(magnitude_downsampled, (w, h), interpolation=cv2.INTER_CUBIC)
        
        # 应用阈值，将低于阈值的区域置零
        magnitude_thresholded = magnitude_smoothed.copy()
        magnitude_thresholded[magnitude_smoothed < threshold] = 0

        # 归一化到 0-255
        mag_min, mag_max = magnitude_thresholded.min(), magnitude_thresholded.max()
        if mag_max > mag_min:
            magnitude_norm = ((magnitude_thresholded - mag_min) / (mag_max - mag_min) * 255).astype(np.uint8)
        else:
            magnitude_norm = np.zeros_like(magnitude_thresholded, dtype=np.uint8)

        # 应用伪彩色映射
        magnitude_colored = cv2.applyColorMap(magnitude_norm, colormap)
        
        # 缩放到固定输出尺寸
        out_img = cv2.resize(magnitude_colored, (400, 400), interpolation=cv2.INTER_CUBIC)

        # 添加边距，使图像在窗口中居中显示
        # margin = 30
        # out_img = cv2.copyMakeBorder(out_img, margin, margin, margin, margin,
        #                              cv2.BORDER_CONSTANT, value=(22, 16, 11))

        return out_img

# 实时图表
npoints = 100 # 曲线缓冲区长度
pg.setConfigOption('background', 'w')
pg.setConfigOption('foreground', 'k')
app = pg.mkQApp("Orisys Plot Example")
win = pg.GraphicsLayoutWidget(show=True, title="Orisys 实时力曲线")
win.setWindowTitle('Orisys 实时力曲线')
win.resize(640, 400)
p1, p2, p3, plot_af, plot_tfx, plot_tfy = create_plot()
data = np.zeros([npoints,2,3])
start_time = None  # Will be set when sensor starts
for i in range(3):
    data[:,0,i] = np.linspace(-npoints,0,npoints)



def main():
    global start_time
    # 将所有输出重定向到日志文件
    script_dir = os.path.dirname(os.path.abspath(__file__))
    log_dir = os.path.join(script_dir, "logs")
    os.makedirs(log_dir, exist_ok=True)
    log_path = os.path.join(log_dir, f"plot_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log")
    sys.stdout = open(log_path, 'w', encoding='utf-8')
    sys.stderr = sys.stdout

    parser = argparse.ArgumentParser()
    parser.add_argument("--video","-v", type=str, default="0", help="视频源：摄像头编号或视频文件路径")
    parser.add_argument("--config","-c", type=str, default="./config/ddjx01.json", help="配置文件名称或路径")
    parser.add_argument("--cal","-cal", type=str, default="./config/ddjx01.npy", help="标定文件路径")
    parser.add_argument("--verbose","-verbose", type=bool, default=True, help="是否输出详细日志")
    args = parser.parse_args()

    # 步骤 1：创建传感器对象（必需）
    # config_name 可以是内置配置名称，也可以是自定义配置文件路径
    # cal_path 为标定文件路径；每台设备建议使用独立的标定文件

    # 输入既可以是摄像头编号，也可以是视频文件路径
    try:
        input = int(args.video)
        is_camera = True
        sensor = orisys.Sensor(input, config_name=args.config, cal_path= args.cal, verbose=args.verbose)
    except:
        is_camera = False
        sensor = orisys.Sensor(args.video, config_name=args.config, cal_path= args.cal, verbose=args.verbose)
    # 下面保留了历史示例写法，便于参考
    # sensor1 = orisys.Sensor("./s1.mp4", config_name="./config/ddjx01.json", cal_path= "./config/ddjx01.npy", cuda=False, verbose=True)

    start_time = time.time()
    is_save = False
    frame_count = 0
    writer = None
    save_path = None
    print("\n按键说明：'q' 退出程序，'s' 开始录像，'e' 结束录像。")
    while True:
        sensor.get_img() # 获取并拼接复眼图像
        
        if is_save and writer is not None and sensor.frame is not None:
            frame_count += 1
            save_img = sensor.frame
            
            # 若为灰度图，则先转换为 BGR 再写入视频
            if len(save_img.shape) == 2 or (len(save_img.shape) == 3 and save_img.shape[2] == 1):
                if len(save_img.shape) == 3:
                    save_img = save_img.squeeze(2)
                save_img = cv2.cvtColor(save_img, cv2.COLOR_GRAY2BGR)
            
            writer.write(save_img)

        sensor.compute_deformation(check_motion=True, threshold=0) # 计算形变与分解结果
        # 读取信息：
        # 向量场：VRAW、VNORMAL、VSHEAR -> 原始 / 法向 / 切向
        # 标量：FNORMAL、FSHEARX、FSHEARY -> 法向力 / X 向切向力 / Y 向切向力
        fps, fn, fx, fy, flow, vnormal, vshear = sensor.read_info(sensor.info.FPS, sensor.info.FNORMAL, sensor.info.FSHEARX, sensor.info.FSHEARY, sensor.info.VRAW, sensor.info.VNORMAL,sensor.info.VSHEAR)
        print(f"FPS={fps:.2f}, 法向力={fn:.4f}, 切向力X={fx:.4f}, 切向力Y={fy:.4f}")
        # 显示形变矢量场
        # arrows  = orisys.util.draw_arrows(sensor.img, flow, threshold=5, grid_spacing=10, arrow_scale=1.0) 
        _bg_bgr = (22, 16, 11)
        _img = sensor.img
        _h, _w = _img.shape[:2]
        black_img = np.empty((_h, _w, 3), dtype=_img.dtype)
        black_img[:, :, 0] = _bg_bgr[0]
        black_img[:, :, 1] = _bg_bgr[1]
        black_img[:, :, 2] = _bg_bgr[2]
        arrows = orisys.util.draw_arrows(
            black_img,
            flow,
            threshold=2,
            grid_spacing=20,
            arrow_scale=0.5,
            below_threshold_color=(160, 160, 160),
            # target_size=int(side),
        )
        arrows = cv2.resize(arrows, (400, 400), interpolation=cv2.INTER_CUBIC)
        arrows = cv2.rotate(arrows, cv2.ROTATE_90_COUNTERCLOCKWISE)
        cv2.imshow("形变矢量场".encode("gbk"), arrows)
        # 显示流场幅值图
        mag = draw_magnitude_map(flow)
        mag = cv2.rotate(mag, cv2.ROTATE_90_COUNTERCLOCKWISE)
        cv2.imshow("流场幅值图".encode("gbk"), mag)

        # 更新实时曲线
        current_time = time.time() - start_time
        for i in range(3):
            data[:-1, 0, i] = data[1:, 0, i]  # 时间轴左移
            data[:-1, 1, i] = data[1:, 1, i]  # 力数据左移
        data[-1, 0, :] = current_time  # 写入当前时间
        data[-1, 1, 0] = fn
        data[-1, 1, 1] = fx
        data[-1, 1, 2] = fy
        update(plot_af, plot_tfx, plot_tfy, data, p1, p2, p3)  # 更新曲线

        key = cv2.waitKey(1)
        if key & 0xFF == ord("q"):
            print("\n正在退出示例程序...")
            break
        if key == ord("s"):
            save_path = f"./data/pulse_test.mp4"
            # 确保输出目录存在
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            # 从当前帧获取视频尺寸
            save_img = sensor.frame

            if save_img is not None:
                # 如为灰度图，先转换为 BGR 以获取视频尺寸
                temp_img = save_img.copy()
                if len(temp_img.shape) == 2 or (len(temp_img.shape) == 3 and temp_img.shape[2] == 1):
                    if len(temp_img.shape) == 3:
                        temp_img = temp_img.squeeze(2)
                    temp_img = cv2.cvtColor(temp_img, cv2.COLOR_GRAY2BGR)
                
                frame_height, frame_width = temp_img.shape[:2]
                writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), 30, (frame_width, frame_height))
                if writer.isOpened():
                    is_save = True
                    print(f"开始录像：{save_path}")
                else:
                    print(f"无法创建录像文件：{save_path}")
                    writer = None
            else:
                print("当前没有可用帧，无法确定录像尺寸。")
        if key == ord("e"):
            is_save = False
            if writer is not None:
                writer.release()
                writer = None
            if save_path is not None:
                print(f"录像结束：{save_path}，总帧数：{frame_count}")
            else:
                print(f"录像结束，总帧数：{frame_count}")
            frame_count = 0
    
    sensor.disconnect() # 断开传感器连接并释放资源
    


if __name__ == '__main__':
    main()
