src/train.py

import os
import re
import sys
from collections import defaultdict
from typing import cast
import pandas as pd
import numpy as np
import inspect
import random

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import (
    T5Tokenizer,
    T5ForConditionalGeneration,
    TrainingArguments,
    Trainer,
    TrainerCallback,
)
from transformers.optimization import Adafactor, AdafactorSchedule
from peft import get_peft_model, LoraConfig, TaskType
import matplotlib.pyplot as plt


print("=" * 80)
print("PROTRANS LORA FINE-TUNING: CONTRASTIVE + MLM LOSS (ContraMLM v1)")
print("=" * 80)

# ============================================================================
# CONFIGURATION
# ============================================================================

print("\n" + "=" * 80)
print("CONFIGURATION")
print("=" * 80)

# Model configuration
MODEL_NAME = "Rostlab/prot_t5_xl_uniref50"

# LoRA configuration
LORA_R = 32
LORA_ALPHA = 64
LORA_DROPOUT = 0.1
LORA_TARGET_MODULES = ["q", "k", "v", "o"]
LORA_TASK_TYPE = TaskType.SEQ_2_SEQ_LM

# Training configuration
#   BATCH_SIZE = number of proteins drawn from the dataset per forward pass.
#   For each non-orphan, the collator samples 1 positive on-the-fly, so the
#   actual forward-pass batch has BATCH_SIZE..2*BATCH_SIZE unique proteins.
BATCH_SIZE = 32
GRADIENT_ACCUMULATION_STEPS = 2
NUM_EPOCHS = 2
MAX_LENGTH = 512
NOISE_DENSITY = 0.15

# Contrastive loss — Contrastive (Con) with full adjacency matrix
#   Total loss = (1 - CONTRASTIVE_LAMBDA) * MLM_loss + CONTRASTIVE_LAMBDA * Con_loss
#   A per-batch adjacency matrix (N×N) is built from the VISEQ pair graph.
#   Every known positive pair in the batch contributes to the numerator;
#   false negatives are impossible by construction (adj built from ground-truth graph).
#   CONTRASTIVE_LAMBDA is the convex-combination weight for the contrastive term.
CONTRASTIVE_LAMBDA = 0.2
CONTRASTIVE_TEMPERATURE = 0.1   # lower → sharper distribution → harder loss

# Curriculum settings (same flags as Default_v2, applied to MLM component only)
USE_LOSS_CLIPPING_CURRICULUM = False
NUM_STAGES = 10
KEEP_FRACTION_START = 0.20
KEEP_FRACTION_END = 1.00
LARGEST = False  # False = keep easiest losses first

version = "v1_1"

OUTPUT_DIR = f"./runs/protrans_XL_Full_lora_envhog_ContraMLM_{version}"

# Data paths
FASTA_FILE = "./data/envhog_phrog2/envhog_filtered_proteins.fasta"
CSV_FILE   = "./data/envhog_phrog2/envhog_phrog2__low_thr_enriched_final.csv"
PAIRS_FILE = "./data/envhog_phrog2/all_positive_viseq_pairs.csv"

# Sample caps (applied to the pair list; existing ~168 K pairs are well below cap)
MAX_TRAIN_SAMPLES = 400000
MAX_EVAL_SAMPLES  = 50000

print(f"Model:                    {MODEL_NAME}")
print(f"LoRA rank:                {LORA_R}")
print(f"Batch size (drawn):       {BATCH_SIZE}  ({BATCH_SIZE}–{2*BATCH_SIZE} unique proteins per step)")
print(f"Gradient accumulation:    {GRADIENT_ACCUMULATION_STEPS}")
print(f"Epochs:                   {NUM_EPOCHS}")
print(f"Max sequence length:      {MAX_LENGTH}")
print(f"Noise density (MLM):      {NOISE_DENSITY}")
print(f"Contrastive lambda:       {CONTRASTIVE_LAMBDA}")
print(f"Contrastive temperature:  {CONTRASTIVE_TEMPERATURE}")
print(f"Loss clipping curriculum: {USE_LOSS_CLIPPING_CURRICULUM}")
print(f"Output directory:         {OUTPUT_DIR}")

# ============================================================================
# CHECK PYTORCH AND GPU
# ============================================================================

print("\n" + "=" * 80)
print("GPU STATUS")
print("=" * 80)

