src/train.py

"""
Multimodal PC Fault Detection - Training Script
=================================================
Loads dataset from Ellaft/pc-fault-real-dataset (HuggingFace Hub).
Uses OGM-GE anti-modality-collapse + auxiliary heads + asymmetric LRs.

Usage:
  python train.py                                    # Full training
  python train.py --quick_test --no_push             # Quick 2-epoch test
  python train.py --eval_robustness                  # + missing modality test
  python train.py --mode visual_only --no_push       # Unimodal ablation
  python train.py --mode audio_only --no_push        # Unimodal ablation

All 6 ablations:
  python run_ablations.py
  python run_ablations.py --quick_test
"""

import os, sys, json, argparse, time
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.optim import AdamW
from torch.optim.lr_scheduler import OneCycleLR
from sklearn.metrics import accuracy_score, f1_score, confusion_matrix, precision_recall_fscore_support

from config import ExperimentConfig, FAULT_CLASSES, NUM_CLASSES
from models import create_model, get_processors, OGMGEModulator
from dataset_v2 import BuiltDataset as PCFaultDataset, multimodal_collate_fn


def compute_metrics(preds, labels, class_names=FAULT_CLASSES):
    accuracy = accuracy_score(labels, preds)
    precision, recall, f1, support = precision_recall_fscore_support(
        labels, preds, average=None, labels=range(len(class_names)), zero_division=0)
    macro_f1 = f1_score(labels, preds, average="macro", zero_division=0)
    weighted_f1 = f1_score(labels, preds, average="weighted", zero_division=0)
    conf_matrix = confusion_matrix(labels, preds, labels=range(len(class_names)))
    metrics = {"accuracy": accuracy, "macro_f1": macro_f1, "weighted_f1": weighted_f1,
               "confusion_matrix": conf_matrix.tolist(), "per_class": {}}
    for i, name in enumerate(class_names):
        metrics["per_class"][name] = {"precision": precision[i], "recall": recall[i],
            "f1": f1[i], "support": int(support[i])}
    return metrics


