import click
from typing import Optional

import numpy as np
import threading

from ouster.sdk import client, open_source
from ouster.sdk.util import resolve_metadata
import ouster.sdk.util.pose_util as pu
from ouster.sdk.viz import Label

from ouster.sdk.viz.accumulators import LidarScanVizAccumulators
from ouster.sdk.viz.accumulators_config import (LidarScanVizAccumulatorsConfig,
                                                MAP_MAX_POINTS_NUM,
                                                MAP_SELECT_RATIO)
from ouster.sdk.viz.model import LidarScanVizModel

# import ouster.sdk.viz.scans_accum as scans_accum_module

from ptudes.utils import (make_point_viz, spin, estimate_apex_dolly,
                          map_points_num, prune_trajectory, read_metadata_json,
                          read_packet_source, pose_scans_from_nc_gt,
                          read_newer_college_gt)

from ptudes.fly import (FlyingState, FState, BuildingState,
                        CameraTransitionState, CoursingState)

@click.command(name="flyby")
@click.argument(
    'file',
    required=True,
    type=click.Path(exists=True))
@click.option(
    '-m',
    '--meta',
    required=False,
    type=click.Path(exists=True, dir_okay=False, readable=True),
    help="Metadata for PCAP, required if automatic metadata resolution fails")
@click.option('--kitti-poses',
              required=False,
              type=click.Path(exists=True, dir_okay=False, readable=True),
              help="Poses file in Kitti format, one pose per scan "
              "(can be generated by kiss-icp)")
@click.option('--nc-gt-poses',
              required=False,
              type=click.Path(exists=True, dir_okay=False, readable=True),
              help="Poses file in Newer College Dataset (ground truth) format")
@click.option('-r',
              '--rate',
              type=float,
              default=1.0,
              help="Playback rate: 0.1, 0.25, 0.5, 1.0, 1.5, 2.0, 3.0")
@click.option(
    "--accum-map-ratio",
    type=float,
    help="Ratio of random points of every scan to add to an overall map")
