"""
CHIẾN LƯỢC TOÀN DIỆN ĐẠT >95% ACCURACY
=========================================
Kết hợp 5 chiến lược song song:
1. Hybrid CNN+XGBoost: Trích xuất 2D features từ CNN nhẹ -> XGBoost
2. Rich Feature Engineering: Thống kê pixel + texture + temporal -> XGBoost  
3. Lightweight ResNet: ResNet-18 nhẹ, không upsample lãng phí
4. Stacking Ensemble: Kết hợp tất cả mô hình
5. StratifiedKFold: Cross-validation chống overfit
"""
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import torchvision.models as models
import torchvision.transforms as T

import joblib
import numpy as np
import os
import 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
import warnings
warnings.filterwarnings('ignore')

# ===== 1. LOAD DATA =====
def load_data():
    data = joblib.load('dataset_cache/training_data_2d_temporal.joblib')
    X, y = data['X'], data['y']
    X = X.astype(np.float32)
    print(f"Loaded data: X={X.shape}, y={y.shape}")
    print(f"Labels unique: {np.unique(y)}")
    
    # Remove invalid labels (label -1 = HT_code 0, which is invalid)
    valid_mask = y >= 0
    # Remove all-zero patches
    non_zero_mask = X.reshape(X.shape[0], -1).sum(axis=1) != 0
    mask = valid_mask & non_zero_mask
    X, y = X[mask], y[mask]
    print(f"After cleanup: X={X.shape}, y={y.shape} (removed {(~mask).sum()} bad samples)")
    
    # Remap labels to 0..N-1
    unique_labels = sorted(np.unique(y).tolist())
    label_map = {lbl: i for i, lbl in enumerate(unique_labels)}
    y_mapped = np.array([label_map[l] for l in y])
    print(f"Remapped labels: {np.unique(y_mapped)}")
    for lbl in np.unique(y_mapped):
        print(f"  Class {lbl}: {(y_mapped==lbl).sum()} samples")
    return X, y_mapped, len(unique_labels)

# ===== 2. RICH FEATURE ENGINEERING =====
def extract_rich_features(X):
    """
    Từ mỗi patch (24, 16, 16) trích xuất hàng trăm features thống kê.
    Channels: [B02,B03,B04,B08,NDVI,NDWI] x 4 timesteps
    """
    N = X.shape[0]
    all_features = []
    
    band_names = ['B02','B03','B04','B08','NDVI','NDWI']
    
    for i in range(N):
        patch = X[i]  # (24, 16, 16)
        feats = []
        
        # Per-channel statistics cho mỗi timestep
        for t in range(4):
            for b in range(6):
                ch = patch[t*6 + b]  # (16, 16)
                feats.extend([
                    np.mean(ch), np.std(ch), np.median(ch),
                    np.min(ch), np.max(ch),
                    np.percentile(ch, 25), np.percentile(ch, 75),
                    # Skewness và kurtosis
                    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)),
                    # Entropy approximation
                    float(-np.sum(np.abs(ch/np.sum(np.abs(ch)+1e-10)) * np.log(np.abs(ch/np.sum(np.abs(ch)+1e-10))+1e-10))),
                ])
        
        # Temporal change features: sự thay đổi giữa các timestep
        for b in range(6):
            vals_over_time = []
            for t in range(4):
                vals_over_time.append(np.mean(patch[t*6 + b]))
            vals = np.array(vals_over_time)
            feats.extend([
                np.std(vals),           # Temporal variability
                np.max(vals) - np.min(vals),  # Range over time
                vals[-1] - vals[0] if len(vals) > 1 else 0,  # Trend
                np.mean(np.abs(np.diff(vals))) if len(vals) > 1 else 0,  # Mean absolute change
            ])
        
        # Cross-band ratios (trung bình qua thời gian)
        for t in range(4):
            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
            feats.extend([
                b08/b04,  # NIR/Red ratio
                b03/b04,  # Green/Red ratio
                (b08-b04)/(b08+b04),  # NDVI recompute
                (b03-b08)/(b03+b08),  # NDWI recompute
                b02/b08,  # Blue/NIR
            ])
        
        # Spatial texture features (Gradient magnitude)
        for t in range(4):
            for b_idx in [3, 4]:  # B08 and NDVI
                ch = patch[t*6 + b_idx]
                # Sobel-like 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 (texture)
                from scipy.ndimage import uniform_filter
                local_mean = uniform_filter(ch, size=3)
                local_var = uniform_filter(ch**2, size=3) - local_mean**2
                feats.extend([
                    grad_mag,
                    np.mean(local_var),
                    np.std(local_var),
                ])
        
        # Center pixel vs edge pixels
        for t in range(4):
            for b_idx in [3, 4]:  # B08 and NDVI
                ch = patch[t*6 + b_idx]
                center = ch[6:10, 6:10].mean()
                edge = np.concatenate([ch[0,:], ch[-1,:], ch[:,0], ch[:,-1]]).mean()
                feats.append(center - edge)
        
        all_features.append(feats)
    
    features = np.array(all_features, dtype=np.float32)
    features = np.nan_to_num(features, nan=0.0, posinf=1e6, neginf=-1e6)
    print(f"Extracted {features.shape[1]} rich features per sample")
    return features

