main_retraining.py

import argparse
import json
import toml
import torch
import torch.nn as nn
import torch.optim as optim
import pickle
import os
import copy
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import sys
from torch.utils.data import DataLoader, Dataset, Subset, ConcatDataset

sys.path.append("/home/go68zas/Documents/coverage_guided_fuzzing_classifiers/")

from utils.util import obj
from fuzzer.model_structure import metrics_param, model_structure
from metrics_utils.metrics_dataloader import metrics_dataloader

class CustomEarlyStopping:
    def __init__(self, patience=5, min_delta=1e-4, crossover_threshold=1.10, verbose=True):
        self.patience = patience
        self.min_delta = min_delta
        self.crossover_threshold = crossover_threshold
        self.verbose = verbose
        
        self.crossover_reached = False
        self.val_loss_at_crossover = None
        self.best_train_loss = np.inf
        self.flatline_counter = 0
        self.early_stop = False
        
        # --- NEW: Track the reason ---
        self.stop_reason = "Max Epochs Reached" 

    def __call__(self, train_loss, val_loss, epoch):
        # Condition 1: Crossover Logic
        if not self.crossover_reached:
            if train_loss < val_loss:
                self.crossover_reached = True
                self.val_loss_at_crossover = val_loss
                if self.verbose:
                    print(f"\n[EarlyStop] Crossover detected at Epoch {epoch}...")
        else:
            limit = self.val_loss_at_crossover * self.crossover_threshold
            if val_loss > limit:
                # --- NEW: Set Reason ---
                self.stop_reason = "Crossover Limit Exceeded"
                self.early_stop = True
                if self.verbose: print(f"\n[EarlyStop] STOP: {self.stop_reason}")
                return True

        # Condition 2: Flatline Logic
        if train_loss < (self.best_train_loss - self.min_delta):
            self.best_train_loss = train_loss
            self.flatline_counter = 0
        else:
            self.flatline_counter += 1
            if self.flatline_counter >= self.patience:
                # --- NEW: Set Reason ---
                self.stop_reason = f"Training Flatlined ({self.patience} epochs)"
                self.early_stop = True
                if self.verbose: print(f"\n[EarlyStop] STOP: {self.stop_reason}")
                return True
        
        return False


def get_unseen_fuzz_dataset(full_fuzz_ds, used_indices):
    total_indices = set(range(len(full_fuzz_ds)))
    used_set = set(used_indices)
    unseen_indices = list(total_indices - used_set)
    return Subset(full_fuzz_ds, unseen_indices)


def load_prioritized_indices(config, method_name, tcp_alg, voxel_size=10, iteration=0):
    """
    Loads indices for a specific TCP algorithm.
    tcp_alg options: 'our_method', 'deep_gini', 'random'
    Filename format: {tcp_alg}_prio_list_{voxel_size}_{iteration}.pickle
    """
    filename = f"{tcp_alg}_prio_list_{voxel_size}_{iteration}.pickle"
    
    # Path: index_list/{dataset}/{fuzz_method}/{filename}
    fpath = os.path.join("index_list", config.data, method_name, filename)

    if not os.path.exists(fpath):
        # Fallback check: sometimes names might vary slightly, helpful for debugging
        raise FileNotFoundError(f"Prioritized index file not found for {tcp_alg}:\n{fpath}")

    with open(fpath, "rb") as f:
        data = pickle.load(f)
    
    # Usually data is a tuple or list, indices are often the first element
    # Adjust this if deep_gini/random pickle structure differs from our_method
    return data

# ===================================================================
# Helper: Evaluation & Loss
# ===================================================================
def evaluate_model(model, dataset, batch_size=256, device='cuda'):
    model.eval()
    if len(dataset) == 0: return 0.0
    
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
    correct = 0; total = 0
    
    with torch.no_grad():
        for inputs, classes in loader:
            inputs = inputs.to(device)
            labels = torch.tensor(
                [c[0] if isinstance(c, (list, np.ndarray)) else c for c in classes],
                dtype=torch.long, device=device
            )
            outputs, _ = model(inputs)
            preds = outputs.argmax(dim=1)
            correct += (preds == labels).sum().item()
            total += labels.size(0)
    return 100.0 * correct / total


