import pandas as pd
from prophet import Prophet
from prophet.diagnostics import cross_validation

class GeoMagModel:
    """
    A class for forecasting geomagnetic data using the Prophet model.
    """

    def __init__(self, changepoint_prior_scale=0.1, weekly_seasonality=True):
        """
        Initialize the GeoMagModel with Prophet configuration.

        Args:
            changepoint_prior_scale (float): Controls flexibility of changepoint detection.
            weekly_seasonality (bool): Whether to include weekly seasonality.
        """
        self.changepoint_prior_scale = changepoint_prior_scale
        self.weekly_seasonality = weekly_seasonality
        self.model = None

    @staticmethod
    def prepare_data(df):
        """
        Prepare the DataFrame for Prophet by renaming columns.

        Args:
            df (pd.DataFrame): Input DataFrame with 'timestamp' and 'Dst' columns.

        Returns:
            pd.DataFrame: A DataFrame with 'ds' (timestamp) and 'y' (Dst) columns.
        """
        if not {'timestamp', 'Dst'}.issubset(df.columns):
            raise ValueError("DataFrame must contain 'timestamp' and 'Dst' columns.")
        return df.rename(columns={'timestamp': 'ds', 'Dst': 'y'})

    def train(self, df):
        """
        Train the Prophet model using the provided DataFrame.

        Args:
            df (pd.DataFrame): DataFrame prepared for Prophet with 'ds' and 'y' columns.
        """
        df = self.prepare_data(df)

        self.model = Prophet(
            interval_width=0.7, 
            changepoint_prior_scale=self.changepoint_prior_scale
        )

        if self.weekly_seasonality:
            self.model.add_seasonality(name='weekly', period=7, fourier_order=5, prior_scale=10)

        self.model.fit(df)

    def forecast(self, periods=1, freq='H'):
        """
        Generate future forecasts with the trained model.

        Args:
            periods (int): Number of future periods to forecast. Default is 1.
            freq (str): Frequency of the forecast (e.g., 'H' for hours). Default is 'H'.

        Returns:
            pd.DataFrame: Forecast DataFrame with columns: 'ds', 'yhat', 'yhat_lower', 'yhat_upper'.
        """
        if self.model is None:
            raise ValueError("Model is not trained. Call train() before forecast().")

        future_dates = self.model.make_future_dataframe(periods=periods, freq=freq)
        forecast = self.model.predict(future_dates)
        return forecast[['ds', 'yhat', 'yhat_lower', 'yhat_upper']]

    def cross_validate(self, initial='36 hours', period='12 hours', horizon='1 hours'):
        """
        Perform cross-validation on the trained model.

        Args:
            initial (str): Initial training period for cross-validation.
            period (str): Frequency of making predictions.
            horizon (str): Forecast horizon for each prediction.

        Returns:
            pd.DataFrame: Cross-validation results with metrics.
        """
        if self.model is None:
            raise ValueError("Model is not trained. Call train() before cross_validate().")

        return cross_validation(self.model, initial=initial, period=period, horizon=horizon)

# Example Usage:
# geomag = GeoMagModel()
# df = pd.read_csv("geomagnetic_data.csv")
# geomag.train(df)
# forecast = geomag.forecast(periods=24, freq='H')
# print(forecast)