"""
CODI + Direct Soft Matching (Ablation: No Prototypes)

Inherits from CODI_CT and replaces PMPSLatentSupervisor with
DirectMatchSupervisor. The forward() method is fully reused — only
the pmps_module is swapped at init time.

Usage:
    from src.model_direct_match import CODI_DirectMatch
    model = CODI_DirectMatch(model_args, training_args, lora_config)
"""

from .model import CODI_CT
from .direct_match_loss import DirectMatchSupervisor


class CODI_DirectMatch(CODI_CT):
    """CODI with direct soft matching (no prototypes).

    Ablation variant: replaces the prototype-mediated PMPS module with
    DirectMatchSupervisor, which computes soft many-to-many matching
    directly in the projected embedding space via cosine similarity.

    All other components (base model, LoRA, projector, PSA, diversity
    loss, distillation loss, CE losses) remain identical.
    """

    def __init__(self, model_args, training_args, lora_config):
        super().__init__(model_args, training_args, lora_config)

        if self.pmps_loss_factor > 0:
            del self.pmps_module
            self.pmps_module = DirectMatchSupervisor(
                hidden_dim=self.dim,
                feat_dim=training_args.pmps_feat_dim,
                hidden_mlp=training_args.pmps_hidden_mlp,
                match_temperature=training_args.pmps_match_temperature,
                use_projector=training_args.pmps_use_projector,
                position_bias_strength=training_args.pmps_position_bias_max,
            )