#!/usr/bin/env python
# coding: utf-8
import os
import json
import joblib
import numpy as np
from statsmodels.tsa.statespace.sarimax import SARIMAX
from ndvi_data_loader import NDVITimeSeriesDataset

print("=" * 70)
print("🚀 Training Statistical Model (SARIMA) for NDVI (REAL DATA & CPU)")
print("=" * 70)

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

# Flatten data for ARIMA (1D series)
timeseries = []
for x, y in dataset:
    timeseries.append(y.item())

print(f"\n[TRAIN] Bắt đầu Training SARIMA với {len(timeseries)} điểm dữ liệu...")
# Use a simple ARIMA (1, 1, 1)
model = SARIMAX(timeseries, order=(1, 1, 1))
results = model.fit(disp=False)
mse = np.mean(results.resid ** 2)

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

with open(os.path.join(model_dir, "ndvi_statistical_real_info.json"), "w") as f:
    json.dump({
        "model_type": "Statistical SARIMA (Real Data & CPU)",
        "target": "NDVI",
        "rmse": float(mse**0.5),
        "mae": float(np.mean(np.abs(results.resid)))
    }, f, indent=2)
print("[SAVE] Model info saved.")