import gradio as gr
import yfinance as yf
from pypfopt.discrete_allocation import DiscreteAllocation, get_latest_prices
from pypfopt import EfficientFrontier
from pypfopt import risk_models
from pypfopt import expected_returns
from pypfopt import plotting
import copy
import numpy as np
import pandas as pd
import plotly.express as px
import matplotlib.pyplot as plt
from datetime import datetime
import datetime

def plot_cum_returns(data, title):    
    daily_cum_returns = 1 + data.dropna().pct_change()
    daily_cum_returns = daily_cum_returns.cumprod()*100
    fig = px.line(daily_cum_returns, title=title)
    return fig

def plot_efficient_frontier_and_max_sharpe(mu, S): 
    # Optimize portfolio for max Sharpe ratio and plot it out with efficient frontier curve
    ef = EfficientFrontier(mu, S)
    fig, ax = plt.subplots(figsize=(6,4))
    ef_max_sharpe = copy.deepcopy(ef)
    plotting.plot_efficient_frontier(ef, ax=ax, show_assets=False)
    # Find the max sharpe portfolio
    ef_max_sharpe.max_sharpe(risk_free_rate=0.02)
    ret_tangent, std_tangent, _ = ef_max_sharpe.portfolio_performance()
    ax.scatter(std_tangent, ret_tangent, marker="*", s=100, c="r", label="Max Sharpe")
    # Generate random portfolios
    n_samples = 1000
    w = np.random.dirichlet(np.ones(ef.n_assets), n_samples)
    rets = w.dot(ef.expected_returns)
    stds = np.sqrt(np.diag(w @ ef.cov_matrix @ w.T))
    sharpes = rets / stds
    ax.scatter(stds, rets, marker=".", c=sharpes, cmap="viridis_r")
    # Output
    ax.legend()
    return fig

def output_results(start_date, end_date, tickers_string):
    tickers = tickers_string.split(',')
    
    # Get Stock Prices
    stocks_df = yf.download(tickers, start=start_date, end=end_date)['Adj Close']
    
    # Plot Individual Stock Prices
    fig_indiv_prices = px.line(stocks_df, title='Price of Individual Stocks')
        
    # Plot Individual Cumulative Returns
    fig_cum_returns = plot_cum_returns(stocks_df, 'Cumulative Returns of Individual Stocks Starting with $100')
    
    # Calculatge and Plot Correlation Matrix between Stocks
    corr_df = stocks_df.corr().round(2)
    fig_corr = px.imshow(corr_df, text_auto=True, title = 'Correlation between Stocks')

    # Calculate expected returns and sample covariance matrix for portfolio optimization later
    mu = expected_returns.mean_historical_return(stocks_df)
    S = risk_models.sample_cov(stocks_df)

    # Plot efficient frontier curve
    fig_efficient_frontier = plot_efficient_frontier_and_max_sharpe(mu, S)

    # Get optimized weights
    ef = EfficientFrontier(mu, S)
    ef.max_sharpe(risk_free_rate=0.02)
    weights = ef.clean_weights()
    expected_annual_return, annual_volatility, sharpe_ratio = ef.portfolio_performance()
    
    expected_annual_return, annual_volatility, sharpe_ratio = '{}%'.format((expected_annual_return*100).round(2)), \
    '{}%'.format((annual_volatility*100).round(2)), \
    '{}%'.format((sharpe_ratio*100).round(2))
    
    weights_df = pd.DataFrame.from_dict(weights, orient = 'index')
    weights_df = weights_df.reset_index()
    weights_df.columns = ['Tickers', 'Weights']

    # Calculate returns of portfolio with optimized weights
    stocks_df['Optimized Portfolio'] = 0
    for ticker, weight in weights.items():
        stocks_df['Optimized Portfolio'] += stocks_df[ticker]*weight

    # Plot Cumulative Returns of Optimized Portfolio
    fig_cum_returns_optimized = plot_cum_returns(stocks_df['Optimized Portfolio'], 'Cumulative Returns of Optimized Portfolio Starting with $100')

    return  fig_cum_returns_optimized, weights_df, fig_efficient_frontier, fig_corr,   \
            expected_annual_return, annual_volatility, sharpe_ratio, fig_indiv_prices, fig_cum_returns


with gr.Blocks() as app:
    with gr.Row():
        gr.HTML("<h1>Bohmian's Stock Portfolio Optimizer</h1>")
    
    with gr.Row():
        start_date = gr.Textbox("2013-01-01", label="Start Date")
        end_date = gr.Textbox(datetime.datetime.now().date(), label="End Date")
    
    with gr.Row():        
        tickers_string = gr.Textbox("MA,META,V,AMZN,JPM,BA", 
                                    label='Enter all stock tickers to be included in portfolio separated \
                                    by commas WITHOUT spaces, e.g. "MA,META,V,AMZN,JPM,BA"')
        btn = gr.Button("Get Optimized Portfolio")
       
    with gr.Row():
        gr.Markdown("Optimizied Portfolio Metrics")
        
    with gr.Row():
        expected_annual_return = gr.Text(label="Expected Annual Return")
        annual_volatility = gr.Text(label="Annual Volatility")
        sharpe_ratio = gr.Text(label="Sharpe Ratio")            
   
    with gr.Row():        
        fig_cum_returns_optimized = gr.Plot(label="Cumulative Returns of Optimized Portfolio (Starting Price of $100)")
        weights_df = gr.DataFrame(label="Optimized Weights of Each Ticker")
        
    with gr.Row():
        fig_efficient_frontier = gr.Plot(label="Efficient Frontier")
        fig_corr = gr.Plot(label="Correlation between Stocks")
    
    with gr.Row():
        fig_indiv_prices = gr.Plot(label="Price of Individual Stocks")
        fig_cum_returns = gr.Plot(label="Cumulative Returns of Individual Stocks Starting with $100")

    btn.click(fn=output_results, inputs=[start_date, end_date, tickers_string], 
              outputs=[fig_cum_returns_optimized, weights_df, fig_efficient_frontier, fig_corr,   \
            expected_annual_return, annual_volatility, sharpe_ratio, fig_indiv_prices, fig_cum_returns])

app.launch()