api.py

import os
import pandas as pd
import numpy as np
import io
import base64
import traceback
from flask import Flask, jsonify, request
from flask_cors import CORS
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

from data_loader import DataLoader
from feature_engineering import FeatureEngineer
from forecast_model import ForecastModel
from event_module import EventModule
from green_score import GreenScore
from explainability import ExplainabilityModule

app = Flask(__name__)
CORS(app)

# ─── Global In-Memory State ────────────────────────────────────────────────────
state = {
    'model':         None,
    'feature_names': None,
    'processed_df':  None,
    'raw_df':        None,
    'train_data':    None,
    'metrics':       None,
    'is_trained':    False,
}

DEFAULT_API_KEY = 'b736fbf9d44443ee9dd83589303e75a8'

# ─── Helpers ───────────────────────────────────────────────────────────────────

def _get_shap_base64():
    if not state['is_trained']:
        return None
    try:
        explainer = ExplainabilityModule(state['model'].model)
        fig = explainer.plot_shap_summary(state['train_data'])
        buf = io.BytesIO()
        fig.savefig(buf, format='png', bbox_inches='tight', dpi=100)
        buf.seek(0)
        img_b64 = base64.b64encode(buf.read()).decode('utf-8')
        plt.close(fig)
        return img_b64
    except Exception:
        return None


def _rolling_forecast_to_date(future_price, future_holiday, future_promo, target_end_date):
    """
    Recursive multi-step forecast.

    Starts from the last row in processed_df, predicts weekly until target_end_date.
    Each week's predicted sales become Lag_1 for the next week — true rolling forecast.

    Returns list of dicts:
      { date, predicted_sales, conf_low, conf_high, month, week_label }
    """
    processed_df = state['processed_df']
    feature_names = state['feature_names']
    model = state['model']

    last_row = processed_df.iloc[-1]
    last_actual_date = pd.to_datetime(last_row['Date'])

    # Seed from actual history
    recent_sales = list(processed_df['Weekly_Sales_Units'].iloc[-4:].astype(float))
    lag_1 = float(last_row['Weekly_Sales_Units'])
    lag_2 = float(last_row.get('Lag_1', last_row['Weekly_Sales_Units']))

    # Static features (from last known row, overriding price/flags from user)
    base_feats = {
        'Unit_Price_USD':          float(future_price),
        'Unit_Cost_USD':           float(last_row.get('Unit_Cost_USD', 0)),
        'Current_Inventory_Units': float(last_row.get('Current_Inventory_Units', 0)),
        'Transport_Distance_km':   float(last_row.get('Transport_Distance_km', 0)),
        'Fuel_Price_Index':        float(last_row.get('Fuel_Price_Index', 0)),
        'Holiday_Flag':            1 if future_holiday else 0,
        'Promotion_Flag':          1 if future_promo else 0,
    }
    # Fill any remaining model features from last known row
    for feat in feature_names:
        if feat not in base_feats and feat not in ('Month', 'Week_Number', 'Lag_1', 'Lag_2', 'Rolling_Mean_4'):
            val = last_row.get(feat, 0)
            base_feats[feat] = float(val) if val is not None else 0.0

    forecasts = []
    current_date = last_actual_date + pd.Timedelta(weeks=1)

    while current_date.date() <= target_end_date:
        rolling_mean_4 = float(np.mean(recent_sales[-4:])) if len(recent_sales) >= 4 else float(np.mean(recent_sales))

        input_data = {
            **base_feats,
            'Month':          int(current_date.month),
            'Week_Number':    int(current_date.isocalendar().week),
            'Lag_1':          lag_1,
            'Lag_2':          lag_2,
            'Rolling_Mean_4': rolling_mean_4,
        }

        param_df   = pd.DataFrame([input_data])[feature_names]
        prediction = float(model.predict(param_df)[0])
        conf_low, conf_high = model.get_confidence_interval(prediction)

        forecasts.append({
            'date':            str(current_date.date()),
            'predicted_sales': round(prediction, 2),
            'conf_low':        round(float(conf_low), 2),
            'conf_high':       round(float(conf_high), 2),
            'month':           int(current_date.month),
            'week_label':      current_date.strftime('Week of %b %d'),
        })

        # Roll forward state
        lag_2 = lag_1
        lag_1 = prediction
        recent_sales.append(prediction)
        if len(recent_sales) > 8:
            recent_sales = recent_sales[-8:]

        current_date += pd.Timedelta(weeks=1)

    return forecasts


# ─── Routes ────────────────────────────────────────────────────────────────────

@app.route('/api/status', methods=['GET'])
def status():
    return jsonify({'is_trained': state['is_trained']})


