from turtle import width
import gradio as gr
import yfinance as yf
from retry import retry
from tqdm.auto import tqdm
import pandas as pd
import time
import mplfinance as mpf
from pypfopt import risk_models, expected_returns
from pypfopt import plotting
from pypfopt import EfficientFrontier, objective_functions
import numpy as np
from warnings import filterwarnings

filterwarnings("ignore")

data = pd.DataFrame()


def load_data(ticker_list="", start_year=2021):

    data = yf.download(
        list(map(lambda x: x.strip().upper(), ticker_list.split(','))),
        start=f'{start_year}-01-01',
        interval='1d',
        keepna=True,
        period="5y",
        rounding=True,
    )['Close'].reset_index().fillna(method='ffill')

    data.Date = pd.to_datetime(data.Date)
    return data


def _prep_evaluation_report(
        hist_data,
        backtest_days,
        investment_size, weights
):
    allocation_df = pd.DataFrame({k: [weights[k]] for k in weights}).T
    testing_data = hist_data[-backtest_days:].copy()

    allocation_df.columns = ['weights']
    init_df = testing_data.iloc[0]
    end_df = testing_data.iloc[-1]
    allocation_df = allocation_df.merge(
        init_df, left_index=True, right_index=True)

    allocation_df = allocation_df.merge(
        end_df, left_index=True, right_index=True)

    allocation_df.columns = ['weights', 'cost_price', 'sell_price']

    allocation_df['amount_allocation'] = allocation_df.weights * \
        investment_size

    allocation_df['unit_pnl'] = allocation_df.sell_price - \
        allocation_df.cost_price

    allocation_df['units'] = allocation_df.weights * \
        investment_size // allocation_df.cost_price

    allocation_df['value'] = allocation_df.units * allocation_df.cost_price

    allocation_df['pnl'] = allocation_df.units * allocation_df.unit_pnl

    profit = np.round(allocation_df.pnl.sum(), 2)
    cagr = np.round(
        (
            (
                (investment_size + profit)/investment_size
            )**(
                1/(backtest_days/365)
            ) - 1
        )*100,
        2
    )

    instrumnets = testing_data.columns.tolist()
    unit_map = allocation_df['units'].to_dict()

    for instrument in instrumnets:
        testing_data[instrument] -= testing_data[instrument].iloc[0]
        testing_data[instrument] *= unit_map.get(instrument, 0)

    return (
        allocation_df.round(2).reset_index(),
        profit,
        cagr,
        gr.BarPlot(
            pd.DataFrame(
                testing_data.sum(axis=1),
                columns=['Portfolio Value']
            ).reset_index(),
            x='Date',
            y='Portfolio Value',
        ),
    )


def re_optimize(
        hist_data: pd.DataFrame,
        backtest_days: int,
        investment_size: int
):
    _, allocation_df, _ = find_weights(
        hist_data,
        backtest_days,
        investment_size
    )

    return find_weights(
        hist_data.drop(
            columns=allocation_df[allocation_df.units == 0]['index']
        ),
        backtest_days,
        investment_size
    )


def find_weights(
        hist_data: pd.DataFrame,
        backtest_days: int,
        investment_size: int
):
    training_data = hist_data[:-backtest_days].copy().set_index(
        'Date'
    ).fillna(0).astype(float)

    mu = expected_returns.capm_return(training_data)

    S = risk_models.semicovariance(training_data)
    initial_weights = np.array(
        [1/len(training_data.columns)] * len(training_data.columns))

    ef = EfficientFrontier(mu, S)

    ef.add_objective(
        objective_functions.transaction_cost,
        w_prev=initial_weights,
        k=0.02
    )
    ef.add_objective(objective_functions.L2_reg)

    ef.max_sharpe()
    weights = ef.clean_weights()
    (
        expected_annual_return,
        annual_volatility,
        sharpe_ratio
    ) = ef.portfolio_performance()

    (
        pnl_df,
        profit,
        cagr,
        pnl_graph
    ) = _prep_evaluation_report(
        hist_data.set_index('Date'),
        backtest_days,
        investment_size,
        weights
    )

    return (
        f'''
> ## Portfolio Performance
|Expected Annual Return|Annual Volatility|Sharpe Ratio|Profit|Profit %|CAGR|
|:---:|:---:|:---:|:---:|:---:|:---:|
|{expected_annual_return*100:.2f}|{annual_volatility*100:.2f}%|{sharpe_ratio:.2f}|{profit:.2f}|{profit/investment_size*100:.2f}%|{cagr:.2f}|''',
        pnl_df.sort_values('pnl', ascending=False),
        pnl_graph
    )


