"""
01.train_ODC.py
Chuyển đổi từ 01.train_ODC.ipynb
Land Use Classification Training - Sentinel-2 / ODC
"""

import os
import sys
import importlib
import traceback
import numpy as np
import xarray as xr
from datetime import datetime

import new_import_ODC
from new_import_ODC import *  # load_train_data, save_model, notebook_utils, ...

# ============================================================
#  HYPERPARAMETERS - Chỉnh sửa tại đây
# ============================================================

# --- Dask cluster ---
DASK_N_WORKERS = 4

# --- Dữ liệu Sentinel-2 ---
DATE_RANGE         = ("2023-03-01", "2023-12-31")
LONGITUDE_RANGE    = (105.5, 106.4)
LATITUDE_RANGE     = (9.2, 10.0)
NUM_SCENES         = 1          # Số scene cần load (None = tất cả)

# --- Cache ---
CACHE_DIR  = "dataset_cache"
CACHE_FILE = f"{CACHE_DIR}/sentinel2_timeseries_40scenes.nc"

# --- Training data ---
TRAIN_PATH = "train/ST_training_data_updated_1130points_new.shp"

# --- Features sử dụng để train ---
AVAILABLE_FEATURES = [
    'ndvi_mean', 'ndvi_min', 'ndvi_max', 'ndvi_std', 'ndvi_range',
    'ndwi_mean', 'ndbi_mean', 'evi_mean',
]

# --- Train / Test split ---
TEST_SIZE    = 0.2
RANDOM_STATE = 42

# --- Random Forest hyperparameters ---
RF_N_ESTIMATORS    = 200
RF_MAX_DEPTH       = 30
RF_MIN_SAMPLES_SPLIT = 5
RF_N_JOBS          = -1
RF_VERBOSE         = 1

# --- Output model ---
MODEL_FILENAME = "model_land_use_odc.joblib"

# --- Label mapping ---
LABEL_MAPPING = {
    "Lua tom":      "0",
    "Lua":          "1",
    "CHN":          "2",
    "CLN":          "3",
    "TS":           "4",
    "Song":         "5",
    "Dat xay dung": "6",
    "Rung":         "7",
}

# ============================================================
#  END HYPERPARAMETERS
# ============================================================


def setup_imports():
    """Load custom modules."""
    import new_import_ODC
    importlib.reload(new_import_ODC)
    print("✅ All modules loaded successfully")
    return new_import_ODC


def setup_dask_and_datacube():
    """Khởi động Dask cluster và kết nối Datacube."""
    from dask.distributed import Client, LocalCluster
    import datacube

    print("✅ AWS credentials loaded from environment variables")

    cluster = LocalCluster(n_workers=DASK_N_WORKERS)
    client  = Client(cluster)
    print("✅ Dask cluster initialized")
    print(f"   Cluster: {cluster}")

    # Cấu hình S3 access cho rasterio/GDAL (bắt buộc để đọc COGs từ S3)
    configure_s3_access(aws_unsigned=False, requester_pays=True, client=client)
    print("✅ S3 access configured (requester_pays=True)")

    try:
        dc = datacube.Datacube()
        print("✅ Datacube connected (metadata only)")
    except Exception as e:
        print(f"⚠️  Datacube connection not critical: {e}")
        dc = None

    print("\n" + "=" * 70)
    return client, cluster, dc


def get_scene_metadata(dc):
    """Lấy metadata Sentinel-2 từ datacube."""
    print("=" * 70)
    print("GETTING SENTINEL-2 SCENE METADATA")
    print("=" * 70)

    try:
        print(f"\n[1] Loading metadata from datacube...")
        datasets = list(dc.find_datasets(product='s2_l2a', time=DATE_RANGE))
        print(f"    ✅ Found {len(datasets)} scenes")

        if datasets:
            selected = datasets[0]
            print(f"\n[2] Selected scene: {selected.metadata.label}")
            scene_datetime = selected.time.begin if hasattr(selected.time, 'begin') else selected.time
            print(f"    Date: {scene_datetime}")

            print(f"\n[3] Available bands:")
            for name, measurement in selected.measurements.items():
                print(f"    - {name}: {measurement['path'][:80]}")

    except Exception as e:
        print(f"❌ Error: {e}")
        traceback.print_exc()

    print("=" * 70)


