import torch


def Camera(K=None, hw=None, Twc=None, Tcw=None, D=None, geometry=None):
    """Factory function: auto-detect camera model from K shape and return the right subclass.
    K: intrinsics tensor (shape determines model), or None for default pinhole.
    hw: (H, W) image resolution. Twc/Tcw: [B,4,4] extrinsics (provide one, not both).
    geometry: override auto-detection ('pinhole', 'ftheta', 'fisheye', 'plumb_bob', 'omnimmt')."""

    if hw is None:
        hw = (50, 200)

    if isinstance(hw, torch.Tensor):
        hw = hw.shape[2:] # If it's tensor, take last two dimensions

    if Twc is None and Tcw is None: # Create dummy extrinsics
        Tcw = torch.eye(4).unsqueeze(0)

    if K is None:
        if Twc is not None: dtype, device = Twc.dtype, Twc.device
        if Tcw is not None: dtype, device = Tcw.dtype, Tcw.device
        if Twc is None and Tcw is None: dtype = device = None
        K = torch.tensor([
            [hw[1]/2, 0, hw[1]/2],
            [0, hw[0]/2, hw[0]/2],
            [0,       0,   1.0  ],
        ], dtype=dtype, device=device).unsqueeze(0)
        geometry = 'pinhole'

    # Auto-detect geometry from K shape (flat vector dim or 3x3 matrix).
    if K.dim() == 3 and K.shape[-2:] == (3, 3):
        geometry = 'pinhole'
    elif K.shape[1] == 4:
        geometry = 'pinhole'
    elif K.shape[1] == 25:
        geometry = 'omnimmt'
    elif K.shape[1] == 14:
        geometry = 'ftheta'
    elif K.shape[1] == 8:
        geometry = 'fisheye'
    elif K.shape[1] in [9, 16]:
        geometry = 'plumb_bob'

    if geometry in ['pinhole']:
        from anydata.geometry.cameras.pinhole import CameraPinhole
        return CameraPinhole(K, hw=hw, Twc=Twc, Tcw=Tcw)
    elif geometry in ['ftheta']:
        from anydata.geometry.cameras.ftheta import CameraFTheta
        return CameraFTheta(K, hw=hw, Twc=Twc, Tcw=Tcw)
    elif geometry in ['fisheye','plumb_bob','distorted']:
        from anydata.geometry.cameras.distorted import CameraDistorted
        return CameraDistorted(K, hw=hw, Twc=Twc, Tcw=Tcw)
    elif geometry in ['omnimmt', 'omni']:
        from anydata.geometry.cameras.omnimmt import CameraOmniMMT
        return CameraOmniMMT(K, hw=hw, Twc=Twc, Tcw=Tcw)
    else:
        raise ValueError(f'Unknown camera geometry {geometry!r} (K.shape={K.shape if K is not None else None})')


def Camera_from_list(cams):
    # TODO: this loses subclass-specific state (D for PlumbBob, Komni params, etc.)
    #  and may pick the wrong subclass once K is normalized to [B,3,3].
    #  Should delegate to type(cams[0]).from_list(cams) instead.
    """Batch cameras from a list by concatenating along batch dim."""
    if not cams:
        raise ValueError('Empty camera list')
    K = torch.cat([c.K for c in cams], 0)
    Tcw = torch.cat([c.Tcw for c in cams], 0)
    return Camera(K=K, hw=cams[0].hw, Tcw=Tcw)