{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2cd3de2c",
   "metadata": {},
   "source": [
    "# Process Data & Train PyTorch CNN Model (Local Machine)\n",
    "Xử lý dữ liệu và train model trên máy cá nhân\n",
    "\n",
    "**Quy trình:**\n",
    "1. Load raw data từ server (NetCDF)\n",
    "2. Cloud masking → NDVI → Fill NaN → Monthly aggregation\n",
    "3. Chuẩn bị training data + augmentation\n",
    "4. Train PyTorch CNN model\n",
    "5. Evaluate & save model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4c7760f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "import importlib\n",
    "import sys\n",
    "import os\n",
    "\n",
    "# Import custom functions\n",
    "import new_import_ODC\n",
    "\n",
    "importlib.reload(new_import_ODC)\n",
    "from new_import_ODC import *\n",
    "\n",
    "# Setup matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "print(\"✅ Libraries imported successfully\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02136110",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Load raw data from NetCDF files\n",
    "print(\"📂 Loading raw satellite data from NetCDF files...\\n\")\n",
    "\n",
    "data_dir = \"data_for_training\"\n",
    "\n",
    "# Load Sentinel-2\n",
    "print(\"📡 Loading Sentinel-2...\")\n",
    "s2_path = os.path.join(data_dir, \"sentinel2_raw.nc\")\n",
    "data_s2 = xr.open_dataset(s2_path)\n",
    "print(f\"   ✅ Shape: {data_s2.dims}\")\n",
    "print(f\"   Bands: {list(data_s2.data_vars.keys())}\")\n",
    "\n",
    "# Load Sentinel-1\n",
    "print(\"\\n📡 Loading Sentinel-1...\")\n",
    "s1_path = os.path.join(data_dir, \"sentinel1_raw.nc\")\n",
    "data_s1 = xr.open_dataset(s1_path)\n",
    "print(f\"   ✅ Shape: {data_s1.dims}\")\n",
    "print(f\"   Bands: {list(data_s1.data_vars.keys())}\")\n",
    "\n",
    "print(\"\\n✅ All raw data loaded\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ed7f66d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "## Process Sentinel-2: Cloud masking\n",
    "print(\"☁️  Applying cloud mask (SCL band)...\\n\")\n",
    "\n",
    "# Extract raw data\n",
    "data = data_s2\n",
    "\n",
    "# Apply cloud mask using SCL band\n",
    "result = mask_clean(data)\n",
    "print(f\"✅ Cloud mask applied\")\n",
    "print(f\"   Shape: {result.dims}\")\n",
    "\n",
    "# Compute to ensure data is loaded\n",
    "result = result.compute()\n",
    "print(f\"✅ Data computed to memory\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e5b24a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Calculate NDVI from cloud-masked data\n",
    "print(\"🌱 Calculating NDVI...\\n\")\n",
    "\n",
    "ds1 = calculate_indices(result, index=\"NDVI\", satellite_mission=\"s2\")\n",
    "ndvi = ds1[\"NDVI\"]\n",
    "\n",
    "print(f\"✅ NDVI calculated\")\n",
    "print(f\"   Shape: {ndvi.shape}\")\n",
    "print(f\"   Value range: [{ndvi.min().values:.3f}, {ndvi.max().values:.3f}]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52c156ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Fill missing values with seasonal interpolation\n",
    "print(\"🔧 Filling missing values (cloud pixels)...\\n\")\n",
    "\n",
    "time_split = [\n",
    "    slice(\"2022-09-01\", \"2023-01-01\"),\n",
    "    slice(\"2023-01-01\", \"2023-05-01\"),\n",
    "    slice(\"2023-05-01\", \"2023-07-01\"),\n",
    "    slice(\"2023-07-01\", \"2023-10-01\"),\n",
    "]\n",
    "\n",
    "fill_nan_ndvi = fill_nan(ndvi, time_split)\n",
    "print(f\"✅ Missing values filled\")\n",
    "print(f\"   NaN pixels remaining: {fill_nan_ndvi.isna().sum().values}\")\n",
    "print(f\"   Valid pixels: {(~fill_nan_ndvi.isna()).sum().values}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adb728bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "## Monthly aggregation of NDVI and S1 data\n",
    "print(\"📊 Aggregating to monthly averages...\\n\")\n",
    "\n",
    "# NDVI monthly average\n",
    "print(\"   - NDVI monthly...\")\n",
    "average_ndvi = fill_nan_ndvi.resample(time=\"1M\").mean()\n",
    "average_ndvi = average_ndvi.compute()\n",
    "\n",
    "# S1 monthly average\n",
    "print(\"   - Sentinel-1 VH/VV monthly...\")\n",
    "average_vh = calculate_average(data_s1['VH'], time_pattern='1M')\n",
    "average_vv = calculate_average(data_s1['VV'], time_pattern='1M')\n",
    "\n",
    "print(f\"\\n✅ Monthly aggregation complete\")\n",
    "print(f\"   NDVI shape: {average_ndvi.shape}\")\n",
    "print(f\"   VH shape: {average_vh.shape}\")\n",
    "print(f\"   VV shape: {average_vv.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2d57a12",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Load training data and prepare datasets\n",
    "print(\"📋 Loading training data...\\n\")\n",
    "\n",
    "train_shp_dir = os.path.join(data_dir, \"train_data\")\n",
    "train_shp_path = os.path.join(train_shp_dir, \"ST_training data_updated_1130points_new.shp\")\n",
    "\n",
    "print(f\"   - Loading from {train_shp_path}...\")\n",
    "train = load_train_data(train_shp_path)\n",
    "print(f\"   ✅ Loaded {len(train)} training points\")\n",
    "print(f\"   Classes: {train['Class'].unique()}\")\n",
    "\n",
    "# Define label mapping\n",
    "label_mapping = {\n",
    "    \"Lua tom\": \"0\",\n",
    "    \"Lua\": \"1\",\n",
    "    \"CHN\": \"2\",\n",
    "    \"CLN\": \"3\",\n",
    "    \"TS\": \"4\",\n",
    "    \"Song\": \"5\",\n",
    "    \"Dat xay dung\": \"6\",\n",
    "    \"Rung\": \"7\",\n",
    "}\n",
    "\n",
    "print(f\"\\n   Label mapping: {label_mapping}\")\n",
    "\n",
    "# Prepare datasets (extract S2 + S1 values at training points)\n",
    "print(f\"\\n   - Extracting features at training points...\")\n",
    "datasets = get_data_sen1_and_sen2(train, average_ndvi, average_vh, average_vv)\n",
    "\n",
    "print(f\"✅ Training dataset prepared\")\n",
    "print(f\"   Features shape: {datasets[0].shape if hasattr(datasets[0], 'shape') else 'N/A'}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bfc71f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Split training/validation/test data\n",
    "print(\"✂️  Splitting data into train/validation/test...\\n\")\n",
    "\n",
    "X_train, X_val, X_test, y_train, y_val, y_test = split_train_data(\n",
    "    train, label_mapping, datasets\n",
    ")\n",
    "\n",
    "print(f\"   Train set: {X_train.shape[0]} samples\")\n",
    "print(f\"   Val set:   {X_val.shape[0]} samples\")\n",
    "print(f\"   Test set:  {X_test.shape[0]} samples\")\n",
    "\n",
    "print(f\"\\n✅ Data split complete\")\n",
    "print(f\"   X_train shape: {X_train.shape}\")\n",
    "print(f\"   X_val shape: {X_val.shape}\")\n",
    "print(f\"   X_test shape: {X_test.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35379fa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "## Train PyTorch CNN Model\n",
    "print(\"🤖 Training PyTorch CNN model...\\n\")\n",
    "\n",
    "# Train the model using the function from new_import_ODC\n",
    "model = train_cnn_pytorch(\n",
    "    X_train, X_val, \n",
    "    y_train, y_val,\n",
    "    epochs=50,\n",
    "    batch_size=32,\n",
    "    learning_rate=0.001,\n",
    "    patience=10\n",
    ")\n",
    "\n",
    "print(f\"\\n✅ Model training complete\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f37dccf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Evaluate model on test set\n",
    "print(\"📊 Evaluating model on test set...\\n\")\n",
    "\n",
    "from sklearn.metrics import accuracy_score, classification_report, confusion_matrix\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "# Get predictions on test set\n",
    "device = torch.device('cpu')\n",
    "model = model.to(device)\n",
    "model.eval()\n",
    "\n",
    "with torch.no_grad():\n",
    "    X_test_tensor = torch.FloatTensor(X_test).to(device)\n",
    "    y_pred_probs = model(X_test_tensor).cpu().numpy()\n",
    "    y_pred_test = np.argmax(y_pred_probs, axis=1)\n",
    "\n",
    "test_accuracy = accuracy_score(y_test, y_pred_test)\n",
    "print(f\"✅ Test Accuracy: {test_accuracy:.4f} ({test_accuracy*100:.2f}%)\")\n",
    "\n",
    "print(f\"\\n📈 Classification Report:\\n\")\n",
    "print(classification_report(y_test, y_pred_test, \n",
    "                          target_names=list(label_mapping.keys())))\n",
    "\n",
    "print(f\"\\n🔲 Confusion Matrix:\\n\")\n",
    "print(confusion_matrix(y_test, y_pred_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e84b7206",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Save trained model\n",
    "print(\"💾 Saving trained model...\\n\")\n",
    "\n",
    "model_path = \"model_cnn_pytorch_local.pth\"\n",
    "torch.save(model.state_dict(), model_path)\n",
    "print(f\"   ✅ Model saved to {model_path}\")\n",
    "\n",
    "# Also save as PyTorch checkpoint with metadata\n",
    "checkpoint = {\n",
    "    'model_state_dict': model.state_dict(),\n",
    "    'accuracy': test_accuracy,\n",
    "    'label_mapping': label_mapping,\n",
    "    'num_classes': len(label_mapping),\n",
    "    'input_features': X_train.shape[1]\n",
    "}\n",
    "\n",
    "checkpoint_path = \"model_cnn_pytorch_local_checkpoint.pth\"\n",
    "torch.save(checkpoint, checkpoint_path)\n",
    "print(f\"   ✅ Checkpoint saved to {checkpoint_path}\")\n",
    "\n",
    "print(f\"\\n✅ Model training & evaluation complete!\")\n",
    "print(f\"   Test Accuracy: {test_accuracy*100:.2f}%\")\n",
    "print(f\"   Model ready for prediction on full spatial extent\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}