# ===== 3. LIGHTWEIGHT CNN =====
class LightCNN(nn.Module):
    """CNN nhẹ thiết kế riêng cho 16x16 patches - KHÔNG upsample"""
    def __init__(self, in_channels=24, num_classes=5):
        super().__init__()
        self.features = nn.Sequential(
            # Block 1: 16x16 -> 8x8
            nn.Conv2d(in_channels, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.GELU(),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.GELU(),
            nn.MaxPool2d(2),
            nn.Dropout2d(0.1),
            
            # Block 2: 8x8 -> 4x4
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.GELU(),
            nn.Conv2d(128, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.GELU(),
            nn.MaxPool2d(2),
            nn.Dropout2d(0.1),
            
            # Block 3: 4x4 -> 2x2
            nn.Conv2d(128, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.GELU(),
            nn.Conv2d(256, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.GELU(),
            nn.MaxPool2d(2),
            nn.Dropout2d(0.2),
        )
        
        # Squeeze and Excitation
        self.se = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(256, 64),
            nn.GELU(),
            nn.Linear(64, 256),
            nn.Sigmoid()
        )
        
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(256, 128),
            nn.GELU(),
            nn.Dropout(0.5),
            nn.Linear(128, num_classes)
        )
        
        self.embedding_head = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
        )
    
    def get_embedding(self, x):
        """Get 256-dim embedding for hybrid approach"""
        f = self.features(x)
        se_w = self.se(f).unsqueeze(-1).unsqueeze(-1)
        f = f * se_w
        return self.embedding_head(f)
    
    def forward(self, x):
        f = self.features(x)
        se_w = self.se(f).unsqueeze(-1).unsqueeze(-1)
        f = f * se_w
        return self.classifier(f)

# ===== 4. TRAIN LIGHTWEIGHT CNN =====
def train_light_cnn(X, y, num_classes, epochs=300, lr=3e-4):
    print("\n" + "="*60)
    print("STRATEGY 1: Lightweight CNN (no upsampling)")
    print("="*60)
    
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")
    
    # Data augmentation
    def augment_batch(x):
        if np.random.random() > 0.5:
            x = torch.flip(x, [2])
        if np.random.random() > 0.5:
            x = torch.flip(x, [3])
        if np.random.random() > 0.5:
            k = np.random.randint(1, 4)
            x = torch.rot90(x, k, [2, 3])
        # Random noise
        if np.random.random() > 0.5:
            noise = torch.randn_like(x) * 0.02
            x = x + noise
        # Mixup
        return x
    
    train_X = torch.FloatTensor(X_train)
    train_y = torch.LongTensor(y_train)
    test_X = torch.FloatTensor(X_test).to(device)
    test_y = torch.LongTensor(y_test)
    
    model = LightCNN(in_channels=X.shape[1], num_classes=num_classes).to(device)
    
    # Class weights
    class_counts = np.bincount(y_train, minlength=num_classes)
    weights = 1.0 / (class_counts + 1)
    weights = torch.FloatTensor(weights / weights.sum() * num_classes).to(device)
    
    criterion = nn.CrossEntropyLoss(weight=weights, label_smoothing=0.1)
    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=50, T_mult=2, eta_min=1e-6)
    
    best_acc = 0
    best_state = None
    patience = 0
    
    for epoch in range(epochs):
        model.train()
        # Shuffle
        perm = torch.randperm(len(train_X))
        train_loss = 0
        n_batches = 0
        
        for i in range(0, len(train_X), 32):
            idx = perm[i:i+32]
            bx = train_X[idx].to(device)
            by = train_y[idx].to(device)
            
            # Augmentation
            bx = augment_batch(bx)
            
            optimizer.zero_grad()
            out = model(bx)
            loss = criterion(out, by)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            train_loss += loss.item()
            n_batches += 1
        
        scheduler.step()
        
        model.eval()
        with torch.no_grad():
            out = model(test_X)
            preds = out.argmax(dim=1).cpu().numpy()
        acc = accuracy_score(test_y.numpy(), 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"  Epoch {epoch+1}/{epochs} Loss={train_loss/n_batches:.4f} Acc={acc:.4f} 🌟")
            if acc >= 0.95:
                print("  🎯 >95% reached!")
                break
        else:
            patience += 1
            if (epoch+1) % 20 == 0:
                print(f"  Epoch {epoch+1}/{epochs} Loss={train_loss/n_batches:.4f} Acc={acc:.4f} (patience={patience})")
        
        if patience >= 60:
            print(f"  Early stop at epoch {epoch+1}")
            break
    
    if best_state:
        model.load_state_dict(best_state)
    
    print(f"  ✅ LightCNN best acc: {best_acc:.4f}")
    return model, best_acc, X_test, y_test

