import os
import json
import shutil
from pathlib import Path
import numpy as np
from PIL import Image
from tqdm import tqdm

class HexPlaneDatasetConverter:
    def __init__(self, source_dir, cams_dir, output_dir, target_size=(800, 800), 
                 flip_coordinates=True):
        """
        Args:
            source_dir: Directory containing cam_XX folders with images
            cams_dir: Directory containing cam_XX.txt files with parameters
            output_dir: Output directory for the converted dataset
            target_size: Target image size for saved files (width, height)
            flip_coordinates: Apply coordinate system conversion for NeRF
        """
        self.source_dir = Path(source_dir)
        self.cams_dir = Path(cams_dir)
        self.output_dir = Path(output_dir)
        self.target_size = target_size
        self.flip_coordinates = flip_coordinates
        
        # Create output directories
        self.train_dir = self.output_dir / "train"
        self.test_dir = self.output_dir / "test"
        self.train_dir.mkdir(parents=True, exist_ok=True)
        self.test_dir.mkdir(parents=True, exist_ok=True)
        
    def parse_camera_file(self, cam_file):
        """Parse camera parameter text file"""
        with open(cam_file, 'r') as f:
            lines = f.readlines()
        
        # Parse extrinsic matrix (4x4)
        extrinsic = []
        line_idx = 0
        for line in lines:
            if 'extrinsic' in line.lower():
                line_idx = lines.index(line) + 1
                break
        
        for i in range(4):
            row = [float(x) for x in lines[line_idx + i].strip().split()]
            extrinsic.append(row)
        extrinsic = np.array(extrinsic)
        
        # Parse intrinsic matrix (3x3)
        intrinsic = []
        for line in lines:
            if 'intrinsic' in line.lower():
                line_idx = lines.index(line) + 1
                break
        
        for i in range(3):
            row = [float(x) for x in lines[line_idx + i].strip().split()]
            intrinsic.append(row)
        intrinsic = np.array(intrinsic)
        
        return intrinsic, extrinsic
    
    def convert_to_nerf_convention(self, extrinsic):
        """
        Convert camera extrinsic to NeRF/OpenGL convention
        NeRF: +X right, +Y up, -Z forward
        """
        # Convert from world-to-camera to camera-to-world
        c2w = np.linalg.inv(extrinsic)
        
        if self.flip_coordinates:
            # Apply coordinate system transformation
            # This flips from OpenCV to OpenGL/NeRF convention
            flip_mat = np.array([
                [1, 0, 0, 0],
                [0, -1, 0, 0],
                [0, 0, -1, 0],
                [0, 0, 0, 1]
            ])
            c2w = c2w @ flip_mat
        
        return c2w
    
    def get_camera_params(self, intrinsic, original_size):
        """Calculate camera parameters using ORIGINAL dimensions"""
        fx = intrinsic[0, 0]
        fy = intrinsic[1, 1]
        cx = intrinsic[0, 2]
        cy = intrinsic[1, 2]
        
        w, h = original_size
        
        # Calculate camera angles
        camera_angle_x = 2 * np.arctan(w / (2 * fx))
        camera_angle_y = 2 * np.arctan(h / (2 * fy))
        
        return {
            'camera_angle_x': float(camera_angle_x),
            'camera_angle_y': float(camera_angle_y),
            'fl_x': float(fx),
            'fl_y': float(fy),
            'cx': float(cx),
            'cy': float(cy),
            'w': int(w),
            'h': int(h)
        }
    
    def calculate_scene_bounds(self, all_extrinsics):
        """Calculate scene bounding box from camera positions"""
        camera_positions = []
        for extrinsic in all_extrinsics:
            c2w = self.convert_to_nerf_convention(extrinsic)
            cam_pos = c2w[:3, 3]
            camera_positions.append(cam_pos)
        
        camera_positions = np.array(camera_positions)
        
        print(f"\nCamera positions:")
        for i, pos in enumerate(camera_positions):
            print(f"  Camera {i}: {pos}")
        
        # Calculate bounds with margin
        min_bounds = camera_positions.min(axis=0)
        max_bounds = camera_positions.max(axis=0)
        center = (min_bounds + max_bounds) / 2
        
        # Calculate range for each axis
        ranges = max_bounds - min_bounds
        max_range = ranges.max()
        
        # Make it a cube with 50% margin
        extent = max_range * 1.5
        
        scene_min = center - extent / 2
        scene_max = center + extent / 2
        
        return scene_min.tolist(), scene_max.tolist()
    
    def process_images(self):
        """Process all images and generate transforms files"""
        train_frames = []
        test_frames = []
        all_extrinsics = []
        
        # Get all camera folders
        cam_folders = sorted([d for d in self.source_dir.iterdir() 
                             if d.is_dir() and d.name.startswith('cam_')])
        
        print(f"Found {len(cam_folders)} camera folders")
        
        # Get first image to determine original size and mode
        first_cam = cam_folders[0]
        first_image_path = sorted(list(first_cam.glob('*.png')))[0]
        first_image = Image.open(first_image_path)
        original_size = first_image.size
        original_mode = first_image.mode
        
        print(f"Original image size: {original_size}")
        print(f"Original image mode: {original_mode}")
        print(f"Saving images as: {self.target_size}")
        print(f"Coordinate flip: {self.flip_coordinates}")
        
        # Process each camera
        for cam_folder in cam_folders:
            cam_name = cam_folder.name
            cam_id = cam_name.split('_')[1]
            
            # Load camera parameters
            cam_file = self.cams_dir / f"{cam_name}.txt"
            if not cam_file.exists():
                print(f"Warning: Camera file {cam_file} not found, skipping {cam_name}")
                continue
            
            intrinsic, extrinsic = self.parse_camera_file(cam_file)
            all_extrinsics.append(extrinsic)
            
            # Convert to NeRF convention
            transform_matrix = self.convert_to_nerf_convention(extrinsic)
            
            # Get all images for this camera
            image_files = sorted(list(cam_folder.glob('*.png')))
            print(f"\nProcessing {cam_name}: {len(image_files)} images")
            
            # Process each image
            for img_idx, img_path in enumerate(tqdm(image_files, desc=f"{cam_name}")):
                # Every 8th image goes to test
                is_test = (img_idx % 8 == 0)
                
                # Load image
                img = Image.open(img_path)
                
                # Handle different image modes
                if img.mode == 'L':  # Grayscale
                    # Convert grayscale to RGB
                    img = img.convert('RGB')
                elif img.mode == 'RGBA':
                    # Composite RGBA onto white background
                    background = Image.new('RGB', img.size, (255, 255, 255))
                    background.paste(img, mask=img.split()[3])
                    img = background
                elif img.mode != 'RGB':
                    img = img.convert('RGB')
                
                # Resize image
                img_resized = img.resize(self.target_size, Image.LANCZOS)
                
                # Convert to numpy array and add alpha channel
                img_array = np.array(img_resized, dtype=np.uint8)
                h, w = img_array.shape[:2]
                
                # Add full opacity alpha channel
                alpha_channel = np.full((h, w, 1), 255, dtype=np.uint8)
                img_rgba = np.concatenate([img_array, alpha_channel], axis=2)
                
                # Convert back to PIL
                img_rgba_pil = Image.fromarray(img_rgba, 'RGBA')
                
                # Create output filename
                frame_number = img_path.stem
                output_filename = f"cam_{cam_id}_frame_{frame_number}.png"
                
                # Determine save path
                if is_test:
                    output_path = self.test_dir / output_filename
                    file_path = f"test/{output_filename[:-4]}"
                else:
                    output_path = self.train_dir / output_filename
                    file_path = f"train/{output_filename[:-4]}"
                
                # Save as RGBA PNG
                img_rgba_pil.save(output_path, 'PNG')
                
                # Calculate time
                time = img_idx / max(len(image_files) - 1, 1)
                
                # Create frame entry
                frame_entry = {
                    "file_path": file_path,
                    "transform_matrix": transform_matrix.tolist(),
                    "time": float(time)
                }
                
                if is_test:
                    test_frames.append(frame_entry)
                else:
                    train_frames.append(frame_entry)
        
        # Calculate scene bounds
        scene_min, scene_max = self.calculate_scene_bounds(all_extrinsics)
        
        # Get camera parameters
        first_cam_file = self.cams_dir / f"{cam_folders[0].name}.txt"
        intrinsic, _ = self.parse_camera_file(first_cam_file)
        cam_params = self.get_camera_params(intrinsic, original_size)
        
        # Create transforms files
        transforms_train = {
            **cam_params,
            "frames": train_frames
        }
        
        transforms_test = {
            **cam_params,
            "frames": test_frames
        }
        
        # Save JSON files
        with open(self.output_dir / "transforms_train.json", 'w') as f:
            json.dump(transforms_train, f, indent=4)
        
        with open(self.output_dir / "transforms_test.json", 'w') as f:
            json.dump(transforms_test, f, indent=4)
        
        # Verify saved image
        test_img_path = list(self.train_dir.glob('*.png'))[0]
        test_img = Image.open(test_img_path)
        test_array = np.array(test_img)
        
        print(f"\n{'='*60}")
        print(f"✓ Dataset conversion complete!")
        print(f"{'='*60}")
        print(f"  Train images: {len(train_frames)}")
        print(f"  Test images: {len(test_frames)}")
        print(f"  Saved image mode: {test_img.mode}")
        print(f"  Saved image shape: {test_array.shape}")
        print(f"  Alpha channel range: {test_array[:,:,3].min()}-{test_array[:,:,3].max()}")
        print(f"\nScene bounds:")
        print(f"  Min: {scene_min}")
        print(f"  Max: {scene_max}")
        print(f"\n{'='*60}")
        print(f"⚠️  Update your config:")
        print(f"{'='*60}")
        print(f"data:")
        print(f"  datadir: \"./{self.output_dir.name}\"")
        print(f"  scene_bbox_min: {scene_min}")
        print(f"  scene_bbox_max: {scene_max}")
        print(f"  downsample: {original_size[0]/self.target_size[0]:.1f}")
        print(f"\nmodel:")
        print(f"  N_voxel_init: 262144")
        print(f"{'='*60}\n")
        
        # Save debug info
        debug_info = {
            "original_size": original_size,
            "original_mode": original_mode,
            "target_size": self.target_size,
            "coordinate_flip": self.flip_coordinates,
            "scene_bounds": {"min": scene_min, "max": scene_max},
            "num_train": len(train_frames),
            "num_test": len(test_frames)
        }
        with open(self.output_dir / "debug_info.json", 'w') as f:
            json.dump(debug_info, f, indent=4)

# ========== USAGE ==========
if __name__ == "__main__":
    SOURCE_DIR = "project1\myimages"
    CAMS_DIR = "project1\cams"
    OUTPUT_DIR = "project1\ds6"
    
    # Try with coordinate flip first
    converter = HexPlaneDatasetConverter(
        source_dir=SOURCE_DIR,
        cams_dir=CAMS_DIR,
        output_dir=OUTPUT_DIR,
        target_size=(800, 800),
        flip_coordinates=True  # Set to False if still white
    )
    
    converter.process_images()
