train.py

# file: train.py
"""Training eines DQN-Agents für BESS-Dispatch mit Konfigurations-Unterstützung.

Dieses Modul implementiert das Training eines Deep Q-Network (DQN) Agenten
für die Optimierung des Battery Energy Storage System (BESS) Dispatch.
"""

import argparse
import json
import signal
import sys
from pathlib import Path
from datetime import datetime
from typing import Optional, Dict, Any

import numpy as np
import torch as th
import torch.nn as nn
from stable_baselines3 import DQN
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.utils import set_random_seed
from stable_baselines3.common.callbacks import (
    CheckpointCallback,
    EvalCallback,
    CallbackList,
)

from config_loader import load_config, Config, save_config
from data_loader import load_data, split_data, get_data_statistics, MarketData
from market_env import BessMultiMarketEnv
from utils.logging_config import setup_logging, get_logger

# Logger für dieses Modul
logger = get_logger("train")

# Globale Variable für Graceful Shutdown
_shutdown_requested = False
_current_model = None
_current_config = None


def signal_handler(signum, frame):
    """Handler für SIGINT/SIGTERM - ermöglicht Graceful Shutdown.

    Bei Abbruch wird das aktuelle Modell als Checkpoint gespeichert.
    """
    global _shutdown_requested
    if _shutdown_requested:
        logger.warning("Zweites Interrupt-Signal empfangen, beende sofort...")
        sys.exit(1)

    _shutdown_requested = True
    logger.warning("Shutdown angefordert, speichere Checkpoint...")

    if _current_model is not None and _current_config is not None:
        try:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            checkpoint_dir = Path(_current_config.paths.checkpoints)
            checkpoint_dir.mkdir(parents=True, exist_ok=True)

            checkpoint_path = checkpoint_dir / f"bess_dqn_interrupted_{timestamp}.zip"
            _current_model.save(str(checkpoint_path))
            logger.info(f"Checkpoint gespeichert: {checkpoint_path}")

            config_path = checkpoint_dir / f"config_interrupted_{timestamp}.yaml"
            save_config(_current_config, str(config_path))
        except Exception as e:
            logger.error(f"Fehler beim Speichern des Checkpoints: {e}")

    sys.exit(0)


def get_activation_fn(name: str):
    """Hole Aktivierungsfunktion nach Namen.

    Args:
        name: Name der Aktivierungsfunktion (z.B. "ReLU", "LeakyReLU").

    Returns:
        PyTorch Aktivierungsfunktion-Klasse.
    """
    activations = {
        "ReLU": nn.ReLU,
        "LeakyReLU": nn.LeakyReLU,
        "Tanh": nn.Tanh,
        "ELU": nn.ELU,
        "GELU": nn.GELU,
    }
    return activations.get(name, nn.LeakyReLU)


