from app.ml.features.lag_features import add_lag_features
from app.ml.features.rolling_features import add_rolling_features
from app.ml.features.time_features import add_time_features
from app.ml.features.regional_features import add_regional_features
from app.ml.features.battery_features import add_battery_features
from typing import List, Optional
import pandas as pd


def build_price_features(
    df: pd.DataFrame,
    price_col: str = "real_time_price",
    lags: Optional[List[int]] = None,
    windows: Optional[List[int]] = None,
    regions: Optional[List[str]] = None,
    include_time: bool = True,
    include_regional: bool = True,
) -> pd.DataFrame:
    if lags is None:
        lags = [1, 5, 10, 15, 30, 60]

    if windows is None:
        windows = [5, 10, 15, 30, 60]

    result = df.copy()

    if price_col in result.columns:
        result = add_lag_features(result, price_col, lags)
        result = add_rolling_features(result, price_col, windows)

    if include_time and "timestamp" in result.columns:
        result = add_time_features(result)

    if include_regional and regions:
        result = add_regional_features(result, regions)

    return result


def build_battery_features(
    df: pd.DataFrame,
    price_df: pd.DataFrame,
    battery_col: str = "charge_level_mwh",
    capacity_col: str = "capacity_mwh",
    timestamp_col: str = "timestamp",
    battery_id_col: str = "battery_id",
) -> pd.DataFrame:
    result = df.copy()
    result = add_battery_features(result, price_df, battery_col, capacity_col, timestamp_col, battery_id_col)
    return result


__all__ = ["build_price_features", "build_battery_features"]