@app.route('/api/upload_and_train', methods=['POST'])
def upload_and_train():
    """
    Accepts either:
      - multipart file upload (key="file")
      - or falls back to default local data files
    Trains the model and returns metrics + data preview.
    """
    file_path = None

    if 'file' in request.files:
        f = request.files['file']
        ext = f.filename.rsplit('.', 1)[-1].lower()
        file_path = f'temp_data.{ext}'
        f.save(file_path)
    else:
        for fp in ['temp_data.csv', 'temp_data.xlsx', 'Walmart.csv', '../Walmart.csv']:
            if os.path.exists(fp):
                file_path = fp
                break

    if not file_path or not os.path.exists(file_path):
        return jsonify({'error': 'No data file found. Please upload a file.'}), 400

    try:
        loader = DataLoader(file_path)
        raw_df = loader.load_data()

        fe = FeatureEngineer()
        processed_df = fe.create_features(raw_df)
        X_train, X_test, y_train, y_test, feature_names = fe.split_data(processed_df)

        model = ForecastModel()
        model.train(X_train[feature_names], y_train)
        metrics = model.evaluate(X_test[feature_names], y_test)

        state.update({
            'model':         model,
            'feature_names': feature_names,
            'processed_df':  processed_df,
            'raw_df':        raw_df,
            'train_data':    X_train[feature_names],
            'metrics':       metrics,
            'is_trained':    True,
        })

        preview = raw_df.tail(5).to_dict(orient='records')
        for row in preview:
            for k, v in row.items():
                if hasattr(v, 'item'):
                    row[k] = v.item()
                elif isinstance(v, float) and np.isnan(v):
                    row[k] = None

        last_date = str(raw_df['Date'].max().date())

        return jsonify({
            'status':     'success',
            'rows':       len(raw_df),
            'date_range': {'start': str(raw_df['Date'].min().date()), 'end': last_date},
            'metrics':    {k: round(float(v), 2) for k, v in metrics.items()},
            'preview':    preview,
            'columns':    raw_df.columns.tolist(),
            'avg_price':  round(float(raw_df['Unit_Price_USD'].mean()), 2),
            'avg_inventory': round(float(raw_df['Current_Inventory_Units'].mean()), 2),
            'last_data_date': last_date,
        })

    except Exception as e:
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500


