import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import joblib
import os
import json
from sklearn.model_selection import StratifiedKFold, train_test_split
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import StandardScaler
from xgboost import XGBClassifier
from lightgbm import LGBMClassifier
from sklearn.ensemble import ExtraTreesClassifier, VotingClassifier
from scipy.ndimage import uniform_filter
import warnings
warnings.filterwarnings('ignore')

def load_and_clean():
    data = joblib.load('dataset_cache/training_data_fusion_32ch.joblib')
    X, y = data['X'].astype(np.float32), data['y']
    
    # Valid mask based on S2 data (channels 0:6). S2 data has 6 channels per timestep.
    # Total channels = 32 (4 timesteps * 8 channels)
    # Timestep 0 S2 channels = X[:, 0:6]
    valid = (y >= 0) & (X.reshape(X.shape[0], -1).sum(1) != 0)
    X, y = X[valid], y[valid]
    
    unique = sorted(np.unique(y).tolist())
    lmap = {l:i for i,l in enumerate(unique)}
    y = np.array([lmap[l] for l in y])
    print(f"Clean FUSION data: {X.shape}, {len(unique)} classes, {[int((y==i).sum()) for i in range(len(unique))]}")
    return X, y, len(unique)

def extract_features_fusion(X):
    """
    Extract features from 32-channel Fusion data (S2 + S1).
    Per timestep (8 channels):
    0-3: S2 B02, B03, B04, B08
    4-5: S2 NDVI, NDWI
    6-7: S1 VV, VH
    """
    N = X.shape[0]
    all_feats = []
    
    for i in range(N):
        patch = X[i]  # (32, 16, 16)
        feats = []
        
        # Valid timesteps for S2
        valid_ts = []
        for t in range(4):
            block_s2 = patch[t*8 : t*8+6]
            if np.abs(block_s2).sum() > 1e-6:
                valid_ts.append(t)
                
        if not valid_ts:
            valid_ts = [0]
            
        # === A. Per-valid-timestep features for S2 ===
        per_ts_stats_s2 = {b: [] for b in range(6)}
        for t in valid_ts:
            for b in range(6):
                ch = patch[t*8 + b]
                per_ts_stats_s2[b].append([
                    np.mean(ch), np.std(ch), np.median(ch),
                    np.min(ch), np.max(ch),
                    np.percentile(ch, 10), np.percentile(ch, 90),
                ])
                
        for b in range(6):
            stats = np.array(per_ts_stats_s2[b])
            feats.extend(stats.mean(axis=0).tolist())
            feats.extend(stats.std(axis=0).tolist())
        
        # === B. Sentinel-1 Features (Radar always penetrates clouds, so use all 4 timesteps) ===
        per_ts_stats_s1 = {b: [] for b in range(2)}
        for t in range(4):
            vv = patch[t*8 + 6]
            vh = patch[t*8 + 7]
            # Handle potential zeros if S1 was missing
            if np.abs(vv).sum() > 1e-6:
                per_ts_stats_s1[0].append([
                    np.mean(vv), np.std(vv), np.median(vv), np.max(vv), np.percentile(vv, 90)
                ])
                per_ts_stats_s1[1].append([
                    np.mean(vh), np.std(vh), np.median(vh), np.max(vh), np.percentile(vh, 90)
                ])
                
                # S1 specific: VH/VV ratio
                ratio = (vh + 1e-6) / (vv + 1e-6)
                feats.extend([np.mean(ratio), np.std(ratio), np.median(ratio)])
            else:
                feats.extend([0.0] * 3)
                
        for b in range(2):
            if len(per_ts_stats_s1[b]) > 0:
                stats = np.array(per_ts_stats_s1[b])
                feats.extend(stats.mean(axis=0).tolist())
                feats.extend(stats.std(axis=0).tolist())
            else:
                feats.extend([0.0] * 10)
                
        # === C. Spatial Texture (Radar Texture is very important!) ===
        for t in valid_ts[:2]:
            for b_idx in [3, 4]: # NIR, NDVI
                ch = patch[t*8 + b_idx]
                gx = np.diff(ch, axis=1); gy = np.diff(ch, axis=0)
                grad_mag = np.sqrt(np.mean(gx**2) + np.mean(gy**2))
                lm = uniform_filter(ch, size=3); lv = uniform_filter(ch**2, size=3) - lm**2
                feats.extend([grad_mag, np.mean(lv), np.std(lv)])
        
        # Radar Texture (VH, VV)
        for b_idx in [6, 7]:
            ch = patch[0*8 + b_idx] # Just use timestep 0 for Radar texture
            gx = np.diff(ch, axis=1); gy = np.diff(ch, axis=0)
            grad_mag = np.sqrt(np.mean(gx**2) + np.mean(gy**2))
            lm = uniform_filter(ch, size=3); lv = uniform_filter(ch**2, size=3) - lm**2
            feats.extend([grad_mag, np.mean(lv), np.std(lv)])
            
        # Pad S2 texture if needed
        needed = 2 * 2 * 3
        got = min(len(valid_ts), 2) * 2 * 3
        feats.extend([0.0] * (needed - got))
        
        # === D. Flat pixel features from best timestep (t=0) for ALL channels ===
        best_t = valid_ts[0]
        for b in range(8):
            ch = patch[best_t*8 + b]
            feats.extend(ch.flatten().tolist())
            
        all_feats.append(feats)
        
    features = np.array(all_feats, dtype=np.float32)
    features = np.nan_to_num(features, nan=0.0, posinf=1e6, neginf=-1e6)
    print(f"Extracted {features.shape[1]} fusion features per sample")
    return features