print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available:  {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA device:     {torch.cuda.get_device_name(0)}")
    print(f"Number of GPUs:  {torch.cuda.device_count()}")

# ============================================================================
# LOAD DATA
# ============================================================================

print("\n" + "=" * 80)
print("LOADING DATA")
print("=" * 80)

for _f in [FASTA_FILE, CSV_FILE, PAIRS_FILE]:
    if not os.path.exists(_f):
        print(f"ERROR: File not found: {_f}")
        sys.exit(1)

# --- 1. FASTA sequences ---
print("Reading FASTA sequences...")
fasta_seqs = {}          # {envhog_id: raw_aa_sequence}
_cur_id  = None
_cur_seq = []
with open(FASTA_FILE) as fh:
    for line in fh:
        line = line.rstrip()
        if line.startswith(">"):
            if _cur_id is not None:
                fasta_seqs[_cur_id] = "".join(_cur_seq)
            _cur_id  = line[1:].split()[0]
            _cur_seq = []
        else:
            _cur_seq.append(line)
if _cur_id is not None:
    fasta_seqs[_cur_id] = "".join(_cur_seq)
print(f"  Sequences in FASTA: {len(fasta_seqs):,}")

# --- 2. Protein metadata CSV ---
print("Reading protein metadata CSV...")
meta_df = pd.read_csv(CSV_FILE)
print(f"  Loaded {len(meta_df):,} rows | columns: {meta_df.columns.tolist()}")

# Keep only proteins that have a sequence in the FASTA
meta_df = meta_df[meta_df["ENVHOG"].isin(fasta_seqs)].reset_index(drop=True)
print(f"  After FASTA intersection: {len(meta_df):,} proteins retained")

# --- 3. Build lookup structures ---
envhog_to_viseq    = dict(zip(meta_df["ENVHOG"], meta_df["VISEQ"]))
viseq_to_proteins  = defaultdict(list)   # viseq  → [envhog_id, ...]
for row in meta_df.itertuples(index=False):
    viseq_to_proteins[row.VISEQ].append(row.ENVHOG)

# --- 4. Load positive VISEQ pairs ---
print("Reading positive VISEQ pairs CSV...")
pairs_df = pd.read_csv(PAIRS_FILE)
print(f"  Loaded {len(pairs_df):,} positive VISEQ pairs")

# Build bidirectional map: viseq → [list of positive viseqs]
positive_viseq_map = defaultdict(list)
for row in pairs_df.itertuples(index=False):
    positive_viseq_map[row.viseq_A].append(row.viseq_B)
    positive_viseq_map[row.viseq_B].append(row.viseq_A)
print(f"  VISEQs with cross-VISEQ positives: {len(positive_viseq_map):,}")

# Convert map values to sets for O(1) lookup during false-negative filtering
positive_viseq_set = {k: set(v) for k, v in positive_viseq_map.items()}

# ============================================================================
# BUILD PROTEIN POOL
# ============================================================================

print("\n" + "=" * 80)
print("BUILDING PROTEIN POOL")
print("=" * 80)


def prepare_t5_seq(seq: str) -> str:
    """Remove gaps, replace rare AAs, space-separate for T5 tokeniser."""
    seq = seq.replace(" ", "")
    seq = re.sub(r"[UZOB]", "X", seq)
    return " ".join(list(seq))


def sample_positive_protein(anchor_envhog: str, anchor_viseq: str):
    """
    Return one positive protein for the given anchor, or None if none exists.

    Positive candidates:
      1. Other proteins sharing the same VISEQ.
      2. A protein from a cross-VISEQ positive (drawn from positive_viseq_map).

    When both options are available, one is chosen at random to expose the model
    to both types of similarity signal across the training epoch.
    """
    same_viseq  = [p for p in viseq_to_proteins[anchor_viseq] if p != anchor_envhog]
    cross_viseqs = positive_viseq_map.get(anchor_viseq, [])

    has_same  = bool(same_viseq)
    has_cross = bool(cross_viseqs)

    if not has_same and not has_cross:
        return None   # orphan: no confirmed positive exists

    if has_same and has_cross:
        strategy = random.choice(["same", "cross"])
    elif has_same:
        strategy = "same"
    else:
        strategy = "cross"

    if strategy == "same":
        return random.choice(same_viseq)

    # Cross-VISEQ: pick a random positive VISEQ, then a random protein from it
    pos_viseq = random.choice(cross_viseqs)
    pos_prots = viseq_to_proteins.get(pos_viseq, [])
    if pos_prots:
        return random.choice(pos_prots)
    # The cross-VISEQ has no protein in our filtered dataset → fall back to same
    if same_viseq:
        return random.choice(same_viseq)
    return None


# Flat list of ALL proteins (orphans and non-orphans alike).
# Positives are sampled on-the-fly in the collator; no pre-building of pairs needed.
all_proteins = []  # list of (seq_t5, viseq, envhog_id)

n_no_fasta = 0
for row in meta_df.itertuples(index=False):
    envhog  = row.ENVHOG
    viseq   = row.VISEQ
    raw_seq = fasta_seqs.get(envhog)
    if raw_seq is None:
        n_no_fasta += 1
        continue
    all_proteins.append((prepare_t5_seq(raw_seq), viseq, envhog))

print(f"  Total proteins in pool:        {len(all_proteins):,}")
print(f"  Skipped (no FASTA sequence):   {n_no_fasta:,}")
print()
print("  Orphan proteins automatically get adj_matrix rows of all-False:")
print("  they contribute to MLM only, excluded from the Con mean.")

if len(all_proteins) == 0:
    print("ERROR: No proteins found. Check data paths.")
    sys.exit(1)

# ============================================================================
# SPLIT DATA
# ============================================================================

print("\n" + "=" * 80)
print("SPLITTING DATA")
print("=" * 80)

random.shuffle(all_proteins)
n_val_proteins = max(1, int(len(all_proteins) * 0.1))
val_proteins   = all_proteins[:n_val_proteins]
train_proteins = all_proteins[n_val_proteins:]

if MAX_TRAIN_SAMPLES and len(train_proteins) > MAX_TRAIN_SAMPLES:
    train_proteins = train_proteins[:MAX_TRAIN_SAMPLES]
if MAX_EVAL_SAMPLES and len(val_proteins) > MAX_EVAL_SAMPLES:
    val_proteins = val_proteins[:MAX_EVAL_SAMPLES]

print(f"Training proteins:   {len(train_proteins):,}")
print(f"Validation proteins: {len(val_proteins):,}")

# ============================================================================
# DATASET
# ============================================================================

print("\n" + "=" * 80)
print("BUILDING DATASET")
print("=" * 80)


class ProteinGraphDataset(torch.utils.data.Dataset):
    """
    A flat pool of all proteins. Each item is a single protein.
    The PairGraphCollator samples positives on-the-fly and builds the
    per-batch adjacency matrix for Contrastive loss.
    """

    def __init__(self, proteins):
        # proteins: list of (seq_t5, viseq, envhog_id)
        self.items = [
            {"seq": s, "viseq": v, "envhog_id": e}
            for s, v, e in proteins
        ]

    def __len__(self):
        return len(self.items)

    def __getitem__(self, idx):
        return self.items[idx]


train_dataset = ProteinGraphDataset(train_proteins)
val_dataset   = ProteinGraphDataset(val_proteins)

print(f"Train dataset: {len(train_dataset):,} proteins")
print(f"Val dataset:   {len(val_dataset):,} proteins")
print(f"Each batch of {BATCH_SIZE} drawn proteins → {BATCH_SIZE}–{2*BATCH_SIZE} unique proteins after positive sampling")

# ============================================================================
# LOAD TOKENIZER AND MODEL
# ============================================================================

print("\n" + "=" * 80)
print("LOADING TOKENIZER AND MODEL")
print("=" * 80)

tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME, do_lower_case=False, legacy=True)
print(f"Tokenizer loaded: {MODEL_NAME}")

