"""
action_tokenizer.py

Extension class; wraps base LLM/VLM tokenizer with logic to discretize and tokenize continuous robot actions.
"""

from typing import List, Union

import numpy as np
from transformers import PreTrainedTokenizerBase
from pathlib import Path
import yaml


class ActionTokenizer:
    def __init__(
        self, tokenizer: PreTrainedTokenizerBase, bins: int = 256, min_action: int = -1, max_action: int = 1
    ) -> None:
        """
        Discretizes continuous robot actions into N bins per dimension and maps to the least used tokens.

        NOTE =>> by default, assumes a BPE-style tokenizer akin to the LlamaTokenizer, where *the least used tokens*
                 appear at the end of the vocabulary!

        :param tokenizer: Base LLM/VLM tokenizer to extend.
        :param bins: Number of bins for each continuous value; we'll adopt a uniform binning strategy.
        :param min_action: Minimum action value (for clipping, setting lower bound on bin interval).
        :param max_action: Maximum action value (for clipping, setting upper bound on bin interval).
        """
        self.tokenizer, self.n_bins, self.min_action, self.max_action = tokenizer, bins, min_action, max_action

        # Create Uniform Bins + Compute Bin Centers
        self.bins = self.get_bins(min_action, max_action, self.n_bins, use_norm_bins=False)
        self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0

        # [Contract] Set "action_token_begin_idx" based on `self.tokenizer.vocab_size - (self.n_bins + 1)`
        #   =>> Assumes we're always overwriting the final `n_bins` tokens of the vocabulary!
        self.action_token_begin_idx: int = int(self.tokenizer.vocab_size - (self.n_bins + 1))

    def __call__(self, action: np.ndarray) -> Union[str, List[str]]:
        """Clip & bin actions to *the last `n_bins` tokens* of the vocabulary (e.g., tokenizer.vocab[-256:])."""
        action = np.clip(action, a_min=float(self.min_action), a_max=float(self.max_action))
        discretized_action = np.digitize(action, self.bins)

        # Handle single element vs. batch
        if len(discretized_action.shape) == 1:
            return self.tokenizer.decode(list(self.tokenizer.vocab_size - discretized_action))
        else:
            return self.tokenizer.batch_decode((self.tokenizer.vocab_size - discretized_action).tolist())

    def decode_token_ids_to_actions(self, action_token_ids: np.ndarray) -> np.ndarray:
        """
        Returns continuous actions for discrete action token IDs.

        NOTE =>> Because of the way the actions are discretized w.r.t. the bins (and not the bin centers), the
                 digitization returns bin indices between [1, # bins], inclusive, when there are actually only
                 (# bins - 1) bin intervals.

                 Therefore, if the digitization returns the last possible index, we map this to the last bin interval.

        EXAMPLE =>> Let's say self._bins has 256 values. Then self._bin_centers has 255 values. Digitization returns
                    indices between [1, 256]. We subtract 1 from all indices so that they are between [0, 255]. There
                    is still one index (i==255) that would cause an out-of-bounds error if used to index into
                    self._bin_centers. Therefore, if i==255, we subtract 1 from it so that it just becomes the index of
                    the last bin center. We implement this simply via clipping between [0, 255 - 1].
        """
        discretized_actions = self.tokenizer.vocab_size - action_token_ids
        discretized_actions = np.clip(discretized_actions - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1)

        return self.bin_centers[discretized_actions]

    @property
    def vocab_size(self) -> int:
        return self.n_bins
    
    def get_bins(self, min_action, max_action, n_bins, use_norm_bins=False):
        if use_norm_bins:
            a = np.linspace( 0.5,   1.0,   15, endpoint=True)
            b = np.linspace( 0.25,  0.5,   40, endpoint=False)
            c = np.linspace( 0.1,   0.25,  42, endpoint=False)
            d = np.linspace(-0.1,   0.1,   62, endpoint=False)
            e = np.linspace(-0.25, -0.1,   42, endpoint=False)
            f = np.linspace(-0.5,  -0.25,  40, endpoint=False)
            g = np.linspace(-1,    -0.5,   15, endpoint=False)
            return np.concatenate((g, f, e, d, c, b, a), axis=0)
        else:
            return np.linspace(min_action, max_action, n_bins)


