# app.py
import gradio as gr
import numpy as np
import pandas as pd
import yfinance as yf
import plotly.express as px

# Default stock options
default_stocks = ["RELIANCE.NS", "TCS.NS", "HDFCBANK.NS", "INFY.NS", "ITC.NS"]

def run_analysis(tickers, risk_free, n_portfolios):
    if len(tickers) < 2:
        raise gr.Error("Please select at least 2 stocks")
    
    # Download data
    data = yf.download(tickers, period="3y")['Close']
    
    # Calculate returns
    returns = data.pct_change().dropna()
    annual_returns = returns.mean() * 252
    cov_matrix = returns.cov() * 252
    
    # Generate portfolios
    results = []
    for _ in range(n_portfolios):
        weights = np.random.random(len(tickers))
        weights /= np.sum(weights)
        ret = np.dot(weights, annual_returns)
        risk = np.sqrt(weights.T @ cov_matrix @ weights)
        sharpe = (ret - risk_free/100) / risk
        results.append([ret, risk, sharpe, weights])
    
    results_df = pd.DataFrame(results, columns=['Return', 'Risk', 'Sharpe', 'Weights'])
    
    # Visualization
    fig = px.scatter(
        results_df, 
        x='Risk', 
        y='Return', 
        color='Sharpe',
        title=f"Risk-Return Profile of {n_portfolios} Random Portfolios",
        hover_data={'Weights': False}
    )
    
    # Highlight optimal portfolios
    max_sharpe = results_df.iloc[results_df['Sharpe'].idxmax()]
    min_risk = results_df.iloc[results_df['Risk'].idxmin()]
    
    fig.add_scatter(
        x=[max_sharpe['Risk']],
        y=[max_sharpe['Return']],
        mode='markers',
        marker=dict(size=15, symbol='star', color='gold'),
        name='Max Sharpe'
    )
    
    fig.add_scatter(
        x=[min_risk['Risk']],
        y=[min_risk['Return']],
        mode='markers',
        marker=dict(size=15, symbol='x', color='red'),
        name='Min Risk'
    )
    
    # Prepare allocation tables
    max_sharpe_df = pd.DataFrame({
        'Stock': tickers,
        'Weight': max_sharpe['Weights']
    }).set_index('Stock')
    
    min_risk_df = pd.DataFrame({
        'Stock': tickers,
        'Weight': min_risk['Weights']
    }).set_index('Stock')
    
    return (
        fig,
        f"Max Sharpe Portfolio - Return: {max_sharpe['Return']:.2%}, Risk: {max_sharpe['Risk']:.2%}, Sharpe: {max_sharpe['Sharpe']:.2f}",
        max_sharpe_df.style.format({'Weight': '{:.2%}'}).to_html(),
        f"Min Risk Portfolio - Return: {min_risk['Return']:.2%}, Risk: {min_risk['Risk']:.2%}, Sharpe: {min_risk['Sharpe']:.2f}",
        min_risk_df.style.format({'Weight': '{:.2%}'}).to_html()
    )

# Define Gradio interface
with gr.Blocks(title="Indian Stock Portfolio Optimizer") as demo:
    gr.Markdown("# 📊 Indian Stock Portfolio Optimizer")
    
    with gr.Row():
        with gr.Column():
            tickers = gr.CheckboxGroup(
                label="Select Indian Stocks (NSE)",
                choices=default_stocks,
                value=default_stocks[:2]
            )
            risk_free = gr.Slider(
                label="Risk-Free Rate (%)",
                minimum=0.0,
                maximum=10.0,
                value=6.0,
                step=0.1
            )
            n_portfolios = gr.Slider(
                label="Number of Portfolios",
                minimum=100,
                maximum=1000,
                value=500,
                step=100
            )
            submit_btn = gr.Button("Run Analysis")
        
        with gr.Column():
            plot = gr.Plot(label="Efficient Frontier")
            max_sharpe_title = gr.Markdown()
            max_sharpe_table = gr.HTML()
            min_risk_title = gr.Markdown()
            min_risk_table = gr.HTML()
    
    submit_btn.click(
        fn=run_analysis,
        inputs=[tickers, risk_free, n_portfolios],
        outputs=[plot, max_sharpe_title, max_sharpe_table, min_risk_title, min_risk_table]
    )

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