# ===== 5. HYBRID CNN + XGBOOST =====
def train_hybrid(X, y, cnn_model, num_classes):
    print("\n" + "="*60)
    print("STRATEGY 2: Hybrid CNN embeddings + XGBoost")
    print("="*60)
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    cnn_model = cnn_model.to(device)
    cnn_model.eval()
    
    # Extract CNN embeddings
    with torch.no_grad():
        embeddings = []
        for i in range(0, len(X), 64):
            batch = torch.FloatTensor(X[i:i+64]).to(device)
            emb = cnn_model.get_embedding(batch)
            embeddings.append(emb.cpu().numpy())
    cnn_features = np.concatenate(embeddings, axis=0)
    print(f"  CNN embeddings: {cnn_features.shape}")
    
    # Extract rich handcrafted features
    rich_features = extract_rich_features(X)
    
    # Combine
    combined = np.concatenate([cnn_features, rich_features], axis=1)
    print(f"  Combined features: {combined.shape}")
    
    # Standardize
    scaler = StandardScaler()
    combined = scaler.fit_transform(combined)
    
    X_train, X_test, y_train, y_test = train_test_split(
        combined, y, test_size=0.2, random_state=42, stratify=y
    )
    
    # XGBoost with tuned params
    xgb = XGBClassifier(
        n_estimators=500,
        max_depth=8,
        learning_rate=0.05,
        subsample=0.8,
        colsample_bytree=0.8,
        min_child_weight=3,
        gamma=0.1,
        reg_alpha=0.1,
        reg_lambda=1.0,
        tree_method='hist', device='cuda',
        eval_metric='mlogloss',
        random_state=42,
        use_label_encoder=False
    )
    xgb.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=False)
    xgb_acc = accuracy_score(y_test, xgb.predict(X_test))
    print(f"  ✅ Hybrid XGBoost acc: {xgb_acc:.4f}")
    
    return xgb, scaler, xgb_acc, combined, X_test, y_test

