"""
CODI + PMPS Model

This module extends the original CODI model with a PMPS (Swapping Assignments
between Views) supervision loss that provides direct supervision for latent
embeddings by aligning them, in a shared cluster space, with the teacher
task's CoT-token hidden states. See ``pmps/PMPS_LATENT_SUPERVISION.md`` for
the full derivation; the module is mirrored at ``CODI_pmps/src/pmps_loss.py``.

Key changes from codi_ori/src/model.py:
  1. Added PMPS-related training arguments (pmps_loss_factor, pmps_temperature, etc.)
  2. In forward(): collect latent hidden states at each iteration step
  3. In forward(): extract CoT-position hidden states from teacher output
  4. Compute PMPS loss between latent embeddings and CoT hidden states
  5. Add PMPS loss to the total training objective
"""

import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import random
from dataclasses import dataclass, field
from typing import Optional
from peft import get_peft_model, PeftModel, PeftConfig
from torch.nn.functional import gelu
import math
from safetensors.torch import load_file, safe_open
from transformers.modeling_outputs import ModelOutput
import json
import glob
import copy

from .pmps_loss import PMPSLatentSupervisor

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def load_safetensors_state_dict(ckpt_dir: str) -> dict:
    """Load a state dict from a checkpoint directory, supporting both single-file
    and sharded safetensors formats (e.g. model-00001-of-00003.safetensors).

    Args:
        ckpt_dir: Path to the directory containing model weights.

    Returns:
        A merged state dict containing all model parameters.
    """
    single_file = os.path.join(ckpt_dir, "model.safetensors")
    if os.path.isfile(single_file):
        return load_file(single_file)

    # Check for sharded safetensors via the index json
    index_file = os.path.join(ckpt_dir, "model.safetensors.index.json")
    if os.path.isfile(index_file):
        with open(index_file, "r") as f:
            index = json.load(f)
        weight_map = index.get("metadata", {}).get("weight_map", index.get("weight_map", {}))
        shard_files = sorted(set(weight_map.values()))
        state_dict = {}
        for shard_name in shard_files:
            shard_path = os.path.join(ckpt_dir, shard_name)
            shard_state = load_file(shard_path)
            state_dict.update(shard_state)
        return state_dict

    # Fallback: try to auto-detect sharded files (model-XXXXX-of-YYYYY.safetensors)
    shard_pattern = os.path.join(ckpt_dir, "model-*-of-*.safetensors")
    shard_files = sorted(glob.glob(shard_pattern))
    if shard_files:
        state_dict = {}
        for shard_path in shard_files:
            shard_state = load_file(shard_path)
            state_dict.update(shard_state)
        return state_dict

    # Last fallback: pytorch format
    pytorch_file = os.path.join(ckpt_dir, "pytorch_model.bin")
    if os.path.isfile(pytorch_file):
        return torch.load(pytorch_file, map_location="cpu")

    raise FileNotFoundError(
        f"No model weights found in {ckpt_dir}. "
        "Expected model.safetensors, sharded safetensors, or pytorch_model.bin."
    )


@dataclass
class ModelArguments:
    model_name_or_path: str = field(default="mistralai/Mistral-7B-Instruct-v0.2")
    separate_decoder_name: str = field(default="")
    lora_r: int = field(default=128, metadata={"help": "lora rank"})
    lora_dropout: float = field(default=0.05, metadata={"help": "lora dropout"})
    full_precision: bool = field(default=True, metadata={"help": "whether use int4 for the base model"})
    train: bool = field(
        default=True,
        metadata={"help": "if true, the model ckpt will be initialized for training; else, it's for inference"},
    )
    lora_init: bool = field(
        default=False,
        metadata={"help": "True: Use zero and gaussian initialization; False: Load adapters from LoftQ in HF hub."},
    )
    token: Optional[str] = field(
        default=None,
        metadata={"help": "HF token to access to private models, e.g., meta-llama"},
    )
    adapter_name_or_path: Optional[str] = field(
        default=None,
        metadata={"help": "Path to the LoRA adapter. Used in evaluation or resuming from the checkpoint."},
    )
    lora_alpha: int = field(
        default=16,
        metadata={"help": "LoftQ does not require this config. Used for QLoRA."},
    )
    ckpt_dir: Optional[str] = field(default=None, metadata={"help": "checkpoint dir for inference."})


@dataclass
class DataArguments:
    data_name: str = field(default=None, metadata={"help": "Path to the training data."})
    # data_path: Optional[str] = field(default=None, metadata={"help": "Path to the dataset file (e.g., MATH dataset JSON file)."})
    debug_data: bool = field(
        default=False,
        metadata={"help": "Enable debug dataset to quickly verify the training process"},
    )
    batch_size: int = field(default=1, metadata={"help": "batch size during inference"})


