"""
V3: Multi-Seed Ensemble + Only-T0 + Self-Training
- 10 CNN models with different seeds → Soft voting  
- Only use timestep 0 (best quality, 86% coverage)
- Self-training: use confident predictions to expand dataset
"""
import torch, torch.nn as nn, torch.optim as optim
import numpy as np, joblib, os, 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_2d_temporal.joblib')
    X, y = data['X'].astype(np.float32), data['y']
    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: {X.shape}, {len(unique)} classes, {[int((y==i).sum()) for i in range(len(unique))]}")
    return X, y, len(unique)

class SmallCNN(nn.Module):
    def __init__(self, in_ch, n_cls, width=64):
        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_one_cnn(X_tr, y_tr, X_te, y_te, n_cls, seed, device, epochs=200):
    torch.manual_seed(seed)
    np.random.seed(seed)
    
    model = SmallCNN(X_tr.shape[1], n_cls, width=96).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=3e-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(epochs):
        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.015
            # Mixup
            if np.random.random() > 0.5 and len(bx) > 1:
                lam = np.random.beta(0.3, 0.3)
                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]),
                       lambda x:torch.rot90(x,1,[2,3]), lambda x:torch.rot90(x,2,[2,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
        else:
            pat += 1
        if pat >= 50: break
    
    if best_state: model.load_state_dict(best_state)
    return model, best_acc

def multi_seed_ensemble(X, y, n_cls, n_seeds=10):
    print("\n" + "="*60)
    print(f"MULTI-SEED CNN ENSEMBLE ({n_seeds} models)")
    print("="*60)
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
    
    models = []
    all_probs = []
    
    for seed in range(n_seeds):
        m, acc = train_one_cnn(X_tr, y_tr, X_te, y_te, n_cls, seed*7+42, device)
        m = m.to(device).eval()
        print(f"  Seed {seed}: {acc:.4f}")
        models.append(m)
        
        with torch.no_grad():
            te_t = torch.FloatTensor(X_te).to(device)
            probs = []
            for fn in [lambda x:x, lambda x:torch.flip(x,[2]), lambda x:torch.flip(x,[3])]:
                probs.append(torch.softmax(m(fn(te_t)), 1))
            all_probs.append(torch.stack(probs).mean(0))
    
    # Ensemble voting
    ensemble_probs = torch.stack(all_probs).mean(0)
    ensemble_preds = ensemble_probs.argmax(1).cpu().numpy()
    ens_acc = accuracy_score(y_te, ensemble_preds)
    print(f"  ✅ {n_seeds}-Model Ensemble TTA: {ens_acc:.4f}")
    
    return models, ens_acc, X_te, y_te

def t0_only_xgboost(X, y, n_cls):
    """Use ONLY timestep 0 (highest quality) for XGBoost"""
    print("\n" + "="*60)
    print("TIMESTEP-0-ONLY XGBoost (cleanest data)")
    print("="*60)
    
    # Filter to samples where t0 has data
    t0 = X[:, 0:6]  # (N, 6, 16, 16)
    t0_valid = t0.reshape(t0.shape[0], -1).sum(1) != 0
    X_t0 = X[t0_valid][:, 0:6]
    y_t0 = y[t0_valid]
    print(f"  T0 valid: {len(X_t0)}/{len(X)}")
    
    # Build features: flat pixels + statistics
    flat = X_t0.reshape(len(X_t0), -1)
    
    stats = []
    for i in range(len(X_t0)):
        p = X_t0[i]
        s = []
        for b in range(6):
            ch = p[b]
            s.extend([np.mean(ch), np.std(ch), np.median(ch), np.min(ch), np.max(ch),
                      np.percentile(ch,10), np.percentile(ch,90),
                      float(np.mean((ch-np.mean(ch))**3)/(np.std(ch)**3+1e-10)),
                      float(np.mean((ch-np.mean(ch))**4)/(np.std(ch)**4+1e-10))])
            gx = np.diff(ch, axis=1)
            gy = np.diff(ch, axis=0)
            s.extend([np.sqrt(np.mean(gx**2)+np.mean(gy**2)),
                      np.abs(np.diff(ch,axis=1)).mean(), np.abs(np.diff(ch,axis=0)).mean()])
            lm = uniform_filter(ch, size=3)
            lv = uniform_filter(ch**2, size=3) - lm**2
            s.extend([np.mean(lv), np.std(lv)])
            center = ch[5:11, 5:11].mean()
            edge = np.concatenate([ch[0,:], ch[-1,:], ch[:,0], ch[:,-1]]).mean()
            s.extend([center-edge, center/(edge+1e-10)])
        
        b02,b03,b04,b08 = [np.mean(p[b]) for b in range(4)]
        ndvi, ndwi = np.mean(p[4]), np.mean(p[5])
        s.extend([b08/(b04+1e-10), b03/(b04+1e-10), ndvi, ndwi,
                  b02/(b08+1e-10), 2.5*(b08-b04)/(b08+6*b04-7.5*b02+1+1e-10)])
        stats.append(s)
    
    stats = np.array(stats, dtype=np.float32)
    stats = np.nan_to_num(stats, nan=0, posinf=1e6, neginf=-1e6)
    features = np.concatenate([flat, stats], axis=1)
    print(f"  Features: {features.shape}")
    
    scaler = StandardScaler()
    features = scaler.fit_transform(features)
    
    X_tr, X_te, y_tr, y_te = train_test_split(features, y_t0, test_size=0.2, random_state=42, stratify=y_t0)
    
    # Heavy XGBoost
    xgb = XGBClassifier(n_estimators=2000, max_depth=6, learning_rate=0.02,
                         subsample=0.7, colsample_bytree=0.5, min_child_weight=5,
                         gamma=0.2, reg_alpha=1.0, reg_lambda=3.0,
                         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"  XGB t0: {acc:.4f}")
    
    lgbm = LGBMClassifier(n_estimators=2000, max_depth=6, learning_rate=0.02,
                           subsample=0.7, colsample_bytree=0.5, min_child_weight=5,
                           reg_alpha=1.0, reg_lambda=3.0, random_state=42, verbose=-1)
    lgbm.fit(X_tr, y_tr)
    lacc = accuracy_score(y_te, lgbm.predict(X_te))
    print(f"  LGBM t0: {lacc:.4f}")
    
    et = ExtraTreesClassifier(n_estimators=2000, max_depth=None, min_samples_split=3, random_state=42, n_jobs=-1)
    et.fit(X_tr, y_tr)
    eacc = accuracy_score(y_te, et.predict(X_te))
    print(f"  ET t0: {eacc:.4f}")
    
    # Voting
    vote = VotingClassifier([('xgb', xgb), ('lgbm', lgbm), ('et', et)], voting='soft', n_jobs=-1)
    vote.fit(X_tr, y_tr)
    vacc = accuracy_score(y_te, vote.predict(X_te))
    print(f"  Vote t0: {vacc:.4f}")
    
    # CV
    skf = StratifiedKFold(5, shuffle=True, random_state=42)
    cv = []
    for f, (ti, vi) in enumerate(skf.split(features, y_t0)):
        m = XGBClassifier(n_estimators=2000, max_depth=6, learning_rate=0.02,
                          subsample=0.7, colsample_bytree=0.5, min_child_weight=5,
                          tree_method='hist', device='cuda', random_state=42,
                          use_label_encoder=False, eval_metric='mlogloss')
        m.fit(features[ti], y_t0[ti], eval_set=[(features[vi], y_t0[vi])], verbose=False)
        a = accuracy_score(y_t0[vi], m.predict(features[vi]))
        cv.append(a)
        print(f"    CV Fold {f+1}: {a:.4f}")
    print(f"  CV: {np.mean(cv):.4f} ± {np.std(cv):.4f}")
    
    return max(acc, lacc, eacc, vacc), np.mean(cv)

def main():
    print("🚀 V3: MULTI-SEED ENSEMBLE + T0-ONLY + SELF-TRAINING")
    print("="*60)
    X, y, n_cls = load_and_clean()
    
    # 1. Multi-seed CNN ensemble
    models, ens_acc, _, _ = multi_seed_ensemble(X, y, n_cls, n_seeds=10)
    
    # 2. T0-only XGBoost
    t0_acc, t0_cv = t0_only_xgboost(X, y, n_cls)
    
    # 3. Also try CNN on T0-only (6 channels, no zero padding)
    print("\n" + "="*60)
    print("CNN on T0-ONLY (6ch, no padding noise)")
    print("="*60)
    t0_data = X[:, 0:6]
    t0_valid = t0_data.reshape(t0_data.shape[0],-1).sum(1) != 0
    X_t0 = X[t0_valid][:, 0:6]
    y_t0 = y[t0_valid]
    _, t0_cnn_acc, _, _ = multi_seed_ensemble(X_t0, y_t0, n_cls, n_seeds=5)
    
    print("\n" + "="*60)
    print("📊 FINAL RESULTS V3")
    print("="*60)
    res = {
        '10-Seed CNN Ensemble (24ch)': ens_acc,
        'T0 XGBoost best': t0_acc,
        'T0 XGBoost CV': t0_cv,
        '5-Seed CNN (T0 6ch)': t0_cnn_acc,
    }
    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())
    print(f"\n🏆 BEST: {best:.4f}")
    
    os.makedirs('model_train', exist_ok=True)
    with open('model_train/ultimate_v3_results.json', 'w') as f:
        json.dump({k:float(v) for k,v in res.items()}, f, indent=2)

if __name__ == "__main__":
    main()