"""Identical collector/learner action conversion in actual joint radians."""
from __future__ import annotations
import numpy as np
from .data import ActionContract, BASE_ACTION_DIM, RESIDUAL_ACTION_DIM, vector


class ResidualActionAdapter:
    def __init__(self, contract: ActionContract):
        self.contract = contract
        self.indices = contract.indices
        self.scales_rad = np.array(contract.residual_scales_rad, dtype=np.float32)
        self.base_scales = np.array(contract.action_scale, dtype=np.float32)
        self.action_increments = self.scales_rad / self.base_scales[self.indices]

    def combine(self, base_action, residual):
        base = np.clip(vector(base_action, BASE_ACTION_DIM), -10.0, 10.0)
        residual = np.clip(vector(residual, RESIDUAL_ACTION_DIM), -1.0, 1.0)
        combined = base.copy()
        combined[self.indices] += self.action_increments * residual
        combined = np.clip(combined, -10.0, 10.0)
        executed_rad = (combined[self.indices] - base[self.indices]) * self.base_scales[self.indices]
        return combined, executed_rad

    def combine_tensor(self, base_action, residual):
        import torch
        indices = torch.as_tensor(self.indices, device=base_action.device)
        scales = torch.as_tensor(self.action_increments, device=base_action.device)
        combined = base_action.clamp(-10.0, 10.0).clone()
        combined[:, indices] = combined[:, indices] + residual.clamp(-1.0, 1.0) * scales
        return combined.clamp(-10.0, 10.0)