import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.ensemble import IsolationForest, RandomForestClassifier
from sklearn.metrics import precision_recall_fscore_support, roc_auc_score
from sklearn.preprocessing import StandardScaler
import xgboost as xgb
from collections import defaultdict
import warnings
warnings.filterwarnings('ignore')

class FraudDataGenerator:
    """Generate synthetic fraud transaction data with realistic properties."""
    
    def __init__(self, n_transactions=100000, fraud_rate=0.005):
        self.n = n_transactions
        self.fraud_rate = fraud_rate
        
    def generate(self):
        np.random.seed(42)
        n_fraud = int(self.n * self.fraud_rate)
        n_legit = self.n - n_fraud
        
        legit = {
            'amount': np.random.lognormal(4.5, 1.2, n_legit),
            'hour': np.random.choice(range(24), n_legit, p=self._hour_weights(False)),
            'day_of_week': np.random.randint(0, 7, n_legit),
            'merchant_category': np.random.randint(0, 20, n_legit),
            'distance_from_home': np.random.exponential(15, n_legit),
            'num_transactions_24h': np.random.poisson(3, n_legit),
            'avg_amount_30d': np.random.lognormal(4.2, 0.8, n_legit),
            'account_age_days': np.random.exponential(365, n_legit),
            'is_international': np.random.binomial(1, 0.05, n_legit),
            'is_card_not_present': np.random.binomial(1, 0.3, n_legit),
            'fraud': np.zeros(n_legit, dtype=int)
        }
        
        fraud = {
            'amount': np.random.lognormal(6.5, 1.5, n_fraud),
            'hour': np.random.choice(range(24), n_fraud, p=self._hour_weights(True)),
            'day_of_week': np.random.randint(0, 7, n_fraud),
            'merchant_category': np.random.randint(0, 20, n_fraud),
            'distance_from_home': np.random.exponential(150, n_fraud),
            'num_transactions_24h': np.random.poisson(8, n_fraud),
            'avg_amount_30d': np.random.lognormal(4.2, 0.8, n_fraud),
            'account_age_days': np.random.exponential(180, n_fraud),
            'is_international': np.random.binomial(1, 0.4, n_fraud),
            'is_card_not_present': np.random.binomial(1, 0.8, n_fraud),
            'fraud': np.ones(n_fraud, dtype=int)
        }
        
        df_legit = pd.DataFrame(legit)
        df_fraud = pd.DataFrame(fraud)
        df = pd.concat([df_legit, df_fraud]).reset_index(drop=True)
        
        return df.sample(frac=1, random_state=42).reset_index(drop=True)
    
    def _hour_weights(self, is_fraud):
        hours = np.zeros(24)
        if is_fraud:
            hours[0:6] = 0.15
            hours[6:12] = 0.08
            hours[12:18] = 0.1
            hours[18:24] = 0.17
        else:
            hours[0:6] = 0.02
            hours[6:12] = 0.12
            hours[12:18] = 0.15
            hours[18:24] = 0.11
        return hours / hours.sum()

class IsolationForestDetector:
    """Isolation Forest for unsupervised anomaly detection."""
    
    def __init__(self, contamination=0.01, n_estimators=100):
        self.model = IsolationForest(
            contamination=contamination,
            n_estimators=n_estimators,
            random_state=42
        )
        self.scaler = StandardScaler()
        
    def fit(self, X):
        X_scaled = self.scaler.fit_transform(X)
        self.model.fit(X_scaled)
        return self
    
    def predict(self, X):
        X_scaled = self.scaler.transform(X)
        scores = -self.model.score_samples(X_scaled)
        return scores
    
    def get_anomalies(self, X, threshold=None):
        scores = self.predict(X)
        if threshold is None:
            threshold = np.percentile(scores, 99)
        return scores > threshold, scores

