import numpy as np
import pandas as pd

def util(rets,w,A,gamma,eta):
    u=(1-gamma)/gamma*(A*(w[0]*rets[0]+w[1]*rets[1]+w[2]*rets[2])/(1-gamma)+eta)**gamma
    # print(u)
    return u

w_rf=[]
w_s=[]
w_b=[]
gammas=[]
etas=[]
percs=(0,0)
A=100
s_rets = np.exp(np.random.normal(.06,.55,size=10000))
b_rets = np.exp(np.random.normal(0,.4,size=10000))
rf_rets=[1.03]*len(s_rets)
for r in range(100):
    # print(r)
    gamma=np.random.uniform(low=-1,high=.999)
    eta=np.random.uniform(low=-gamma*15-5,high=20)
    gammas.append(round(gamma,4))
    etas.append(round(eta,4))
    best_port = None
    eU=[]
    for i in range(21):
        if i>1:
            if best_port[0]<(i-1)/20:
                break

        for k in range((21-i)):
            w=(5*i/100,5*k/100,5*(20-i-k)/100)
            df=pd.DataFrame({"rf":rf_rets,"stock":s_rets,"bond":b_rets})
            df_np=df.to_numpy()
            df['util']=np.apply_along_axis(util,1,df_np,w,A,gamma,eta)

            eU.append((w[0],w[1],w[2],np.mean(df['util'])))
            if not best_port:
                best_port=eU[-1]

            if eU[-1][3]>best_port[3]:
                best_port=eU[-1]

    w_rf.append(best_port[0])
    w_s.append(best_port[1])
    w_b.append(best_port[2])

    print(r, round(gamma,4), round(eta,4), w_rf[-1], w_s[-1], w_b[-1])

dfP=pd.DataFrame({
    "gamma":gammas,
    "eta":etas,
    "rf":w_rf,
    "stock":w_s,
    "bond":w_b,
    }
)

print(dfP.to_string())