model.py

from supabase_client import supabase
import os
import random
import numpy as np
import pandas as pd
from datetime import datetime, timedelta
from dotenv import load_dotenv
from supabase import create_client, Client
from sklearn.preprocessing import LabelEncoder, MinMaxScaler
import tensorflow as tf
from tensorflow.keras.models import Model, load_model
from tensorflow.keras.layers import Input, LSTM, Dense
import uuid
import joblib
from pathlib import Path

# ===============================
# 3. CONSTANTS & ENCODING
# ===============================

MODEL_PATH = "/tmp/global_model.h5"
DATA_PATH = "/tmp/initial_data.csv"
SCALER_PATH = "/tmp/scaler.pkl"

REGIONS = ['urban', 'rural']
SEASONS = ['winter', 'spring', 'summer', 'autumn']
EVENTS = ['normal', 'fasting', 'guests', 'sickness', 'travel', 'meal_off']

REGION_MULTIPLIER = {'urban': 1.0, 'rural': 1.1}
SEASON_MULTIPLIER = {'winter': 1.1, 'spring': 1.0, 'summer': 0.9, 'autumn': 1.0}
EVENT_MULTIPLIER = {'normal': 1.0, 'fasting': 0.7, 'guests': 1.3, 'sickness': 0.5, 'travel': 0.3, 'meal_off': 0.2}

BASE_CONSUMPTION = {
    'rice': {'adult_male': 0.3, 'adult_female': 0.25, 'child': 0.15},
    'milk': {'adult_male': 0.2, 'adult_female': 0.18, 'child': 0.3},
    'potato': {'adult_male': 0.25, 'adult_female': 0.2, 'child': 0.15},
    'onion': {'adult_male': 0.1, 'adult_female': 0.1, 'child': 0.05}
}

le_region = LabelEncoder().fit(REGIONS)
le_season = LabelEncoder().fit(SEASONS)
le_event = LabelEncoder().fit(EVENTS)
le_product = None
scaler = None

# ===============================
# 4. UTILITIES & DATA FUNCTIONS
# ===============================

def get_season(date):
    month = date.month
    if month in [12, 1, 2]: return 'winter'
    elif month in [3, 4, 5]: return 'spring'
    elif month in [6, 7, 8]: return 'summer'
    return 'autumn'

def generate_family():
    return {
        'adult_male': random.randint(1, 3),
        'adult_female': random.randint(1, 3),
        'child': random.randint(0, 3)
    }, random.choice(REGIONS)

def calculate_base_consumption(fam, region, season, event, product):
    base = BASE_CONSUMPTION.get(product, {'adult_male': 0.1, 'adult_female': 0.1, 'child': 0.05})
    total = sum(base[k]*fam.get(k, 0) for k in base)
    total *= REGION_MULTIPLIER.get(region, 1.0)
    total *= SEASON_MULTIPLIER.get(season, 1.0)
    total *= EVENT_MULTIPLIER.get(event, 1.0)
    total *= np.random.normal(1, 0.05)
    return max(total, 0.01)

def generate_data(products, families=3, days=180):
    data = []
    for _ in range(families):
        fam, region = generate_family()
        start = datetime.today() - timedelta(days=days)
        for d in range(days):
            date = start + timedelta(days=d)
            season = get_season(date)
            event = random.choices(EVENTS, weights=[70,5,5,5,5,10], k=1)[0]
            for prod in products:
                cons = calculate_base_consumption(fam, region, season, event, prod)
                stock = random.uniform(1.0, 10.0)
                pred_days = stock / cons if cons > 0 else 1
                act_days = pred_days * np.random.normal(1, 0.1)
                finish_error = int(act_days - pred_days)
                data.append({
                    'date': date.strftime('%Y-%m-%d'),
                    'product': prod, 'region': region, 'season': season, 'event': event,
                    'adult_male': fam['adult_male'], 'adult_female': fam['adult_female'], 'child': fam['child'],
                    'consumption': cons,
                    'stock_quantity': stock,
                    'finish_error': finish_error, 'finish_days': act_days
                })
    return pd.DataFrame(data)