class LightCNN_32ch(nn.Module):
    def __init__(self, in_ch=32, n_cls=7, width=96):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_ch, width, 3, padding=1), nn.BatchNorm2d(width), nn.GELU(),
            nn.Conv2d(width, width, 3, padding=1), nn.BatchNorm2d(width), nn.GELU(),
            nn.MaxPool2d(2), nn.Dropout2d(0.05),
            nn.Conv2d(width, width*2, 3, padding=1), nn.BatchNorm2d(width*2), nn.GELU(),
            nn.Conv2d(width*2, width*2, 3, padding=1), nn.BatchNorm2d(width*2), nn.GELU(),
            nn.MaxPool2d(2), nn.Dropout2d(0.1),
            nn.Conv2d(width*2, width*4, 3, padding=1), nn.BatchNorm2d(width*4), nn.GELU(),
            nn.AdaptiveAvgPool2d(1), nn.Flatten(),
        )
        self.head = nn.Sequential(
            nn.Linear(width*4, width*2), nn.GELU(), nn.Dropout(0.3),
            nn.Linear(width*2, n_cls)
        )
    def forward(self, x): return self.head(self.net(x))
    def embed(self, x): return self.net(x)

def train_cnn_fusion(X, y, n_cls, seed=42):
    print("\n" + "="*60)
    print("32-CHANNELS FUSION CNN")
    print("="*60)
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    torch.manual_seed(seed)
    np.random.seed(seed)
    
    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.2, random_state=seed, stratify=y)
    
    model = LightCNN_32ch(in_ch=32, n_cls=n_cls, width=128).to(device)
    cc = np.bincount(y_tr, minlength=n_cls)
    w = torch.FloatTensor((1.0/(cc+1)) / (1.0/(cc+1)).sum() * n_cls).to(device)
    
    crit = nn.CrossEntropyLoss(weight=w, label_smoothing=0.1)
    opt = optim.AdamW(model.parameters(), lr=4e-4, weight_decay=0.02)
    sched = optim.lr_scheduler.CosineAnnealingWarmRestarts(opt, T_0=40, T_mult=2, eta_min=1e-6)
    
    tr_t = torch.FloatTensor(X_tr)
    tr_y = torch.LongTensor(y_tr)
    te_t = torch.FloatTensor(X_te).to(device)
    
    best_acc, best_state, pat = 0, None, 0
    for ep in range(300):
        model.train()
        perm = torch.randperm(len(tr_t))
        for i in range(0, len(tr_t), 32):
            idx = perm[i:i+32]
            bx = tr_t[idx].to(device)
            by = tr_y[idx].to(device)
            if np.random.random() > 0.5: bx = torch.flip(bx, [2])
            if np.random.random() > 0.5: bx = torch.flip(bx, [3])
            if np.random.random() > 0.5: bx = torch.rot90(bx, np.random.randint(1,4), [2,3])
            bx = bx + torch.randn_like(bx) * 0.02
            
            # Mixup
            if np.random.random() > 0.5 and len(bx) > 1:
                lam = np.random.beta(0.4, 0.4)
                i2 = torch.randperm(bx.size(0))
                bx = lam*bx + (1-lam)*bx[i2]
                oh1 = torch.zeros(by.size(0), n_cls, device=device).scatter_(1, by.unsqueeze(1), 1)
                oh2 = torch.zeros(by.size(0), n_cls, device=device).scatter_(1, by[i2].unsqueeze(1), 1)
                out = model(bx)
                loss = (-(lam*oh1 + (1-lam)*oh2) * torch.log_softmax(out,1)).sum(1).mean()
            else:
                loss = crit(model(bx), by)
                
            opt.zero_grad(); loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
        sched.step()
        
        model.eval()
        with torch.no_grad():
            probs = []
            for fn in [lambda x:x, lambda x:torch.flip(x,[2]), lambda x:torch.flip(x,[3])]:
                probs.append(torch.softmax(model(fn(te_t)), 1))
            avg = torch.stack(probs).mean(0)
            preds = avg.argmax(1).cpu().numpy()
            
        acc = accuracy_score(y_te, preds)
        if acc > best_acc:
            best_acc = acc; best_state = {k:v.cpu().clone() for k,v in model.state_dict().items()}; pat = 0
            print(f"  Ep {ep+1} Fusion-Acc={acc:.4f} 🌟")
        else:
            pat += 1
        if pat >= 60: break
    
    model.load_state_dict(best_state)
    return model, best_acc