use_bf16     = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
model_dtype  = torch.bfloat16 if use_bf16 else torch.float32

model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME, torch_dtype=model_dtype)
model.config.use_cache = False
if hasattr(model, "gradient_checkpointing_enable"):
    try:
        model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    except TypeError:
        model.gradient_checkpointing_enable()

print(f"Full Encoder-Decoder Model loaded: {MODEL_NAME}")
print(f"Total parameters: {model.num_parameters():,}")

# ============================================================================
# CONFIGURE AND APPLY LORA
# ============================================================================

print("\n" + "=" * 80)
print("CONFIGURING LORA")
print("=" * 80)

lora_config = LoraConfig(
    r=LORA_R,
    lora_alpha=LORA_ALPHA,
    target_modules=LORA_TARGET_MODULES,
    lora_dropout=LORA_DROPOUT,
    bias="none",
    task_type=LORA_TASK_TYPE,
)

model = get_peft_model(model, lora_config)


def _get_input_embedding_layer(model_obj):
    getter = getattr(model_obj, "get_input_embeddings", None)
    if callable(getter):
        return getter()
    base_model = getattr(model_obj, "base_model", None)
    if base_model is not None:
        base_getter = getattr(base_model, "get_input_embeddings", None)
        if callable(base_getter):
            return base_getter()
    raise AttributeError("Could not resolve input embedding layer for model")


def _make_inputs_require_grad(module, inputs, output):
    output.requires_grad_(True)


embedding_layer = cast(nn.Embedding, _get_input_embedding_layer(model))
embedding_layer.register_forward_hook(_make_inputs_require_grad)

model.print_trainable_parameters()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model  = model.to(device)
print(f"Model moved to device: {device}")

# ============================================================================
# DATA COLLATOR
# ============================================================================

print("\n" + "=" * 80)
print("PREPARING DATA COLLATOR")
print("=" * 80)