@dataclass
class TrainingArguments(transformers.TrainingArguments):
    cache_dir: Optional[str] = field(default=None)
    optim: str = field(default="adamw_torch")
    model_max_length: int = field(
        default=28000,
        metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."},
    )
    restore_from: str = field(
        default="",
        metadata={"help": "The checkpoint that should be restored from for fine-tuning"},
    )
    per_device_train_batch_size: int = field(default=1)
    per_device_eval_batch_size: int = field(default=1)
    expt_name: str = field(default="default", metadata={"help": "Experiment name"})
    icot_train_path: str = field(
        default="/users/k24020023/efficient_cot/icae/code/coconut/icot_gsm8k/train.txt",
        metadata={"help": "The training data path"},
    )
    num_latent: int = field(default=5, metadata={"help": "The number of latent for training or inference."})
    use_lora: bool = field(default=True, metadata={"help": "Use lora or not."})
    greedy: bool = field(default=False, metadata={"help": "Greedy decoding during inference."})
    exp_mode: bool = field(default=False, metadata={"help": "Use partial number of data. for debugging."})
    exp_data_num: int = field(default=10000, metadata={"help": "The number of data used in exp mode"})
    use_prj: bool = field(default=False, metadata={"help": "Use a prj module after the llm for latent generation."})
    prj_dim: int = field(default=2048, metadata={"help": "The hidden dim of the projection module."})
    prj_dropout: float = field(default=0.0, metadata={"help": "Dropout ratio of the projection module."})
    prj_no_ln: bool = field(default=False, metadata={"help": "Remove the Layer Norm layer for the projection module."})
    distill_loss_div_std: bool = field(
        default=False, metadata={"help": "Divide the distillation loss by a std for normallisation."}
    )
    distill_loss_type: str = field(
        default="smooth_l1", metadata={"help": "Specify the distillation loss. Use smoothL1 by default."}
    )
    distill_loss_factor: float = field(default=1.0, metadata={"help": "A multiplier of the distillation loss."})
    ref_loss_factor: float = field(default=1.0, metadata={"help": "A multiplier of the reference CE loss."})
    inf_latent_iterations: int = field(default=1)
    inf_num_iterations: int = field(default=5, metadata={"help": "Run multiple times during inference"})
    remove_eos: bool = field(default=False, metadata={"help": "Do not add <eos> as a delimiter to split QA."})
    use_latent_reasoning: bool = field(
        default=True,
        metadata={
            "help": "If True (default), perform latent reasoning at inference time. "
                    "If False, skip latent iterations and let the model directly generate an explicit CoT."
        },
    )
    print_ref_model_stats: bool = field(default=False, metadata={"help": "Print some stats for the teacher task."})
    include_last_cot: bool = field(
        default=False, metadata={"help": "Include the last CoT step in the training data."}
    )
    fix_attn_mask: bool = field(default=False, metadata={"help": "Correct a bug about attention mask."})
    log_full: bool = field(default=False, metadata={"help": "Log all losses."})
    print_loss: bool = field(default=True)
    max_token_num: int = field(default=1000, metadata={"help": "Limit the longest data to avoid OOM."})

    # ===== Probabilistic Latent Embedding Arguments (NEW) =====
    enable_top_k: bool = field(
        default=False,
        metadata={"help": "If True, use top-k tokens for the weighted embedding in latent reasoning."},
    )
    latent_top_k: int = field(
        default=10,
        metadata={"help": "Number of top-k tokens to use for the weighted embedding in latent reasoning."},
    )
    latent_top_k_temperature: float = field(
        default=1.0,
        metadata={
            "help": "Temperature for the softmax over top-k logits. "
                    "Lower values make the distribution sharper (closer to argmax); "
                    "higher values make it smoother (closer to uniform)."
        },
    )

    # ===== PMPS Loss Arguments (NEW) =====
    pmps_loss_factor: float = field(
        default=0.1,
        metadata={"help": "Weight of the PMPS loss in the total objective. Set to 0 to disable PMPS loss."},
    )
    pmps_temperature: float = field(
        default=0.1,
        metadata={"help": "Temperature for the prototype softmax (lower = sharper predictions)."},
    )
    pmps_match_temperature: float = field(
        default=0.3,
        metadata={"help": "Temperature for the cluster-space soft matching between latent and CoT tokens."},
    )
    pmps_epsilon: float = field(
        default=0.05,
        metadata={"help": "Sinkhorn-Knopp regularisation parameter."},
    )
    pmps_sinkhorn_iterations: int = field(
        default=3,
        metadata={"help": "Number of Sinkhorn-Knopp iterations."},
    )
    pmps_nmb_prototypes: int = field(
        default=24,
        metadata={"help": "Number of shared prototypes K. Recommended: 2~8 x num_latent."},
    )
    pmps_feat_dim: int = field(
        default=128,
        metadata={"help": "Projection feature dim (only used when pmps_use_projector=True)."},
    )
    pmps_hidden_mlp: int = field(
        default=512,
        metadata={"help": "Projection MLP hidden dim (only used when pmps_use_projector=True)."},
    )
    pmps_use_projector: bool = field(
        default=True,
        metadata={"help": "If True, use a shared MLP projection head before scoring against prototypes."},
    )
    pmps_freeze_prototypes_steps: int = field(
        default=0,
        metadata={"help": "Freeze the prototype layer for the first N optimisation steps."},
    )
    pmps_layer: str = field(
        default="last",
        metadata={
            "help": "Which layer's hidden states to use for PMPS. "
            "'last' = last layer only, 'all' = average over all layers, "
            "or an integer for a specific layer index."
        },
    )
    pmps_start_epoch: float = field(
        default=0.0,
        metadata={"help": "Epoch threshold to start adding the PMPS loss into the total loss."},
    )

    # ===== Progressive Sequential PMPS Arguments (NEW) =====
    pmps_position_bias_max: float = field(
        default=0.0,
        metadata={
            "help": "Maximum positional bias strength for sequential PMPS matching. "
                    "0 = disabled (original free matching). Recommended: 3.0~8.0. "
                    "Higher values enforce stronger sequential correspondence between "
                    "latent steps and CoT token positions."
        },
    )
    pmps_position_bias_schedule: str = field(
        default="cosine",
        metadata={
            "help": "Schedule for decaying position bias strength over training. "
                    "'cosine' = cosine annealing from max to 0 (recommended), "
                    "'linear' = linear decay from max to 0, "
                    "'constant' = no decay (fixed strength throughout training)."
        },
    )

    # ===== Latent Diversity Regularisation (NEW) =====
    diversity_loss_factor: float = field(
        default=0.0,
        metadata={
            "help": "Weight of the latent diversity regularisation loss. "
                    "Penalises high cosine similarity between different latent "
                    "embeddings to prevent functional collapse. "
                    "0 = disabled. Recommended: 0.05~0.3."
        },
    )

    # ===== Wandb Arguments =====
    use_wandb: bool = field(
        default=False,
        metadata={"help": "Enable Weights & Biases logging."},
    )
    wandb_api_key: str = field(
        default="",
        metadata={"help": "Wandb API key for login. Will be read from YAML config."},
    )
    wandb_project: str = field(
        default="codi-pmps",
        metadata={"help": "Wandb project name."},
    )
    wandb_run_name: Optional[str] = field(
        default=None,
        metadata={"help": "Wandb run name. Defaults to expt_name if not set."},
    )