def create_env(data: MarketData, config: Config, seed: int = 0) -> BessMultiMarketEnv:
    """Erstelle Environment aus Daten und Config.

    Args:
        data: Marktdaten für das Environment.
        config: Konfigurationsobjekt mit allen Parametern.
        seed: Random Seed für Reproduzierbarkeit.

    Returns:
        Konfiguriertes BessMultiMarketEnv.
    """
    dt_hours = data.dt_hours if data.dt_hours is not None else config.env.dt_hours
    if data.dt_hours is not None and abs(data.dt_hours - config.env.dt_hours) > 1e-6:
        logger.warning(
            "Config dt_hours=%.4f weicht von Datenauflösung %.4f ab. Verwende Datenauflösung.",
            config.env.dt_hours,
            data.dt_hours,
        )

    # b_max = P_max * dt / C
    b_max = config.battery.p_max_mw * dt_hours / config.battery.capacity_mwh

    return BessMultiMarketEnv(
        price=data.price,
        fr_signal=data.fr_signal,
        reserve_price=data.reserve_price,
        demand=data.demand,
        re_gen=data.re_gen,
        timestamps=data.timestamps,
        temperature=data.temperature,
        dt_hours=dt_hours,
        c_min=config.battery.soc_min,
        c_max=config.battery.soc_max,
        c_init=config.battery.soc_init,
        b_max=b_max,
        delta_fr=config.env.delta_fr,
        delta_load=config.env.delta_load,
        alpha_d=config.degradation.alpha_d,
        beta=config.degradation.beta,
        stack_K=config.env.stack_k if config.features.observation_stacking else 1,
        task=config.task,
        seed=seed,
        eta_charge=config.battery.eta_charge,
        eta_discharge=config.battery.eta_discharge,
        capacity_mwh=config.battery.capacity_mwh,
        p_max_mw=config.battery.p_max_mw,
        arbitrage_weight=config.reward.arbitrage_weight,
        degradation_weight=config.reward.degradation_weight,
        fr_penalty_weight=config.reward.fr_penalty_weight,
        reserve_weight=config.reward.reserve_weight,
        load_penalty_weight=config.reward.load_penalty_weight,
        enable_reserve_penalty=config.reward.enable_reserve_penalty,
        reserve_penalty_factor=config.reward.reserve_penalty_factor,
        enable_efficiency=config.market.enable_efficiency,
        enable_bid_ask_spread=config.market.enable_bid_ask_spread,
        spread_percent=config.market.spread_percent,
        enable_market_impact=config.market.enable_market_impact,
        impact_coefficient=config.market.impact_coefficient,
        price_min=config.market.price_min,
        price_max=config.market.price_max,
        n_power_levels=config.env.n_power_levels,
        n_reserve_levels=config.env.n_reserve_levels,
        reserve_max_fraction=config.env.reserve_max_fraction,
        time_encoding=config.features.time_encoding,
        use_calendar_time=config.features.use_calendar_time,
        enable_calendar=config.degradation.enable_calendar,
        calendar_rate=config.degradation.calendar_rate,
        calendar_soc_factor=config.degradation.calendar_soc_factor,
        enable_crate=config.degradation.enable_crate,
        crate_exponent=config.degradation.crate_exponent,
        crate_ref=config.degradation.crate_ref,
        enable_temperature=config.degradation.enable_temperature,
        temperature_ref_celsius=config.degradation.temperature_ref_celsius,
        activation_energy=config.degradation.activation_energy,
        use_economic_degradation=config.degradation.use_economic_degradation,
        degradation_reference_dod=config.degradation.reference_dod,
        degradation_cycle_life=config.degradation.cycle_life,
        replacement_cost_eur_per_mwh=config.degradation.replacement_cost_eur_per_mwh,
        end_of_life_capacity_fraction=config.degradation.end_of_life_capacity_fraction,
    )


def create_model(env, config: Config) -> DQN:
    """Erstelle DQN-Modell aus Config.

    Args:
        env: Gymnasium Environment (oder VecEnv).
        config: Konfigurationsobjekt mit Agent-Parametern.

    Returns:
        Konfiguriertes DQN-Modell.
    """
    policy_kwargs = dict(
        activation_fn=get_activation_fn(config.agent.network.activation),
        net_arch=config.agent.network.hidden_layers,
    )

    # TensorBoard-Pfad
    tb_log = config.paths.tensorboard if config.paths.tensorboard else None

    model = DQN(
        "MlpPolicy",
        env,
        learning_rate=config.agent.learning_rate,
        buffer_size=config.agent.buffer_size,
        learning_starts=config.agent.learning_starts,
        batch_size=config.agent.batch_size,
        gamma=config.agent.gamma,
        train_freq=config.agent.train_freq,
        target_update_interval=config.agent.target_update_interval,
        exploration_fraction=config.agent.exploration_fraction,
        exploration_initial_eps=config.agent.exploration_initial_eps,
        exploration_final_eps=config.agent.exploration_final_eps,
        gradient_steps=config.agent.gradient_steps,
        policy_kwargs=policy_kwargs,
        verbose=1,
        seed=config.training.seed,
        tensorboard_log=tb_log,
        device="auto",
    )

    return model


