"""
Reusable training utilities for the JSpace language model.

The goal is to keep training scalable without hard-wiring it to one script:
sampling, validation, checkpointing, replay, and EWC all live behind a single
session object that can be reused by CLI tools, tests, and future workers.
"""
from __future__ import annotations

from contextlib import contextmanager
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Callable

import torch
import torch.nn.functional as F

from .language_data import CharTokenizer
from .child_data import ChildDialogExample, format_child_prompt
from .language_model import (
    EWCOptimizer,
    ExperienceReplay,
    ExpertPlasticity,
    JSpaceLanguageModel,
    LanguageConfig,
)


@dataclass
class LanguageTrainingConfig:
    seq_len: int = 64
    batch_size: int = 8
    lr: float = 1e-3
    weight_decay: float = 0.0
    ewc_lambda: float = 0.05
    max_grad_norm: float = 0.5
    replay_capacity: int = 500
    replay_batch_size: int = 4
    replay_weight: float = 0.5
    consolidate_every: int = 50
    consolidate_samples: int = 10
    validate_every: int = 50
    validate_batches: int = 4
    save_every: int = 100
    train_fraction: float = 0.98
    use_euler_during_train: bool = True


@contextmanager
def expert_integration_mode(model, use_rk4: bool):
    """Temporarily set expert integration mode."""
    original = [expert.use_rk4 for expert in model.experts]
    for expert in model.experts:
        expert.use_rk4 = use_rk4
    try:
        yield
    finally:
        for expert, enabled in zip(model.experts, original):
            expert.use_rk4 = enabled


class TokenBatchSampler:
    """Random contiguous sampler over a token stream with a train/val split."""

    def __init__(
        self,
        token_ids: list[int],
        seq_len: int = 64,
        train_fraction: float = 0.98,
    ):
        if not token_ids:
            raise ValueError("token_ids must not be empty")
        self.seq_len = seq_len
        min_len = seq_len + 2
        if len(token_ids) < min_len:
            repeats = min_len // len(token_ids) + 1
            token_ids = (token_ids * repeats)[:min_len]

        split = int(len(token_ids) * train_fraction)
        split = min(max(split, min_len), len(token_ids))
        self.train_tokens = token_ids[:split]
        self.val_tokens = token_ids[split:] if len(token_ids) - split >= min_len else token_ids[:split]

    def sample(self, batch_size: int, split: str = "train", device: str = "cpu") -> torch.Tensor:
        tokens = self.train_tokens if split == "train" else self.val_tokens
        max_start = max(1, len(tokens) - self.seq_len - 1)
        rows = []
        for _ in range(batch_size):
            start = torch.randint(0, max_start, (1,)).item()
            rows.append(tokens[start:start + self.seq_len])
        return torch.tensor(rows, dtype=torch.long, device=device)


class ChatBatchSampler:
    """Samples prompt/answer examples and masks loss to answer tokens only."""

    def __init__(
        self,
        examples: list[ChildDialogExample],
        tokenizer: CharTokenizer,
        seq_len: int = 96,
        train_fraction: float = 0.95,
    ):
        if not examples:
            raise ValueError("examples must not be empty")
        self.tokenizer = tokenizer
        self.seq_len = seq_len
        split = int(len(examples) * train_fraction)
        split = min(max(1, split), len(examples))
        self.train_examples = examples[:split]
        self.val_examples = examples[split:] or examples[:split]

    def sample(
        self,
        batch_size: int,
        split: str = "train",
        device: str = "cpu",
    ) -> tuple[torch.Tensor, torch.Tensor]:
        examples = self.train_examples if split == "train" else self.val_examples
        token_rows = []
        mask_rows = []
        for _ in range(batch_size):
            idx = torch.randint(0, len(examples), (1,)).item()
            token_row, mask_row = self.encode_example(examples[idx])
            token_rows.append(token_row)
            mask_rows.append(mask_row)
        return (
            torch.tensor(token_rows, dtype=torch.long, device=device),
            torch.tensor(mask_rows, dtype=torch.float32, device=device),
        )

    def encode_example(self, example: ChildDialogExample) -> tuple[list[int], list[float]]:
        prompt = format_child_prompt(example.user)
        answer = example.assistant + "\n"
        prompt_ids = self.tokenizer.encode(prompt)
        answer_ids = self.tokenizer.encode(answer)
        token_ids = (prompt_ids + answer_ids)[:self.seq_len]
        actual_len = len(token_ids)

        if actual_len < 2:
            token_ids = (token_ids + [0, 0])[:self.seq_len]
            actual_len = len(token_ids)

        padded = token_ids + [0] * max(0, self.seq_len - len(token_ids))
        padded = padded[:self.seq_len]

        answer_start = min(len(prompt_ids), self.seq_len)
        mask = []
        for target_pos in range(1, self.seq_len):
            is_answer = answer_start <= target_pos < actual_len
            mask.append(1.0 if is_answer else 0.0)
        return padded, mask


