"""Training orchestration logic for Flash-ANSR models.

This module provides the `Trainer` class, which encapsulates the full
training/validation loop as well as utilities for configuration driven
initialisation and experiment logging.
"""

import copy
import os
from collections import Counter
from typing import Any, Literal, Sequence

import torch

from torch.optim.lr_scheduler import LRScheduler
from torch import nn

from tqdm import tqdm

from flash_ansr.model import FlashANSRModel
from flash_ansr.data import FlashANSRDataset
from flash_ansr.utils.config_io import load_config, save_config, unfold_config
from flash_ansr.utils.paths import substitute_root_path
from flash_ansr.eval.metrics.token_prediction import correct_token_predictions_at_k, reciprocal_rank
from flash_ansr.train.optimizers import get_optimizer
from flash_ansr.train.schedules import pw_linear_schedule

# ---------------------------------------------------------------------------
# Compatibility patches
# ---------------------------------------------------------------------------

# Pydantic 2.x emits UnsupportedFieldAttributeWarning when libraries such as
# wandb still pass ``repr``/``frozen`` to ``Field``. Strip these keywords before
# wandb imports its autogenerated models to keep logs clean without suppressing
# warnings globally.
try:  # pragma: no cover - defensive import guard
    import pydantic
    from pydantic import Field as _ORIGINAL_PYDANTIC_FIELD
    from pydantic import fields as _pydantic_fields
except Exception:  # noqa: BLE001 - best-effort shim; fall back silently
    _ORIGINAL_PYDANTIC_FIELD = None  # type: ignore[assignment]
else:
    assert _ORIGINAL_PYDANTIC_FIELD is not None

    def _field_without_unused_attrs(*args: Any, **kwargs: Any) -> Any:
        kwargs.pop('repr', None)
        kwargs.pop('frozen', None)
        return _ORIGINAL_PYDANTIC_FIELD(*args, **kwargs)

    pydantic.Field = _field_without_unused_attrs  # type: ignore[attr-defined]

    _orig_fieldinfo_init = _pydantic_fields.FieldInfo.__init__

    def _fieldinfo_init_without_unused_attrs(self: Any, *args: Any, **kwargs: Any) -> None:
        kwargs.pop('repr', None)
        kwargs.pop('frozen', None)
        _orig_fieldinfo_init(self, *args, **kwargs)

    _pydantic_fields.FieldInfo.__init__ = _fieldinfo_init_without_unused_attrs  # type: ignore[method-assign]

import wandb


def _enable_tf32_precision() -> None:
    if not torch.cuda.is_available():
        return

    try:
        device_capabilities = [torch.cuda.get_device_capability(i) for i in range(torch.cuda.device_count())]
    except RuntimeError:
        # CUDA context is not ready; defer enabling.
        return

    if not any(major >= 8 for major, _ in device_capabilities):
        # TF32 support starts with Ampere (compute capability 8.0).
        return

    matmul_backend = getattr(torch.backends.cuda, 'matmul', None)
    conv_backend = getattr(torch.backends.cudnn, 'conv', None)

    try:
        if matmul_backend is not None and hasattr(matmul_backend, 'fp32_precision'):
            matmul_backend.fp32_precision = 'tf32'
        if conv_backend is not None and hasattr(conv_backend, 'fp32_precision'):
            conv_backend.fp32_precision = 'tf32'
    except AttributeError:
        # Older Torch builds without the new API still honour the legacy helper.
        torch.set_float32_matmul_precision('high')


_enable_tf32_precision()


