# SPDX-FileCopyrightText: 2026 bartzbeielstein
# SPDX-License-Identifier: AGPL-3.0-or-later

"""
Task demo: compare baseline, covariate, and custom LightGBM forecasts against ground truth.

This script executes the baseline N-to-1 task, the covariate-enhanced N-to-1
pipeline, and a custom LightGBM model with optimized hyperparameters, then loads
the ground truth from ~/spotforecast2_data/data_test.csv using the safety-critical
load_actual_combined function from spotforecast2_safe, and plots Actual vs
Predicted using Plotly.

The plot includes:
    - Actual combined values (ground truth)
    - Baseline combined prediction (n2n_predict)
    - Covariate combined prediction (n2n_predict_with_covariates, default LGBM)
    - Custom LightGBM combined prediction (optimized hyperparameters, Europe/Berlin tz)

Safety-Critical Features:
    - Uses load_actual_combined from spotforecast2_safe for validated data loading
    - DemoConfig provides immutable configuration with sensible defaults
    - Path objects ensure cross-platform compatibility
    - Comprehensive error handling with file existence checks

Examples:
    Run the demo:

    >>> python tasks/task_demo.py

    Force training (case-insensitive boolean):

    >>> python tasks/task_demo.py --force_train false

    Save the plot as a single HTML file (default: task_demo_plot.html):

    >>> python tasks/task_demo.py --html

    Save to a specific path:

    >>> python tasks/task_demo.py --html results/plot.html
"""

from __future__ import annotations

import argparse
import warnings
from pathlib import Path
from typing import Optional

from lightgbm import LGBMRegressor

from spotforecast2_safe.processing.agg_predict import agg_predict
from spotforecast2_safe.processing.n2n_predict import n2n_predict
from spotforecast2_safe.processing.n2n_predict_with_covariates import (
    n2n_predict_with_covariates,
)
from spotforecast2_safe.manager.tools import _parse_bool
from spotforecast2_safe.manager.datasets import DemoConfig, load_actual_combined
from spotforecast2.manager.plotter import plot_actual_vs_predicted

warnings.simplefilter("ignore")


