import pandas as pd
import os
from tqdm import tqdm
import pyarrow.parquet as pq
import numpy as np
import json

# データセットのディレクトリ
data_dir = "D:\\Waymo"

# 各フォルダのパス
folders = ["testing"]

# データの読み込み関数
def load_parquet_files(folder_path, chunk_size=10000):
    lidar_files = [f for f in os.listdir(folder_path) if f.endswith('.parquet')]
    data_frames = []
    for file in tqdm(lidar_files, desc=f"Loading files from {folder_path}"):
        file_path = os.path.join(folder_path, file)
        try:
            # pyarrowを使用してチャンクで読み込む
            table = pq.ParquetFile(file_path)
            # チャンクごとにデータを読み込む
            for batch in table.iter_batches(batch_size=chunk_size):
                df = batch.to_pandas()
                data_frames.append(df)
        except Exception as e:
            print(f"Error reading {file_path}: {e}")
    return data_frames

# ポイントクラウドデータの前処理関数
def preprocess_point_cloud(data_frames):
    point_clouds = []
    for df in data_frames:
        # カラム名を表示して確認
        print(df.columns)
        # 必要なカラムを抽出
        range_image_str = df['[LiDARComponent].range_image_return1.values'].values[0]
        # range_imageが文字列として保存されている場合の処理
        if isinstance(range_image_str, str):
            range_image = np.array(json.loads(range_image_str))
        else:
            range_image = np.array(range_image_str)
        # range_imageの内容を確認
        print(f"range_image content: {range_image[:10]}")  # 最初の10要素を表示
        # range_imageの形状を確認
        print(f"range_image shape: {range_image.shape}")
        # 1次元配列の場合、2次元配列に変換
        if range_image.ndim == 1:
            range_image = range_image.reshape(-1, 4)
        # x, y, z, intensityを抽出
        if range_image.ndim == 2 and range_image.shape[1] >= 4:
            x = range_image[:, 0]
            y = range_image[:, 1]
            z = range_image[:, 2]
            intensity = range_image[:, 3]
            points = np.vstack((x, y, z, intensity)).T
            point_clouds.append(points)
        else:
            print("Invalid range_image shape")
    return point_clouds

# 各フォルダのデータを読み込む
for folder in folders:
    folder_path = os.path.join(data_dir, folder, "lidar")
    if os.path.exists(folder_path):
        # フォルダ内のすべてのファイルを読み込む
        data_frames = load_parquet_files(folder_path)
        print(f"Loaded {len(data_frames)} files from {folder_path}")

        # ポイントクラウドデータの前処理
        point_clouds = preprocess_point_cloud(data_frames)
        print(f"Processed {len(point_clouds)} point clouds")

# 例として、最初のポイントクラウドデータを表示
if point_clouds:
    print(point_clouds[0][:5])  # 最初の5ポイントを表示

import pandas as pd
import os
import pyarrow.parquet as pq
import numpy as np
import json

# データセットのディレクトリ
data_dir = "D:\\e2258\\Waymo"

# 特定のファイル名
specific_file = "10504764403039842352_460_000_480_000.parquet"

# データの読み込み関数
def load_parquet_file(file_path):
    try:
        # pyarrowを使用してデータを読み込む
        table = pq.read_table(file_path)
        df = table.to_pandas()
        return df
    except Exception as e:
        print(f"Error reading {file_path}: {e}")
        return None

# ポイントクラウドデータの前処理関数
def preprocess_point_cloud(df):
    point_clouds = []
    # カラム名を表示して確認
    print(df.columns)
    # 必要なカラムを抽出
    range_image_str = df['[LiDARComponent].range_image_return1.values'].values[0]
    # range_imageが文字列として保存されている場合の処理
    if isinstance(range_image_str, str):
        range_image = np.array(json.loads(range_image_str))
    else:
        range_image = np.array(range_image_str)
    # range_imageの内容を確認
    print(f"range_image content: {range_image[:10]}")  # 最初の10要素を表示
    # range_imageの形状を確認
    print(f"range_image shape: {range_image.shape}")
    # 1次元配列の場合、2次元配列に変換
    if range_image.ndim == 1:
        range_image = range_image.reshape(-1, 4)
    # 無効なデータをフィルタリング
    valid_points = range_image[range_image[:, 0] != -1]
    # x, y, z, intensityを抽出
    if valid_points.ndim == 2 and valid_points.shape[1] >= 4:
        x = valid_points[:, 0]
        y = valid_points[:, 1]
        z = valid_points[:, 2]
        intensity = valid_points[:, 3]
        # 有効なポイントのみを抽出
        valid_points = np.vstack((x, y, z, intensity)).T
        valid_points = valid_points[valid_points[:, 3] != -1]
        point_clouds.append(valid_points)
    else:
        print("Invalid range_image shape")
    return point_clouds

# 特定のファイルのパス
file_path = os.path.join(data_dir, "testing", "lidar", specific_file)

# ファイルを読み込む
df = load_parquet_file(file_path)
if df is not None:
    print(f"Loaded file: {file_path}")

    # ポイントクラウドデータの前処理
    point_clouds = preprocess_point_cloud(df)
    print(f"Processed {len(point_clouds)} point clouds")

    # 例として、最初のポイントクラウドデータを表示
    if point_clouds:
        print(point_clouds[0][:5])  # 最初の5ポイントを表示
else:
    print(f"Failed to load file: {file_path}")