"""
===============================================================================
SIZINTI HARİTASI (LEAKAGE MAP): Sızıntının Ağ Üzerinde Nasıl Aktığını Göster
===============================================================================

HİÇ YAPILMAMIŞ ÖZGÜN KATKI:

Random split yaptığınızda test düğümlerinin komşuları eğitim setinde olabiliyor.
GraphSAGE gibi GNN'ler komşuluktan bilgi topladığı için, bu komşuluk üzerinden
geleceğin bilgisi geçmişe SIZIYOR.

Biz bunu ilk kez ölçüyor ve haritasını çıkarıyoruz:

1. KOMŞULUK SIZINTI ORANI: Her test düğümü için komşularının yüzde kaçı 
   eğitim setinde? (Yüksek = çok sızıntı)

2. SIZINTI-BAŞARI KORELASYONU: Komşu sızıntısı yüksek olan düğümler daha mı 
   doğru tahmin ediliyor? (Korelasyon varsa = sızıntı kanıtı)

3. TEMPORAL SIZINTI: Test düğümünün komşuları farklı bir zaman diliminden mi?
   (Farklı zaman = temporal bilgi sızıntısı)

4. YAPI DUYARLILIĞI: Komşu sızıntısı kaldırıldığında (test düğümlerinin 
   eğitim komşularıyla bağları kesildiğinde) performans ne kadar düşüyor?
   (Düşüş = model sızıntıya bağımlıydı)
===============================================================================
"""

import os, json, warnings
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import f1_score, roc_auc_score
from sklearn.model_selection import train_test_split
import lightgbm as lgb
warnings.filterwarnings('ignore')
np.random.seed(42)

FIGDIR = '/app/figures_proof'
OUTDIR = '/app/results_proof'

# ─── VERİ ───
feat_df = pd.read_csv('/app/data/elliptic_txs_features.csv', header=None)
class_df = pd.read_csv('/app/data/elliptic_txs_classes.csv')
edge_df = pd.read_csv('/app/data/elliptic_txs_edgelist.csv')

txids = feat_df.iloc[:, 0].values
timesteps = feat_df.iloc[:, 1].values.astype(int)
features = feat_df.iloc[:, 2:].values.astype(np.float32)
N = len(txids)

id_map = {tid: i for i, tid in enumerate(txids)}
label_map = {'1': 1, '2': 0, 'unknown': -1}
labels = np.array([label_map[str(c)] for c in class_df['class'].values])

# Komşuluk sözlüğü oluştur
print("Komşuluk sözlüğü oluşturuluyor...")
neighbors = {i: set() for i in range(N)}
for _, row in edge_df.iterrows():
    s = id_map.get(row['txId1'])
    d = id_map.get(row['txId2'])
    if s is not None and d is not None:
        neighbors[s].add(d)
        neighbors[d].add(s)
print(f"  Toplam düğüm: {N}, ortalama komşu: {np.mean([len(v) for v in neighbors.values()]):.1f}")

labeled_mask = labels >= 0
X = features[labeled_mask]
y = labels[labeled_mask]
ts = timesteps[labeled_mask]
labeled_indices = np.where(labeled_mask)[0]  # Global indeksler

# ─── RANDOM SPLIT ───
idx_all = np.arange(len(y))
train_idx, test_idx = train_test_split(idx_all, test_size=0.2, random_state=42, stratify=y)
train_global = set(labeled_indices[train_idx])
test_global = labeled_indices[test_idx]

# ─── Model eğit ───
scaler = StandardScaler()
X_tr = scaler.fit_transform(X[train_idx])
X_te = scaler.transform(X[test_idx])
model = lgb.LGBMClassifier(n_estimators=300, max_depth=10, scale_pos_weight=10, random_state=42, n_jobs=-1, verbose=-1)
model.fit(X_tr, y[train_idx])
pred = model.predict(X_te)
proba = model.predict_proba(X_te)[:, 1]
correct = (pred == y[test_idx])

print(f"\nRandom split F1: {f1_score(y[test_idx], pred):.4f}")

