# ===================================================================
# MODULE 7, LESSON 1: EDGE AI AND IOT IN BANKING
# ===================================================================

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.ensemble import IsolationForest
from sklearn.preprocessing import StandardScaler
import joblib
import time
import hashlib
import json
import warnings
warnings.filterwarnings('ignore')

# For TinyML simulation, we'll use quantisation
try:
    import tensorflow as tf
    from tensorflow import keras
    TENSORFLOW_AVAILABLE = True
except ImportError:
    TENSORFLOW_AVAILABLE = False
    print("TensorFlow not installed. Install with: pip install tensorflow")

# Set style
sns.set_style("whitegrid")
np.random.seed(42)

print("="*70)
print("EDGE AI AND IOT IN BANKING – ATM ANOMALY DETECTION")
print("="*70)

# ----------------------------------------------------------------
# PART A: GENERATE ATM TRANSACTION DATA
# ----------------------------------------------------------------

def generate_atm_data(n_transactions=10000, anomaly_rate=0.05):
    """
    Generate synthetic ATM transaction data.
    Includes normal transactions and anomalies.
    """
    n_normal = int(n_transactions * (1 - anomaly_rate))
    n_anomaly = n_transactions - n_normal
    
    # Normal transactions: small amounts, normal times, typical patterns
    amounts_normal = np.random.lognormal(3.5, 0.6, n_normal).clip(10, 500)
    times_normal = np.random.normal(12, 4, n_normal).clip(0, 23)  # Hours (0-23)
    weekday_normal = np.random.choice([0, 1, 2, 3, 4, 5, 6], n_normal, p=[0.15, 0.15, 0.15, 0.15, 0.15, 0.1, 0.15])
    duration_normal = np.random.gamma(2, 1, n_normal).clip(0.1, 10)
    balance_after_normal = np.random.gamma(5, 200, n_normal).clip(50, 10000)
    
    # Anomalies: high amounts, unusual times, multiple transactions
    amounts_anomaly = np.random.lognormal(6, 1.2, n_anomaly).clip(1000, 50000)
    times_anomaly = np.random.normal(3, 2, n_anomaly).clip(0, 23)  # Early morning
    weekday_anomaly = np.random.choice([0, 5, 6], n_anomaly, p=[0.3, 0.4, 0.3])  # Weekends
    duration_anomaly = np.random.gamma(0.5, 1, n_anomaly).clip(0.1, 10)
    balance_after_anomaly = np.random.gamma(2, 100, n_anomaly).clip(0, 5000)
    
    # Combine
    amounts = np.concatenate([amounts_normal, amounts_anomaly])
    times = np.concatenate([times_normal, times_anomaly])
    weekdays = np.concatenate([weekday_normal, weekday_anomaly])
    durations = np.concatenate([duration_normal, duration_anomaly])
    balances = np.concatenate([balance_after_normal, balance_after_anomaly])
    labels = np.concatenate([np.zeros(n_normal), np.ones(n_anomaly)])
    
    # Shuffle
    idx = np.random.permutation(n_transactions)
    amounts = amounts[idx]
    times = times[idx]
    weekdays = weekdays[idx]
    durations = durations[idx]
    balances = balances[idx]
    labels = labels[idx]
    
    # Create DataFrame
    df = pd.DataFrame({
        'amount': amounts,
        'time': times,
        'weekday': weekdays,
        'duration': durations,
        'balance_after': balances,
        'is_anomaly': labels
    })
    return df

# Generate data
df_atm = generate_atm_data(10000, anomaly_rate=0.05)
print("ATM Transaction Data Generated:")
print(f"  Total transactions: {len(df_atm)}")
print(f"  Anomalies: {df_atm['is_anomaly'].sum()} ({df_atm['is_anomaly'].mean()*100:.2f}%)")

# Visualise
fig, axes = plt.subplots(2, 2, figsize=(14, 10))

ax = axes[0, 0]
ax.hist(df_atm[df_atm['is_anomaly'] == 0]['amount'], bins=50, alpha=0.5, label='Normal', color='blue')
ax.hist(df_atm[df_atm['is_anomaly'] == 1]['amount'], bins=50, alpha=0.5, label='Anomaly', color='red')
ax.set_xlabel('Transaction Amount ($)')
ax.set_ylabel('Frequency')
ax.set_title('Transaction Amounts')
ax.legend()
ax.grid(True, alpha=0.3)