def train_hybrid_fusion(X, y, cnn_model, n_cls):
    print("\n" + "="*60)
    print("HYBRID FUSION: CNN embed + S1/S2 Rich features + XGBoost")
    print("="*60)
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    cnn_model = cnn_model.to(device).eval()
    
    with torch.no_grad():
        embs = []
        for i in range(0, len(X), 64):
            b = torch.FloatTensor(X[i:i+64]).to(device)
            embs.append(cnn_model.embed(b).cpu().numpy())
    cnn_feat = np.concatenate(embs)
    
    rich = extract_features_fusion(X)
    combined = np.concatenate([cnn_feat, rich], axis=1)
    print(f"  Final Feature Vector: {combined.shape}")
    
    scaler = StandardScaler()
    combined = scaler.fit_transform(combined)
    
    X_tr, X_te, y_tr, y_te = train_test_split(combined, y, test_size=0.2, random_state=42, stratify=y)
    
    xgb = XGBClassifier(n_estimators=1500, max_depth=7, learning_rate=0.02,
                         subsample=0.8, colsample_bytree=0.5, min_child_weight=3,
                         tree_method='hist', device='cuda',
                         random_state=42, use_label_encoder=False, eval_metric='mlogloss')
    xgb.fit(X_tr, y_tr, eval_set=[(X_te, y_te)], verbose=False)
    acc = accuracy_score(y_te, xgb.predict(X_te))
    print(f"  ✅ Hybrid Fusion Acc: {acc:.4f}")
    
    # K-Fold CV
    skf = StratifiedKFold(5, shuffle=True, random_state=42)
    cv_accs = []
    for fold, (ti, vi) in enumerate(skf.split(combined, y)):
        m = XGBClassifier(n_estimators=1500, max_depth=7, learning_rate=0.02,
                          subsample=0.8, colsample_bytree=0.5,
                          tree_method='hist', device='cuda',
                          random_state=42, use_label_encoder=False, eval_metric='mlogloss')
        m.fit(combined[ti], y[ti], eval_set=[(combined[vi], y[vi])], verbose=False)
        a = accuracy_score(y[vi], m.predict(combined[vi]))
        cv_accs.append(a)
        print(f"    Fold {fold+1}: {a:.4f}")
        
    cv_mean = np.mean(cv_accs)
    print(f"  ✅ CV Mean: {cv_mean:.4f} ± {np.std(cv_accs):.4f}")
    return acc, cv_mean

def main():
    print("🚀 V4: TÍCH HỢP RADAR SENTINEL-1 (32-CHANNELS FUSION)")
    print("="*60)
    
    X, y, n_cls = load_and_clean()
    
    cnn_model, cnn_acc = train_cnn_fusion(X, y, n_cls)
    print(f"\n✅ CNN Fusion best: {cnn_acc:.4f}")
    
    hyb_acc, hyb_cv = train_hybrid_fusion(X, y, cnn_model, n_cls)
    
    print("\n" + "="*60)
    print("📊 FINAL RESULTS V4 (WITH RADAR)")
    print("="*60)
    res = {
        'CNN Fusion (32ch)': cnn_acc,
        'Hybrid Fusion (CNN+XGB)': hyb_acc,
        'Hybrid Fusion CV': hyb_cv,
    }
    for n, a in sorted(res.items(), key=lambda x:-x[1]):
        mk = "🏆" if a>=0.95 else "✅" if a>=0.90 else "📈"
        print(f"  {mk} {n}: {a:.4f}")
        
    best = max(res.values())
    if best >= 0.95:
        print(f"\n🎉 THÀNH CÔNG VƯỢT MỐC 95%! BEST: {best:.4f}")
    else:
        print(f"\n🏆 BEST: {best:.4f}")
        
    os.makedirs('model_train', exist_ok=True)
    with open('model_train/ultimate_v4_fusion_results.json', 'w') as f:
        json.dump({k:float(v) for k,v in res.items()}, f, indent=2)

if __name__ == "__main__":
    main()