main.py

"""
Genomic Sequence Classification Pipeline: Host vs. Pathogen (Zero-Shot)
Author: Gemini CLI Agent
Role: Expert Bioinformatician & PyTorch ML Engineer
"""

import os
import torch
import numpy as np
import argparse
from data_utils import DNADataProcessor, load_split_data, prepare_datasets
from models import XGBoostBaseline, PureCNN, CNNBiLSTM
from trainer import train_torch_model, measure_inference_latency, get_predictions
from visualizer import generate_diagnostic_panel

def main():
    # --- Configuration ---
    parser = argparse.ArgumentParser(description="Genomic Sequence Classification Pipeline")
    parser.add_argument("--train", action="store_true", help="Force retraining of models even if weights exist")
    parser.add_argument("--epochs", type=int, default=20, help="Number of training epochs")
    parser.add_argument("--mutation_rate", type=float, default=0.05, help="Rate of dynamic mutation injection")
    args = parser.parse_args()

    BASE_DIR = "/home/louai/Documents/pathogen_detection/data"
    WEIGHTS_DIR = "weights"
    WINDOW_SIZE = 250
    STRIDE = 125
    BATCH_SIZE = 64
    EPOCHS = args.epochs
    MUTATION_RATE = args.mutation_rate
    
    os.makedirs(WEIGHTS_DIR, exist_ok=True)
    
    print("Initializing DNA Data Processor...")
    processor = DNADataProcessor(window_size=WINDOW_SIZE, stride=STRIDE)
    
    # --- Step 1: Data Loading & Chunking ---
    print("\n--- Step 1: Data Loading & Synthetic Chunking ---")
    splits = load_split_data(BASE_DIR, processor)
    
    # --- Step 2: Encoding ---
    print("\n--- Step 2: Data Encoding (Tabular & Spatial) ---")
    data = prepare_datasets(splits, processor)
    
    X_tab_train, y_train = data['X_tabular_train'], data['y_train']
    X_tab_test, y_test = data['X_tabular_test'], data['y_test']
    X_spa_train = data['X_spatial_train']
    X_spa_test = data['X_spatial_test']
    
    print(f"Train samples: {len(y_train)}, Test samples: {len(y_test)}")
    
    # Storage for results
    results = {}
    
    # --- Step 3: Core Architectures & Training/Loading ---
    print("\n--- Step 3: Loading or Training 3-Tier Spectrum Models ---")
    
    # Model A: XGBoost Baseline
    model_a_path = os.path.join(WEIGHTS_DIR, "xgboost_baseline.json")
    model_a = XGBoostBaseline()
    
    num_pos = np.sum(y_train == 1)
    num_neg = np.sum(y_train == 0)
    scale_pos_weight = num_neg / num_pos if num_pos > 0 else 1.0
    
    if os.path.exists(model_a_path) and not args.train:
        print(f"Loading Model A from {model_a_path}...")
        model_a.load(model_a_path)
    else:
        print("Training Model A: XGBoost Baseline...")
        model_a.train(X_tab_train, y_train, scale_pos_weight=scale_pos_weight)
        model_a.save(model_a_path)
        print(f"Model A saved to {model_a_path}")
    
    # Model B: Pure Convolutional
    model_b_path = os.path.join(WEIGHTS_DIR, "pure_cnn.pth")
    model_b = PureCNN()
    if os.path.exists(model_b_path) and not args.train:
        print(f"Loading Model B from {model_b_path}...")
        model_b.load(model_b_path)
    else:
        print(f"Training Model B: Pure CNN (1D) with mutation_rate={MUTATION_RATE}...")
        time_b = train_torch_model(model_b, X_spa_train, y_train, epochs=EPOCHS, batch_size=BATCH_SIZE, mutation_rate=MUTATION_RATE)
        print(f"Model B Training Time: {time_b:.2f}s")
        model_b.save(model_b_path)
        print(f"Model B saved to {model_b_path}")
    
    # Model C: Recurrent Hybrid
    model_c_path = os.path.join(WEIGHTS_DIR, "cnn_bilstm.pth")
    model_c = CNNBiLSTM()
    if os.path.exists(model_c_path) and not args.train:
        print(f"Loading Model C from {model_c_path}...")
        model_c.load(model_c_path)
    else:
        print(f"Training Model C: CNN-BiLSTM with mutation_rate={MUTATION_RATE}...")
        time_c = train_torch_model(model_c, X_spa_train, y_train, epochs=EPOCHS, batch_size=BATCH_SIZE, mutation_rate=MUTATION_RATE)
        print(f"Model C Training Time: {time_c:.2f}s")
        model_c.save(model_c_path)
        print(f"Model C saved to {model_c_path}")
    
    # --- Step 4: Zero-Shot Evaluation & Latency Tracking ---
    print("\n--- Step 4: Zero-Shot Evaluation & Latency Tracking ---")
    
    # Evaluate Model A
    results['XGBoost'] = {
        'probs': get_predictions(model_a, X_tab_test, is_torch=False),
        'latency': measure_inference_latency(model_a, X_tab_test, is_torch=False)
    }
    
    # Evaluate Model B
    results['PureCNN'] = {
        'probs': get_predictions(model_b, X_spa_test, is_torch=True),
        'latency': measure_inference_latency(model_b, X_spa_test, is_torch=True)
    }
    
    # Evaluate Model C
    results['CNN-BiLSTM'] = {
        'probs': get_predictions(model_c, X_spa_test, is_torch=True),
        'latency': measure_inference_latency(model_c, X_spa_test, is_torch=True)
    }
    
    for name, res in results.items():
        print(f"{name} -> Inference Latency (10k): {res['latency']:.2f} ms")

    # --- Step 5: Multi-Model Visualization ---
    print("\n--- Step 5: Generating Multi-Model Diagnostic Panel ---")
    generate_diagnostic_panel(
        results=results,
        X_spatial_test=X_spa_test,
        y_test=y_test,
        model_c=model_c,
        save_path="genomics_diagnostic_panel.png"
    )

if __name__ == "__main__":
    main()