class TransactionGNN(nn.Module):
    """Graph Neural Network for fraud detection on transaction graphs."""
    
    def __init__(self, input_dim=16, hidden_dim=64, n_layers=3):
        super().__init__()
        self.layers = nn.ModuleList()
        self.layers.append(nn.Linear(input_dim, hidden_dim))
        for _ in range(n_layers - 1):
            self.layers.append(nn.Linear(hidden_dim, hidden_dim))
        
        self.attention = nn.MultiheadAttention(hidden_dim, num_heads=4, batch_first=True)
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim // 2, 1),
            nn.Sigmoid()
        )
        
    def forward(self, x, edge_index=None):
        h = x
        for i, layer in enumerate(self.layers):
            h = layer(h)
            if i < len(self.layers) - 1:
                h = F.relu(h)
                h = F.dropout(h, p=0.2, training=self.training)
        
        if edge_index is not None:
            h = h.unsqueeze(0)
            h, _ = self.attention(h, h, h)
            h = h.squeeze(0)
        
        return self.classifier(h)

class AutoencoderDetector(nn.Module):
    """Variational Autoencoder for anomaly detection."""
    
    def __init__(self, input_dim=10, latent_dim=4):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 32),
            nn.ReLU(),
            nn.Linear(32, 16),
            nn.ReLU()
        )
        self.mu = nn.Linear(16, latent_dim)
        self.logvar = nn.Linear(16, latent_dim)
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 16),
            nn.ReLU(),
            nn.Linear(16, 32),
            nn.ReLU(),
            nn.Linear(32, input_dim)
        )
        
    def encode(self, x):
        h = self.encoder(x)
        return self.mu(h), self.logvar(h)
    
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def decode(self, z):
        return self.decoder(z)
    
    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        recon = self.decode(z)
        return recon, mu, logvar
    
    def anomaly_score(self, x):
        recon, mu, logvar = self.forward(x)
        recon_loss = F.mse_loss(recon, x, reduction='none').sum(dim=1)
        kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1)
        return recon_loss + kl_loss

class StreamingFraudDetector:
    """Online learning fraud detector with concept drift detection."""
    
    def __init__(self, window_size=1000, drift_threshold=0.05):
        self.window_size = window_size
        self.drift_threshold = drift_threshold
        self.score_buffer = []
        self.label_buffer = []
        self.model = RandomForestClassifier(
            n_estimators=50, max_depth=10, random_state=42
        )
        self.is_fitted = False
        self.drift_detector = ADWINDriftDetector()
        
    def partial_fit(self, X_batch, y_batch):
        self.score_buffer.extend(X_batch.tolist())
        self.label_buffer.extend(y_batch.tolist())
        
        if len(self.score_buffer) > self.window_size * 2:
            self.score_buffer = self.score_buffer[-self.window_size * 2:]
            self.label_buffer = self.label_buffer[-self.window_size * 2:]
        
        if len(self.score_buffer) >= self.window_size:
            X_train = np.array(self.score_buffer[-self.window_size:])
            y_train = np.array(self.label_buffer[-self.window_size:])
            self.model.fit(X_train, y_train)
            self.is_fitted = True
    
    def predict(self, X):
        if not self.is_fitted:
            return np.zeros(len(X))
        return self.model.predict_proba(X)[:, 1]
    
    def detect_drift(self, recent_scores):
        return self.drift_detector.detect(recent_scores)

class ADWINDriftDetector:
    """ADWIN concept drift detector."""
    
    def __init__(self, delta=0.002):
        self.delta = delta
        self.window = []
        
    def detect(self, new_value):
        self.window.append(new_value)
        
        if len(self.window) < 30:
            return False
        
        n = len(self.window)
        mid = n // 2
        
        w0 = np.array(self.window[:mid])
        w1 = np.array(self.window[mid:])
        
        mu0, mu1 = np.mean(w0), np.mean(w1)
        n0, n1 = len(w0), len(w1)
        
        epsilon = np.sqrt((1.0 / (2.0 * min(n0, n1))) * np.log(4.0 / self.delta))
        
        if abs(mu0 - mu1) >= epsilon:
            self.window = self.window[mid:]
            return True
        
        return False