def print_trainable_parameters(model):
    trainable_parameters = 0
    all_param = 0
    for _, param in model.named_parameters():
        all_param += param.numel()
        if param.requires_grad:
            trainable_parameters += param.numel()
    print(
        f"trainable params: {trainable_parameters} || all params: {all_param} || trainable%: {100 * trainable_parameters / all_param}"
    )


def freeze_model(model):
    for _, param in model.named_parameters():
        param.requires_grad = False


class CODI_CT(torch.nn.Module):
    """
    CODI model enhanced with PMPS-style supervision on latent embeddings.

    Architecture overview:
      1. Encode question → initial latent embedding
      2. Iterate latent embeddings num_latent times (implicit CoT)
      3. Decode answer from the final latent state
      4. Teacher task: forward pass with explicit CoT tokens
      5. Losses:
         - Student CE loss (answer prediction)
         - Teacher CE loss (CoT + answer prediction)
         - Distillation loss (align answer-position hidden states)
         - PMPS loss (NEW: align latent embeddings with CoT hidden states
           in a shared cluster space)

    NOTE: the class name is preserved as ``CODI_CT`` to avoid breaking
    existing import paths in ``train.py`` / ``test.py``.
    """

    def __init__(self, model_args, training_args, lora_config):
        super().__init__()
        self.model_args = model_args
        self.training_args = training_args
        self.model_name = model_args.model_name_or_path
        model_wrapper_class = AutoModelForCausalLM

        if model_args.full_precision:
            self.codi = model_wrapper_class.from_pretrained(
                self.model_name,
                torch_dtype=(torch.bfloat16 if training_args.bf16 else torch.float32),
                resume_download=True,
            )
        else:
            self.codi = model_wrapper_class.from_pretrained(
                self.model_name,
                torch_dtype=(torch.bfloat16 if training_args.bf16 else torch.float32),
                resume_download=True,
                quantization_config=transformers.BitsAndBytesConfig(
                    load_in_4bit=True,
                    bnb_4bit_compute_dtype=torch.bfloat16,
                    bnb_4bit_use_double_quant=False,
                    bnb_4bit_quant_type="nf4",
                ),
            )

        ori_vocab_size = self.codi.config.vocab_size
        self.training = self.model_args.train

        # Special tokens to enclose the latent embeddings
        self.pad_token_id = ori_vocab_size
        self.bot_id = ori_vocab_size + 1
        self.eot_id = ori_vocab_size + 2

        self.codi.resize_token_embeddings(ori_vocab_size + 3)

        self.dim = self.codi.config.hidden_size
        self.num_latent = training_args.num_latent
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)

        # LoRA
        if training_args.use_lora:
            self.codi = get_peft_model(self.codi, lora_config)

        # Projection Layer
        self.use_prj = training_args.use_prj
        self.prj_no_ln = training_args.prj_no_ln
        if training_args.use_prj:
            self.prj = nn.Sequential(
                nn.Dropout(training_args.prj_dropout),
                nn.Linear(self.dim, training_args.prj_dim),
                nn.GELU(),
                nn.Linear(training_args.prj_dim, self.dim),
            )
            if not self.prj_no_ln:
                self.prj.add_module("ln", nn.LayerNorm(self.dim))

        # ===== Probabilistic Latent Embedding Config (NEW) =====
        self.enable_top_k = training_args.enable_top_k
        self.latent_top_k = training_args.latent_top_k
        self.latent_top_k_temperature = training_args.latent_top_k_temperature
        # ===== PMPS Loss Module (NEW) =====
        self.pmps_loss_factor = training_args.pmps_loss_factor
        self.pmps_layer = training_args.pmps_layer
        self.pmps_freeze_prototypes_steps = training_args.pmps_freeze_prototypes_steps
        # Progressive Sequential PMPS config
        self.pmps_position_bias_max = training_args.pmps_position_bias_max
        self.pmps_position_bias_schedule = training_args.pmps_position_bias_schedule
        # Latent diversity regularisation
        self.diversity_loss_factor = training_args.diversity_loss_factor
        if self.pmps_loss_factor > 0:
            self.pmps_module = PMPSLatentSupervisor(
                hidden_dim=self.dim,
                feat_dim=training_args.pmps_feat_dim,
                hidden_mlp=training_args.pmps_hidden_mlp,
                nmb_prototypes=training_args.pmps_nmb_prototypes,
                temperature=training_args.pmps_temperature,
                match_temperature=training_args.pmps_match_temperature,
                epsilon=training_args.pmps_epsilon,
                sinkhorn_iterations=training_args.pmps_sinkhorn_iterations,
                use_projector=training_args.pmps_use_projector,
                position_bias_strength=training_args.pmps_position_bias_max,
            )

        # Losses
        self.print_loss = training_args.print_loss
        self.ref_loss_factor = training_args.ref_loss_factor

        # Cross Entropy Loss
        self.loss_fct = nn.CrossEntropyLoss(ignore_index=-100)

        # Distillation Loss
        self.distill_loss_div_std = training_args.distill_loss_div_std
        self.distill_loss_type = training_args.distill_loss_type
        self.distill_loss_factor = training_args.distill_loss_factor
        if self.distill_loss_type == "smooth_l1":
            self.distill_loss_fct = nn.SmoothL1Loss()
        elif self.distill_loss_type == "l2":
            self.distill_loss_fct = nn.MSELoss()
        else:
            raise NotImplementedError

        # General
        self.fix_attn_mask = training_args.fix_attn_mask

        if self.tokenizer.pad_token_id is None:
            self.tokenizer.add_special_tokens({"pad_token": "[PAD]"})
            self.tokenizer.pad_token_id = self.pad_token_id

        if self.training:
            self.init()

    def get_embd(self, model, model_name):
        try:
            if "pythia" in model_name:
                return model.get_base_model().gpt_neox.embed_in
            elif "gpt2" in model_name:
                try:
                    return model.get_base_model().transformer.wte
                except Exception:
                    return model.transformer.wte
            else:
                try:
                    return model.get_base_model().model.embed_tokens
                except Exception:
                    return model.model.embed_tokens
        except AttributeError:
            if "pythia" in model_name:
                return model.gpt_neox.embed_in
            raise NotImplementedError

    def init(self):
        print_trainable_parameters(self)
        if self.training_args.restore_from is not None and self.training_args.restore_from != "":
            print(f"Loading from the pretrained checkpoint: {self.training_args.restore_from}...")
            restore_path = self.training_args.restore_from
            # Support both single file path and directory (sharded) path
            if os.path.isdir(restore_path):
                state_dict = load_safetensors_state_dict(restore_path)
            elif os.path.isfile(restore_path):
                state_dict = load_file(restore_path)
            else:
                raise FileNotFoundError(f"Checkpoint path not found: {restore_path}")
            self.load_state_dict(state_dict)
            print(f"Finished loading from {self.training_args.restore_from}")

    def _compute_current_bias_strength(self, step_ratio):
        """Compute current positional bias strength based on training progress.

        Implements curriculum learning: the bias starts strong (encouraging
        strict sequential matching) and decays over training (allowing free
        matching once the model has learned basic sequential structure).

        Args:
            step_ratio: float in [0, 1], current_step / total_steps.

        Returns:
            float: current positional bias strength.
        """
        max_strength = self.pmps_position_bias_max
        if max_strength <= 0:
            return 0.0

        schedule = self.pmps_position_bias_schedule
        ratio = min(max(step_ratio if step_ratio is not None else 0.0, 0.0), 1.0)

        if schedule == "cosine":
            return max_strength * 0.5 * (1 + math.cos(math.pi * ratio))
        elif schedule == "linear":
            return max_strength * (1.0 - ratio)
        elif schedule == "constant":
            return max_strength
        else:
            return max_strength

    def _extract_hidden_states_at_layer(self, hidden_states_tuple, layer_spec):
        """
        Extract hidden states from a specific layer or average across layers.

        Args:
            hidden_states_tuple: tuple of (num_layers+1,) tensors, each [batch, seq_len, hidden_size]
            layer_spec: "last", "all", or an integer string like "6"

        Returns:
            hidden_states: [batch, seq_len, hidden_size]
        """
        if layer_spec == "last":
            return hidden_states_tuple[-1]
        elif layer_spec == "all":
            stacked = torch.stack(hidden_states_tuple[1:], dim=0)  # skip embedding layer
            return stacked.mean(dim=0)
        else:
            layer_idx = int(layer_spec)
            return hidden_states_tuple[layer_idx]
        

    # def _compute_topk_weighted_embedding(self, logits, embed_layer):
    #     """
    #     Compute the latent embedding via top-k softmax weighted sum of token embeddings.

    #     Supports both 2D and 3D input:
    #       - 2D: [batch, vocab_size] → returns [batch, 1, hidden_size]
    #       - 3D: [batch, seq_len, vocab_size] → returns [batch, seq_len, hidden_size]

    #     Args:
    #         logits: [batch, vocab_size] or [batch, seq_len, vocab_size]
    #         embed_layer: the embedding layer (nn.Embedding) to look up token embeddings.

    #     Returns:
    #         latent_embd: [batch, 1, hidden_size] (2D input) or [batch, seq_len, hidden_size] (3D input)
    #     """
    #     top_k = self.latent_top_k
    #     temperature = self.latent_top_k_temperature

    #     is_2d = (logits.dim() == 2)
    #     if is_2d:
    #         # Expand to 3D for unified processing: [batch, 1, vocab_size]
    #         logits = logits.unsqueeze(1)

    #     # logits: [batch, seq_len, vocab_size]
    #     batch_size, seq_len, vocab_size = logits.shape

    #     # Select top-k tokens by logit value along vocab dim
    #     top_k_logits, top_k_indices = torch.topk(logits, k=top_k, dim=-1)  # [batch, seq_len, k]

    #     # Softmax over the top-k logits with temperature scaling
    #     top_k_weights = F.softmax(top_k_logits / temperature, dim=-1)  # [batch, seq_len, k]

    #     # Look up embeddings for the top-k tokens: [batch, seq_len, k, hidden_size]
    #     top_k_embeds = embed_layer(top_k_indices)

    #     # Weighted sum via einsum: [batch, seq_len, hidden_size]
    #     latent_embd = torch.einsum('bsk,bskd->bsd', top_k_weights, top_k_embeds)

    #     return latent_embd

    def _compute_topk_weighted_embedding(self, logits_last_token, embed_layer):
        """
        Compute the latent embedding via top-k softmax weighted sum of token embeddings.

        Args:
            logits_last_token: [batch, vocab_size] — logits at the last token position.
            embed_layer: the embedding layer (nn.Embedding) to look up token embeddings.

        Returns:
            latent_embd: [batch, 1, hidden_size] — the weighted-sum embedding.
        """
        # breakpoint()
        top_k = self.latent_top_k
        temperature = self.latent_top_k_temperature

        # Select top-k tokens by logit value
        top_k_logits, top_k_indices = torch.topk(logits_last_token, k=top_k, dim=-1)  # [batch, k]

        # Softmax over the top-k logits with temperature scaling
        top_k_weights = F.softmax(top_k_logits / temperature, dim=-1)  # [batch, k]

        # Look up embeddings for the top-k tokens: [batch, k, hidden_size]
        top_k_embeds = embed_layer(top_k_indices)

        # Weighted sum: [batch, 1, k] @ [batch, k, hidden_size] -> [batch, 1, hidden_size]
        latent_embd = torch.bmm(top_k_weights.unsqueeze(1), top_k_embeds)

        return latent_embd

    def _compute_all_weighted_embedding(self, logits_last_token, embed_layer):
        """
        Compute the latent embedding via full softmax weighted sum of all token embeddings.

        Args:
            logits_last_token: [batch, vocab_size] — logits at the last token position.
            embed_layer: the embedding layer (nn.Embedding) to look up token embeddings.

        Returns:
            latent_embd: [batch, 1, hidden_size] — the weighted-sum embedding.
        """
        # breakpoint()
        temperature = self.latent_top_k_temperature
        vocab_size = logits_last_token.size(-1)
        all_indices = torch.arange(vocab_size, device=logits_last_token.device)
        all_weights = F.softmax(logits_last_token / temperature, dim=-1)  # [batch, vocab_size]
        all_embeds = embed_layer(all_indices)

        # Weighted sum: [batch, 1, k] @ [batch, k, hidden_size] -> [batch, 1, hidden_size]
        latent_embd = torch.matmul(all_weights, all_embeds).unsqueeze(1)

        return latent_embd

    def forward(
        self,
        encoder_input_ids: torch.LongTensor = None,
        decoder_input_ids: torch.LongTensor = None,
        ref_input_ids: torch.LongTensor = None,
        labels: Optional[torch.LongTensor] = None,
        encoder_attention_mask: Optional[torch.LongTensor] = None,
        ref_answer_position: Optional[torch.LongTensor] = None,
        model_answer_position: Optional[torch.LongTensor] = None,
        ref_attention_mask: Optional[torch.LongTensor] = None,
        ref_labels: torch.LongTensor = None,
        cot_start_position: Optional[torch.LongTensor] = None,  # NEW: start of CoT in ref_input_ids
        cot_end_position: Optional[torch.LongTensor] = None,  # NEW: end of CoT in ref_input_ids
        step: int = None,
        step_ratio: float = None,
    ):
        if not self.fix_attn_mask:
            ref_attention_mask = None

        # latent_hidden_states_list collects hidden states from each latent iteration step only
        # (not from the question encoding step, which is not a latent reasoning step)
        latent_hidden_states_list = []

        # Get the embedding layer (used for top-k weighted embedding computation)
        embed_layer = self.get_embd(self.codi, self.model_name)

        # ========== Step 1: Encode the question ==========
        past_key_values = None
        outputs = self.codi(input_ids=encoder_input_ids, use_cache=True, output_hidden_states=True, past_key_values=past_key_values, attention_mask=encoder_attention_mask)
        past_key_values = outputs.past_key_values

        if self.enable_top_k:
            logits_last = outputs.logits[:, -1, :]  # [batch, vocab_size]
            if self.latent_top_k == 0:
                latent_embd = self._compute_all_weighted_embedding(logits_last, embed_layer)
            else:
                latent_embd = self._compute_topk_weighted_embedding(logits_last, embed_layer)  # [batch, 1, hidden_size]
        else:
            latent_embd = outputs.hidden_states[-1][:, -1, :].unsqueeze(1)  # as the next input

        if self.use_prj:
            latent_embd = self.prj(latent_embd)

        len_pred_loss = 0
        dynamic_mask = None
        if self.fix_attn_mask:
            dynamic_mask = torch.ones((encoder_attention_mask.size(0), self.num_latent), device=ref_labels.device)

        # Iterate over the latent embeddings
        distill_loss_total = 0
        ce_loss_total = 0

        with torch.no_grad():
            ref_outputs = self.codi(input_ids=ref_input_ids, output_hidden_states=True, attention_mask=ref_attention_mask)
        ref_outputs_with_grad = self.codi(input_ids=ref_input_ids, output_hidden_states=True, attention_mask=ref_attention_mask)
        
        # Formatting for deprecated exps
        ref_outputs_list = [ref_outputs] 
        ref_input_ids = [ref_input_ids] 

        # Process the position tensor
        # Normalise the position definition 
        if "llama" in self.model_name.lower() or "qwen" in self.model_name.lower(): # there is one more token standing for " " 
            model_answer_position = model_answer_position + 1
            ref_answer_position = ref_answer_position + 1
        # For DEBUG
        if self.training_args.print_ref_model_stats:
            for i, (ref_inputs, ref_outputs) in enumerate(zip(ref_input_ids, ref_outputs_list)):
                # evalutae the reference model
                if len(ref_outputs_list) > 1:
                    pos = ref_answer_position[i]
                else:
                    pos = ref_answer_position
                ref_probs = torch.nn.functional.softmax(ref_outputs.logits, dim=-1)
                input_positions = (pos-1).unsqueeze(1).unsqueeze(1).expand(-1, -1, ref_probs.size(2))
                ref_probs_at_positions = ref_probs.gather(1, input_positions)
                probe_positions_positions = pos.unsqueeze(1)
                probe_positions = ref_inputs.gather(1, probe_positions_positions).unsqueeze(1)
                ref_probs_of_target = ref_probs_at_positions.gather(2, probe_positions)
                print(f'stage{i}: mean of the prob of the target token: {ref_probs_of_target.mean()}')

        model_answer_position = model_answer_position - 1
        ref_answer_position = ref_answer_position - 1

        # ========== Step 3: Iterate latent embeddings & collect hidden states ==========
        num_latent = self.num_latent
        logits = None  # initialise to avoid NameError when num_latent == 0
        ce_loss = 0
        distill_loss = 0

        if self.num_latent != 0:
            for i in range(num_latent):
                # Implicit CoT generation
                outputs = self.codi(inputs_embeds=latent_embd, use_cache=True, output_hidden_states=True, past_key_values=past_key_values)
                past_key_values = outputs.past_key_values

                if self.enable_top_k:
                    logits_last = outputs.logits[:, -1, :]  # [batch, vocab_size]
                    if self.latent_top_k == 0:
                        latent_embd = self._compute_all_weighted_embedding(logits_last, embed_layer)
                    else:
                        latent_embd = self._compute_topk_weighted_embedding(logits_last, embed_layer)  # [batch, 1, hidden_size]
                else:
                    latent_embd = outputs.hidden_states[-1][:, -1, :].unsqueeze(1) # [batch, 1, hidden_size]
                
                if self.use_prj:
                    latent_embd = self.prj(latent_embd)

                # NEW: Collect the latent hidden state at this step for PMPS loss.
                # We use the model's hidden_size representation (not the vocab-prob),
                # which is what PMPS expects.
                if self.pmps_loss_factor > 0:
                    if self.enable_top_k:
                        latent_hidden_states_list.append(latent_embd.squeeze(1))  # [batch, hidden_size]
                    else:
                        latent_hidden = self._extract_hidden_states_at_layer(outputs.hidden_states, self.pmps_layer)[:, -1, :]  # [batch, hidden_size]
                        latent_hidden_states_list.append(latent_hidden)
                    

                # Calculate distillation loss and CE loss at the last latent step
                if i == num_latent - 1: # the last latent embedding
                    embds = self.get_embd(self.codi, self.model_name)(decoder_input_ids)

                    if dynamic_mask is not None: # Prevent attending the paddings
                        decoder_mask = torch.ones((embds.size(0), embds.size(1)), dtype=torch.bool).to(dynamic_mask)
                        dynamic_mask = torch.cat((encoder_attention_mask, dynamic_mask, decoder_mask), dim=1)
                        dynamic_mask = dynamic_mask.bool()
                    # Student task's output
                    outputs = self.codi(inputs_embeds=embds, use_cache=True, output_hidden_states=True, past_key_values=past_key_values, attention_mask=dynamic_mask) 
                    # Teacher task's output
                    ref_outputs = ref_outputs_list[0]

                    # Distillation loss between the teacher's logits and the student's logits 
                    distill_loss = 0
                    for j, (out, ref_out) in enumerate(zip(outputs.hidden_states, ref_outputs.hidden_states)):
                        ref_selected = ref_out.gather(1, ref_answer_position.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, ref_out.size(-1)))
                        out_selected = out.gather(1, model_answer_position.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, out.size(-1)))

                        distill_loss_tmp = self.distill_loss_fct(out_selected, ref_selected.detach())
                        if self.distill_loss_div_std:
                            if self.distill_loss_type == "l2":
                                distill_loss_tmp /= ref_selected.std()
                            distill_loss_tmp /= ref_selected.std()
                        distill_loss += distill_loss_tmp

                    distill_loss /= len(outputs.hidden_states)

                    if self.print_loss:
                        print(f"latent{i}: distill_loss={distill_loss}")

                    distill_loss_total += distill_loss

                    # CE loss for the student task
                    if i == num_latent - 1:
                        logits = outputs.logits
                        effective_logits = logits[:, :-1, :]
                        effective_logits = effective_logits.reshape(-1, logits.size(-1))
                        target_ids = labels[:, 1:].reshape(-1)       #remove eot token 
                        ce_loss = self.loss_fct(effective_logits, target_ids)
                        ce_loss_total += ce_loss

        # ========== Step 4: Teacher CE loss ==========
        ref_ce_loss = 0
        ref_logits = ref_outputs_with_grad.logits
        effective_ref_logits = ref_logits[:, :-1, :]
        effective_ref_logits = effective_ref_logits.reshape(-1, ref_logits.size(-1))
        ref_target_ids = ref_labels[:, 1:].reshape(-1)
        ref_ce_loss = self.loss_fct(effective_ref_logits, ref_target_ids)
        ref_ce_loss *= self.ref_loss_factor

        # ========== Step 5: PMPS Loss (NEW) ==========
        pmps_loss_value = 0
        pmps_log = {}
        if (
            self.pmps_loss_factor > 0
            and len(latent_hidden_states_list) > 0
            and cot_start_position is not None
        ):
            # Stack latent hidden states: [batch, num_latent, hidden_size]
            latent_hiddens = torch.stack(latent_hidden_states_list, dim=1)

            # Extract CoT hidden states from teacher output (detached: teacher
            # is the supervision target, not a gradient source for PMPS).
            ref_hidden = self._extract_hidden_states_at_layer(ref_outputs_list[0].hidden_states, self.pmps_layer).detach()  # [batch, ref_seq_len, hidden_size]
            
            ref_logits = ref_outputs_list[0].logits.detach()
    
            batch_size = ref_hidden.size(0)
            cot_hidden_list = []
            cot_lens = []
            for batch_idx in range(batch_size):
                start = cot_start_position[batch_idx].item()
                end = cot_end_position[batch_idx].item()
                # Autoregressive offset: hidden_states[i] corresponds to the
                # state that predicts token i+1, so [start-1 : end-1] are the
                # states that produce CoT tokens C_0 ... C_{c-1}.
                if self.enable_top_k:
                    cot_logits = ref_logits[batch_idx, start - 1 : end - 1, :]
                    with torch.no_grad():
                        if self.latent_top_k == 0:
                            cot_hidden = self._compute_all_weighted_embedding(cot_logits, embed_layer).squeeze(1)
                        else:
                            cot_hidden = self._compute_topk_weighted_embedding(cot_logits, embed_layer).squeeze(1)
                else:
                    cot_hidden = ref_hidden[batch_idx, start - 1 : end - 1, :]  # [cot_len_i, hidden_size]
                cot_hidden_list.append(cot_hidden)
                cot_lens.append(cot_hidden.size(0))

            max_cot_len = max(cot_lens) if cot_lens else 0
            if max_cot_len > 0:
                # Pad cot hiddens to a unified [batch, max_cot_len, hidden_size]
                # and build a 0/1 mask so that the PMPS module can ignore
                # padded positions (which would otherwise pollute the
                # Sinkhorn equipartition).
                cot_padded = ref_hidden.new_zeros(batch_size, max_cot_len, ref_hidden.size(-1))
                cot_mask = ref_hidden.new_zeros(batch_size, max_cot_len)
                for batch_idx, cot_h in enumerate(cot_hidden_list):
                    L = cot_h.size(0)
                    if L > 0:
                        cot_padded[batch_idx, :L] = cot_h
                        cot_mask[batch_idx, :L] = 1.0

                latent_mask = ref_hidden.new_ones(batch_size, latent_hiddens.size(1))

                # Compute current positional bias strength (curriculum)
                current_bias_strength = self._compute_current_bias_strength(step_ratio)

                self.pmps_module.normalize_prototypes()
                pmps_loss_value, pmps_log = self.pmps_module(
                    latent_hidden=latent_hiddens,
                    cot_hidden=cot_padded,
                    latent_mask=latent_mask,
                    cot_mask=cot_mask,
                    position_bias_strength_override=current_bias_strength,
                )
                pmps_log["position_bias_strength"] = current_bias_strength

        # ========== Step 5b: Latent Diversity Regularisation (NEW) ==========
        diversity_loss_value = 0
        latent_avg_cosine_sim = 0
        if (
            self.diversity_loss_factor > 0
            and len(latent_hidden_states_list) > 1
        ):
            # latent_hiddens: (B, N_l, D) — may already be stacked above
            if 'latent_hiddens' not in dir():
                latent_hiddens = torch.stack(latent_hidden_states_list, dim=1)
            # L2-normalise each latent embedding
            latent_normed = F.normalize(latent_hiddens, dim=-1, p=2)  # (B, N_l, D)
            # Cosine similarity matrix between all pairs of latent tokens
            sim_matrix = torch.bmm(latent_normed, latent_normed.transpose(1, 2))  # (B, N_l, N_l)
            # Minimise off-diagonal elements (encourage orthogonality)
            N_l = latent_hiddens.size(1)
            eye = torch.eye(N_l, device=sim_matrix.device, dtype=sim_matrix.dtype).unsqueeze(0)  # (1, N_l, N_l)
            off_diag = sim_matrix - eye  # zero out diagonal
            diversity_loss_value = off_diag.pow(2).mean()
            # Diagnostic: average off-diagonal cosine similarity
            with torch.no_grad():
                off_diag_mask = 1.0 - eye
                latent_avg_cosine_sim = (sim_matrix * off_diag_mask).sum() / off_diag_mask.sum().clamp(min=1)

        # ========== Step 6: Combine all losses ==========
        distill_loss *= self.distill_loss_factor
        distill_loss_total *= self.distill_loss_factor

        # Weight PMPS loss
        weighted_pmps_loss = (
            pmps_loss_value * self.pmps_loss_factor
            if isinstance(pmps_loss_value, torch.Tensor)
            else 0
        )

        # Weight diversity loss
        weighted_diversity_loss = (
            diversity_loss_value * self.diversity_loss_factor
            if isinstance(diversity_loss_value, torch.Tensor)
            else 0
        )

        if self.print_loss:
            pmps_display = (
                pmps_log.get("pmps_loss").item()
                if isinstance(pmps_log.get("pmps_loss", None), torch.Tensor)
                else 0
            )
            diversity_display = (
                diversity_loss_value.item()
                if isinstance(diversity_loss_value, torch.Tensor)
                else diversity_loss_value
            )
            ce_display = ce_loss.item() if isinstance(ce_loss, torch.Tensor) else ce_loss
            distill_display = (
                distill_loss.item() if isinstance(distill_loss, torch.Tensor) else distill_loss
            )
            ref_ce_display = (
                ref_ce_loss.item() if isinstance(ref_ce_loss, torch.Tensor) else ref_ce_loss
            )
            print(
                f"ce_loss={ce_display:.4f} | distill_loss={distill_display:.4f} | "
                f"ref_ce_loss={ref_ce_display:.4f} | pmps_loss={pmps_display:.4f} | "
                f"diversity_loss={diversity_display:.4f}"
            )

        # Total loss WITHOUT PMPS (used during the warmup phase before pmps_start_epoch)
        no_pmps_loss = ce_loss_total + distill_loss_total + ref_ce_loss

        # Final total loss including PMPS + diversity
        total_loss = no_pmps_loss + weighted_pmps_loss + weighted_diversity_loss

        # ========== Step 7: Build the output dict ==========
        # All scalar fields are converted to python floats so the Trainer
        # logger / wandb can serialise them. Tensors are kept for the
        # ``loss`` and ``no_pmps_loss`` so that gradients flow correctly
        # through the optimiser.
        def _scalar(x):
            return x.item() if isinstance(x, torch.Tensor) else float(x)

        output_dict = {
            "loss": total_loss,
            "no_pmps_loss": no_pmps_loss,
            "ce_loss": _scalar(ce_loss_total),
            "distill_loss": _scalar(distill_loss_total),
            "ref_ce_loss": _scalar(ref_ce_loss),
            "pmps_loss": _scalar(pmps_log.get("pmps_loss", 0)),
            "pmps_loss_c2l": _scalar(pmps_log.get("pmps_loss_c2l", 0)),
            "pmps_loss_l2c": _scalar(pmps_log.get("pmps_loss_l2c", 0)),
            "code_entropy_latent": _scalar(pmps_log.get("code_entropy_latent", 0)),
            "code_entropy_cot": _scalar(pmps_log.get("code_entropy_cot", 0)),
            "match_entropy_c2l": _scalar(pmps_log.get("match_entropy_c2l", 0)),
            "match_order_correlation": _scalar(pmps_log.get("match_order_correlation", 0)),
            "position_bias_strength": _scalar(pmps_log.get("position_bias_strength", 0)),
            "diversity_loss": _scalar(diversity_loss_value),
            "latent_avg_cosine_sim": _scalar(latent_avg_cosine_sim),
        }
        return output_dict