"""
CHIẾN LƯỢC V2: Tập trung vào timestep 0 (chất lượng tốt nhất)
+ Pixel-level XGBoost + Spatial features + Stacking
+ CNN với masking zeros
+ TTA (Test-Time Augmentation)
"""
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import joblib
import os, json
from sklearn.model_selection import StratifiedKFold, train_test_split
from sklearn.metrics import accuracy_score, classification_report
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import (
    RandomForestClassifier, GradientBoostingClassifier,
    StackingClassifier, VotingClassifier, ExtraTreesClassifier
)
from xgboost import XGBClassifier
from lightgbm import LGBMClassifier
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 data: {X.shape}, {len(unique)} classes")
    return X, y, len(unique)

def extract_features_v2(X):
    """
    Chiến lược mới: Chỉ dùng timestep có dữ liệu thật.
    Tính features theo từng timestep rồi lấy max/mean/std qua thời gian.
    """
    N = X.shape[0]
    all_feats = []
    
    for i in range(N):
        patch = X[i]  # (24, 16, 16)
        feats = []
        
        # Xác định timestep nào có dữ liệu (không phải toàn zero)
        valid_ts = []
        for t in range(4):
            block = patch[t*6:(t+1)*6]
            if np.abs(block).sum() > 1e-6:
                valid_ts.append(t)
        
        if not valid_ts:
            valid_ts = [0]
        
        # === A. Per-valid-timestep features ===
        per_ts_stats = {b: [] for b in range(6)}
        
        for t in valid_ts:
            for b in range(6):
                ch = patch[t*6 + b]
                per_ts_stats[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),
                ])
        
        # Aggregate across valid timesteps
        for b in range(6):
            stats = np.array(per_ts_stats[b])
            feats.extend(stats.mean(axis=0).tolist())  # Mean of stats
            feats.extend(stats.std(axis=0).tolist())    # Variability of stats
            if len(stats) > 1:
                feats.extend((stats[-1] - stats[0]).tolist())  # Trend
            else:
                feats.extend([0.0]*7)
        
        # === B. Band ratios (averaged over valid timesteps) ===
        ratio_lists = {k: [] for k in ['nir_red', 'grn_red', 'ndvi', 'ndwi', 'blu_nir', 'evi']}
        for t in valid_ts:
            b02 = np.mean(patch[t*6+0]) + 1e-10
            b03 = np.mean(patch[t*6+1]) + 1e-10
            b04 = np.mean(patch[t*6+2]) + 1e-10
            b08 = np.mean(patch[t*6+3]) + 1e-10
            ratio_lists['nir_red'].append(b08/b04)
            ratio_lists['grn_red'].append(b03/b04)
            ratio_lists['ndvi'].append((b08-b04)/(b08+b04))
            ratio_lists['ndwi'].append((b03-b08)/(b03+b08))
            ratio_lists['blu_nir'].append(b02/b08)
            ratio_lists['evi'].append(2.5*(b08-b04)/(b08+6*b04-7.5*b02+1+1e-10))
        
        for k, v in ratio_lists.items():
            v = np.array(v)
            feats.extend([v.mean(), v.std(), v.max()-v.min()])
        
        # === C. Spatial texture features (B08 and NDVI only) ===
        for t in valid_ts[:2]:  # max 2 timesteps
            for b_idx in [3, 4]:
                ch = patch[t*6 + b_idx]
                # Gradient
                gx = np.diff(ch, axis=1)
                gy = np.diff(ch, axis=0)
                grad_mag = np.sqrt(np.mean(gx**2) + np.mean(gy**2))
                # Local variance
                lm = uniform_filter(ch, size=3)
                lv = uniform_filter(ch**2, size=3) - lm**2
                # GLCM-like: pixel value differences
                h_diff = np.abs(np.diff(ch, axis=1)).mean()
                v_diff = np.abs(np.diff(ch, axis=0)).mean()
                # Homogeneity
                feats.extend([
                    grad_mag, np.mean(lv), np.std(lv),
                    h_diff, v_diff,
                    np.mean(np.abs(ch - np.mean(ch))),  # MAD
                ])
        # Pad if fewer valid timesteps
        needed = 2 * 2 * 6
        got = min(len(valid_ts), 2) * 2 * 6
        feats.extend([0.0] * (needed - got))
        
        # === D. Center vs edge ===
        for t in valid_ts[:2]:
            for b_idx in [3, 4]:
                ch = patch[t*6 + b_idx]
                center = ch[5:11, 5:11].mean()
                edge = np.concatenate([ch[0,:], ch[-1,:], ch[:,0], ch[:,-1]]).mean()
                feats.extend([center - edge, center / (edge + 1e-10)])
        needed_d = 2 * 2 * 2
        got_d = min(len(valid_ts), 2) * 2 * 2
        feats.extend([0.0] * (needed_d - got_d))
        
        # === E. Number of valid timesteps as feature ===
        feats.append(len(valid_ts))
        
        # === F. Flat pixel features from best timestep (t=0) ===
        best_t = valid_ts[0]
        for b in range(6):
            ch = patch[best_t*6 + 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]} features per sample")
    return features

