"""
V6: Exhaustive Hyperparameter Tuning for Maximum Accuracy
- Multi-seed CNN ensembles for better embeddings
- Optuna-style manual grid search on XGBoost/LightGBM/ExtraTrees
- Stacking instead of simple Voting
- Feature selection to remove noise
"""
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, RepeatedStratifiedKFold
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import StandardScaler
from sklearn.feature_selection import SelectKBest, f_classif
from xgboost import XGBClassifier
from lightgbm import LGBMClassifier
from sklearn.ensemble import ExtraTreesClassifier, StackingClassifier, RandomForestClassifier, GradientBoostingClassifier
from sklearn.linear_model import LogisticRegression
from scipy.ndimage import uniform_filter
import itertools
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 = (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"Data: {X.shape}, {len(unique)} classes, dist={[int((y==i).sum()) for i in range(len(unique))]}")
    return X, y, len(unique)

def extract_features_fusion(X):
    N = X.shape[0]
    all_feats = []
    for i in range(N):
        patch = X[i]
        feats = []
        valid_ts = [t for t in range(4) if np.abs(patch[t*8:t*8+6]).sum() > 1e-6]
        if not valid_ts: valid_ts = [0]
        
        # S2 per-band stats
        for t in valid_ts:
            for b in range(6):
                ch = patch[t*8 + b]
                feats.extend([np.mean(ch), np.std(ch), np.median(ch), np.min(ch), np.max(ch),
                              np.percentile(ch, 10), np.percentile(ch, 25), np.percentile(ch, 75), np.percentile(ch, 90),
                              np.mean(ch > np.mean(ch))])
        # Pad to fixed length (4 timesteps * 6 bands * 10 stats = 240)
        needed = 4 * 6 * 10
        feats.extend([0.0] * (needed - len(feats)))
        
        # S1 per-band stats + ratios
        for t in range(4):
            vv, vh = patch[t*8+6], patch[t*8+7]
            if np.abs(vv).sum() > 1e-6:
                ratio = (vh+1e-6)/(vv+1e-6)
                diff = vv - vh
                feats.extend([np.mean(vv), np.std(vv), np.median(vv), np.max(vv), np.percentile(vv, 90),
                              np.mean(vh), np.std(vh), np.median(vh), np.max(vh), np.percentile(vh, 90),
                              np.mean(ratio), np.std(ratio), np.median(ratio), np.min(ratio), np.max(ratio),
                              np.mean(diff), np.std(diff)])
            else:
                feats.extend([0.0] * 17)
        
        # Temporal variance (S2)
        for b in range(6):
            ts_means = [np.mean(patch[t*8+b]) for t in valid_ts]
            feats.extend([np.std(ts_means) if len(ts_means) > 1 else 0.0,
                          np.max(ts_means) - np.min(ts_means) if len(ts_means) > 1 else 0.0])
        
        # Temporal variance (S1)
        for b_offset in [6, 7]:
            ts_means = [np.mean(patch[t*8+b_offset]) for t in range(4) if np.abs(patch[t*8+b_offset]).sum() > 1e-6]
            feats.extend([np.std(ts_means) if len(ts_means) > 1 else 0.0,
                          np.max(ts_means) - np.min(ts_means) if len(ts_means) > 1 else 0.0])
        
        # Spatial texture
        for t in valid_ts[:2]:
            for b_idx in [3, 4, 6, 7]:  # NIR, NDVI, VV, VH
                ch = patch[t*8 + b_idx] if b_idx < 6 else patch[valid_ts[0]*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
                entropy_approx = -np.mean(np.abs(lv) * np.log(np.abs(lv) + 1e-10))
                feats.extend([grad_mag, np.mean(lv), np.std(lv), entropy_approx])
        needed_tex = 2 * 4 * 4
        got_tex = min(len(valid_ts), 2) * 4 * 4
        feats.extend([0.0] * (needed_tex - got_tex))
        
        # Flat pixels from best timestep
        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)
    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(X, y, n_cls, seed=42, width=128, epochs=300):
    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=width).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(epochs):
        model.train()
        perm = torch.randperm(len(tr_t))
        for i in range(0, len(tr_t), 32):
            idx = perm[i:i+32]
            bx, by = tr_t[idx].to(device), 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
            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)
                loss = (-(lam*oh1 + (1-lam)*oh2) * torch.log_softmax(model(bx),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 = [torch.softmax(model(fn(te_t)), 1) for fn in [lambda x:x, lambda x:torch.flip(x,[2]), lambda x:torch.flip(x,[3])]]
            preds = torch.stack(probs).mean(0).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 >= 60: break
    model.load_state_dict(best_state)
    return model, best_acc

def get_cnn_embeddings(X, models, device):
    all_embs = []
    for model in models:
        model = 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(model.embed(b).cpu().numpy())
        all_embs.append(np.concatenate(embs))
    return np.concatenate(all_embs, axis=1)

def run_hyperparameter_search(combined, y, n_cls):
    print("\n" + "="*60)
    print("🔬 EXHAUSTIVE HYPERPARAMETER SEARCH")
    print("="*60)
    
    scaler = StandardScaler()
    combined_scaled = scaler.fit_transform(combined)
    
    skf = StratifiedKFold(5, shuffle=True, random_state=42)
    
    # ===== CONFIG SPACE =====
    configs = [
        # Config 1: XGB Deep trees
        {"name": "XGB-deep", "model": lambda: XGBClassifier(
            n_estimators=2000, max_depth=9, learning_rate=0.01, subsample=0.75, colsample_bytree=0.4,
            min_child_weight=2, gamma=0.1, reg_alpha=0.5, reg_lambda=1.5,
            tree_method='hist', device='cuda', random_state=42, use_label_encoder=False, eval_metric='mlogloss')},
        # Config 2: XGB Shallow wide
        {"name": "XGB-shallow", "model": lambda: XGBClassifier(
            n_estimators=3000, max_depth=5, learning_rate=0.008, subsample=0.85, colsample_bytree=0.35,
            min_child_weight=5, gamma=0.2, reg_alpha=1.0, reg_lambda=2.0,
            tree_method='hist', device='cuda', random_state=42, use_label_encoder=False, eval_metric='mlogloss')},
        # Config 3: XGB Balanced
        {"name": "XGB-balanced", "model": lambda: XGBClassifier(
            n_estimators=2500, max_depth=7, learning_rate=0.015, subsample=0.8, colsample_bytree=0.45,
            min_child_weight=3, gamma=0.05, reg_alpha=0.3, reg_lambda=1.0,
            tree_method='hist', device='cuda', random_state=42, use_label_encoder=False, eval_metric='mlogloss')},
        # Config 4: LGBM Tuned
        {"name": "LGBM-tuned", "model": lambda: LGBMClassifier(
            n_estimators=2000, max_depth=8, learning_rate=0.015, subsample=0.8, colsample_bytree=0.45,
            min_child_samples=5, reg_alpha=0.5, reg_lambda=1.0, num_leaves=63,
            random_state=42, verbosity=-1)},
        # Config 5: LGBM Conservative
        {"name": "LGBM-conservative", "model": lambda: LGBMClassifier(
            n_estimators=3000, max_depth=6, learning_rate=0.008, subsample=0.75, colsample_bytree=0.35,
            min_child_samples=10, reg_alpha=1.0, reg_lambda=2.0, num_leaves=31,
            random_state=42, verbosity=-1)},
        # Config 6: ExtraTrees Deep
        {"name": "ETC-deep", "model": lambda: ExtraTreesClassifier(
            n_estimators=2000, max_depth=20, max_features='sqrt', min_samples_leaf=2,
            random_state=42, n_jobs=-1)},
        # Config 7: RandomForest
        {"name": "RF-tuned", "model": lambda: RandomForestClassifier(
            n_estimators=2000, max_depth=15, max_features='sqrt', min_samples_leaf=3,
            random_state=42, n_jobs=-1)},
        # Config 8: GradientBoosting (sklearn)
        {"name": "GBT-sklearn", "model": lambda: GradientBoostingClassifier(
            n_estimators=500, max_depth=5, learning_rate=0.05, subsample=0.8,
            min_samples_leaf=5, random_state=42)},
    ]
    
    results = {}
    for cfg in configs:
        cv_accs = []
        for fold, (ti, vi) in enumerate(skf.split(combined_scaled, y)):
            m = cfg["model"]()
            if hasattr(m, 'eval_set'):
                m.fit(combined_scaled[ti], y[ti], eval_set=[(combined_scaled[vi], y[vi])], verbose=False)
            else:
                m.fit(combined_scaled[ti], y[ti])
            a = accuracy_score(y[vi], m.predict(combined_scaled[vi]))
            cv_accs.append(a)
        mean_acc = np.mean(cv_accs)
        results[cfg["name"]] = (mean_acc, np.std(cv_accs), cv_accs)
        mk = "🏆" if mean_acc >= 0.95 else "✅" if mean_acc >= 0.93 else "📈"
        print(f"  {mk} {cfg['name']}: {mean_acc:.4f} ± {np.std(cv_accs):.4f}  (folds: {[f'{a:.3f}' for a in cv_accs]})")
    
    # ===== STACKING ENSEMBLE =====
    print("\n--- Stacking Ensemble ---")
    
    best_3 = sorted(results.items(), key=lambda x: -x[1][0])[:3]
    print(f"  Top-3 base models: {[b[0] for b in best_3]}")
    
    # Build stacking with top models
    base_estimators = []
    for cfg in configs:
        if cfg["name"] in [b[0] for b in best_3]:
            base_estimators.append((cfg["name"], cfg["model"]()))
    
    stacking_configs = [
        {"name": "Stack-LR", "meta": LogisticRegression(C=1.0, max_iter=1000, random_state=42)},
        {"name": "Stack-XGB", "meta": XGBClassifier(n_estimators=200, max_depth=3, learning_rate=0.1,
                                                      tree_method='hist', device='cuda', random_state=42,
                                                      use_label_encoder=False, eval_metric='mlogloss')},
    ]
    
    for scfg in stacking_configs:
        stack = StackingClassifier(estimators=base_estimators, final_estimator=scfg["meta"],
                                   cv=3, stack_method='predict_proba', n_jobs=-1)
        cv_accs = []
        for fold, (ti, vi) in enumerate(skf.split(combined_scaled, y)):
            stack_clone = StackingClassifier(estimators=[(n, cfg["model"]()) for cfg in configs for n in [cfg["name"]] if n in [b[0] for b in best_3]],
                                             final_estimator=scfg["meta"], cv=3, stack_method='predict_proba', n_jobs=-1)
            stack_clone.fit(combined_scaled[ti], y[ti])
            a = accuracy_score(y[vi], stack_clone.predict(combined_scaled[vi]))
            cv_accs.append(a)
        mean_acc = np.mean(cv_accs)
        results[scfg["name"]] = (mean_acc, np.std(cv_accs), cv_accs)
        mk = "🏆" if mean_acc >= 0.95 else "✅" if mean_acc >= 0.93 else "📈"
        print(f"  {mk} {scfg['name']}: {mean_acc:.4f} ± {np.std(cv_accs):.4f}  (folds: {[f'{a:.3f}' for a in cv_accs]})")
    
    # ===== FEATURE SELECTION + BEST MODEL =====
    print("\n--- Feature Selection ---")
    for k_feat in [500, 800, 1200, 1500, 2000]:
        selector = SelectKBest(f_classif, k=min(k_feat, combined_scaled.shape[1]))
        X_sel = selector.fit_transform(combined_scaled, y)
        cv_accs = []
        for fold, (ti, vi) in enumerate(skf.split(X_sel, y)):
            m = XGBClassifier(n_estimators=2500, max_depth=7, learning_rate=0.015, subsample=0.8, colsample_bytree=0.45,
                              min_child_weight=3, gamma=0.05, reg_alpha=0.3, reg_lambda=1.0,
                              tree_method='hist', device='cuda', random_state=42, use_label_encoder=False, eval_metric='mlogloss')
            m.fit(X_sel[ti], y[ti])
            a = accuracy_score(y[vi], m.predict(X_sel[vi]))
            cv_accs.append(a)
        mean_acc = np.mean(cv_accs)
        mk = "🏆" if mean_acc >= 0.95 else "✅" if mean_acc >= 0.93 else "📈"
        print(f"  {mk} XGB k={k_feat}: {mean_acc:.4f} ± {np.std(cv_accs):.4f}  (folds: {[f'{a:.3f}' for a in cv_accs]})")
        results[f"XGB-feat{k_feat}"] = (mean_acc, np.std(cv_accs), cv_accs)
    
    return results

def main():
    print("🚀 V6: EXHAUSTIVE HYPERPARAMETER TUNING")
    print("="*60)
    
    X, y, n_cls = load_and_clean()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # Train multi-seed CNN ensemble for richer embeddings
    print("\n--- Training Multi-Seed CNN Ensemble ---")
    models = []
    for seed in [42, 123, 777]:
        m, acc = train_cnn(X, y, n_cls, seed=seed, width=128)
        print(f"  Seed {seed}: CNN Acc = {acc:.4f}")
        models.append(m)
    
    # Get combined embeddings from all CNN seeds
    cnn_feat = get_cnn_embeddings(X, models, device)
    print(f"  Multi-seed CNN embedding: {cnn_feat.shape}")
    
    rich = extract_features_fusion(X)
    combined = np.concatenate([cnn_feat, rich], axis=1)
    print(f"  Total features: {combined.shape}")
    
    results = run_hyperparameter_search(combined, y, n_cls)
    
    # Final summary
    print("\n" + "="*60)
    print("📊 LEADERBOARD")
    print("="*60)
    sorted_results = sorted(results.items(), key=lambda x: -x[1][0])
    for rank, (name, (mean, std, folds)) in enumerate(sorted_results, 1):
        mk = "🏆" if mean >= 0.95 else "✅" if mean >= 0.93 else "📈"
        print(f"  #{rank} {mk} {name}: {mean:.4f} ± {std:.4f}")
    
    best_name, (best_mean, best_std, best_folds) = sorted_results[0]
    print(f"\n🏆 CHAMPION: {best_name} = {best_mean:.4f}")
    if best_mean >= 0.95:
        print("🎉 VƯỢT MỐC 95%!")
    
    os.makedirs('model_train', exist_ok=True)
    with open('model_train/v6_tuning_results.json', 'w') as f:
        json.dump({k: {"mean": float(v[0]), "std": float(v[1]), "folds": [float(x) for x in v[2]]} for k, v in results.items()}, f, indent=2)

if __name__ == "__main__":
    main()