"""Learn four action values per pose from the arm's own transitions.

Run: python3 classes/ml/labs/robot-arm/q_table.py
Seed 7, 300 episodes, at most 400 moves each. Evaluation freezes Q.
"""
import numpy as np
from world import HOME, GAMMA, step, evaluate


def train(seed=7, episodes=300, observe=None):
    rng = np.random.default_rng(seed)
    q = np.zeros((19, 19, 4))
    lengths, first = [], None
    records = []
    for episode in range(episodes):
        state = HOME
        for move in range(400):
            action = int(rng.integers(4)) if rng.random() < 0.2 else int(np.argmax(q[state] + 1e-9 * rng.random(4)))
            after, reward, done = step(state, action)
            target = reward + (0 if done else GAMMA * q[after].max())
            if reward and first is None:
                first = {"episode": episode + 1, "record": [*state, action, reward, *after, done],
                         "previous": records[-1] if records else None}
            q[state][action] += 0.5 * (target - q[state][action])
            records.append([*state, action, reward, *after, done])
            state = after
            if done:
                break
        lengths.append(move + 1)
        if observe:
            observe(episode + 1, q)
    return {"seed": seed, "episodes": episodes, "lengths": lengths, "first": first,
            "evaluation": evaluate(q), "q": q}


if __name__ == "__main__":
    result = train()
    print("First success episode:", result["first"]["episode"])
    print("Training moves:", sum(result["lengths"]))
    print("Evaluation:", {k: v for k, v in result["evaluation"].items() if k != "starts"})