class LightCNN(nn.Module):
    def __init__(self, in_ch=24, n_cls=7):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(in_ch, 96, 3, padding=1), nn.BatchNorm2d(96), nn.GELU(),
            nn.Conv2d(96, 96, 3, padding=1), nn.BatchNorm2d(96), nn.GELU(),
            nn.MaxPool2d(2), nn.Dropout2d(0.1),
            
            nn.Conv2d(96, 192, 3, padding=1), nn.BatchNorm2d(192), nn.GELU(),
            nn.Conv2d(192, 192, 3, padding=1), nn.BatchNorm2d(192), nn.GELU(),
            nn.MaxPool2d(2), nn.Dropout2d(0.1),
            
            nn.Conv2d(192, 384, 3, padding=1), nn.BatchNorm2d(384), nn.GELU(),
            nn.Conv2d(384, 384, 3, padding=1), nn.BatchNorm2d(384), nn.GELU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.head = nn.Sequential(
            nn.Flatten(), nn.Linear(384, 192), nn.GELU(),
            nn.Dropout(0.5), nn.Linear(192, n_cls)
        )
        self.embed = nn.Sequential(nn.Flatten())
    
    def get_embedding(self, x):
        return self.embed(self.features(x))
    
    def forward(self, x):
        return self.head(self.features(x))

def train_cnn_with_tta(X, y, n_cls):
    print("\n" + "="*60)
    print("CNN + TTA (Test-Time Augmentation)")
    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)
    
    model = LightCNN(in_ch=X.shape[1], n_cls=n_cls).to(device)
    
    cc = np.bincount(y_tr, minlength=n_cls)
    w = 1.0 / (cc + 1)
    w = torch.FloatTensor(w / w.sum() * n_cls).to(device)
    
    criterion = nn.CrossEntropyLoss(weight=w, label_smoothing=0.1)
    optimizer = optim.AdamW(model.parameters(), lr=5e-4, weight_decay=0.01)
    scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=30, 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 = 0
    best_state = None
    patience = 0
    
    for ep in range(300):
        model.train()
        perm = torch.randperm(len(tr_t))
        loss_sum = 0
        nb = 0
        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)
            
            # Augmentation
            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])
            if np.random.random() > 0.3: bx = bx + torch.randn_like(bx) * 0.01
            
            # Mixup
            if np.random.random() > 0.5:
                lam = np.random.beta(0.4, 0.4)
                idx2 = torch.randperm(bx.size(0))
                bx = lam * bx + (1 - lam) * bx[idx2]
                by_oh = torch.zeros(by.size(0), n_cls, device=device)
                by_oh.scatter_(1, by.unsqueeze(1), 1)
                by2_oh = torch.zeros(by.size(0), n_cls, device=device)
                by2_oh.scatter_(1, by[idx2].unsqueeze(1), 1)
                target_oh = lam * by_oh + (1 - lam) * by2_oh
                out = model(bx)
                loss = (-target_oh * torch.log_softmax(out, dim=1)).sum(dim=1).mean()
            else:
                out = model(bx)
                loss = criterion(out, by)
            
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            loss_sum += loss.item()
            nb += 1
        
        scheduler.step()
        
        # TTA evaluation
        model.eval()
        with torch.no_grad():
            preds_all = []
            for aug_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]),
            ]:
                out = model(aug_fn(te_t))
                preds_all.append(torch.softmax(out, dim=1))
            
            avg_pred = torch.stack(preds_all).mean(dim=0)
            preds = avg_pred.argmax(dim=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()}
            patience = 0
            print(f"  Ep {ep+1} Loss={loss_sum/nb:.4f} TTA-Acc={acc:.4f} 🌟")
            if acc >= 0.95:
                print("  🎯 >95% REACHED!")
                break
        else:
            patience += 1
            if (ep+1) % 30 == 0:
                print(f"  Ep {ep+1} Loss={loss_sum/nb:.4f} TTA-Acc={acc:.4f} (pat={patience})")
        
        if patience >= 80:
            print(f"  Early stop ep {ep+1}")
            break
    
    if best_state: model.load_state_dict(best_state)
    model = model.to(device)
    print(f"  ✅ CNN+TTA best: {best_acc:.4f}")
    return model, best_acc, X_te, y_te

