#!/usr/bin/env python2.7

from datetime import datetime, timedelta

from epidata.EpidataLiteStreamingContext import EpidataLiteStreamingContext
import sys
import numpy
import pandas
import sqlite3
from sqlite3 import Error

import unittest

ec = EpidataLiteStreamingContext()

database = "/Users/lujiajun/Documents/epidata-intern/data/epidata_development.db"

def create_connection(db_file):
    conn = None
    try:
        conn = sqlite3.connect(db_file, detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES)
    except Error as e:
        print(e)

    return conn

con = create_connection(database)

class Base(unittest.TestCase):

    def assertEqualRows(self, one, two):
        self.maxDiff = None
        # if not isinstance(one, Column):
        self.assertEqual(one, two)

    def assertEqualDataFrames(self, one, two):
        self.assertEqual(one.count(), two.count())
        for i, j in zip(one, two):
            self.assertEqualRows(i, j)

class EpidataContextTests(Base):
    # def test_simple_query_test(self):
    #     # create a database connection
    #     cur = con.cursor()
    #     cur.execute("select * from measurements_original where key2 = 'Test-1'")
    #     rows = cur.fetchall()
    #
    #     for row in rows:
    #         print(row)

    def test_simpleTest(self):

        ec.printSomething('god')
        trans = ec.createTransformations("Identity", ["Meas-1"], {})
        print(trans)
        ec.createStream("measurements_original", "measurements_cleansed", trans)
        print("Start stream")
        ec.startStream()
        print("Enter 'Q' to stop streaming")
        # data = input("Enter a number to quit")
        # while data:
        #     data = input("Enter a number to quit")
        ec.stopStream()

if __name__ == "__main__":
    # NOTE Fixtures are added externally, by IPythonSpec.scala.
    unittest.main()


import yfinance as yf

class IndexArbitrageNative:
    def __init__(self, basket_tickers, basket_weights):
        self.basket_tickers = basket_tickers
        self.basket_weights = basket_weights

    def calculate_fair_value(self, spy_price, basket_prices):
        # Calculate the fair value of the index futures contract
        fair_value = sum(price * weight for price, weight in zip(basket_prices, self.basket_weights))
        return fair_value

    def check_arbitrage_opportunity(self):
        # Fetch the latest prices of the basket of stocks and SPY
        basket_data = yf.download(self.basket_tickers, period='1d')['Adj Close']
        spy_data = yf.download('SPY', period='1d')['Adj Close']

        # Extract the latest prices from the fetched data
        basket_prices = [basket_data[ticker][-1] for ticker in self.basket_tickers]
        spy_price = spy_data[-1]

        # Calculate the fair value of the index futures contract
        fair_value = self.calculate_fair_value(spy_price, basket_prices)

        # Check if there is an arbitrage opportunity
        if spy_price < fair_value:
            return {
                'opportunity': 'buy_spy',
                'spy_price': spy_price,
                'basket_prices': basket_prices
            }
        elif spy_price > fair_value:
            return {
                'opportunity': 'sell_spy',
                'spy_price': spy_price,
                'basket_prices': basket_prices
            }
        else:
            return {
                'opportunity': 'no_arbitrage',
                'spy_price': spy_price,
                'basket_prices': basket_prices
            }


from fastapi import FastAPI, Query
from datetime import datetime
from pytz import utc

app = FastAPI()

@app.get("/api")
async def get_info(
    slack_name: str = Query(..., title="Slack Name"),
    track: str = Query(..., title="Track"),
):
    # Get current day of the week
    current_day = datetime.now(utc).strftime('%A')

    # Get current UTC time
    current_utc_time = datetime.now(utc).strftime('%Y-%m-%dT%H:%M:%SZ')

    # GitHub file and repo URLs
    github_file_url = "https://github.com/preciousajorgba/HNGX/blob/main/task1.py"
    github_repo_url = "https://github.com/preciousajorgba/HNGX"

    # Response JSON
    response_data = {
        "slack_name": slack_name,
        "current_day": current_day,
        "utc_time": current_utc_time,
        "track": track,
        "github_file_url": github_file_url,
        "github_repo_url": github_repo_url,
        "status_code": 200
    }

    return response_data

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)


import datetime
import journal_scan
import os
import finance_tracker
import pirate_massacre
import threading
import time
import web_server
from pytz import UTC

pirate_scanner = None
finance_scanner = None
book_mark = None

def all_scanners():
    return [pirate_scanner, finance_scanner]

mutex = threading.Lock()

# There appears to be a bug in the bookmark system. Until that's fixed we'll rebuild
# the complete scanner every time a request is made
rebuild_scanner_on_request = False

def build_scanners():
    global pirate_scanner
    global finance_scanner
    global book_mark
    global mutex

    mutex.acquire()
    pirate_scanner = pirate_massacre.PirateMassacreScanner()
    finance_scanner = finance_tracker.FinanceTracker()
    mutex.release()

    # Wing missions last around 7 days, so we only need to scan that far back
    start_date = datetime.datetime.now(UTC) - datetime.timedelta(8)
    book_mark = journal_scan.scan_journal_files_in_date_range(all_scanners(), start_date, None)

    print("Scanners rebuilt.")

def strfdelta(tdelta, fmt):
    d = {"days": tdelta.days}
    d["hours"], rem = divmod(tdelta.seconds, 3600)
    d["minutes"], d["seconds"] = divmod(rem, 60)
    return fmt.format(**d)

def build_remaining_time_str(expiry):
    now = UTC.localize(datetime.datetime.utcnow())
    if(expiry < now):
        return 'EXPIRED!'
    else:
        message = "Expires in "
        remaining = expiry - now
        if remaining.days > 0:
            message += strfdelta(remaining, '{days}D {hours}H {minutes}MIN')
        else:
            message += strfdelta(remaining, '{hours}H {minutes}MIN')
        return message

rep_classes = {
    "HOSTILE": "FactionRepBad",
    "UNFRIENDLY": "FactionRepBad",
    "NEUTRAL": "FactionRepNeutral",
    "FRIENDLY": "FactionRepGood",
    "ALLIED": "FactionRepGood"
}

def get_reputation_class(rep_string):
    if rep_string not in rep_classes:
        return "FactionRepNeutral"
    return rep_classes[rep_string]

class PirateMassacreHandler(web_server.Handler):
    def get_handler(self, request):
        if request == "/mission_report":
            return self.get_mission_report
        elif request == "/finance_report":
            return self.get_finance_report
        elif request == "/rebuild":
            return self.get_rebuild
        return None

    def get_mission_report(self):
        if pirate_scanner == None:
            return

        if rebuild_scanner_on_request:
            build_scanners()

        global mutex
        mutex.acquire()
        report = pirate_scanner.build_report()
        mutex.release()

        self.send_response(200)
        self.send_header("Content-type", "text/html")
        self.end_headers()

        content = ""

        content += f'