def save_language_checkpoint(
    path: str | Path,
    model: JSpaceLanguageModel,
    config: LanguageConfig,
    tokenizer: CharTokenizer,
    trainer_state: dict | None = None,
    metadata: dict | None = None,
):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save({
        "model": model.state_dict(),
        "config": config,
        "tokenizer_chars": tokenizer.chars,
        "trainer": trainer_state or {},
        "metadata": metadata or {},
    }, path)


class LanguageTrainingSession:
    """Stateful trainer for scalable language-model training."""

    def __init__(
        self,
        model: JSpaceLanguageModel,
        model_config: LanguageConfig,
        tokenizer: CharTokenizer,
        train_config: LanguageTrainingConfig | None = None,
        device: str = "cpu",
    ):
        self.model = model.to(device)
        self.model_config = model_config
        self.tokenizer = tokenizer
        self.train_config = train_config or LanguageTrainingConfig()
        self.device = device
        self.optimizer = EWCOptimizer(
            self.model,
            lr=self.train_config.lr,
            ewc_lambda=self.train_config.ewc_lambda,
            max_grad_norm=self.train_config.max_grad_norm,
            weight_decay=self.train_config.weight_decay,
        )
        self.replay_buffer = ExperienceReplay(
            capacity=self.train_config.replay_capacity,
            seq_len=self.train_config.seq_len,
        )
        self.plasticity = ExpertPlasticity(num_experts=model_config.num_experts)
        self.global_step = 0
        self.history: list[dict] = []

    def learn_batch(self, token_seq: torch.Tensor) -> dict:
        token_seq = token_seq.to(self.device)
        self.model.train()
        logits, info = self.model(token_seq)
        loss = self.next_token_loss(logits, token_seq)

        replay_loss = torch.tensor(0.0, device=self.device)
        replay_seq = None
        if self.train_config.replay_weight > 0:
            replay_seq = self.replay_buffer.sample(self.train_config.replay_batch_size)
        if replay_seq is not None:
            replay_seq = replay_seq.to(self.device)
            replay_logits, _ = self.model(replay_seq)
            replay_loss = self.next_token_loss(replay_logits, replay_seq)

        total_task_loss = loss + self.train_config.replay_weight * replay_loss
        total_loss = self.optimizer.step(total_task_loss)

        self.plasticity.update(info["alpha"].detach(), token_seq.detach())
        self.replay_buffer.push(token_seq.detach().cpu())

        return {
            "loss": float(loss.item()),
            "replay_loss": float(replay_loss.item()),
            "total_loss": float(total_loss),
            "w_norm_mean": float(info["w_norm"].mean().item()),
            "expert_usage": self.plasticity.usage.tolist(),
        }

    def learn_masked_batch(self, token_seq: torch.Tensor, loss_mask: torch.Tensor) -> dict:
        token_seq = token_seq.to(self.device)
        loss_mask = loss_mask.to(self.device)
        self.model.train()
        logits, info = self.model(token_seq)
        loss = self.next_token_loss(logits, token_seq, loss_mask)
        total_loss = self.optimizer.step(loss)
        self.plasticity.update(info["alpha"].detach(), token_seq.detach())
        self.replay_buffer.push(token_seq.detach().cpu())
        return {
            "loss": float(loss.item()),
            "replay_loss": 0.0,
            "total_loss": float(total_loss),
            "w_norm_mean": float(info["w_norm"].mean().item()),
            "expert_usage": self.plasticity.usage.tolist(),
        }

    def next_token_loss(
        self,
        logits: torch.Tensor,
        token_seq: torch.Tensor,
        loss_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        pred = logits[:, :-1]
        target = token_seq[:, 1:]
        if loss_mask is None:
            return F.cross_entropy(
                pred.reshape(-1, self.model_config.vocab_size),
                target.reshape(-1),
            )
        per_token = F.cross_entropy(
            pred.reshape(-1, self.model_config.vocab_size),
            target.reshape(-1),
            reduction="none",
        )
        mask = loss_mask.reshape(-1).to(per_token.device).float()
        return (per_token * mask).sum() / mask.sum().clamp_min(1.0)

    @torch.no_grad()
    def evaluate(self, sampler: TokenBatchSampler) -> float:
        was_training = self.model.training
        self.model.eval()
        losses = []
        for _ in range(max(1, self.train_config.validate_batches)):
            token_seq = sampler.sample(
                self.train_config.batch_size,
                split="val",
                device=self.device,
            )
            logits, _ = self.model(token_seq)
            loss = F.cross_entropy(
                logits[:, :-1].reshape(-1, self.model_config.vocab_size),
                token_seq[:, 1:].reshape(-1),
            )
            losses.append(loss.item())
        if was_training:
            self.model.train()
        return float(sum(losses) / len(losses))

    @torch.no_grad()
    def evaluate_chat(self, sampler: ChatBatchSampler) -> float:
        was_training = self.model.training
        self.model.eval()
        losses = []
        for _ in range(max(1, self.train_config.validate_batches)):
            token_seq, loss_mask = sampler.sample(
                self.train_config.batch_size,
                split="val",
                device=self.device,
            )
            logits, _ = self.model(token_seq)
            losses.append(self.next_token_loss(logits, token_seq, loss_mask).item())
        if was_training:
            self.model.train()
        return float(sum(losses) / len(losses))

    def fit_text(
        self,
        text: str,
        max_steps: int,
        checkpoint_path: str | Path | None = None,
        on_progress: Callable[[dict], None] | None = None,
    ) -> list[dict]:
        return self.fit_tokens(
            self.tokenizer.encode(text),
            max_steps=max_steps,
            checkpoint_path=checkpoint_path,
            on_progress=on_progress,
        )

    def fit_tokens(
        self,
        token_ids: list[int],
        max_steps: int,
        checkpoint_path: str | Path | None = None,
        on_progress: Callable[[dict], None] | None = None,
    ) -> list[dict]:
        sampler = TokenBatchSampler(
            token_ids,
            seq_len=self.train_config.seq_len,
            train_fraction=self.train_config.train_fraction,
        )
        use_rk4 = not self.train_config.use_euler_during_train
        with expert_integration_mode(self.model, use_rk4=use_rk4):
            for _ in range(max_steps):
                batch = sampler.sample(
                    self.train_config.batch_size,
                    split="train",
                    device=self.device,
                )
                stats = self.learn_batch(batch)
                self.global_step += 1
                stats["step"] = self.global_step

                if (
                    self.train_config.validate_every > 0
                    and self.global_step % self.train_config.validate_every == 0
                ):
                    stats["val_loss"] = self.evaluate(sampler)

                if (
                    self.train_config.consolidate_every > 0
                    and self.global_step % self.train_config.consolidate_every == 0
                ):
                    self.optimizer.consolidate(
                        batch.detach(),
                        n_samples=self.train_config.consolidate_samples,
                    )

                self.history.append(stats)
                if on_progress:
                    on_progress(stats)

                if (
                    checkpoint_path is not None
                    and self.train_config.save_every > 0
                    and self.global_step % self.train_config.save_every == 0
                ):
                    self.save_checkpoint(checkpoint_path)

        if checkpoint_path is not None:
            self.save_checkpoint(checkpoint_path)
        return self.history

    def fit_chat_examples(
        self,
        examples: list[ChildDialogExample],
        max_steps: int,
        checkpoint_path: str | Path | None = None,
        on_progress: Callable[[dict], None] | None = None,
    ) -> list[dict]:
        sampler = ChatBatchSampler(
            examples,
            self.tokenizer,
            seq_len=self.train_config.seq_len,
            train_fraction=self.train_config.train_fraction,
        )
        use_rk4 = not self.train_config.use_euler_during_train
        with expert_integration_mode(self.model, use_rk4=use_rk4):
            for _ in range(max_steps):
                batch, mask = sampler.sample(
                    self.train_config.batch_size,
                    split="train",
                    device=self.device,
                )
                stats = self.learn_masked_batch(batch, mask)
                self.global_step += 1
                stats["step"] = self.global_step

                if (
                    self.train_config.validate_every > 0
                    and self.global_step % self.train_config.validate_every == 0
                ):
                    stats["val_loss"] = self.evaluate_chat(sampler)

                if (
                    self.train_config.consolidate_every > 0
                    and self.global_step % self.train_config.consolidate_every == 0
                ):
                    self.optimizer.consolidate(
                        batch.detach(),
                        n_samples=self.train_config.consolidate_samples,
                    )

                self.history.append(stats)
                if on_progress:
                    on_progress(stats)

                if (
                    checkpoint_path is not None
                    and self.train_config.save_every > 0
                    and self.global_step % self.train_config.save_every == 0
                ):
                    self.save_checkpoint(checkpoint_path, metadata={"curriculum": "child"})

        if checkpoint_path is not None:
            self.save_checkpoint(checkpoint_path, metadata={"curriculum": "child"})
        return self.history

    def state_dict(self) -> dict:
        return {
            "global_step": self.global_step,
            "optimizer": self.optimizer.state_dict(),
            "replay_buffer": list(self.replay_buffer.buffer),
            "plasticity": {
                "usage": self.plasticity.usage,
                "expert_specialization": self.plasticity.expert_specialization,
            },
            "train_config": asdict(self.train_config),
            "history": self.history[-200:],
        }

    def load_state_dict(self, state: dict):
        self.global_step = int(state.get("global_step", 0))
        if "optimizer" in state:
            self.optimizer.load_state_dict(state["optimizer"])
        replay = state.get("replay_buffer", [])
        self.replay_buffer.buffer.clear()
        for seq in replay:
            self.replay_buffer.push(seq)
        plasticity = state.get("plasticity", {})
        if "usage" in plasticity:
            self.plasticity.usage = plasticity["usage"].detach().cpu()
        if "expert_specialization" in plasticity:
            self.plasticity.expert_specialization = plasticity["expert_specialization"]
        self.history = list(state.get("history", []))

    def save_checkpoint(self, path: str | Path, metadata: dict | None = None):
        save_language_checkpoint(
            path,
            self.model,
            self.model_config,
            self.tokenizer,
            trainer_state=self.state_dict(),
            metadata=metadata,
        )

    def load_checkpoint_state(self, path: str | Path):
        ckpt = torch.load(path, map_location=self.device, weights_only=False)
        trainer_state = ckpt.get("trainer")
        if trainer_state:
            self.load_state_dict(trainer_state)