"""Action-chunk execution and normalization checks; no learned model/checkpoint."""
import json
import math


def rollout(chunk_size):
    x, actions, errors = 0.0, [], []
    for step in range(80):
        if not actions:
            predicted = x
            for _ in range(chunk_size):
                action = max(-2.0, min(2.0, 3.0*(1.0-predicted)))
                actions.append(action)
                predicted += 0.1*action
        x += 0.1*actions.pop(0)
        if step == 25:
            x += 0.5  # an unpredicted external displacement
        errors.append(abs(x-1.0))
    return {"chunk_size": chunk_size, "final_error": abs(x-1.0),
            "post_disturbance_error_sum": sum(errors[26:41])}


def denormalize(values, lows, highs):
    if len(values) != len(lows) or len(values) != len(highs):
        raise ValueError("action/statistics dimension mismatch")
    if any(not math.isfinite(x) for x in values+lows+highs):
        raise ValueError("non-finite input")
    if any(high <= low for low, high in zip(lows, highs)):
        raise ValueError("invalid normalization statistics")
    if any(abs(x) > 1 for x in values):
        raise ValueError("normalized action out of range")
    return [low + (v+1)*0.5*(high-low) for v, low, high in zip(values,lows,highs)]


if __name__ == "__main__":
    one, eight = rollout(1), rollout(8)
    restored = denormalize([0.0, 1.0, -1.0], [-0.1, -0.1, -0.1], [0.1, 0.1, 0.1])
    assert restored == [0.0, 0.1, -0.1]
    assert one["post_disturbance_error_sum"] < eight["post_disturbance_error_sum"]
    rejected = 0
    for bad in ([float("nan"), 0.0, 0.0], [2.0, 0.0, 0.0], [0.0]):
        try:
            denormalize(bad, [-0.1]*3, [0.1]*3)
        except ValueError:
            rejected += 1
    assert rejected == 3
    print(json.dumps({"execution_model": "ideal-1D-integrator-no-learned-model",
                      "feedback_every_step": one, "execute_eight_steps": eight,
                      "denormalized_delta_position_m": restored,
                      "invalid_actions_rejected": rejected}, indent=2))
