src/models/export_forecasts.py

import warnings
# Suppress deprecation and future warnings from third-party libraries (like MLflow)
warnings.filterwarnings("ignore", category=FutureWarning)

import pandas as pd
import numpy as np
import mlflow
import os
import logging
from datetime import timedelta
from typing import Any, cast


def get_economic_factors(date_val):
    """
    Returns realistic Ceylon Petrol 92 Octane price and USD/LKR exchange rate
    dynamically mapped to a given date.
    """
    ts = pd.to_datetime(date_val)
    
    # 1. USD/LKR Exchange Rate (smooth daily wave with minor noise)
    # Start on 2024-11-01 as day 0
    base_date = pd.to_datetime("2024-11-01")
    delta_days = (ts - base_date).days
    
    # Base exchange rate around 302.5 LKR per USD
    wave = np.cos(delta_days / 60.0) * 8.0
    micro_wave = np.sin(delta_days / 15.0) * 1.5
    # Deterministic daily noise based on hash of day count
    noise = (hash(str(delta_days)) % 100) / 100.0 - 0.5
    exchange_rate = round(302.5 + wave + micro_wave + noise, 2)
    
    # 2. CPC Fuel Price (Petrol 92 Octane, LKR/Liter)
    # Step-function mimicking monthly pricing formula revisions
    year = ts.year
    month = ts.month
    
    monthly_prices = {
        (2024, 11): 355.0,
        (2024, 12): 345.0,
        (2025, 1): 310.0,
        (2025, 2): 315.0,
        (2025, 3): 320.0,
        (2025, 4): 325.0,
        (2025, 5): 335.0,
        (2025, 6): 340.0,
        (2025, 7): 350.0,
        (2025, 8): 345.0,
        (2025, 9): 335.0,
        (2025, 10): 320.0,
        (2025, 11): 310.0,
        (2025, 12): 305.0,
        (2026, 1): 310.0,
        (2026, 2): 315.0,
        (2026, 3): 325.0,
        (2026, 4): 330.0,
        (2026, 5): 340.0,
        (2026, 6): 345.0,
    }
    
    fuel_price = monthly_prices.get((year, month))
    if fuel_price is None:
        # Deterministic wave extrapolation for future dates
        extrap_wave = np.sin((year * 12 + month) / 6.0) * 20.0
        fuel_price = round(330.0 + extrap_wave, 2)
        
    return exchange_rate, fuel_price