ax = axes[0, 1]
ax.hist(df_atm[df_atm['is_anomaly'] == 0]['time'], bins=24, alpha=0.5, label='Normal', color='blue')
ax.hist(df_atm[df_atm['is_anomaly'] == 1]['time'], bins=24, alpha=0.5, label='Anomaly', color='red')
ax.set_xlabel('Time of Day (Hours)')
ax.set_ylabel('Frequency')
ax.set_title('Transaction Times')
ax.legend()
ax.grid(True, alpha=0.3)

ax = axes[1, 0]
ax.scatter(df_atm['time'], df_atm['amount'], c=df_atm['is_anomaly'], 
           cmap='coolwarm', alpha=0.5, s=10)
ax.set_xlabel('Time of Day (Hours)')
ax.set_ylabel('Transaction Amount ($)')
ax.set_title('Amount vs Time of Day')
ax.grid(True, alpha=0.3)

ax = axes[1, 1]
ax.scatter(df_atm['duration'], df_atm['amount'], c=df_atm['is_anomaly'], 
           cmap='coolwarm', alpha=0.5, s=10)
ax.set_xlabel('Transaction Duration (min)')
ax.set_ylabel('Transaction Amount ($)')
ax.set_title('Amount vs Duration')
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('atm_transactions.png', dpi=300)
plt.show()

# ----------------------------------------------------------------
# PART B: EDGE-BASED ANOMALY DETECTION MODEL (ISOLATION FOREST)
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART B: Training Edge-Based Anomaly Detection Model")
print("-"*60)

# Features for model
features = ['amount', 'time', 'weekday', 'duration', 'balance_after']
X = df_atm[features]
y = df_atm['is_anomaly']

# Standardise (for Isolation Forest, scaling isn't required but helps)
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# Train Isolation Forest
model_if = IsolationForest(contamination=0.05, random_state=42, n_estimators=100)
model_if.fit(X_scaled)

# Predictions (1 = normal, -1 = anomaly)
preds = model_if.predict(X_scaled)
pred_labels = (preds == -1).astype(int)

# Performance
from sklearn.metrics import classification_report, confusion_matrix
print("Performance on Training Data:")
print(classification_report(y, pred_labels, target_names=['Normal', 'Anomaly']))

# Confusion matrix
cm = confusion_matrix(y, pred_labels)
fig, ax = plt.subplots(figsize=(6, 5))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
            xticklabels=['Normal', 'Anomaly'],
            yticklabels=['Normal', 'Anomaly'])
ax.set_title('Confusion Matrix – Isolation Forest')
plt.tight_layout()
plt.savefig('confusion_matrix.png', dpi=300)
plt.show()

# ----------------------------------------------------------------
# PART C: MODEL QUANTISATION (FOR EDGE DEPLOYMENT)
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART C: Model Quantisation for Edge Deployment")
print("-"*60)

# Convert Isolation Forest to a lightweight format for edge
# For demonstration, we'll save the model and demonstrate quantisation
# using a small neural network (more realistic for edge)

# Build a small neural network for comparison
if TENSORFLOW_AVAILABLE:
    # Build a simple model
    model_nn = keras.Sequential([
        keras.layers.Dense(16, activation='relu', input_shape=(5,)),
        keras.layers.Dense(8, activation='relu'),
        keras.layers.Dense(1, activation='sigmoid')
    ])
    model_nn.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
    
    # Train
    X_train = X_scaled[:8000]
    y_train = y[:8000]
    X_val = X_scaled[8000:]
    y_val = y[8000:]
    
    model_nn.fit(X_train, y_train, epochs=20, batch_size=64, 
                 validation_data=(X_val, y_val), verbose=0)
    
    # Convert to TensorFlow Lite (quantised)
    converter = tf.lite.TFLiteConverter.from_keras_model(model_nn)
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    converter.target_spec.supported_types = [tf.float16]  # FP16 quantisation
    tflite_model = converter.convert()
    
    # Save quantised model
    with open('model_quantised.tflite', 'wb') as f:
        f.write(tflite_model)
    
    # Model sizes
    import os
    fp32_size = os.path.getsize('model_quantised.tflite') / 1024
    print(f"Quantised model size: {fp32_size:.2f} KB")
    print("  Suitable for deployment on edge devices (microcontrollers).")
else:
    print("TensorFlow not available. Quantisation step skipped.")

# For Isolation Forest, we can save a compressed version
# Use reduced precision (float16) for storage
import pickle
import struct

def compress_isolation_forest(model):
    """Compress Isolation Forest model for edge deployment."""
    # Extract tree structures
    compressed = {
        'n_estimators': len(model.estimators_),
        'max_samples': model.max_samples_,
        'contamination': model.contamination,
    }
    trees = []
    for tree in model.estimators_:
        tree_data = {
            'tree': tree.tree_,  # Could compress further
            'feature': tree.tree_.feature,
            'threshold': tree.tree_.threshold,
            'n_node_count': tree.tree_.node_count,
            'children_left': tree.tree_.children_left,
            'children_right': tree.tree_.children_right,
        }
        trees.append(tree_data)
    compressed['trees'] = trees
    return compressed

