from pathlib import Path
from av2.map.map_api import ArgoverseStaticMap

from typing import Any, Dict, List, Optional, Tuple, Union
from pyarrow import feather
import os
import numpy as np
import numpy.typing as npt
import pandas as pd
import mayavi.mlab as mlab


NDArrayNumber = npt.NDArray["np.number[Any]"]
NDArrayBool = npt.NDArray[np.bool_]
NDArrayFloat = npt.NDArray[np.float64]
NDArrayByte = npt.NDArray[np.uint8]
NDArrayInt = npt.NDArray[np.int64]
NDArrayObject = npt.NDArray[np.object_]

# <Copyright 2022, Argo AI, LLC. Released under the MIT license.>

"""Unit tests for HD maps."""

argoverse_dataset_dir = "/mnt/hgfs/share-folder/data/argoverse/Argoverse_dataset_sample/train_000/sensor/train"
arg_id = "00a6ffc1-6ce9-3bc3-a060-6006e9893a1a"
log_map_dir = Path(argoverse_dataset_dir) / arg_id / "map"
log_lidar_dir = Path(argoverse_dataset_dir) / arg_id / "sensors" / "lidar"
lidar_id = "315967376859506000"
lidar_file_type = ".feather"

def process_all_files_name(dir:str):
    for root, dirs, files in os.walk(dir):
        print('root_dir:', root)  #当前路径
        print('sub_dirs:', dirs)   #子文件夹
        print('files:', files)     #文件名称，返回list类型

def process_all_files_name2(dir:str):
    camera_dir = dir / arg_id / "sensors" / "cameras"
    filelist = os.listdir(camera_dir)
    for cur_camera_dir in filelist:
        cur_camera_files = os.listdir(camera_dir / cur_camera_dir)
        print("=========================================================================")
        print("=========================================================================")
        print("[current camera]:%s" %(cur_camera_dir))
        for cur_file in cur_camera_files:
            cur_file_path = camera_dir / cur_camera_dir / cur_file
            #os.rename(cur_file_path, str(cur_file_path).split(".")[0])
            #print("[old]:%s      ------>    [new]:%s" %(cur_file_path, str(cur_file_path).split(".")[0]))
            os.rename(cur_file_path, str(cur_file_path) + ".jpg")
            print("[old]:%s      ------>    [new]:%s" %(cur_file_path, str(cur_file_path) + ".jpg"))


def read_feather(path: Path, columns: Optional[Tuple[str, ...]] = None) -> pd.DataFrame:
    """Read Apache Feather data from a .feather file.

    AV2 uses .feather to serialize much of its data. This function handles the deserialization
    process and returns a `pandas` DataFrame with rows corresponding to the records and the
    columns corresponding to the record attributes.

    Args:
        path: Source data file (e.g., 'lidar.feather', 'calibration.feather', etc.)
        columns: Tuple of columns to load for the given record. Defaults to None.

    Returns:
        (N,len(columns)) Apache Feather data represented as a `pandas` DataFrame.
    """
    data: pd.DataFrame = feather.read_feather(path, columns=columns)
    return data


def read_lidar_sweep(fpath: Path, attrib_spec: str = "xyz") -> NDArrayFloat:
    """Load a point cloud file from a filepath.

    Args:
        fpath: path to a .feather file
        attrib_spec: string of C characters, each char representing a desired point attribute
            x -> point x-coord
            y -> point y-coord
            z -> point z-coord

        The following attributes are not loaded:
            intensity -> point intensity/reflectance
            laser_number -> laser number of laser from which point was returned
            offset_ns -> nanosecond timestamp offset per point, from sweep timestamp.

    Returns:
        Array of shape (N, C). If attrib_str is invalid, `None` will be returned

    Raises:
        ValueError: If any element of `attrib_spec` is not in (x, y, z, intensity, laser_number, offset_ns).
    """
    possible_attributes = ["x", "y", "z", "intensity", "laser_number", "offset_ns"]
    if not all([a in possible_attributes for a in attrib_spec]):
        raise ValueError("Attributes must be in (x, y, z, intensity, laser_number, offset_ns).")

    sweep_df = read_feather(fpath)

    # return only the requested point attributes
    sweep: NDArrayFloat = sweep_df[list(attrib_spec)].to_numpy().astype(np.float64)
    return sweep



