from core.finetuning import DPOTrainer, torch, pad_sequence, F

# ---------------------------------------------------------------
#  Add Adaptive Gradient Clipping to *every* parameter in-place
# ---------------------------------------------------------------
def attach_agc(model, clip: float = 0.01, eps: float = 1e-3):
    """
    Registers a per-parameter hook that applies Brock et al.’s
    Adaptive Gradient Clipping:

        ||g||₂ > clip * (||θ||₂ + eps)  →  g ← g * (threshold / ||g||₂)

    Works with params in fp32, bf16, or bitsandbytes int4.
    """

    def _agc_hook(grad, param):
        #print('agc hook')
        if grad is None:
            return grad
        param_norm = param.detach().norm()            # ||θ||
        grad_norm  = grad.norm()                      # ||g||
        max_norm   = clip * (param_norm + eps)
        if grad_norm > max_norm:
            grad = grad * (max_norm / (grad_norm + 1e-6))
        return grad

    for p in model.parameters():
        if p.requires_grad:
            p.register_hook(lambda g, p=p: _agc_hook(g, p))


# ---------------------------------------------------------------------
# Early-stopping on any logged scalar (loss, chosen_win, etc.)
# ---------------------------------------------------------------------
from transformers.trainer_callback import TrainerCallback


class ThresholdStop(TrainerCallback):
    """
    Stop training immediately when `monitor` crosses `threshold`.

    If `higher_is_better` is True  → stop when metric >= threshold.  
    If False                       → stop when metric <= threshold.
    """
    def __init__(self, monitor: str, threshold: float, higher_is_better: bool):
        self.monitor           = monitor
        self.threshold         = threshold
        self.higher_is_better  = higher_is_better

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is None or self.monitor not in logs:
            return
        value = logs[self.monitor]
        stop  = (value >= self.threshold) if self.higher_is_better else (value <= self.threshold)
        if stop:
            control.should_training_stop = True
            print(f"[ThresholdStop] {self.monitor}={value:.4f} "
                  f"crossed {'≥' if self.higher_is_better else '≤'} "
                  f"{self.threshold} – stopping.")
            
class EarlyStoppingByMetric(TrainerCallback):
    """
    Stop training when a monitored metric has stopped improving.

    Args
    ----
    monitor: str
        Key that appears in the `logs` dict (e.g. "loss", "chosen_win").
    higher_is_better: bool
        True  → metric should increase (e.g. chosen_win)  
        False → metric should decrease (e.g. loss / pref_loss)
    patience: int
        How many *log events* with no improvement to wait before stopping.
    min_delta: float
        Minimum change that counts as an improvement.
    """
    def __init__(self,
                 monitor: str,
                 higher_is_better: bool,
                 patience: int = 10,
                 min_delta: float = 0.0):
        self.monitor          = monitor
        self.higher_is_better = higher_is_better
        self.patience         = patience
        self.min_delta        = min_delta
        self.best             = None
        self.counter          = 0                     # events since last improv.

    # ── invoked every time trainer logs metrics ─────────────────────
    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is None or self.monitor not in logs:
            return

        current = logs[self.monitor]

        # first observation
        if self.best is None:
            self.best = current
            return

        # compute signed improvement
        if self.higher_is_better:
            improvement = current - self.best
        else:
            improvement = self.best - current

        # has the metric improved “enough”?
        if improvement > self.min_delta:
            self.best    = current
            self.counter = 0
        else:
            self.counter += 1
            if self.counter >= self.patience:
                #  signal the Trainer to halt
                control.should_training_stop = True
                print(f"[EarlyStopping] '{self.monitor}' plateaued "
                      f"(best={self.best:.5f}) – stopping training.")