# Compress
compressed_model = compress_isolation_forest(model_if)
print(f"Model compression: {len(compressed_model['trees'])} trees stored.")

# ----------------------------------------------------------------
# PART D: SIMULATE EDGE DEPLOYMENT
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART D: Simulating Edge Deployment")
print("-"*60)

class EdgeAnomalyDetector:
    """
    Simulate an edge-based anomaly detection device.
    """
    def __init__(self, model, scaler, threshold=0.5):
        self.model = model
        self.scaler = scaler
        self.threshold = threshold
        self.transaction_count = 0
        self.anomaly_count = 0
        self.running = False
    
    def start(self):
        """Start the edge device."""
        self.running = True
        print("Edge device started. Monitoring transactions...")
    
    def stop(self):
        """Stop the edge device."""
        self.running = False
        print(f"Edge device stopped. Processed {self.transaction_count} transactions.")
        print(f"Anomalies detected: {self.anomaly_count} ({self.anomaly_count/self.transaction_count*100:.2f}%)")
    
    def process_transaction(self, transaction):
        """
        Process a single transaction on the edge.
        Returns: (is_anomaly, score, latency_ms)
        """
        if not self.running:
            return None
        
        start_time = time.time()
        self.transaction_count += 1
        
        # Preprocess
        features = ['amount', 'time', 'weekday', 'duration', 'balance_after']
        X_input = np.array([[transaction[f] for f in features]])
        X_scaled = self.scaler.transform(X_input)
        
        # Predict (Isolation Forest)
        score = self.model.decision_function(X_scaled)[0]  # Anomaly score (lower = more anomalous)
        # Invert to get anomaly probability-like value
        score_normalised = 1 / (1 + np.exp(-score))  # Sigmoid transform
        is_anomaly = score_normalised < self.threshold
        
        # Convert to binary label
        is_anomaly_bool = bool(is_anomaly)
        
        if is_anomaly_bool:
            self.anomaly_count += 1
        
        latency_ms = (time.time() - start_time) * 1000
        
        return {
            'is_anomaly': is_anomaly_bool,
            'score': score_normalised,
            'latency_ms': latency_ms,
            'transaction_count': self.transaction_count
        }

# Create edge detector
edge_detector = EdgeAnomalyDetector(model_if, scaler, threshold=0.4)
edge_detector.start()

# Simulate live stream of transactions
print("\nSimulating live transaction stream...")
transaction_stream = df_atm.sample(n=100, random_state=42)

for idx, tx in transaction_stream.iterrows():
    transaction = tx[['amount', 'time', 'weekday', 'duration', 'balance_after']].to_dict()
    result = edge_detector.process_transaction(transaction)
    if result['is_anomaly']:
        print(f"⚠ ANOMALY DETECTED! Transaction #{result['transaction_count']}")
        print(f"   Amount: ${tx['amount']:.2f}, Time: {tx['time']:.1f}h, Score: {result['score']:.3f}")
        print(f"   Latency: {result['latency_ms']:.2f}ms")

edge_detector.stop()

# ----------------------------------------------------------------
# PART E: EDGE AI PERFORMANCE COMPARISON
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART E: Edge AI Performance Analysis")
print("-"*60)

# Simulate latency comparison
cloud_latency = np.random.normal(150, 50, 1000)  # ms
edge_latency = np.random.normal(5, 2, 1000)      # ms

fig, ax = plt.subplots(figsize=(10, 6))
ax.hist(cloud_latency, bins=30, alpha=0.5, label='Cloud AI', color='blue')
ax.hist(edge_latency, bins=30, alpha=0.5, label='Edge AI', color='green')
ax.axvline(np.mean(cloud_latency), color='blue', linestyle='--', label=f'Cloud Mean: {np.mean(cloud_latency):.0f}ms')
ax.axvline(np.mean(edge_latency), color='green', linestyle='--', label=f'Edge Mean: {np.mean(edge_latency):.0f}ms')
ax.set_xlabel('Latency (ms)')
ax.set_ylabel('Frequency')
ax.set_title('Edge AI vs Cloud AI – Latency Comparison')
ax.legend()
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('latency_comparison.png', dpi=300)
plt.show()

