#!/usr/bin/env python3
# -*- coding:utf-8 -*-
"""
Description: Based on https://www.lfs.net/programmer/lyt
"""

from pathlib import Path
from typing import List, Tuple, Union, cast, Dict
from struct import Struct
from enum import IntEnum

import numpy as np



class ConeTypes(IntEnum):
    """
    Enum for all possible cone types
    """

    UNKNOWN = 0
    YELLOW = 1
    RIGHT = 1
    BLUE = 2
    LEFT = 2
    ORANGE_SMALL = 3
    START_FINISH_AREA = 3
    ORANGE_BIG = 4
    START_FINISH_LINE = 4


HEADER_STRUCT = Struct("6sBBhBB")
BLOCK_STRUCT = Struct("2h4B")

LytObjectIndexToConeType: Dict[int, ConeTypes] = {
    25: ConeTypes.UNKNOWN,
    29: ConeTypes.YELLOW,
    30: ConeTypes.YELLOW,
    23: ConeTypes.BLUE,
    24: ConeTypes.BLUE,
    27: ConeTypes.ORANGE_BIG,
    20: ConeTypes.ORANGE_SMALL,
}   

def split_header_blocks(data: bytes) -> Tuple[bytes, bytes]:
    """
    Split the content of the lyt file into header and block. This split is easy because
    the header has a fixed size

    Args:
        data (bytes): The content of the lyt file

    Returns:
        Tuple[bytes, bytes]: The header and the block
    """
    return data[: HEADER_STRUCT.size], data[HEADER_STRUCT.size :]


def verify_lyt_header(header_data: bytes) -> None:
    """
    Parse the header and perform some sanity checks suggested by the LFS documentation

    Args:
        header_data (bytes): The header bytes of the `.lyt` file
    """

    header = cast(
        Tuple[bytes, int, int, int, int, int], HEADER_STRUCT.unpack(header_data)
    )

    file_type, version, revision, _, _, _ = header
    assert file_type == b"LFSLYT"
    assert version <= 0, version
    # revision allowed up to 252
    # https://www.lfs.net/forum/thread/96153-LYT-revision-252-in-W-patch
    assert revision <= 252, revision


def extract_cone_lists(blocks_data: bytes) -> List[List[Tuple[float, float]]]:
    """
    Extract the cone object positions from the object blocks bytes of a lyt file

    Args:
        blocks_data (bytes): The data in the lyt file that is not the header

    Returns:
        List[List[Tuple[int, int]]]: The cone positions split by cone type
    """
    decoded_blocks = BLOCK_STRUCT.iter_unpack(blocks_data)
    all_cones_per_type: List[List[Tuple[float, float]]] = [[] for _ in ConeTypes]

    # cone_info:
    for cone_info in decoded_blocks:
        obj_x, obj_y, _, _, lyt_obj_idx, _ = cast(
            Tuple[int, int, int, int, int, int], cone_info
        )

        try:
            cone_type = LytObjectIndexToConeType[lyt_obj_idx]
        except KeyError:
            # not a cone
            continue

        # the stored x,y pos is multiplied by
        # 16 in the file so we need to convert it back
        # (and cast to a float by using real div)
        obj_x_meters = obj_x / 16
        obj_y_meters = obj_y / 16
        all_cones_per_type[cone_type].append((obj_x_meters, obj_y_meters))
    return all_cones_per_type


def load_lyt_file(filename: Union[Path, str]) -> List[np.ndarray]:
    """
    Load a `.lyt` file and return the positions of the cone objects inside it split
    according to `ConeTypes`

    Args:
        filename (Path): The path to the `.lyt` file

    Returns:
        List[np.ndarray]: A list of 2d np.ndarrays representing the cone positions of
        for all cone types
    """
    if isinstance(filename, str):
        filename = Path(filename)
    assert filename.is_file(), filename
    assert filename.suffix == ".lyt", filename
    data = filename.read_bytes()
    header_data, blocks_data = split_header_blocks(data)
    verify_lyt_header(header_data)

    all_cones_per_type = extract_cone_lists(blocks_data)

    all_cones_per_type_arrays = [
        np.array(cone_list) for cone_list in all_cones_per_type
    ]

    return all_cones_per_type_arrays


def main():
    # parse arguments
    import argparse
    import json

    # one argument for the lyt file
    parser = argparse.ArgumentParser(description="Parse a lyt file")
    parser.add_argument("lyt_file", type=str, help="The lyt file to parse")

    args = parser.parse_args()

    # load the lyt file
    lyt_file = Path(args.lyt_file)
    cone_positions = [x.reshape(-1,2) for x in load_lyt_file(lyt_file)]


    center = np.mean(cone_positions[ConeTypes.ORANGE_BIG], axis=0)

    cone_positions_centered = [x - center for x in cone_positions]
    

    cones_dict = {
        "unknown": cone_positions_centered[ConeTypes.UNKNOWN].tolist(),
        "yellow": cone_positions_centered[ConeTypes.YELLOW].tolist(),
        "blue": cone_positions_centered[ConeTypes.BLUE].tolist(),
        "orange_small": cone_positions_centered[ConeTypes.ORANGE_SMALL].tolist(),
        "orange_big": cone_positions_centered[ConeTypes.ORANGE_BIG].tolist(),
    }

    # print the cones as json
    print(json.dumps(cones_dict, indent=1))

if __name__ == "__main__":
    main()