def setup_callbacks(config: Config, eval_env) -> CallbackList:
    """Erstelle Callbacks für Training.

    Args:
        config: Konfigurationsobjekt mit Pfaden und Training-Parametern.
        eval_env: Environment für Evaluation während des Trainings.

    Returns:
        CallbackList mit allen konfigurierten Callbacks.
    """
    callbacks = []

    # Checkpoint-Callback
    checkpoint_dir = Path(config.paths.checkpoints)
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    checkpoint_callback = CheckpointCallback(
        save_freq=config.training.save_freq,
        save_path=str(checkpoint_dir),
        name_prefix="bess_dqn",
        save_replay_buffer=False,
        save_vecnormalize=False,
    )
    callbacks.append(checkpoint_callback)

    # Evaluation-Callback
    if eval_env is not None:
        eval_callback = EvalCallback(
            eval_env,
            best_model_save_path=str(checkpoint_dir),
            log_path=str(Path(config.paths.logs)),
            eval_freq=config.training.save_freq,
            n_eval_episodes=config.training.eval_episodes,
            deterministic=True,
        )
        callbacks.append(eval_callback)

    return CallbackList(callbacks)


def run_evaluation(
    model: DQN,
    env: BessMultiMarketEnv,
    n_episodes: int = 1,
) -> Dict[str, Any]:
    """Führe Evaluation durch und sammle KPIs.

    Args:
        model: Trainiertes DQN-Modell.
        env: Environment für Evaluation.
        n_episodes: Anzahl der Evaluationsepisoden.

    Returns:
        Dictionary mit aggregierten Evaluationsergebnissen.
    """
    results = []

    for _ in range(n_episodes):
        obs, _ = env.reset()
        kpis = {
            "total_reward": 0.0,
            "economic_reward": 0.0,
            "he_cost": 0.0,  # Energiekosten
            "fr_penalty": 0.0,  # FR-Abweichung
            "deg_cost": 0.0,  # Degradation
            "reserve_reward": 0.0,  # Reserve-Vergütung
            "reserve_penalty": 0.0,
            "load_penalty": 0.0,
            "throughput_mwh": 0.0,
            "steps": 0,
            "soc_history": [],
            "price_history": [],
            "action_history": [],
        }

        done = False
        while not done:
            action, _ = model.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, info = env.step(int(action))
            done = terminated or truncated

            kpis["total_reward"] += reward
            kpis["economic_reward"] += info["economic_reward"]
            kpis["he_cost"] += info["he_cost"]
            kpis["fr_penalty"] += info["fr_penalty"]
            kpis["deg_cost"] += info["deg_cost"]
            kpis["reserve_reward"] += info["reserve_reward"]
            kpis["reserve_penalty"] += info["reserve_penalty"]
            kpis["load_penalty"] += info["load_penalty"]
            kpis["throughput_mwh"] += info["throughput_mwh"]
            kpis["steps"] += 1
            kpis["soc_history"].append(info["soc"])
            kpis["action_history"].append(info["b"])

        results.append(kpis)

    # Aggregiere über Episoden
    summary = {
        "mean_reward": np.mean([r["total_reward"] for r in results]),
        "std_reward": np.std([r["total_reward"] for r in results]),
        "mean_economic_reward": np.mean([r["economic_reward"] for r in results]),
        "mean_he_cost": np.mean([r["he_cost"] for r in results]),
        "mean_fr_penalty": np.mean([r["fr_penalty"] for r in results]),
        "mean_deg_cost": np.mean([r["deg_cost"] for r in results]),
        "mean_reserve_reward": np.mean([r["reserve_reward"] for r in results]),
        "mean_reserve_penalty": np.mean([r["reserve_penalty"] for r in results]),
        "mean_load_penalty": np.mean([r["load_penalty"] for r in results]),
        "mean_throughput_mwh": np.mean([r["throughput_mwh"] for r in results]),
        "episodes": results,
    }

    return summary


def find_latest_checkpoint(checkpoint_dir: str) -> Optional[str]:
    """Finde den neuesten Checkpoint im Verzeichnis.

    Args:
        checkpoint_dir: Verzeichnis mit Checkpoints.

    Returns:
        Pfad zum neuesten Checkpoint oder None.
    """
    checkpoint_path = Path(checkpoint_dir)
    if not checkpoint_path.exists():
        return None

    checkpoints = list(checkpoint_path.glob("bess_dqn_*.zip"))
    if not checkpoints:
        return None

    # Sortiere nach Änderungszeit
    latest = max(checkpoints, key=lambda p: p.stat().st_mtime)
    return str(latest)


