import numpy as np
import torch
from shap_e.models.nn.camera import DifferentiableCameraBatch, DifferentiableProjectiveCamera
from shap_e.util.collections import AttrDict

def create_custom_cameras(size: int, device: torch.device, azimuths=None, elevations=None, fov_degrees=30, distance=3.0) -> DifferentiableCameraBatch:
    """
    Create custom camera angles for rendering.
    
    Args:
        size: The width and height of the rendered image.
        device: The device to put the camera parameters on.
        azimuths: List of azimuth angles in degrees.
        elevations: List of elevation angles in degrees.
        fov_degrees: Field of view in degrees.
        distance: Distance from the origin.
    
    Returns:
        A DifferentiableCameraBatch containing the specified cameras.
    """
    if azimuths is None:
        azimuths = [0]
    if elevations is None:
        elevations = [0]
    
    origins = []
    xs = []
    ys = []
    zs = []
    
    for azimuth, elevation in zip(azimuths, elevations):
        # Convert to radians
        azimuth_rad = np.deg2rad(azimuth)
        elevation_rad = np.deg2rad(elevation)
        
        # Calculate camera position
        x_pos = distance * np.cos(elevation_rad) * np.sin(azimuth_rad)
        y_pos = distance * np.cos(elevation_rad) * np.cos(azimuth_rad)
        z_pos = distance * np.sin(elevation_rad)
        
        # Camera origin (position)
        origin = np.array([x_pos, y_pos, z_pos])
        
        # Camera z-axis (looking at the origin)
        z = -origin / np.linalg.norm(origin)
        
        # Camera x-axis (right)
        x = np.array([np.cos(azimuth_rad + np.pi/2), np.sin(azimuth_rad + np.pi/2), 0.0])
        x = x - np.dot(x, z) * z  # Make orthogonal to z
        x = x / np.linalg.norm(x)  # Normalize
        
        # Camera y-axis (up, computed as z cross x)
        y = np.cross(z, x)
        y = y / np.linalg.norm(y)  # Normalize
        
        origins.append(origin)
        xs.append(x)
        ys.append(y)
        zs.append(z)
    
    # Convert from radians to the appropriate x and y fov
    fov_rad = np.deg2rad(fov_degrees)
    
    return DifferentiableCameraBatch(
        shape=(1, len(origins)),
        flat_camera=DifferentiableProjectiveCamera(
            origin=torch.from_numpy(np.stack(origins, axis=0)).float().to(device),
            x=torch.from_numpy(np.stack(xs, axis=0)).float().to(device),
            y=torch.from_numpy(np.stack(ys, axis=0)).float().to(device),
            z=torch.from_numpy(np.stack(zs, axis=0)).float().to(device),
            width=size,
            height=size,
            x_fov=fov_rad,
            y_fov=fov_rad,
        ),
    )