# ===== 6. PURE RICH FEATURES + ENSEMBLE =====
def train_rich_ensemble(X, y, num_classes):
    print("\n" + "="*60)
    print("STRATEGY 3: Rich Features + Stacking Ensemble")
    print("="*60)
    
    rich_features = extract_rich_features(X)
    scaler = StandardScaler()
    rich_features = scaler.fit_transform(rich_features)
    
    X_train, X_test, y_train, y_test = train_test_split(
        rich_features, y, test_size=0.2, random_state=42, stratify=y
    )
    
    # Multiple base learners
    models_dict = {
        'XGBoost': XGBClassifier(
            n_estimators=500, max_depth=8, learning_rate=0.05,
            subsample=0.8, colsample_bytree=0.8, min_child_weight=3,
            tree_method='hist', device='cuda', eval_metric='mlogloss',
            random_state=42, use_label_encoder=False
        ),
        'LightGBM': LGBMClassifier(
            n_estimators=500, max_depth=8, learning_rate=0.05,
            subsample=0.8, colsample_bytree=0.8, min_child_weight=3,
            random_state=42, verbose=-1
        ),
        'ExtraTrees': ExtraTreesClassifier(
            n_estimators=500, max_depth=None, min_samples_split=5,
            random_state=42, n_jobs=-1
        ),
        'RandomForest': RandomForestClassifier(
            n_estimators=500, max_depth=None, min_samples_split=5,
            random_state=42, n_jobs=-1
        ),
        'GBM': GradientBoostingClassifier(
            n_estimators=300, max_depth=6, learning_rate=0.05,
            subsample=0.8, random_state=42
        ),
    }
    
    results = {}
    for name, model in models_dict.items():
        model.fit(X_train, y_train)
        acc = accuracy_score(y_test, model.predict(X_test))
        results[name] = acc
        print(f"  {name}: {acc:.4f}")
    
    # Stacking ensemble
    estimators = [(name, model) for name, model in models_dict.items() if name != 'GBM']
    stacking = StackingClassifier(
        estimators=estimators,
        final_estimator=XGBClassifier(
            n_estimators=200, max_depth=4, learning_rate=0.05,
            tree_method='hist', device='cuda', random_state=42, use_label_encoder=False
        ),
        cv=5, n_jobs=-1
    )
    stacking.fit(X_train, y_train)
    stack_acc = accuracy_score(y_test, stacking.predict(X_test))
    print(f"  Stacking Ensemble: {stack_acc:.4f}")
    
    # Voting ensemble
    voting = VotingClassifier(
        estimators=[(name, model) for name, model in models_dict.items()],
        voting='soft', n_jobs=-1
    )
    voting.fit(X_train, y_train)
    vote_acc = accuracy_score(y_test, voting.predict(X_test))
    print(f"  Voting Ensemble: {vote_acc:.4f}")
    
    results['Stacking'] = stack_acc
    results['Voting'] = vote_acc
    
    best_name = max(results, key=results.get)
    best_acc = results[best_name]
    print(f"  ✅ Best ensemble: {best_name} = {best_acc:.4f}")
    
    return stacking, voting, results, scaler, X_test, y_test

# ===== 7. CROSS-VALIDATION =====
def cross_validate_best(X_features, y, best_model_fn):
    print("\n" + "="*60)
    print("STRATEGY 4: 5-Fold Stratified Cross-Validation")
    print("="*60)
    
    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    fold_accs = []
    
    for fold, (train_idx, test_idx) in enumerate(skf.split(X_features, y)):
        X_tr, X_te = X_features[train_idx], X_features[test_idx]
        y_tr, y_te = y[train_idx], y[test_idx]
        
        model = best_model_fn()
        model.fit(X_tr, y_tr)
        acc = accuracy_score(y_te, model.predict(X_te))
        fold_accs.append(acc)
        print(f"  Fold {fold+1}: {acc:.4f}")
    
    mean_acc = np.mean(fold_accs)
    std_acc = np.std(fold_accs)
    print(f"  ✅ CV Mean: {mean_acc:.4f} ± {std_acc:.4f}")
    return mean_acc, std_acc

