import gradio as gr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import gurobipy as gp
from gurobipy import GRB


# =====================================================
# DATA
# =====================================================

df = pd.read_csv(
    "dow_jones_close_prices_2014_to_present3.csv"
)

df.set_index("Date", inplace=True)

returns_df = df.pct_change().dropna()


# =====================================================
# OPTIMIZER
# =====================================================

def utility_frontier(
    df_returns,
    selected_assets=None,
    gammas=None,
    allow_short=False
):

    if selected_assets:
        df = df_returns[selected_assets].copy()
    else:
        df = df_returns.copy()

    n = df.shape[1]

    mu = df.mean().values
    Sigma = df.cov().values

    risks = []
    returns = []
    utilities = []
    weights_list = []

    for gamma in gammas:

        model = gp.Model()
        model.setParam("OutputFlag", 0)

        if allow_short:

            w = model.addVars(
                n,
                lb=-GRB.INFINITY,
                name="w"
            )

        else:

            w = model.addVars(
                n,
                lb=0,
                ub=1,
                name="w"
            )

        expected_return = gp.quicksum(
            mu[i] * w[i]
            for i in range(n)
        )

        variance = gp.quicksum(
            Sigma[i, j] * w[i] * w[j]
            for i in range(n)
            for j in range(n)
        )

        model.setObjective(
            expected_return - gamma * variance,
            GRB.MAXIMIZE
        )

        model.addConstr(
            gp.quicksum(w[i] for i in range(n)) == 1
        )

        model.optimize()

        weights = np.array(
            [w[i].X for i in range(n)]
        )

        port_return = weights @ mu

        port_var = (
            weights.T @ Sigma @ weights
        )

        port_risk = np.sqrt(port_var)

        utility = (
            port_return
            - gamma * port_var
        )

        risks.append(port_risk)
        returns.append(port_return)
        utilities.append(utility)
        weights_list.append(weights)

    frontier = pd.DataFrame({
        "Gamma": gammas,
        "Risk": risks,
        "Return": returns,
        "Utility": utilities
    })

    weights_df = pd.DataFrame(
        weights_list,
        columns=df.columns
    )

    return frontier, weights_df


# =====================================================
# CALLBACK
# =====================================================

def optimize(
    assets,
    gamma,
    allow_short
):

    if len(assets) < 2:
        raise gr.Error(
            "Select at least two assets."
        )

    frontier, weights_df = utility_frontier(
        returns_df,
        selected_assets=assets,
        gammas=np.linspace(
            0.01,
            gamma,
            100
        ),
        allow_short=allow_short
    )

    # ---------------------------------
    # Frontier plot
    # ---------------------------------

    fig1, ax1 = plt.subplots(
        figsize=(8, 5)
    )

    ax1.plot(
        frontier["Risk"],
        frontier["Return"],
        marker="o"
    )

    asset_returns = (
        returns_df[assets]
        .mean()
    )

    asset_risks = (
        returns_df[assets]
        .std()
    )

    ax1.scatter(
        asset_risks,
        asset_returns,
        s=100
    )

    for asset in assets:

        ax1.annotate(
            asset,
            (
                asset_risks[asset],
                asset_returns[asset]
            )
        )

    ax1.set_title(
        "Efficient Frontier"
    )

    ax1.set_xlabel(
        "Risk"
    )

    ax1.set_ylabel(
        "Return"
    )

    ax1.grid(True)

    # ---------------------------------
    # Best utility portfolio
    # ---------------------------------

    best_idx = frontier[
        "Utility"
    ].idxmax()

    best_weights = (
        weights_df.iloc[best_idx]
    )

    # ---------------------------------
    # Weight plot
    # ---------------------------------

    fig2, ax2 = plt.subplots(
        figsize=(8, 5)
    )

    best_weights.sort_values(
        ascending=False
    ).plot.bar(ax=ax2)

    ax2.set_title(
        "Portfolio Weights"
    )

    ax2.set_ylabel(
        "Weight"
    )

    ax2.grid(True)

    weights_table = (
        best_weights
        .sort_values(
            ascending=False
        )
        .reset_index()
    )

    weights_table.columns = [
        "Asset",
        "Weight"
    ]

    return (
        fig1
 #       fig2,
  #      weights_table
    )


# =====================================================
# UI
# =====================================================

with gr.Blocks() as demo:

    gr.Markdown(
        "# Portfolio Optimization"
    )

    with gr.Row():

        with gr.Column(scale=1):

            asset_selector = gr.CheckboxGroup(
                choices=returns_df.columns.tolist(),
                value=returns_df.columns[:5].tolist(),
                label="Assets"
            )

            gamma_slider = gr.Slider(
                minimum=0.1,
                maximum=100,
                value=100,
                step=0.1,
                label="Maximum Gamma"
            )

            # short_box = gr.Checkbox(
            #     label="Allow Short Selling"
            # )

            run_btn = gr.Button(
                "Optimize"
            )

        with gr.Column(scale=2):

            frontier_plot = gr.Plot(
                label="Efficient Frontier"
            )

            # weight_plot = gr.Plot(
            #     label="Weights"
            # )

            # weight_table = gr.Dataframe(
            #     label="Weights Table"
            # )

    run_btn.click(
        fn=optimize,
        inputs=[
            asset_selector,
            gamma_slider,
            # short_box
        ],
        outputs=[
            frontier_plot
            # weight_plot,
            # weight_table
        ]
    )

demo.launch()