class Trainer:
    """Manage end-to-end training for a ``FlashANSRModel``.

    Notes
    -----
    Parameters mirror the components required to run training while being
    flexible enough to support configuration driven instantiation via
    :meth:`Trainer.from_config`.
    """

    def __init__(
            self,
            model: FlashANSRModel,
            optimizer: torch.optim.Optimizer,
            amp_dtype: torch.dtype,
            scaler: torch.amp.GradScaler,
            lr_scheduler: LRScheduler | None,
            batch_size: int,
            train_dataset: FlashANSRDataset,
            val_dataset: FlashANSRDataset,
            gradient_accumulation_steps: int = 1,
            num_workers: int | None = None,
            config: dict[str, Any] = None) -> None:
        """Create a fully configured trainer instance.

        Parameters
        ----------
        model : FlashANSRModel
            The model to optimise.
        optimizer : torch.optim.Optimizer
            Optimiser responsible for updating model parameters.
        amp_dtype : torch.dtype
            Mixed precision dtype used with autocast.
        scaler : torch.amp.GradScaler
            Gradient scaler used for automatic mixed precision.
        lr_scheduler : LRScheduler or None
            Optional learning rate scheduler.
        batch_size : int
            Number of samples processed per training step.
        train_dataset : FlashANSRDataset
            Dataset providing training batches.
        val_dataset : FlashANSRDataset
            Dataset providing validation batches.
        gradient_accumulation_steps : int, default=1
            Number of micro-batches per optimizer step.
        num_workers : int or None, optional
            Number of worker processes for data generation.
        config : dict[str, Any] or None, optional
            Trainer configuration metadata (used for logging).
        """

        self.model = model
        self.optimizer = optimizer
        self.amp_dtype = amp_dtype
        self.scaler = scaler
        self.lr_scheduler = lr_scheduler
        self.batch_size = batch_size
        self.train_dataset = train_dataset
        self.val_dataset = val_dataset
        self.gradient_accumulation_steps = gradient_accumulation_steps
        self.config = config or {}
        self.num_workers = num_workers
        self.device = torch.device("cpu")  # Updated during ``run``
        default_worker_preprocess = self.train_dataset.preprocessor is not None
        if isinstance(self.config, dict) and 'worker_preprocess' in self.config:
            self.worker_preprocess = bool(self.config.get('worker_preprocess'))
        else:
            self.worker_preprocess = default_worker_preprocess

        # Metrics and Loss Functions
        self.metrics_ignore_index = self.model.tokenizer["<pad>"]
        self.cross_entropy_loss = nn.CrossEntropyLoss(ignore_index=self.metrics_ignore_index)

        self.total_pflops = 0.0
        self.encoder_parameters = sum(p.numel() for p in self.model.encoder.parameters() if p.requires_grad)
        self.decoder_parameters = sum(p.numel() for p in self.model.decoder.parameters() if p.requires_grad)
        tokenizer_mapping = getattr(self.model.tokenizer, "token2idx", {})
        if not isinstance(tokenizer_mapping, dict):
            tokenizer_mapping = {}
        self._prompt_token_ids: dict[str, int | None] = {
            "complexity": tokenizer_mapping.get("<complexity>"),
        }
        self.prompt_combo_counts: Counter[tuple] = Counter()
        self.prompt_total_samples = 0
        self.decoder_max_seq_len = self._resolve_decoder_max_seq_len()

    @classmethod
    def from_config(cls, config: dict[str, Any] | str) -> "Trainer":
        """Instantiate a trainer from a config dictionary or YAML file path.

        Parameters
        ----------
        config : dict[str, Any] or str
            Dictionary of trainer settings or path to a YAML configuration file.

        Returns
        -------
        Trainer
            Trainer populated using the provided configuration.
        """
        config_ = load_config(config)

        if "trainer" in config_.keys():
            config_ = config_["trainer"]

        # Handle relative model paths
        if (
            isinstance(config, str)
            and isinstance(config_["model"], str)
            and config_["model"].startswith('.')
            and not os.path.isabs(config_["model"])
            and not os.path.normpath(config_["model"]).startswith(os.path.normpath(os.path.dirname(config)))
        ):
            # Only rewrite relative model paths that have not already been expanded
            config_["model"] = os.path.join(os.path.dirname(config), config_["model"])

        print(f'Creating model from {config_["model"]}')
        model = FlashANSRModel.from_config(config_["model"])

        print(f"Loading optimizer with config {config_['optimizer']}")
        optimizer = get_optimizer(
            config_['optimizer']['name'],
            params=model.parameters(),
            **config_['optimizer'].get('kwargs', {}),
        )

        # On CPU we avoid float16 autocast to prevent overflow; prefer bf16 when available.
        cpu_bf16_supported = bool(getattr(torch.cpu, "is_bf16_supported", lambda: False)())
        if torch.cuda.is_available():
            amp_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
        else:
            amp_dtype = torch.bfloat16 if cpu_bf16_supported else torch.float32

        scaler = torch.amp.GradScaler(enabled=amp_dtype == torch.float16)
        print(f'Using amp_dtype={amp_dtype}, GradScaler enabled={scaler.is_enabled()}')

        print(f'Loading lr_scheduler with config {config_["lr_schedule"]}')
        lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
            optimizer=optimizer,
            lr_lambda=lambda step: pw_linear_schedule(step, config_['lr_schedule'])  # Pass the step and the points that define the piecewise linear schedule
        )

        num_workers = config_.get("num_workers", None)
        print(f'Using num_workers={num_workers} for data generation')
        print(f'Loading train_dataset with config {config_["train_dataset"]}')
        train_dataset = FlashANSRDataset.from_config(config_["train_dataset"])

        print(f'Loading val_dataset with config {config_["val_dataset"]}')
        if isinstance(config_["val_dataset"], str):
            resolved_config_path = substitute_root_path(config_["val_dataset"])
            if os.path.isfile(resolved_config_path):
                val_dataset = FlashANSRDataset.from_config(resolved_config_path)
            elif os.path.isdir(resolved_config_path):
                _, val_dataset = FlashANSRDataset.load(directory=resolved_config_path)
            else:
                raise ValueError(f"val_dataset must be a file or directory, not {resolved_config_path}")
        elif isinstance(config_["val_dataset"], dict):
            val_dataset = FlashANSRDataset.from_config(config_["val_dataset"])
        else:
            raise ValueError(f"val_dataset must be a dict or path to a file or directory, not {config_['val_dataset']}")

        config_.setdefault('preprocess', False)

        return cls(
            model=model,
            optimizer=optimizer,
            amp_dtype=amp_dtype,
            scaler=scaler,
            lr_scheduler=lr_scheduler,
            batch_size=config_['batch_size'],
            train_dataset=train_dataset,
            val_dataset=val_dataset,
            gradient_accumulation_steps=config_.get("gradient_accumulation_steps", 1),
            num_workers=num_workers,
            config=config_,
        )

    def _setup_training_state(self, device: str, verbose: bool) -> None:
        """Move the model to ``device`` and reset training state.

        Parameters
        ----------
        device : str
            Torch device string used for training (for example ``"cuda"``).
        verbose : bool
            If ``True`` emit human readable status messages.
        """
        self.device = torch.device(device)
        self.model.to(self.device)

        self.total_pflops = 0.0

        if verbose:
            config_value = os.environ.get('PYTORCH_CUDA_ALLOC_CONF')
            if config_value:
                print(f"PYTORCH_CUDA_ALLOC_CONF is set to: '{config_value}'")
            else:
                print("PYTORCH_CUDA_ALLOC_CONF is not set.")

    def _resolve_decoder_max_seq_len(self) -> int:
        max_seq_len = getattr(self.model, 'decoder_max_seq_len', None)
        if max_seq_len is None:
            decoder = getattr(self.model, 'decoder', None)
            rope = getattr(decoder, 'rope', None) if decoder is not None else None
            max_seq_len = getattr(rope, 'max_seq_len', None) if rope is not None else None

        if max_seq_len is None:
            raise ValueError("Unable to determine decoder max sequence length from model.")

        if hasattr(max_seq_len, 'item'):
            max_seq_len = max_seq_len.item()

        return int(max_seq_len)

    def run_training(
            self,
            steps: int,
            preprocess: bool = False,
            device: str = "cpu",
            compile_mode: str | None = None,
            checkpoint_interval: int | None = None,
            checkpoint_directory: str | None = None,
            validate_interval: int | None = None,
            validate_size: int | None = None,
            validate_batch_size: int = 128,
            preprocess_in_worker: bool | None = None,
            resume_from: str | None = None,
            resume_step: int | None = None,
            verbose: bool = False) -> FlashANSRModel:
        """Execute the core training loop.

        Parameters
        ----------
        steps : int
            Number of optimisation steps to perform.
        preprocess : bool, default=False
            Whether to run dataset preprocessing before each batch.
        device : str, default='cpu'
            Target device identifier passed to :func:`torch.device`.
        compile_mode : str or None, optional
            Torch.compile mode for model compilation.
        checkpoint_interval : int or None, optional
            Save checkpoints every ``n`` steps when provided.
        checkpoint_directory : str or None, optional
            Directory for persisted checkpoints.
        validate_interval : int or None, optional
            Frequency (in steps) to run validation.
        validate_size : int or None, optional
            Limit the number of validation examples processed.
        validate_batch_size : int, default=128
            Batch size used for validation passes.
        preprocess_in_worker : bool or None, optional
            When ``True`` run dataset preprocessing inside producer workers
            instead of the main process. Defaults to the trainer configuration
            when ``None`` (which enables worker preprocessing whenever the
            dataset has a preprocessor).
        resume_from : str or None, optional
            Path to a checkpoint directory from which to resume.
        resume_step : int or None, optional
            Explicitly set the global step to resume from. Overrides any step
            inferred from the checkpoint.
        verbose : bool, default=False
            If ``True`` display progress bars and additional logs.

        Returns
        -------
        FlashANSRModel
            The trained model instance.
        """

        worker_preprocess = self.worker_preprocess if preprocess_in_worker is None else bool(preprocess_in_worker)
        if not preprocess:
            worker_preprocess = False

        try:
            self._setup_training_state(device, verbose=verbose)

            step_offset = 0
            if resume_from is not None:
                step_offset = self._load_checkpoint(resume_from, device=device)

            if compile_mode is not None:
                self.model = torch.compile(self.model, mode=compile_mode)

            if resume_step is not None:
                step_offset = int(resume_step)

            if step_offset < 0:
                raise ValueError("resume_step must be non-negative")

            remaining_steps = max(0, steps - step_offset)
            if verbose:
                print(f"Resuming from step {step_offset} with {remaining_steps} steps remaining (target={steps})")

            if remaining_steps == 0:
                return self.model

            pbar = tqdm(range(remaining_steps), disable=not verbose, smoothing=0, desc="Training")
            for local_step, batch in enumerate(
                    self.train_dataset.iterate(
                        steps=remaining_steps,
                        batch_size=self.batch_size,
                        preprocess=preprocess,
                        preprocess_in_worker=worker_preprocess,
                        num_workers=self.num_workers,
                        max_seq_len=self.decoder_max_seq_len)):
                global_step = step_offset + local_step
                self._train_step(batch, global_step, preprocess, do_optimizer_step=True)

                pbar.update(1)

                human_step = global_step + 1
                if validate_interval is not None and ((human_step % validate_interval == 0) or (human_step == steps)):
                    self._validate_step(human_step, validate_size, validate_batch_size, preprocess, worker_preprocess, verbose)

                if checkpoint_interval is not None and checkpoint_directory is not None and (human_step % checkpoint_interval == 0):
                    self._save_checkpoint(human_step, checkpoint_directory)

            pbar.close()
            return self.model

        except Exception:
            self.train_dataset.shutdown()
            self.val_dataset.shutdown()
            raise

    def run(
            self,
            project_name: str,
            entity: str,
            name: str,
            steps: int,
            preprocess: bool = False,
            device: str = "cpu",
            compile_mode: str | None = None,
            checkpoint_interval: int | None = None,
            checkpoint_directory: str | None = None,
            validate_interval: int | None = None,
            validate_size: int | None = None,
            validate_batch_size: int = 128,
            wandb_mode: Literal['online', 'offline', 'disabled'] = 'online',
            wandb_watch_log: Literal['gradients', 'parameters', 'all'] | None = None,
            wandb_watch_log_freq: int = 1000,
            preprocess_in_worker: bool | None = None,
            num_workers: int | None = None,
            resume_from: str | None = None,
            resume_step: int | None = None,
            verbose: bool = False) -> FlashANSRModel:
        """Train the model while managing the experiment lifecycle via W&B.

        Parameters
        ----------
        project_name : str
            Name of the Weights & Biases project.
        entity : str
            W&B entity (team or user) under which to log the run.
        name : str
            Run name displayed in W&B.
        steps : int
            Number of training iterations to execute.
        preprocess : bool, default=False
            Whether to preprocess dataset batches.
        device : str, default='cpu'
            Torch device string (for example ``"cpu"`` or ``"cuda"``).
        compile_mode : str or None, optional
            Torch.compile mode, if supported.
        checkpoint_interval : int or None, optional
            Step interval for checkpointing.
        checkpoint_directory : str or None, optional
            Directory where checkpoints are written when enabled.
        validate_interval : int or None, optional
            Step cadence for validation.
        validate_size : int or None, optional
            Maximum number of validation samples, if limited.
        validate_batch_size : int, default=128
            Batch size to use during validation.
        wandb_mode : {'online', 'offline', 'disabled'}, default='online'
            W&B initialisation mode.
        wandb_watch_log : {'gradients', 'parameters', 'all'} or None, optional
            W&B ``watch`` setting controlling what to log.
        wandb_watch_log_freq : int, default=1000
            Frequency (steps) for W&B gradient/parameter logging.
        preprocess_in_worker : bool or None, optional
            When ``True`` run preprocessing inside dataset workers. Defaults to
            the trainer configuration when omitted (which enables worker
            preprocessing when available).
        num_workers : int or None, optional
            Number of worker processes for data generation.
        resume_from : str or None, optional
            Path to a checkpoint directory from which to resume training.
        resume_step : int or None, optional
            Override the inferred global step when resuming from ``resume_from``.
        verbose : bool, default=False
            When ``True`` print progress information to stdout.

        Returns
        -------
        FlashANSRModel
            The trained model instance.
        """

        if verbose:
            print(f"Training model ({self.model.n_params:,} parameters) for {steps:,} steps on device {device}")

        if num_workers is not None:
            print(f'Overriding trainer num_workers to {num_workers}')
            self.num_workers = num_workers

        wandb_config = unfold_config(copy.deepcopy(self.config))
        wandb_config.update({"steps": steps, "device": device, "verbose": verbose, "num_workers": num_workers})

        with wandb.init(config=wandb_config, project=project_name, entity=entity, name=name, mode=wandb_mode):  # type: ignore
            if wandb_mode != 'disabled':
                wandb.watch(self.model, log=wandb_watch_log, log_freq=wandb_watch_log_freq)  # type: ignore
                if verbose and wandb_watch_log is not None:
                    print(f'Watching model with wandb log={wandb_watch_log} at frequency {wandb_watch_log_freq}')

            return self.run_training(
                steps=steps,
                preprocess=preprocess,
                device=device,
                compile_mode=compile_mode,
                checkpoint_interval=checkpoint_interval,
                checkpoint_directory=checkpoint_directory,
                validate_interval=validate_interval,
                validate_size=validate_size,
                validate_batch_size=validate_batch_size,
                preprocess_in_worker=preprocess_in_worker,
                resume_from=resume_from,
                resume_step=resume_step,
                verbose=verbose
            )

    def _update_total_pflops(self, encoder_tokens: int, decoder_tokens: int, batch_size: int) -> None:
        """Accumulate an estimate of total pico floating-point operations.

        Parameters
        ----------
        encoder_tokens : int
            Number of encoder tokens processed in the current step.
        decoder_tokens : int
            Number of decoder tokens processed in the current step.
        batch_size : int
            Size of the current batch (prior to micro batching).
        """
        self.total_pflops += 6 * (self.encoder_parameters * encoder_tokens + self.decoder_parameters * decoder_tokens) * batch_size * 1e-15

    def _apply_prompt_mask(self, batch: dict[str, torch.Tensor]) -> None:
        """Mask loss contributions for prompt tokens when available."""
        prompt_mask = batch.get('prompt_mask')
        if prompt_mask is None:
            return

        if not isinstance(prompt_mask, torch.Tensor):
            prompt_mask = torch.as_tensor(prompt_mask, device=self.device, dtype=torch.bool)
            batch['prompt_mask'] = prompt_mask
        else:
            if prompt_mask.dtype != torch.bool:
                prompt_mask = prompt_mask.to(dtype=torch.bool)
            batch['prompt_mask'] = prompt_mask.to(device=self.device)

        labels = batch.get('labels')
        if labels is None or not isinstance(labels, torch.Tensor):
            return

        if labels.device != self.device:
            labels = labels.to(self.device)
            batch['labels'] = labels

        target_mask = batch['prompt_mask'][..., 1:]
        if target_mask.shape[-1] > labels.shape[-1]:
            target_mask = target_mask[..., :labels.shape[-1]]
        elif target_mask.shape[-1] < labels.shape[-1]:
            labels = labels[..., :target_mask.shape[-1]]
            batch['labels'] = labels

        labels[target_mask] = self.metrics_ignore_index

    @staticmethod
    def _canonicalize_prompt_terms(terms: Any) -> tuple[tuple[str, ...], ...]:
        """Return a hashable representation of prompt terms."""
        if not isinstance(terms, Sequence) or isinstance(terms, (str, bytes)):
            if terms is None or terms == []:
                return tuple()
            return ((str(terms),),)

        canonical_terms: list[tuple[str, ...]] = []
        for term in terms:
            if isinstance(term, Sequence) and not isinstance(term, (str, bytes)):
                canonical_terms.append(tuple(str(t) for t in term))
            elif term is not None:
                canonical_terms.append((str(term),))
        return tuple(sorted(canonical_terms))

    def _update_prompt_statistics(self, batch: dict[str, Any]) -> None:
        """Track prompt metadata usage statistics for diagnostics."""
        prompt_metadata = batch.get('prompt_metadata')
        if not prompt_metadata or not isinstance(prompt_metadata, Sequence):
            return

        for metadata in prompt_metadata:
            if not isinstance(metadata, dict):
                continue

            allowed = self._canonicalize_prompt_terms(metadata.get('allowed_terms'))
            include = self._canonicalize_prompt_terms(metadata.get('include_terms'))
            exclude = self._canonicalize_prompt_terms(metadata.get('exclude_terms'))
            key = (allowed, include, exclude)

            self.prompt_combo_counts[key] += 1
            self.prompt_total_samples += 1

    def _train_step(self, batch: dict[str, torch.Tensor], step: int, preprocess: bool, do_optimizer_step: bool = True) -> None:
        """Perform a single optimisation step with optional gradient accumulation.

        Parameters
        ----------
        batch : dict[str, torch.Tensor]
            Batched tensors produced by the training dataset.
        step : int
            Zero-indexed global training step.
        preprocess : bool
            Whether to preprocess batch data prior to collation.
        do_optimizer_step : bool, default=True
            When ``False`` skip the optimiser update (useful for gradient
            accumulation across micro-batches).
        """
        self.model.train()
        self.optimizer.zero_grad()

        total_ce_loss = 0.0
        total_loss = 0.0

        # Split the batch into micro-batches to support gradient accumulation
        for acc_step in range(self.gradient_accumulation_steps):
            micro_batch_size = len(batch['x_tensors']) // self.gradient_accumulation_steps
            micro_batch = {k: v[acc_step * micro_batch_size:(acc_step + 1) * micro_batch_size] for k, v in batch.items()}
            micro_batch = self.train_dataset.collate(micro_batch, device=self.device)
            if preprocess:
                self._update_prompt_statistics(micro_batch)
            self._apply_prompt_mask(micro_batch)

            data_tensor = torch.cat([micro_batch['x_tensors'], micro_batch['y_tensors']], dim=-1)

            if self.amp_dtype == torch.float32:
                autocast_context = torch.autocast(self.device.type, dtype=self.amp_dtype, enabled=False)
            else:
                autocast_context = torch.autocast(self.device.type, dtype=self.amp_dtype)

            with autocast_context:
                logits = self.model(micro_batch['input_ids'], data_tensor, input_num=micro_batch.get('input_num', None), data_attn_mask=micro_batch['data_attn_mask'].to(self.device))
                flat_logits = logits[:, :-1].reshape(-1, logits.shape[-1])
                flat_labels = micro_batch['labels'].reshape(-1)
                valid_labels = flat_labels != self.metrics_ignore_index
                if not valid_labels.any():
                    ce_loss = torch.zeros((), device=self.device, dtype=logits.dtype)
                else:
                    ce_loss = self.cross_entropy_loss(flat_logits, flat_labels)

                # Force every parameter to contribute to the loss so that gradient
                # tracking tools (e.g. wandb.watch) do not encounter ``None`` gradients
                # for frozen or unused tensors.
                param_sum = sum(p.sum() for p in self.model.parameters())
                zero_loss = 0.0 * param_sum

                loss = ce_loss / self.gradient_accumulation_steps + zero_loss

            # If the loss is nan or inf, stop the training
            if not torch.isfinite(loss):
                raise ValueError(f"Loss is {loss.item()}, stopping training")

            self.scaler.scale(loss).backward()
            total_ce_loss += ce_loss.item()
            total_loss += loss.item() * self.gradient_accumulation_steps

        # Perform the optimizer step
        self.scaler.unscale_(self.optimizer)

        self._update_total_pflops(encoder_tokens=data_tensor.shape[1], decoder_tokens=micro_batch['input_ids'].shape[1], batch_size=len(batch['x_tensors']))

        total_gradient_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), 2.0)
        if do_optimizer_step:
            self.scaler.step(self.optimizer)
            self.scaler.update()

            # Log metrics and update scheduler after the optimizer step
            self._log_metrics(step, flat_logits, flat_labels, total_ce_loss, total_loss, total_gradient_norm)
            if self.lr_scheduler is not None:
                self.lr_scheduler.step()

    def _validate_step(self, step: int, size: int | None, batch_size: int, preprocess: bool, worker_preprocess: bool, verbose: bool) -> None:
        """Evaluate the model on the validation split and log aggregate metrics.

        Parameters
        ----------
        step : int
            Global training step at which validation is triggered.
        size : int or None
            Optional limit on the number of validation examples.
        batch_size : int
            Number of samples per validation batch.
        preprocess : bool
            Whether to preprocess validation batches.
        worker_preprocess : bool
            Whether preprocessing should run inside dataset workers.
        verbose : bool
            If ``True`` display a validation progress bar.
        """
        self.model.eval()

        val_ce_loss = 0.0
        val_mrr = 0.0
        val_acc_at_1 = 0.0
        total_items = 0
        total_batches = 0

        with torch.no_grad():
            if size is None:
                steps = len(self.val_dataset) // batch_size
            else:
                steps = size // batch_size

            pbar = tqdm(total=steps, leave=False, position=1, disable=not verbose, desc="Validating", smoothing=0.0)
            for batch in self.val_dataset.iterate(
                    size=size,
                    batch_size=batch_size,
                    preprocess=preprocess,
                    preprocess_in_worker=worker_preprocess,
                    num_workers=self.num_workers,
                    max_seq_len=self.decoder_max_seq_len):
                batch = self.val_dataset.collate(batch, device=self.device)
                self._apply_prompt_mask(batch)
                data_tensor = torch.cat([batch['x_tensors'], batch['y_tensors']], dim=-1)

                with torch.autocast(device_type=self.device.type, dtype=self.amp_dtype):
                    logits = self.model(batch['input_ids'], data_tensor, input_num=batch.get('input_num', None), data_attn_mask=batch['data_attn_mask'].to(self.device))
                    flat_logits = logits[:, :-1].reshape(-1, logits.shape[-1])
                    flat_labels = batch['labels'].reshape(-1)
                    valid_labels = flat_labels != self.metrics_ignore_index
                    if not valid_labels.any():
                        ce_loss = torch.zeros((), device=self.device, dtype=logits.dtype)
                    else:
                        ce_loss = self.cross_entropy_loss(flat_logits, flat_labels)

                    # Accumulate metrics for each batch
                    val_ce_loss += ce_loss.item() * flat_labels.shape[0]
                    total_items += flat_labels.shape[0]

                # Filter out ignored indices for metric calculation
                valid_indices = flat_labels != self.metrics_ignore_index
                if valid_indices.any():
                    val_mrr += reciprocal_rank(flat_logits[valid_indices], flat_labels[valid_indices])
                    val_acc_at_1 += correct_token_predictions_at_k(flat_logits[valid_indices], flat_labels[valid_indices], k=1)
                    total_batches += 1

                pbar.update(1)
            pbar.close()

        # Calculate average metrics
        avg_val_ce_loss = val_ce_loss / total_items if total_items > 0 else 0.0
        avg_val_mrr = val_mrr / total_batches if total_batches > 0 else 0.0
        avg_val_acc_at_1 = val_acc_at_1 / total_batches if total_batches > 0 else 0.0

        # Log averaged validation metrics
        self._log_validation_metrics(step, avg_val_ce_loss, avg_val_mrr, avg_val_acc_at_1)

    def _save_checkpoint(self, step: int, checkpoint_directory: str) -> None:
        """Persist model weights, optimiser state, and config for ``step``.

        Parameters
        ----------
        step : int
            Training step associated with the persisted checkpoint.
        checkpoint_directory : str
            Directory into which the checkpoint will be written.
        """
        save_directory = os.path.join(checkpoint_directory, f"checkpoint_{step}")
        self.model.save(directory=save_directory, errors='ignore')
        torch.save(self.optimizer.state_dict(), os.path.join(save_directory, "optimizer.pt"))
        if self.lr_scheduler is not None:
            torch.save(self.lr_scheduler.state_dict(), os.path.join(save_directory, "lr_scheduler.pt"))
        if self.scaler is not None:
            torch.save(self.scaler.state_dict(), os.path.join(save_directory, "scaler.pt"))
        torch.save({"step": step, "total_pflops": self.total_pflops}, os.path.join(save_directory, "training_state.pt"))
        save_config(
            load_config(self.config, resolve_paths=True),
            directory=save_directory,
            filename='train.yaml',
            reference='relative',
            recursive=True,
            resolve_paths=True)
        print(f"Checkpoint saved at {save_directory}")

    def _infer_resume_step(self, checkpoint_directory: str) -> int | None:
        """Infer the resume step from the checkpoint directory name."""
        base = os.path.basename(os.path.normpath(checkpoint_directory))
        prefix = "checkpoint_"
        if base.startswith(prefix):
            suffix = base[len(prefix):]
            if suffix.isdigit():
                return int(suffix)
        return None

    def _load_checkpoint(self, checkpoint_directory: str, device: str) -> int:
        """Load training state from ``checkpoint_directory`` and return the step."""

        checkpoint_directory = substitute_root_path(checkpoint_directory)
        if not os.path.isdir(checkpoint_directory):
            raise FileNotFoundError(f"Checkpoint directory not found: {checkpoint_directory}")

        model_state_path = os.path.join(checkpoint_directory, "state_dict.pt")
        optimizer_state_path = os.path.join(checkpoint_directory, "optimizer.pt")
        scheduler_state_path = os.path.join(checkpoint_directory, "lr_scheduler.pt")
        scaler_state_path = os.path.join(checkpoint_directory, "scaler.pt")
        training_state_path = os.path.join(checkpoint_directory, "training_state.pt")

        map_location = torch.device(device)

        self.model.load_state_dict(torch.load(model_state_path, map_location=map_location))
        self.optimizer.load_state_dict(torch.load(optimizer_state_path, map_location=map_location))

        resume_step: int | None = None
        if os.path.isfile(training_state_path):
            training_state = torch.load(training_state_path, map_location=map_location)
            if isinstance(training_state, dict):
                step_value = training_state.get("step")
                if step_value is not None:
                    try:
                        resume_step = int(step_value)
                    except (TypeError, ValueError):
                        resume_step = None

                self.total_pflops = float(training_state.get("total_pflops", self.total_pflops))

        if resume_step is None:
            resume_step = self._infer_resume_step(checkpoint_directory)

        if self.lr_scheduler is not None:
            if os.path.isfile(scheduler_state_path):
                self.lr_scheduler.load_state_dict(torch.load(scheduler_state_path, map_location=map_location))
            else:
                if resume_step is None:
                    print("Warning: lr_scheduler state missing and resume step unknown; set resume_step explicitly to align LR schedule.")
                else:
                    self.lr_scheduler.last_epoch = resume_step - 1
                    print(f"Warning: lr_scheduler state missing; reconstructing schedule at step {resume_step}.")

        if self.scaler is not None:
            if os.path.isfile(scaler_state_path):
                self.scaler.load_state_dict(torch.load(scaler_state_path, map_location=map_location))
            else:
                print("Warning: scaler state missing; starting with a fresh GradScaler.")

        if resume_step is None:
            raise ValueError(
                "Unable to infer resume step; provide resume_step explicitly or name the checkpoint directory checkpoint_<step>.")

        return resume_step

    def _log_metrics(self, step: int, logits: torch.Tensor, labels: torch.Tensor, ce_loss: float, total_loss: float, total_gradient_norm: torch.Tensor) -> None:
        """Submit training metrics for the current batch to Weights & Biases.

        Parameters
        ----------
        step : int
            Global training step the metrics correspond to.
        logits : torch.Tensor
            Model logits for the current batch.
        labels : torch.Tensor
            Ground-truth labels aligned with ``logits``.
        ce_loss : float
            Cross-entropy loss for the batch.
        total_loss : float
            Total loss (including auxiliary terms) for the batch.
        total_gradient_norm : torch.Tensor
            Gradient norm measured after clipping.
        """

        log_data = {
            "total_gradient_norm": total_gradient_norm.item(),
            "train_ce_loss": ce_loss,
            "train_loss": total_loss,
            "lr": self.optimizer.param_groups[0]['lr'],
            "train_mean_reciprocal_rank": reciprocal_rank(logits, labels, ignore_index=self.metrics_ignore_index),
            "train_correct_token_predictions_at_1": correct_token_predictions_at_k(logits, labels, k=1, ignore_index=self.metrics_ignore_index),
            "total_pflops": self.total_pflops,
        }

        wandb.log(log_data, step=step)  # type: ignore

    def _log_validation_metrics(self, step: int, val_ce_loss: float, val_mrr: float, val_acc_at_1: float) -> None:
        """Submit aggregated validation metrics to Weights & Biases.

        Parameters
        ----------
        step : int
            Global training step the metrics correspond to.
        val_ce_loss : float
            Mean validation cross-entropy loss.
        val_mrr : float
            Mean reciprocal rank on the validation set.
        val_acc_at_1 : float
            Accuracy at one token prediction on the validation set.
        """
        log_data = {
            "val_ce_loss": val_ce_loss,
            "val_mean_reciprocal_rank": val_mrr,
            "val_correct_token_predictions_at_1": val_acc_at_1,
        }
        wandb.log(log_data, step=step)  # type: ignore
