"""Imitation learning on the same arm: copy an expert's commands with ordinary PyTorch.

Run: python3 classes/ml/labs/robot-arm/imitation.py

The expert is the full-model reference from world.py. It stands in for a person who knows the way, so its
demonstrations are synthetic. The policy network never sees a reward: it learns which command the expert
chose at each demonstrated pose (behavioral cloning). Evaluation is the same frozen 60-command test as the
Q-table and Q-network workshops.
"""
import numpy as np
import torch
from torch import nn
from world import STARTS, GOAL, HOME, step, encode, evaluate, reference

EXPERT_Q = reference()


def expert(state):
    """The expert's command at a pose. Ties go to the first command in the list, so labels never flip."""
    return int(np.argmax(EXPERT_Q[state]))


def demonstrate(start):
    """One demonstration episode: the expert's (pose, command) pairs from `start` to the cube."""
    state, pairs = tuple(start), []
    while not GOAL[state]:
        action = expert(state)
        pairs.append((state, action))
        state = step(state, action)[0]
    return pairs


def demonstrations(count, seed, near_home=False):
    """`count` episodes from random starts: anywhere (broad) or among the 40 starts nearest home (narrow)."""
    rng = np.random.default_rng(seed)
    pool = sorted(STARTS, key=lambda s: abs(s[0] - HOME[0]) + abs(s[1] - HOME[1]))[:40] if near_home else STARTS
    starts = [pool[i] for i in rng.choice(len(pool), count, replace=False)]
    return [demonstrate(s) for s in starts]


def make_policy(hidden=16):
    return nn.Sequential(nn.Linear(2, hidden), nn.ReLU(), nn.Linear(hidden, 4))   # 2 angles in, 4 logits out


def train(pairs, hidden=16, lr=1.0, epochs=300, seed=0, validation=(), observe=None):
    torch.set_num_threads(1)
    torch.manual_seed(seed)
    policy = make_policy(hidden)
    optimizer = torch.optim.SGD(policy.parameters(), lr=lr)
    loss_fn = nn.CrossEntropyLoss()
    poses = torch.tensor(encode([p for p, a in pairs]))
    labels = torch.tensor([a for p, a in pairs])
    order = torch.Generator().manual_seed(seed)
    for epoch in range(epochs):
        for batch in torch.randperm(len(labels), generator=order).split(32):
            logits = policy(poses[batch])
            loss = loss_fn(logits, labels[batch])     # cross-entropy against the expert's command
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        if observe:
            observe(epoch + 1, policy, accuracy(policy, pairs), accuracy(policy, validation))
    return policy


def logits_table(policy):
    """The policy's four logits at every one of the 19 × 19 poses."""
    with torch.no_grad():
        grid = torch.tensor(encode([[i, j] for i in range(19) for j in range(19)]))
        return policy(grid).numpy().reshape(19, 19, 4)


def accuracy(policy, pairs):
    if not pairs:
        return None
    table = logits_table(policy)
    return float(np.mean([np.argmax(table[p]) == a for p, a in pairs]))


def dagger(pairs, rounds=4, starts_per_round=10, seed=0, **train_args):
    """DAgger: run the policy, ask the expert for the command at every new pose it visited, retrain."""
    rng = np.random.default_rng(1000 + seed)
    pairs = list(pairs)
    labelled = {p for p, a in pairs}
    history = []
    for r in range(rounds + 1):
        policy = train(pairs, seed=seed, **train_args)
        table = logits_table(policy)
        history.append({"round": r, "labels": len(pairs), "passed": evaluate(table)["successes"], "asked": []})
        if r == rounds:
            return policy, history
        visited = set()
        for i in rng.choice(len(STARTS), starts_per_round, replace=False):
            state = STARTS[i]
            for _ in range(60):                       # the policy drives; the expert only labels
                visited.add(state)
                state = step(state, int(np.argmax(table[state])))[0]
                if GOAL[state]:
                    break
        new = sorted(s for s in visited if s not in labelled)
        history[-1]["asked"] = new                    # the poses the expert was asked about this round
        pairs += [(s, expert(s)) for s in new]
        labelled.update(new)


if __name__ == "__main__":
    DRAW = 3        # the demonstration draw the page shows: the lower median of draws 0–9 (tools/run_imitation.py)
    episodes = demonstrations(20, seed=DRAW)
    train_pairs = [pair for e in episodes[:16] for pair in e]
    val_pairs = [pair for e in episodes[16:] for pair in e]
    print(f"Expert from home: {len(demonstrate(HOME))} commands. Demonstration draw {DRAW}: 20 episodes, "
          f"{len(train_pairs)} training pairs (16 episodes), {len(val_pairs)} validation pairs (4 episodes)")
    policy = train(train_pairs, seed=DRAW, validation=val_pairs)
    table = logits_table(policy)
    result = evaluate(table)
    print(f"Training accuracy {accuracy(policy, train_pairs):.3f}, validation accuracy {accuracy(policy, val_pairs):.3f}")
    print(f"Frozen test: home {'reached in ' + str(result['home']['moves']) + ' commands' if result['home']['success'] else 'failed'}; "
          f"starts passed {result['successes']} / {result['trials']}")