"""The existing hg-robot arm, made into a small reusable environment.

Angles are on a 10-degree grid; links are 10 cm. A forbidden destination
rejects the command. This checks poses, not the continuous motion between them.
No dynamics or physical contact simulation is performed.
"""
import numpy as np

ANGLES = np.arange(0, 181, 10)
ELBOW = np.arange(-180, 1, 10)
ACTIONS = np.array([[1, 0], [-1, 0], [0, 1], [0, -1]])
ACTION_NAMES = ["Shoulder +10°", "Shoulder −10°", "Elbow +10°", "Elbow −10°"]
HOME = (12, 6)
TARGET = np.array([14.0, 4.0])
GAMMA = 0.9


def forward(shoulder, elbow):
    a, b = np.deg2rad([shoulder, shoulder + elbow])
    joint = 10 * np.array([np.cos(a), np.sin(a)])
    return joint, joint + 10 * np.array([np.cos(b), np.sin(b)])


def forbidden(shoulder, elbow):
    joint, hand = forward(shoulder, elbow)
    fraction = np.linspace(0, 1, 60)[:, None]
    points = np.concatenate([fraction * joint, joint + fraction * (hand - joint)])
    below = (points[:, 1] < -1e-9).any()
    barrier = ((points[:, 0] >= 9) & (points[:, 0] <= 10.5)
               & (points[:, 1] >= 0) & (points[:, 1] <= 5)).any()
    return bool(below or barrier)


BLOCKED = np.array([[forbidden(a, b) for b in ELBOW] for a in ANGLES])
DISTANCE = np.array([[np.linalg.norm(forward(a, b)[1] - TARGET) for b in ELBOW] for a in ANGLES])
GOAL = (DISTANCE < 1.5) & ~BLOCKED
STARTS = [tuple(map(int, x)) for x in np.argwhere(~BLOCKED & ~GOAL)]


def step(state, action):
    i, j = np.array(state) + ACTIONS[action]
    if not (0 <= i < 19 and 0 <= j < 19) or BLOCKED[i, j]:
        return tuple(state), 0.0, False
    return (int(i), int(j)), float(GOAL[i, j]), bool(GOAL[i, j])


def encode(states):
    return np.asarray(states, dtype=np.float32) / 9 - 1


def rollout(values, start=HOME, limit=60):
    state = tuple(start)
    path = [state]
    rejected = 0
    for _ in range(limit):
        action = int(np.argmax(values[state]))
        after, reward, done = step(state, action)
        rejected += after == state
        path.append(after)
        state = after
        if done:
            return {"path": path, "success": True, "moves": len(path) - 1, "rejected": rejected}
    return {"path": path, "success": False, "moves": limit, "rejected": rejected}


def evaluate(values):
    trials = [rollout(values, s) for s in STARTS]
    return {"trials": len(trials), "successes": sum(t["success"] for t in trials),
            "home": rollout(values), "starts": [t["success"] for t in trials]}


def reference():
    """Full-model value iteration: an evaluation baseline, never training labels."""
    values = np.zeros((19, 19))
    for _ in range(100):
        q = np.zeros((19, 19, 4))
        for state in STARTS:
            for action in range(4):
                after, reward, done = step(state, action)
                q[state][action] = reward + (0 if done else GAMMA * values[after])
        new = q.max(axis=2)
        if np.max(np.abs(new - values)) < 1e-10:
            break
        values = new
    return q


if __name__ == "__main__":
    q = reference()
    print(f"{len(STARTS)} nonterminal starts; {int(GOAL.sum())} goal poses")
    print("Reference from home:", rollout(q))
    print("Home value:", q[HOME].max())