# ===== 8. FLAT FEATURES + XGBOOST (baseline comparison) =====
def train_flat_xgboost(X, y):
    print("\n" + "="*60)
    print("STRATEGY 5: Flat pixel features + XGBoost (sanity check)")
    print("="*60)
    
    X_flat = X.reshape(X.shape[0], -1)
    print(f"  Flat features: {X_flat.shape}")
    
    scaler = StandardScaler()
    X_flat = scaler.fit_transform(X_flat)
    
    X_train, X_test, y_train, y_test = train_test_split(
        X_flat, y, test_size=0.2, random_state=42, stratify=y
    )
    
    xgb = XGBClassifier(
        n_estimators=500, max_depth=8, learning_rate=0.05,
        subsample=0.8, colsample_bytree=0.8,
        tree_method='hist', device='cuda', eval_metric='mlogloss',
        random_state=42, use_label_encoder=False
    )
    xgb.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=False)
    acc = accuracy_score(y_test, xgb.predict(X_test))
    print(f"  ✅ Flat XGBoost acc: {acc:.4f}")
    return xgb, acc

# ===== MAIN =====
def main():
    print("🚀 CHIẾN LƯỢC TOÀN DIỆN ĐẠT >95% ACCURACY")
    print("="*60)
    
    X, y, num_classes = load_data()
    
    # Strategy 5: Flat baseline
    flat_xgb, flat_acc = train_flat_xgboost(X, y)
    
    # Strategy 1: Lightweight CNN
    cnn_model, cnn_acc, _, _ = train_light_cnn(X, y, num_classes)
    
    # Strategy 2: Hybrid CNN + XGBoost  
    hybrid_xgb, hybrid_scaler, hybrid_acc, combined_features, _, _ = train_hybrid(X, y, cnn_model, num_classes)
    
    # Strategy 3: Rich Features + Stacking Ensemble
    stacking, voting, ensemble_results, rich_scaler, _, _ = train_rich_ensemble(X, y, num_classes)
    
    # Strategy 4: Cross-validate the best
    rich_features = extract_rich_features(X)
    rich_features_scaled = StandardScaler().fit_transform(rich_features)
    
    cv_mean, cv_std = cross_validate_best(
        rich_features_scaled, y,
        lambda: XGBClassifier(
            n_estimators=500, max_depth=8, learning_rate=0.05,
            subsample=0.8, colsample_bytree=0.8,
            tree_method='hist', device='cuda', random_state=42, use_label_encoder=False
        )
    )
    
    # ===== SUMMARY =====
    print("\n" + "="*60)
    print("📊 TỔNG KẾT KẾT QUẢ")
    print("="*60)
    all_results = {
        'Flat XGBoost (baseline)': flat_acc,
        'LightCNN': cnn_acc,
        'Hybrid CNN+XGBoost': hybrid_acc,
    }
    all_results.update({f'Ensemble {k}': v for k, v in ensemble_results.items()})
    all_results['CV Mean (XGBoost rich)'] = cv_mean
    
    for name, acc in sorted(all_results.items(), key=lambda x: -x[1]):
        marker = "🏆" if acc >= 0.95 else "✅" if acc >= 0.90 else "📈"
        print(f"  {marker} {name}: {acc:.4f}")
    
    best_name = max(all_results, key=all_results.get)
    best_acc = all_results[best_name]
    print(f"\n🏆 BEST: {best_name} = {best_acc:.4f}")
    
    # Save best model
    os.makedirs('land_classification_model', exist_ok=True)
    os.makedirs('model_train', exist_ok=True)
    
    info = {
        "all_results": {k: float(v) for k, v in all_results.items()},
        "best_model": best_name,
        "best_accuracy": float(best_acc),
        "cv_mean": float(cv_mean),
        "cv_std": float(cv_std),
    }
    with open('model_train/ultimate_results.json', 'w') as f:
        json.dump(info, f, indent=2)
    
    print(f"\n✅ Kết quả đã được lưu vào model_train/ultimate_results.json")
    
    if best_acc >= 0.95:
        print("🎯🎯🎯 ĐÃ ĐẠT MỤC TIÊU >95% ACCURACY! 🎯🎯🎯")
    else:
        print(f"⚠️ Chưa đạt 95%. Best = {best_acc:.4f}. Cần thêm dữ liệu hoặc feature engineering.")

if __name__ == "__main__":
    main()