import json
import numpy as np
import logging
from pathlib import Path
from config import TEST_DATA_DIR

logger = logging.getLogger(__name__)

def create_test_data_structure(progress_callback=None):
    """Create sample camera extrinsics data for testing"""
    
    if progress_callback:
        progress_callback(0.0, desc="Creating test data structure...")
    
    # Create directories
    data_dir = Path(f"{TEST_DATA_DIR}/cameras")
    videos_dir = Path(f"{TEST_DATA_DIR}/videos")
    data_dir.mkdir(parents=True, exist_ok=True)
    videos_dir.mkdir(parents=True, exist_ok=True)
    
    camera_file = data_dir / "camera_extrinsics.json"
    
    # Skip if file already exists
    if camera_file.exists():
        logger.info(f"✓ Camera extrinsics already exist at {camera_file}")
        
        if progress_callback:
            progress_callback(1.0, desc="Test data structure already exists")
            
        return
    
    if progress_callback:
        progress_callback(0.3, desc="Generating camera extrinsics data...")
    
    # Generate sample camera data
    camera_data = {}
    
    # Create 81 frames with 10 camera trajectories each
    for frame_idx in range(81):
        frame_key = f"frame{frame_idx}"
        camera_data[frame_key] = {}
        
        for cam_idx in range(1, 11):  # Camera types 1-10
            # Create a sample camera matrix (this is just an example - replace with actual logic if needed)
            # In reality, these would be calculated based on specific camera movement patterns
            
            # Create a base identity matrix
            base_matrix = np.eye(4)
            
            # Add some variation based on frame and camera type
            # This is a simplistic example - real camera movements would be more complex
            if cam_idx == 1:  # Pan Right
                base_matrix[0, 3] = 0.01 * frame_idx  # Move right over time
            elif cam_idx == 2:  # Pan Left
                base_matrix[0, 3] = -0.01 * frame_idx  # Move left over time
            elif cam_idx == 3:  # Tilt Up
                # Rotate around X-axis
                angle = 0.005 * frame_idx
                base_matrix[1, 1] = np.cos(angle)
                base_matrix[1, 2] = -np.sin(angle)
                base_matrix[2, 1] = np.sin(angle)
                base_matrix[2, 2] = np.cos(angle)
            elif cam_idx == 4:  # Tilt Down
                # Rotate around X-axis (opposite direction)
                angle = -0.005 * frame_idx
                base_matrix[1, 1] = np.cos(angle)
                base_matrix[1, 2] = -np.sin(angle)
                base_matrix[2, 1] = np.sin(angle)
                base_matrix[2, 2] = np.cos(angle)
            elif cam_idx == 5:  # Zoom In
                base_matrix[2, 3] = -0.01 * frame_idx  # Move forward over time
            elif cam_idx == 6:  # Zoom Out
                base_matrix[2, 3] = 0.01 * frame_idx  # Move backward over time
            elif cam_idx == 7:  # Translate Up (with rotation)
                base_matrix[1, 3] = 0.01 * frame_idx  # Move up over time
                angle = 0.003 * frame_idx
                base_matrix[0, 0] = np.cos(angle)
                base_matrix[0, 2] = np.sin(angle)
                base_matrix[2, 0] = -np.sin(angle)
                base_matrix[2, 2] = np.cos(angle)
            elif cam_idx == 8:  # Translate Down (with rotation)
                base_matrix[1, 3] = -0.01 * frame_idx  # Move down over time
                angle = -0.003 * frame_idx
                base_matrix[0, 0] = np.cos(angle)
                base_matrix[0, 2] = np.sin(angle)
                base_matrix[2, 0] = -np.sin(angle)
                base_matrix[2, 2] = np.cos(angle)
            elif cam_idx == 9:  # Arc Left (with rotation)
                angle = 0.005 * frame_idx
                radius = 2.0
                base_matrix[0, 3] = -radius * np.sin(angle)
                base_matrix[2, 3] = -radius * np.cos(angle) + radius
                # Rotate to look at center
                look_angle = angle + np.pi
                base_matrix[0, 0] = np.cos(look_angle)
                base_matrix[0, 2] = np.sin(look_angle)
                base_matrix[2, 0] = -np.sin(look_angle)
                base_matrix[2, 2] = np.cos(look_angle)
            elif cam_idx == 10:  # Arc Right (with rotation)
                angle = -0.005 * frame_idx
                radius = 2.0
                base_matrix[0, 3] = -radius * np.sin(angle)
                base_matrix[2, 3] = -radius * np.cos(angle) + radius
                # Rotate to look at center
                look_angle = angle + np.pi
                base_matrix[0, 0] = np.cos(look_angle)
                base_matrix[0, 2] = np.sin(look_angle)
                base_matrix[2, 0] = -np.sin(look_angle)
                base_matrix[2, 2] = np.cos(look_angle)
            
            # Format the matrix as a string (as expected by the app)
            # Format: [row1] [row2] [row3] [row4] with spaces between values
            matrix_str = ' '.join([f"[{' '.join([str(base_matrix[i, j]) for j in range(4)])}]" for i in range(4)]) + ' '
            
            camera_data[frame_key][f"cam{cam_idx:02d}"] = matrix_str
    
    if progress_callback:
        progress_callback(0.7, desc="Saving camera extrinsics data...")
    
    # Save camera extrinsics to JSON file
    with open(camera_file, 'w') as f:
        json.dump(camera_data, f, indent=2)
    
    logger.info(f"Created sample camera extrinsics at {camera_file}")
    logger.info(f"Created directory for example videos at {videos_dir}")
    
    if progress_callback:
        progress_callback(1.0, desc="Test data structure created successfully!")