import warnings
import numpy as np
from matplotlib import pyplot as plt
import pandas as pd
import seaborn as sns
import yfinance as yf
from sklearn import linear_model
from sklearn.naive_bayes import GaussianNB
from sklearn.svm import SVC
from sklearn.ensemble import RandomForestClassifier
from sklearn.neural_network import MLPClassifier
import datetime
import backtrader as bt
from backtrader.feeds import PandasData
import backtrader.analyzers as btanalyzers

warnings.filterwarnings("ignore")
plt.style.use('seaborn-colorblind')

ticker = 'TSLA'
start = datetime.datetime(2000, 1, 1)
end = datetime.datetime(2030, 12, 28)
stock = yf.download(ticker, progress=True, actions=True, start=start, end=end)

stock = stock['Adj Close']
stock = pd.DataFrame(stock)

stock.head()

stock.rename(columns = {"Adj Close": ticker}, inplace=True)

stock['returns'] = np.log(stock/stock.shift(1))
stock.dropna(inplace=True)

stock['direction'] = np.sign(stock['returns']).astype(int)

stock.head(10)

fig, ax = plt.subplots(2, 1, sharex=True, figsize=(12, 6))
ax[0].plot(stock[ticker], label = f'{ticker} Adj Close')
ax[0].set(title=f'{ticker} Closing Price', ylabel='Price')
ax[0].grid(True)
ax[0].legend()
ax[1].plot(stock['returns'], label = 'Daily Returns')
ax[1].set(title=f'{ticker} Daily Returns', ylabel='Returns')
ax[1].grid(True)
ax[1].legend()

lags= [1,2,3,4,5]
cols = []
for lag in lags:
    col = f'rtn_lag{lag}'
    stock[col] = stock['returns'].shift(lag)
    cols.append(col)
stock.dropna(inplace=True)
stock.head(2)

def create_bins(data, bins=[0]):
    global cols_bin
    cols_bin = []
    for col in cols:
        col_bin = col+'_bin'
        data[col_bin] = np.digitize(data[col], bins=bins)
        cols_bin.append(col_bin)

create_bins(stock)

models = {
    'log_reg': linear_model.LogisticRegression(),
    'gauss_nb': GaussianNB(),
    'svm': SVC(),
    'random_forest:': RandomForestClassifier(max_depth=10, n_estimators=100),
    'MLP': MLPClassifier(max_iter=100)
}

def fit_models(data):
    nfit = {model: models[model].fit(data[cols_bin], data['direction']) for model in models.keys()}

def derive_positions(data):
    for model in models.keys():
        data['pos_' + model] = models[model].predict(data[cols_bin])

def evaluate_strats(data):
    global strategy_rtn
    strategy_rtn = []
    for model in models.keys():
        col = 'strategy_' + model
        data[col] = data['pos_' + model] * data['returns']
        strategy_rtn.append(col)
    strategy_rtn.insert(0, 'returns')

fit_models(stock)
derive_positions((stock))
evaluate_strats(stock)
stock[strategy_rtn].sum().apply(np.exp)
ax = stock[strategy_rtn].cumsum().apply(np.exp).plot(figsize=(12, 6), title='ML return Comp')
ax.set_ylabel('Cum Returns')
ax.grid(True)

price = yf.download(ticker, progress=True, action=True, start=start, end=end)
price.head()

predictons = stock['strategy_svm']

predictions = pd.DataFrame(predictons)

predictions.rename(columns = {'strategy_rtn': 'predicted'}, inplace=True)

price = predictions.join(price, how='right').dropna()

price.head()

predictions.head()

OHLCH=['open', 'high', 'low', 'close', 'volume']

class SignalData(PandasData):
    cols = OHLCH +['predicted']
    lines = tuple(cols)
    params = {c: -1 for c in cols}
    params.update({'datetime': None})
    params = tuple(params.items())

cols = OHLCH + ['predicted']
lines = tuple(cols)
params = {c: -1 for c in cols}
params.update({'datetime': None})
params = tuple(params.items())

print(params)

class MLStrategy(bt.Strategy):
    params = dict(

    )
    def __init__(self):
        self.data_predicted = self.datas[0].predicted
        self.data_open = self.datas[0].open
        self.data_close = self.datas[0].close

        self.order=None
        self.price=None
        self.comm=None

    def log(self, txt):
        """Logging function"""
        dt = self.datas[0].datetime.date(0).isoformat()
        print(f'{dt}, {txt}')

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            # order already submitted/accepted - no action required
            return
        # report executed order
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f'BUY EXECUTED --- Price: {order.executed.price:.2f}, Cost: {order.executed.value:.2f},Commision: {order.executed.comm:.2f}'
                )
                self.price = order.executed.price
                self.comm = order.executed.comm
            else:
                self.log(f'SELL EXECUTED --- Price: {order.executed.price:.2f}, Cost: {order.executed.value:.2f}, Commission: {order.executed.comm:.2f}'
                )
        # report failed order
        elif order.status in [order.Canceled, order.Margin,
                              order.Rejected]:
            self.log('Order Failed')
        # set no pending order
        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed():
            return
        self.log(f'OPERATION RESULT --- Gross: {trade.pnl:.2f}, Net: {trade.pnlcomm:.2f}')

    def next_open(self):
        if not self.position:
            if self.data_predicted>0:
                size = int((self.broker.getcash()/self.datas[0].open)*0.9)
                self.log(f"BUY CREATED - SIZE: {size}, CASH: {self.broker.getcash():.2f}, Open: {self.data_open[0]}, Close: {self.data_close[0]}")
                self.buy(size=size)

        else:
            if self.data_predicted < 0:
                self.log(f'SELL CREATED - SIZE: {self.position.size}')
                self.close()

data = SignalData(dataname=price)

cerebro = bt.Cerebro(cheat_on_open=True)
cerebro.addstrategy(MLStrategy)
cerebro.adddata(data, name=ticker)
cerebro.broker.setcash(100000.0)
cerebro.broker.setcommission(commission=0.00025)

cerebro.addanalyzer(btanalyzers.SharpeRatio, _name='sharpe')
cerebro.addanalyzer(btanalyzers.DrawDown, _name='drawdown')
cerebro.addanalyzer(btanalyzers.Returns, _name='returns')
print('Start Cash: %.2f' % cerebro.broker.getvalue())
back = cerebro.run()
print('Ending Cash: %.2f' % cerebro.broker.getvalue())

par_list = [[
    x.analyers.returns.get_analylis()['rtot'],
    x.analyers.returns.get_analylis()['rnorn100'],
    x.analyers.drawdown.get_analylis()['max']['drawndown'],
    x.analyers.sharpe.get_analylis()['sharperatio']
] for x in back]

par_df = pd.DataFrame(par_list, columns = ['Total Return', 'APR', 'Drawdown', 'Sharpe'])
print(par_df)

cerebro.plot(style='style')

stock.head()
print(stock)
plt.show()