import os
import torch
import torch.nn.functional as F
import pandas as pd
import numpy as np
import xgboost as xgb
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, f1_score, precision_score, recall_score, confusion_matrix
import networkx as nx

# Setup Python path and import configs
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from src.config_loader import get_config
from src.ml.gnn_trainer import GraphSAGEFraudDetector

cfg = get_config()['ml']

def eval_xgboost():
    print("=== XGBoost Evaluation ===")
    import pickle
    with open("data/checkpoints/features.pkl", "rb") as f:
        df = pickle.load(f)
    feature_cols = [c for c in cfg['feature_cols'] if c in df.columns]
    
    X = df[feature_cols].copy()
    y = df['fraud_flag'].copy()
    
    X_temp, X_hold, y_temp, y_hold = train_test_split(
        X, y, test_size=cfg['val_size'], random_state=cfg['random_state'], stratify=y
    )
    
    model = xgb.XGBClassifier()
    model.load_model(cfg['model_path'])
    
    dtest = xgb.DMatrix(X_hold)
    y_proba = model.get_booster().predict(dtest)
    y_pred = (y_proba > 0.5).astype(int)
    
    print(f"AUC: {roc_auc_score(y_hold, y_proba):.4f}")
    print(f"F1:  {f1_score(y_hold, y_pred, zero_division=0):.4f}")
    print(f"Precision: {precision_score(y_hold, y_pred, zero_division=0):.4f}")
    print(f"Recall: {recall_score(y_hold, y_pred, zero_division=0):.4f}")
    print(f"Confusion Matrix:\n{confusion_matrix(y_hold, y_pred)}")

def eval_gnn():
    print("\n=== GraphSAGE Evaluation ===")
    import pickle
    with open("data/checkpoints/features.pkl", "rb") as f:
        df = pickle.load(f)
    feature_cols = [c for c in cfg['feature_cols'] if c in df.columns]
    
    # Load model to get x_mean and x_std
    checkpoint = torch.load(cfg['gnn_model_path'], map_location='cpu', weights_only=False)
    model = GraphSAGEFraudDetector(checkpoint['in_channels'], checkpoint['hidden_channels'])
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    
    x_mean = checkpoint['x_mean']
    x_std = checkpoint['x_std']
    
    node_list = sorted(list(df['account']))
    node_to_idx = {n: i for i, n in enumerate(node_list)}
    
    df_sorted = df.set_index('account').reindex(node_list)
    x_numpy = df_sorted[feature_cols].fillna(0).values.astype(float)
    y_numpy = df_sorted['fraud_flag'].fillna(0).values.astype(int)
    
    x_scaled = (x_numpy - x_mean) / (x_std + 1e-8)
    x_scaled = np.nan_to_num(x_scaled, nan=0.0)
    
    x = torch.tensor(x_scaled, dtype=torch.float)
    y = torch.tensor(y_numpy, dtype=torch.long)
    
    # Load edges
    edge_df = pd.read_parquet("data/checkpoints/graph_edges.parquet")
    src = edge_df['source'].map(node_to_idx)
    tgt = edge_df['target'].map(node_to_idx)
    valid = src.notna() & tgt.notna()
    src = src[valid].values.astype(int)
    tgt = tgt[valid].values.astype(int)
    all_src = np.concatenate([src, tgt])
    all_tgt = np.concatenate([tgt, src])
    edge_index = torch.tensor(np.stack([all_src, all_tgt]), dtype=torch.long)
    
    with torch.no_grad():
        out = model(x, edge_index)
        probs = F.softmax(out, dim=1)[:, 1].cpu().numpy()
        preds = (probs >= 0.5).astype(int)
    
    # Since GNN uses random split, we just evaluate on the whole graph to see overall performance
    print("Note: GNN split wasn't seeded/saved. Evaluating on full graph.")
    print(f"AUC: {roc_auc_score(y_numpy, probs):.4f}")
    print(f"F1:  {f1_score(y_numpy, preds, zero_division=0):.4f}")
    print(f"Precision: {precision_score(y_numpy, preds, zero_division=0):.4f}")
    print(f"Recall: {recall_score(y_numpy, preds, zero_division=0):.4f}")
    print(f"Confusion Matrix:\n{confusion_matrix(y_numpy, preds)}")

def eval_hybrid():
    print("\n=== Hybrid Stack Evaluation ===")
    from src.ml.hybrid_predictor import hybrid_predictor
    import pickle
    with open("data/checkpoints/features.pkl", "rb") as f:
        df = pickle.load(f)
    
    # Load graph
    edge_df = pd.read_parquet("data/checkpoints/graph_edges.parquet")
    edge_df['timestamp'] = pd.to_datetime(edge_df['timestamp'], errors='coerce').fillna(pd.Timestamp("2000-01-01"))
    G = nx.MultiDiGraph()
    for row in edge_df.itertuples(index=False):
        G.add_edge(row.source, row.target, amount=row.amount, timestamp=row.timestamp)
    
    # Let's evaluate on the holdout split from XGBoost (just as a representative test set)
    y = df['fraud_flag']
    _, df_hold = train_test_split(df, test_size=cfg['val_size'], random_state=cfg['random_state'], stratify=y)
    
    accounts = df_hold['account'].tolist()
    
    # Run hybrid prediction
    scores_dict = hybrid_predictor.predict(G, accounts)
    
    y_true = df_hold.set_index('account').loc[accounts]['fraud_flag'].values
    probs = np.array([scores_dict.get(a, 0.0) for a in accounts])
    preds = (probs >= 0.5).astype(int)
    
    print("Evaluating on XGBoost Holdout set.")
    print(f"AUC: {roc_auc_score(y_true, probs):.4f}")
    print(f"F1:  {f1_score(y_true, preds, zero_division=0):.4f}")
    print(f"Precision: {precision_score(y_true, preds, zero_division=0):.4f}")
    print(f"Recall: {recall_score(y_true, preds, zero_division=0):.4f}")
    print(f"Confusion Matrix:\n{confusion_matrix(y_true, preds)}")

