"""Task plugin interface: the only place where MiniJEV touches a domain.

A task turns its native states into (a) compact arrays for storage and
(b) model inputs for one of MiniJEV's state encoders. Everything else (model,
training, calibration, export, runtime) is task-agnostic.

Plugins live in `environments/<name>/task.py` and expose `make_task(cfg)`.
"""
from __future__ import annotations

import importlib
from typing import Any, Sequence

import numpy as np


class Task:
    name: str = "task"
    #: which MiniJEV state encoder consumes this task's inputs ("grid", "text", ...)
    state_encoder: str = "grid"
    default_question: str = ""
    default_choices: tuple[str, ...] = ()

    def __init__(self, cfg: dict | None = None):
        self.cfg = cfg or {}

    # labels -----------------------------------------------------------------
    #: stored label -> label shown to the model (e.g. Snake "UP" -> "move up"); empty = identity
    @property
    def choice_alias(self) -> dict[str, str]:
        return {}

    def model_label(self, label: str) -> str:
        return self.choice_alias.get(label, label)

    def action_of(self, label: str) -> str:
        """Model-facing label -> native action (inverse of choice_alias)."""
        inv = {v: k for k, v in self.choice_alias.items()}
        return inv.get(label, label)

    # storage format --------------------------------------------------------
    def states_to_arrays(self, states: Sequence[Any]) -> dict[str, np.ndarray]:
        raise NotImplementedError

    def encode_arrays(self, arrays: dict[str, np.ndarray], idx: np.ndarray) -> dict[str, np.ndarray]:
        """Model inputs for rows `idx` of stored arrays (vectorised numpy)."""
        raise NotImplementedError

    # convenience ------------------------------------------------------------
    def encode_states(self, states: Sequence[Any]) -> dict[str, np.ndarray]:
        arrays = self.states_to_arrays(states)
        return self.encode_arrays(arrays, np.arange(len(states)))

    def encoder_class(self):
        """Optional task-specific StateEncoder class (must keep the generic encoder's parameters)."""
        return None

    def encoder_config(self) -> dict:
        """Hyper-parameters the state encoder needs from the task (channels, ...)."""
        return {}

    def example_states(self, n: int = 4, seed: int = 0) -> list[Any]:
        """A few valid states (used for ONNX export tracing and parity checks)."""
        raise NotImplementedError


BUILTIN = {"text": "minijev.text_task", "log": "minijev.text_task"}


def get_task(name: str, cfg: dict | None = None) -> Task:
    mod = importlib.import_module(BUILTIN.get(name) or f"environments.{name.replace('-', '_')}.task")
    return mod.make_task(cfg or {})


# ----------------------------------------------------------------- ragged utils
def ragged_pack(seqs: Sequence[Sequence[int]] | Sequence[np.ndarray], dtype=np.int32) -> tuple[np.ndarray, np.ndarray]:
    lengths = np.fromiter((len(s) for s in seqs), dtype=np.int64, count=len(seqs))
    offsets = np.zeros(len(seqs) + 1, dtype=np.int64)
    np.cumsum(lengths, out=offsets[1:])
    flat = np.concatenate([np.asarray(s, dtype=dtype) for s in seqs]) if len(seqs) else np.zeros(0, dtype)
    return flat.astype(dtype, copy=False), offsets


def ragged_take(flat: np.ndarray, offsets: np.ndarray, idx: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Gather rows `idx` of a ragged array; returns (flat, lengths)."""
    starts = offsets[idx]
    lengths = offsets[idx + 1] - starts
    total = int(lengths.sum())
    if total == 0:
        return flat[:0], lengths
    rel = np.arange(total) - np.repeat(np.cumsum(lengths) - lengths, lengths)
    return flat[np.repeat(starts, lengths) + rel], lengths