# Copyright 2026 Toyota Research Institute.  All rights reserved.

import logging
import os
import subprocess
from typing import TYPE_CHECKING, Dict, List

import numpy as np
import pandas as pd
import torch
from scipy.spatial.transform import Rotation

from anydata.converters.utils import fill_metadata, parse_dst_seq, run
from anydata.dataloaders.Unified import UnifiedDataset
from anydata.geometry.camera import Camera
from anydata.utils.write import write_json, write_seq, rgb_ext, data_ext

if TYPE_CHECKING:
    from physical_ai_av.calibration import CameraIntrinsics as CI
    from physical_ai_av.calibration import SensorExtrinsics as SE
    from physical_ai_av.dataset import PhysicalAIAVDatasetInterface
    from physical_ai_av.egomotion import EgomotionState as ES
    from physical_ai_av.utils.camera_models import FThetaCameraModel as FTM
    from physical_ai_av.utils.interpolation import Interpolator
    from physical_ai_av.video import SeekVideoReader as SVR

# Suppress HTTP request logging from urllib3 and huggingface_hub
logging.getLogger("urllib3").setLevel(logging.WARNING)
logging.getLogger("huggingface_hub").setLevel(logging.WARNING)
logging.getLogger("physical_ai_av.dataset").setLevel(logging.ERROR)
logging.getLogger("requests").setLevel(logging.WARNING)
logging.getLogger("httpx").setLevel(logging.WARNING)

DATASET = None

#######################################################

def pose_to_matrix(pos: np.ndarray, rotation: np.ndarray) -> torch.Tensor:
    """Convert 7D pose [x, y, z, qx, qy, qz, qw] to 4x4 transformation matrix"""
    matrix = torch.zeros((*pos.shape[:-1], 4, 4), dtype=torch.float32)
    matrix[..., :3, :3] = torch.from_numpy(Rotation.from_quat(rotation).as_matrix()).float()
    matrix[..., :3, 3] = torch.from_numpy(pos).float()
    matrix[..., 3, 3] = 1.0
    return matrix


#######################################################


def get_physical_ai_dataset():
    """Initialize PhysicalAI dataset on first use"""
    global DATASET
    if DATASET is None:
        try:
            from physical_ai_av.dataset import PhysicalAIAVDatasetInterface
            revision = "26.03"
            DATASET = PhysicalAIAVDatasetInterface(revision, confirm_download_threshold_gb=float("inf"))
        except ImportError as e:
            message = "physical_ai_av is required. Install with: pip install -e physical_ai_av"
            raise ImportError(message) from e
    return DATASET


def select_random_device(index: int = None):
    """
    Selects a random device (CPU or a random available GPU) for PyTorch.
    """
    nvcc_installed = subprocess.run(['which', 'nvcc'], capture_output=True, text=True, check=True)
    if torch.cuda.is_available() and nvcc_installed.stdout.strip():
        idx = index if index is not None else np.random.randint(0, torch.cuda.device_count())
        return torch.device(f"cuda:{idx % torch.cuda.device_count()}")
    return torch.device("cpu")


def get_sequences(args):
    api = get_physical_ai_dataset()
    cams = list(getattr(getattr(api.features, "CAMERA"), "ALL"))
    egomotion = api.feature_presence["egomotion"].values[:,None]
    lidar = api.feature_presence['lidar_top_360fov'].values[:,None]
    camera = api.feature_presence[cams].values
    camera_intrinsics = api.feature_presence['camera_intrinsics'].values[:,None]
    sensor_extrinsics = api.feature_presence['sensor_extrinsics'].values[:,None]
    required = np.concatenate([egomotion, lidar, camera, camera_intrinsics, sensor_extrinsics], -1)
    valid = np.all(required, axis=-1)
    valid_indices = np.where(valid)[0]
    return valid_indices

#######################################################