# ═════════════════════════════════════════════════════════════════════════
# ÖLÇÜM 1: KOMŞULUK SIZINTI ORANI
# ═════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print("ÖLÇÜM 1: KOMŞULUK SIZINTI ORANI")
print("Her test düğümünün komşularının yüzde kaçı eğitim setinde?")
print("=" * 70)

leakage_ratios = []
neighbor_counts = []
train_neighbor_ts_diffs = []

for i, global_idx in enumerate(test_global):
    nbrs = neighbors[global_idx]
    if len(nbrs) == 0:
        leakage_ratios.append(0.0)
        neighbor_counts.append(0)
        train_neighbor_ts_diffs.append(0.0)
        continue
    
    train_nbrs = nbrs & train_global
    ratio = len(train_nbrs) / len(nbrs)
    leakage_ratios.append(ratio)
    neighbor_counts.append(len(nbrs))
    
    # Zamansal mesafe: test düğümü ile eğitim komşuları arası
    test_ts = timesteps[global_idx]
    if len(train_nbrs) > 0:
        nbr_ts = [timesteps[n] for n in train_nbrs]
        avg_diff = np.mean(np.abs(np.array(nbr_ts) - test_ts))
        train_neighbor_ts_diffs.append(avg_diff)
    else:
        train_neighbor_ts_diffs.append(0.0)

leakage_ratios = np.array(leakage_ratios)
neighbor_counts = np.array(neighbor_counts)
train_neighbor_ts_diffs = np.array(train_neighbor_ts_diffs)

print(f"  Ortalama sızıntı oranı: {leakage_ratios.mean():.3f}")
print(f"  Median sızıntı oranı:  {np.median(leakage_ratios):.3f}")
print(f"  Sızıntısız düğüm (0.0): {(leakage_ratios == 0).sum()} ({(leakage_ratios == 0).mean()*100:.1f}%)")
print(f"  Tam sızıntı (1.0):      {(leakage_ratios == 1).sum()} ({(leakage_ratios == 1).mean()*100:.1f}%)")

# ═════════════════════════════════════════════════════════════════════════
# ÖLÇÜM 2: SIZINTI-BAŞARI KORELASYONU
# ═════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print("ÖLÇÜM 2: SIZINTI-BAŞARI KORELASYONU")
print("Komşu sızıntısı yüksek düğümler daha mı kolay tahmin ediliyor?")
print("=" * 70)

# Sızıntı oranı gruplarına göre performans
bins = [0.0, 0.2, 0.4, 0.6, 0.8, 1.01]
bin_labels = ['0-20%', '20-40%', '40-60%', '60-80%', '80-100%']
leak_bins = pd.cut(leakage_ratios, bins=bins, labels=bin_labels, include_lowest=True)

print(f"\n  Sızıntı grubuna göre doğruluk:")
bin_results = []
for bl in bin_labels:
    mask = (leak_bins == bl)
    if mask.sum() < 10:
        continue
    acc = correct[mask].mean()
    n = mask.sum()
    n_ill = y[test_idx][mask].sum()
    f1_bin = f1_score(y[test_idx][mask], pred[mask], zero_division=0) if len(np.unique(y[test_idx][mask])) > 1 else 0
    bin_results.append({'bin': bl, 'accuracy': acc, 'f1': f1_bin, 'n': n, 'n_illicit': n_ill})
    print(f"    Sızıntı {bl:>8s}: doğruluk={acc:.4f}, F1={f1_bin:.4f} (n={n}, illicit={n_ill})")

# Korelasyon
corr = np.corrcoef(leakage_ratios[neighbor_counts > 0], 
                    correct[neighbor_counts > 0].astype(float))[0, 1]
print(f"\n  Sızıntı oranı ↔ doğru tahmin korelasyonu: r = {corr:.4f}")
if abs(corr) > 0.05:
    print(f"  ⚠️ Pozitif korelasyon: Sızıntı yüksek → tahmin daha kolay")