def list_checkpoints(checkpoint_dir: str) -> list:
    """Liste alle verfügbaren Checkpoints auf.

    Args:
        checkpoint_dir: Verzeichnis mit Checkpoints.

    Returns:
        Liste von Checkpoint-Pfaden sortiert nach Datum.
    """
    checkpoint_path = Path(checkpoint_dir)
    if not checkpoint_path.exists():
        return []

    checkpoints = list(checkpoint_path.glob("bess_dqn_*.zip"))
    return sorted(checkpoints, key=lambda p: p.stat().st_mtime, reverse=True)


def train(config: Config, resume_from: Optional[str] = None) -> str:
    """Haupttraining-Funktion.

    Args:
        config: Konfigurationsobjekt.
        resume_from: Pfad zu Checkpoint zum Fortsetzen.
                    "latest" findet automatisch den neuesten Checkpoint.

    Returns:
        Pfad zum gespeicherten Modell.

    Raises:
        RuntimeError: Bei Fehlern während des Trainings.
    """
    global _current_model, _current_config

    # Seed setzen
    set_random_seed(config.training.seed)

    # Verzeichnisse erstellen
    for path_attr in ["checkpoints", "logs", "tensorboard"]:
        Path(getattr(config.paths, path_attr)).mkdir(parents=True, exist_ok=True)

    # Daten laden
    logger.info("=" * 60)
    logger.info("BESS DQN Training")
    logger.info("=" * 60)

    data = load_data(config)
    stats = get_data_statistics(data)
    logger.info(f"Daten: {data}")
    logger.info(f"Preis: {stats['price']['mean']:.1f} +/- {stats['price']['std']:.1f} EUR/MWh")
    logger.info(f"       [{stats['price']['min']:.1f}, {stats['price']['max']:.1f}]")

    # Train/Val Split
    train_data, val_data = split_data(data, train_ratio=0.8)
    logger.info(f"Training: {len(train_data)} Schritte, Validation: {len(val_data)} Schritte")

    # Environments erstellen
    def make_train_env():
        return create_env(train_data, config, seed=config.training.seed)

    def make_eval_env():
        return create_env(val_data, config, seed=config.training.seed + 1)

    train_env = DummyVecEnv([make_train_env])
    eval_env = make_eval_env()

    # Resume-Handling
    if resume_from == "latest":
        resume_from = find_latest_checkpoint(config.paths.checkpoints)
        if resume_from:
            logger.info(f"Neuester Checkpoint gefunden: {resume_from}")
        else:
            logger.warning("Kein Checkpoint gefunden, starte neues Training")
    elif resume_from == "list":
        checkpoints = list_checkpoints(config.paths.checkpoints)
        if checkpoints:
            logger.info("Verfügbare Checkpoints:")
            for i, cp in enumerate(checkpoints[:10]):  # Max 10 anzeigen
                logger.info(f"  {i+1}. {cp.name}")
        else:
            logger.info("Keine Checkpoints gefunden")
        return ""

    # Modell erstellen oder laden
    if resume_from and Path(resume_from).exists():
        logger.info(f"Lade Checkpoint: {resume_from}")
        model = DQN.load(resume_from, env=train_env)
    else:
        logger.info("Erstelle neues Modell...")
        model = create_model(train_env, config)

    # Für Graceful Shutdown
    _current_model = model
    _current_config = config

    logger.info(f"Netzwerk: {config.agent.network.hidden_layers}")
    logger.info(f"Lernrate: {config.agent.learning_rate}")
    logger.info(f"Buffer: {config.agent.buffer_size}")
    logger.info(f"Gamma: {config.agent.gamma}")

    # Callbacks
    callbacks = setup_callbacks(config, eval_env)

    # Training
    logger.info(f"Starte Training ({config.training.total_timesteps} Schritte)...")
    logger.info("-" * 60)

    try:
        model.learn(
            total_timesteps=config.training.total_timesteps,
            callback=callbacks,
            log_interval=config.training.log_interval,
            progress_bar=True,
        )
    except KeyboardInterrupt:
        logger.warning("Training durch Benutzer unterbrochen")
    except Exception as e:
        logger.error(f"Fehler während des Trainings: {e}")
        # Notfall-Checkpoint speichern
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        emergency_path = Path(config.paths.checkpoints) / f"bess_dqn_error_{timestamp}.zip"
        model.save(str(emergency_path))
        logger.info(f"Notfall-Checkpoint gespeichert: {emergency_path}")
        raise

    # Finales Modell speichern
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    final_path = Path(config.paths.checkpoints) / f"bess_dqn_final_{timestamp}.zip"
    model.save(str(final_path))
    logger.info(f"Modell gespeichert: {final_path}")

    # Config speichern
    config_path = Path(config.paths.checkpoints) / f"config_{timestamp}.yaml"
    save_config(config, str(config_path))

    # Finale Evaluation
    logger.info("=" * 60)
    logger.info("Finale Evaluation")
    logger.info("=" * 60)

    eval_results = run_evaluation(model, eval_env, n_episodes=config.training.eval_episodes)
    logger.info(f"Reward: {eval_results['mean_reward']:.2f} +/- {eval_results['std_reward']:.2f}")
    logger.info(f"Oekonomischer Gewinn: {eval_results['mean_economic_reward']:.2f}")
    logger.info(f"Energiekosten: {eval_results['mean_he_cost']:.2f}")
    logger.info(f"FR-Strafe: {eval_results['mean_fr_penalty']:.2f}")
    logger.info(f"Degradation: {eval_results['mean_deg_cost']:.4f}")
    logger.info(f"Reserve-Vergütung: {eval_results['mean_reserve_reward']:.2f}")
    logger.info(f"Reserve-Penalty: {eval_results['mean_reserve_penalty']:.2f}")
    logger.info(f"Load-Penalty: {eval_results['mean_load_penalty']:.2f}")
    logger.info(f"Durchsatz: {eval_results['mean_throughput_mwh']:.2f} MWh")

    # Ergebnisse speichern
    results_path = Path(config.paths.logs) / f"results_{timestamp}.json"
    with open(results_path, "w") as f:
        # Numpy-Arrays für JSON konvertieren
        serializable_results = {
            k: v for k, v in eval_results.items()
            if k != "episodes"
        }
        json.dump(serializable_results, f, indent=2)

    return str(final_path)


