"""
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
import json
import torch

class ActionTokenizer:
    def __init__(
        self, tokenizer: PreTrainedTokenizerBase, bins: int = 256, min_action: int = -1, max_action: int = 1, use_norm_bins: bool = False
    ) -> 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)
        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 + 1e-5)
        
        # 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:
        print("use statistics")
        # 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
        # print("not use statistics")
        actions = sentence.split(" ")
        actions = [float(action) for action in actions]
        actions_lang = action_tokenizer(actions)

    return actions_lang
def encode_actions_forpipper(sentence, action_tokenizer, statistics=None):
    assert statistics is not None, "Statistics file path is required."
    assert Path(statistics).exists(), "Statistical file not found."

    # 读取 JSON
    with open(statistics, "r") as f:
        config = json.load(f)

    # 从 JSON 中提取动作最小值 / 最大值
    act_min_bound = np.array(config["action"]["min"])
    act_max_bound = np.array(config["action"]["max"])

    # 支持 Tensor 或字符串输入
    if isinstance(sentence, torch.Tensor):
        actions = sentence.cpu().numpy()
    elif isinstance(sentence, list):
        actions = np.array(sentence)
    elif isinstance(sentence, str):
        actions = np.array([float(a) for a in sentence.strip().split()])
    else:
        raise TypeError("Unsupported type for `sentence`")

    action_dim = len(act_min_bound)

    # 如果 actions 是多个动作拼接（如两帧），自动扩展
    if len(actions) != action_dim:
        assert len(actions) % action_dim == 0, f"动作长度 {len(actions)} 不是基础维度 {action_dim} 的整数倍"
        repeat_factor = len(actions) // action_dim
        act_min_bound = np.tile(act_min_bound, repeat_factor)
        act_max_bound = np.tile(act_max_bound, repeat_factor)

    # Min-max 归一化到 [-1, 1]
    normed_actions = (actions - act_min_bound) / (act_max_bound - act_min_bound + 1e-5)
    normed_actions = normed_actions * 2 - 1

    # Tokenizer 处理
    actions_lang = action_tokenizer(normed_actions)
    return actions_lang
def decode_actions_forpipper(normed_actions, statistics=None):
    assert statistics is not None, "Statistics file path is required."
    assert Path(statistics).exists(), "Statistical file not found."

    # 读取 JSON 文件，提取动作范围
    with open(statistics, "r") as f:
        config = json.load(f)

    act_min_bound = np.array(config["action"]["min"])
    act_max_bound = np.array(config["action"]["max"])

    # 如果传入的是 Tensor / list / str，都统一转换为 numpy
    if isinstance(normed_actions, torch.Tensor):
        normed_actions = normed_actions.cpu().numpy()
    elif isinstance(normed_actions, list):
        normed_actions = np.array(normed_actions)
    elif isinstance(normed_actions, str):
        normed_actions = np.array([float(a) for a in normed_actions.strip().split()])

    action_dim = len(act_min_bound)

    # 处理重复动作帧情况（如拼接了多帧）
    if len(normed_actions) != action_dim:
        assert len(normed_actions) % action_dim == 0, \
            f"动作长度 {len(normed_actions)} 不是基础维度 {action_dim} 的整数倍"
        repeat_factor = len(normed_actions) // action_dim
        act_min_bound = np.tile(act_min_bound, repeat_factor)
        act_max_bound = np.tile(act_max_bound, repeat_factor)

    # 反归一化：[-1, 1] → [min, max]
    normed_actions = (normed_actions + 1) / 2
    actions = normed_actions * (act_max_bound - act_min_bound + 1e-5) + act_min_bound

    return actions
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, "The observation of the robot is mismatch."
    robot_obs = [float(item) for item in robot_obs]

    # 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 + 1e-5)
    
    # 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

def encode_robot_obs_forpipper(sentence, action_tokenizer, statistics=None):
    assert Path(statistics).exists(), "Statistical file not found."

    with open(statistics, "r") as f:
        stats = json.load(f)

    # 读取 observation.state 部分的 mean 和 std（注意：不是 robot_obs，而是 observation.state）
    obs_mean = np.array(stats["observation.state"]["mean"])
    obs_std = np.array(stats["observation.state"]["std"])

    # ✅ 支持 tensor / list / str
    if isinstance(sentence, torch.Tensor):
        robot_obs = sentence.cpu().numpy().tolist()
    elif isinstance(sentence, list):
        robot_obs = sentence
    elif isinstance(sentence, str):
        robot_obs = [float(x) for x in sentence.strip().split()]
    else:
        raise TypeError("Unsupported input type for `sentence`")

    assert len(robot_obs) == len(obs_mean), f"Expected {len(obs_mean)} obs values, got {len(robot_obs)}"

    # 归一化：根据 3σ 原则缩放至 [-1, 1]
    robot_obs = (np.array(robot_obs) - obs_mean) / (obs_std + 1e-5)
    robot_obs = robot_obs / 3.0  # 缩放至 [-1, 1]

    # 使用 action_tokenizer 进行离散化（或文本化）
    robot_obs_lang = action_tokenizer(robot_obs)
    return robot_obs_lang

