main.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import pandas as pd
import numpy as np
import lightgbm as lgb
from typing import List, Dict, Any
import datetime

app = FastAPI()

class SheetData(BaseModel):
    data: List[Dict[str, Any]]
    days: int = 17  # Default updated to 17

def create_features_for_df(df_input, target_col):
    df = df_input.copy()
    
    # 1. Date Features
    df['month'] = df['date'].dt.month
    df['day_of_week'] = df['date'].dt.dayofweek
    df['day_of_year'] = df['date'].dt.dayofyear
    
    # 2. Time Series Features
    df['lag_1'] = df[target_col].shift(1)
    df['lag_7'] = df[target_col].shift(7)
    df['rolling_mean_7'] = df[target_col].shift(1).rolling(window=7).mean()
    
    return df

@app.get("/")
def home():
    return {"status": "PNW Predictor is Running"}

@app.post("/predict_forecast")
def predict_forecast(payload: SheetData):
    try:
        # 1. Prepare Initial Data
        df_history = pd.DataFrame(payload.data)
        df_history['date'] = pd.to_datetime(df_history['date'])
        df_history = df_history.sort_values('date').reset_index(drop=True)
        
        targets = ['food', 'coal', 'oil', 'uranium', 'lead', 'iron', 'bauxite', 
                   'gasoline', 'munitions', 'steel', 'aluminum', 'credits']
        
        # Train models
        models = {}
        feature_cols = ['month', 'day_of_week', 'day_of_year', 'lag_1', 'lag_7', 'rolling_mean_7']
        
        for target in targets:
            df_feat = create_features_for_df(df_history, target)
            train_data = df_feat.dropna(subset=feature_cols + [target])
            X = train_data[feature_cols]
            y = train_data[target]
            
            model = lgb.LGBMRegressor(n_estimators=150, learning_rate=0.05, random_state=42, n_jobs=1)
            model.fit(X, y)
            models[target] = model

        # 2. Recursive Forecasting Loop
        df_current = df_history.copy()
        future_predictions = []
        start_date = df_current['date'].max()
        
        # Loop for the requested number of days (default 17)
        for step in range(1, payload.days + 1):
            next_date = start_date + datetime.timedelta(days=step)
            next_row = {'date': next_date}
            
            for target in targets:
                # Add temp row to calculate rolling features correctly
                temp_row = pd.DataFrame({'date': [next_date]})
                df_temp = pd.concat([df_current, temp_row], ignore_index=True)
                
                df_feat = create_features_for_df(df_temp, target)
                X_future = df_feat.iloc[[-1]][feature_cols]
                
                pred_value = models[target].predict(X_future)[0]
                
                # --- CHANGE: Convert to Integer ---
                next_row[target] = int(round(pred_value))
            
            df_current = pd.concat([df_current, pd.DataFrame([next_row])], ignore_index=True)
            
            response_row = next_row.copy()
            response_row['date'] = next_row['date'].strftime('%Y-%m-%d')
            future_predictions.append(response_row)

        return {
            "status": "success",
            "forecast_start": (start_date + datetime.timedelta(days=1)).strftime('%Y-%m-%d'),
            "forecast": future_predictions
        }

    except Exception as e:
        import traceback
        traceback.print_exc()
        return {"status": "error", "detail": str(e)}