def main():
    """Haupteinstiegspunkt für das Training."""
    parser = argparse.ArgumentParser(description="BESS DQN Training")
    parser.add_argument(
        "--config",
        type=str,
        default=None,
        help="Pfad zur Config-Datei (default: configs/default.yaml)",
    )
    parser.add_argument(
        "--data",
        type=str,
        default=None,
        help="Pfad zur Datendatei (überschreibt config.data.path)",
    )
    parser.add_argument(
        "--resume",
        type=str,
        default=None,
        help="Pfad zum Checkpoint ('latest' für neuesten, 'list' zum Auflisten)",
    )
    parser.add_argument(
        "--timesteps",
        type=int,
        default=None,
        help="Trainingsschritte (überschreibt config)",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Random Seed (überschreibt config)",
    )
    parser.add_argument(
        "--log-level",
        type=str,
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        help="Logging-Level",
    )
    parser.add_argument(
        "--log-file",
        action="store_true",
        help="Log auch in Datei schreiben",
    )

    args = parser.parse_args()

    # Logging einrichten
    import logging
    level = getattr(logging, args.log_level)
    setup_logging(
        level=level,
        log_dir="logs" if args.log_file else None,
    )

    # Signal-Handler für Graceful Shutdown
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    # Config laden
    config = load_config(args.config)

    # CLI-Argumente überschreiben Config
    if args.data:
        config.data.path = args.data
    if args.timesteps:
        config.training.total_timesteps = args.timesteps
    if args.seed:
        config.training.seed = args.seed

    # Training starten
    train(config, resume_from=args.resume)


if __name__ == "__main__":
    main()