# Build amino acid token IDs for random-replacement masking
_aa_set = set()
for _aa in "ACDEFGHIKLMNPQRSTVWY":
    for _tid in tokenizer.encode(_aa, add_special_tokens=False):
        if _tid not in (tokenizer.unk_token_id, tokenizer.eos_token_id, tokenizer.pad_token_id):
            _aa_set.add(_tid)
AA_TOKENS = list(_aa_set)
if not AA_TOKENS:
    AA_TOKENS = list(range(3, tokenizer.vocab_size))


class PairGraphCollator:
    """
    Collates BATCH_SIZE drawn protein items into a graph-structured batch.

    On-the-fly positive sampling:
      For each drawn protein that has at least one known positive, sample 1
      positive protein and add it to the pool (deduplicated by envhog_id).

    Adjacency matrix A (N×N):
      A[i,j] = True  iff  proteins i and j are known positives:
                          same VISEQ  OR  cross-VISEQ positive in the pair graph.
      Diagonal is always False (self-loops excluded).
      Orphans (no positives anywhere) have all-False rows → contribute to MLM
      only, automatically excluded from the Con mean by pos_count == 0.

    Output keys:
      all_input_ids      : (N, L) — BART-masked token ids
      all_attention_mask : (N, L)
      all_labels         : (N, L) — original tokens; padding → -100
      adj_matrix         : Python list[list[bool]] (N×N), passed through as-is
    """

    def __init__(self, tokenizer, mlm_probability=0.15, pad_to_multiple_of=8):
        self.tokenizer         = tokenizer
        self.mlm_probability   = mlm_probability
        self.pad_to_multiple_of = pad_to_multiple_of

        mask_id = tokenizer.mask_token_id
        if mask_id is None:
            mask_id = tokenizer.convert_tokens_to_ids("<extra_id_0>")
        self.mask_token_id = mask_id
        self.pad_token_id  = tokenizer.pad_token_id
        self.eos_token_id  = tokenizer.eos_token_id

    def _apply_bart_mask(self, input_tensor, attention_tensor):
        """Apply BART-style masking: 90% → <mask>, 10% → random amino acid."""
        corrupted   = input_tensor.clone()
        special_ids = {self.pad_token_id, self.eos_token_id}

        prob_matrix = torch.full(input_tensor.shape, self.mlm_probability)
        for sid in special_ids:
            prob_matrix[input_tensor == sid] = 0.0
        prob_matrix[attention_tensor == 0] = 0.0

        mask_positions    = torch.bernoulli(prob_matrix).bool()
        replace_with_mask = torch.bernoulli(
            torch.full(mask_positions.shape, 0.9)
        ).bool() & mask_positions
        corrupted[replace_with_mask] = self.mask_token_id

        replace_with_random = mask_positions & ~replace_with_mask
        n_random = int(replace_with_random.sum().item())
        if n_random > 0:
            corrupted[replace_with_random] = torch.tensor(
                random.choices(AA_TOKENS, k=n_random), dtype=torch.long
            )
        return corrupted

    def __call__(self, features):
        # features: list of BATCH_SIZE dicts {seq, viseq, envhog_id}

        # Deduplicate drawn proteins by envhog_id (rare but possible)
        pool_by_id = {}
        for f in features:
            eid = f["envhog_id"]
            if eid not in pool_by_id:
                pool_by_id[eid] = f

        # For each drawn protein, sample 1 positive and add if not already in pool
        for f in list(pool_by_id.values()):
            pos_eid = sample_positive_protein(f["envhog_id"], f["viseq"])
            if pos_eid is not None and pos_eid not in pool_by_id:
                pos_viseq = envhog_to_viseq.get(pos_eid)
                if pos_viseq is not None:
                    pool_by_id[pos_eid] = {
                        "seq":       prepare_t5_seq(fasta_seqs[pos_eid]),
                        "viseq":     pos_viseq,
                        "envhog_id": pos_eid,
                    }

        pool       = list(pool_by_id.values())  # N = 8..16 unique proteins
        N          = len(pool)
        all_seqs   = [p["seq"]   for p in pool]
        all_viseqs = [p["viseq"] for p in pool]

        # Build adjacency matrix (N×N) from the VISEQ pair graph
        adj = []
        for i in range(N):
            vi      = all_viseqs[i]
            pos_set = positive_viseq_set.get(vi, set()) | {vi}
            adj.append([
                (j != i and all_viseqs[j] in pos_set)
                for j in range(N)
            ])

        # Tokenise all N sequences
        encoding  = self.tokenizer(
            all_seqs, truncation=True, max_length=MAX_LENGTH, add_special_tokens=True
        )
        all_ids   = encoding["input_ids"]
        all_masks = encoding["attention_mask"]

        # Pad to longest (aligned to pad_to_multiple_of)
        max_len = max(len(ids) for ids in all_ids)
        if self.pad_to_multiple_of:
            max_len = (
                (max_len + self.pad_to_multiple_of - 1)
                // self.pad_to_multiple_of
                * self.pad_to_multiple_of
            )

        pad_id       = self.pad_token_id
        padded_ids   = []
        padded_masks = []
        for ids, mask in zip(all_ids, all_masks):
            pad_len = max_len - len(ids)
            padded_ids.append(ids + [pad_id] * pad_len)
            padded_masks.append(mask + [0]   * pad_len)

        input_tensor = torch.tensor(padded_ids,   dtype=torch.long)
        attn_tensor  = torch.tensor(padded_masks, dtype=torch.long)

        # Labels for MLM: original tokens; padding positions → -100
        labels = input_tensor.clone()
        labels[attn_tensor == 0] = -100

        corrupted = self._apply_bart_mask(input_tensor, attn_tensor)

        return {
            "adj_matrix":         adj,          # Python list[list[bool]], passed through as-is
            "all_input_ids":      corrupted,
            "all_attention_mask": attn_tensor,
            "all_labels":         labels,
        }