def is_new_product(product, existing_df):
    return product not in existing_df['product'].unique()

def generate_and_append_new_product(product, csv_path=DATA_PATH):
    print(f"Generating synthetic data for new product: {product}")
    new_df = generate_data([product], families=3, days=180)
    if os.path.exists(csv_path):
        existing = pd.read_csv(csv_path)
        combined = pd.concat([existing, new_df], ignore_index=True)
    else:
        combined = new_df
    combined.to_csv(csv_path, index=False)
    print(f"Product '{product}' added to dataset.")

def create_user_if_not_exists(user_id, user_name="test_user", email="test@example.com"):
    """Create a user if they don't exist in the database"""
    try:
        response = supabase.table("users").select("user_id").eq("user_id", user_id).execute()
        if response.data:
            print(f"User {user_id} already exists")
            return True
        new_user = {
            "user_id": user_id,
            "user_name": user_name,
            "email": email,
            "region": "urban",
            "adult_male": 1,
            "adult_female": 1,
            "child": 0
        }
        response = supabase.table("users").insert(new_user).execute()
        print(f"Created new user: {user_id}")
        return True
    except Exception as e:
        print(f"Error creating user: {e}")
        return False

def fetch_feedback_for_user(user_id):
    response = supabase.table("feedback_data").select("*").eq("user_id", user_id).execute()
    return pd.DataFrame(response.data) if response.data else pd.DataFrame()

def insert_feedback(user_id, df):
    if not create_user_if_not_exists(user_id):
        print("Failed to create user, cannot insert feedback")
        return
    df['user_id'] = user_id
    response = supabase.table("feedback_data").insert(df.to_dict(orient="records")).execute()
    print("Feedback inserted successfully")

def get_product_id(product_name):
    res = supabase.table("products").select("product_id").eq("product_name", product_name).execute()
    if res.data:
        return res.data[0]['product_id']
    new_id = str(uuid.uuid4())
    supabase.table("products").insert({
        'product_id': new_id,
        'product_name': product_name,
        'unit': 'unit'
    }).execute()
    return new_id

def store_predictions(user_id, predictions, user_input):
    rows = []
    for product, result in predictions.items():
        product_id = get_product_id(product)
        row = {
            'user_id': user_id,
            'product_id': product_id,
            'prediction_date': datetime.today().date().isoformat(),
            'predicted_consumption': result['predicted_consumption'],
            'predicted_finish_days': result['predicted_finish_days'],
            'predicted_finish_date': result['predicted_finish_date'],
            'predicted_error': result['predicted_finish_error'],
            'stock_quantity': user_input['stock'][product],
        }
        rows.append(row)
    supabase.table("prediction_outputs").insert(rows).execute()
    print("✅ Predictions stored successfully")

# ===============================
# 5. MODELING
# ===============================

def prepare_data(df, products):
    global le_product, scaler
    le_product = LabelEncoder().fit(products)
    df['region_enc'] = le_region.transform(df['region'])
    df['season_enc'] = le_season.transform(df['season'])
    df['event_enc'] = le_event.transform(df['event'])
    df['product_enc'] = le_product.transform(df['product'])

    df['date'] = pd.to_datetime(df['date'])
    df = df.sort_values('date')

    if scaler is None:
        scaler_local = MinMaxScaler()
        df[['adult_male','adult_female','child','consumption','stock_quantity']] = scaler_local.fit_transform(df[['adult_male','adult_female','child','consumption','stock_quantity']])
    else:
        scaler_local = scaler
        df[['adult_male','adult_female','child','consumption','stock_quantity']] = scaler_local.transform(df[['adult_male','adult_female','child','consumption','stock_quantity']])

    X, y1, y2, y3 = [], [], [], []
    seq_len = 7
    for p in df['product_enc'].unique():
        sub = df[df['product_enc'] == p].reset_index(drop=True)
        feats = sub[['adult_male','adult_female','child','consumption','stock_quantity','region_enc','season_enc','event_enc','product_enc']].values
        c = sub['consumption'].values
        err = sub['finish_error'].values
        days = sub['finish_days'].values
        for i in range(len(sub) - seq_len):
            X.append(feats[i:i + seq_len])
            y1.append(c[i + seq_len])
            y2.append(err[i + seq_len])
            y3.append(days[i + seq_len])
    return np.array(X), np.array(y1), np.array(y2), np.array(y3), scaler_local