def float64_to_float32(data_dict: dict | list | np.ndarray) -> dict | np.ndarray:
    """Convert all values in the dictionary to numpy arrays."""
    if isinstance(data_dict, dict): 
        return {k: float64_to_float32(v) for k, v in data_dict.items()}
    elif isinstance(data_dict, list): 
        return [float64_to_float32(v) for v in data_dict]
    elif isinstance(data_dict, np.ndarray) and data_dict.dtype == np.float64: 
        return data_dict.astype(np.float32)
    return data_dict


def download_clip_data(api: "PhysicalAIAVDatasetInterface", clip_id: int, save_path: str = None):
    try:
        import DracoPy
    except ImportError as e:
        raise ImportError("DracoPy is required. Install with: pip install -e DracoPy") from e
    
    save_path = f'{save_path}.npz' if save_path is not None else None
    if save_path is not None and os.path.exists(save_path):
        return np.load(save_path, allow_pickle=True)["data"][()]

    sensor_extrinsics: "SE" = api.get_clip_feature(clip_id, "sensor_extrinsics", maybe_stream=True)
    camera_intrinsics: "CI" = api.get_clip_feature(clip_id, "camera_intrinsics", maybe_stream=True)
    egomotion_data: "Interpolator" = api.get_clip_feature(clip_id, "egomotion", maybe_stream=True)
    used_sensors = list(camera_intrinsics.camera_models) + ["lidar_top_360fov"]

    lidar_data = api.get_clip_feature(clip_id, "lidar_top_360fov", maybe_stream=True)
    pcdf: pd.DataFrame = lidar_data["pointclouds"]
    point_clouds = [DracoPy.decode(depc).points for depc in pcdf["draco_encoded_pointcloud"]]
    lidar_timestamps = pcdf["spin_end_timestamp"].values
    num_steps = len(lidar_timestamps)

    ego_states = {"pos": np.zeros((num_steps, 3)), "rotation": np.zeros((num_steps, 4))}
    for i, t in enumerate(lidar_timestamps):
        ego_state: "ES" = egomotion_data(t)
        ego_states["pos"][i] = ego_state.pose.translation
        ego_states["rotation"][i] = ego_state.pose.rotation.as_quat()

    extrinsics = {}
    for sensor in used_sensors:
        transform = sensor_extrinsics.sensor_poses[sensor]
        extrinsics[sensor] = dict(pos=transform.translation, rotation=transform.rotation.as_quat())

    intrinsics, camera_images = {}, {}
    for cam in camera_intrinsics.camera_models:
        ftm: "FTM" = camera_intrinsics.camera_models[cam]
        intrinsics[cam] = dict(cxy=ftm.principal_point, fw_poly=ftm.th2r.coef, bw_poly=ftm.r2th.coef)
        camera_data: "SVR" = api.get_clip_feature(clip_id, cam, maybe_stream=True)
        timestep_diffs = np.abs(lidar_timestamps[:, None] - camera_data.timestamps[None, :])
        closest_camera_indices = np.argmin(timestep_diffs, axis=1)
        camera_images[cam] = camera_data.decode_images_from_frame_indices(closest_camera_indices)

    scenario = float64_to_float32({
        "sensor_extrinsics": extrinsics,
        "camera_intrinsics": intrinsics,
        "time_series": {
            "timestamps": lidar_timestamps,
            "ego_states": ego_states,
            "point_clouds": point_clouds,
            "camera_images": camera_images,
        }
    })

    if save_path is not None:
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        np.savez(save_path, data=scenario)
    return scenario


def parse_sequence(seq, args):
    return parse_dst_seq(seq, args)