class MultimodalTrainer:
    def __init__(self, model, train_dataset, val_dataset, config, device,
                 use_ogm=True, ogm_alpha=0.3, ogm_noise_sigma=0.1,
                 visual_lr_multiplier=3.0, audio_lr_multiplier=0.5):
        self.model = model.to(device)
        self.device, self.config = device, config
        self.use_ogm = use_ogm and (model.mode == "multimodal")
        self.ogm = OGMGEModulator(alpha=ogm_alpha, noise_sigma=ogm_noise_sigma) if self.use_ogm else None
        if self.use_ogm:
            print(f"[Trainer] OGM-GE enabled: alpha={ogm_alpha}, noise_sigma={ogm_noise_sigma}")

        self.train_loader = DataLoader(train_dataset, batch_size=config.per_device_train_batch_size,
            shuffle=True, collate_fn=multimodal_collate_fn, num_workers=2, pin_memory=True, drop_last=True)
        self.val_loader = DataLoader(val_dataset, batch_size=config.per_device_eval_batch_size,
            shuffle=False, collate_fn=multimodal_collate_fn, num_workers=2, pin_memory=True)

        param_groups = self._get_param_groups(visual_lr_multiplier, audio_lr_multiplier)
        self.optimizer = AdamW(param_groups, weight_decay=config.weight_decay)
        total_steps = max(len(self.train_loader) * config.num_epochs // config.gradient_accumulation_steps, 1)
        self.scheduler = OneCycleLR(self.optimizer, max_lr=[pg["lr"] for pg in param_groups],
            total_steps=total_steps, pct_start=config.warmup_ratio, anneal_strategy="cos")
        self.scaler = torch.amp.GradScaler("cuda") if config.fp16 and device.type == "cuda" else None
        self.best_metric, self.best_epoch = 0.0, 0
        self.history = {"train_loss": [], "val_loss": [], "val_accuracy": [], "val_macro_f1": [],
            "train_loss_fusion": [], "train_loss_visual": [], "train_loss_audio": [],
            "ogm_visual_conf": [], "ogm_audio_conf": [], "ogm_coeff_visual": [], "ogm_coeff_audio": []}

    def _get_param_groups(self, visual_lr_mult, audio_lr_mult):
        visual_p, audio_p, fusion_p = [], [], []
        for name, p in self.model.named_parameters():
            if not p.requires_grad: continue
            if "visual_branch" in name: visual_p.append(p)
            elif "audio_branch" in name: audio_p.append(p)
            else: fusion_p.append(p)
        base_lr = self.config.lora_learning_rate
        groups = []
        if visual_p: groups.append({"params": visual_p, "lr": base_lr * visual_lr_mult, "name": "visual"})
        if audio_p: groups.append({"params": audio_p, "lr": base_lr * audio_lr_mult, "name": "audio"})
        if fusion_p: groups.append({"params": fusion_p, "lr": base_lr, "name": "fusion"})
        for g in groups: print(f"  [Param Group] {g.get('name','?')}: {len(g['params'])} tensors, lr={g['lr']:.2e}")
        if not groups: raise ValueError("No trainable parameters!")
        return groups

    def train_epoch(self, epoch):
        self.model.train()
        tot_loss, tot_fus, tot_vis, tot_aud, n = 0., 0., 0., 0., 0
        ogm_vc, ogm_ac, ogm_cv, ogm_ca = [], [], [], []
        self.optimizer.zero_grad()
        for bi, batch in enumerate(self.train_loader):
            pv, av, labels = batch["pixel_values"].to(self.device), batch["audio_values"].to(self.device), batch["labels"].to(self.device)
            if self.scaler:
                with torch.amp.autocast("cuda"):
                    out = self.model(pixel_values=pv, audio_values=av, labels=labels)
                    loss = out["loss"] / self.config.gradient_accumulation_steps
                self.scaler.scale(loss).backward()
            else:
                out = self.model(pixel_values=pv, audio_values=av, labels=labels)
                loss = out["loss"] / self.config.gradient_accumulation_steps
                loss.backward()
            tot_loss += loss.item() * self.config.gradient_accumulation_steps; n += 1
            if "loss_fusion" in out: tot_fus += out["loss_fusion"]; tot_vis += out["loss_visual"]; tot_aud += out["loss_audio"]
            if self.use_ogm and "visual_logits" in out:
                cv, ca, st = self.ogm.compute_modulation_coefficients(out["visual_logits"], out["audio_logits"], labels)
                ogm_vc.append(st["visual_conf"]); ogm_ac.append(st["audio_conf"]); ogm_cv.append(cv); ogm_ca.append(ca)
            if (bi + 1) % self.config.gradient_accumulation_steps == 0:
                if self.scaler: self.scaler.unscale_(self.optimizer)
                if self.use_ogm and ogm_cv: self.ogm.apply_gradient_modulation(self.model, ogm_cv[-1], ogm_ca[-1])
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.max_grad_norm)
                if self.scaler: self.scaler.step(self.optimizer); self.scaler.update()
                else: self.optimizer.step()
                self.scheduler.step(); self.optimizer.zero_grad()
            if (bi + 1) % self.config.logging_steps == 0 or bi == 0:
                msg = f"  [Epoch {epoch+1}] Step {bi+1}/{len(self.train_loader)} | Loss: {tot_loss/n:.4f}"
                if "loss_fusion" in out: msg += f" | fus:{tot_fus/n:.4f} vis:{tot_vis/n:.4f} aud:{tot_aud/n:.4f}"
                if ogm_cv: msg += f" | OGM cv:{ogm_cv[-1]:.3f} ca:{ogm_ca[-1]:.3f}"
                print(msg)
        stats = {"train_loss": tot_loss/max(n,1), "loss_fusion": tot_fus/max(n,1), "loss_visual": tot_vis/max(n,1), "loss_audio": tot_aud/max(n,1)}
        if ogm_vc: stats.update({"ogm_visual_conf": np.mean(ogm_vc), "ogm_audio_conf": np.mean(ogm_ac), "ogm_coeff_visual": np.mean(ogm_cv), "ogm_coeff_audio": np.mean(ogm_ca)})
        return stats

    @torch.no_grad()
    def evaluate(self, modality_mask=None):
        self.model.eval()
        preds, lbls, tot_loss, n = [], [], 0., 0
        for batch in self.val_loader:
            pv, av, labels = batch["pixel_values"].to(self.device), batch["audio_values"].to(self.device), batch["labels"].to(self.device)
            if modality_mask:
                if modality_mask.get("visual", 1.0) == 0.0: pv = torch.zeros_like(pv)
                if modality_mask.get("audio", 1.0) == 0.0: av = torch.zeros_like(av)
            out = self.model(pixel_values=pv, audio_values=av, labels=labels)
            tot_loss += out["loss"].item(); n += 1
            preds.extend(out["logits"].argmax(dim=-1).cpu().numpy()); lbls.extend(labels.cpu().numpy())
        m = compute_metrics(np.array(preds), np.array(lbls))
        m["val_loss"] = tot_loss / max(n, 1)
        return m

    def train(self):
        print(f"\n{'='*60}\nTraining: mode={self.model.mode}, epochs={self.config.num_epochs}, device={self.device}")
        print(f"OGM-GE: {'ON' if self.use_ogm else 'OFF'}")
        if self.model.mode == "multimodal": print(f"λ_visual={self.model.lambda_visual}, λ_audio={self.model.lambda_audio}")
        print(f"{'='*60}\n")
        for epoch in range(self.config.num_epochs):
            t0 = time.time()
            ts = self.train_epoch(epoch)
            vm = self.evaluate()
            print(f"\n[Epoch {epoch+1}/{self.config.num_epochs}] ({time.time()-t0:.1f}s)")
            msg = f"  Train Loss: {ts['train_loss']:.4f}"
            if ts.get("loss_fusion", 0) > 0: msg += f" (fus={ts['loss_fusion']:.4f} vis={ts['loss_visual']:.4f} aud={ts['loss_audio']:.4f})"
            print(msg)
            print(f"  Val Loss: {vm['val_loss']:.4f} | Acc: {vm['accuracy']:.4f} | F1: {vm['macro_f1']:.4f}")
            if "ogm_visual_conf" in ts: print(f"  OGM-GE: v_conf={ts['ogm_visual_conf']:.4f} a_conf={ts['ogm_audio_conf']:.4f} | cv={ts['ogm_coeff_visual']:.4f} ca={ts['ogm_coeff_audio']:.4f}")
            for k in ["train_loss","val_loss","val_accuracy","val_macro_f1","train_loss_fusion","train_loss_visual","train_loss_audio"]: self.history[k].append(ts.get(k, vm.get(k, 0)))
            if "ogm_visual_conf" in ts:
                for k in ["ogm_visual_conf","ogm_audio_conf","ogm_coeff_visual","ogm_coeff_audio"]: self.history[k].append(ts[k])
            if vm[self.config.metric_for_best_model] > self.best_metric:
                self.best_metric, self.best_epoch = vm[self.config.metric_for_best_model], epoch + 1
                os.makedirs(self.config.output_dir, exist_ok=True)
                torch.save({"model_state_dict": self.model.state_dict(), "epoch": epoch+1, "metrics": vm},
                    os.path.join(self.config.output_dir, "best_model.pt"))
                print(f"  ✓ Best model saved (F1={self.best_metric:.4f})")
        print(f"\nDone. Best epoch={self.best_epoch}, F1={self.best_metric:.4f}")
        return self.history

    def run_robustness_evaluation(self):
        print("\n=== Missing Modality Robustness ===")
        results = {}
        for name, mask in [("both", None), ("visual_only", {"visual": 1.0, "audio": 0.0}), ("audio_only", {"visual": 0.0, "audio": 1.0})]:
            m = self.evaluate(modality_mask=mask)
            results[name] = {"accuracy": m["accuracy"], "macro_f1": m["macro_f1"]}
            print(f"  {name:15s}: Acc={m['accuracy']:.4f} F1={m['macro_f1']:.4f}")
            for cls, cm in m["per_class"].items(): print(f"    {cls:25s} P:{cm['precision']:.3f} R:{cm['recall']:.3f} F1:{cm['f1']:.3f}")
        return results


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", default="multimodal", choices=["multimodal", "visual_only", "audio_only"])
    parser.add_argument("--finetune", default="lora", choices=["lora", "full", "linear_probe"])
    parser.add_argument("--epochs", type=int); parser.add_argument("--batch_size", type=int)
    parser.add_argument("--lr", type=float); parser.add_argument("--fusion", default="concat")
    parser.add_argument("--modality_dropout", type=float); parser.add_argument("--output_dir", type=str)
    parser.add_argument("--hub_model_id", type=str); parser.add_argument("--no_push", action="store_true")
    parser.add_argument("--eval_robustness", action="store_true"); parser.add_argument("--quick_test", action="store_true")
    parser.add_argument("--no_ogm", action="store_true")
    parser.add_argument("--ogm_alpha", type=float); parser.add_argument("--ogm_noise_sigma", type=float)
    parser.add_argument("--lambda_visual", type=float); parser.add_argument("--lambda_audio", type=float)
    parser.add_argument("--visual_lr_mult", type=float); parser.add_argument("--audio_lr_mult", type=float)
    args = parser.parse_args()

    config = ExperimentConfig()
    config.train.mode, config.train.finetune_method, config.model.fusion_type = args.mode, args.finetune, args.fusion
    if args.epochs: config.train.num_epochs = args.epochs
    if args.batch_size: config.train.per_device_train_batch_size = args.batch_size
    if args.lr: config.train.learning_rate = config.train.lora_learning_rate = args.lr
    if args.modality_dropout is not None: config.model.modality_dropout_p = args.modality_dropout
    if args.output_dir: config.train.output_dir = args.output_dir
    if args.hub_model_id: config.train.hub_model_id = args.hub_model_id
    if args.no_push: config.train.push_to_hub = False
    if args.quick_test:
        config.train.num_epochs, config.train.per_device_train_batch_size = 2, 4
        config.train.per_device_eval_batch_size, config.train.gradient_accumulation_steps = 4, 1
        config.train.logging_steps = 2
    if args.finetune != "lora": config.lora.enabled = False

    ogm_a = args.ogm_alpha if args.ogm_alpha is not None else config.ogm_alpha
    ogm_s = args.ogm_noise_sigma if args.ogm_noise_sigma is not None else config.ogm_noise_sigma
    lv = args.lambda_visual if args.lambda_visual is not None else config.lambda_visual
    la = args.lambda_audio if args.lambda_audio is not None else config.lambda_audio
    vlr = args.visual_lr_mult if args.visual_lr_mult is not None else config.visual_lr_multiplier
    alr = args.audio_lr_mult if args.audio_lr_mult is not None else config.audio_lr_multiplier

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.manual_seed(config.train.seed); np.random.seed(config.train.seed)

    print(f"\n{'='*60}\nMultimodal PC Fault Detection\n{'='*60}")
    print(f"Mode: {args.mode} | Finetune: {args.finetune} | Device: {device}")
    print(f"Dataset: {config.hub_dataset}")
    print(f"OGM-GE: {'OFF' if args.no_ogm else 'ON'} (α={ogm_a}, σ={ogm_s})")
    print(f"λ_visual={lv}, λ_audio={la}, LR mult: visual={vlr}x audio={alr}x")

    vit_proc, ast_ext = get_processors(config.model)
    train_ds = PCFaultDataset(config.data, config.model, "train", vit_proc, ast_ext, True,
        source="hub", hub_dataset=config.hub_dataset)
    val_ds = PCFaultDataset(config.data, config.model, "val", vit_proc, ast_ext, False,
        source="hub", hub_dataset=config.hub_dataset)

    model = create_model(config.model, config.lora, mode=args.mode, finetune_method=args.finetune,
        use_ogm=not args.no_ogm, lambda_visual=lv, lambda_audio=la)
    trainer = MultimodalTrainer(model, train_ds, val_ds, config.train, device,
        use_ogm=not args.no_ogm, ogm_alpha=ogm_a, ogm_noise_sigma=ogm_s,
        visual_lr_multiplier=vlr, audio_lr_multiplier=alr)

    history = trainer.train()
    final = trainer.evaluate()
    print(f"\nFinal: Acc={final['accuracy']:.4f} F1={final['macro_f1']:.4f}")
    for cls, m in final["per_class"].items(): print(f"  {cls:25s} P:{m['precision']:.3f} R:{m['recall']:.3f} F1:{m['f1']:.3f} N:{m['support']}")

    robustness = trainer.run_robustness_evaluation() if args.eval_robustness and args.mode == "multimodal" else None

    os.makedirs(config.train.output_dir, exist_ok=True)
    results = {"experiment": config.experiment_name, "mode": args.mode, "finetune": args.finetune,
        "dataset": config.hub_dataset, "ogm_ge": not args.no_ogm,
        "final_metrics": {"accuracy": final["accuracy"], "macro_f1": final["macro_f1"],
            "weighted_f1": final["weighted_f1"], "per_class": final["per_class"], "confusion_matrix": final["confusion_matrix"]},
        "history": history, "best_epoch": trainer.best_epoch, "best_metric": trainer.best_metric}
    if robustness: results["robustness"] = robustness
    with open(os.path.join(config.train.output_dir, "results.json"), "w") as f: json.dump(results, f, indent=2)

    if config.train.push_to_hub:
        try:
            from huggingface_hub import HfApi, login
            login(token=os.environ.get("HF_TOKEN"))
            HfApi().upload_folder(folder_path=config.train.output_dir, repo_id=config.train.hub_model_id,
                repo_type="model", commit_message=f"Training: {config.experiment_name}")
            print(f"✓ Pushed to https://huggingface.co/{config.train.hub_model_id}")
        except Exception as e: print(f"✗ Push failed: {e}")

if __name__ == "__main__":
    main()