data_collator = PairGraphCollator(
    tokenizer=tokenizer,
    mlm_probability=NOISE_DENSITY,
    pad_to_multiple_of=8,
)
print("Data collator: PairGraphCollator — on-the-fly positive sampling + Con adjacency matrix")
print(f"  Each batch: {BATCH_SIZE} drawn proteins → {BATCH_SIZE}–{2*BATCH_SIZE} unique proteins after positive sampling")
print("  Adjacency matrix built from VISEQ pair graph; orphan rows are all-False")

# ============================================================================
# CONTRASTIVE LOSS HELPERS
# ============================================================================


def mean_pool(hidden_states, attention_mask):
    """
    Mean-pool encoder last hidden states over non-padding token positions.
    hidden_states  : (N, seq_len, hidden_dim)
    attention_mask : (N, seq_len)   — 1 for real tokens, 0 for padding
    returns        : (N, hidden_dim)
    """
    mask = attention_mask.unsqueeze(-1).float()
    return (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)


def con_loss(z, adj, temperature):
    """
    Contrastive Loss (Con) using a prebuilt adjacency matrix.

    z   : (N, D) — unit-normalized encoder embeddings (F.normalize applied before calling)
    adj : (N, N) bool — adj[i,j] = True means i and j are known positives
    temperature : scalar

    For each protein i that has at least one positive in the batch:
      loss_i = -1/|P(i)| * sum_{j in P(i)} [sim(i,j)/T - log(sum_{k≠i} exp(sim(i,k)/T))]

    Returns the mean over proteins with at least one positive.
    Returns 0 (no grad) if no protein has a positive in the batch.
    """
    N   = z.size(0)
    dev = z.device

    # Full pairwise similarity matrix, temperature-scaled
    sim = torch.matmul(z, z.T) / temperature               # (N, N)

    # Mask diagonal so it does not contribute to the denominator
    self_mask  = torch.eye(N, dtype=torch.bool, device=dev)
    sim_masked = sim.masked_fill(self_mask, float("-inf"))

    # log-sum-exp over all k≠i  →  log denominator for each anchor i
    log_denom = torch.logsumexp(sim_masked, dim=1)         # (N,)

    # log p(j | i) = sim[i,j]/T - log_denom[i]   for each j
    log_prob  = sim - log_denom.unsqueeze(1)               # (N, N)

    # Number of positives per protein
    n_positives   = adj.float().sum(dim=1)                 # (N,)
    has_positive  = n_positives > 0                        # (N,) bool

    if not has_positive.any():
        return torch.tensor(0.0, device=dev, requires_grad=True)

    # Per-anchor loss: -1/|P(i)| * sum_{j: adj[i,j]} log_prob[i,j]
    pos_log_sum   = (adj.float() * log_prob).sum(dim=1)    # (N,)
    per_anchor    = -pos_log_sum / n_positives.clamp(min=1) # (N,)

    return per_anchor[has_positive].mean()



# ============================================================================
# CURRICULUM LOSS (applied to MLM component — identical to Default_v2)
# ============================================================================

curriculum_state = {"global_step": 0, "max_steps": 1, "phase": "train"}

# Shared state for component loss logging.
# Written by ContraMLMTrainer.compute_loss at every training step;
# read by ComponentLossLogCallback.on_log to inject into the Trainer log dict.
_component_losses: dict = {"mlm_loss": None, "con_loss": None}


def curriculum_keep_fraction(progress):
    current_stage = int(np.floor(progress * NUM_STAGES))
    if current_stage >= NUM_STAGES:
        return KEEP_FRACTION_END
    stage_size = (KEEP_FRACTION_END - KEEP_FRACTION_START) / (NUM_STAGES - 1)
    return KEEP_FRACTION_START + (current_stage * stage_size)