def encode_actions_real(sentence, action_tokenizer, statistics=None):
    if statistics is not None:
        assert Path(statistics).exists(), "Statistical file not found."
        with open(statistics, "r") as file_:
            config = yaml.safe_load(file_)
        act_std = np.array(config["act_min_bound"])
        act_mean = np.array(config["act_max_bound"])

        act_mean = np.tile(act_mean, 5)
        act_std = np.tile(act_std, 5)

        actions = sentence.split(" ")
        actions = [float(action) for action in actions]
        actions = (np.array(actions) - act_mean) / act_std
        
        # 3 sigma clipping
        actions = actions / 3
        actions = np.clip(actions, a_min=-1, a_max=1)
        actions_lang = action_tokenizer(actions)
    return actions_lang


def encode_actions(sentence, action_tokenizer, statistics=None):
    if statistics is not None:
        # use the absolute action
        assert Path(statistics).exists(), "Statistical file not found."
        with open(statistics, "r") as file_:
            config = yaml.safe_load(file_)
        act_min_bound = np.array(config["act_min_bound"])
        act_max_bound = np.array(config["act_max_bound"])
        act_min_bound = np.tile(act_min_bound, 5)
        act_max_bound = np.tile(act_max_bound, 5)

        actions = sentence.split(" ")
        actions = [float(action) for action in actions]
        actions = (np.array(actions) - act_min_bound) / (act_max_bound - act_min_bound)
        actions = actions * 2 - 1
        actions_lang = action_tokenizer(actions)
    else:
        # use the relative action
        actions = sentence.split(" ")
        actions = [float(action) for action in actions]
        actions_lang = action_tokenizer(actions)
    return actions_lang


def denormalize_actions(actions: np.float32, statistics=None):
    assert Path(statistics).exists(), "Statistical file not found."
    with open(statistics, "r") as file_:
        config = yaml.safe_load(file_)
    act_min_bound = np.array(config["act_min_bound"])
    act_max_bound = np.array(config["act_max_bound"])

    times = len(actions) // len(act_min_bound)
    act_min_bound = np.tile(act_min_bound, times)
    act_max_bound = np.tile(act_max_bound, times)
    actions = (actions + 1) / 2
    actions = actions * (act_max_bound - act_min_bound) + act_min_bound
    return actions


def denormalize_actions_real(actions: np.float32, statistics=None):
    assert Path(statistics).exists(), "Statistical file not found."
    with open(statistics, "r") as file_:
        config = yaml.safe_load(file_)
    act_std = np.array(config["act_min_bound"])
    act_mean = np.array(config["act_max_bound"])

    act_mean = np.tile(act_mean, 5)
    act_std = np.tile(act_std, 5)
    actions = actions * 3
    actions = actions * act_std + act_mean
    return actions


def encode_robot_obs(sentence, action_tokenizer, statistics=None):
    assert Path(statistics).exists(), "Statistical file not found."
    with open(statistics, "r") as file_:
        config = yaml.safe_load(file_)
    
    robot_obs_mean = np.array(config["robot_obs"][0]["mean"])
    robot_obs_std = np.array(config["robot_obs"][0]["std"])

    robot_obs = sentence.split(" ")
    assert len(robot_obs) == 15 or len(robot_obs) == 75, "The observation of the robot is mismatch."
    robot_obs = [float(item) for item in robot_obs]
    if len(robot_obs) == 75:
        robot_obs_mean = np.tile(robot_obs_mean, 5)
        robot_obs_std = np.tile(robot_obs_std, 5)

    # According to 3 sigma principle, value is in the range of (-3, 3).
    robot_obs = (np.array(robot_obs) - robot_obs_mean) / (robot_obs_std + 0.000001)
    
    # Convert value to the range of (-1, 1).
    robot_obs = robot_obs / 3.0

    robot_obs_lang = action_tokenizer(robot_obs)
    return robot_obs_lang