def check_cache():
    """Kiểm tra dataset cache. Trả về (data, use_cache)."""
    print("=" * 70)
    print("CHECKING FOR CACHED DATASET")
    print("=" * 70)

    use_cache = False
    data      = None

    if os.path.exists(CACHE_FILE):
        print(f"\n✅ Cache file found: {CACHE_FILE}")
        file_size_gb = os.path.getsize(CACHE_FILE) / (1024 ** 3)
        print(f"   File size: {file_size_gb:.2f} GB")

        try:
            print(f"\n🔄 Loading dataset from cache...")
            data = xr.open_dataset(CACHE_FILE)
            print(f"✅ Dataset loaded from cache!")
            print(f"   Total scenes: {len(data['time'])}")
            print(f"   Variables:    {len(data.data_vars)}")
            print(f"   Dimensions:   {dict(data.dims)}")
            print(f"\n   ⏭️  Skipping S3 download (using cached data)")
            use_cache = True
        except Exception as e:
            print(f"❌ Error loading cache: {e}")
            print(f"   Will download fresh data from S3")
    else:
        print(f"\n⏳ Cache file not found: {CACHE_FILE}")
        print(f"   Will download from S3 and save cache")
        print(f"   (Next run will use cache automatically)")

    print("=" * 70)
    return data, use_cache


