import os
import glob
import time
import json
import itertools
import numpy as np
import joblib
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, accuracy_score

from train_module import SwinUNetClassifier

def load_data():
    cache_files = glob.glob('dataset_cache/training_data_*.joblib')
    if not cache_files:
        raise FileNotFoundError("No cache files found in dataset_cache/")
    
    # Get the latest cache file
    cache_file = max(cache_files, key=os.path.getctime)
    print(f"Loading data from {cache_file}...")
    data = joblib.load(cache_file)
    features = data['features']
    labels = data['labels']
    
    # Map labels to 0..N-1
    unique_labels = sorted(list(np.unique(labels)))
    label_map = {lbl: idx for idx, lbl in enumerate(unique_labels)}
    mapped_labels = np.array([label_map[l] for l in labels])
    
    return features, mapped_labels, unique_labels

def train_evaluate(features, labels, embed_dim, lr, weight_decay, epochs, patience, device):
    X_train, X_test, y_train, y_test = train_test_split(features, labels, test_size=0.2, random_state=42)
    
    n_features = X_train.shape[1]
    n_classes = len(np.unique(labels))
    
    model = SwinUNetClassifier(n_features, n_classes, embed_dim=embed_dim).to(device)
    
    X_train_t = torch.FloatTensor(X_train)
    y_train_t = torch.LongTensor(y_train)
    X_test_t = torch.FloatTensor(X_test)
    y_test_t = torch.LongTensor(y_test)
    
    train_dataset = TensorDataset(X_train_t, y_train_t)
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    
    # Class weights
    class_counts = np.bincount(y_train)
    class_weights = 1.0 / (class_counts + 1e-6)
    class_weights = class_weights / class_weights.sum() * len(class_counts)
    class_weights_t = torch.FloatTensor(class_weights).to(device)
    
    criterion = nn.CrossEntropyLoss(weight=class_weights_t)
    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)
    
    best_acc = 0.0
    patience_counter = 0
    best_model_state = None
    
    model.train()
    for epoch in range(epochs):
        model.train()
        for batch_X, batch_y in train_loader:
            batch_X, batch_y = batch_X.to(device), batch_y.to(device)
            optimizer.zero_grad()
            outputs = model(batch_X)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()
            
        scheduler.step()
        
        # Eval
        model.eval()
        with torch.no_grad():
            outputs = model(X_test_t.to(device))
            _, preds = torch.max(outputs, 1)
            acc = accuracy_score(y_test, preds.cpu().numpy())
            
        if acc > best_acc:
            best_acc = acc
            best_model_state = model.state_dict()
            patience_counter = 0
        else:
            patience_counter += 1
            
        if patience_counter >= patience:
            break
            
    # Restore best
    if best_model_state:
        model.load_state_dict(best_model_state)
        
    return model, best_acc, X_test_t, y_test

def main():
    print("🚀 BẮT ĐẦU TÌM KIẾM SIÊU THAM SỐ CHO SWIN-UNET")
    features, labels, unique_labels = load_data()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    
    param_grid = {
        'embed_dim': [64, 128, 256, 512],
        'lr': [1e-3, 5e-4, 1e-4],
        'weight_decay': [0.01, 0.001],
        'epochs': [200, 500]
    }
    
    keys = param_grid.keys()
    combinations = list(itertools.product(*(param_grid[k] for k in keys)))
    
    best_global_acc = 0.0
    best_params = None
    best_model = None
    
    # Ensure dir exists
    os.makedirs('land_classification_model', exist_ok=True)
    os.makedirs('model_train', exist_ok=True)
    
    for i, values in enumerate(combinations):
        params = dict(zip(keys, values))
        print(f"\n[{i+1}/{len(combinations)}] Training with params: {params}")
        
        model, acc, X_test_t, y_test = train_evaluate(
            features, labels, 
            embed_dim=params['embed_dim'], 
            lr=params['lr'], 
            weight_decay=params['weight_decay'], 
            epochs=params['epochs'], 
            patience=30, 
            device=device
        )
        
        print(f"Test Accuracy: {acc:.4f}")
        
        if acc > best_global_acc:
            best_global_acc = acc
            best_params = params
            best_model = model
            
            print(f"🌟 NEW BEST ACCURACY: {acc:.4f}")
            
            if acc >= 0.95:
                print("🎯 ĐẠT MỤC TIÊU >95%! DỪNG TÌM KIẾM.")
                break
                
    if best_model is not None:
        model_path = 'land_classification_model/model_swin-unet_optimized_95.joblib'
        best_model = best_model.cpu()
        joblib.dump(best_model, model_path)
        print(f"\n✅ Đã lưu mô hình tốt nhất (Acc: {best_global_acc:.4f}) vào {model_path}")
        print(f"Cấu hình tốt nhất: {best_params}")
        
        # generate report
        best_model.eval()
        with torch.no_grad():
            outputs = best_model(X_test_t)
            _, preds = torch.max(outputs, 1)
            clf_rep = classification_report(y_test, preds.cpu().numpy(), output_dict=True)
            
        info = {
            "model_type": "swin-unet",
            "test_accuracy": float(best_global_acc),
            "params": {"n_estimators": best_params['epochs'], "max_depth": best_params['embed_dim']},
            "classification_report": clf_rep
        }
        with open('model_train/model_swin-unet_auto_info.json', 'w') as f:
            json.dump(info, f, indent=2)

if __name__ == "__main__":
    main()