def evaluate_loss(model, loader, criterion, device):
    model.eval()
    total_loss = 0.0
    with torch.no_grad():
        for inputs, classes in loader:
            inputs = inputs.to(device)
            labels = torch.tensor(
                [c[0] if isinstance(c, (list, np.ndarray)) else c for c in classes],
                dtype=torch.long, device=device
            )
            outputs, _ = model(inputs)
            loss = criterion(outputs, labels)
            total_loss += loss.item()
    return total_loss / len(loader)


def evaluate_all_sets(model, datasets, device, prefix=""):
    results = {}
    print(f"\n--- Evaluation ({prefix}) ---")
    for name, ds in datasets.items():
        acc = evaluate_model(model, ds, device=device)
        results[name] = acc
        print(f"{name:25s}: {acc:.2f}%")
    return results

# ===================================================================
# Plotting
# ===================================================================
def plot_retraining_metrics(train_losses, val_losses, val_accs, save_path):
    epochs = range(1, len(train_losses) + 1)
    plt.figure(figsize=(12, 5))
    
    plt.subplot(1, 2, 1)
    plt.plot(epochs, train_losses, 'b-', label="Train Loss")
    plt.plot(epochs, val_losses, 'r-', label="Val Loss")
    plt.title("Loss"); plt.xlabel("Epoch"); plt.legend(); plt.grid(True, linestyle='--', alpha=0.7)

    plt.subplot(1, 2, 2)
    plt.plot(epochs, val_accs, 'g-', label="Val Acc")
    plt.title("Accuracy"); plt.xlabel("Epoch"); plt.legend(); plt.grid(True, linestyle='--', alpha=0.7)

    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()


def main_retrain(model, config, prio_dataset, val_dataset,
                 save_name_prefix, # e.g. "mnist/nbc_ourmethod_S1"
                 batch_size=32, num_epochs=400, 
                 lr=1e-4,
                 # New params for the custom stopper
                 stopper_patience=5, stopper_min_delta=1e-4, stopper_crossover=1.10):

    print(f"\n========== RETRAINING: {os.path.basename(save_name_prefix)} ==========")
    device = config.device
    model.to(device)

    train_loader = DataLoader(prio_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)

    # Initialize Custom Early Stopper
    early_stopper = CustomEarlyStopping(
        patience=stopper_patience,
        min_delta=stopper_min_delta,
        crossover_threshold=stopper_crossover,
        verbose=True
    )
    
    train_losses, val_losses, val_accs = [], [], []
    final_epoch = 0

    # ================= TRAINING LOOP =================
    for epoch in range(num_epochs):
        final_epoch = epoch + 1
        model.train()
        running_loss = 0.0

        for inputs, classes in train_loader:
            inputs = inputs.to(device)
            # Handle list/array labels if necessary
            labels = torch.tensor(
                [c[0] if isinstance(c, (list, np.ndarray)) else c for c in classes],
                dtype=torch.long, device=device
            )
            
            optimizer.zero_grad()
            outputs, _ = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        # --- Metrics ---
        train_loss = running_loss / len(train_loader)
        val_loss = evaluate_loss(model, val_loader, criterion, device) 
        current_val_acc = evaluate_model(model, val_dataset, batch_size, device)
        
        train_losses.append(train_loss)
        val_losses.append(val_loss)
        val_accs.append(current_val_acc)

        print(f"[Epoch {epoch+1:02d}] Train={train_loss:.4f} | Val Loss={val_loss:.4f} | Val Acc={current_val_acc:.2f}%")

        # --- Logic: Early Stopping Only ---
        # No more safety restoration. We trust the stopper.
        if early_stopper(train_loss, val_loss, epoch):
            print(f"Stopping criteria met at Epoch {epoch+1}")
            break

    # ================= TEARDOWN =================
    # Save Model (Current state is the final state)
    save_dir = os.path.dirname(save_name_prefix)
    if save_dir:
        os.makedirs(save_dir, exist_ok=True)
    
    torch.save(model.state_dict(), f"{save_name_prefix}.pth")
    plot_retraining_metrics(train_losses, val_losses, val_accs, f"{save_name_prefix}.png")

    # Return stats for your Excel report
    return {
        "epochs": final_epoch,
        "reason": early_stopper.stop_reason
    }


# ===================================================================
# Main Execution
# ===================================================================
def parse_experiment_path(experiment_path):
    # exp_name = os.path.basename(experiment_path)
    parts = experiment_path.split("_")
    return parts[-4], parts[-2] 