def load_satellite_data(dc, data, use_cache):
    """Download / load Sentinel-2 data và tính spectral indices."""
    import rasterio
    from scipy import ndimage

    print("=" * 70)
    print("LOADING SENTINEL-2 DATA FROM S3 COGs (RASTERIO) - OPTIMAL ACCURACY")
    print("=" * 70)

    ndvi = None

    try:
        if use_cache and data is not None:
            print(f"\n✅ Using cached dataset - skipping download!")
            print(f"   Variables: {len(data.data_vars)}")
            print(f"   Shape:     {data.dims}")

        else:
            # --- Download from S3 ---
            print(f"\n📥 Downloading from S3...")
            datasets = list(dc.find_datasets(product='s2_l2a', time=DATE_RANGE))

            if not datasets:
                raise ValueError("No datasets found for date range")

            print(f"\n📦 Found {len(datasets)} available scenes")
            print(f"   Date range: {DATE_RANGE[0]} to {DATE_RANGE[1]}")

            num_scenes = NUM_SCENES if NUM_SCENES is not None else len(datasets)
            print(f"\n[LOADING] Loading {num_scenes} scenes with ALL available bands...")
            print(f"   (Keeping NATIVE resolution - NO upsampling/magnification)")

            all_data_dict = {}
            failed_scenes = []
            scene_dates   = []

            first_scene          = datasets[0]
            all_available_bands  = list(first_scene.measurements.keys())
            print(f"   Available bands: {all_available_bands}")

            for scene_idx in range(num_scenes):
                selected       = datasets[scene_idx]
                scene_label    = selected.metadata.label
                scene_datetime = selected.time.begin if hasattr(selected.time, 'begin') else selected.time
                scene_dates.append(scene_datetime)

                if scene_idx % 5 == 0 or scene_idx == 0 or scene_idx == num_scenes - 1:
                    print(f"\n   [{scene_idx + 1:2d}/{num_scenes}] {scene_label} ({scene_datetime.date()})")

                scene_data_dict = {}

                for band_name in all_available_bands:
                    if band_name in selected.measurements:
                        band_path = selected.measurements[band_name]['path']
                        try:
                            with rasterio.open(band_path) as src:
                                scene_data_dict[band_name] = src.read(1)
                        except Exception as e:
                            if scene_idx % 5 == 0:
                                print(f"      ⚠️ Error loading {band_name}: {str(e)[:30]}")
                            failed_scenes.append((scene_idx, scene_label, band_name, str(e)))

                if scene_data_dict:
                    all_data_dict[scene_idx] = scene_data_dict
                    if scene_idx % 5 == 0 or scene_idx == num_scenes - 1:
                        print(f"      ✅ {len(scene_data_dict)} bands loaded")
                else:
                    failed_scenes.append((scene_idx, scene_label, "all", "No bands loaded"))

            if not all_data_dict:
                raise ValueError("Could not load any bands from any scene")

            print(f"\n✅ Successfully loaded {len(all_data_dict)} scenes!")
            if failed_scenes:
                print(f"⚠️  Failed to load {len(failed_scenes)} band instances (will be skipped)")

            # --- Resolution normalization ---
            print(f"\n[RESOLUTION NORMALIZATION] Aligning all bands to native resolution...")

            ref_resolution = max_size = 0
            max_band = None

            for s_idx in all_data_dict:
                for b_name, b_data in all_data_dict[s_idx].items():
                    sz = b_data.shape[0]
                    if sz > max_size:
                        max_size = sz
                        ref_resolution = sz
                        max_band = b_name

            print(f"   Reference resolution: {max_size}×{max_size} pixels (native {max_band})")

            resampled_count = 0
            for s_idx in all_data_dict:
                for b_name in list(all_data_dict[s_idx]):
                    arr = all_data_dict[s_idx][b_name]
                    if arr.shape[0] != ref_resolution:
                        scale = ref_resolution / arr.shape[0]
                        order = 0 if b_name == 'scl' else 1
                        resampled = ndimage.zoom(arr, scale, order=order)
                        all_data_dict[s_idx][b_name] = resampled
                        if s_idx == 0:
                            print(f"      Resampling {b_name}: {arr.shape[0]}×{arr.shape[0]} "
                                  f"→ {resampled.shape[0]}×{resampled.shape[0]}")
                        resampled_count += 1

            print(f"✅ Resolution normalization complete! ({resampled_count} bands resampled)")

            # --- Spectral indices ---
            print(f"\n[SPECTRAL INDICES] Calculating spectral indices for each scene...")
            indices_count = 0

            for s_idx in all_data_dict:
                sd = all_data_dict[s_idx]
                try:
                    if 'nir' in sd and 'red' in sd:
                        nir, red = sd['nir'].astype(float), sd['red'].astype(float)
                        sd['ndvi'] = ((nir - red) / (nir + red + 1e-8)).astype(np.float32)
                        indices_count += 1

                    if 'b11' in sd and 'nir' in sd:
                        swir, nir = sd['b11'].astype(float), sd['nir'].astype(float)
                        sd['ndbi'] = ((swir - nir) / (swir + nir + 1e-8)).astype(np.float32)
                        indices_count += 1

                    if 'nir' in sd and 'b11' in sd:
                        nir, swir = sd['nir'].astype(float), sd['b11'].astype(float)
                        sd['ndwi'] = ((nir - swir) / (nir + swir + 1e-8)).astype(np.float32)
                        indices_count += 1

                    if 'nir' in sd and 'red' in sd and 'blue' in sd:
                        nir  = sd['nir'].astype(float)
                        red  = sd['red'].astype(float)
                        blue = sd['blue'].astype(float)
                        sd['evi'] = (2.5 * (nir - red) / (nir + 6 * red - 7.5 * blue + 1)).astype(np.float32)
                        indices_count += 1
                except Exception:
                    pass

            print(f"✅ Calculated {indices_count} spectral indices per scene")

            # --- Stack along time ---
            print(f"\n[STACKING] Stacking {len(all_data_dict)} scenes to time-series...")

            data_vars  = {}
            band_names = list(all_data_dict[0].keys())

            for b_name in band_names:
                arr_list = [
                    all_data_dict[si][b_name]
                    for si in sorted(all_data_dict)
                    if b_name in all_data_dict[si]
                ]
                if arr_list:
                    data_vars[b_name] = (['time', 'y', 'x'], np.stack(arr_list, axis=0))

            first_arr = list(all_data_dict[0].values())[0]
            y_size, x_size = first_arr.shape

            data = xr.Dataset(
                data_vars,
                coords={
                    'time': np.arange(len(all_data_dict)),
                    'x':    np.arange(x_size),
                    'y':    np.arange(y_size),
                }
            )

            # --- Temporal features ---
            print(f"\n[TEMPORAL FEATURES] Computing temporal features...")
            tf_added = 0

            if 'ndvi' in data.data_vars:
                ts = data['ndvi']
                data['ndvi_min']   = ts.min(dim='time');   tf_added += 1
                data['ndvi_max']   = ts.max(dim='time');   tf_added += 1
                data['ndvi_mean']  = ts.mean(dim='time');  tf_added += 1
                data['ndvi_range'] = data['ndvi_max'] - data['ndvi_min']; tf_added += 1
                data['ndvi_std']   = ts.std(dim='time');   tf_added += 1

            for b_name in ['ndbi', 'ndwi', 'evi']:
                if b_name in data.data_vars:
                    data[f'{b_name}_mean'] = data[b_name].mean(dim='time')
                    tf_added += 1

            print(f"✅ Added {tf_added} temporal/aggregate features")

            # --- Save cache ---
            print(f"\n[CACHE] Saving dataset to cache...")
            try:
                os.makedirs(CACHE_DIR, exist_ok=True)
                data.to_netcdf(CACHE_FILE, engine='netcdf4')
                cache_size = os.path.getsize(CACHE_FILE) / (1024 ** 3)
                print(f"✅ Dataset saved to cache: {CACHE_FILE}")
                print(f"   Cache size: {cache_size:.2f} GB")
            except Exception as e:
                print(f"⚠️  Error saving cache: {e}")

            print(f"\n✅ Dataset created!")
            print(f"   🎬 Scenes: {len(all_data_dict)}")
            print(f"   📊 Variables: {len(data.data_vars)}")
            print(f"   🖼️  Size: {x_size} × {y_size} px")
            print(f"   ⏰ {scene_dates[0].date()} → {scene_dates[-1].date()}")

        # --- Extract NDVI ---
        print(f"\n[NDVI EXTRACTION] Extracting NDVI for model training...")
        if 'ndvi_mean' in data.data_vars:
            ndvi = data['ndvi_mean']
            print(f"✅ NDVI extracted (mean across time) - shape: {ndvi.shape}")
        elif 'ndvi' in data.data_vars:
            ndvi = data['ndvi'].isel(time=0)
            print(f"✅ NDVI extracted (first time step) - shape: {ndvi.shape}")
        else:
            print(f"❌ NDVI not found in dataset")

    except Exception as e:
        print(f"❌ Error: {e}")
        traceback.print_exc()
        data = None

    print("=" * 70)
    return data, ndvi


