from ..base import BaseModel


class Datum(BaseModel):
    model_input: list[int]
    """输入 token ids，为数字列表"""

    loss_fn_inputs: dict[str, list[float] | list[int]]
    """损失函数，包含的参数有 weights, target_tokens 等"""

    def get_data(self, loss_fn: str = "cross_entropy") -> dict:
        """校验参数，并设置为传入给训练服务的参数模型"""
        if loss_fn == "cross_entropy":
            return self._get_data_cross_entropy()
        elif loss_fn == "importance_sampling":
            return self._get_data_importance_sampling()
        elif loss_fn == "ppo":
            return self._get_data_ppo()
        else:
            raise ValueError(f"不支持的损失函数类型: {loss_fn}")

    def _get_data_cross_entropy(self):
        if "target_tokens" not in self.loss_fn_inputs:
            raise ValueError("必须传入参数 target_tokens")
        if len(self.model_input) != len(self.loss_fn_inputs["target_tokens"]):
            raise ValueError("输入 token 和目标 token 的长度需要相同")
        weights = None
        if "weights" in self.loss_fn_inputs:
            weights = self.loss_fn_inputs["weights"]
        return {
            "input_ids": self.model_input,
            "target_ids": self.loss_fn_inputs["target_tokens"],
            "weights": weights,
        }

    def _get_data_importance_sampling(self):
        if "target_tokens" not in self.loss_fn_inputs:
            raise ValueError("必须传入参数 target_tokens")
        if "logprobs" not in self.loss_fn_inputs:
            raise ValueError("必须传入参数 logprobs")
        if "advantages" not in self.loss_fn_inputs:
            raise ValueError("必须传入参数 advantages")
        if (
            len(self.model_input) != len(self.loss_fn_inputs["target_tokens"])
            or len(self.model_input) != len(self.loss_fn_inputs["logprobs"])
            or len(self.model_input) != len(self.loss_fn_inputs["advantages"])
        ):
            raise ValueError(
                "输入 token, 目标 token, logprobs 和 advantages 的长度需要相同"
            )

        return {
            "input_ids": self.model_input,
            "target_ids": self.loss_fn_inputs["target_tokens"],
            "logprobs": self.loss_fn_inputs["logprobs"],
            "advantages": self.loss_fn_inputs["advantages"],
        }

    def _get_data_ppo(self):
        return self._get_data_importance_sampling()
