import marimo

__generated_with = "0.6.22"
app = marimo.App()


@app.cell
def __():
    import numpy as np
    import cvxpy as cp
    return cp, np


@app.cell
def __(mo, np):
    # Define constants, B, n, p, mu, Sigma, S

    rng = np.random.default_rng(0xEE364A)

    B = 100
    n = 100
    p = rng.uniform(0, 2, n) ** 2 + 0.5
    mu = rng.normal(0.01, 1, n)
    Sigma = (Z := rng.normal(0, 1 / np.sqrt(n), (n, n))).T @ Z + 0.1 * np.eye(n)
    S = 15

    mo.md("This cell defines constants `B`, `n`, `p`, `mu`, `Sigma`, `S`")
    return B, S, Sigma, Z, mu, n, p, rng


@app.cell
def __(mu):
    mu.shape
    return


@app.cell
def __(mo):
    mo.md(
        """
        Define your optimization variables.

        *Hint: it should be cp.Variable*
        """
    )
    return


@app.cell
def __(cp, n):
    x = cp.Variable(n)
    return x,


@app.cell(hide_code=True)
def __(mo):
    mo.md(
        """
        Define the objective. 

        *Hint: it should be either cp.Maximize or cp.Minimize*
        """
    )
    return


@app.cell
def __(cp, mu, x):
    obj = cp.Maximize(mu.T @ x)
    return obj,


@app.cell
def __(mo):
    mo.md("Create a standard deviation constraint")
    return


@app.cell
def __(S, Sigma, cp, x):
    stddev_constraint = cp.quad_form(x, Sigma) <= S ** 2
    return stddev_constraint,


@app.cell(hide_code=True)
def __(mo):
    mo.md("Add your budget constraint")
    return


@app.cell
def __(B, p, x):
    budget_constraint = p @ x == B
    return budget_constraint,


@app.cell
def __(mo):
    mo.md(rf"Add your long only constraint")
    return


@app.cell
def __(x):
    long_only_constraint = x >= 0
    return long_only_constraint,


@app.cell(hide_code=True)
def __(mo):
    mo.md("Add your asset diversity constraint")
    return


@app.cell
def __(B, cp, p, x):
    asset_diversity_constraint = cp.multiply(x, p) <= 0.05 * B
    return asset_diversity_constraint,


@app.cell(hide_code=True)
def __(mo):
    mo.md("Make and run your CVXPY problem!")
    return


@app.cell
def __(
    asset_diversity_constraint,
    budget_constraint,
    cp,
    long_only_constraint,
    obj,
    stddev_constraint,
):
    constr = [
        stddev_constraint,
        budget_constraint,
        long_only_constraint,
        asset_diversity_constraint,
    ]

    prob = cp.Problem(obj, constr)
    prob.solve()
    return constr, prob


@app.cell
def __():
    import marimo as mo
    return mo,


if __name__ == "__main__":
    app.run()