def train_or_load(model, config, prio_dataset, val_dataset, save_name_prefix, lr, stopper_patience):
    """
    Checks if model exists. 
    If yes: Loads weights and returns placeholder stats.
    If no: Trains, saves weights, and returns actual stats.
    """
    pth_path = f"{save_name_prefix}.pth"
    
    if os.path.exists(pth_path):
        print(f"\n[INFO] Found existing model: {pth_path}")
        print("       Skipping training and loading weights...")
        
        # Load the existing weights
        model.load_state_dict(torch.load(pth_path, map_location=config.device))
        
        # Return placeholder stats so the report generation doesn't crash
        return {"epochs": "Existing", "reason": "Loaded from disk"}
    
    else:
        # Model doesn't exist, so we train it
        return main_retrain(
            model, config, prio_dataset, val_dataset, 
            save_name_prefix=save_name_prefix,
            lr=lr, 
            stopper_patience=stopper_patience 
        )


def arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument("-config_file", default="config/cifar100.toml")
    return parser.parse_args()

def get_fresh_model():
    if config.model not in model_structure: raise NotImplementedError
    m = model_structure[config.model]()
    m.load_state_dict(torch.load(m.model_path, map_location=device))
    m.to(device); m.eval()
    return m
    
if __name__ == "__main__":
    splits_to_analyze = ["cc_nc", "cc_kmnc", "cc_nbc", "cc_lsc", "cc_dsc", "cc_lscd"] #  "cc_kmnc", "cc_nbc", "cc_lsc", "cc_dsc", "cc_lscd"

    experiment_paths = {
        "cc_nc": "/srv/vivek/experiment_data_eq_samples/cifar100_1_nc_tp_full/", 
        "cc_kmnc": "/srv/vivek/experiment_data_eq_samples/cifar100_1_kmnc_tp_full/",
        "cc_nbc": "/srv/vivek/experiment_data_eq_samples/cifar100_1_nbc_tp_full/",
        "cc_lsc": "/srv/vivek/experiment_data_eq_samples/cifar100_1_lsc_tp_full/",
        "cc_dsc": "/srv/vivek/experiment_data_eq_samples/cifar100_1_dsc_tp_full/",
        "cc_lscd": "/srv/vivek/experiment_data_eq_samples/cifar100_1_lscd_tp_full/",
    }
    args = arguments()        
    config = json.loads(json.dumps(toml.load(args.config_file)), object_hook=obj)
    device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
    config.device = device

    final_results_matrix = []
 
    # 1. BASELINE 
    baseline_scores = {}
    model_base = get_fresh_model()
    clean_train_ds = metrics_dataloader(config, type="org", mode="train")
    clean_val_ds   = metrics_dataloader(config, type="org", mode="val")
    clean_test_ds  = metrics_dataloader(config, type="org", mode="test")
    
    base_acc_clean = evaluate_model(model_base, clean_test_ds, device=device)
    baseline_scores["test"] = base_acc_clean
    print(f"\n[BASELINE] Clean Test: {base_acc_clean:.2f}%")
    
    for eval_split in splits_to_analyze:
        eval_path = experiment_paths[eval_split]
        eval_ds = metrics_dataloader(config, type="crashes", mode=eval_split, experiment_path=eval_path)
        acc = evaluate_model(model_base, eval_ds, device=device)
        baseline_scores[eval_split] = acc
        print(f"Baseline Acc on {eval_split}: {acc:.2f}%")


    for train_split in splits_to_analyze:
        print(f"\n>>> [RETRAINING] on Split: {train_split.upper()}")
        experiment_path = experiment_paths[train_split]

        print("\n" + "#" * 80)
        print(f"Experiment: {experiment_path}")
        print("#" * 80)

        dataset_name, method_name = parse_experiment_path(experiment_path)

        # config.experiment_path = [experiment_path]
        config.input_size = (config.input_channels, config.input_width, config.input_height)

        # Dynamic Hyperparams
        lr = 1e-4; tolerance = 2.0; patience = 5
        if dataset_name == "svhn":
            lr = 1e-5; tolerance = 5.0

        # Load Standard Sets
        train_full_fuzz_ds  = metrics_dataloader(config, type="crashes", mode=train_split, experiment_path=experiment_path)

        # 2. STRATEGY 2: FULL FUZZ
        # Your code checks if it exists. If you have S2_Full, this will just load it (Fast).
        print("\n>>> Strategy 2 (Check/Load)...")
        model_s2 = get_fresh_model()
        s2_save_name = f"retrained_models/{config.data}/{method_name}_S2_Full"
        s2_stats = train_or_load(model_s2, config, train_full_fuzz_ds, clean_val_ds, 
                                 save_name_prefix=s2_save_name, lr=lr, stopper_patience=patience)
        
        s2_acc_clean = evaluate_model(model_s2, clean_test_ds, device=device)
        gain_clean_test = s2_acc_clean - baseline_scores["test"]
        print(f"Evaluating on Original Test Dataset | Acc: {s2_acc_clean:.2f}% | Gain: {gain_clean_test:+.2f}%")
        
        final_results_matrix.append({
            "Trained_On": train_split,
            "Tested_On": "Original Test Data",
            "Baseline_Acc": baseline_scores["test"],
            "Retrained_Acc": s2_acc_clean,
            "Accuracy_Gain": gain_clean_test,
            "Stop_Reason": s2_stats["reason"]
        })
        
        for eval_split in splits_to_analyze:
                # if eval_split == train_split:
                #     continue   
                eval_path = experiment_paths[eval_split]
                unseen_ds = metrics_dataloader(config, type="crashes", mode=eval_split, experiment_path=eval_path)
            
                retrained_acc = evaluate_model(model_s2, unseen_ds, device=device)
                
                # Calculation of Gain
                gain = retrained_acc - baseline_scores[eval_split]
                
                print(f"Evaluating on {eval_split:8s} | Acc: {retrained_acc:.2f}% | Gain: {gain:+.2f}%")
                
                final_results_matrix.append({
                    "Trained_On": train_split,
                    "Tested_On": eval_split,
                    "Baseline_Acc": baseline_scores[eval_split],
                    "Retrained_Acc": retrained_acc,
                    "Accuracy_Gain": gain,
                    "Stop_Reason": s2_stats["reason"]
                })

        # # 3. TCP LOOP (ONLY Our Method and DeepGini)
        # for tcp_alg in TCP_ALGORITHMS:
        #     print(f"\nProcessing TCP Algorithm: {tcp_alg.upper()}")

        #     # A. Load Indices
        #     try:
        #         prio_data = load_prioritized_indices(config, method_name, tcp_alg, voxel_size=10, iteration=0)
        #         prio_indices = prio_data[0]
        #     except FileNotFoundError as e:
        #         print(f"Skipping {tcp_alg}: {e}")
        #         continue

        #     # B. Create Subsets
        #     prio_ds = Subset(full_fuzz_ds, prio_indices)
            
        #     # --------------------------------------------------------------
        #     # Strategy 1: Prioritized Only
        #     # --------------------------------------------------------------
        #     s1_name = f"{method_name}_{tcp_alg}_S1_PrioOnly"
        #     s1_save_path = f"retrained_models/{config.data}/{s1_name}"
        #     model_s1 = get_fresh_model()
            
        #     # === CRITICAL: This will find your EXISTING S1 file and skip training ===
        #     print(f"Checking S1 for {tcp_alg}...")
        #     s1_stats = train_or_load(model_s1, config, prio_ds, clean_val_ds, 
        #                              save_name_prefix=s1_save_path, 
        #                              lr=lr, stopper_patience=patience)

        #     # --------------------------------------------------------------
        #     # Strategy 3: Mixed (THIS IS WHAT WILL TRAIN)
        #     # --------------------------------------------------------------
        #     # Construct Mixed Dataset
        #     num_replay = min(len(prio_ds), len(clean_train_ds))
        #     replay_idx = np.random.choice(len(clean_train_ds), size=num_replay, replace=False)
        #     clean_subset = Subset(clean_train_ds, replay_idx)
        #     mixed_ds = ConcatDataset([prio_ds, clean_subset])
                
        #     s3_name = f"{method_name}_{tcp_alg}_S3_Mixed"
        #     s3_save_path = f"retrained_models/{config.data}/{s3_name}"
        #     model_s3 = get_fresh_model()

        #     # === CRITICAL: This will NOT find the file and START training ===
        #     print(f"Starting S3 (Mixed) for {tcp_alg}...")
        #     s3_stats = train_or_load(model_s3, config, mixed_ds, clean_val_ds, 
        #                              save_name_prefix=s3_save_path, 
        #                              lr=lr, stopper_patience=patience)
    
    df = pd.DataFrame(final_results_matrix)
    df.to_csv(f"retrained_models/retraining_results_{config.data}.csv", index=False)       
    print("\n[DONE] Missing models trained.")