with gr.Blocks(title='Portfolio Optimizer') as demo:
    # with gr.Blocks() as demo:

    gr.Label('Portfolio Optimizer')

    with gr.Accordion('About'):

        gr.Markdown(
            '''
    # Portfolio Optimizer

    This is a simple portfolio optimizer that uses the Efficient Frontier
    to optimize the portfolio weights. The optimizer uses the past `N` years of
    data to optimize the portfolio weights. The optimizer also provides a
    backtest of the portfolio for the last `X` days.

    The optimizer also provides a PnL graph for the backtest period.

    > `Note`: This optimizer does not claim any future outlook or guarantee
    any returns. This is just for educational purposes and uses
    [PyPortfolioOpt](https://pyportfolioopt.readthedocs.io/) as an optimizer.

    ## How to use the optimizer?

    1. Enter the TickerCodes separated by a comma.
    2. Select the number of days for backtesting.
    3. Select the investment amount.
    4. Select the start year for analysis.
    5. Click on `Load Data` to load the data.
    6. Click on `Optimize` to optimize the portfolio weights.

    The optimizer uses yFinance to fetch the data.
    Kindly use the TickerCodes that are available on `Yahoo Finance`.

    '''
        )

    with gr.Row():
        with gr.Column():
            stocks = gr.Textbox(
                label="TickerCodes",
                show_label=True,
                interactive=True
            )
            backtest_days = gr.Slider(
                minimum=1,
                maximum=200,
                value=90,
                label="Backtest Days",
                interactive=True
            )
            investment_size = gr.Slider(
                minimum=10_000,
                maximum=10_00_000,
                value=100_000,
                interactive=True,
                step=10_000,
                label="Investment Amount"
            )
            start_year = gr.Slider(
                minimum=2020,
                maximum=2024,
                interactive=True,
                value=2021,
                label="Analysis Start Year"
            )
            with gr.Row():
                fetch_data = gr.Button("Load Data")
                # optimize = gr.Button("Optimize")
                optimize = gr.Button("Optimize")

    with gr.Row():
        with gr.Column():
            pnl_img = gr.BarPlot(
                title='PnL Graph',
                # interactive=True,
                # height=400,
            )
            performance = gr.Markdown()
            pnl_statement = gr.DataFrame(
                pd.DataFrame(columns=[
                    'instrument',
                    'weights',
                    'cost_price',
                    'sell_price',
                    'amount_allocation',
                    'unit_pnl',
                    'units',
                    'value',
                    'pnl'
                ]))
            historical_data = gr.DataFrame(
                pd.DataFrame(columns=['Date', 'Close'])
            )

    fetch_data.click(
        load_data,
        [stocks, start_year],
        [historical_data]
    )
    # optimize.click(
    #     find_weights,
    #     [historical_data, backtest_days, investment_size],
    #     [performance, pnl_statement]
    # )
    optimize.click(
        re_optimize,
        [historical_data, backtest_days, investment_size],
        [performance, pnl_statement, pnl_img]
    )

    gr.Examples(
        examples=[
            [','.join(
                ['ADANIENT.BO',
                 'ADANIPORTS.NS',
                 'APOLLOHOSP.NS',
                 'ASIANPAINT.NS',
                 'AXISBANK.BO',
                 'BAJAJ-AUTO.NS',
                 'BAJAJFINSV.NS',
                 'BAJFINANCE.NS',
                 'BHARTIARTL.BO',
                 'BPCL.NS',
                 'BRITANNIA.NS',
                 'CIPLA.NS',
                 'COALINDIA.NS',
                 'DIVISLAB.NS',
                 'DRREDDY.BO',
                 'EICHERMOT.BO',
                 'GRASIM.NS',
                 'HCLTECH.NS',
                 'HDFCBANK.NS',
                 'HDFCLIFE.BO',
                 'HEROMOTOCO.BO',
                 'HINDALCO.BO',
                 'HINDUNILVR.BO',
                 'ICICIBANK.NS',
                 'INDUSINDBK.NS',
                 'INFY.NS',
                 'ITC.NS',
                 'JSWSTEEL.BO',
                 'KOTAKBANK.BO',
                 'LT.BO',
                 'LTIM.BO',
                 'M&M.NS',
                 'MARUTI.NS',
                 'NTPC.BO',
                 'ONGC.NS',
                 'POWERGRID.BO',
                 'RELIANCE.NS',
                 'SBILIFE.BO',
                 'SBIN.NS',
                 'SHRIRAMFIN.BO',
                 'SUNPHARMA.NS',
                 'TATACONSUM.BO',
                 'TATAMOTORS.BO',
                 'TATASTEEL.BO',
                 'TCS.NS',
                 'TECHM.BO',
                 'TITAN.BO',
                 'ULTRACEMCO.NS',
                 'WIPRO.BO'
                 ]
            ),
                90,
                100000,
                2021
            ]
        ],
        inputs=[stocks, backtest_days, investment_size, start_year]
    )

demo.launch(debug=True, )