import pandas as pd
import numpy as np
import torch
import os
from sklearn.metrics import precision_recall_curve, f1_score, auc, recall_score, precision_score
from train_hmoc_ensemble import Spectra1DCNN, pad_or_truncate
from sklearn.model_selection import StratifiedKFold

DATA_PATH = "./data/cda_qm_spectra_pre2008277_lvl3_train.parquet"
MODELS_DIR = "./hmoc_binary_classifier/models"
METRICS_OUT = "./hmoc_binary_classifier/hmoc_metrics.parquet"
TARGET_LEN = 1000

def get_best_threshold_and_metrics(y_true, y_probs):
    precision, recall, thresholds = precision_recall_curve(y_true, y_probs)
    # Avoid div by zero in f1 calculation
    f1_scores = 2 * (precision * recall) / (precision + recall + 1e-8)
    best_idx = np.argmax(f1_scores)
    best_thresh = thresholds[best_idx] if best_idx < len(thresholds) else 0.5
    
    # Alternatively find recall at 95% precision if we want
    # For general "best F1"
    final_preds = (y_probs >= best_thresh).astype(int)
    
    return {
        "Best_Threshold": best_thresh,
        "Precision": precision_score(y_true, final_preds, zero_division=0),
        "Recall": recall_score(y_true, final_preds, zero_division=0),
        "F1_Score": f1_score(y_true, final_preds, zero_division=0),
        "AUPRC": auc(recall, precision)
    }

def main():
    print("Loading labeled training data for evaluation...")
    df = pd.read_parquet(DATA_PATH)
    exclude_classes = ['?', 'Noise']
    df = df[~df['class'].isin(exclude_classes)].copy()
    df['label'] = (df['class'] == 'HMOC').astype(int)
    
    X = np.stack(df['spectrum'].apply(lambda s: pad_or_truncate(np.array(s), TARGET_LEN)).values)
    X_min = X.min(axis=1, keepdims=True)
    X_max = X.max(axis=1, keepdims=True)
    X = (X - X_min) / (X_max - X_min + 1e-8)
    y = df['label'].values
    
    X_tensor = torch.tensor(X, dtype=torch.float32).unsqueeze(1)
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(device)
    
    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    ensemble_seeds = [42, 100, 2026, 777, 888]
    
    model_records = []
    
    for seed_idx, seed in enumerate(ensemble_seeds):
        fold = seed_idx % 5
        _, val_idx = list(skf.split(X, y))[fold]
        
        X_val, y_val = X_tensor[val_idx], y[val_idx]
        
        # Load Model
        model_path = os.path.join(MODELS_DIR, f"hmoc_model_seed_{seed}.pth")
        model = Spectra1DCNN(TARGET_LEN)
        model.load_state_dict(torch.load(model_path, map_location=device))
        model.to(device)
        model.eval()
        
        with torch.no_grad():
            outputs = torch.sigmoid(model(X_val.to(device))).cpu().numpy().flatten()
            
        metrics = get_best_threshold_and_metrics(y_val, outputs)
        
        record = {
            "Model_Name": f"Spectra1DCNN_Fold{fold}_Seed{seed}",
            "Architecture": "3x Conv1D (32->64->128) + GAP + FC(64->1) + Dropout(0.4)",
            "Loss_Function": "FocalLoss(alpha=0.85, gamma=2.0)",
            "Optimizer": "AdamW(lr=1e-3, weight_decay=1e-4)",
            "Batch_Size": 64,
            "Sequence_Length": TARGET_LEN,
            "Validation_Set_Size": len(y_val),
            **metrics
        }
        model_records.append(record)
        
    # Aggregate ensemble metrics (Full pseudo out-of-fold predictions)
    # Note: simple unweighted average for now over the folds
    
    metrics_df = pd.DataFrame(model_records)
    metrics_df.to_parquet(METRICS_OUT, index=False)
    print(f"Metrics saved to {METRICS_OUT}")
    print(metrics_df)

if __name__ == "__main__":
    main()