import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0" 
import cv2
import torch
import torchvision.transforms.functional as F
import torchvision.transforms as T
from torchvision.models.optical_flow import raft_large, raft_small
import numpy as np
import time
import orisys

def preprocess(batch, resize=None):
    """
    Preprocess image batch for RAFT model.
    
    Args:
        batch: Tensor of shape (N, C, H, W)
        resize: Optional tuple (height, width) to resize. If None, keeps original size.
    
    Returns:
        Preprocessed tensor
    """
    transform_list = [
        T.ConvertImageDtype(torch.float32),
        T.Normalize(mean=0.5, std=0.5),  # map [0, 1] into [-1, 1]
    ]
    
    if resize is not None:
        transform_list.append(T.Resize(size=resize))
    
    transforms = T.Compose(transform_list)
    batch = transforms(batch)
    return batch
    
def setup_simple_raft():
    """
    Setup a simpler version using available models
    """
    model = raft_large(pretrained=True, progress=False).to('cuda')
    model = model.eval()
    # Alternative: Use LiteFlowNet (lighter alternative)
    print("RAFT loaded successfully")
    return model

def track_raft(prev_frame, frame, model, device):

    gray = frame.copy()
    start_time = time.time()
    prev_tensor = torch.from_numpy(prev_frame.transpose(2, 0, 1).astype(np.float32) / 255.0)
    curr_tensor = torch.from_numpy(gray.transpose(2, 0, 1).astype(np.float32) / 255.0)
    prev_tensor = preprocess(prev_tensor.unsqueeze(0), resize=None).to(device)
    curr_tensor = preprocess(curr_tensor.unsqueeze(0), resize=None).to(device)
    with torch.no_grad():
        flow = model(prev_tensor, curr_tensor)
        predicted_flows = flow[-1]

    flow_np = predicted_flows[0].cpu().numpy()  # or .squeeze(0)
    
    # Transpose to (H, W, 2) format expected by draw_arrows
    flow_np = flow_np.transpose(1, 2, 0)  # (2, H, W) -> (H, W, 2)
    return flow_np

def main_simple():
    # Initialize camera
    # sensor = orisys.Sensor('s1.mp4', config_name="./config/prototype_of3.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True)
    model = setup_simple_raft()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
    
    prev_gray = None
    fps_buffer = []

    
    while True:

        ret, frame = cap.read()
        if not ret:
            break
        dx = 300
        dy = 0
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame1 = frame[420:660, 840:1080, :]
        frame2 = frame[420+dy:660+dy, 840+dx:1080+dx, :]
        frame3 = frame[420-dy:660-dy, 840-dx:1080-dx, :]

        frame = np.concatenate([frame3, frame1, frame2], axis=1)
        # cv2.imshow("frame", frame3)
        # cv2.waitKey(0)
        if prev_gray is None:
            prev_gray = frame.copy()
        gray = frame
        # Compute optical flow
        start_time = time.time()
        flow_np = track_raft(prev_gray, gray, model, device)
        fps = 1.0 / (time.time() - start_time)
        fps_buffer.append(fps)
        if len(fps_buffer) > 30:
            fps_buffer.pop(0)
            
            avg_fps = np.mean(fps_buffer)
            print(f"avg_fps: {avg_fps}")
            
            
            # # Remove batch dimension: (1, 2, H, W) -> (2, H, W)
            # flow_np = predicted_flows[0].cpu().numpy()  # or .squeeze(0)
            
            # # Transpose to (H, W, 2) format expected by draw_arrows
            # flow_np = flow_np.transpose(1, 2, 0)  # (2, H, W) -> (H, W, 2)

            # Flow is already in pixel coordinates, no denormalization needed
            # Keep as float32 for draw_arrows (it can handle float values)
            print(f"Flow numpy shape: {flow_np.shape}, dtype: {flow_np.dtype}")
            print(f"Flow range: x=[{flow_np[:,:,0].min():.2f}, {flow_np[:,:,0].max():.2f}], "
                  f"y=[{flow_np[:,:,1].min():.2f}, {flow_np[:,:,1].max():.2f}]")

            arrows = orisys.util.draw_arrows(
                gray,
                flow_np,
                threshold=2,
                grid_spacing=40,
                arrow_scale=1.0
            )
            
            arrows = cv2.resize(arrows, (1200, 400))
            cv2.imshow('Optical Flow', arrows)
        
        # prev_gray = gray.copy()
        
        #步骤 3.5: 键盘控制
        key = cv2.waitKey(1) & 0xFF
        if key == ord("q"):
            print("\n退出程序...")
            break
        if key == ord("r"):
            # 重置光流追踪器（当追踪漂移时使用）
            prev_gray = None
            print("传感器已重置")
    
if __name__ == "__main__":
    main_simple()