def process_sequence(proc_index: int, seq: str, dst: str, args):
    ### Parse destination path
    api = get_physical_ai_dataset()
    clip_id = api.clip_index.index[seq]
    src = os.path.join(args.src, f"{seq:06d}_{clip_id}")

    device = select_random_device(proc_index)
    data = download_clip_data(api, clip_id, save_path=None if args.official else src)
    intrinsics_all: Dict[str, Dict[str, np.ndarray]] = data['camera_intrinsics']
    extrinsics_all: Dict[str, Dict[str, np.ndarray]] = data['sensor_extrinsics']
    camera_images_all: Dict[str, np.ndarray] = data['time_series']['camera_images']

    pointclouds: List[torch.Tensor] = []
    lidar_transform = pose_to_matrix(**extrinsics_all["lidar_top_360fov"]).to(device)
    for pc_np in data['time_series']['point_clouds']:
        pc = torch.from_numpy(pc_np).unsqueeze(0).permute(0, 2, 1).to(device)
        pc4d = torch.cat([pc, torch.ones_like(pc[:, :1])], dim=1)
        pc_world = lidar_transform.unsqueeze(0).bmm(pc4d)[:, :3]
        pointclouds.append(pc_world.cpu())

    framerate = 10
    resolution, num_frames = {}, {}
    rgbs_ext, depth_ext, lowdim_ext = rgb_ext(args.storage), data_ext(args.storage), "npz"
    for cam, ftm in intrinsics_all.items():
        scaling_factor_xy = np.ones([2], dtype=np.float32)
        intrinsics_np = np.hstack([scaling_factor_xy, ftm["cxy"], ftm["fw_poly"], ftm["bw_poly"]])
        intrinsics = torch.from_numpy(intrinsics_np.astype(np.float32)).unsqueeze(0).float()
        extrinsics_local = pose_to_matrix(**extrinsics_all[cam]).unsqueeze(0).to(device)

        rgbs = torch.from_numpy(camera_images_all[cam])
        camera = Camera(K=intrinsics, Tcw=extrinsics_local.cpu(), hw=rgbs.shape[1:3], geometry='ftheta')
        depths = torch.cat([camera.project_pointcloud(pc) for pc in pointclouds]).squeeze(1)

        resolution[cam] = {"rgb": rgbs.shape[1:3], "depth": depths[0].shape[-2:]}
        num_frames[cam] = {"rgb": rgbs.shape[0], "depth": depths.shape[0]}
        timesteps = np.arange(rgbs.shape[0])

        rgb_array = rgbs.cpu()
        depth_array = depths.cpu().numpy()

        ego_states = data['time_series']['ego_states']
        ego_transforms = pose_to_matrix(ego_states["pos"], ego_states["rotation"]).to(device)
        all_extrinsics_np = (ego_transforms @ extrinsics_local).cpu().numpy()
        all_intrinsics_np = intrinsics.repeat(ego_transforms.shape[0], 1).cpu().numpy()
        lowdim = dict(timestep=timesteps, extrinsics=all_extrinsics_np, intrinsics=all_intrinsics_np)

        write_seq(f'{dst}/{UnifiedDataset.RGB_FOLDER}/{cam}', rgb_array, rgbs_ext, args.storage, fps=framerate)
        write_seq(f'{dst}/{UnifiedDataset.DEPTH_FOLDER}/{cam}', dict(depth=depth_array), depth_ext, args.storage, bits=18)
        write_seq(f'{dst}/{UnifiedDataset.LOWDIM_FOLDER}/{cam}', lowdim, lowdim_ext, args.storage)

    ############ METADATA 
    filename = f'{dst}/metadata.json'
    seq_metadata = fill_metadata(
        args=args,
        info=dict(name='PhysicalAI', tags=['real', 'dynamic', 'driving'], raw_id=clip_id),
        labels=["rgb", "intrinsics", "extrinsics", "depth"],
        cameras=list(camera_images_all.keys()),
        resolution=resolution,
        num_frames=num_frames,
        framerate=framerate,
        rgb=dict(extension=rgbs_ext),
        lowdim=dict(extension=lowdim_ext),
        intrinsics=dict(model=camera.geometry),
        extrinsics=dict(transform='cam2world',metric=True),
        depth=dict(extension=depth_ext,metric=True,sparse=True),
        semantic=None,
        action=None,
        language=None,
        specific=None,
    )
    write_json(filename, seq_metadata)

    return dst

#######################################################

if __name__ == '__main__':
    converter = os.path.basename(__file__)
    run(converter, get_sequences, parse_sequence, process_sequence)

#######################################################