class FTPOTrainer(DPOTrainer):
    """
    Trainer for final token preference optimisation (ftpo).
    Replaces TRL’s standard loss with a log-ratio on the **last**
    autoregressive position.
    """

    # ──────────────────────────────────────────────────────────────
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.remove_unused_columns = False
        self.data_collator = self.ftpo_collator                    # override

    # ──────────────────────────────────────────────────────────────
    @staticmethod
    def _get_proj(model):
        """
        Return the output-projection module in a model-agnostic way.
        Falls back to `lm_head` if `get_output_embeddings()` is None.
        """
        proj = model.get_output_embeddings()
        if proj is None:
            proj = getattr(model, "lm_head", None)
        if proj is None:
            raise AttributeError(
                "Model lacks both get_output_embeddings() and lm_head."
            )
        return proj

    # ──────────────────────────────────────────────────────────────
    def ftpo_collator(self, features):
        """
        Left-pads every prompt to `self.args.max_length`, so the last real
        token is always at position -1.  That lets the loss read logits
        with a single slice ([:, -1, :]).
        """
        pad_id   = self.padding_value
        max_len  = self.args.max_length
        batch_sz = len(features)

        # ── build [B, L] prompt tensor ───────────────────────────────
        prompt_ids   = torch.full((batch_sz, max_len), pad_id, dtype=torch.long)
        attention_ms = torch.zeros_like(prompt_ids, dtype=torch.bool)

        for i, feat in enumerate(features):
            seq = torch.tensor(feat["prompt_ids"], dtype=torch.long)
            if seq.size(0) > max_len:
                seq = seq[-max_len:]                    # truncate left if over-long
            prompt_ids[i, -seq.size(0):] = seq          # left-pad
            attention_ms[i, -seq.size(0):] = True

        # ── universal fields ─────────────────────────────────────────
        batch = dict(
            prompt_ids      = prompt_ids,
            attention_mask  = attention_ms,
            rejected_token_id = torch.tensor([f["rejected_token_id"] for f in features]),
        )

        # ── ftpo vs single-token branch ────────────────────────
        max_c = max(len(f["chosen_ids"]) for f in features)
        chosen_pad = torch.full((batch_sz, max_c), pad_id, dtype=torch.long)
        chosen_mask = torch.zeros_like(chosen_pad, dtype=torch.bool)
        for i, f in enumerate(features):
            ids = torch.tensor(f["chosen_ids"], dtype=torch.long)
            chosen_pad [i, :ids.size(0)] = ids
            chosen_mask[i, :ids.size(0)] = True
        batch.update(chosen_ids = chosen_pad,
                    chosen_mask = chosen_mask)

        return batch

    def compute_loss(self, model, inputs, return_outputs=False, **_):
        # We use 2 separate MSE loss terms
        # 1. A lightly applied tokenwise MSE loss applied to only the target tokens
        lambda_mse_target   = getattr(self, "lambda_mse_target", 0.05)  # strength
        tau_mse_target      = getattr(self, "tau_mse_target", 1.0)      # grace region (zero cost movement)

        # 2. A strongly applied tokenwise MSE loss applied to the remaining (non-target) vocab
        lambda_mse = getattr(self, "lambda_mse", 0.4) # how strongly the remaining vocab (other than chosen/rejected) is tethered to reference via mse loss

        # loss contribution is clipped if (chosen - rejected) logits delta is above this
        clip_epsilon_logits  = getattr(self, "clip_epsilon_logits", 2) 

        USE_MSE_LOSS=True # tether all the logits other than the ones we are interested in moving to the reference

        # ── unpack ---------------------------------------------------------
        device   = next(model.parameters()).device            # works for DP / DDP
        ids   = inputs["prompt_ids"].to(device)               # [B,L]
        attn  = inputs["attention_mask"].to(device)           # [B,L]
        B, L  = ids.shape
        
        seq_len  = attn.sum(1)
        pad_off  = (L - seq_len).unsqueeze(1)
        arange_L = torch.arange(L, device=ids.device).unsqueeze(0)
        pos_full = (arange_L - pad_off).clamp(min=0)
        pos_full = pos_full.masked_fill(attn == 0, 0)

        outputs = model(
            ids,
            attention_mask=attn,
            position_ids=pos_full,
            token_type_ids=torch.zeros_like(ids),
            use_cache=False,
            return_dict=True,
        )
        
        logits_last = outputs.logits[:, -1, :]  # [B, V]
        logp_all = F.log_softmax(logits_last, dim=-1)  # [B, V]

        ch_ids  = inputs["chosen_ids"].to(device)       
        ch_mask = inputs["chosen_mask"].to(device)      
        rejected     = inputs["rejected_token_id"].to(device)  
        logp_bad = logp_all.gather(-1, rejected.unsqueeze(-1)).squeeze(-1)  

        batch_rows = torch.arange(B, device=logp_all.device).unsqueeze(1)
        delta_tok = logits_last[batch_rows, ch_ids] - logits_last.gather(
            -1, rejected.unsqueeze(-1)
        )
        weights = (
            torch.clamp(
                (clip_epsilon_logits - delta_tok) / clip_epsilon_logits,
                0.0,
                1.0,
            )
            * ch_mask
        )

        tau = 1.0
        gap = clip_epsilon_logits - delta_tok
        per_tok_loss = F.softplus(gap / tau)

        chosen_counts = ch_mask.sum(dim=-1).clamp(min=1)
        pref_loss = ((per_tok_loss * weights).sum(dim=-1) / chosen_counts).mean()

        extra_metrics = {}

        if USE_MSE_LOSS:
            with torch.no_grad():
                if self.ref_model is None:
                    with self.null_ref_context():
                        ref_logits_last = model(
                            ids, attention_mask=attn, position_ids=pos_full,
                            token_type_ids=torch.zeros_like(ids),
                            use_cache=False, return_dict=True,
                        ).logits[:, -1, :]
                else:
                    ref_logits_last = self.ref_model(
                        ids, attention_mask=attn, position_ids=pos_full,
                        token_type_ids=torch.zeros_like(ids),
                        use_cache=False, return_dict=True,
                    ).logits[:, -1, :]

            tether_mask = torch.ones_like(logits_last, dtype=torch.bool)
            rows = torch.arange(B, device=ch_ids.device).unsqueeze(1).expand_as(ch_ids)
            tether_mask[rows[ch_mask], ch_ids[ch_mask]] = False
            tether_mask.scatter_(1, rejected.unsqueeze(-1), False)

            diff        = logits_last - ref_logits_last
            mse_elem_raw = (tether_mask * diff.pow(2)).sum() / tether_mask.sum()

            tgt_mask = torch.zeros_like(logits_last, dtype=torch.bool)
            rows = torch.arange(B, device=ch_ids.device).unsqueeze(1).expand_as(ch_ids)
            tgt_mask[rows[ch_mask], ch_ids[ch_mask]] = True
            tgt_mask.scatter_(1, rejected.unsqueeze(-1), True)
            
            if lambda_mse_target:
                diff_tok = logits_last - ref_logits_last          
                diff_tok = diff_tok * tgt_mask
                excess_tok = torch.clamp(diff_tok.abs() - tau_mse_target, min=0.0)
                mse_target_raw = (excess_tok.pow(2)).sum() / tgt_mask.sum()
            else:
                mse_target_raw = logits_last.new_tensor(0.0)

            mse_loss = (
                lambda_mse         * mse_elem_raw   
                + lambda_mse_target   * mse_target_raw
            )
            loss    = pref_loss + mse_loss

            extra_metrics.update({
                "mse_elem"        : mse_elem_raw.detach(),
                "mse_tgt_tokenwise"         : mse_target_raw.detach(),
            })

        else:
            loss = pref_loss

        lp_chosen = logp_all.gather(-1, ch_ids)               
        lp_bad    = logp_bad.unsqueeze(-1)                    

        wins_tok  = (lp_chosen > lp_bad) & ch_mask            
        frac_win  = wins_tok.float().sum(-1) / ch_mask.sum(-1).clamp(min=1e-8)  
        chosen_win = frac_win.mean().detach()                 

        active_delta = delta_tok[ch_mask]
        active_weights = weights[ch_mask]
        margin_win = (
            ((delta_tok >= clip_epsilon_logits) & ch_mask).float().sum()
            / ch_mask.float().sum().clamp(min=1e-8)
        ).detach()
        mean_delta = active_delta.mean().detach()
        median_delta = active_delta.median().detach()
        active_weight = active_weights.mean().detach()

        metrics = {
            "pref_loss":  pref_loss.detach(),
            "chosen_win": chosen_win,
            "margin_win": margin_win,
            "mean_delta": mean_delta,
            "median_delta": median_delta,
            "active_weight": active_weight,
            **extra_metrics,
        }
        self.store_metrics(metrics, train_eval="train")

        if return_outputs:
            return loss, metrics
        return loss




    # ----------------------------------------------------------
    def _prepare_dataset(self, dataset, *args, **_):
        return dataset