@click.option("--start-scan", type=int, default=0, help="Start scan number")
@click.option("--end-scan", type=int, help="End scan number, inclusive")
def ptudes_flyby(file: str, meta: Optional[str], kitti_poses: Optional[str],
                 nc_gt_poses: Optional[str], rate: float,
                 accum_map_ratio: Optional[float], start_scan: int,
                 end_scan: Optional[int]) -> None:
    """Flyby 3d visualizer over the lidar scans with poses.

    Data is provided via FILE in Ouster raw packets formats: PCAP or BAG with
    lidar/imu packets.
    
    NOTE: --start-scan/--end-scan works as expected for poses files that have
    all poses starting from the beginning of the data file, if chunk is started
    and saved not from the beginning, make sure to use --nc-gt-poses option,
    otherwise it has some unexpected quirks (sorry I don't have time to
    fix/restructure it now :()
    """
    meta = resolve_metadata(file, meta)
    if not meta:
        raise click.ClickException(
            "File not found, please specify a metadata file with `-m`")
    print(f"Reading metadata from: {meta}")
    info = read_metadata_json(meta)

    packet_source = read_packet_source(file, meta=info)
    scans_source = client.Scans(packet_source)

    # TODO: remove double read of poses files, jugling around some API
    # limitations
    if kitti_poses:
        poses = pu.load_kitti_poses(kitti_poses)
        scans_num = len(poses)
        start_scan = start_scan if start_scan < scans_num else 0
        end_scan = (end_scan if end_scan is not None and end_scan < scans_num
                    and end_scan >= start_scan else scans_num - 1)
        scans = pu.pose_scans_from_kitti(scans_source, kitti_poses)
    elif nc_gt_poses:
        gts_poses = read_newer_college_gt(nc_gt_poses)
        scans_num = len(gts_poses)
        start_scan = start_scan if start_scan < scans_num else 0
        end_scan = (end_scan if end_scan is not None and end_scan < scans_num
                    and end_scan >= start_scan else scans_num - 1)
        # Move all poses to the origin (i.e. --start-scan wise) (easier to
        # compare various GTs)
        pose0_inv = np.linalg.inv(gts_poses[start_scan][1])
        gts_poses = [(t, pose0_inv @ p) for t, p in gts_poses]
        scans = pose_scans_from_nc_gt(scans_source, nc_gt_poses=gts_poses)
    else:
        raise click.ClickException(
            "Required one of --kitti-poses or --nc-gt-poses, but none was set."
        )

    MAP_MAX_POINTS_NUM = 1500000
    # estimate accum map ratio for the densest map
    estimated_map_ratio = False
    if accum_map_ratio is None:
        # TODO: account for the column window when it's not 360
        pts_per_scan = (info.format.pixels_per_column *
                        info.format.columns_per_frame)
        pts_total = (end_scan - start_scan + 1) * pts_per_scan
        accum_map_ratio = min(1.0, MAP_MAX_POINTS_NUM / pts_total)
        estimated_map_ratio = True
        click.echo(f"Estimated accum map ratio: {accum_map_ratio}")

    rates = [0.1, 0.25, 0.5, 1.0, 1.5, 2.0, 3.0]
    try:
        rate_ind = rates.index(rate)
    except ValueError:
        click.echo(
            f"WARNING: {rate = } is not found in {rates}, using rate = 1.0")
        rate_ind = rates.index(1.0)

    point_viz = make_point_viz(title="Flyby")

    accum_config = LidarScanVizAccumulatorsConfig(
        accum_max_num=0,
        accum_min_dist_meters=0,
        accum_min_dist_num=1,
        map_enabled=True,
        map_select_ratio=accum_map_ratio,
        map_max_points=MAP_MAX_POINTS_NUM)

    scan_viz_model = LidarScanVizModel([scans_source.metadata],
                                       _img_aspect_ratio=0)

    scans_accum = LidarScanVizAccumulators(scan_viz_model, point_viz,
                                           accum_config, threading.Lock())

    # initialize flyby osd
    flyby_osd = Label("", 1, 1, align_right=True)
    point_viz.add(flyby_osd)

    pause = False
    osd_enabled = True

    def handle_keys(ctx, key, mods) -> bool:
        nonlocal pause, rate_ind, osd_enabled
        if key == 32:
            pause = not pause
        elif key == ord('.') and mods == 1:
            rate_ind = (rate_ind + 1) % len(rates)
        elif key == ord(',') and mods == 1:
            rate_ind = (rate_ind + len(rates) - 1) % len(rates)
        elif key == ord('O'):
            osd_enabled = not osd_enabled
            scans_accum.toggle_osd(osd_enabled)
            scans_accum.draw(update=False)
        return True

    point_viz.push_key_handler(handle_keys)

    # TODO: min_max/res_poses logic are sketchy and smells here, needs a
    # refactoring. (some weirdness come from using kitti vs nc gt formats
    # with/wo start/end scan params)
    # min and max point in the cloud to calculate dolly
    min_max = np.zeros((3, 2))
    # used for camera coursing pruned trajectory
    res_poses = []

    def make_osd_str() -> str:
        pause_str = "" if not pause else " (PAUSED)"
        osd_str = f"map of scans: {start_scan} - {end_scan}"
        osd_str += f"\nmap num points: {map_points_num(scans_accum)} " + (
            "(O)" if scans_accum._ma._map_overflow else "")
        osd_str += f"\nscans downsample ratio: {accum_map_ratio:.03f}"
        osd_str += f"\nstate: {fstate.name.replace('_', '..')}"
        osd_str += f"\nplayback: {rates[rate_ind]} {pause_str}"
        return osd_str

    def create_state(flying_state: FlyingState) -> FState:
        nonlocal min_max, res_poses
        if flying_state == FlyingState.BUILDING:
            return BuildingState(scans,
                                 scans_accum,
                                 start_scan=start_scan,
                                 end_scan=end_scan,
                                 min_max=min_max,
                                 poses=res_poses,
                                 next_state=FlyingState.TO_THE_BEGINNING)
        elif flying_state == FlyingState.TO_THE_BEGINNING:
            return CameraTransitionState(name="TO_THE_BEGINNING",
                                         target=res_poses[0][1],
                                         velocity=0.15,
                                         next_state=FlyingState.COURSING)
        elif flying_state == FlyingState.COURSING:
            pruned_traj_poses = prune_trajectory(res_poses)
            coursing_traj_eval = pu.TrajectoryEvaluator(pruned_traj_poses,
                                                        time_bounds=1.0)
            return CoursingState(coursing_traj_eval,
                                 velocity=5.0,
                                 next_state=FlyingState.TO_THE_APEX)
        elif flying_state == FlyingState.TO_THE_APEX:
            dolly = estimate_apex_dolly(min_max, point_viz.camera.get_fov())
            return CameraTransitionState(
                name="TO_THE_APEX",
                pitch=0,
                yaw=140,
                dolly=dolly,
                velocity=0.15,
                next_state=FlyingState.TO_THE_BEGINNING)

    fstate: Optional[FState] = create_state(FlyingState.BUILDING)

    last_ts = -1
    def on_update(pviz, tick_ts) -> None:
        nonlocal last_ts, fstate, flyby_osd
        t = tick_ts

        # initialization step
        if last_ts < 0:
            last_ts = t
            return

        dt = tick_ts - last_ts
        dt = rates[rate_ind] * dt
        if pause:
            dt = 0
        last_ts = t

        new_state = fstate.update(dt, pviz)
        if new_state is not None:
            fstate = create_state(new_state)

        if osd_enabled:
            flyby_osd.set_text(make_osd_str())
        else:
            flyby_osd.set_text("")


    spin(point_viz,
         on_update,
         period=0.0333)
