"""PyTorch companion to learning_lab.py; CPU only. PyTorch is optional.

Keep both files in the same folder. See the article for validation status.
"""
import json
import numpy as np
import torch
from torch import nn
from torch.distributions import Categorical
from learning_lab import (ACTIONS, GAMMA, HORIZON, TRAIN_STARTS, demonstrations,
                          evaluate, transition)


def main():
    torch.manual_seed(7)
    torch.set_num_threads(1)
    rng = np.random.default_rng(7)
    # An embedding gives one trainable row of action logits per discrete state.
    bc, policy = nn.Embedding(11, 3), nn.Embedding(11, 3)
    nn.init.zeros_(bc.weight)
    nn.init.zeros_(policy.weight)
    rows = [(s+5, a) for episode, s, a in demonstrations(7) if episode < 96]
    states = torch.tensor([r[0] for r in rows], dtype=torch.long)
    actions = torch.tensor([r[1] for r in rows], dtype=torch.long)
    optimizer = torch.optim.SGD(bc.parameters(), lr=0.8)
    for _ in range(300):
        loss = nn.functional.cross_entropy(bc(states), actions)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    optimizer = torch.optim.SGD(policy.parameters(), lr=0.3)
    baseline = np.zeros(11)
    for _ in range(300):
        losses = []
        for _ in range(32):
            state, episode = int(rng.choice(TRAIN_STARTS)), []
            old_baseline = baseline.copy()
            for _ in range(HORIZON):
                distribution = Categorical(logits=policy(torch.tensor(state+5)))
                action_index = distribution.sample()
                nxt, reward, done = transition(state, int(ACTIONS[int(action_index)]), rng)
                episode.append((state, distribution.log_prob(action_index), reward))
                state = nxt
                if done:
                    break
            total, returns = 0.0, []
            for _, _, reward in reversed(episode):
                total = reward + GAMMA*total
                returns.append(total)
            for (state, logp, _), total in zip(episode, reversed(returns)):
                losses.append(-logp*(total-old_baseline[state+5]))
                baseline[state+5] += 0.05*(total-baseline[state+5])
        optimizer.zero_grad()
        (torch.stack(losses).sum()/32).backward()
        optimizer.step()
    # Tensor/NumPy RNGs differ, so do not expect identical training trajectories.
    with torch.no_grad():
        result = {"torch": torch.__version__, "device": "cpu",
                  "bc": evaluate(bc.weight.detach().cpu().numpy(), 10007),
                  "reinforce": evaluate(policy.weight.detach().cpu().numpy(), 10007)}
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()