def load_training_data():
    """Load training shapefile và trả về GeoDataFrame."""
    print("=" * 70)
    print("TRAINING DATA SETUP")
    print("=" * 70)
    print(f"\n[1] Loading training data: {TRAIN_PATH}")

    try:
        train = load_train_data(TRAIN_PATH)
        print(f"    ✅ Loaded {len(train)} training points")
        print(f"    Columns: {list(train.columns)}")
    except Exception as e:
        print(f"    ❌ Error: {e}")
        train = None

    print(f"\n[2] Label mapping:")
    for label, code in LABEL_MAPPING.items():
        print(f"    {code}: {label}")

    print("\n" + "=" * 70)
    return train


def train_model(train, data):
    """Extract features, split, train RandomForest, evaluate."""
    from sklearn.model_selection import train_test_split
    from sklearn.ensemble import RandomForestClassifier
    from sklearn.metrics import accuracy_score, classification_report

    print("=" * 70)
    print("LAND USE CLASSIFICATION TRAINING")
    print("=" * 70)

    model = X_train = X_test = y_train = y_test = y_pred = accuracy = None
    features_to_use = class_names = []

    if train is None or data is None:
        print("❌ Missing training data or satellite data")
        return model, None, None, None, None, None, None, features_to_use, class_names

    print("\n[1] Extracting features from satellite data...")

    try:
        features_to_use = [f for f in AVAILABLE_FEATURES if f in data.data_vars]

        if not features_to_use:
            print("    ❌ No spectral features found in dataset!")
            print("    Available variables:", list(data.data_vars))
            return model, None, None, None, None, None, None, features_to_use, class_names

        print(f"    Using {len(features_to_use)} features: {features_to_use}")

        X, y = [], []
        for _, point in train.iterrows():
            try:
                feature_vec = [
                    float(data[fn].sel(x=point.geometry.x, y=point.geometry.y, method='nearest').values)
                    for fn in features_to_use
                ]
                label = LABEL_MAPPING[point.Hientrang]
                if not np.isnan(feature_vec).any():
                    X.append(feature_vec)
                    y.append(int(label))
            except Exception:
                continue

        if not X:
            print("    ❌ No samples extracted")
            return model, None, None, None, None, None, None, features_to_use, class_names

        X = np.array(X)
        y = np.array(y)
        print(f"    ✅ {len(X)} samples × {X.shape[1]} features")

        print(f"\n    Feature statistics:")
        for i, fn in enumerate(features_to_use):
            print(f"      {fn:15s}: mean={X[:, i].mean():.3f}, std={X[:, i].std():.3f}")

        # Split
        print(f"\n[2] Splitting data ({int((1-TEST_SIZE)*100)}-{int(TEST_SIZE*100)})...")
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=TEST_SIZE, random_state=RANDOM_STATE, stratify=y
        )
        print(f"    Train: {len(X_train)}, Test: {len(X_test)}")

        unique, counts = np.unique(y_train, return_counts=True)
        print(f"\n    Class distribution in training set:")
        for cls, count in zip(unique, counts):
            cls_name = [k for k, v in LABEL_MAPPING.items() if v == str(cls)][0]
            print(f"      {cls}: {cls_name:15s} - {count:4d} ({count/len(y_train)*100:.1f}%)")

        # Train
        print(f"\n[3] Training Random Forest...")
        model = RandomForestClassifier(
            n_estimators=RF_N_ESTIMATORS,
            max_depth=RF_MAX_DEPTH,
            min_samples_split=RF_MIN_SAMPLES_SPLIT,
            random_state=RANDOM_STATE,
            n_jobs=RF_N_JOBS,
            verbose=RF_VERBOSE,
        )
        model.fit(X_train, y_train)

        y_pred   = model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        print(f"\n    ✅ Training accuracy: {model.score(X_train, y_train)*100:.2f}%")
        print(f"    ✅ Testing accuracy:  {accuracy*100:.2f}%")

        # Feature importance
        print(f"\n    Feature importance:")
        importances = model.feature_importances_
        for rank, idx in enumerate(np.argsort(importances)[::-1]):
            print(f"      {rank+1}. {features_to_use[idx]:15s}: {importances[idx]:.4f}")

        # Classification report
        class_names = [k for k, v in sorted(LABEL_MAPPING.items(), key=lambda x: x[1])]
        print(f"\n[4] Classification Report:")
        print(classification_report(y_test, y_pred, target_names=class_names, zero_division=0))

    except Exception as e:
        print(f"    ❌ Error: {e}")
        traceback.print_exc()
        model = None

    print("\n" + "=" * 70)
    return model, X_train, X_test, y_train, y_test, y_pred, accuracy, features_to_use, class_names