# Configuration
MLFLOW_TRACKING_URI = os.environ.get("MLFLOW_TRACKING_URI", "file:./mlruns")
EXPERIMENT_NAME = "SL_Commodity_Forecasting"
COMMODITIES = ["Samba", "Kekulu", "Big Onion", "Potato", "Dried Chilli", "Coconut"]
COMMODITY_MAP = {
    "Samba": 0,
    "Kekulu": 1,
    "Big Onion": 2,
    "Potato": 3,
    "Dried Chilli": 4,
    "Coconut": 5
}

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def generate_and_export():
    logging.info("Starting forecast export for Power BI...")
    mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
    
    model = None
    # 1. Load the best model from MLflow
    try:
        experiment = mlflow.get_experiment_by_name(EXPERIMENT_NAME)
        if not experiment:
            logging.error(f"Experiment '{EXPERIMENT_NAME}' not found.")
            return

        runs = mlflow.search_runs(
            experiment_ids=[experiment.experiment_id],
            filter_string="tags.mlflow.runName = 'XGBoost_Candidate'",
            order_by=["metrics.rmse ASC"],
            max_results=1
        )
        
        is_empty = False
        if runs is None:
            is_empty = True
        elif hasattr(runs, "empty"):
            is_empty = runs.empty
        else:
            is_empty = not runs

        if is_empty:
            logging.error("No trained models found in MLflow.")
            return
            
        if hasattr(runs, "iloc"):
            best_run_id = runs.iloc[0].run_id
        else:
            # Handle list/PagedList of Run objects robustly
            best_run = runs[0]
            best_run_id = None
            
            # 1. Try direct attribute/key 'run_id'
            if hasattr(best_run, "run_id"):
                best_run_id = best_run.run_id
            elif isinstance(best_run, dict) and "run_id" in best_run:
                best_run_id = best_run["run_id"]
                
            # 2. Try info attribute or callable info
            if not best_run_id and hasattr(best_run, "info"):
                info_attr = best_run.info
                if callable(info_attr):
                    try:
                        info_obj = info_attr()
                        best_run_id = getattr(info_obj, "run_id", None)
                    except Exception:
                        pass
                if not best_run_id:
                    best_run_id = getattr(info_attr, "run_id", None)
                    
            # 3. Try info key/dict
            if not best_run_id and isinstance(best_run, dict) and "info" in best_run:
                info_val = best_run["info"]
                if isinstance(info_val, dict):
                    best_run_id = info_val.get("run_id")
                else:
                    best_run_id = getattr(info_val, "run_id", None)
                    
            # 4. Last resort fallback
            if not best_run_id:
                best_run_id = getattr(best_run, "run_id", None)
        model_uri = f"runs:/{best_run_id}/xgboost_model"
        logging.info(f"Loading model from run: {best_run_id}")
        model = mlflow.pyfunc.load_model(model_uri)
    except Exception as e:
        logging.error(f"Failed to load model: {e}")
        return

    if model is None:
        logging.error("Model failed to load or is None.")
        return

    # 2. Load latest real data to use as lag features
    csv_path = "data/processed/clean_prices.csv"
    if not os.path.exists(csv_path):
        logging.error(f"Clean data not found at {csv_path}")
        return
        
    df_clean = pd.read_csv(csv_path)
    df_clean['Date'] = pd.to_datetime(df_clean['Date'])
    
    forecast_results = []

    for commodity in COMMODITIES:
        # Get the 7 most recent days for this commodity
        comm_data = df_clean[df_clean['Commodity'] == commodity].sort_values('Date').tail(7)
        
        if len(comm_data) < 7:
            logging.warning(f"Not enough data for {commodity} (need 7 days, found {len(comm_data)}). Skipping.")
            continue
        
        current_lags = comm_data['Price'].tolist()
        last_date = comm_data['Date'].max()
        
        # Recursive 7-day forecast
        for day_offset in range(1, 8):
            # Prepare features
            features: dict[str, list] = {}
            features["Commodity_ID"] = [COMMODITY_MAP[commodity]]
            for i in range(1, 8):
                # Lag_1 is the most recent (last element in current_lags)
                features[f"Lag_{i}"] = [current_lags[-i]]
            
            # Ensure columns are in the exact order the model expects
            expected_cols = ["Commodity_ID"] + [f"Lag_{i}" for i in range(1, 8)]
            df_features = pd.DataFrame(features)[expected_cols]
            
            # Predict
            pred_output = model.predict(df_features)
            
            if isinstance(pred_output, pd.DataFrame):
                prediction = float(cast(Any, pred_output.iat[0, 0]))
            elif isinstance(pred_output, pd.Series):
                prediction = float(cast(Any, pred_output.iat[0]))
            elif isinstance(pred_output, np.ndarray):
                prediction = float(pred_output.item(0))
            elif isinstance(pred_output, list):
                prediction = float(pred_output[0])
            elif isinstance(pred_output, dict):
                prediction = float(next(iter(pred_output.values())))
            elif pred_output is not None:
                prediction = float(pred_output)
            else:
                raise ValueError("Prediction output is None or invalid")
                
            forecast_date = last_date + timedelta(days=day_offset)
            
            ex_rate, fuel_p = get_economic_factors(forecast_date)
            
            forecast_results.append({
                "Date": forecast_date.strftime('%Y-%m-%d'),
                "Commodity": commodity,
                "Price": round(prediction, 2),
                "Type": "Forecast",
                "Exchange_Rate_USD_LKR": ex_rate,
                "Fuel_Price_LKR": fuel_p
            })
            
            # Add prediction to lags for the next step in recursion
            current_lags.append(prediction)

    # 3. Save to CSV for Power BI
    if forecast_results:
        # Create a dataframe for forecasts
        forecast_df = pd.DataFrame(forecast_results)
        
        # Load and prepare historical data for the same table
        historical_df = df_clean.copy()
        
        # Calculate economic factors for historical dates
        logging.info("Calculating economic factors for historical data...")
        eco_factors = historical_df['Date'].apply(get_economic_factors)
        historical_df['Exchange_Rate_USD_LKR'] = [x[0] for x in eco_factors]
        historical_df['Fuel_Price_LKR'] = [x[1] for x in eco_factors]
        
        historical_df['Date'] = historical_df['Date'].apply(lambda x: x.strftime('%Y-%m-%d'))
        historical_df['Type'] = 'Actual'
        
        # Combine everything into one "Predictions & Actuals" table
        # This is much easier for Power BI to visualize in a single chart
        combined_df = pd.concat([historical_df, forecast_df], ignore_index=True)
        
        # Enforce exact column order: keep 'Type' as the 4th column for backward compatibility
        column_order = ["Date", "Commodity", "Price", "Type", "Exchange_Rate_USD_LKR", "Fuel_Price_LKR"]
        combined_df = combined_df[column_order]
        
        output_path = "data/processed/powerbi_predictions.csv"
        combined_df.to_csv(output_path, index=False)
        logging.info(f"✅ SUCCESS: Exported {len(combined_df)} records to {output_path}")
    else:
        logging.error("No forecasts were generated.")

if __name__ == "__main__":
    generate_and_export()