from transformers import pipeline
import joblib
import pandas as pd
import numpy as np
from datetime import datetime, timedelta

class PotatoPricePredictor:
    def __init__(self):
        self.model = joblib.load('best_potato_price_model_Ridge.joblib')
        self.sentiment_analyzer = pipeline("sentiment-analysis")

    def preprocess(self, data):
        df = pd.DataFrame([data])
        df['Date'] = pd.to_datetime(df['Date'])
        
        df['DayOfWeek'] = df['Date'].dt.dayofweek
        df['Month'] = df['Date'].dt.month
        df['Quarter'] = df['Date'].dt.quarter
        df['Year'] = df['Date'].dt.year
        
        df['Events_Sentiment'] = df['Events'].apply(lambda x: self.sentiment_analyzer(x)[0]['score'] if x else 0)
        df['Impacts_Sentiment'] = df['Impacts'].apply(lambda x: self.sentiment_analyzer(x)[0]['score'] if x else 0)
        
        return df

    def predict(self, data):
        processed_data = self.preprocess(data)
        features = ['ArrivalQuantity', 'Temperature', 'Humidity', 'Wind direction',
                    'Events_Sentiment', 'Impacts_Sentiment', 'DayOfWeek', 'Month', 'Quarter', 'Year',
                    'PriceLag1', 'PriceLag7', 'PriceRollingMean7', 'PriceRollingStd7', 'PrevWeekAvgPrice']
        
        X = processed_data[features]
        prediction = self.model.predict(X)
        
        return {'predicted_price': float(prediction[0])}

    def predict_future(self, days=30):
        last_date = datetime.now().date()
        future_dates = [last_date + timedelta(days=i) for i in range(1, days + 1)]
        
        future_prices = []
        last_price = 50  # You may want to adjust this initial value

        for date in future_dates:
            data = {
                'Date': date.strftime('%Y-%m-%d'),
                'ArrivalQuantity': 1000,  # You may want to randomize or adjust these values
                'Temperature': 25,
                'Humidity': 60,
                'Wind direction': 180,
                'Events': 'Normal day',
                'Impacts': 'No significant impacts',
                'PriceLag1': last_price,
                'PriceLag7': last_price,
                'PriceRollingMean7': last_price,
                'PriceRollingStd7': 2,
                'PrevWeekAvgPrice': last_price
            }
            
            prediction = self.predict(data)
            future_prices.append(prediction['predicted_price'])
            last_price = prediction['predicted_price']

        return {'future_prices': [{'date': date.strftime('%Y-%m-%d'), 'price': price} for date, price in zip(future_dates, future_prices)]}

predictor = PotatoPricePredictor()

def query(payload):
    if payload.get('predict_future'):
        days = payload.get('days', 30)
        return predictor.predict_future(days)
    else:
        return predictor.predict(payload)