def save_trained_model(model, X_train, X_test, y_train, y_test,
                       y_pred, accuracy, features_to_use, class_names):
    """Lưu model + metadata bằng ModelManager."""
    from sklearn.metrics import classification_report

    print("=" * 70)
    print("MODEL SAVING")
    print("=" * 70)

    if model is None:
        print("❌ No model to save")
        return

    print("\n🔄 Saving model with metadata...")
    try:
        y = np.concatenate([y_train, y_test])
        metadata = {
            "timestamp":             datetime.now().isoformat(),
            "data_source":           "Local S3 ODC (Open Data Cube)",
            "collections":           ["sentinel-2-l2a"],
            "features":              features_to_use,
            "feature_mode":          "extended",
            "training_samples":      len(X_train),
            "testing_samples":       len(X_test),
            "test_size":             TEST_SIZE,
            "train_accuracy":        float(model.score(X_train, y_train)),
            "test_accuracy":         float(accuracy),
            "model_type":            "random_forest",
            "device":                "cpu",
            "n_estimators":          RF_N_ESTIMATORS,
            "max_depth":             RF_MAX_DEPTH,
            "min_samples_split":     RF_MIN_SAMPLES_SPLIT,
            "learning_rate":         None,
            "cnn_epochs":            None,
            "n_features":            X_train.shape[1],
            "n_classes":             len(np.unique(y)),
            "class_names":           list(LABEL_MAPPING.keys()),
            "classification_report": classification_report(
                y_test, y_pred, target_names=class_names,
                output_dict=True, zero_division=0
            ),
            "bbox":       None,
            "time_range": f"{DATE_RANGE[0]}/{DATE_RANGE[1]}",
            "resolution": 10,
            "notes":      "Land Use Classification (8 classes) trained from 01.train_ODC.py",
        }

        save_model(MODEL_FILENAME, model, metadata=metadata, label_encoder=None)

        print(f"✅ Model saved → model_train/{MODEL_FILENAME}")
        print(f"   Train Accuracy: {metadata['train_accuracy']*100:.2f}%")
        print(f"   Test  Accuracy: {metadata['test_accuracy']*100:.2f}%")
        print(f"   Classes:        {metadata['n_classes']}")
    except Exception as e:
        print(f"❌ Error saving model: {e}")
        traceback.print_exc()

    print("=" * 70)