# ═════════════════════════════════════════════════════════════════════════
# ÖLÇÜM 3: TEMPORAL SIZINTI
# ═════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print("ÖLÇÜM 3: TEMPORAL SIZINTI")
print("Test düğümünün eğitim komşuları farklı zaman diliminden mi?")
print("=" * 70)

# Test düğümü ile eğitim komşuları arasındaki zamansal mesafe
has_train_nbrs = train_neighbor_ts_diffs > 0
print(f"  Ortalama zamansal mesafe (test ↔ eğitim komşu): {train_neighbor_ts_diffs[has_train_nbrs].mean():.2f} timestep")
print(f"  Aynı timestep'ten komşu oranı: {(train_neighbor_ts_diffs == 0).sum() / max(has_train_nbrs.sum(), 1) * 100:.1f}%")

# TEMPORAL SPLIT İLE KARŞILAŞTIR
print("\n  Temporal split'te sızıntı oranı:")
cutoff = 39
temp_train_global = set(np.where((labels >= 0) & (timesteps <= cutoff))[0])
temp_test_global = np.where((labels >= 0) & (timesteps > cutoff))[0]

temp_leak = []
for global_idx in temp_test_global:
    nbrs = neighbors[global_idx]
    if len(nbrs) == 0:
        temp_leak.append(0.0)
        continue
    temp_leak.append(len(nbrs & temp_train_global) / len(nbrs))

temp_leak = np.array(temp_leak)
print(f"  Random split ortalama sızıntı oranı:   {leakage_ratios.mean():.3f}")
print(f"  Temporal split ortalama sızıntı oranı: {temp_leak.mean():.3f}")
print(f"  FARK: Random split {leakage_ratios.mean() / max(temp_leak.mean(), 0.001):.1f}x daha fazla sızıntı!")

# ═════════════════════════════════════════════════════════════════════════
# ÖLÇÜM 4: YAPI DUYARLILIĞI TESTİ
# ═════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print("ÖLÇÜM 4: YAPI DUYARLILIĞI (Sızıntıya Bağımlılık)")
print("Aggregated özellikler (komşuluk bilgisi) çıkarılınca ne oluyor?")
print("=" * 70)

# Özellik 0-93: local (düğümün kendi özellikleri)
# Özellik 94-165: aggregated (komşuluk bilgisi)
X_local = features[labeled_mask][:, :93]  # Sadece local özellikler
X_agg = features[labeled_mask][:, 93:]     # Sadece aggregated özellikler

# Local ile eğit (sızıntıdan etkilenmez)
scaler_l = StandardScaler()
X_tr_l = scaler_l.fit_transform(X_local[train_idx])
X_te_l = scaler_l.transform(X_local[test_idx])
model_l = lgb.LGBMClassifier(n_estimators=300, max_depth=10, scale_pos_weight=10, random_state=42, n_jobs=-1, verbose=-1)
model_l.fit(X_tr_l, y[train_idx])
pred_l = model_l.predict(X_te_l)
f1_local = f1_score(y[test_idx], pred_l)

# Aggregated ile eğit (komşuluk = sızıntı kanalı)
scaler_a = StandardScaler()
X_tr_a = scaler_a.fit_transform(X_agg[train_idx])
X_te_a = scaler_a.transform(X_agg[test_idx])
model_a = lgb.LGBMClassifier(n_estimators=300, max_depth=10, scale_pos_weight=10, random_state=42, n_jobs=-1, verbose=-1)
model_a.fit(X_tr_a, y[train_idx])
pred_a = model_a.predict(X_te_a)
f1_agg = f1_score(y[test_idx], pred_a)

# Tüm özelliklerle
f1_all = f1_score(y[test_idx], pred)

# Temporal split ile aynı testi yap
tr_t = np.where((labels[labeled_mask] >= 0) & (ts <= cutoff))[0]
te_t = np.where((labels[labeled_mask] >= 0) & (ts > cutoff))[0]