def build_model(input_shape):
    inp = Input(shape=input_shape)  # (7, 9)
    x = LSTM(64)(inp)
    x = Dense(32, activation='relu')(x)
    out1 = Dense(1, name='daily_consumption_output')(x)
    out2 = Dense(1, name='finish_error_output')(x)
    out3 = Dense(1, name='finish_days_output')(x)
    model = Model(inputs=inp, outputs=[out1, out2, out3])
    model.compile(optimizer='adam', loss=tf.keras.losses.MeanSquaredError())
    return model

def load_or_train_model():
    global scaler
    model_path = MODEL_PATH

    if not os.path.exists(DATA_PATH):
        pd.DataFrame(columns=[
            'date', 'product', 'region', 'season', 'event',
            'adult_male', 'adult_female', 'child',
            'consumption', 'stock_quantity', 'finish_error', 'finish_days'
        ]).to_csv(DATA_PATH, index=False)

    df = pd.read_csv(DATA_PATH)
    if df.empty:
        df = generate_data(list(BASE_CONSUMPTION.keys()))
        df.to_csv(DATA_PATH, index=False)

    products = df['product'].unique().tolist()

    X, y1, y2, y3, scaler_obj = prepare_data(df, products)
    scaler = scaler_obj

    if os.path.exists(model_path):
        os.remove(model_path)

    model = build_model((X.shape[1], X.shape[2]))
    model.fit(
        X,
        {
            'daily_consumption_output': y1,
            'finish_error_output': y2,
            'finish_days_output': y3
        },
        epochs=10,
        batch_size=32,
        validation_split=0.1
    )
    model.save(model_path)
    joblib.dump(scaler, SCALER_PATH)

    print("✅ Model trained and saved.")
    print("✅ Scaler saved to scaler.pkl.")

    return model

def retrain_model_with_feedback(user_id):
    base_df = pd.read_csv(DATA_PATH)
    feedback_df = fetch_feedback_for_user(user_id)
    if feedback_df.empty:
        print("No feedback found for user:", user_id)
        return {"message": f"No feedback found for user {user_id}"}
    
    combined_df = pd.concat([base_df, feedback_df], ignore_index=True)
    X, y1, y2, y3, scaler_obj = prepare_data(combined_df, combined_df['product'].unique().tolist())
    model = build_model((X.shape[1], X.shape[2]))
    model.fit(
        X,
        {'daily_consumption_output': y1,
         'finish_error_output': y2,
         'finish_days_output': y3},
        epochs=10,
        batch_size=32,
        validation_split=0.1,
    )
    model.save(MODEL_PATH)
    joblib.dump(scaler_obj, SCALER_PATH)
    print("Retrained model saved.")
    return {"message": f"Retraining complete for user {user_id}"}

# ===============================
# 6. PREDICTION
# ===============================