def train_ensemble_v2(X, y, n_cls):
    print("\n" + "="*60)
    print("RICH FEATURES V2 + ENSEMBLE")
    print("="*60)
    
    feats = extract_features_v2(X)
    scaler = StandardScaler()
    feats = scaler.fit_transform(feats)
    
    X_tr, X_te, y_tr, y_te = train_test_split(feats, y, test_size=0.2, random_state=42, stratify=y)
    
    models = {
        'XGB': XGBClassifier(n_estimators=1000, max_depth=7, learning_rate=0.03,
                             subsample=0.8, colsample_bytree=0.6, min_child_weight=3,
                             gamma=0.1, reg_alpha=0.5, reg_lambda=2.0,
                             tree_method='hist', device='cuda',
                             random_state=42, use_label_encoder=False, eval_metric='mlogloss'),
        'LGBM': LGBMClassifier(n_estimators=1000, max_depth=7, learning_rate=0.03,
                               subsample=0.8, colsample_bytree=0.6, min_child_weight=3,
                               reg_alpha=0.5, reg_lambda=2.0, random_state=42, verbose=-1),
        'ET': ExtraTreesClassifier(n_estimators=1000, max_depth=None, min_samples_split=3,
                                   min_samples_leaf=1, random_state=42, n_jobs=-1),
        'RF': RandomForestClassifier(n_estimators=1000, max_depth=None, min_samples_split=3,
                                     min_samples_leaf=1, random_state=42, n_jobs=-1),
    }
    
    results = {}
    for name, m in models.items():
        if name in ['XGB']:
            m.fit(X_tr, y_tr, eval_set=[(X_te, y_te)], verbose=False)
        else:
            m.fit(X_tr, y_tr)
        acc = accuracy_score(y_te, m.predict(X_te))
        results[name] = acc
        print(f"  {name}: {acc:.4f}")
    
    # Soft voting
    vote = VotingClassifier([(n, m) for n, m in models.items()], voting='soft', n_jobs=-1)
    vote.fit(X_tr, y_tr)
    vacc = accuracy_score(y_te, vote.predict(X_te))
    results['Vote'] = vacc
    print(f"  Voting: {vacc:.4f}")
    
    # Cross-validate best
    print("\n  5-Fold CV:")
    skf = StratifiedKFold(5, shuffle=True, random_state=42)
    cv_accs = []
    for fold, (ti, vi) in enumerate(skf.split(feats, y)):
        m = XGBClassifier(n_estimators=1000, max_depth=7, learning_rate=0.03,
                          subsample=0.8, colsample_bytree=0.6, min_child_weight=3,
                          tree_method='hist', device='cuda',
                          random_state=42, use_label_encoder=False, eval_metric='mlogloss')
        m.fit(feats[ti], y[ti], eval_set=[(feats[vi], y[vi])], verbose=False)
        a = accuracy_score(y[vi], m.predict(feats[vi]))
        cv_accs.append(a)
        print(f"    Fold {fold+1}: {a:.4f}")
    cv_mean = np.mean(cv_accs)
    cv_std = np.std(cv_accs)
    print(f"  CV: {cv_mean:.4f} ± {cv_std:.4f}")
    
    return results, cv_mean, cv_std, feats, scaler

def train_hybrid_v2(X, y, cnn_model, n_cls):
    print("\n" + "="*60)
    print("HYBRID V2: CNN embed + 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.get_embedding(b).cpu().numpy())
    cnn_feat = np.concatenate(embs)
    
    rich = extract_features_v2(X)
    combined = np.concatenate([cnn_feat, rich], axis=1)
    print(f"  Combined: {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=1000, max_depth=8, learning_rate=0.03,
                         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 V2: {acc:.4f}")
    
    # 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=1000, max_depth=8, learning_rate=0.03,
                          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"  CV: {np.mean(cv_accs):.4f} ± {np.std(cv_accs):.4f}")
    
    return acc, np.mean(cv_accs)

def main():
    print("🚀 CHIẾN LƯỢC V2: TOÀN DIỆN ĐẠT >95%")
    print("="*60)
    
    X, y, n_cls = load_and_clean()
    
    # 1. CNN with TTA
    cnn_model, cnn_acc, _, _ = train_cnn_with_tta(X, y, n_cls)
    
    # 2. Rich features ensemble
    ens_results, cv_mean, cv_std, _, _ = train_ensemble_v2(X, y, n_cls)
    
    # 3. Hybrid
    hyb_acc, hyb_cv = train_hybrid_v2(X, y, cnn_model, n_cls)
    
    # Summary
    print("\n" + "="*60)
    print("📊 KẾT QUẢ TỔNG HỢP V2")
    print("="*60)
    all_res = {'CNN+TTA': cnn_acc, 'Hybrid V2': hyb_acc, 'Hybrid CV': hyb_cv, 'Ens CV': cv_mean}
    all_res.update({f'Ens_{k}': v for k, v in ens_results.items()})
    
    for n, a in sorted(all_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(all_res, key=all_res.get)
    print(f"\n🏆 BEST: {best} = {all_res[best]:.4f}")
    
    os.makedirs('model_train', exist_ok=True)
    with open('model_train/ultimate_v2_results.json', 'w') as f:
        json.dump({k: float(v) for k, v in all_res.items()}, f, indent=2)

if __name__ == "__main__":
    main()