"""ACT, opened up: the policy LeRobot trains, run on its own simulated task (two arms hand over a cube).

Run (the lerobot environment, see README.md; the trained policies come from act_runs.sh):
    MUJOCO_GL=egl ~/.virtualenvs/lerobot/bin/python classes/ml/labs/lerobot/act_tour.py

The first run downloads the dataset (lerobot/aloha_sim_transfer_cube_human) into ~/.cache/huggingface/lerobot.
"""
from pathlib import Path

import gymnasium as gym
import gym_aloha  # noqa: F401  (registers the ALOHA tasks with gymnasium)
import numpy as np
import torch
from lerobot.configs import PreTrainedConfig
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.envs.utils import preprocess_observation
from lerobot.policies.act.modeling_act import ACTPolicy
from lerobot.policies.factory import make_pre_post_processors

REPO = "lerobot/aloha_sim_transfer_cube_human"
RUNS = Path.home() / "lerobot-runs/w4b"
PUBLISHED = RUNS / "published_migrated"                     # lerobot/act_aloha_sim_transfer_cube_human, 100,000 steps
OURS = RUNS / "train_act/checkpoints/010000/pretrained_model"  # our own run, 10,000 steps


def load(path, **overrides):
    """The policy and its two processors: normalise on the way in, un-normalise on the way out.
    `overrides` change settings that need no retraining, such as n_action_steps."""
    config = PreTrainedConfig.from_pretrained(path)
    for key, value in overrides.items():
        setattr(config, key, value)
    policy = ACTPolicy.from_pretrained(path, config=config).eval()
    pre, post = make_pre_post_processors(policy.config, pretrained_path=str(path))
    return policy, pre, post


def parts(policy):
    """Parameter count of each named part of the network."""
    counts = {}
    for name, p in policy.model.named_parameters():
        part = name.split(".")[0]
        counts[part] = counts.get(part, 0) + p.numel()
    return counts


def training_batch(policy, size=8, seed=0):
    """One batch as the training script builds it: each frame with its next 100 actions."""
    dataset = LeRobotDataset(REPO, video_backend="pyav",
                             delta_timestamps={"action": [i / 50 for i in range(policy.config.chunk_size)]})
    loader = torch.utils.data.DataLoader(dataset, batch_size=size, shuffle=True, generator=torch.Generator().manual_seed(seed))
    return next(iter(loader))


def tokens(policy, batch):
    """What the transformer encoder reads: one latent token, one joint token, and the image as a grid of tokens."""
    features = policy.model.backbone(batch["observation.images.top"])["feature_map"]   # ResNet-18's last layer
    return {"image": tuple(batch["observation.images.top"].shape), "feature_map": tuple(features.shape),
            "image_tokens": features.shape[2] * features.shape[3], "all_tokens": 2 + features.shape[2] * features.shape[3]}


def loss_by_hand(policy, batch, seed=0):
    """ACT's training loss written out, next to LeRobot's own policy.forward on the same batch."""
    policy.train()
    batch = dict(batch, **{"observation.images": [batch["observation.images.top"]]})
    torch.manual_seed(seed)
    predicted, (mu, log_sigma_x2) = policy.model(batch)                 # z is sampled inside, from mu and sigma
    keep = ~batch["action_is_pad"].unsqueeze(-1)                        # chunk entries past the episode's end don't count
    l1 = ((batch["action"] - predicted).abs() * keep).sum() / (keep.sum() * predicted.shape[-1])
    kl = (-0.5 * (1 + log_sigma_x2 - mu ** 2 - log_sigma_x2.exp())).sum(-1).mean()
    torch.manual_seed(seed)
    lerobot_loss, _ = policy.forward(batch)
    policy.eval()
    return {"l1": l1.item(), "kl": kl.item(), "loss": (l1 + policy.config.kl_weight * kl).item(), "lerobot": lerobot_loss.item()}


@torch.no_grad()
def z_effect(policy, batch, seed=0):
    """Does the style variable z change the chunk? Compare z = 0 (what the robot uses) with z drawn at random."""
    batch = dict(batch, **{"observation.images": [batch["observation.images.top"]]})
    model, gen = policy.model, torch.Generator().manual_seed(seed)
    model.eval()                                                        # eval mode: the model sets z = 0
    zero = model(batch)[0]
    shifts = []
    for _ in range(4):
        z = torch.randn(zero.shape[0], policy.config.latent_dim, generator=gen)
        swap = model.encoder_latent_input_proj.register_forward_hook(
            lambda layer, _, __: torch.nn.functional.linear(z.to(layer.weight.device), layer.weight, layer.bias))
        shifts.append((model(batch)[0] - zero).abs().mean().item())   # the same frames, a random z instead of zeros
        swap.remove()
    return {"chunk_size": zero.abs().mean().item(), "shift_random_z": float(np.mean(shifts))}


@torch.no_grad()
def rollout(policy, pre, post, seed, steps=400, keep_chunks=False):
    """One closed-loop episode in the simulator: observe, act, repeat, 50 times a second for 8 seconds."""
    env = gym.make("gym_aloha/AlohaTransferCube-v0", obs_type="pixels_agent_pos", render_mode="rgb_array", max_episode_steps=steps)
    observation, _ = env.reset(seed=seed)
    policy.reset()
    actions, rewards, chunks, frames = [], [], [], []
    for t in range(steps):
        obs = pre(preprocess_observation(observation))
        if keep_chunks and (policy.config.temporal_ensemble_coeff is not None or len(policy._action_queue) == 0):
            chunks.append((t, post(policy.predict_action_chunk(obs))[0].cpu().numpy()))   # the chunk this call writes
        action = post(policy.select_action(obs))[0].cpu().numpy()
        if t % 10 == 0:
            frames.append(observation["pixels"]["top"])
        observation, reward, terminated, truncated, info = env.step(action)
        actions.append(action)
        rewards.append(reward)
        if terminated or truncated:
            break
    env.close()
    return {"success": bool(max(rewards) == 4), "stage": int(max(rewards)), "actions": np.array(actions),
            "chunks": chunks, "frames": frames}


def cube_start(seed):
    """Where the simulator puts the cube for a given start (the same rule gym_aloha uses)."""
    return np.random.RandomState(seed).uniform([0.0, 0.4], [0.2, 0.6])


if __name__ == "__main__":
    torch.set_num_threads(8)
    policy, pre, post = load(PUBLISHED)
    for part, n in parts(policy).items():
        print(f"  {part:36s} {n:12,d}")
    print(f"total {sum(parts(policy).values()):,} parameters")

    batch = training_batch(policy)
    for key in ("observation.images.top", "observation.state", "action", "action_is_pad"):
        print(f"  batch {key:24s} {tuple(batch[key].shape)}")
    normalised = pre(dict(batch))
    print("tokens:", tokens(policy, normalised))
    print("loss:", loss_by_hand(policy, normalised))
    print("z:", z_effect(policy, normalised))

    episode = rollout(policy, pre, post, seed=1000, keep_chunks=True)
    print(f"start 1000 (cube at {cube_start(1000).round(3)}): stage {episode['stage']}, success {episode['success']}, "
          f"chunks written at steps {[t for t, _ in episode['chunks']]}")