import gradio as gr
import yfinance as yf
import requests
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
import time

# Polymarket GraphQL API endpoint
POLYMARKET_API = "https://api.polymarket.com/graphql"

# Function to fetch Polymarket data
def fetch_polymarket_data(search_term="S&P"):
    try:
        query = """
        query {
          markets(first: 5, searchTerm: "%s") {
            edges {
              node {
                id
                question
                outcomes {
                  name
                  price
                }
              }
            }
          }
        }
        """ % search_term
        response = requests.post(POLYMARKET_API, json={"query": query}, timeout=10)
        if response.status_code != 200:
            return None
        data = response.json()
        markets = data["data"]["markets"]["edges"]
        if not markets:
            return None
        
        # Parse the first relevant market
        for market in markets:
            node = market["node"]
            outcomes = node["outcomes"]
            if len(outcomes) >= 2:
                return {
                    "question": node["question"],
                    "outcomes": {outcome["name"]: float(outcome["price"]) for outcome in outcomes}
                }
        return None
    except Exception as e:
        return None

# Function to fetch Yahoo Finance data with retry
def fetch_yahoo_data(ticker, retries=3, delay=2):
    for attempt in range(retries):
        try:
            stock = yf.download(ticker, period="1y", auto_adjust=False, progress=False)
            if stock.empty or len(stock) < 2:
                return None, None, None, f"No data returned for ticker '{ticker}'. It may be invalid or lack sufficient history."
            daily_returns = stock["Close"].pct_change().dropna()
            if daily_returns.empty:
                return None, None, None, f"No valid returns calculated for ticker '{ticker}'. Insufficient price data."
            mu = daily_returns.mean() * 252  # Annualized drift
            sigma = daily_returns.std() * np.sqrt(252)  # Annualized volatility
            last_price = stock["Close"][-1]  # Use most recent unadjusted Close
            return mu, sigma, last_price, None
        except Exception as e:
            error_msg = f"Attempt {attempt + 1}/{retries} failed for ticker '{ticker}': {str(e)}"
            if attempt < retries - 1:
                time.sleep(delay)  # Wait before retrying
            else:
                return None, None, None, error_msg
    return None, None, None, f"Failed to fetch data for '{ticker}' after {retries} attempts."

# Monte Carlo Simulation with GBM
def monte_carlo_simulation(S0, mu, sigma, T, N, sims, risk_factor, pm_data):
    dt = 1 / 252  # Daily time step
    steps = int(T * 252)
    sim_paths = np.zeros((sims, steps + 1))
    sim_paths[:, 0] = S0

    # Adjust drift based on Polymarket probabilities
    if pm_data and "outcomes" in pm_data:
        outcomes = pm_data["outcomes"]
        bullish_prob = outcomes.get("Yes", 0.5) if "Yes" in outcomes else 0.5
        bearish_prob = 1 - bullish_prob
        mu_bull = mu * 1.2 * risk_factor
        mu_bear = mu * -0.5 * risk_factor
        mu_adjusted = bullish_prob * mu_bull + bearish_prob * mu_bear
    else:
        mu_adjusted = mu * risk_factor

    for t in range(1, steps + 1):
        Z = np.random.standard_normal(sims)
        sim_paths[:, t] = sim_paths[:, t-1] * np.exp(
            (mu_adjusted - 0.5 * sigma**2) * dt + sigma * np.sqrt(dt) * Z
        )
    
    return sim_paths

# Main simulation function
def run_simulation(investment, ticker, horizon, num_sims, risk_factor):
    # Validate inputs
    if not ticker or investment <= 0 or horizon <= 0 or num_sims <= 0:
        return None, "Please provide valid inputs: positive investment, horizon, and simulations, and a ticker."

    # Fetch data
    mu, sigma, S0, error_msg = fetch_yahoo_data(ticker)
    if mu is None:
        return None, error_msg
    
    pm_data = fetch_polymarket_data("S&P")
    
    # Run Monte Carlo simulation
    sim_paths = monte_carlo_simulation(S0, mu, sigma, horizon, num_sims, num_sims, risk_factor, pm_data)
    final_values = sim_paths[:, -1] * (investment / S0)  # Scale to investment amount
    
    # Create histogram
    fig, ax = plt.subplots(figsize=(8, 6))
    ax.hist(final_values, bins=50, color="skyblue", edgecolor="black")
    ax.set_title(f"Distribution of Final Investment Value ({ticker})")
    ax.set_xlabel("Final Value ($)")
    ax.set_ylabel("Frequency")
    plt.tight_layout()
    
    # Calculate stats
    mean_val = np.mean(final_values)
    min_val = np.min(final_values)
    max_val = np.max(final_values)
    std_val = np.std(final_values)
    
    # Prepare summary text
    summary = f"Market Data (Yahoo Finance):\n"
    summary += f"- Drift (μ): {mu:.4f} (based on unadjusted Close)\n"
    summary += f"- Volatility (σ): {sigma:.4f}\n"
    summary += f"- Last Close Price: ${S0:.2f}\n\n"
    if pm_data:
        summary += f"Polymarket Data:\n- Question: {pm_data['question']}\n"
        for outcome, prob in pm_data["outcomes"].items():
            summary += f"- {outcome}: {prob*100:.1f}%\n"
    else:
        summary += "Polymarket Data: No relevant market found or API unavailable.\n"
    summary += f"\nSimulation Results ({num_sims} runs):\n"
    summary += f"- Mean Final Value: ${mean_val:.2f}\n"
    summary += f"- Min Final Value: ${min_val:.2f}\n"
    summary += f"- Max Final Value: ${max_val:.2f}\n"
    summary += f"- Std Dev: ${std_val:.2f}"
    
    return fig, summary

# Gradio UI
with gr.Blocks(title="Investment Simulation Platform") as demo:
    gr.Markdown("# Investment Decision Simulation Platform")
    gr.Markdown("Simulate investment outcomes using Yahoo Finance data (unadjusted) and Polymarket probabilities.")
    
    with gr.Row():
        with gr.Column():
            investment = gr.Number(label="Investment Amount ($)", value=10000)
            ticker = gr.Textbox(label="Ticker Symbol (e.g., AAPL, ^GSPC)", value="AAPL")
            horizon = gr.Number(label="Investment Horizon (Years)", value=1)
            num_sims = gr.Number(label="Number of Simulations", value=1000)
            risk_factor = gr.Slider(0.5, 2.0, value=1.0, label="Risk Factor")
            submit_btn = gr.Button("Run Simulation")
        
        with gr.Column():
            plot_output = gr.Plot(label="Final Value Distribution")
            text_output = gr.Textbox(label="Simulation Summary", lines=12)
    
    submit_btn.click(
        fn=run_simulation,
        inputs=[investment, ticker, horizon, num_sims, risk_factor],
        outputs=[plot_output, text_output]
    )

if __name__ == "__main__":
    demo.launch()