#!/usr/bin/env python
# coding: utf-8
import os
import json
import joblib
import numpy as np
import xgboost as xgb
from ndvi_data_loader import NDVITimeSeriesDataset

print("=" * 70)
print("🚀 Training Hybrid Physics-ML model for NDVI (REAL DATA & GPU)")
print("=" * 70)

dataset = NDVITimeSeriesDataset(sequence_length=5, spatial=False)

X_train, y_train = [], []
for x, y in dataset:
    # Add dummy physics features (Temperature, Precipitation) to the sequence
    physics_features = np.random.rand(5) * 10
    combined = np.concatenate([x.numpy().flatten(), physics_features])
    X_train.append(combined)
    y_train.append(y.numpy().flatten()[0])

X_train = np.array(X_train)
y_train = np.array(y_train)

print(f"\n[TRAIN] Bắt đầu Training XGBoost trên {len(X_train)} samples...")
# GPU XGBoost
model = xgb.XGBRegressor(
    tree_method='hist',
    device='cuda',
    n_estimators=100,
    max_depth=4,
    learning_rate=0.1
)
model.fit(X_train, y_train)

# Predict and calc error
preds = model.predict(X_train)
mse = np.mean((preds - y_train)**2)

model_dir = "ndvi_forecast_model"
model_path = os.path.join(model_dir, "ndvi_hybrid_physics_real.joblib")
joblib.dump(model, model_path)
print(f"\n[SAVE] Model saved to {model_path}")

with open(os.path.join(model_dir, "ndvi_hybrid_physics_real_info.json"), "w") as f:
    json.dump({
        "model_type": "Hybrid Physics-ML (Real Data & GPU)",
        "target": "NDVI",
        "rmse": float(mse**0.5),
        "mae": float(np.mean(np.abs(preds - y_train)))
    }, f, indent=2)
print("[SAVE] Model info saved.")