print(f"Edge AI latency: {np.mean(edge_latency):.2f}ms")
print(f"Cloud AI latency: {np.mean(cloud_latency):.2f}ms")
print(f"Speedup: {np.mean(cloud_latency)/np.mean(edge_latency):.1f}x")
print(f"Privacy benefit: Data stays on the ATM (no sensitive data transmitted).")

# ----------------------------------------------------------------
# PART F: EDGE AI DEPLOYMENT ROADMAP
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART F: Edge AI Deployment Roadmap for Banking")
print("-"*60)

roadmap = {
    "Phase 1 – Assessment (3-6 months)": {
        "Actions": [
            "Identify use cases with clear ROI (ATM fraud detection, branch analytics).",
            "Assess existing IoT infrastructure.",
            "Evaluate edge AI platforms (TensorFlow Lite, NVIDIA Jetson, Coral).",
            "Develop PoC for one use case."
        ]
    },
    "Phase 2 – Pilot (6-12 months)": {
        "Actions": [
            "Deploy edge AI in a small set of ATMs/branches.",
            "Monitor performance, latency, and accuracy.",
            "Integrate with existing monitoring and alerting systems.",
            "Develop edge model update pipeline."
        ]
    },
    "Phase 3 – Scale (12-24 months)": {
        "Actions": [
            "Roll out across all ATMs and branches.",
            "Expand to other use cases (wearables, IoT collateral monitoring).",
            "Establish edge AI Centre of Excellence.",
            "Integrate with cloud AI for advanced analytics."
        ]
    },
    "Phase 4 – Maturity (24+ months)": {
        "Actions": [
            "Continuous model improvement with federated learning.",
            "Explore 5G-enabled edge AI.",
            "Invest in custom edge AI hardware.",
            "Develop ecosystem partnerships."
        ]
    }
}

for phase, details in roadmap.items():
    print(f"\n{phase}:")
    for action in details['Actions']:
        print(f"  • {action}")

# ----------------------------------------------------------------
# PART G: REGULATORY AND SECURITY CONSIDERATIONS
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART G: Regulatory and Security Considerations")
print("-"*60)

print("""
Edge AI in Banking – Key Considerations:

1. Privacy and Data Protection:
   - GDPR/CCPA compliance: Data processed on edge, not transmitted.
   - Data minimisation: Only essential data used.
   - Anonymisation: Even local data should be anonymised where possible.

2. Security:
   - Device security: ATMs and IoT devices must be physically and digitally secure.
   - Model integrity: Prevent tampering with edge models.
   - Secure boot: Ensure only authorised firmware runs.
   - Encryption: Encrypt data at rest and in transit.

3. Regulatory Compliance:
   - Model validation: Edge models must be validated per SR 11-7.
   - Auditability: Track edge decisions and update logs.
   - Explainability: SHAP/LIME may not run on edge, so provide cloud-based explanation fallback.

4. Operational Risk:
   - Model drift: Edge models must be updated regularly.
   - Device failure: Redundancy and failover procedures.
   - Network dependency: Edge AI should operate offline.

5. Ethical Considerations:
   - Bias in edge models can lead to unfair outcomes.
   - Transparency: Customers should know when AI is used.
   - Human oversight: Critical decisions should have human review.

Best Practices:
  - Conduct regular security audits of edge devices.
  - Implement remote model update capabilities (over-the-air updates).
  - Monitor edge device health and performance.
  - Maintain a hybrid approach: Edge for real-time, cloud for advanced analytics.
""")

# ----------------------------------------------------------------
# PART H: SUMMARY AND RECOMMENDATIONS
# ----------------------------------------------------------------

print("\n" + "="*70)
print("PART H: Summary and Recommendations")
print("="*70)

print("""
Edge AI in Banking – Key Takeaways:

1. Edge AI processes data locally on devices (ATMs, wearables, sensors).
2. Benefits: Low latency (sub-10ms), privacy, offline operation, bandwidth efficiency.
3. Key use cases: ATM fraud detection, smart branches, predictive maintenance, wearables.
4. TinyML enables AI on microcontrollers with limited resources.
5. Model quantisation (INT8, FP16) reduces model size and accelerates inference.
6. Challenges: Device constraints, security, model updates, regulatory compliance.
7. Roadmap: Assess → Pilot → Scale → Maturity.

Recommendations:
  - Start with a high-impact, low-risk use case (e.g., ATM anomaly detection).
  - Build a PoC with an existing edge AI platform (TensorFlow Lite, NVIDIA Jetson).
  - Establish a cross-functional team (IoT, data science, security, operations).
  - Invest in edge AI capabilities – it's the future of real-time banking.
""")

print("="*70)
print("END OF LESSON 1 – MODULE 7")
print("="*70)