import bt
import matplotlib.pyplot as plt

import pandas as pd

import quantcommon
import streamlit_quant.strategy.multi_factor as multi_factor
import streamlit_quant.strategy.magic_formula as magic_formula


qc = quantcommon.QuantCommon()
mf = multi_factor.get_multi_factor_top(qc, 20)
magic_formula = magic_formula.get_magic_formula_top(20)

codes = ','.join(magic_formula['종목코드'].array)
price = qc.get_price_list_by_code(codes)

# price = price.set_index(['날짜'])
# price.rename(columns={"날짜": "Date"})
price["Date"] = pd.to_datetime(price["날짜"])

pivot_df = price.pivot(index="Date", columns="종목코드", values="종가")

# print(pivot_df.tail)

strategy = bt.Strategy("Asset_EW", [
    bt.algos.SelectAll(),    # 모든 데이터 사용
    bt.algos.WeighEqually(), # 동일 비중 투자
    bt.algos.RunMonthly(),   # 매 월말 리밸런싱
    bt.algos.Rebalance()     # 계산된 비중에 따라 리밸런싱
])

# 가격 데이터 중 시작 시점이 모두 다르므로, dropna() 함수를 통해 NA를 모두 제거하여 시작 시점을 맞춤
pivot_df.dropna(inplace=True)

# 백테스트 생성
backtest = bt.Backtest(strategy, pivot_df)

# 백테스트 실행
result = bt.run(backtest)
# prices: 누적 수익률이 데이터프레임 형태로 나타나며, 시작 시점을 100으로 환산하여 계산
# to_returns: 수익률 계산
# print(result.prices.to_returns())

result.plot(figsize=(10, 6), legend=False)
plt.show()

result.display()