"""Neural Q-learning on the same arm, using ordinary PyTorch.

Run: python3 classes/ml/labs/robot-arm/q_network.py
Actions collect transitions; optimizer steps learn from replay batches.
Evaluation freezes weights and uses the largest estimate, without exploration.
"""
import numpy as np
import torch
from torch import nn
from world import HOME, GAMMA, step, encode, evaluate


def make_network(initialization="custom"):
    model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 4))
    if initialization in ("custom", "scaled-hidden"):
        nn.init.normal_(model[0].weight, 0, 0.7)
        nn.init.zeros_(model[0].bias)
    if initialization in ("custom", "tiny-output"):
        nn.init.uniform_(model[2].weight, -0.003, 0.003)
        nn.init.zeros_(model[2].bias)
    return model


def train(seed=0, lr=0.3, episodes=300, observe=None,
          initialization="custom", target_period=100):
    torch.set_num_threads(1)
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    network = make_network(initialization)
    target_network = network if target_period == 1 else make_network(initialization)
    target_network.load_state_dict(network.state_dict())
    optimizer = torch.optim.SGD(network.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
    memory, lengths, losses = [], [], []
    interactions, updates, first_success = 0, 0, None
    if observe:
        observe(0, network, interactions, updates)
    for episode in range(episodes):
        state = HOME
        total_loss = 0
        for move in range(400):
            with torch.no_grad():
                values = network(torch.tensor(encode(state))).numpy()
            action = int(rng.integers(4)) if rng.random() < 0.2 else int(np.argmax(values + 1e-9 * rng.random(4)))
            after, reward, done = step(state, action)
            memory.append([*state, action, reward, *after, done])
            interactions += 1
            if reward and first_success is None:
                first_success = {"episode": episode + 1, "move": interactions}
            for _ in range(2):
                indices = rng.integers(len(memory), size=min(32, len(memory)))
                batch = np.asarray([memory[k] for k in indices], dtype=np.float32)
                states = torch.tensor(encode(batch[:, :2]))
                actions = torch.tensor(batch[:, 2:3], dtype=torch.long)
                rewards = torch.tensor(batch[:, 3])
                next_states = torch.tensor(encode(batch[:, 4:6]))
                terminal = torch.tensor(batch[:, 6])
                with torch.no_grad():
                    next_values = target_network(next_states).max(dim=1).values
                    targets = rewards + GAMMA * (1 - terminal) * next_values
                predictions = network(states).gather(1, actions).squeeze(1)
                loss = loss_fn(predictions, targets)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                updates += 1
                total_loss += loss.item()
            if target_period > 1 and interactions % target_period == 0:
                target_network.load_state_dict(network.state_dict())
            state = after
            if done:
                break
        lengths.append(move + 1)
        losses.append(total_loss / (2 * (move + 1)))
        if observe:
            observe(episode + 1, network, interactions, updates)
    with torch.no_grad():
        all_states = torch.tensor(encode([[i, j] for i in range(19) for j in range(19)]))
        q = network(all_states).numpy().reshape(19, 19, 4)
    result = {"seed": seed, "lr": lr, "optimizer": "SGD", "initialization": initialization,
              "target_period": target_period, "episodes": episodes, "moves": interactions, "updates": updates,
              "first_success": first_success, "lengths": lengths, "losses": losses,
              "evaluation": evaluate(q)}
    return network, result


if __name__ == "__main__":
    network, result = train()
    print("Training moves:", result["moves"])
    print("Home test:", result["evaluation"]["home"])
    print("Starts passed:", result["evaluation"]["successes"], "/", result["evaluation"]["trials"])