# pts_mode='sphere'
def draw_lidar(
    pc,
    color=None,
    fig=None,
    bgcolor=(0, 0, 0),
    pts_scale=0.3,
    pts_mode="sphere",
    pts_color=None,
    color_by_intensity=False,
    pc_label=False,
):
    """ Draw lidar points
    Args:
        pc: numpy array (n,3) of XYZ
        color: numpy array (n) of intensity or whatever
        fig: mayavi figure handler, if None create new one otherwise will use it
    Returns:
        fig: created or used fig
    """
    # ind = (pc[:,2]< -1.65)
    # pc = pc[ind]
    pts_mode = "point"
    print("====================", pc.shape)
    if fig is None:
        fig = mlab.figure(
            figure=None, bgcolor=bgcolor, fgcolor=None, engine=None, size=(1600, 1000)
        )
    if color is None:
        color = pc[:, 2]
    if pc_label:
        color = pc[:, 4]
    if color_by_intensity:
        color = pc[:, 2]

    mlab.points3d(
        pc[:, 0],
        pc[:, 1],
        pc[:, 2],
        color,
        color=pts_color,
        mode=pts_mode,
        colormap="gnuplot",
        scale_factor=pts_scale,
        figure=fig,
    )

    # draw origin
    mlab.points3d(0, 0, 0, color=(1, 1, 1), mode="sphere", scale_factor=0.2)

    # draw axis
    axes = np.array(
        [[2.0, 0.0, 0.0, 0.0], [0.0, 2.0, 0.0, 0.0], [0.0, 0.0, 2.0, 0.0]],
        dtype=np.float64,
    )
    mlab.plot3d(
        [0, axes[0, 0]],
        [0, axes[0, 1]],
        [0, axes[0, 2]],
        color=(1, 0, 0),
        tube_radius=None,
        figure=fig,
    )
    mlab.plot3d(
        [0, axes[1, 0]],
        [0, axes[1, 1]],
        [0, axes[1, 2]],
        color=(0, 1, 0),
        tube_radius=None,
        figure=fig,
    )
    mlab.plot3d(
        [0, axes[2, 0]],
        [0, axes[2, 1]],
        [0, axes[2, 2]],
        color=(0, 0, 1),
        tube_radius=None,
        figure=fig,
    )

    # draw fov (todo: update to real sensor spec.)
    fov = np.array(
        [[20.0, 20.0, 0.0, 0.0], [20.0, -20.0, 0.0, 0.0]], dtype=np.float64  # 45 degree
    )

    mlab.plot3d(
        [0, fov[0, 0]],
        [0, fov[0, 1]],
        [0, fov[0, 2]],
        color=(1, 1, 1),
        tube_radius=None,
        line_width=1,
        figure=fig,
    )
    mlab.plot3d(
        [0, fov[1, 0]],
        [0, fov[1, 1]],
        [0, fov[1, 2]],
        color=(1, 1, 1),
        tube_radius=None,
        line_width=1,
        figure=fig,
    )

    # draw square region
    TOP_Y_MIN = -20
    TOP_Y_MAX = 20
    TOP_X_MIN = 0
    TOP_X_MAX = 40
    #TOP_Z_MIN = -2.0
    #TOP_Z_MAX = 0.4

    x1 = TOP_X_MIN
    x2 = TOP_X_MAX
    y1 = TOP_Y_MIN
    y2 = TOP_Y_MAX
    mlab.plot3d(
        [x1, x1],
        [y1, y2],
        [0, 0],
        color=(0.5, 0.5, 0.5),
        tube_radius=0.1,
        line_width=1,
        figure=fig,
    )
    mlab.plot3d(
        [x2, x2],
        [y1, y2],
        [0, 0],
        color=(0.5, 0.5, 0.5),
        tube_radius=0.1,
        line_width=1,
        figure=fig,
    )
    mlab.plot3d(
        [x1, x2],
        [y1, y1],
        [0, 0],
        color=(0.5, 0.5, 0.5),
        tube_radius=0.1,
        line_width=1,
        figure=fig,
    )
    mlab.plot3d(
        [x1, x2],
        [y2, y2],
        [0, 0],
        color=(0.5, 0.5, 0.5),
        tube_radius=0.1,
        line_width=1,
        figure=fig,
    )

    # mlab.orientation_axes()
    mlab.view(
        azimuth=180,
        elevation=70,
        focalpoint=[12.0909996, -1.04700089, -2.03249991],
        distance=62.0,
        figure=fig,
    )
    return fig




if __name__ == "__main__":
    
    lidar_fpath = log_lidar_dir / (lidar_id + lidar_file_type)

    lidar_points = read_lidar_sweep(lidar_fpath, attrib_spec=("x", "y", "z", "intensity", "laser_number", "offset_ns"))
    #avm = ArgoverseStaticMap.from_map_dir(log_map_dirpath=log_map_dir, build_raster=False)
    print(("All point num: ", lidar_points.shape[0]))
    print("lidar points", lidar_points.shape)

    fig = mlab.figure(
        figure=None, bgcolor=(0, 0, 0), fgcolor=None, engine=None, size=(1000, 500)
    )
    draw_lidar(lidar_points, fig=fig, color_by_intensity=True)

    mlab.show()
    #process_all_files_name2(Path(argoverse_dataset_dir))
    print("Check OK")