Tracking {report["MissionCount"]}/20 missions for a potential {"{:,}".format(report["TotalReward"])} credit reward. ({"{:,}".format(report["TotalWingReward"])} credits in wing missions.)
'
        for system_entry in report['MissionsBySystem']:
            system_name = system_entry['Name']
            system = report['Systems'][system_name]
            content += (f'
{system.name}
')
            for faction_entry in system_entry["Factions"]:
                faction_name = faction_entry['Name']
                faction = report['Factions'][faction_name]

                rep_string = faction.get_reputation_string()
                rep_class = get_reputation_class(rep_string)

                icons = ""
                if faction.government == 'Anarchy':
                    icons += '🏴☠' # pirate flag
                if faction.state in ["War", "CivilWar"]:
                    icons += '⚔' # crossed swords

                content += f'
{faction.name}{icons}
{rep_string}

'
                missions = faction_entry["Missions"]
                if len(missions) > 0:
                    for mission in missions:
                        kills = mission.total_kills - mission.remaining_kills
                        content += '
'

                        message = 'Kill {} for {:,} credits'.format(mission.target_faction, mission.reward)
                        if mission.wing:
                            message += ' WING '
                        message += ' — '

                        if mission.remaining_kills > 0:
                            message += f' {kills}/{mission.total_kills} ({mission.remaining_kills} remain)'
                            message += ' — '
                            message += build_remaining_time_str(mission.expiry)
                        else:
                            message += ' DONE!'

                        content += message
                        content += '
'
                else:
                    content += '
'
                    content += 'None'
                    if faction_entry['HasMissionsInOtherSystem']:
                        content +=' (have mission in other system)'
                    content += '
'

        self.wfile.write(bytes(content, "utf-8"))

    def get_finance_report(self):
        if finance_scanner == None:
            return

        global mutex
        mutex.acquire()
        report = finance_scanner.build_report_json()
        mutex.release()

        self.send_response(200)
        self.send_header("Content-type", "text/json")
        self.end_headers()

        content = report
        self.wfile.write(bytes(content, "utf-8"))

    def get_rebuild(self):
        build_scanners()

        self.send_response(200)
        self.send_header("Content-type", "text/plain")
        self.end_headers()

        content = "Scanners Rebuilt"
        self.wfile.write(bytes(content, "utf-8"))

def main_loop():
    if rebuild_scanner_on_request:
        return

    global mutex
    global pirate_scanner
    global book_mark

    while True:
        time.sleep(0.5)

        mutex.acquire()
        book_mark = journal_scan.resume_from_book_mark(
            all_scanners(),
            book_mark)
        mutex.release()

if __name__ == "__main__":
    args = web_server.parse_args("Run a web server showing a live stream of your pirate massacre missions grouped by system.")

    build_scanners()

    public_path = os.path.join(
        os.path.dirname(
            os.path.realpath(__file__)
        ), "Public")

    web_server.start(args, PirateMassacreHandler, public_path)

    main_loop()


from h2o_wave import ui, data, Q
from .common import global_nav
from .synthetic_data import *

async def show_grey_dashboard(q: Q):
    q.page['meta'] = ui.meta_card(box='', layouts=[
        ui.layout(
            breakpoint='xs',
            min_width='800px',
            zones=[
                ui.zone('header', size='76px'),
                ui.zone('title'),
                ui.zone('body', size='1000px', zones=[
                    ui.zone('top', direction=ui.ZoneDirection.ROW, size='25%'),
                    ui.zone('middle', direction=ui.ZoneDirection.ROW, size='25%'),
                    ui.zone('middle2', direction=ui.ZoneDirection.ROW, size='25%'),
                    ui.zone('bottom', direction=ui.ZoneDirection.ROW, size='25%'),
                ]),
                ui.zone('footer', size='80px'),
            ]
        )
    ])

    q.page['header'] = ui.header_card(box='header', title='H2O Wave Demo', subtitle='Grey Dashboard',
                                      image='https://wave.h2o.ai/img/h2o-logo.svg',
                                      items=[ui.tabs(name='Dashboards', value='#dashboards/grey',
                                                     items=global_nav),])
    q.page['section'] = ui.section_card(
        box='title',
        title=next(sample_title),
        subtitle=next(sample_caption),
        items=[
            ui.label(label='Start:'),
            ui.date_picker(name='target_date', label='', value='2020-12-20'),
            ui.label(label='End:'),
            ui.date_picker(name='target_date', label='', value='2020-12-25'),
        ],
    )

    stock_dates = generate_time_series(10000)
    stock_prices = generate_random_walk()

    q.page['small'] = ui.small_stat_card(
        box=ui.box('top', order=1),
        title=next(sample_term),
        value=next(sample_dollars),
    )
    q.page['small_series'] = ui.small_series_stat_card(
        box=ui.box('top', order=2),
        title=next(sample_term),
        value=next(sample_dollars),
        plot_category='date',
        plot_value='price',
        plot_data=data(
            fields=['date', 'price'],
            rows=[(next(stock_dates), next(stock_prices)) for i in range(30)],
            pack=True,
        ),
    )
    q.page['small_series_interval'] = ui.small_series_stat_card(
        box=ui.box('top', order=3),
        title=next(sample_term),
        value=next(sample_dollars),
        plot_category='date',
        plot_value='price',
        plot_type=ui.SmallSeriesStatCardPlotType.INTERVAL,
        plot_data=data(
            fields=['date', 'price'],
            rows=[(next(stock_dates), next(stock_prices)) for i in range(30)],
            pack=True,
        ),
    )
    q.page['wide_series'] = ui.wide_series_stat_card(
        box=ui.box('middle', order=1),
        title=next(sample_term),
        value=next(sample_dollars),
        aux_value=next(sample_percent),
        plot_category='date',
        plot_value='price',
        plot_data=data(
            fields=['date', 'price'],
            rows=[(next(stock_dates), next(stock_prices)) for i in range(30)],
            pack=True,
        ),
    )
    q.page['wide_bar'] = ui.wide_bar_stat_card(
        box=ui.box('middle', order=2),
        title=next(sample_term),
        value=next(sample_dollars),
        aux_value=next(sample_percent),
        progress=random.random(),
    )
    q.page['wide_gauge'] = ui.wide_gauge_stat_card(
        box=ui.box('middle', order=3),
        title=next(sample_term),
        value=next(sample_dollars),
        aux_value=next(sample_percent),
        progress=random.random(),
    )
    q.page['tall_series'] = ui.tall_series_stat_card(
        box=ui.box('middle2', order=1),
        title=next(sample_term),
        value=next(sample_dollars),
        aux_value=next(sample_percent),
        plot_category='date',
        plot_value='price',
        plot_data=data(
            fields=['date', 'price'],
            rows=[(next(stock_dates), next(stock_prices)) for i in range(30)],
            pack=True,
        ),
    )
    q.page['tall_gauge'] = ui.tall_gauge_stat_card(
        box=ui.box('middle2', order=2),
        title=next(sample_term),
        value=next(sample_dollars),
        aux_value=next(sample_percent),
        progress=random.random(),
    )
    q.page['large'] = ui.large_stat_card(
        box=ui.box('bottom', order=1),
        title=next(sample_term),
        value=next(sample_dollars),
        aux_value=next(sample_percent),
        caption=next(sample_caption),
    )
    q.page['large_bar'] = ui.large_bar_stat_card(
        box=ui.box('bottom', order=2),
        title=next(sample_term),
        value=next(sample_dollars),
        value_caption=next(sample_term),
        aux_value=next(sample_dollars),
        aux_value_caption=next(sample_term),
        progress=random.random(),
        caption=next(sample_caption),
    )

    q.page['footer'] = ui.footer_card(box='footer', caption='(c) 2021 H2O.ai. All rights reserved.')

    await q.page.save()


from collections import deque
def solution(time, r):

    answer = []
    wait = deque()
    pyo = deque()

    t = 0
    while t <= max(time):
        for idx,i in enumerate(time):
            if i == t:
                wait.append(idx)
                pyo.append(r[idx])
        if len(wait) == 1:
            answer.append(wait.popleft())

        for idx2,w in enumerate(wait):
            if min(pyo) == r[w]:
                wait.pop(idx2)
        t+=1
    return answer

solution([0,3,1,0],[0,1,2,3])
# [0,1,3,0]	[0,1,2,3]	[0, 1, 3, 2]
# [7,6,8,1]	[0,1,2,3]	[3, 1, 0, 2]


#!/usr/bin/python

import sys
import os

file = sys.argv[1]
output = sys.argv[2]

import libtorrent as lt

fs = lt.file_storage()
lt.add_files( fs, file )

ct = lt.create_torrent( fs, 5000 )

#ct.add_tracker("http://tracker.openbittorrent.com/announce", 0)
#ct.add_tracker("udp://tracker.openbittorrent.com:80/announce", 0)
ct.add_tracker("http://tracker.publicbt.com/announce", 0)
ct.add_tracker("udp://tracker.publicbt.com:80/announce", 0)

ct.set_creator("iftd")

lt.set_piece_hashes( ct, os.path.dirname(file))

bt_str = lt.bencode( ct.generate() )

fd = open(output, "wb" )
fd.write( bt_str )
fd.close()


# -*- coding: utf-8 -*-
"""
Created on Sun May 17 18:39:27 2020

Implementing a simple VAE.
Following the tutorial at: https://graviraja.github.io/vanillavae/#

@author: Jack
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

import umap
import seaborn as sns

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

import sklearn
from sklearn.preprocessing import MinMaxScaler
from sklearn.cluster import KMeans
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import confusion_matrix

import itertools

#%%

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

transforms = transforms.Compose([transforms.ToTensor()])
train_ds = datasets.MNIST('./data', train=True, download=True,
                          transform=transforms)

test_ds = datasets.MNIST('./data', train=False, download=True,
        transform=transforms
    )

''' TODO: remove - the DataLoader does scaling
scaler = MinMaxScaler()
scaler.fit(train_ds.data.view(-1,28*28))
scaled = scaler.transform(train_ds.data.view(-1,28*28))
train_ds_sc = torch.Tensor(scaled).view(-1,28,28)

test_ds_sc = torch.Tensor(scaler.transform(test_ds.data.view(-1,28*28)))
test_ds_sc = test_ds_sc.view(-1,28,28)
'''

class Encoder(nn.Module):

    def __init__(self, input_dim, hidden_dim1, hidden_dim2, hidden_dim3, latent_dim):

        super().__init__()
        self.linear1 = nn.Linear(input_dim, hidden_dim1)
        self.linear2 = nn.Linear(hidden_dim1, hidden_dim2)
        self.linear3 = nn.Linear(hidden_dim2, hidden_dim3)
        self.mu = nn.Linear(hidden_dim3, latent_dim)
        self.var = nn.Linear(hidden_dim3, latent_dim)

    def forward(self, data):

        hidden1 = F.relu(self.linear1(data))  #Not sure about using ReLu
        hidden2 = F.relu(self.linear2(hidden1))
        hidden3 = F.relu(self.linear3(hidden2))
        self.z_mu = self.mu(hidden3)
        self.z_var = self.var(hidden3)

        return self.z_mu, self.z_var

class Decoder(nn.Module):

    def __init__(self, latent_dim, hidden_dim3, hidden_dim2, hidden_dim1, out_dim):

        super().__init__()
        self.linear1 = nn.Linear(latent_dim, hidden_dim3)
        self.linear2 = nn.Linear(hidden_dim3, hidden_dim2)
        self.linear3 = nn.Linear(hidden_dim2, hidden_dim1)
        self.out = nn.Linear(hidden_dim1, out_dim)

    def forward(self, data):

        hidden1 = F.relu(self.linear1(data))
        hidden2 = F.relu(self.linear2(hidden1))
        hidden3 = F.relu(self.linear3(hidden2))
        pred = F.relu(self.out(hidden3))  #Not sure about sigmoid output

        return pred

class VAE(nn.Module):

    def __init__(self, encoder, decoder):

        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, data):

        self.z_mu, self.z_var = self.encoder(data)

        #Reparameterize:
        std = torch.exp(self.z_var/2)

        e = torch.randn_like(std)   #Samples a N(0,1) in the shape of std
        latent_sample = e.mul(std).add_(self.z_mu)   #Does elementwise multiplication and addition

        prediction = self.decoder(latent_sample)

        return prediction, self.z_mu, self.z_var

    def sample(self, num=1):
        z = torch.randn(num, z_dim)
        output = self.decoder(z)
        return output

in_dim = 28*28
hidden_dim1 = 256
hidden_dim2 = 128
z_dim = 12
lr = 0.001
batches = 50

train_it = DataLoader(train_ds, batch_size=batches, shuffle=True, drop_last=True)
test_it = DataLoader(test_ds, batch_size=batches, drop_last=True)

#KL divergence between N(mu,var) and N(0,1)
def KL_div(mu, var):
    mu_sq = pow(mu,2)
    sig_sq = torch.exp(var) #doesnt need /2
    loss = 0.5 * torch.sum(mu_sq - var - 1 + sig_sq )
    return loss

def get_loss(input_x, output_x, mu, logvar):
    recon_loss = F.mse_loss(output_x, input_x, reduction='sum')
    kl_loss = KL_div(mu, logvar)
    return recon_loss + kl_loss

def train(mod, e):

    mod.train()
    train_loss = 0
    loss_list = []
    average_loss = 999
    #while average_loss > 40:
    percent = 0
    for i, (data, target) in enumerate(train_it):
        data = data.view(-1,28*28)
        optimizer.zero_grad()

        output, mu, logvar = mod(data)
        loss = get_loss(data, output, mu, logvar)

        loss.backward()
        optimizer.step()
        train_loss += loss.item()
        loss_list.append(loss.item())
        percent = i*batches*100/len(train_it.dataset)
        #print(loss.item())
        if percent % 20 == 0:
            print(f"{percent}% , {loss.item()}")
        #if i*batches % len(train_it.dataset) == 0:
            #print(f"{100*i/len(train_it.dataset)}% - Loss = {loss}")
    average_loss = train_loss/len(train_it.dataset)
    print(f"====> Epoch {e} - Average train loss = {average_loss}")

    return train_loss, loss_list

def test(mod):

    mod.eval()

    test_loss = 0
    mu_list = torch.Tensor()
    logvar_list = torch.Tensor()
    output_list = torch.Tensor()

    with torch.no_grad():
        for i, (data, target) in enumerate(test_it):

            data = data.view(-1, 28*28)
            data = data.to(device)

            sample, z_mu, z_var = mod(data)
            mu_list = torch.cat([mu_list, z_mu], dim=0)
            logvar_list = torch.cat([logvar_list, z_var], dim=0)
            output_list = torch.cat([output_list, sample], dim=0)
            loss = get_loss(data, sample, z_mu, z_var)
            test_loss += loss.item()

    return test_loss, output_list, mu_list, logvar_list

num_tests = 5
epochs = 2

test_output_scores = []
test_latent_scores = []
loss_lists = []
fit_sizes = [15,20,30,40,50,60,70,80,90,100,150,200,250,300,350,400,450,500,600,700,
             800,900,1000,2000,4000,6000,8000,8250,8500,8600,8700,8800,8900,9000]

for t in range(num_tests):

    best_loss = 1e10
    wins = 0

    all_losses = []

    dim_list = [28*28, 256, 128, 64, 12]
    enc = Encoder(*dim_list)
    dim_list.reverse()
    dec = Decoder(*dim_list)
    model = VAE(enc, dec)
    optimizer = optim.Adam(model.parameters(), lr=lr)

    for e in range(epochs):

        train_loss, loss_list = train(model, e)
        all_losses.append(loss_list)
        test_loss, outputs, z_mus, z_vars = test(model)

        train_loss /= len(train_ds)
        test_loss /= len(test_ds)
        print(f"Epoch {e} - Train loss = {round(train_loss,2)}, Test loss = {round(test_loss,2)}")

        if test_loss < best_loss:
            best_loss = test_loss
        else:
            wins += 1

        if wins > 3:
            print(f"{e} epochs completed.")
            break

    loss_lists.append(all_losses)

    num = 15
    #size = 9000
    test_size = 1000
    test_loss, outputs, z_mus, z_vars = test(model)

    score_list = []
    latent_score_list = []

    for size in fit_sizes:

        knn_output = KNeighborsClassifier(n_neighbors=num)
        knn_output.fit(outputs[0:size].detach().view(-1,28*28), test_ds.targets[0:size])
        score = knn_output.score(outputs[size:size+test_size].detach().view(-1,28*28),
                                 test_ds.targets[size:size+test_size])

        knn_latent = KNeighborsClassifier(n_neighbors=num)
        knn_latent.fit(z_mus[0:size].detach(), test_ds.targets[0:size])
        latent_score = knn_latent.score(z_mus[size:size+test_size].detach(),
                                        test_ds.targets[size:size+test_size])

        score_list.append(score)
        latent_score_list.append(latent_score)

    test_output_scores.append(score_list)
    test_latent_scores.append(latent_score_list)
    print(f"Test {t} done.")

test_output_scores = pd.DataFrame(test_output_scores)
test_latent_scores = pd.DataFrame(test_latent_scores)
loss_lists = pd.DataFrame(loss_lists)

mean_outputs = test_output_scores.mean(axis=0)
mean_latent = test_latent_scores.mean(axis=0)
std_output = test_output_scores.std(axis=0)
std_latent = test_latent_scores.std(axis=0)

x = loss_lists.loc[0][0]
x += loss_lists.loc[0][1]

y = loss_lists.loc[1][0]
y += loss_lists.loc[1][1]

plt.scatter([i for i in range(len(x))], x)
plt.scatter([i for i in range(len(y))], y)

fig = plt.figure()
ax1 = fig.add_subplot(111)
ax1.errorbar(fit_sizes, mean_outputs, yerr=std_output,
             c='b', label='Output space')
ax1.errorbar(fit_sizes, mean_latent, yerr=std_latent,
             c='r', label='Latent space')
plt.xlabel("Number of supervised examples")
plt.ylabel("Mean accuracy")
plt.title("Effect of supervision on Architecture 1")
plt.legend(loc='middle right');
#plt.savefig("saved_plots/VAE_MNIST_arch1_AccuracyVSnum_08-06.png")

test_loss, outputs, z_mus, z_vars = test(model)

knn_latent = KNeighborsClassifier(n_neighbors=num)
knn_latent.fit(z_mus[0:size].detach(), test_ds.targets[0:size])
predictions = knn_latent.predict(z_mus[size:size+test_size].detach())

conf_mat = confusion_matrix(test_ds.targets[size:size+test_size], predictions)
plt.imshow(conf_mat, cmap=plt.cm.Blues)

def plot_confusion_matrix(cm, classes, normalize=False, title='Confusion matrix',
                          cmap=plt.cm.Blues, save_name=None):
    """
    This function prints and plots the confusion matrix.
    Normalization can be applied by setting `normalize=True`.
    """
    plt.figure(figsize=(9,7))
    plt.imshow(cm, interpolation='nearest', cmap=cmap)
    plt.title(title)
    if not normalize:
        plt.colorbar()
    tick_marks = np.arange(len(classes))
    plt.xticks(tick_marks, classes)
    plt.yticks(tick_marks, classes)

    if normalize:
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
        cm = cm.round(2)

    thresh = cm.max() / 2.
    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
        plt.text(j, i, cm[i, j],
                 horizontalalignment="center",
                 color="white" if cm[i, j] > thresh else "black")
    #plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label')
    if save_name != None:
        plt.savefig(f"{save_name}")
        print("Saving...")

plot_confusion_matrix(conf_mat, classes=range(10), normalize=True,
                      save_name="saved_plots/VAE_MNIST_confusion-matrix-normalised_9kSupervisedKNN_08-06.png")

'''
sample = model.sample(1)
img = sample.view(28,28).data
plt.imshow(img, cmap='gray')

losses = []
for i in range(len(all_losses)):
    for j in range(len(all_losses[i])):
        losses.append(all_losses[i][j])

plt.scatter([i for i in range(len(losses))], losses)

losses, outputs, mus, logvars = test(model)
plt.imshow(outputs[0].view(28,28), cmap='gray')
plt.imshow(test_ds.data[0].view(28,28), cmap='gray')
'''

#%%

#Attempts at UMAP, with trained model:

'''
#Plot the separations in the data
fit = umap.UMAP(random_state=42)
u = fit.fit_transform(train_ds.data.view(-1,28*28))
fig,ax = plt.subplots(figsize=(10,8))
cmap = plt.get_cmap('gist_rainbow', 10)
cax = ax.scatter(u[:,0], u[:,1], s=5, c=train_ds.targets, cmap=cmap,
                 vmin=0, vmax=train_ds.targets.max())
fig.colorbar(cax)#, extend='min')
#plt.savefig(fname="saved_plots/MNIST_UMAP.png", quality=100)
'''

'''
#test_data = model.z_mu.detach().numpy() Needed since z_mu requires grad
fit2 = umap.UMAP(n_neighbors=50)
emb = fit2.fit_transform(z_mu_tensor)

fig,ax = plt.subplots(figsize=(10,8))
cmap = plt.get_cmap('Spectral', 10)
cax = ax.scatter(emb[:,0], emb[:,1], s=1, c=targets, cmap=cmap,
                 vmin=0, vmax=train_ds.targets.max())
fig.colorbar(cax)
plt.title('Latent mean-vector space of the MNIST dataset', fontsize=14)
plt.savefig(fname=f'saved_plots/latent_dim/MNIST_scaled_{z_dim}.png', quality=100)
'''


import csv
import ijson  # type: ignore
from pylegend._typing import (
    PyLegendSequence,
    PyLegendOptional,
    TYPE_CHECKING
)
if TYPE_CHECKING:
    from pylegend.core.tds.tds_frame import PyLegendTdsFrame  # pragma: no cover
from pylegend.core.tds.result_handler.result_handler import ResultHandler
from pylegend.core.request.response_reader import ResponseReader

__all__: PyLegendSequence[str] = [
    "ToCsvFileResultHandler"
]

class ToCsvFileResultHandler(ResultHandler[None]):

    def __init__(self, file_or_writer, parse_float_as_decimal: bool = False) -> None:  # type: ignore
        self.__file: PyLegendOptional[str] = file_or_writer if isinstance(file_or_writer, str) else None
        self.__csv_writer = None if isinstance(file_or_writer, str) else file_or_writer
        self.__parse_float_as_decimal = parse_float_as_decimal

    def handle_result(self, frame: "PyLegendTdsFrame", result: ResponseReader) -> None:
        if self.__file is None:
            self.__write_result(self.__csv_writer, frame, result, self.__parse_float_as_decimal)
        else:
            with open(self.__file, 'w', newline='') as csvfile:
                csv_writer = csv.writer(csvfile)
                self.__write_result(csv_writer, frame, result, self.__parse_float_as_decimal)

    @staticmethod
    def __write_result(  # type: ignore
            csv_writer,
            frame,
            result,
            parse_as_decimal
    ) -> None:
        csv_writer.writerow([col.get_name() for col in frame.columns()])
        for next_row in ijson.items(result, "result.rows.item.values", use_float=not parse_as_decimal):
            csv_writer.writerow(next_row)
        return


import dataclasses
from typing import Any, Optional

from flask import jsonify

from ...exceptions.exceptions import ResourceExistsException
from ..models.video_model import Video
from ..repositories.unit_of_work import BaseUnitfWork
from .querie_mixin import QueryMixin
from .use_case import UseCase

class AddVideoUseCase(UseCase):
    def __init__(self, unit_of_work: Optional[BaseUnitfWork] = None) -> None:
        super().__init__(unit_of_work)

    def execute(self, data: dict[str, Any]) -> dict[str, Any]:
        with self.unit_of_work as uow:
            video = Video(
                video_id=data["video_id"],
                video_title=data["video_title"],
                channel_title=data["channel_title"],
                video_description=data["video_description"],
                video_thumbnail=data["video_thumbnail"],
                video_duration=data["video_duration"],
                views_count=data["views_count"],
                likes_count=data["likes_count"],
                comments_count=data["comments_count"],
                published_at=data["published_at"],
                channel_id=data["channel_id"],
            )
            uow.repository.add(video)
        return dataclasses.asdict(video)

class GetVideoUseCase(UseCase):
    def __init__(self, unit_of_work: Optional[BaseUnitfWork] = None) -> None:
        super().__init__(unit_of_work)

    def execute(self, data: dict[str, Any]) -> dict[str, Any]:
        with self.unit_of_work as uow:
            video_id = data["video_id"]
            video = uow.repository.get_by_id(video_id)
        return dataclasses.asdict(video)

class DeleteVideoUseCase(UseCase):
    def __init__(self, unit_of_work: Optional[BaseUnitfWork] = None) -> None:
        super().__init__(unit_of_work)

    def execute(self, data: dict[str, Any]) -> dict[str, Any]:
        with self.unit_of_work as uow:
            video_id = data["video_id"]
            video = uow.repository.get_by_id(video_id)
            uow.repository.delete(video_id)
        return dataclasses.asdict(video)

class UpdateVideoUseCase(UseCase):
    def __init__(self, unit_of_work: Optional[BaseUnitfWork] = None) -> None:
        super().__init__(unit_of_work)

    def execute(self, data: dict[str, Any]) -> dict[str, Any]:
        with self.unit_of_work as uow:
            id = int(data["id"])
            video = uow.repository.get_by_id(id)
            if data.get("channel_title"):
                video.channel_title = data.get("channel_title")
            if data.get("comments_count"):
                video.comments_count = data.get("comments_count")
            if data.get("likes_count"):
                video.likes_count = data.get("likes_count")
            if data.get("video_description"):
                video.video_description = data.get("video_description")
            if data.get("video_duration"):
                video.video_duration = data.get("video_duration")
            if data.get("video_id"):
                video.video_id = data.get("video_id")
            if data.get("video_thumbnail"):
                video.video_thumbnail = data.get("video_thumbnail")
            if data.get("video_title"):
                video.video_title = data.get("video_title")
            if data.get("views_count"):
                video.views_count = data.get("views_count")
            if data.get("published_at"):
                video.published_at = data.get("published_at")
            uow.repository.update(video)
        return dataclasses.asdict(video)

class GetVideosUseCase(UseCase):
    def __init__(self, unit_of_work: Optional[BaseUnitfWork] = None) -> None:
        super().__init__(unit_of_work)

    def execute(self, data: dict[str, Any]) -> dict[str, Any]:
        with self.unit_of_work as uow:
            sort_field = "id"
            limit = 10
            offset = 0
            sort_order = "ASC"
            if data.get("sort_field"):
                sort_field = data.get("sort_field")
            if data.get("limit"):
                limit = data.get("limit")
            if data.get("offset"):
                offset = data.get("offset")
            if data.get("sort_order"):
                sort_order = data.get("sort_order")
            videos = uow.repository.list_all(sort_field=sort_field, limit=limit, offset=offset, sort_order=sort_order)
        return jsonify(videos)

class AddManyVideoUseCase(UseCase):
    def __init__(self, unit_of_work: Optional[BaseUnitfWork] = None) -> None:
        super().__init__(unit_of_work)

    def execute(self, data: dict[str, Any]) -> dict[str, Any]:
        with self.unit_of_work as uow:
            video_data_list = data["videos"]
            videos_added: int = 0
            for video_data in video_data_list:
                print(video_data)
                video = Video(
                    video_id=video_data["video_id"],
                    video_title=video_data["video_title"],
                    video_description=video_data["video_description"],
                    video_thumbnail=video_data["video_thumbnail"],
                    video_duration=video_data["video_duration"],
                    views_count=video_data["views_count"],
                    likes_count=video_data["likes_count"],
                    comments_count=video_data["comments_count"],
                    published_at=video_data["published_at"],
                    channel_id=video_data["channel_id"],
                )
                try:
                    uow.repository.add(video)
                    videos_added += 1
                except ResourceExistsException:
                    pass
        return {"Videos Added": videos_added}

class QueryVideoUseCase(QueryMixin, UseCase):
    def __init__(self, unit_of_work: Optional[BaseUnitfWork] = None) -> None:
        super().__init__(unit_of_work)

    def execute(self, data: dict[str, Any]) -> dict[str, Any]:
        with self.unit_of_work as uow:
            videos = []
            video_data_list = uow.repository.query(self.generate_query(data))
            if video_data_list:
                for video_data in video_data_list:
                    d = {}
                    for i, item in enumerate(video_data):
                        d[data["fields"][i]] = item
                    videos.append(d)
        return jsonify(videos)


import concurrent
import logging
import os
import re
import sys
import tarfile
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from logging import Logger
from typing import Any, Callable, Dict, List

from spark_rapids_tools import CspEnv
from spark_rapids_pytools.cloud_api.sp_types import get_platform, \
    ClusterBase, DeployMode, NodeHWInfo
from spark_rapids_pytools.common.prop_manager import YAMLPropertiesContainer
from spark_rapids_pytools.common.sys_storage import FSUtil, FileVerifier
from spark_rapids_pytools.common.utilities import ToolLogging, Utils, ToolsSpinner
from spark_rapids_pytools.rapids.rapids_job import RapidsJobPropContainer
from spark_rapids_pytools.rapids.tool_ctxt import ToolContext

@dataclass
class RapidsTool(object):
    """
    A generic class that represents a RAPIDS plugin tool.
    :param platform_type: the type of platform associated with the current execution.
    :param cluster: name of the cluster on which the application will be running
    :param output_folder: location to store the output of the execution
    :param config_path: location of the configuration file of the current tool
    :param wrapper_options: dictionary containing options specific to the wrapper tool execution.
    :param rapids_options: dictionary containing the options to be passed as CLI arguments to the RAPIDS Accelerator.
    :param name: the name of the tool
    :param ctxt: context manager for the current tool execution.
    :param logger: the logger instant associated to the current tool.
    """
    platform_type: CspEnv
    cluster: str = None
    output_folder: str = None
    config_path: str = None
    wrapper_options: dict = field(default_factory=dict)
    rapids_options: dict = field(default_factory=dict)
    name: str = field(default=None, init=False)
    ctxt: ToolContext = field(default=None, init=False)
    logger: Logger = field(default=None, init=False)

    def pretty_name(self):
        return self.name.capitalize()

    def get_exec_cluster(self) -> ClusterBase:
        return self.ctxt.get_ctxt('execCluster')

    def is_remote_cluster_execution(self) -> bool:
        """
        used to verify whether a rapids tool runs on a remote cluster submission.
        This does not include the serverlessMode
        :return: True when the tool needs to have a remote cluster established
        """
        return self.ctxt.get_deploy_mode() == DeployMode.REMOTE_CLUSTER

    def requires_remote_folder(self) -> bool:
        """
        used to verify whether a rapids tool running remotely has a defined remote path to generate the
        rapids tool output.
        :return: True when the tool needs to have a remote cluster folder
        """
        deply_mode: DeployMode = self.ctxt.get_deploy_mode()
        return deply_mode.requires_remote_storage()

    def requires_cluster_connection(self) -> bool:
        return self.is_remote_cluster_execution()

    def timeit(timed_item: str):  # pylint: disable=no-self-argument
        def decorator(func_cb: Callable):
            def wrapper(self, *args, **kwargs):
                start_time = time.monotonic()
                func_cb(self, *args, **kwargs)  # pylint: disable=not-callable
                end_time = time.monotonic()
                self.logger.info('Total Execution Time: %s => %s seconds', timed_item,
                                 f'{(end_time-start_time):,.3f}')
            return wrapper
        return decorator

    def phase_banner(phase_name: str,  # pylint: disable=no-self-argument
                     enable_prologue: bool = True,
                     enable_epilogue: bool = True):
        def decorator(func_cb: Callable):
            def wrapper(self, *args, **kwargs):
                try:
                    if enable_prologue:
                        self.logger.info('******* [%s]: Starting *******', phase_name)
                    func_cb(self, *args, **kwargs)  # pylint: disable=not-callable
                    if enable_epilogue:
                        self.logger.info('======= [%s]: Finished =======', phase_name)
                except Exception:    # pylint: disable=broad-except
                    logging.exception('%s. Raised an error in phase [%s]\n',
                                      self.pretty_name(),
                                      phase_name)
                    sys.exit(1)
            return wrapper
        return decorator

    def __post_init__(self):
        # when debug is set to true set it in the environment.
        self.logger = ToolLogging.get_and_setup_logger(f'rapids.tools.{self.name}')

    def _check_environment(self) -> None:
        self.ctxt.platform.setup_and_validate_env()

    def _process_output_args(self):
        self.logger.debug('Processing Output Arguments')
        # make sure that output_folder is being absolute
        if self.output_folder is None:
            self.output_folder = Utils.get_rapids_tools_env('OUTPUT_DIRECTORY', os.getcwd())
        self.output_folder = FSUtil.get_abs_path(self.output_folder)
        self.logger.debug('Root directory of local storage is set as: %s', self.output_folder)
        self.ctxt.set_local_workdir(self.output_folder)
        self.ctxt.load_prepackaged_resources()

    def _process_rapids_args(self):
        pass

    def _process_custom_args(self):
        pass

    def _process_job_submission_args(self):
        pass

    @phase_banner('Process-Arguments')
    def _process_arguments(self):
        # 0- process the output location
        self._process_output_args()
        # 1- process any arguments to be passed to the RAPIDS tool
        self._process_rapids_args()
        # 2- we need to process the arguments of the CLI
        self._process_custom_args()
        # 3- process submission arguments
        self._process_job_submission_args()

    @phase_banner('Initialization')
    def _init_tool(self):
        self._init_ctxt()
        self._check_environment()

    def _init_ctxt(self):
        if self.config_path is None:
            self.config_path = Utils.resource_path(f'{self.name}-conf.yaml')
        self.ctxt = ToolContext(platform_cls=get_platform(self.platform_type),
                                platform_opts=self.wrapper_options.get('platformOpts'),
                                prop_arg=self.config_path,
                                name=self.name)

    def _run_rapids_tool(self):
        # 1- copy dependencies to remote server
        # 2- prepare the arguments
        # 3- create a submission job
        # 4- execute
        pass

    @phase_banner('Execution')
    def _execute(self):
        """
        Phase representing actual execution of the wrapper command.
        """
        self._run_rapids_tool()

    def _process_output(self):
        pass

    def _delete_local_dep_folder(self):
        # clean_up the local dependency folder
        local_dep_folder = self.ctxt.get_local_work_dir()
        if self.ctxt.platform.storage.resource_exists(local_dep_folder):
            self.ctxt.platform.storage.remove_resource(local_dep_folder)

    def _delete_remote_dep_folder(self):
        # clean up the remote dep folder first
        remote_dep_folder = self.ctxt.get_remote('depFolder')
        if self.ctxt.platform.storage.resource_exists(remote_dep_folder):
            # delete the folder. Note that for dataproc this will also delete the parent directory when it is empty
            self.ctxt.platform.storage.remove_resource(remote_dep_folder)

    def _download_remote_output_folder(self):
        # download the output folder in to the local one with overriding
        remote_output_folder = self.ctxt.get_remote('workDir')
        # for dataproc it is possible that the entire directory has been deleted when it is empty
        if self.ctxt.platform.storage.resource_exists(remote_output_folder):
            local_folder = self.ctxt.get_local('outputFolder')
            self.ctxt.platform.storage.download_resource(remote_output_folder, local_folder)

    def _download_output(self):
        self._delete_local_dep_folder()
        # clean up the remote dep folder first
        self._delete_remote_dep_folder()
        # download the output folder in to the local one with overriding
        self._download_remote_output_folder()

    @phase_banner('Generating Report Summary',
                  enable_epilogue=False)
    def _finalize(self):
        print(Utils.gen_str_header(f'{self.pretty_name().upper()} Report',
                                   ruler='_',
                                   line_width=100))
        self._write_summary()

    def _write_summary(self):
        pass

    @phase_banner('Archiving Tool Output')
    def _archive_phase(self):
        self._archive_results()

    def _archive_results(self):
        pass

    @phase_banner('Collecting-Results')
    def _collect_result(self):
        """
        Following a successful run, collect and process data as needed
        :return:
        """
        self._download_output()
        self._process_output()

    @phase_banner('Connecting to Execution Cluster')
    def _connect_to_execution_cluster(self):
        """
        Connecting to execution cluster
        :return:
        """
        if self.requires_cluster_connection():
            self.logger.info('%s requires the execution cluster %s to be running. '
                             'Establishing connection to cluster',
                             self.pretty_name(),
                             self.cluster)
            exec_cluster = self.ctxt.platform.connect_cluster_by_name(self.cluster)
            self.ctxt.set_ctxt('execCluster', exec_cluster)
            self._verify_exec_cluster()
        else:
            self.logger.info('%s requires no execution cluster. Skipping phase', self.pretty_name())

    def _handle_non_running_exec_cluster(self, err_msg: str) -> None:
        self.logger.warning(err_msg)

    def _verify_exec_cluster(self):
        # For remote job we should fail once we find that the cluster is not actually running
        exec_cluster = self.get_exec_cluster()
        if exec_cluster and exec_cluster.is_cluster_running():
            return
        # For remote cluster mode, the execution cluster must be running
        if not exec_cluster:
            msg = 'An error initializing the execution cluster'
        else:
            msg = f'Remote execution Cluster [{exec_cluster.get_name()}] is not active. ' \
                  f'The execution cluster should be in RUNNING state'
        self._handle_non_running_exec_cluster(msg)

    def launch(self):
        with ToolsSpinner(in_debug_mode=ToolLogging.is_debug_mode_enabled()):
            self._init_tool()
            self._connect_to_execution_cluster()
            self._process_arguments()
            self._execute()
            self._collect_result()
            self._archive_phase()
            self._finalize()

    def _report_tool_full_location(self) -> str:
        pass

    def _report_results_are_empty(self):
        return [f'The {self.pretty_name()} tool did not generate any output. Nothing to display.']

    def _generate_section_lines(self, sec_conf: dict) -> List[str]:
        all_lines = sec_conf['content'].get('lines')
        if all_lines:
            return all_lines
        return None

    def _generate_section_content(self, sec_conf: dict) -> List[str]:
        sec_title = sec_conf.get('sectionName')
        rep_lines = []
        if sec_title:
            rep_lines.append(Utils.gen_report_sec_header(sec_title, title_width=20))
        if sec_conf.get('content'):
            headers = sec_conf['content'].get('header')
            if headers:
                rep_lines.extend(headers)
            all_lines = self._generate_section_lines(sec_conf)
            if all_lines:
                rep_lines.extend(all_lines)
        return rep_lines

    def _generate_platform_report_sections(self) -> List[str]:
        section_arr = self.ctxt.platform.configs.get_value_silent('wrapperReporting',
                                                                  self.name,
                                                                  'sections')
        if section_arr:
            rep_lines = []
            for curr_sec in section_arr:
                required_flag = curr_sec.get('requiresBoolFlag')
                # if section requires a condition that was not enabled the section is skipped
                if not required_flag or self.ctxt.get_ctxt(required_flag):
                    rep_lines.extend(self._generate_section_content(curr_sec))
            return rep_lines
        return None

    def _calculate_spark_settings(self, worker_info: NodeHWInfo) -> dict:
        """
        Calculate the cluster properties that we need to append to the /etc/defaults of the spark
        if necessary.
        :param worker_info: the hardware info as extracted from the worker. Note that we assume
                            that all the workers have the same configurations.
        :return: dictionary containing 7 spark properties to be set by default on the cluster.
        """
        num_gpus = worker_info.gpu_info.num_gpus
        gpu_mem = worker_info.gpu_info.gpu_mem
        num_cpus = worker_info.sys_info.num_cpus
        cpu_mem = worker_info.sys_info.cpu_mem

        config_path = Utils.resource_path('cluster-configs.yaml')
        constants = YAMLPropertiesContainer(prop_arg=config_path).get_value('clusterConfigs', 'constants')
        executors_per_node = num_gpus
        num_executor_cores = max(1, num_cpus // executors_per_node)
        gpu_concurrent_tasks = min(constants.get('maxGpuConcurrent'), gpu_mem // constants.get('gpuMemPerTaskMB'))
        # account for system overhead
        usable_worker_mem = max(0, cpu_mem - constants.get('systemReserveMB'))
        executor_container_mem = usable_worker_mem // executors_per_node
        # reserve 10% of heap as memory overhead
        max_executor_heap = max(0, int(executor_container_mem * (1 - constants.get('heapOverheadFraction'))))
        # give up to 2GB of heap to each executor core
        executor_heap = min(max_executor_heap, constants.get('heapPerCoreMB') * num_executor_cores)
        executor_mem_overhead = int(executor_heap * constants.get('heapOverheadFraction'))
        # use default for pageable_pool to add to memory overhead
        pageable_pool = constants.get('defaultPageablePoolMB')
        # pinned memory uses any unused space up to 4GB
        pinned_mem = min(constants.get('maxPinnedMemoryMB'),
                         executor_container_mem - executor_heap - executor_mem_overhead - pageable_pool)
        executor_mem_overhead += pinned_mem + pageable_pool
        res = {
            'spark.executor.cores': num_executor_cores,
            'spark.executor.memory': f'{executor_heap}m',
            'spark.executor.memoryOverhead': f'{executor_mem_overhead}m',
            'spark.rapids.sql.concurrentGpuTasks': gpu_concurrent_tasks,
            'spark.rapids.memory.pinnedPool.size': f'{pinned_mem}m',
            'spark.sql.files.maxPartitionBytes': f'{constants.get("maxSqlFilesPartitionsMB")}m',
            'spark.task.resource.gpu.amount': 1 / num_executor_cores,
            'spark.rapids.shuffle.multiThreaded.reader.threads': num_executor_cores,
            'spark.rapids.shuffle.multiThreaded.writer.threads': num_executor_cores,
            'spark.rapids.sql.multiThreadedRead.numThreads': max(20, num_executor_cores)
        }
        return res

@dataclass
class RapidsJarTool(RapidsTool):
    """
    A wrapper class to represent wrapper commands that require RAPIDS jar file.
    """

    def _process_jar_arg(self):
        tools_jar_url = self.wrapper_options.get('toolsJar')
        if tools_jar_url is None:
            tools_jar_url = self.ctxt.get_rapids_jar_url()
        # download the jar
        jar_path = self.ctxt.platform.storage.download_resource(tools_jar_url,
                                                                self.ctxt.get_local_work_dir(),
                                                                fail_ok=False,
                                                                create_dir=True)
        self.logger.info('RAPIDS accelerator jar is downloaded to work_dir %s', jar_path)
        # get the jar file name and add it to the tool args
        jar_file_name = FSUtil.get_resource_name(jar_path)
        self.ctxt.add_rapids_args('jarFileName', jar_file_name)
        self.ctxt.add_rapids_args('jarFilePath', jar_path)

    def __accept_tool_option(self, option_key: str) -> bool:
        defined_tool_options = self.ctxt.get_value_silent('sparkRapids', 'cli', 'toolOptions')
        if defined_tool_options is not None:
            if option_key not in defined_tool_options:
                self.logger.warning('Ignoring tool option [%s]. Invalid option.', option_key)
                return False
        return True

    def _process_tool_args_from_input(self) -> list:
        """
        Process the arguments passed from the CLI if any and return a list of strings representing
        the arguments to be passed to the final command running the job. This needs processing
        because we need to verify the arguments and handle hiphens
        :return: list of the rapids arguments added by the user
        """
        arguments_list = []
        self.logger.debug('Processing Rapids plugin Arguments %s', self.rapids_options)
        raw_tool_opts: Dict[str, Any] = {}
        for key, value in self.rapids_options.items():
            if not isinstance(value, bool):
                # a boolean flag, does not need to have its value added to the list
                if isinstance(value, str):
                    # if the argument is multiple word, then protect it with single quotes.
                    if re.search(r'\s|\(|\)|,', value):
                        value = f"'{value}'"
                raw_tool_opts.setdefault(key, []).append(value)
            else:
                if value:
                    raw_tool_opts.setdefault(key, [])
                else:
                    # argument parser removes the "no-" prefix and set the value to false.
                    # we need to restore the original key
                    raw_tool_opts.setdefault(f'no{key}', [])
        for key, value in raw_tool_opts.items():
            self.logger.debug('Processing tool CLI argument.. %s:%s', key, value)
            if len(key) > 1:
                # python forces "_" to "-". we need to reverse that back.
                fixed_key = key.replace('_', '-')
                prefix = '--'
            else:
                # shortcut argument
                fixed_key = key
                prefix = '-'
            if self.__accept_tool_option(fixed_key):
                k_arg = f'{prefix}{fixed_key}'
                if len(value) >= 1:
                    # handle list options
                    for value_entry in value[0:]:
                        arguments_list.append(f'{k_arg}')
                        arguments_list.append(f'{value_entry}')
                else:
                    # this could be a boolean type flag that has no arguments
                    arguments_list.append(f'{k_arg}')
        return arguments_list

    def _process_tool_args(self):
        """
        Process the arguments passed from the CLI if any and return a string representing the
        arguments to be passed to the final command running the job.
        :return:
        """
        self.ctxt.add_rapids_args('rapidsOpts', self._process_tool_args_from_input())

    def _process_dependencies(self):
        """
        For local deployment mode, we need to process the extra dependencies specific to the platform
        :return:
        """
        if 'deployMode' in self.ctxt.platform_opts:
            # process the deployment
            # we need to download the dependencies locally if necessary
            self._download_dependencies()

    def timeit(timed_item: str):  # pylint: disable=no-self-argument
        def decorator(func_cb: Callable):
            def wrapper(self, *args, **kwargs):
                start_time = time.monotonic()
                func_cb(self, *args, **kwargs)  # pylint: disable=not-callable
                end_time = time.monotonic()
                self.logger.info('Total Execution Time: %s => %s seconds', timed_item,
                                 f'{(end_time-start_time):,.3f}')
            return wrapper
        return decorator

    @timeit('Downloading dependencies for local Mode')  # pylint: disable=too-many-function-args
    def _download_dependencies(self):

        def exception_handler(future):
            # Handle any exceptions raised by the task
            exception = future.exception()
            if exception:
                self.logger.error('Error while downloading dependency: %s', exception)

        def cache_single_dependency(dep: dict) -> str:
            """
            Downloads the specified URL and saves it to disk
            """
            start_time = time.monotonic()
            self.logger.info('Checking dependency %s', dep['name'])
            dest_folder = self.ctxt.get_cache_folder()
            resource_file_name = FSUtil.get_resource_name(dep['uri'])
            resource_file = FSUtil.build_path(dest_folder, resource_file_name)
            file_check_dict = {'size': dep['size']}
            signature_file = FileVerifier.get_signature_file(dep['uri'], dest_folder)
            if signature_file is not None:
                file_check_dict['signatureFile'] = signature_file
            algorithm = FileVerifier.get_integrity_algorithm(dep)
            if algorithm is not None:
                file_check_dict['hashlib'] = {
                    'algorithm': algorithm,
                    'hash': dep[algorithm]
                }
            is_created = FSUtil.cache_from_url(dep['uri'], resource_file, file_checks=file_check_dict)
            if is_created:
                self.logger.info('The dependency %s has been downloaded into %s', dep['uri'],
                                 resource_file)
                # check if we need to decompress files
            if dep['type'] == 'archive':
                destination_path = self.ctxt.get_local_work_dir()
                with tarfile.open(resource_file, mode='r:*') as tar:
                    tar.extractall(destination_path)
                    tar.close()
                dep_item = FSUtil.remove_ext(resource_file_name)
                if dep.get('relativePath') is not None:
                    dep_item = FSUtil.build_path(dep_item, dep.get('relativePath'))
                dep_item = FSUtil.build_path(destination_path, dep_item)
            else:
                # copy the jar into dependency folder
                dep_item = self.ctxt.platform.storage.download_resource(resource_file,
                                                                        self.ctxt.get_local_work_dir())
            end_time = time.monotonic()
            self.logger.info('Completed downloading of dependency [%s] => %s seconds',
                             dep['name'],
                             f'{(end_time-start_time):,.3f}')
            return dep_item

        def cache_all_dependencies(dep_arr: List[dict]):
            """
            Create a thread pool and download specified urls
            """
            futures_list = []
            results = []
            with ThreadPoolExecutor(max_workers=4) as executor:
                for dep in dep_arr:
                    futures = executor.submit(cache_single_dependency, dep)
                    futures.add_done_callback(exception_handler)
                    futures_list.append(futures)
                try:
                    # set the timeout to 30 minutes.
                    for future in concurrent.futures.as_completed(futures_list, timeout=1800):
                        result = future.result()
                        results.append(result)
                except Exception as ex:    # pylint: disable=broad-except
                    self.logger.error('Failed to download dependencies %s', ex)
                    raise ex
            return results

        # TODO: Verify the downloaded file by checking their MD5
        deploy_mode = DeployMode.tostring(self.ctxt.get_deploy_mode())
        depend_arr = self.ctxt.platform.configs.get_value_silent('dependencies',
                                                                 'deployMode',
                                                                 deploy_mode)
        if depend_arr:
            dep_list = cache_all_dependencies(depend_arr)
            if any(dep_item is None for dep_item in dep_list):
                raise RuntimeError('Could not download all dependencies. Aborting Executions.')
            self.logger.info('Dependencies are processed as: %s',
                             Utils.gen_joined_str(join_elem='; ',
                                                  items=dep_list))
            self.ctxt.add_rapids_args('javaDependencies', dep_list)

    def _process_rapids_args(self):
        # add a dictionary to hold the rapids arguments
        self._process_jar_arg()
        self._process_dependencies()
        self._process_tool_args()

    def _process_offline_cluster_args(self):
        pass

    def _process_gpu_cluster_args(self, offline_cluster_opts: dict = None) -> bool:
        pass

    def _copy_dependencies_to_remote(self):
        self.logger.info('Skipping preparing remote dependency folder')

    def _prepare_job_arguments(self):
        self._prepare_local_job_arguments()

    def _run_rapids_tool(self):
        # 1- copy dependencies to remote server
        self._copy_dependencies_to_remote()
        # 2- prepare the arguments
        #  2.a -check if the app_id is not none
        self._prepare_job_arguments()
        #
        # 3- create a submission job
        # 4- execute

    def _get_main_cluster_obj(self):
        return self.ctxt.get_ctxt('cpuClusterProxy')

    def _process_eventlogs_args(self):
        eventlog_arg = self.wrapper_options.get('eventlogs')
        if eventlog_arg is None:
            # get the eventlogs from spark properties
            cpu_cluster_obj = self._get_main_cluster_obj()
            if cpu_cluster_obj:
                spark_event_logs = cpu_cluster_obj.get_eventlogs_from_config()
            else:
                self.logger.warning('Eventlogs is not set properly. The property cannot be pulled '
                                    'from cluster because it is not defined')
                spark_event_logs = []
        else:
            if isinstance(eventlog_arg, tuple):
                spark_event_logs = List[eventlog_arg]
            elif isinstance(eventlog_arg, str):
                spark_event_logs = eventlog_arg.split(',')
            else:
                spark_event_logs = eventlog_arg
        if len(spark_event_logs) < 1:
            self.logger.error('Eventlogs list is empty. '
                              'The cluster Spark properties may be missing "spark.eventLog.dir". '
                              'Re-run the command passing "--eventlogs" flag to the wrapper.')
            raise RuntimeError('Invalid arguments. The list of Apache Spark event logs is empty.')
        self.ctxt.set_ctxt('eventLogs', spark_event_logs)

    def _create_migration_cluster(self, cluster_type: str, cluster_arg: str) -> ClusterBase:
        if cluster_arg is None:
            raise RuntimeError(f'The {cluster_type} cluster argument is not set.')
        arg_is_file = self.ctxt.platform.storage.is_file_path(cluster_arg)
        if not arg_is_file:
            self.logger.info('Loading %s cluster properties by name %s. Note that this will fail '
                             'if the cluster was permanently deleted.',
                             cluster_type,
                             cluster_arg)
            # create a cluster by name
            cluster_obj = self.ctxt.platform.connect_cluster_by_name(cluster_arg)
        else:
            self.logger.info('Loading %s cluster properties from file %s',
                             cluster_type,
                             cluster_arg)
            # create cluster by loading properties files
            # download the file to the working directory
            cluster_conf_path = self.ctxt.platform.storage.download_resource(cluster_arg,
                                                                             self.ctxt.get_local_work_dir())
            cluster_obj = self.ctxt.platform.load_cluster_by_prop_file(cluster_conf_path)
        return cluster_obj

    def _gen_output_tree(self) -> List[str]:
        tree_conf = self.ctxt.get_value('local', 'output', 'treeDirectory')
        if tree_conf and tree_conf.get('enabled'):
            level = tree_conf.get('depthLevel')
            indentation = tree_conf.get('indentation', '\t')
            ex_patterns = tree_conf.get('excludedPatterns', {})
            exc_dirs = ex_patterns.get('directories')
            exc_files = ex_patterns.get('files')
            out_folder_path = self.ctxt.get_local('outputFolder')
            out_tree_list = FSUtil.gen_dir_tree(out_folder_path,
                                                depth_limit=level,
                                                indent=indentation,
                                                exec_dirs=exc_dirs,
                                                exec_files=exc_files)
            doc_url = self.ctxt.get_value('sparkRapids', 'outputDocURL')
            out_tree_list.append(f'{indentation}- To learn more about the output details, visit {doc_url}')
            return out_tree_list
        return None

    def _report_tool_full_location(self) -> str:
        if not self._rapids_jar_tool_has_output():
            return None
        out_folder_path = self.ctxt.get_rapids_output_folder()
        res_arr = [Utils.gen_report_sec_header('Output'),
                   f'{self.pretty_name()} tool output: {out_folder_path}']
        out_tree_list = self._gen_output_tree()
        return Utils.gen_multiline_str(res_arr, out_tree_list)

    def _evaluate_rapids_jar_tool_output_exist(self) -> bool:
        """
        Used as a subtask of self._process_output(). this method has the responsibility of
        checking if the tools produced no output and take the necessary action
        :return: True if the tool has generated an output
        """
        rapids_output_dir = self.ctxt.get_rapids_output_folder()
        res = True
        if not self.ctxt.platform.storage.resource_exists(rapids_output_dir):
            res = False
            self.ctxt.set_ctxt('wrapperOutputContent',
                               self._report_results_are_empty())
            self.logger.info('The Rapids jar tool did not generate an output directory')
        self.ctxt.set_ctxt('rapidsOutputIsGenerated', res)
        return res

    def _rapids_jar_tool_has_output(self) -> bool:
        return self.ctxt.get_ctxt('rapidsOutputIsGenerated')

    @timeit('Processing job submission arguments')  # pylint: disable=too-many-function-args
    def _process_job_submission_args(self):
        self._process_local_job_submission_args()

    def _set_remote_folder_for_submission(self, requires_remote_storage: bool) -> dict:
        res = {}
        submission_args = self.wrapper_options.get('jobSubmissionProps')
        # get the root remote folder and make sure it exists
        remote_folder = submission_args.get('remoteFolder')
        # If remote_folder is not specified, then ignore it
        if requires_remote_storage:
            # if the remote storage required and no remote folder specified. then try to assign the
            # tmp storage of the exec_cluster to be used for storage
            archive_enabled = True
            if not remote_folder:
                # get the execCluster
                exec_cluster = self.get_exec_cluster()
                if exec_cluster:
                    remote_folder = exec_cluster.get_tmp_storage()
                    if remote_folder:
                        archive_enabled = False
            self.ctxt.set_ctxt('archiveToRemote', archive_enabled)
        if remote_folder is None:
            # the output is only for local machine
            self.logger.info('No remote output folder specified.')
            if requires_remote_storage:
                raise RuntimeError(f'Remote folder [{remote_folder}] is invalid.')
        else:
            if not self.ctxt.platform.storage.resource_exists(remote_folder):
                raise RuntimeError(f'Remote folder invalid path. [{remote_folder}] does not exist.')
            # now we should make the subdirectory to indicate the output folder,
            # by appending the name of the execution folder
            exec_full_name = self.ctxt.get_ctxt('execFullName')
            remote_workdir = FSUtil.build_url_from_parts(remote_folder, exec_full_name)
            self.ctxt.set_remote('rootFolder', remote_folder)
            self.ctxt.set_remote('workDir', remote_workdir)
            self.logger.info('Remote workdir is set as %s', remote_workdir)
            remote_dep_folder = FSUtil.build_url_from_parts(remote_workdir,
                                                            self.ctxt.get_ctxt('depFolderName'))
            self.ctxt.set_remote('depFolder', remote_dep_folder)
            self.logger.info('Remote dependency folder is set as %s', remote_dep_folder)
        if requires_remote_storage:
            res.update({'outputDirectory': self.ctxt.get_remote('workDir')})
        else:
            # the output folder has to be set any way
            res.update({'outputDirectory': self.ctxt.get_output_folder()})
        return res

    def _process_local_job_submission_args(self):
        job_args = {}
        submission_args = self.wrapper_options.get('jobSubmissionProps')
        job_args.update(self._set_remote_folder_for_submission(self.requires_remote_folder()))
        platform_args = submission_args.get('platformArgs')
        if platform_args is not None:
            processed_platform_args = self.ctxt.platform.cli.build_local_job_arguments(platform_args)
            ctxt_rapids_args = self.ctxt.get_ctxt('rapidsArgs')
            dependencies = ctxt_rapids_args.get('javaDependencies')
            processed_platform_args.update({'dependencies': dependencies})
            job_args['platformArgs'] = processed_platform_args
        self.ctxt.update_job_args(job_args)

    def _init_rapids_arg_list(self) -> List[str]:
        return []

    @timeit('Building Job Arguments and Executing Job CMD')  # pylint: disable=too-many-function-args
    def _prepare_local_job_arguments(self):
        job_args = self.ctxt.get_ctxt('jobArgs')
        # now we can create the job object
        # Todo: For dataproc, this can be autogenerated from cluster name
        rapids_arg_list = self._init_rapids_arg_list()
        ctxt_rapids_args = self.ctxt.get_ctxt('rapidsArgs')
        jar_file_path = ctxt_rapids_args.get('jarFilePath')
        rapids_opts = ctxt_rapids_args.get('rapidsOpts')
        if rapids_opts:
            rapids_arg_list.extend(rapids_opts)
        # add the eventlogs at the end of all the tool options
        rapids_arg_list.extend(self.ctxt.get_ctxt('eventLogs'))
        class_name = self.ctxt.get_value('sparkRapids', 'mainClass')
        rapids_arg_obj = {
            'jarFile': jar_file_path,
            'jarArgs': rapids_arg_list,
            'className': class_name
        }
        platform_args = job_args.get('platformArgs')
        spark_conf_args = {}
        job_properties_json = {
            'outputDirectory': job_args.get('outputDirectory'),
            'rapidsArgs': rapids_arg_obj,
            'sparkConfArgs': spark_conf_args,
            'platformArgs': platform_args
        }
        job_properties = RapidsJobPropContainer(prop_arg=job_properties_json,
                                                file_load=False)
        job_obj = self.ctxt.platform.create_local_submission_job(job_prop=job_properties,
                                                                 ctxt=self.ctxt)
        job_obj.run_job()

    def _archive_results(self):
        self._archive_local_results()

    def _archive_local_results(self):
        remote_work_dir = self.ctxt.get_remote('workDir')
        if remote_work_dir and self._rapids_jar_tool_has_output():
            local_folder = self.ctxt.get_output_folder()
            # TODO make sure it worth issuing the command
            self.ctxt.platform.storage.upload_resource(local_folder, remote_work_dir)


#
# python3 rl/rl_daily/env_buy_only.py
# tensorboard --logdir logs
#

import numpy as np
import gym
from gym import spaces

import pandas as pd
import pandas_ta  # for TA magic

from stable_baselines3 import PPO
from stable_baselines3 import A2C
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.noise import ActionNoise
from gym.envs.registration import EnvSpec
from stable_baselines3.common.env_checker import check_env

from os import path

from plotly.subplots import make_subplots

INITIAL_DEPOSIT = 3000    # The agent starts with $3000 deposit
ORDER_FEES = 1            # $1 for execution of each order
STATE_DATA = 100
PROFIT_PERIOD = 20

ACTIONS = {0: 0, 1: 0, 2: 0}

class ProfitEnv(gym.Env):
    df = None
    current_index = 0
    last_bought_position = 0
    last_bought_price = 0
    stop_loss = 0
    take_profit = 0
    budget = 0
    number_of_shares_bought = 0
    total_deals = 0
    metadata = {'render.modes': ['human']}
    ticker = ''
    mode = ''
    sell_signals = None
    buy_signals = None
    hold_signals = None
    sum_profit = 0
    tickers = None
    initial_tickers = []
    failures = 0
    spec = EnvSpec(id='ProfitEnv-v0', max_episode_steps=2000, reward_threshold=5000)
    total_rewards = 0
    done = False

    _max_episode_steps = 2000  # PROFIT_PERIOD * 10  # get profit in 50 days

    def __init__(self, tickers_list: list = None, mode='learn'):
        super(ProfitEnv, self).__init__()

        self.mode = mode

        self.action_space = spaces.Discrete(2)  # hold or buy only

        # Current price: Low, High, Open, Close
        # EMA3, EMA7, EMA20, EMA70, EMA200, SMA50, SMA200,
        # Bollinger_low, Bollinger_up, RSI, SuperTrend,
        # MACD, MACD_signal, ATR, current profit,
        self.observation_space = spaces.Box(low=-1, high=1, shape=(STATE_DATA, 23), dtype=np.float32)   # shape=(1, 20)

        self.tickers = tickers_list.copy()
        self.initial_tickers = tickers_list.copy()

        self.ticker = self.tickers.pop()
        self.get_data(self.ticker)

        self.reset()

    def prepare_data(self):
        k_period = 100
        k_indicator = {i: 0 for i in range(200)}
        k2_indicator = {i: 0 for i in range(200)}

        try:
            self.df.ta.ema(length=3, append=True, col_names=('EMA3',))
            self.df.ta.ema(length=7, append=True, col_names=('EMA7',))
            self.df.ta.ema(length=20, append=True, col_names=('EMA20',))

            self.df.ta.sma(length=50, append=True, col_names=('SMA50',))
            self.df.ta.sma(length=200, append=True, col_names=('SMA200',))

            self.df.ta.rsi(append=True, col_names=('RSI',))
            self.df['RSI'] /= 100

            self.df.ta.atr(append=True, col_names=('ATR',))
            self.df['ATR'] /= 100

            self.df.ta.bbands(col_names=('L', 'M', 'U', 'B', 'P'), append=True)   # length=34,

            self.df.ta.macd(append=True, col_names=('MACD', 'MACD_hist', 'MACD_signal'))

            self.df.ta.supertrend(append=True, length=10, multiplier=3.0,
                                  col_names=('S_trend', 'S_trend_d', 'S_trend_l', 'S_trend_s',))

            # Calculate our own metrics:

            self.df['K2'] = (self.df['High'] - self.df['Low']) / self.df['L']
            self.df['K'] = (self.df['EMA20'] - self.df['EMA7']) / self.df['Close'] - 0.01

            self.df['diff1'] = self.df['EMA7'] / self.df['SMA50'] - 1
            self.df['diff2'] = (self.df['EMA7'] - self.df['EMA3']) / self.df['L']

            """
            for i in range(200, len(self.df['Close'])):
                # Integral-based metrics
                diff = self.df['SMA200'].values[i - k_period:i] - self.df['EMA7'].values[i - k_period:i]

                prob_up = diff[diff > 0].sum()
                prob_down = abs(diff[diff < 0].sum())

                k_indicator[i] = prob_up / (prob_up + prob_down)

                diff = self.df['EMA7'].values[i - k_period:i] - self.df['EMA3'].values[i - k_period:i]
                prob_up = diff[diff > 0].sum()
                prob_down = abs(diff[diff < 0].sum())

                k2_indicator[i] = prob_up / (prob_up + prob_down)

            self.df['W'] = k_indicator.values()
            self.df['W2'] = k2_indicator.values()
            """
        except Exception as e:
            print(f'WRONG DATA {self.ticker} {e}')
            exit(1)

        # self.df = self.df.replace([np.inf, -np.inf], np.nan)

        low_50 = {i: 0 for i in range(50)}
        high_50 = {i: 0 for i in range(50)}
        for i in range(50, len(self.df['Close'])):
            low_50[i] = self.df[i - 50:i]['Low'].min()
            high_50[i] = self.df[i - 50:i]['High'].max()

        low_30 = {i: 0 for i in range(30)}
        high_30 = {i: 0 for i in range(30)}
        for i in range(30, len(self.df['Close'])):
            low_30[i] = self.df[i - 30:i]['Low'].min()
            high_30[i] = self.df[i - 30:i]['High'].max()

        self.df['low_30'] = low_30.values()
        self.df['low_50'] = low_50.values()
        self.df['high_30'] = high_30.values()
        self.df['high_50'] = high_50.values()

        self.df = self.df[200:].copy().fillna(0)

    def get_state(self):
        states = []

        for i in range(self.current_index - STATE_DATA, self.current_index):
            state = []

            scaled_parameters = [
                'Low', 'High', 'Open', 'Close', 'EMA3', 'EMA7', 'EMA20',
                'SMA50', 'SMA200', 'L', 'U',
            ]

            for parameter in scaled_parameters:
                state.append(self.df[parameter].values[i])

            parameters = ['RSI', 'MACD', 'MACD_signal', 'ATR', 'S_trend_d', 'K', 'K2',
                          'low_30', 'low_50', 'high_30', 'high_50']   # 'S_trend_d', 'MACD', 'MACD_signal', 'ATR',
            for parameter in parameters:
                state.append(self.df[parameter].values[i])

            # Our own moving average crossover for daily trades
            # state.append(self.df['diff1'].values[i])
            # state.append(self.df['diff2'].values[i])

            # state.append(self.df['W'].values[i])
            # state.append(self.df['W2'].values[i])

            state.append(self.last_bought_price)  # add the info about purchased price to the state

            state = np.array(state, dtype=np.float32).flatten()

            states.append(state)

        return np.array(states)   # states[0]  # np.array(states)  # states

    def step(self, action):
        # actions: 0 - hold, 1 - buy, 2 - sell

        reward = 0  # to motivate purchases
        profit = 0

        current_price = self.df['Close'].values[self.current_index]

        if self.number_of_shares_bought > 0:
            if self.df['Low'].values[self.current_index] < self.stop_loss:
                reward = -1

                # Sell by stop loss in 3% down
                self.last_bought_price = 0
                self.last_bought_position = 0
                self.number_of_shares_bought = 0

                self.sell_signals[self.current_index] = True

            elif current_price > self.take_profit:
                profit = (current_price - self.last_bought_price) / self.last_bought_price

                self.last_bought_price = 0
                self.last_bought_position = 0
                self.number_of_shares_bought = 0

                reward = 1

                self.sell_signals[self.current_index] = True

        if action == 1:
            if self.number_of_shares_bought == 0:
                # bet_size = min(self.budget, 20000)  # limit size of every deal

                number_of_shares = 100  #  bet_size // current_price

                self.last_bought_price = current_price
                self.last_bought_position = self.current_index
                self.number_of_shares_bought = number_of_shares
                self.total_deals += 1

                self.buy_signals[self.current_index] = True

                reward = 0.01

        if action == 0:
            self.hold_signals[self.current_index] = True
            reward = 0

        self.current_index += 1
        # (self.current_index > 100 and self.total_rewards < 3) or
        done = bool(self.current_index >= len(self.df['Close']) - PROFIT_PERIOD)
        self.done = done

        if done and len(self.tickers) > 0:
            done = False

            ticker = self.tickers.pop()

            self.get_data(ticker)
            self.prepare_data()
            self.reset(cleanup_budget=False)
        elif not self.tickers and self.mode != 'test':
            self.tickers = self.initial_tickers.copy()

        budget = self.budget + self.last_bought_price * self.number_of_shares_bought
        if budget < 500:
            self.budget = 3000

        if self.mode != 'test':
            self.total_rewards += reward

        ACTIONS[action] += 1

        return self.get_state(), reward, done, {}

    def get_data(self, ticker):
        file_name = f'rl/data/daily_{ticker}.xlsx'
        self.ticker = ticker

        if path.exists(file_name):
            self.df = pd.read_excel(file_name, index_col=0)  # , header=[0, 1]
        else:
            self.df = pd.DataFrame()
            self.df = self.df.ta.ticker(ticker, period='700d', interval='1h')
            self.prepare_data()

            self.df.index = self.df.index.strftime('%m/%d/%Y %H:%M')
            self.df.to_excel(file_name, index=True, header=True)

        scaled_parameters = [
            'Low', 'High', 'Open', 'Close', 'EMA3', 'EMA7', 'EMA20',
            'SMA50', 'SMA200', 'L', 'U'
        ]

        for parameter in scaled_parameters:
            self.df[parameter] = self.df[parameter] / self.df[parameter].values[0]  # self.df['SMA200'] - 1

        # self.df = self.df.replace([np.inf, -np.inf], np.nan)

    def reset(self, cleanup_budget=True):

        print(self.total_rewards, self.current_index)

        if cleanup_budget:
            self.budget = INITIAL_DEPOSIT
        else:
            if self.last_bought_position > 0:
                self.budget += self.last_bought_price * self.number_of_shares_bought

        self.current_index = STATE_DATA  # we start from the Day #5 in the data set to have a history
        self.last_bought_position = 0
        self.last_bought_price = 0
        self.number_of_shares_bought = 0
        self.sum_profit = 0

        self.buy_signals = {i: False for i in range(len(self.df))}
        self.sell_signals = {i: False for i in range(len(self.df))}
        self.hold_signals = {i: False for i in range(len(self.df))}

        self.failures = 0
        self.total_rewards = 0
        self.done = False

        return self.get_state()

    def draw(self):
        df = self.df

        graph = make_subplots(rows=1, cols=1)
        graph.update_layout(title=self.ticker, xaxis_rangeslider_visible=False)

        graph.add_candlestick(open=df['Open'], high=df['High'], low=df['Low'], close=df['Close'])

        graph.add_scatter(y=df['SMA50'], mode='lines', name='SMA50',
                          line={'color': 'magenta'})
        # graph.add_scatter(y=df['SMA200'], mode='lines', name='SMA200',
        #                   line={'color': 'blue', 'width': 3})

        buy_signals = [i for i, (index, row) in enumerate(df.iterrows()) if row['buy_signals'] == True]
        graph.add_scatter(x=buy_signals, y=df['Close'].values[buy_signals], name='Buy Signal', mode='markers',
                          marker_symbol='triangle-up', marker_color='#00FE35', marker_size=15, row=1, col=1)

        sell_signals = [i for i, (index, row) in enumerate(df.iterrows()) if row['sell_signals'] == True]
        graph.add_scatter(x=sell_signals, y=df['Close'].values[sell_signals], name='Sell Signal', mode='markers',
                          marker_symbol='triangle-down', marker_color='#D62728', marker_size=15, row=1, col=1)

        hold_signals = [i for i, (index, row) in enumerate(df.iterrows()) if row['hold_signals'] == True]
        graph.add_scatter(x=hold_signals, y=df['Close'].values[hold_signals], name='HOLD Signal', mode='markers',
                          marker_symbol='hash', marker_color='#555555', marker_size=15, row=1, col=1)

        graph.show()

    def render(self, mode='human', close=False):
        if self.total_deals > 0 and self.total_deals % 100 == 0:
            total_budget = self.budget
            if self.last_bought_position > 0:
                total_budget += self.last_bought_price * self.number_of_shares_bought

            message = f"""
            Current budget: ${total_budget:.1f}
            Profit: {100 * (total_budget / INITIAL_DEPOSIT - 1):.0f}%
            Total Deals: {self.total_deals}
            """
            print(message)

        if self.total_rewards > 50:
            self.df['buy_signals'] = self.buy_signals.values()
            self.df['sell_signals'] = self.sell_signals.values()
            self.df['hold_signals'] = self.hold_signals.values()
            self.draw()

            exit()

with open('smp500.txt', 'r') as f:
    TICKERS = f.readlines()

TICKERS = [t.replace('\n', '') for t in TICKERS if '^' not in t and '/' not in t and '.' not in t]

TICKERS.remove('GOOG')
TICKERS.remove('TSLA')

TICKERS = ['USDJPY=X']
non_vectorized_env = ProfitEnv(TICKERS[:1].copy())  # :100
# wrapped_env = gym.wrappers.Monitor(non_vectorized_env, './results', force=True)
env = DummyVecEnv([lambda: non_vectorized_env])

state = env.reset()

# check_env(non_vectorized_env)

model = PPO('MlpPolicy', env, verbose=False, tensorboard_log='logs')  # learning_rate=0.0001
model.learn(total_timesteps=500000, progress_bar=True)

# non_vectorized_env = ProfitEnv(TICKERS[:1].copy(), mode='test')
# env = gym.wrappers.Monitor(non_vectorized_env, './results', force=True)
env.reset()
# env = DummyVecEnv([lambda: wrapped_env])

evaluate_policy(model, env, n_eval_episodes=1, render=True)


from __future__ import annotations
from typing import TYPE_CHECKING

import random

import aiohttp
import discord
from discord import app_commands
from discord import Embed, File
from discord.ext import commands
from discord.utils import escape_markdown

from utils import Color, Emoji
from utils.xp import Score

if TYPE_CHECKING:
    from tau import Tau

class Social(commands.Cog):
    def __init__(self, bot: Tau):
        self.bot = bot
        self._cd = commands.CooldownMapping.from_cooldown(1, 86400.0, commands.BucketType.user)

    async def fetch_gif(self, query: str) -> str:
        async with aiohttp.ClientSession() as session:
            query = query.replace(' ', '%20')
            filters = 'contentfilter=medium&mediafilter=minimal&limit=1'
            async with session.get(f'https://api.tenor.com/v1/random?q={query}&key={self.bot.conf.tenor_api_key}&{filters}') as response:
                content = await response.json()
                url = content['results'][0]['media'][0]['gif']['url']

            return url

    @app_commands.command(name='boop')
    async def boop(self, interaction: discord.Interaction, user: discord.User):
        '''Boop someone!'''
        gif = await self.fetch_gif('anime boop nose')

        recipient = 'themselves' if interaction.user == user else user.display_name
        embed = (
            Embed(color=Color.secondary)
            .set_author(name=f'{interaction.user.display_name} booped {recipient}!', icon_url=interaction.user.avatar)
            .set_image(url=gif)
        )
        await interaction.response.send_message(embed=embed)

    @app_commands.command(name='hug')
    async def hug(self, interaction: discord.Interaction, user: discord.User):
        '''Hug someone!'''
        gif = await self.fetch_gif('anime hug cute')

        recipient = 'themselves' if interaction.user == user else user.display_name
        embed = (
            Embed(color=Color.secondary)
            .set_author(name=f'{interaction.user.display_name} hugged {recipient}!', icon_url=interaction.user.avatar)
            .set_image(url=gif)
        )
        await interaction.response.send_message(embed=embed)

    @app_commands.command(name='kiss')
    async def kiss(self, interaction: discord.Interaction, user: discord.User):
        '''Kiss someone!'''
        gif = await self.fetch_gif('anime kiss')

        recipient = 'themselves' if interaction.user == user else user.display_name
        embed = (
            Embed(color=Color.pink)
            .set_author(name=f'{interaction.user.display_name} kissed {recipient}!', icon_url=interaction.user.avatar)
            .set_image(url=gif)
        )
        await interaction.response.send_message(embed=embed)

    @app_commands.command(name='pat')
    async def pat(self, interaction: discord.Interaction, user: discord.User):
        '''Headpat someone!'''
        gif = await self.fetch_gif('anime headpat')

        recipient = 'themselves' if interaction.user == user else user.display_name
        embed = (
            Embed(color=Color.secondary)
            .set_author(name=f'{interaction.user.display_name} patted {recipient}!', icon_url=interaction.user.avatar)
            .set_image(url=gif)
        )
        await interaction.response.send_message(embed=embed)

    @app_commands.command(name='slap')
    async def slap(self, interaction: discord.Interaction, user: discord.User):
        '''Slap someone! (not too hard tho)'''
        gif = await self.fetch_gif('anime slap')

        recipient = 'themselves' if interaction.user == user else user.display_name
        embed = (
            Embed(color=Color.red)
            .set_author(name=f'{interaction.user.display_name} slapped {recipient}!', icon_url=interaction.user.avatar)
            .set_image(url=gif)
        )
        await interaction.response.send_message(embed=embed)

    @app_commands.command(name='profile')
    @app_commands.checks.bot_has_permissions(external_emojis=True)
    @app_commands.guild_only()
    async def profile(self, interaction: discord.Interaction, member: discord.Member | None = None):
        '''View user profile'''
        member = member if member is not None else interaction.user
        if member not in self.bot.members:
            member_stats = await self.bot.members.add(member)
        else:
            member_stats = self.bot.members(member)
        embed = (
            Embed(color=Color.secondary)
            .set_author(name=member.display_name, icon_url=member.display_avatar.url)
            .add_field(name='Level', value=f'{member_stats.xp.levels} {Emoji.xp}')
            .add_field(name='Credits', value=f'{member_stats.credits} {Emoji.credits}')
        )
        await interaction.response.send_message(embed=embed)

    @app_commands.command(name='server')
    @app_commands.checks.bot_has_permissions(external_emojis=True)
    @app_commands.guild_only()
    async def server(self, interaction: discord.Interaction):
        '''Display server info'''
        guild = interaction.guild

        bots = 0
        for member in guild.members:
            if member.bot:
                bots += 1

        emoji_count = len(guild.emojis)
        sample = random.sample(guild.emojis, 10 if emoji_count > 10 else emoji_count)
        emojis = ''.join(str(emoji) for emoji in sample)

        banner_url = guild.banner.url if guild.banner is not None else None
        embed = (
            Embed(color=Color.primary)
            .set_author(name=guild.name, icon_url=guild.icon.url)
            .add_field(name=f'Owner {Emoji.owner}', value=f'`{guild.owner}`')
            .add_field(name=f'Server Boosts {Emoji.boost}', value=f'`{len(guild.premium_subscribers)}`')
            .add_field(name=f'Members `{guild.member_count}`', value=f'Users: `{guild.member_count-bots}`\nBots: `{bots}`')
            .add_field(name=f'Emoji `{emoji_count}`', value=emojis)
            .set_footer(text=f'ID: {guild.id}, created')
            .set_image(url=banner_url)
        )
        embed.timestamp = guild.created_at

        await interaction.response.send_message(embed=embed)

    @app_commands.command(name='leaderboard')
    @app_commands.guild_only()
    async def leaderboard(self, interaction: discord.Interaction):
        '''Display leaderboard'''
        highscores: tuple[Score] = await self.bot.members.fetch_highscores()
        embed = (
            Embed(color=Color.primary)
            .set_author(name='Leaderboard', icon_url='attachment://unknown.png')
        )
        for i, score in enumerate(highscores):
            name = escape_markdown(str(score.member))
            embed.add_field(name=f'{i+1}. {name}', value=f'**```yml\nLevel: {score.xp.levels}\nXP: {score.xp.points}```**', inline=i == 0)

        await interaction.response.send_message(embed=embed, file=File('assets/dot.png', 'unknown.png'))

async def setup(bot: Tau):
    await bot.add_cog(Social(bot))