def eval_ensemble():
    print("\n=== Blended Ensemble (XGB + GNN + Hybrid) ===")
    from src.ml.hybrid_predictor import hybrid_predictor
    import pickle
    with open("data/checkpoints/features.pkl", "rb") as f:
        df = pickle.load(f)
    
    # 1. XGBoost Probs
    feature_cols = [c for c in cfg['feature_cols'] if c in df.columns]
    X = df[feature_cols].copy()
    y = df['fraud_flag'].copy()
    _, X_hold, _, y_hold = train_test_split(X, y, test_size=cfg['val_size'], random_state=cfg['random_state'], stratify=y)
    
    xgb_model = xgb.XGBClassifier()
    xgb_model.load_model(cfg['model_path'])
    dtest = xgb.DMatrix(X_hold)
    xgb_probs = xgb_model.get_booster().predict(dtest)
    
    # 2. GNN Probs (Full graph mapped to holdout accounts)
    checkpoint = torch.load(cfg['gnn_model_path'], map_location='cpu', weights_only=False)
    gnn_model = GraphSAGEFraudDetector(checkpoint['in_channels'], checkpoint['hidden_channels'])
    gnn_model.load_state_dict(checkpoint['model_state_dict'])
    gnn_model.eval()
    
    node_list = sorted(list(df['account']))
    node_to_idx = {n: i for i, n in enumerate(node_list)}
    df_sorted = df.set_index('account').reindex(node_list)
    x_numpy = df_sorted[feature_cols].fillna(0).values.astype(float)
    x_scaled = (x_numpy - checkpoint['x_mean']) / (checkpoint['x_std'] + 1e-8)
    x_scaled = np.nan_to_num(x_scaled, nan=0.0)
    
    edge_df = pd.read_parquet("data/checkpoints/graph_edges.parquet")
    src = edge_df['source'].map(node_to_idx)
    tgt = edge_df['target'].map(node_to_idx)
    valid = src.notna() & tgt.notna()
    edge_index = torch.tensor(np.stack([np.concatenate([src[valid].values.astype(int), tgt[valid].values.astype(int)]), 
                                      np.concatenate([tgt[valid].values.astype(int), src[valid].values.astype(int)])]), dtype=torch.long)
    
    with torch.no_grad():
        gnn_out = gnn_model(torch.tensor(x_scaled, dtype=torch.float), edge_index)
        gnn_all_probs = F.softmax(gnn_out, dim=1)[:, 1].cpu().numpy()
    
    gnn_prob_dict = {n: p for n, p in zip(node_list, gnn_all_probs)}
    
    # 3. Hybrid Probs
    edge_df['timestamp'] = pd.to_datetime(edge_df['timestamp'], errors='coerce').fillna(pd.Timestamp("2000-01-01"))
    G = nx.MultiDiGraph()
    for row in edge_df.itertuples(index=False):
        G.add_edge(row.source, row.target, amount=row.amount, timestamp=row.timestamp)
    
    _, df_hold = train_test_split(df, test_size=cfg['val_size'], random_state=cfg['random_state'], stratify=y)
    accounts = df_hold['account'].tolist()
    hybrid_scores_dict = hybrid_predictor.predict(G, accounts)
    
    # Map all to the holdout set
    gnn_probs = np.array([gnn_prob_dict.get(a, 0.0) for a in accounts])
    hybrid_probs = np.array([hybrid_scores_dict.get(a, 0.0) for a in accounts])
    y_true = df_hold.set_index('account').loc[accounts]['fraud_flag'].values
    
    # 4. Blend using the server.py calibration weights (normalized to sum to 1 since we don't have orig heuristic score here)
    # Weights in server.py: XGB(0.20), GNN(0.15), Hybrid(0.25) -> Total ML weight = 0.60
    # Normalized weights: XGB(0.333), GNN(0.25), Hybrid(0.417)
    blended_probs = (xgb_probs * (0.20/0.60)) + (gnn_probs * (0.15/0.60)) + (hybrid_probs * (0.25/0.60))
    blended_preds = (blended_probs >= 0.5).astype(int)
    
    print("Evaluating Ensemble (XGB 33%, GNN 25%, Hybrid 42%) on XGBoost Holdout set.")
    print(f"AUC: {roc_auc_score(y_true, blended_probs):.4f}")
    print(f"F1:  {f1_score(y_true, blended_preds, zero_division=0):.4f}")
    print(f"Precision: {precision_score(y_true, blended_preds, zero_division=0):.4f}")
    print(f"Recall: {recall_score(y_true, blended_preds, zero_division=0):.4f}")
    print(f"Confusion Matrix:\n{confusion_matrix(y_true, blended_preds)}")

if __name__ == "__main__":
    eval_ensemble()