import os
import numpy as np
import open3d as o3d
import trimesh
import random
from tqdm import tqdm
from pathlib import Path

def augment_obj_file(input_path, output_path, n_augmentations=5):
    """
    Augment an OBJ file with random transformations
    Args:
        input_path: Path to input OBJ file
        output_path: Directory to save augmented files
        n_augmentations: Number of augmented copies to create
    """
    # Create output directory if it doesn't exist
    os.makedirs(output_path, exist_ok=True)
    
    # Load the original mesh
    mesh = trimesh.load(input_path)
    original_name = Path(input_path).stem
    
    for i in range(n_augmentations):
        # Create a copy of the original mesh
        augmented_mesh = mesh.copy()
        
        # Random rotation (0-360 degrees around each axis)
        angle_x = np.random.uniform(0, 2*np.pi)
        angle_y = np.random.uniform(0, 2*np.pi)
        angle_z = np.random.uniform(0, 2*np.pi)
        rotation_matrix = trimesh.transformations.euler_matrix(angle_x, angle_y, angle_z)
        augmented_mesh.apply_transform(rotation_matrix)
        
        # Random scaling (0.8-1.2 range)
        scale_factor = np.random.uniform(0.8, 1.2, size=3)
        scale_matrix = np.eye(4)
        scale_matrix[:3, :3] *= scale_factor
        augmented_mesh.apply_transform(scale_matrix)
        
        # Random translation (-0.1 to 0.1 range in each dimension)
        translation = np.random.uniform(-0.1, 0.1, size=3)
        translation_matrix = np.eye(4)
        translation_matrix[:3, 3] = translation
        augmented_mesh.apply_transform(translation_matrix)
        
        # Save the augmented mesh
        output_file = os.path.join(output_path, f"{original_name}_aug_{i}.obj")
        augmented_mesh.export(output_file)

def augment_all_objs(source_dir, target_dir, n_augmentations=5):
    """
    Augment all OBJ files in a directory
    Args:
        source_dir: Directory containing original OBJ files
        target_dir: Directory to save augmented files
        n_augmentations: Number of augmented copies per file
    """
    # Get all OBJ files in source directory
    obj_files = [f for f in os.listdir(source_dir) if f.endswith('.obj')]
    
    print(f"Found {len(obj_files)} OBJ files to augment")
    print(f"Will create {n_augmentations} augmented versions per file")
    
    # Process each file
    for obj_file in tqdm(obj_files, desc="Augmenting OBJ files"):
        input_path = os.path.join(source_dir, obj_file)
        augment_obj_file(input_path, target_dir, n_augmentations)
    
    print(f"Finished! Augmented files saved to: {target_dir}")

if __name__ == "__main__":
    # Configuration
    N_AUGMENTATIONS = 10  # Number of augmented copies per file
    SOURCE_DIR = "/root/shapenet_data/train_mesh_data_under_25kb_1000/train"  # Directory with original OBJ files
    TARGET_DIR = f"/root/shapenet_data/train_mesh_data_under_25kb_1000/train_{N_AUGMENTATIONS}augment"  # Where to save augmented files
    
    # Run augmentation
    augment_all_objs(SOURCE_DIR, TARGET_DIR, N_AUGMENTATIONS)