# app.py
import gradio as gr
import torch
import pandas as pd
from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
import yfinance as yf

# Load the fine-tuned RAG model and tokenizer
tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-base")
retriever = RagRetriever.from_pretrained("facebook/rag-sequence-base", index_name="custom")
model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-base", retriever=retriever)

# Function to fetch and preprocess ICICI Bank data
def fetch_and_preprocess_data():
    # Fetch ICICI Bank data using yfinance
    ticker = "ICICIBANK.NS"
    data = yf.download(ticker, start="2020-01-01", end="2023-01-01")
    
    # Calculate technical indicators
    data['MA_50'] = data['Close'].rolling(window=50).mean()
    data['MA_200'] = data['Close'].rolling(window=200).mean()
    
    return data

# Function to analyze trading data using the RAG model
def analyze_trading_data(question):
    # Fetch and preprocess data
    data = fetch_and_preprocess_data()
    
    # Prepare context for the RAG model
    context = (
        f"ICICI Bank stock data:\n"
        f"Latest Close Price: {data['Close'].iloc[-1]:.2f}\n"
        f"50-Day Moving Average: {data['MA_50'].iloc[-1]:.2f}\n"
        f"200-Day Moving Average: {data['MA_200'].iloc[-1]:.2f}\n"
    )
    
    # Combine question and context
    input_text = f"Question: {question}\nContext: {context}"
    
    # Tokenize the input
    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
    
    # Generate the answer using the RAG model
    outputs = model.generate(inputs['input_ids'])
    
    # Decode the output to get the answer
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return answer

# Gradio interface
iface = gr.Interface(
    fn=analyze_trading_data,
    inputs="text",
    outputs="text",
    title="ICICI Bank Trading Analysis",
    description="Ask any question about ICICI Bank's trading data and get a detailed analysis.",
    examples=[
        "What is the current trend of ICICI Bank stock?",
        "Is the 50-day moving average above the 200-day moving average?",
        "What is the latest closing price of ICICI Bank?"
    ]
)

# Launch the app
iface.launch()