if len(tr_t) > 0 and len(te_t) > 0 and len(np.unique(y[te_t])) > 1:
    # Local - temporal
    scaler_lt = StandardScaler()
    model_lt = lgb.LGBMClassifier(n_estimators=300, max_depth=10, scale_pos_weight=10, random_state=42, n_jobs=-1, verbose=-1)
    model_lt.fit(scaler_lt.fit_transform(X_local[tr_t]), y[tr_t])
    f1_local_temp = f1_score(y[te_t], model_lt.predict(scaler_lt.transform(X_local[te_t])))
    
    # Aggregated - temporal
    scaler_at = StandardScaler()
    model_at = lgb.LGBMClassifier(n_estimators=300, max_depth=10, scale_pos_weight=10, random_state=42, n_jobs=-1, verbose=-1)
    model_at.fit(scaler_at.fit_transform(X_agg[tr_t]), y[tr_t])
    f1_agg_temp = f1_score(y[te_t], model_at.predict(scaler_at.transform(X_agg[te_t])))
    
    # All - temporal  
    scaler_allt = StandardScaler()
    model_allt = lgb.LGBMClassifier(n_estimators=300, max_depth=10, scale_pos_weight=10, random_state=42, n_jobs=-1, verbose=-1)
    model_allt.fit(scaler_allt.fit_transform(X[tr_t]), y[tr_t])
    f1_all_temp = f1_score(y[te_t], model_allt.predict(scaler_allt.transform(X[te_t])))

print(f"\n  RANDOM SPLIT:")
print(f"    Tüm özellikler:           F1 = {f1_all:.4f}")
print(f"    Sadece local (93 feat):   F1 = {f1_local:.4f}")
print(f"    Sadece aggregated (72):   F1 = {f1_agg:.4f}")
print(f"    Aggregated'ın katkısı:    +{f1_all - f1_local:.4f}")

print(f"\n  TEMPORAL SPLIT:")
print(f"    Tüm özellikler:           F1 = {f1_all_temp:.4f}")
print(f"    Sadece local (93 feat):   F1 = {f1_local_temp:.4f}")
print(f"    Sadece aggregated (72):   F1 = {f1_agg_temp:.4f}")
print(f"    Aggregated'ın katkısı:    +{f1_all_temp - f1_local_temp:.4f}")

agg_boost_random = f1_all - f1_local
agg_boost_temporal = f1_all_temp - f1_local_temp
print(f"\n  KILIT BULGU:")
print(f"    Aggregated özelliklerin random split'te ekstra katkısı: +{agg_boost_random:.4f}")
print(f"    Aggregated özelliklerin temporal split'te ekstra katkısı: +{agg_boost_temporal:.4f}")
if agg_boost_random > agg_boost_temporal + 0.01:
    print(f"    ⚠️ Aggregated özellikler random split'te {agg_boost_random/max(agg_boost_temporal, 0.001):.1f}x daha fazla katkı sağlıyor!")
    print(f"    Bu, komşuluk bilgisinin random split'te sızıntı kanalı olduğunun kanıtıdır.")

# ═════════════════════════════════════════════════════════════════════════
# FİGÜRLER
# ═════════════════════════════════════════════════════════════════════════
print("\nFİGÜRLER OLUŞTURULUYOR...")
sns.set_theme(style='whitegrid', font_scale=1.1)

# ── FİGÜR: Sızıntı Haritası ──
fig, axes = plt.subplots(2, 2, figsize=(18, 14))

# 1: Sızıntı oranı dağılımı
axes[0, 0].hist(leakage_ratios, bins=30, color='#FF6B6B', edgecolor='white', alpha=0.8)
axes[0, 0].set_xlabel('Komşuluk Sızıntı Oranı', fontsize=12)
axes[0, 0].set_ylabel('Düğüm Sayısı', fontsize=12)
axes[0, 0].set_title('Her Test Düğümünün Komşuluk Sızıntı Oranı\n(Random Split)', fontsize=13, fontweight='bold')
axes[0, 0].axvline(x=leakage_ratios.mean(), color='black', linewidth=2, linestyle='--',
                   label=f'Ortalama: {leakage_ratios.mean():.2f}')