class FraudEnsemble:
    """Ensemble fraud detector combining multiple models."""
    
    def __init__(self):
        self.models = {}
        self.weights = {}
        
    def add_model(self, name, model, weight=1.0):
        self.models[name] = model
        self.weights[name] = weight
        
    def predict(self, X):
        predictions = {}
        for name, model in self.models.items():
            if hasattr(model, 'predict_proba'):
                pred = model.predict_proba(X)[:, 1]
            else:
                pred = model.predict(X)
            predictions[name] = pred
        
        weighted_sum = np.zeros(len(X))
        total_weight = sum(self.weights.values())
        
        for name, pred in predictions.items():
            weighted_sum += pred * self.weights[name] / total_weight
        
        return weighted_sum
    
    def evaluate(self, X, y_true):
        ensemble_pred = self.predict(X)
        ensemble_auc = roc_auc_score(y_true, ensemble_pred)
        
        model_aucs = {}
        for name, model in self.models.items():
            if hasattr(model, 'predict_proba'):
                pred = model.predict_proba(X)[:, 1]
            else:
                pred = model.predict(X)
            model_aucs[name] = roc_auc_score(y_true, pred)
        
        return {
            'ensemble_auc': ensemble_auc,
            'model_aucs': model_aucs
        }

def train_vae_detector(model, X_normal, epochs=50, lr=0.001):
    """Train VAE on normal transactions only."""
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    X_tensor = torch.FloatTensor(X_normal)
    
    for epoch in range(epochs):
        recon, mu, logvar = model(X_tensor)
        
        recon_loss = F.mse_loss(recon, X_tensor, reduction='sum')
        kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        loss = recon_loss + kl_loss
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch+1}/{epochs}, Loss: {loss.item():.2f}")
    
    return model

# Example usage
if __name__ == "__main__":
    generator = FraudDataGenerator(n_transactions=50000, fraud_rate=0.005)
    data = generator.generate()
    
    X = data.drop('fraud', axis=1)
    y = data['fraud']
    
    split_idx = int(0.8 * len(data))
    X_train, X_test = X[:split_idx], X[split_idx:]
    y_train, y_test = y[:split_idx], y[split_idx:]
    
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)
    
    print(f"Training set: {len(X_train)} transactions, {y_train.sum()} frauds ({y_train.mean()*100:.2f}%)")
    print(f"Test set: {len(X_test)} transactions, {y_test.sum()} frauds ({y_test.mean()*100:.2f}%)")
    
    iso_forest = IsolationForestDetector(contamination=0.01)
    iso_forest.fit(X_train_scaled)
    iso_scores = iso_forest.predict(X_test_scaled)
    iso_auc = roc_auc_score(y_test, iso_scores)
    print(f"\nIsolation Forest AUC: {iso_auc:.4f}")
    
    vae = AutoencoderDetector(input_dim=X_train_scaled.shape[1], latent_dim=4)
    X_normal = X_train_scaled[y_train == 0]
    vae = train_vae_detector(vae, X_normal, epochs=30)
    
    with torch.no_grad():
        vae_scores = vae.anomaly_score(torch.FloatTensor(X_test_scaled)).numpy()
    vae_auc = roc_auc_score(y_test, vae_scores)
    print(f"VAE Anomaly Score AUC: {vae_auc:.4f}")
    
    xgb_model = xgb.XGBClassifier(
        n_estimators=100, max_depth=6, learning_rate=0.1,
        scale_pos_weight=len(y_train[y_train==0]) / len(y_train[y_train==1]),
        random_state=42
    )
    xgb_model.fit(X_train_scaled, y_train)
    xgb_pred = xgb_model.predict_proba(X_test_scaled)[:, 1]
    xgb_auc = roc_auc_score(y_test, xgb_pred)
    precision, recall, f1, _ = precision_recall_fscore_support(y_test, (xgb_pred > 0.5).astype(int))
    print(f"XGBoost AUC: {xgb_auc:.4f}, Precision: {precision[1]:.4f}, Recall: {recall[1]:.4f}, F1: {f1[1]:.4f}")
    
    ensemble = FraudEnsemble()
    ensemble.add_model('xgboost', xgb_model, weight=0.5)
    ensemble.add_model('isoforest', iso_forest, weight=0.25)
    
    ens_results = ensemble.evaluate(X_test_scaled, y_test)
    print(f"\nEnsemble AUC: {ens_results['ensemble_auc']:.4f}")
    for name, auc in ens_results['model_aucs'].items():
        print(f"  {name}: {auc:.4f}")