def main(
    force_train: bool = True,
    html_path: Optional[str] = None,
) -> None:
    """Run the demo, compute predictions for three models, and plot actual vs predicted."""
    DATA_PATH = "~/spotforecast2_data/data_test.csv"
    FORECAST_HORIZON = 24
    CONTAMINATION = 0.01
    WINDOW_SIZE = 72
    LAGS = 24
    TRAIN_RATIO = 0.8
    VERBOSE = True
    SHOW_PROGRESS = True
    FORCE_TRAIN = force_train

    WEIGHTS = [
        1.0,
        1.0,
        -1.0,
        -1.0,
        1.0,
        -1.0,
        1.0,
        1.0,
        1.0,
        -1.0,
        1.0,
    ]

    print("--- Starting task_demo: baseline, covariates, and custom LightGBM ---")

    # --- Baseline predictions ---
    baseline_predictions, _ = n2n_predict(
        columns=None,
        forecast_horizon=FORECAST_HORIZON,
        contamination=CONTAMINATION,
        window_size=WINDOW_SIZE,
        verbose=VERBOSE,
        show_progress=SHOW_PROGRESS,
        force_train=FORCE_TRAIN,
        model_dir="~/spotforecast2_models/task_demo_baseline",
    )

    baseline_combined = agg_predict(baseline_predictions, weights=WEIGHTS)

    # --- Covariate-enhanced predictions ---
    cov_predictions, _, _ = n2n_predict_with_covariates(
        forecast_horizon=FORECAST_HORIZON,
        contamination=CONTAMINATION,
        window_size=WINDOW_SIZE,
        lags=LAGS,
        train_ratio=TRAIN_RATIO,
        verbose=VERBOSE,
        show_progress=SHOW_PROGRESS,
        force_train=FORCE_TRAIN,
        model_dir="~/spotforecast2_models/task_demo_covariates",
    )

    covariates_combined = agg_predict(cov_predictions, weights=WEIGHTS)

    # --- Custom LightGBM predictions (optimized hyperparameters) ---
    custom_lgbm = LGBMRegressor(
        n_estimators=1059,
        learning_rate=0.04191323446625026,
        num_leaves=212,
        min_child_samples=54,
        subsample=0.5014650987802548,
        colsample_bytree=0.6080926628683118,
        random_state=42,
        verbose=-1,
    )
    custom_lgbm_predictions, _, _ = n2n_predict_with_covariates(
        forecast_horizon=FORECAST_HORIZON,
        contamination=CONTAMINATION,
        window_size=WINDOW_SIZE,
        lags=LAGS,
        train_ratio=TRAIN_RATIO,
        timezone="UTC",
        estimator=custom_lgbm,
        verbose=VERBOSE,
        show_progress=SHOW_PROGRESS,
        force_train=FORCE_TRAIN,
        model_dir="~/spotforecast2_models/task_demo_custom_lgbm",
    )
    custom_lgbm_combined = agg_predict(custom_lgbm_predictions, weights=WEIGHTS)

    # --- Debug output ---
    print("\n=== DEBUG INFO ===")
    print(f"Baseline combined shape: {baseline_combined.shape}")
    print(
        f"Baseline index: {baseline_combined.index[0]} to {baseline_combined.index[-1]}"
    )
    print(f"Baseline values (first 5): {baseline_combined.head().values}")
    print(f"\nCovariates combined shape: {covariates_combined.shape}")
    print(
        f"Covariates index: {covariates_combined.index[0]} to {covariates_combined.index[-1]}"
    )
    print(f"Covariates values (first 5): {covariates_combined.head().values}")
    print(f"\nCustom LightGBM combined shape: {custom_lgbm_combined.shape}")
    print(
        f"Custom LightGBM index: {custom_lgbm_combined.index[0]} to {custom_lgbm_combined.index[-1]}"
    )
    print(f"Custom LightGBM values (first 5): {custom_lgbm_combined.head().values}")
    print(
        f"\nAre indices aligned? {(baseline_combined.index == covariates_combined.index).all()}"
    )
    print(
        f"Baseline vs Covariates identical? {(baseline_combined.values == covariates_combined.values).all()}"
    )
    print(
        f"Baseline vs Custom LightGBM identical? {(baseline_combined.values == custom_lgbm_combined.values).all()}"
    )
    print(
        f"Covariates vs Custom LightGBM identical? {(covariates_combined.values == custom_lgbm_combined.values).all()}"
    )
    if not (baseline_combined.values == covariates_combined.values).all():
        diff = baseline_combined - covariates_combined
        print(f"Baseline - Covariates diff stats:\n{diff.describe()}")
    if not (covariates_combined.values == custom_lgbm_combined.values).all():
        diff_lgbm = covariates_combined - custom_lgbm_combined
        print(f"Covariates - Custom LightGBM diff stats:\n{diff_lgbm.describe()}")
    print("==================\n")

    # --- Ground truth ---
    columns = list(baseline_predictions.columns)
    # Use load_actual_combined from spotforecast2_safe with minimal config
    config = DemoConfig(data_path=Path(DATA_PATH).expanduser())
    actual_combined = load_actual_combined(
        config=config,
        columns=columns,
        forecast_horizon=FORECAST_HORIZON,
        weights=WEIGHTS,
    )

    # Align indices to predictions for clean plotting
    actual_combined = actual_combined.reindex(baseline_combined.index)

    # --- Plot ---
    plot_actual_vs_predicted(
        actual_combined=actual_combined,
        baseline_combined=baseline_combined,
        covariates_combined=covariates_combined,
        custom_lgbm_combined=custom_lgbm_combined,
        html_path=html_path,
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run the spotforecast2 demo task.")
    parser.add_argument(
        "--force_train",
        type=_parse_bool,
        default=True,
        help="Force training (true/false, case-insensitive).",
    )
    parser.add_argument(
        "--html",
        nargs="?",
        const="task_demo_plot.html",
        default=None,
        metavar="PATH",
        help="Save the plot as a single self-contained HTML file. Default path: task_demo_plot.html",
    )
    args = parser.parse_args()
    main(force_train=args.force_train, html_path=args.html)