axes[0, 0].legend(fontsize=11)

# 2: Sızıntı oranı vs doğruluk
if len(bin_results) > 0:
    br_df = pd.DataFrame(bin_results)
    axes[0, 1].bar(br_df['bin'], br_df['f1'], color='#4ECDC4', edgecolor='black')
    axes[0, 1].set_xlabel('Komşuluk Sızıntı Oranı', fontsize=12)
    axes[0, 1].set_ylabel('Illicit F1 Score', fontsize=12)
    axes[0, 1].set_title('Sızıntı ↑ = Tahmin Doğruluğu ↑ ?\n(Korelasyon varsa sızıntı kanıtı)', fontsize=13, fontweight='bold')

# 3: Random vs Temporal sızıntı karşılaştırması
data_compare = [
    ('Random Split', leakage_ratios.mean()),
    ('Temporal Split', temp_leak.mean()),
]
axes[1, 0].bar([d[0] for d in data_compare], [d[1] for d in data_compare],
               color=['#FF6B6B', '#4ECDC4'], edgecolor='black')
axes[1, 0].set_ylabel('Ortalama Komşuluk Sızıntı Oranı', fontsize=12)
axes[1, 0].set_title('Random vs Temporal: Komşuluk Sızıntısı', fontsize=13, fontweight='bold')
for i, (name, val) in enumerate(data_compare):
    axes[1, 0].text(i, val + 0.01, f'{val:.3f}', ha='center', fontsize=13, fontweight='bold')

# 4: Yapı duyarlılığı
feat_types = ['Sadece Local\n(93 özellik)', 'Sadece Aggregated\n(72 özellik)', 'Tümü\n(165 özellik)']
random_vals = [f1_local, f1_agg, f1_all]
temporal_vals = [f1_local_temp, f1_agg_temp, f1_all_temp]
x = np.arange(len(feat_types))
axes[1, 1].bar(x - 0.18, random_vals, 0.32, label='Random Split', color='#FF6B6B')
axes[1, 1].bar(x + 0.18, temporal_vals, 0.32, label='Temporal Split', color='#4ECDC4')
axes[1, 1].set_xticks(x)
axes[1, 1].set_xticklabels(feat_types, fontsize=10)
axes[1, 1].set_ylabel('Illicit F1 Score', fontsize=12)
axes[1, 1].set_title('Yapı Duyarlılığı: Aggregated Özellikler\nRandom Split\'te Daha mı Çok Katkı Sağlıyor?', fontsize=13, fontweight='bold')
axes[1, 1].legend(fontsize=11)

plt.suptitle('SIZINTI HARİTASI: Bilgi Ağ Üzerinden Nasıl Sızıyor?', fontsize=16, fontweight='bold', y=1.02)
plt.tight_layout()
plt.savefig(f'{FIGDIR}/fig7_leakage_map.png', dpi=150, bbox_inches='tight')
plt.close()
print("  ✓ Figür 7: Sızıntı Haritası")

# Sonuçları kaydet
results = {
    'avg_leakage_ratio_random': float(leakage_ratios.mean()),
    'avg_leakage_ratio_temporal': float(temp_leak.mean()),
    'leakage_ratio_multiplier': float(leakage_ratios.mean() / max(temp_leak.mean(), 0.001)),
    'correlation_leakage_accuracy': float(corr),
    'f1_local_random': float(f1_local),
    'f1_agg_random': float(f1_agg),
    'f1_all_random': float(f1_all),
    'f1_local_temporal': float(f1_local_temp),
    'f1_agg_temporal': float(f1_agg_temp),
    'f1_all_temporal': float(f1_all_temp),
    'agg_boost_random': float(agg_boost_random),
    'agg_boost_temporal': float(agg_boost_temporal),
}
with open(f'{OUTDIR}/leakage_map_results.json', 'w') as f:
    json.dump(results, f, indent=2)

print("\n✓ Tamamlandı!")