"""Snake as a MiniJEV task: SnakeState -> multi-channel board image.

Channels (board padded by a 1-cell wall border; smaller boards in a batch are
padded with more wall, which is semantically exact):

    0 wall   1 head   2 body (excl. head)   3 tail   4 food
    5 time-until-free: (L - i) / area for body segment i (head=0)
    6 row offset from the head   (r - r_head) / max(H, W)
    7 col offset from the head   (c - c_head) / max(H, W)

Channel 5 lets the model reason about which cells will open up as the snake
moves - the same information the oracle's time-aware search uses. Channels 6-7
make every cell egocentric ("the food is 3 up, 2 left"), which removes the
long learning plateau of discovering relative geometry from absolute positions.
"""
from __future__ import annotations

from typing import Sequence

import numpy as np

from minijev.tasks import Task, ragged_pack, ragged_take

from .env import ACTIONS, SnakeState

N_CHANNELS = 8
DIRS = {a: i for i, a in enumerate(ACTIONS)}


SEMANTIC_CHOICES = {"UP": "move up", "DOWN": "move down", "LEFT": "move left", "RIGHT": "move right"}


class SnakeTask(Task):
    """task.choice_style: "semantic" -> choices "move up/down/left/right" (v2 models);
    "legacy" (default, what pre-v2 checkpoints were trained with) -> "UP/DOWN/LEFT/RIGHT"."""

    name = "snake"
    state_encoder = "grid"
    default_question = "Choose the next move"

    @property
    def choice_alias(self) -> dict[str, str]:
        return SEMANTIC_CHOICES if self.cfg.get("choice_style", "legacy") == "semantic" else {}

    @property
    def default_choices(self) -> tuple[str, ...]:
        return tuple(self.model_label(a) for a in ACTIONS)

    def encoder_class(self):
        from .encoder import SnakeStateEncoder
        return SnakeStateEncoder

    def encoder_config(self) -> dict:
        return {"in_channels": N_CHANNELS}

    def states_to_arrays(self, states: Sequence[SnakeState]) -> dict[str, np.ndarray]:
        bodies = [[r * s.width + c for r, c in s.body] for s in states]
        flat, offsets = ragged_pack(bodies, np.int16)
        return {
            "body": flat,
            "body_offsets": offsets,
            "food": np.array([s.food[0] * s.width + s.food[1] if s.food else -1 for s in states], np.int16),
            "height": np.array([s.height for s in states], np.uint8),
            "width": np.array([s.width for s in states], np.uint8),
            "direction": np.array([DIRS[s.direction] for s in states], np.uint8),
        }

    def example_states(self, n=4, seed=0):
        from .agents import OracleAgent
        from .env import SnakeConfig, SnakeEnv

        states, oracle = [], OracleAgent()
        env = SnakeEnv(SnakeConfig.from_dict({"board_size": self.cfg.get("board_size", 10)}), seed=seed)
        while len(states) < n:
            states.append(env.state())
            for _ in range(7):  # spread samples along the game
                env.step(oracle.act(env.state()))
                if env.done:
                    env.reset()
        return states

    def encode_arrays(self, arrays, idx):
        return {"state_grid": encode_grid(arrays, np.asarray(idx))}


def encode_grid(arrays: dict[str, np.ndarray], idx: np.ndarray) -> np.ndarray:
    B = len(idx)
    hs = arrays["height"][idx].astype(np.int64)
    ws = arrays["width"][idx].astype(np.int64)
    H, W = int(hs.max()) + 2, int(ws.max()) + 2
    x = np.zeros((B, N_CHANNELS, H, W), dtype=np.float32)
    x[:, 0] = 1.0
    for h, w in set(zip(hs.tolist(), ws.tolist())):
        sel = np.nonzero((hs == h) & (ws == w))[0]
        x[sel, 0, 1:h + 1, 1:w + 1] = 0.0

    cells, lengths = ragged_take(arrays["body"], arrays["body_offsets"], idx)
    cells = cells.astype(np.int64)
    sample = np.repeat(np.arange(B), lengths)
    pos = np.arange(len(cells)) - np.repeat(np.cumsum(lengths) - lengths, lengths)  # index from head
    Ls = np.repeat(lengths, lengths)
    w_s = ws[sample]
    r, c = cells // w_s + 1, cells % w_s + 1
    head = pos == 0
    x[sample[head], 1, r[head], c[head]] = 1.0
    x[sample[~head], 2, r[~head], c[~head]] = 1.0
    tail = pos == Ls - 1
    x[sample[tail], 3, r[tail], c[tail]] = 1.0
    x[sample, 5, r, c] = (Ls - pos) / (hs[sample] * w_s)

    hr, hc = np.zeros(B, np.int64), np.zeros(B, np.int64)
    hr[sample[head]], hc[sample[head]] = r[head], c[head]
    scale = np.maximum(hs, ws).astype(np.float32)[:, None, None]
    x[:, 6] = (np.arange(H)[None, :, None] - hr[:, None, None]) / scale
    x[:, 7] = (np.arange(W)[None, None, :] - hc[:, None, None]) / scale

    food = arrays["food"][idx].astype(np.int64)
    has = np.nonzero(food >= 0)[0]
    x[has, 4, food[has] // ws[has] + 1, food[has] % ws[has] + 1] = 1.0
    return x


def make_task(cfg: dict) -> SnakeTask:
    return SnakeTask(cfg)