def cleanup(client, cluster):
    """Đóng Dask client/cluster."""
    print("=" * 70)
    print("CLEANUP")
    print("=" * 70)
    try:
        client.close()
        cluster.close()
        print("✅ Cleanup complete")
    except Exception as e:
        print(f"⚠️  Error during cleanup: {e}")
    print("\n" + "=" * 70)
    print("✅ PIPELINE COMPLETE")
    print("=" * 70)


# ============================================================
#  MAIN
# ============================================================
if __name__ == "__main__":
    # 1. Reload modules
    setup_imports()

    # 2. Dask + Datacube
    client, cluster, dc = setup_dask_and_datacube()

    # 3. Scene metadata
    if dc is not None:
        get_scene_metadata(dc)

    # 4. Check cache
    data, use_cache = check_cache()

    # 5. Load / download satellite data
    data, ndvi = load_satellite_data(dc, data, use_cache)

    # 6. Training data
    train = load_training_data()

    # 7. Train model
    (model, X_train, X_test, y_train, y_test,
     y_pred, accuracy, features_to_use, class_names) = train_model(train, data)

    # 8. Save model
    if model is not None:
        save_trained_model(model, X_train, X_test, y_train, y_test,
                           y_pred, accuracy, features_to_use, class_names)

    # 9. Cleanup
    cleanup(client, cluster)