import yfinance as yf
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
import gradio as gr

def fetch_stock_data(tickers, start_date, end_date):
    data = {}
    for ticker in tickers.split(','):
        stock_data = yf.download(ticker.strip(), start=start_date, end=end_date)
        data[ticker.strip()] = stock_data
    return data

def preprocess_data(data, window_size=5):
    processed_data = {}
    feature_names = [f'Lag_{i}' for i in range(1, window_size + 1)]
    for ticker, df in data.items():
        df['Return'] = df['Adj Close'].pct_change()
        df['Future Return'] = df['Return'].shift(-1)
        df['Target'] = df['Future Return'].apply(lambda x: 1 if x > 0 else 0)
        
        for i in range(1, window_size + 1):
            df[f'Lag_{i}'] = df['Return'].shift(i)
        
        df = df.dropna()
        processed_data[ticker] = (df[feature_names], df['Target'])
    return processed_data, feature_names

def train_model(processed_data):
    models = {}
    for ticker, (X, y) in processed_data.items():
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
        model = LogisticRegression()
        model.fit(X_train, y_train)
        y_pred = model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        models[ticker] = model
        print(f'Model for {ticker} trained with accuracy: {accuracy:.2f}')
    return models

def predict(models, data, feature_names, window_size=5):
    predictions = {}
    for ticker, model in models.items():
        df = data[ticker]
        latest_data = df['Adj Close'].pct_change().dropna().iloc[-window_size:].values.reshape(1, -1)
        if latest_data.shape[1] == window_size:
            latest_data_df = pd.DataFrame(latest_data, columns=feature_names)
            prediction = model.predict(latest_data_df)
            predictions[ticker] = 'Buy' if prediction[0] == 1 else 'Sell'
        else:
            predictions[ticker] = 'Insufficient data'
    return predictions

def stock_prediction_app(tickers, start_date, end_date):
    stock_data = fetch_stock_data(tickers, start_date, end_date)
    processed_data, feature_names = preprocess_data(stock_data)
    models = train_model(processed_data)
    predictions = predict(models, stock_data, feature_names)
    return predictions

# Gradio Interface
inputs = [
    gr.Textbox(lines=2, placeholder="Enter tickers separated by commas (e.g., AAPL, MSFT, GOOGL)", label="Stock Tickers"),
    gr.Textbox(placeholder="Enter start date (YYYY-MM-DD)", label="Start Date"),
    gr.Textbox(placeholder="Enter end date (YYYY-MM-DD)", label="End Date")
]
outputs = gr.JSON()

interface = gr.Interface(
    fn=stock_prediction_app,
    inputs=inputs,
    outputs=outputs,
    title="Stock Prediction App",
    description="Enter stock tickers, start date, and end date to get buy/sell predictions."
)

interface.launch()