def curriculum_loss_from_outputs(outputs, labels):
    logits = outputs.logits
    token_losses = F.cross_entropy(
        logits.view(-1, logits.size(-1)),
        labels.view(-1),
        ignore_index=-100,
        reduction="none",
    )
    valid_tokens = labels.view(-1) != -100

    if USE_LOSS_CLIPPING_CURRICULUM and torch.is_grad_enabled():
        max_steps    = max(1, int(curriculum_state["max_steps"]))
        progress     = min(1.0, float(curriculum_state["global_step"]) / float(max_steps))
        keep_fraction = curriculum_keep_fraction(progress)

        valid_indices = torch.nonzero(valid_tokens, as_tuple=False).squeeze(-1)
        valid_losses  = token_losses[valid_tokens]
        k             = max(1, int(valid_losses.numel() * keep_fraction))
        selected_pos  = torch.topk(valid_losses, k=k, largest=LARGEST).indices
        keep_indices  = valid_indices[selected_pos]
        keep_tokens   = torch.zeros_like(valid_tokens, dtype=torch.bool)
        keep_tokens[keep_indices] = True
    else:
        keep_tokens = valid_tokens

    if not keep_tokens.any():
        raise FloatingPointError("No valid tokens available for loss computation.")

    return token_losses[keep_tokens].mean()


# ============================================================================
# CONFIGURE TRAINING ARGUMENTS
# ============================================================================

print("\n" + "=" * 80)
print("CONFIGURING TRAINING ARGUMENTS")
print("=" * 80)

dataloader_workers = min(8, os.cpu_count() or 1)

training_kwargs = {
    "output_dir":                    OUTPUT_DIR,
    "save_strategy":                 "steps",
    "eval_strategy":                 "steps",
    "save_steps":                    1000,
    "eval_steps":                    1000, # 5000
    "gradient_accumulation_steps":   GRADIENT_ACCUMULATION_STEPS,
    "per_device_train_batch_size":   BATCH_SIZE,
    "per_device_eval_batch_size":    16,
    "num_train_epochs":              NUM_EPOCHS,
    "dataloader_num_workers":        dataloader_workers,
    "dataloader_pin_memory":         True,
    "logging_dir":                   f"{OUTPUT_DIR}/logs",
    "logging_steps":                 100,
    "save_total_limit":              3,
    "fp16":                          False,
    "bf16":                          use_bf16,
    # REQUIRED: our batch dict uses custom keys (all_input_ids, etc.)
    # that are not in the model's forward signature.
    "remove_unused_columns":         False,
    "load_best_model_at_end":        False,
    "report_to":                     "none",
    "push_to_hub":                   False,
}

training_signature = inspect.signature(TrainingArguments.__init__).parameters
if "eval_strategy" not in training_signature:
    training_kwargs.pop("eval_strategy", None)
    training_kwargs["evaluation_strategy"] = "steps"

if "bf16_full_eval" in training_signature:
    training_kwargs["bf16_full_eval"] = use_bf16

training_args = TrainingArguments(**training_kwargs)
print("Training arguments configured")
print(f"  remove_unused_columns = False  (required for custom batch keys)")
print(
    f"  Effective batch: {BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS} drawn proteins "
    f"({BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS}–{2 * BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS} unique proteins per gradient step)"
)

# ============================================================================
# CALLBACKS
# ============================================================================

print("\n" + "=" * 80)
print("INITIALIZING CALLBACKS")
print("=" * 80)


class LiveLossPlotCallback(TrainerCallback):
    def __init__(self, output_dir):
        self.output_dir = output_dir

    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
        if metrics and "eval_loss" in metrics:
            print(f"\n>>> [Step {state.global_step}] Evaluation Loss: {metrics['eval_loss']:.4f} <<<\n")
        self._update_plot(state)

    def on_log(self, args, state, control, logs=None, **kwargs):
        self._update_plot(state)

    def _update_plot(self, state):
        history     = state.log_history
        train_loss  = [x["loss"]         for x in history if "loss"         in x]
        train_steps = [x["step"]         for x in history if "loss"         in x]
        eval_loss   = [x["eval_loss"]    for x in history if "eval_loss"    in x]
        eval_steps  = [x["step"]         for x in history if "eval_loss"    in x]
        mlm_loss    = [x["train_mlm_loss"]    for x in history if "train_mlm_loss"    in x]
        con_loss = [x["train_con_loss"] for x in history if "train_con_loss" in x]
        comp_steps  = [x["step"]         for x in history if "train_mlm_loss" in x]

        if not train_loss:
            return

        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 10), sharex=True)

        # Top panel: total loss + eval loss
        ax1.plot(train_steps, train_loss, label="Train Loss (total)",
                 color="blue", alpha=0.6)
        if eval_loss:
            ax1.plot(eval_steps, eval_loss, label="Eval Loss (MLM only)",
                     color="red", marker="o", linewidth=2)
        ax1.set_ylabel("Loss")
        ax1.grid(True, linestyle="--", alpha=0.6)
        ax1.legend(loc="upper right")
        ax1.set_title("Live Training Loss — ContraMLM v1")

        # Bottom panel: MLM vs Con components
        if mlm_loss:
            ax2.plot(comp_steps, mlm_loss,    label="MLM loss",    color="green",  alpha=0.7)
            ax2.plot(comp_steps, con_loss, label="Con loss", color="orange", alpha=0.7)
            ax2.legend(loc="upper right")
        ax2.set_xlabel("Training Steps")
        ax2.set_ylabel("Component Loss")
        ax2.grid(True, linestyle="--", alpha=0.6)

        plt.tight_layout()
        plt.savefig(os.path.join(self.output_dir, "live_loss_curve.png"), dpi=300)
        plt.close()