def predict_user_input(user_input):
    global le_product, scaler, ml_model
    initial_df = pd.read_csv(DATA_PATH)
    predictions = {}

    for product in user_input['stock'].keys():
        if is_new_product(product, initial_df):
            generate_and_append_new_product(product)
            initial_df = pd.read_csv(DATA_PATH)

        products = initial_df['product'].unique().tolist()
        le_product = LabelEncoder().fit(products)
        product_enc = le_product.transform([product])[0]

        vec = []
        for _ in range(7):
            base = calculate_base_consumption(
                user_input['family'],
                user_input['region'],
                user_input['season'],
                user_input['event'],
                product
            )
            stock_val = user_input['stock'][product]
            raw = [
                user_input['family']['adult_male'],
                user_input['family']['adult_female'],
                user_input['family']['child'],
                base,
                stock_val
            ]
            df_input = pd.DataFrame(
                [raw],
                columns=['adult_male', 'adult_female', 'child', 'consumption', 'stock_quantity']
            )
            raw_scaled = scaler.transform(df_input)[0]

            region_enc = le_region.transform([user_input['region']])[0]
            season_enc = le_season.transform([user_input['season']])[0]
            event_enc = le_event.transform([user_input['event']])[0]

            features = list(raw_scaled) + [region_enc, season_enc, event_enc, product_enc]
            vec.append(features)

        vec = np.array(vec)[np.newaxis, :, :]  # shape = (1, 7, 9)

        y1, y2, y3 = ml_model.predict(vec, verbose=0)

        # -------- INVERSE TRANSFORM START --------
        stock_val = user_input['stock'][product]
        df_temp = pd.DataFrame([[
            user_input['family']['adult_male'],
            user_input['family']['adult_female'],
            user_input['family']['child'],
            0,  # placeholder
            stock_val
        ]], columns=['adult_male', 'adult_female', 'child', 'consumption', 'stock_quantity'])

        scaled_temp = scaler.transform(df_temp)
        scaled_temp[0][3] = y1[0][0]  # Replace scaled 'consumption' with predicted

        unscaled = scaler.inverse_transform(scaled_temp)[0]
        daily = float(unscaled[3])  # actual predicted consumption (inverse-transformed)
        # -------- INVERSE TRANSFORM END --------

        error = float(y2[0][0])
        days = float(y3[0][0])
        finish_date = datetime.today() + timedelta(days=days)

        predictions[product] = {
            'predicted_consumption': round(daily, 3),
            'predicted_finish_days': round(days, 2),
            'predicted_finish_date': finish_date.strftime('%Y-%m-%d'),
            'predicted_finish_error': round(error, 2)
        }

    return predictions


# ===============================
# 8. GLOBAL LOADING FOR FASTAPI
# ===============================

if os.path.exists(MODEL_PATH):
    ml_model = tf.keras.models.load_model(MODEL_PATH)
    print("✅ Loaded trained model for prediction.")
else:
    ml_model = load_or_train_model()

scaler = joblib.load(SCALER_PATH)
print("✅ Loaded scaler.")

# ===============================
# 7. EXAMPLE RUN
# ===============================

if __name__ == "__main__":
    model = load_or_train_model()

    user_input = {
        "family": {"adult_male": 2, "adult_female": 2, "child": 1},
        "region": "urban",
        "season": "summer",
        "event": "normal",
        "stock": {"rice": 5, "milk": 3, "chicken": 4}  # 'chicken' is new
    }

    results = predict_user_input(user_input)
    print(pd.DataFrame(results).T)

    feedback = pd.DataFrame([{
        'date': datetime.today().strftime('%Y-%m-%d'),
        'product': k,
        'region': user_input['region'],
        'season': user_input['season'],
        'event': user_input['event'],
        'adult_male': user_input['family']['adult_male'],
        'adult_female': user_input['family']['adult_female'],
        'child': user_input['family']['child'],
        'consumption': v['predicted_consumption'],
        'finish_error': v['predicted_finish_error'],
        'finish_days': v['predicted_finish_days'],
        'stock_quantity': user_input['stock'][k]
    } for k, v in results.items()])

    user_uuid = str(uuid.uuid5(uuid.NAMESPACE_DNS, "user001"))

    try:
        insert_feedback(user_uuid, feedback)
        store_predictions(user_uuid, results, user_input)
        print("Data insertion completed successfully!")
    except Exception as e:
        print(f"Error during data insertion: {e}")