# app.py
import os
import pandas as pd
import gradio as gr

from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score

from groq import Groq


# =========================
# 1. TRAIN DBS REGRESSION MODEL
# =========================

# Load dataset (make sure DBS_SingDollar.csv is in repo root)
df = pd.read_csv("DBS_SingDollar.csv")

df = df[["DBS", "SGD"]].dropna()

X = df[["SGD"]]
y = df["DBS"]

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

model = LinearRegression()
model.fit(X_train, y_train)

# Optional evaluation (printed in HF logs)
y_pred = model.predict(X_test)
rmse = mean_squared_error(y_test, y_pred) ** 0.5
r2 = r2_score(y_test, y_pred)

print(f"RMSE: {rmse:.4f}")
print(f"R²: {r2:.4f}")


def predict_dbs_price(sgd_rate):
    pred = model.predict([[sgd_rate]])[0]
    return f"Predicted DBS price: {pred:.2f}"


# =========================
# 2. GROQ LLM FUNCTION
# =========================

# HF → Settings → Secrets → GROQ_API_KEY
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

def groq_chat(text):
    completion = client.chat.completions.create(
        model="llama-3.1-8b-instant",
        messages=[{"role": "user", "content": text}]
    )
    return completion.choices[0].message.content


# =========================
# 3. GRADIO UI (TABS)
# =========================

with gr.Blocks(title="DBS Predictor & Groq Chat") as demo:

    gr.Markdown("# 📊 DBS Price Predictor & 🤖 Groq Chat")

    with gr.Tab("DBS Price Predictor"):
        gr.Markdown("Predict DBS share price using SGD exchange rate.")
        sgd_input = gr.Number(label="SGD Exchange Rate")
        dbs_output = gr.Textbox(label="Predicted DBS Price")
        predict_btn = gr.Button("Predict")

        predict_btn.click(
            fn=predict_dbs_price,
            inputs=sgd_input,
            outputs=dbs_output
        )

    with gr.Tab("Groq LLM Chat"):
        gr.Markdown("Chat with LLaMA 3 via Groq API.")
        chat_input = gr.Textbox(label="Enter your prompt", lines=8)
        chat_output = gr.Textbox(label="Response", lines=8)
        chat_btn = gr.Button("Ask")

        chat_btn.click(
            fn=groq_chat,
            inputs=chat_input,
            outputs=chat_output
        )


demo.launch()