class CurriculumStateCallback(TrainerCallback):
    def on_train_begin(self, args, state, control, **kwargs):
        curriculum_state["global_step"] = state.global_step
        curriculum_state["max_steps"]   = state.max_steps if state.max_steps and state.max_steps > 0 else 1
        curriculum_state["phase"]       = "train"

    def on_step_begin(self, args, state, control, **kwargs):
        curriculum_state["global_step"] = state.global_step
        curriculum_state["phase"]       = "train"

    def on_evaluate(self, args, state, control, **kwargs):
        curriculum_state["phase"] = "eval"


# ============================================================================
# TRAINER
# ============================================================================

print("\n" + "=" * 80)
print("INITIALIZING TRAINER")
print("=" * 80)


class AdaFactorTrainer(Trainer):
    def create_optimizer_and_scheduler(self, num_training_steps: int):
        self.optimizer = Adafactor(
            [p for p in self.model.parameters() if p.requires_grad],
            scale_parameter=True,
            relative_step=True,
            warmup_init=True,
            lr=None,
        )
        self.lr_scheduler = AdafactorSchedule(self.optimizer)


class ContraMLMTrainer(AdaFactorTrainer):
    """
    Trainer combining BART-style MLM with Contrastive (Con) loss.

    Batch layout (set by PairGraphCollator):
        all_input_ids / all_attention_mask / all_labels : (N, L)
        adj_matrix : Python list[list[bool]] (N×N)
        N ranges from BATCH_SIZE (all orphans) to 2*BATCH_SIZE (all non-orphans).

    Loss:
        total = (1 - CONTRASTIVE_LAMBDA) * mlm_loss + CONTRASTIVE_LAMBDA * con_loss
        Con is skipped (returns 0) if no protein in the batch has a positive.
    """

    def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
        # adj_matrix is not a model argument; remove it before the standard eval forward pass
        inputs.pop("adj_matrix", None)
        # Remap custom collator keys to the standard model argument names
        if "all_input_ids" in inputs:
            inputs["input_ids"] = inputs.pop("all_input_ids")
        if "all_attention_mask" in inputs:
            inputs["attention_mask"] = inputs.pop("all_attention_mask")
        if "all_labels" in inputs:
            inputs["labels"] = inputs.pop("all_labels")
        return super().prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        # During training: custom keys (all_input_ids, etc.) + adj_matrix present.
        # During eval:     prediction_step remaps keys to standard names, adj_matrix is gone.
        adj_matrix         = inputs.get("adj_matrix", None)   # None during eval
        all_input_ids      = inputs.get("all_input_ids",      inputs.get("input_ids"))
        all_attention_mask = inputs.get("all_attention_mask", inputs.get("attention_mask"))
        all_labels         = inputs.get("all_labels",         inputs.get("labels"))

        # ---- Single encoder-decoder forward pass for all N proteins ----
        outputs = model(
            input_ids=all_input_ids,
            attention_mask=all_attention_mask,
            labels=all_labels,
        )

        # ---- MLM loss (curriculum-aware, same as Default_v2) ----
        mlm_loss = curriculum_loss_from_outputs(outputs, all_labels)

        # ---- Skip Con during eval (no adj_matrix) or all-orphan batch ----
        if adj_matrix is None:
            return (mlm_loss, outputs) if return_outputs else mlm_loss

        # ---- Convert adjacency list to bool tensor ----
        adj_tensor = torch.tensor(adj_matrix, dtype=torch.bool, device=all_input_ids.device)

        if not adj_tensor.any():
            _component_losses["mlm_loss"]    = mlm_loss.detach().float().item()
            _component_losses["con_loss"] = 0.0
            return (mlm_loss, outputs) if return_outputs else mlm_loss

        # ---- Encoder embeddings → mean-pooled, unit-normalized representations ----
        enc_hidden = outputs.encoder_last_hidden_state
        z = mean_pool(enc_hidden, all_attention_mask)

        # ---- Contrastive loss ----
        contrastive_loss = con_loss(
            F.normalize(z, dim=-1), adj_tensor, CONTRASTIVE_TEMPERATURE
        )

        _component_losses["mlm_loss"]    = mlm_loss.detach().float().item()
        _component_losses["con_loss"] = contrastive_loss.detach().float().item()

        total_loss = (1.0 - CONTRASTIVE_LAMBDA) * mlm_loss + CONTRASTIVE_LAMBDA * contrastive_loss
        return (total_loss, outputs) if return_outputs else total_loss

    def log(self, logs):
        # Inject component losses into the log dict BEFORE the base class freezes
        # it into state.log_history — this is the only way they appear in the history
        # that _update_plot reads.
        if _component_losses["mlm_loss"] is not None and "loss" in logs:
            logs["train_mlm_loss"]    = round(_component_losses["mlm_loss"],    6)
            logs["train_con_loss"] = round(_component_losses["con_loss"], 6)
        super().log(logs)