@app.route('/api/forecast', methods=['POST'])
def forecast():
    """
    Rolling multi-week forecast from the last historical data point through
    all remaining weeks of March 2026.

    Each predicted week's sales value is fed back as Lag_1 for the next week,
    making this a true recursive multi-step forecast rather than a single-step
    next-week prediction.

    Inputs (JSON):
      - future_price   float   - price to apply to all forecast weeks
      - future_holiday bool    - holiday flag for forecast weeks
      - future_promo   bool    - promo flag for forecast weeks
      - api_key        str     - NewsAPI key (optional)

    Returns:
      - weekly_forecasts: per-week breakdown through March 2026
      - march_summary:    totals and averages across March weeks
      - chart:            full timeline (8 historical + all forecast weeks)
      - event, green, shap_image
    """
    if not state['is_trained']:
        return jsonify({'error': 'Model not trained yet. Please upload data first.'}), 400

    req         = request.json or {}
    future_price   = float(req.get('future_price', float(state['raw_df']['Unit_Price_USD'].mean())))
    future_holiday = bool(req.get('future_holiday', False))
    future_promo   = bool(req.get('future_promo', False))
    api_key        = req.get('api_key', DEFAULT_API_KEY)

    try:
        processed_df     = state['processed_df']
        last_row         = processed_df.iloc[-1]
        last_actual_date = pd.to_datetime(last_row['Date'])

        # Determine target end date: always project through end of March 2026
        target_end = pd.Timestamp('2026-03-31').date()
        # If data already covers March, still project forward 5 weeks as a fallback
        if last_actual_date.date() > pd.Timestamp('2026-03-31').date():
            target_end = (last_actual_date + pd.Timedelta(weeks=5)).date()

        # ── Generate rolling weekly forecasts ─────────────────────────────────
        weekly_forecasts = _rolling_forecast_to_date(
            future_price, future_holiday, future_promo, target_end
        )

        if not weekly_forecasts:
            return jsonify({
                'error': 'No forecast weeks to generate. Historical data already extends through the target period.'
            }), 400

        # ── Identify March 2026 weeks ─────────────────────────────────────────
        march_weeks = [w for w in weekly_forecasts if w['month'] == 3]
        display_weeks = march_weeks if march_weeks else weekly_forecasts

        # Headline = first March week (or first predicted week)
        headline  = display_weeks[0]
        base_forecast   = headline['predicted_sales']
        conf_low_h      = headline['conf_low']
        conf_high_h     = headline['conf_high']

        # Aggregate March stats
        march_sales        = [w['predicted_sales'] for w in display_weeks]
        march_total        = sum(march_sales)
        march_avg          = march_total / len(march_sales) if march_sales else 0

        # ── Chart data: 8 historical actuals + all forecast weeks ─────────────
        last_n      = processed_df.iloc[-8:]
        chart_dates = [str(d.date()) for d in last_n['Date']]
        chart_sales = [float(x) for x in last_n['Weekly_Sales_Units']]
        chart_types = ['actual'] * len(last_n)

        for w in weekly_forecasts:
            chart_dates.append(w['date'])
            chart_sales.append(w['predicted_sales'])
            # 'bridge' = weeks between last data and March; 'forecast' = March weeks
            chart_types.append('forecast' if w['month'] == 3 else 'bridge')

        # ── Event adjustment ──────────────────────────────────────────────────
        event_mod  = EventModule(api_key)
        news_texts = event_mod.fetch_news()
        impact     = event_mod.compute_event_impact_score(news_texts)

        alpha             = 0.2
        adj_factor        = 1 + (alpha * impact['adjusted_effect'])
        adjusted_forecast = float(base_forecast * adj_factor)
        adjustment_pct    = round(float(impact['adjusted_effect']) * 20, 2)
        march_adj_total   = round(march_total * adj_factor, 2)

        # ── Green score ───────────────────────────────────────────────────────
        transport_km      = float(last_row.get('Transport_Distance_km', 0))
        fuel_index        = float(last_row.get('Fuel_Price_Index', 0))
        current_inventory = float(last_row.get('Current_Inventory_Units', 0))

        gs             = GreenScore()
        green_score    = gs.compute_score(transport_km, fuel_index, current_inventory, adjusted_forecast)
        recommendation = gs.get_recommendation(green_score)

        # ── SHAP ─────────────────────────────────────────────────────────────
        shap_b64 = _get_shap_base64()

        return jsonify({
            'status': 'success',

            # Headline (first March week)
            'future_date':       headline['date'],
            'base_forecast':     round(base_forecast, 2),
            'conf_low':          round(conf_low_h, 2),
            'conf_high':         round(conf_high_h, 2),
            'adjusted_forecast': round(adjusted_forecast, 2),
            'adjustment_pct':    adjustment_pct,

            # Full weekly breakdown
            'weekly_forecasts': weekly_forecasts,

            # March 2026 summary
            'march_summary': {
                'total':          round(march_total, 0),
                'average_weekly': round(march_avg, 0),
                'weeks_count':    len(display_weeks),
                'adjusted_total': round(march_adj_total, 0),
                'last_data_date': str(last_actual_date.date()),
            },

            # Chart
            'chart': {
                'dates': chart_dates,
                'sales': chart_sales,
                'types': chart_types,
            },

            # Event
            'event': {
                'raw_score':     round(float(impact['score']), 3),
                'type':          impact['type'],
                'sensitivity':   impact['sensitivity'],
                'direction':     impact['direction'],
                'interpretation': impact['interpretation'],
                'news':          news_texts[:5] if news_texts else [],
            },

            # Green
            'green': {
                'score':          float(green_score),
                'recommendation': recommendation,
                'transport_km':   transport_km,
                'fuel_index':     fuel_index,
                'inventory':      current_inventory,
            },

            'shap_image': shap_b64,
        })

    except Exception as e:
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500


@app.route('/api/metrics', methods=['GET'])
def get_metrics():
    if not state['is_trained']:
        return jsonify({'error': 'Not trained'}), 400
    m = state['metrics']
    return jsonify({k: round(float(v), 2) for k, v in m.items()})


if __name__ == '__main__':
    # Attempt to preload if a local data file exists
    with app.app_context():
        for fp in ['temp_data.csv', 'temp_data.xlsx', 'Walmart.csv', '../Walmart.csv']:
            if os.path.exists(fp):
                try:
                    loader = DataLoader(fp)
                    raw_df = loader.load_data()
                    fe = FeatureEngineer()
                    processed_df = fe.create_features(raw_df)
                    X_train, X_test, y_train, y_test, feature_names = fe.split_data(processed_df)
                    model = ForecastModel()
                    model.train(X_train[feature_names], y_train)
                    metrics = model.evaluate(X_test[feature_names], y_test)
                    state.update({
                        'model': model, 'feature_names': feature_names,
                        'processed_df': processed_df, 'raw_df': raw_df,
                        'train_data': X_train[feature_names],
                        'metrics': metrics, 'is_trained': True,
                    })
                    print(f"[API] Auto-loaded & trained from {fp}")
                    break
                except Exception as e:
                    print(f"[API] Preload failed: {e}")

    app.run(port=5000, debug=True, use_reloader=False)