import torch
import transformers
import numpy as np
import gradio as gr
import openai
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt
from datetime import datetime, timedelta

# Load Hugging Face model for financial insights generation
generator = transformers.pipeline("text-generation", model="gpt2")

# Sample function for fetching stock market data (you can use libraries like yfinance or Alpha Vantage)
def get_stock_data(symbol, start_date, end_date):
    # This is just a placeholder. In production, use an API to get actual stock data
    dates = np.array([start_date + timedelta(days=i) for i in range((end_date - start_date).days)])
    prices = np.sin(np.linspace(0, 20, len(dates))) * 100 + 1000  # Placeholder for stock prices
    return dates, prices

# Preprocessing stock data for prediction
def preprocess_data(prices):
    scaler = MinMaxScaler(feature_range=(0, 1))
    prices_scaled = scaler.fit_transform(prices.reshape(-1, 1))
    return scaler, prices_scaled

# LSTM-based deep learning model for stock price prediction
class StockPredictor(torch.nn.Module):
    def __init__(self, input_size=1, hidden_size=50, num_layers=1):
        super(StockPredictor, self).__init__()
        self.lstm = torch.nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = torch.nn.Linear(hidden_size, 1)

    def forward(self, x):
        out, _ = self.lstm(x)
        out = self.fc(out[:, -1, :])
        return out

# Function to train the model on historical stock data
def train_model(prices_scaled):
    model = StockPredictor()
    criterion = torch.nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    
    # Prepare data for training
    X_train, y_train = [], []
    for i in range(len(prices_scaled) - 1):
        X_train.append(prices_scaled[i:i+10])  # Using the last 10 days to predict next day's price
        y_train.append(prices_scaled[i+10])
    
    X_train = torch.tensor(np.array(X_train), dtype=torch.float32)
    y_train = torch.tensor(np.array(y_train), dtype=torch.float32)
    
    # Train the model
    for epoch in range(100):
        model.train()
        optimizer.zero_grad()
        outputs = model(X_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()

    return model

# Predict future stock prices
def predict_stock(model, scaler, input_data):
    model.eval()
    with torch.no_grad():
        predictions = model(torch.tensor(input_data, dtype=torch.float32))
    return scaler.inverse_transform(predictions.detach().numpy())

# Generate financial insights using Hugging Face's GPT model
def generate_financial_insight(stock_symbol):
    prompt = f"Generate a detailed financial report for the stock symbol {stock_symbol} based on recent market trends."
    response = generator(prompt, max_length=150, num_return_sequences=1)
    return response[0]['generated_text']

# Gradio interface function
def financial_advisor(symbol):
    start_date = datetime.now() - timedelta(days=365)  # Last year data
    end_date = datetime.now()
    
    # Get stock data (for simplicity, using dummy data here)
    dates, prices = get_stock_data(symbol, start_date, end_date)
    
    # Preprocess data for prediction
    scaler, prices_scaled = preprocess_data(prices)
    
    # Train the model on historical data
    model = train_model(prices_scaled)
    
    # Predict future prices (for simplicity, predicting the next 30 days)
    future_data = prices_scaled[-10:].reshape(1, 10, 1)
    predictions = predict_stock(model, scaler, future_data)
    
    # Generate financial insights
    insight = generate_financial_insight(symbol)
    
    # Plot predicted prices
    plt.plot(dates, prices, label="Historical Prices")
    predicted_dates = [dates[-1] + timedelta(days=i) for i in range(1, 31)]
    plt.plot(predicted_dates, predictions, label="Predicted Prices", linestyle="--")
    plt.xlabel('Date')
    plt.ylabel('Price')
    plt.legend()
    plt.xticks(rotation=45)
    
    # Return results to display in Gradio
    return plt.gcf(), f"Financial Insight for {symbol}: {insight}"

# Launch Gradio interface
gr.Interface(
    fn=financial_advisor, 
    inputs=[gr.Textbox(label="Stock Symbol", placeholder="Enter Stock Symbol (e.g., AAPL)")], 
    outputs=[gr.Plot(), gr.Textbox()]
).launch()