callbacks = [LiveLossPlotCallback(OUTPUT_DIR)]
if USE_LOSS_CLIPPING_CURRICULUM:
    callbacks.append(CurriculumStateCallback())

trainer = ContraMLMTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    data_collator=data_collator,
    callbacks=callbacks,
)

print("ContraMLMTrainer initialised")
print(f"  MLM loss weight:         {1.0 - CONTRASTIVE_LAMBDA:.2f}  (= 1 - CONTRASTIVE_LAMBDA)")
print(f"  Con loss weight:      {CONTRASTIVE_LAMBDA}")
print(f"  Contrastive temperature: {CONTRASTIVE_TEMPERATURE}")
print(f"  Orphans: adj row all-False → auto-excluded from Con mean, MLM only")

# ============================================================================
# TRAIN THE MODEL
# ============================================================================

print("\n" + "=" * 80)
print("STARTING TRAINING")
print("=" * 80)


def find_latest_checkpoint(output_dir):
    if not os.path.isdir(output_dir):
        return None
    latest_path = None
    latest_step = -1
    for entry in os.listdir(output_dir):
        if not entry.startswith("checkpoint-"):
            continue
        step_str = entry.split("checkpoint-")[-1]
        if not step_str.isdigit():
            continue
        full_path = os.path.join(output_dir, entry)
        if not os.path.isdir(full_path):
            continue
        step = int(step_str)
        if step > latest_step:
            latest_step = step
            latest_path = full_path
    return latest_path


def quarantine_rng_state_files(checkpoint_dir):
    moved_files = []
    for entry in os.listdir(checkpoint_dir):
        if not (entry.startswith("rng_state") and entry.endswith(".pth")):
            continue
        src = os.path.join(checkpoint_dir, entry)
        if not os.path.isfile(src):
            continue
        dst = src + ".bak"
        os.replace(src, dst)
        moved_files.append((src, dst))
    return moved_files


try:
    resume_checkpoint = find_latest_checkpoint(OUTPUT_DIR)
    if resume_checkpoint is not None:
        print(f"Resuming training from checkpoint: {resume_checkpoint}")
        moved_rng_files = quarantine_rng_state_files(resume_checkpoint)
        if moved_rng_files:
            print(
                f"Skipped rng_state*.pth files for PyTorch 2.6 compatibility: "
                f"{len(moved_rng_files)} file(s)."
            )
        train_result = trainer.train(resume_from_checkpoint=resume_checkpoint)
    else:
        print("No checkpoint found. Starting training from scratch.")
        train_result = trainer.train()

    print("\n" + "=" * 80)
    print("TRAINING COMPLETED!")
    print("=" * 80)
    print(f"Train loss: {train_result.training_loss:.4f}")
    print(f"Training time: {train_result.metrics['train_runtime']:.2f} seconds")

except Exception as e:
    print(f"\nERROR during training: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)

# ============================================================================
# EVALUATE THE MODEL
# ============================================================================

print("\n" + "=" * 80)
print("EVALUATING MODEL")
print("=" * 80)

try:
    eval_results = trainer.evaluate()
    print("Evaluation Results:")
    for key, value in eval_results.items():
        print(f"  {key}: {value:.4f}")
except Exception as e:
    print(f"ERROR during evaluation: {e}")

# ============================================================================
# SAVE THE MODEL
# ============================================================================

print("\n" + "=" * 80)
print("SAVING MODEL")
print("=" * 80)

lora_output_dir = f"{OUTPUT_DIR}/lora_adapters"
model.save_pretrained(lora_output_dir)
tokenizer.save_pretrained(lora_output_dir)

print(f"LoRA adapters saved to: {lora_output_dir}")

print("\n" + "=" * 80)
print("FINE-TUNING COMPLETE!")
print("=" * 80)