import os.path
import json
from math import ceil
from skimage.transform import resize
import numpy as np
import matplotlib.pyplot as plt
import random
from os import listdir

# In this exercise task you will implement an image generator. Generator objects in python are defined as having a next function.
# This next function returns the next generated object. In our case it returns the input of a neural network each time it gets called.
# This input consists of a batch of images and its corresponding labels.
class ImageGenerator:
    def __init__(self, file_path, label_path, batch_size, image_size, rotation=False, mirroring=False, shuffle=False):
        # Define all members of your generator class object as global members here.
        # These need to include:
        # the batch size
        # the image size
        # flags for different augmentations and whether the data should be shuffled for each epoch
        # Also depending on the size of your data-set you can consider loading all images into memory here already.
        # The labels are stored in json format and can be directly loaded as dictionary.
        # Note that the file names correspond to the dicts of the label dictionary.

        self.class_dict = {0: 'airplane', 1: 'automobile', 2: 'bird', 3: 'cat', 4: 'deer', 5: 'dog', 6: 'frog',
                           7: 'horse', 8: 'ship', 9: 'truck'}
        self.file_path = file_path
        self.label_path = label_path
        self.batch_size = batch_size
        self.image_size = image_size
        self.rotation = rotation
        self.mirroring = mirroring
        self.shuffle = shuffle
        self.files = listdir(self.file_path)
        if self.shuffle:
            random.shuffle(self.files)
        with open(label_path) as f:
            self.labels = json.load(f)
        self.epoch = 0
        self.batch_number = 0
        self.output = None

    def next(self):
        # This function creates a batch of images and corresponding labels and returns them.
        # In this context a "batch" of images just means a bunch, say 10 images that are forwarded at once.
        # Note that your amount of total data might not be divisible without remainder with the batch_size.
        # Think about how to handle such cases
        first_batch_index = self.batch_number*self.batch_size

        #due to the fact that the tests consider that if the last batch ends with the last file, then the epoch has not changed yet
        if first_batch_index == len(self.files):
            self.epoch += 1
            self.batch_number = 0
            if self.shuffle:
                random.shuffle(self.files)

        last_batch_index = (self.batch_number + 1)*self.batch_size
        self.batch_number += 1
        additional_files = 0
        is_changed_epoch = False
        if last_batch_index > len(self.files):
            is_changed_epoch = True
            additional_files = last_batch_index - len(self.files)
        images_names = self.files[first_batch_index:last_batch_index] + self.files[:additional_files]

        if is_changed_epoch:
            self.epoch += 1
            self.batch_number = 0
            if self.shuffle:
                random.shuffle(self.files)

        labels = np.array([self.labels[image_name.split('.')[0]] for image_name in images_names])
        images = np.array(
            [
                self.augment(np.load(os.path.join(self.file_path, image_name)))
                for image_name in images_names])
        self.output = (images.copy(), labels.copy())
        return images, labels

    def augment(self,img):
        # this function takes a single image as an input and performs a random transformation
        # (mirroring and/or rotation) on it and outputs the transformed image
        if self.mirroring and random.choice([True, False]):
            img = np.flip(img, axis=1)
        if self.rotation and random.choice([True, False]):
            angle = random.choice([90, 180, 270])
            img = np.rot90(img, k=angle // 90)

        img = resize(img, self.image_size)
        return img

    def current_epoch(self):
        # return the current epoch number
        return self.epoch

    def class_name(self, x):
        # This function returns the class name for a specific input
        return self.class_dict[x]

    def show(self):
        # In order to verify that the generator creates batches as required, this functions calls next to get a
        # batch of images and labels and visualizes it.
        columns = 3
        rows = ceil(self.batch_size / columns)
        fig, axes = plt.subplots(rows, columns, figsize=(15, 6), dpi=90)
        axes_image = axes.flatten()[:self.batch_size]
        [ax.set_axis_off() for ax in axes.flatten()[self.batch_size:]]
        for i, ax in enumerate(axes_image):
            image = self.output[0][i]
            label = self.output[1][i]
            ax.imshow(image)
            ax.set_title(self.class_dict[label])
            ax.axis('off')
        plt.show()
        return


import numpy as np
import sys
import cv2 as cv
import glob
import matplotlib.pyplot as plt

# termination criteria

def main(image_directory, rows, columns, square_size):
    """
        This function takes images of a chessboard pattern and generates a camera calibration GH_matrix
        The function uses the open CV library built in funtions for camera calibration

        The images should be at a variety of angles  distances from the camera to get a good calibration
    """
    #read in the image files
    images = glob.glob(image_directory + 'cal_image*.npy')
    #Set the interation end criteria [defaults]
    criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    # prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
    objp = np.zeros((columns*rows,3), np.float32)
    objp[:,:2] = np.mgrid[0:rows,0:columns].T.reshape(-1,2)
    objp = objp * square_size # set the spacing of the object points

    # Arrays to store object points and image points from all the images.
    objpoints = [] # 3d point in real world space
    imgpoints = [] # 2d points in image plane.

    print(str(len(images)) + " images found")
    ret_num = 0
    for fname in images:
        img = np.uint8(np.load(fname))
        #show image
        cv.imshow('GrayImage',img)
        cv.waitKey(1000)
        cv.destroyAllWindows()
        # Find the chessboard corners
        ret, corners = cv.findChessboardCorners(img, (columns,rows), None)
        # If found, add object points, image points (after refining them)
        if ret == True:
            objpoints.append(objp)
            corners2 = cv.cornerSubPix(img,corners, (11,11), (-1,-1), criteria)
            imgpoints.append(corners)
            # Draw and display the corners
            cv.drawChessboardCorners(img, (rows,columns), corners2, ret)
            cv.imshow('img', img)
            cv.waitKey(1000)
            ret_num+=1
    cv.destroyAllWindows()

    print(f"{ret_num} patterns matched")
    # Create the calibration matrix
    ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(objpoints, imgpoints, img.shape[::-1], None, None)
    print(mtx)
    # Save the calibration matrix
    np.save(image_directory + 'calibration_matrix.npy',mtx)
    # Save distortion coefficents
    np.save(image_directory + 'distortion_coefficients.npy', dist)

if len(sys.argv) < 5:
    print('Please enter the following arguments:\n'
    '1) Directory of image files\n'
    '2) Number of rows in the pattern\n'
    '3) Number of columns in the pattern\n'
    '4) Dimensions of a square side in the pattern\n')
main(sys.argv[1],int(sys.argv[2]),int(sys.argv[3]),int(sys.argv[4]))


import os
import shutil
from keras.models import Model
from keras.layers import Input, LSTM, Dense
import numpy as np
from sklearn.preprocessing import OneHotEncoder
from keras.callbacks import CSVLogger, ModelCheckpoint, TensorBoard
from keras import optimizers

def one_hot(num_tokens, x):
    x_oh = np.zeros((x.shape[0], x.shape[1], num_tokens))
    for i in range(x.shape[0]):
       for j in range(x.shape[1]):
           x_oh[i, j, x[i,j]] = 1
    return x_oh

def get_params():
    P = {'num_encoder_tokens': 76,
         'num_decoder_tokens': 77,
         'ndata': 10000,
         'encoder_width': 60,
         'batch_size': 100,
         'epochs': 4,
         'validation_split': 0.12,
         'decoder_width': 40,
         'latent_dim': 300,
         'output_dir': 'output_lstm'}
    return P

def clean_dir(P):
    if os.path.isdir(P['output_dir']):
       print("deleting : ", P['output_dir'])
       shutil.rmtree(P['output_dir'])
    if not os.path.exists(P['output_dir']):
       os.makedirs(P['output_dir'])

def get_data(P):
    x = np.random.randint(P['num_encoder_tokens'], size=(P['ndata'], P['encoder_width']))
    y = np.random.randint(P['num_decoder_tokens'], size=(P['ndata'], P['decoder_width']))

    x_train = one_hot(P['num_encoder_tokens'], x)
    y_train = one_hot(P['num_decoder_tokens'], y)

    return x_train, y_train

def get_model(P):

    # Define an input sequence and process it.
    encoder_inputs = Input(shape=(None, P['num_encoder_tokens']))
    encoder = LSTM(P['latent_dim'], return_state=True)
    encoder_outputs, state_h, state_c = encoder(encoder_inputs)
    # We discard `encoder_outputs` and only keep the states.
    encoder_states = [state_h, state_c]

    # Set up the decoder, using `encoder_states` as initial state.
    decoder_inputs = Input(shape=(None, P['num_decoder_tokens']))
    # We set up our decoder to return full output sequences,
    # and to return internal states as well. We don't use the
    # return states in the training model, but we will use them in inference.
    decoder_lstm = LSTM(P['latent_dim'], return_sequences=True, return_state=True)
    decoder_outputs, _, _ = decoder_lstm(decoder_inputs,
                                     initial_state=encoder_states)
    decoder_dense = Dense(P['num_decoder_tokens'], activation='softmax')
    decoder_outputs = decoder_dense(decoder_outputs)

    # Define the model that will turn
    # `encoder_input_data` & `decoder_input_data` into `decoder_target_data`
    model = Model([encoder_inputs, decoder_inputs], decoder_outputs)

    # Define sampling models
    encoder_model = Model(encoder_inputs, encoder_states)

    decoder_state_input_h = Input(shape=(P['latent_dim'],))
    decoder_state_input_c = Input(shape=(P['latent_dim'],))
    decoder_states_inputs = [decoder_state_input_h, decoder_state_input_c]
    decoder_outputs, state_h, state_c = decoder_lstm(decoder_inputs, initial_state=decoder_states_inputs)
    decoder_states = [state_h, state_c]
    decoder_outputs = decoder_dense(decoder_outputs)
    decoder_model = Model( [decoder_inputs] + decoder_states_inputs,  [decoder_outputs] + decoder_states)

    return model, encoder_model, decoder_model

def fit_model(P, model, x_train, y_train):
    # Run training

    encoder_input_data = x_train
    decoder_input_data = y_train[:,:-1]
    decoder_target_data = y_train[:,1:]

    model_checkpoint = ModelCheckpoint(P['output_dir'] + os.sep + 'model.hdf5',
         monitor='val_loss', save_best_only=True, period=1)

    csv_logger = CSVLogger(P['output_dir'] + os.sep + 'history.log')

    tensorboard = TensorBoard(log_dir = P['output_dir'] + os.sep + "tensorboard",
                histogram_freq = 10,
                batch_size = P['batch_size'],
                write_graph = True,
                write_grads = False,
                write_images = False,
                embeddings_freq = 0,
                embeddings_layer_names = None,
                embeddings_metadata = None,
                embeddings_data = None)

    model.compile(optimizer=optimizers.Nadam(lr=0.01), loss='categorical_crossentropy', metrics=['acc'])
    #model.compile(optimizer='rmsprop', loss='categorical_crossentropy', metrics=['acc'])
    h = model.fit([encoder_input_data, decoder_input_data], decoder_target_data,
          batch_size=P['batch_size'],
          epochs=P['epochs'], callbacks= [csv_logger, model_checkpoint, tensorboard],
          validation_split=P['validation_split'])

    return h

if __name__ == "__main__":

    P = get_params()

    clean_dir(P)

    model, encoder_model, decoder_model = get_model(P)
    print(model.summary())

    x_train, y_train = get_data(P)

    h = fit_model(P, model, x_train, y_train)

    scores = model.evaluate([x_train, y_train[:,:-1]], y_train[:,1:], verbose=1)
    print("%s: %.2f%%" % (model.metrics_names[1], scores[1]*100))


import os
import flask
from flask import Flask, json, render_template, request, redirect, url_for
from flask_restful import Resource, Api
from flask.helpers import make_response
from flask_bootstrap import Bootstrap
from flask_cors import CORS
import discord
import threading

from rpi.mpc import Mpc
from rpi.speaker import Speaker
from rpi.radio import WebRadioEncoder

# Configuration region
app = Flask(__name__)
cors = CORS(app, resources={r"/api/*": {"origins": "*"}})

app.secret_key = os.urandom(12)

Bootstrap(app)
app.config['BOOTSTRAP_SERVE_LOCAL'] = True

# MPC declaration region
_mpc = Mpc()

# Speaker declaration region
_speaker = Speaker()

# API region
api = Api(app)

class getAllStation(Resource):
    def get(self):
        resp = flask.Response(WebRadioEncoder().encode(_mpc.webRadios.getAll()))
        return resp

class getVolume(Resource):
    def get(self):
        resp = flask.Response(_mpc.getVolume())
        return resp

class setVolume(Resource):
    def post(self):
        request_data = request.get_json()
        volume = request_data['volume']
        _mpc.volumeChange(volume)
        resp = make_response(json.dumps({'volume':_mpc.getVolume()}), 200)
        return resp

class playRadio(Resource):
    def post(self):
        request_data = request.get_json()
        name = request_data['name']
        url = request_data['url']
        radio = _mpc.webRadios.fromName(name)
        _mpc.play(radio)
        resp = make_response(json.dumps({'name':name, 'url': url}), 200)
        return resp

class stopRadio(Resource):
    def post(self):
        _mpc.stop()
        resp = flask.Response()
        return resp

class getPlayingStation(Resource):
    def get(self):
        radio = _mpc.webRadios.getPlaying()
        if radio != "":
            resp = make_response(json.dumps({'name':radio.name, 'url': radio.url}), 200)
            return resp
        else:
            return ""

class speakerOnOff(Resource):
    def post(self):
        _speaker.OnOff()
        return 'OK'

class speakerVolumeUP(Resource):
    def post(self):
        _speaker.VolumeUP()
        return 'OK'

class speakerVolumeDOWN(Resource):
    def post(self):
        _speaker.VolumeDOWN()
        return 'OK'

api.add_resource(getAllStation, '/api/getAllStation/')
api.add_resource(getVolume, '/api/getVolume/')
api.add_resource(setVolume,'/api/setVolume/')
api.add_resource(playRadio,'/api/playRadio/')
api.add_resource(stopRadio,'/api/stopRadio/')
api.add_resource(speakerOnOff,'/api/speakerOnOff/')
api.add_resource(speakerVolumeUP,'/api/speakerVolumeUP/')
api.add_resource(speakerVolumeDOWN,'/api/speakerVolumeDOWN/')

# Web page region
@app.route('/', methods=['GET', 'POST'])
def hello():

    if request.method == "GET":
        select = request.form.get('selectRadio')

        return render_template('index.html', radios = _mpc.webRadios.getAll(), actualPlay = _mpc.webRadios.getPlaying(), volume=_mpc.getVolume())

    if request.method == "POST":
        select = request.form.get('selectRadio')
        if request.form['button'] == 'Play':
            radio = _mpc.webRadios.fromName(select)
            _mpc.play(radio)
        if request.form['button'] == 'Stop':
            _mpc.stop()
        if request.form['button'] == 'VolumeUp':
            _mpc.volumeChange('+1')
        if request.form['button'] == 'VolumeUpUp':
            _mpc.volumeChange('+10')
        if request.form['button'] == 'VolumeDown':
            _mpc.volumeChange('-1')
        if request.form['button'] == 'VolumeDownDown':
            _mpc.volumeChange('-10')
        if request.form['button'] == 'SpeakerVolumeDown':
            _speaker.VolumeDOWN()
        if request.form['button'] == 'SpeakerVolumeUp':
            _speaker.VolumeUP()
        if request.form['button'] == 'SpeakerOnOff':
            _speaker.OnOff()
            if _speaker.Enable:
                sleep(5)
                os.system('/home/pi/RaspbianWebRadio/bin/bt_connect.sh')

        return redirect(request.referrer)

# Discord region
class DiscordClient(discord.Client):
    async def on_ready(self):
        print('Logged on as', self.user)

    async def on_message(self, message):
        # don't respond to ourselves
        if message.author == self.user:
            return

        if str(message.author) == 'panee#9393':
            if message.content == 'Ping':
                await message.channel.send('pong')

            elif message.content == 'Help':
                comandList = 'Ping\n' + 'Play:\n' +'Radios\n' +  'Playing\n' +'Volume\n'+ 'VolumeUP\n'+ 'VolumeDOWN\n'+ 'SpeakerOnOff\n'+ 'SpeakerVolumeUP\n' +'SpeakerVolumeDOWN\n'
                await message.channel.send(comandList)

            elif 'Play:' in message.content:
                table = message.content.split(':')
                radio = _mpc.webRadios.fromName(table[1])
                if radio is None:
                    await message.channel.send('Radio not exist')
                else:
                    _mpc.stop()
                    _mpc.play(radio)
                    await message.channel.send('Play ' + radio.name)

            elif message.content == 'Radios':
                listRadios = _mpc.webRadios.getAll()
                ret = ''
                for item in listRadios:
                    ret += str(item.name) + '\n'
                await message.channel.send(ret)

            elif message.content == 'Playing':
                radio = _mpc.webRadios.getPlaying()
                await message.channel.send(radio.name)

            elif message.content == 'Volume':
                await message.channel.send(_mpc.getVolume())

            elif message.content == 'VolumeUP':
                _mpc.volumeChange('+10')
                await message.channel.send(_mpc.getVolume())

            elif message.content == 'VolumeDOWN':
                _mpc.volumeChange('-10')
                await message.channel.send(_mpc.getVolume())

            elif message.content == 'SpeakerOnOff':
                _speaker.OnOff()
                await message.channel.send('OK')

            elif message.content == 'SpeakerVolumeUP':
                _speaker.VolumeUP()
                await message.channel.send('OK')

            elif message.content == 'SpeakerVolumeDOWN':
                _speaker.VolumeDOWN()
                await message.channel.send('OK')

            else:
                await message.channel.send('Command not exist')

def readDiscordTokenFromFile():
    with open('/home/pi/RaspbianWebRadio/rpi/token.txt') as f:
        token = f.readline()
    return token

client = DiscordClient()

# Start Flask thread
threading.Thread(target=lambda: app.run(host='0.0.0.0', port=5000)).start()

# Start Discord thread
threading.Thread(client.run(readDiscordTokenFromFile())).start()


from fastapi import FastAPI, WebSocket
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()

origins = [
    "http://localhost",
    "http://localhost:3000",
    "http://127.0.0.1:3000",
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/")
async def root():
    return {"message": "Hello World"}

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    # while True:
    #     data = await websocket.receive_text


#!/usr/bin/env python
def configuration(parent_package='', top_path=None):
    from numpy.distutils.misc_util import Configuration
    config = Configuration('wave_1d_fd_abc', parent_package, top_path)
    config.add_extension(name='oneway', sources=['oneway.f90'], extra_f90_compile_args=['-Wall', '-Wextra', '-pedantic', '-Werror', '-march=native', '-O2', '-std=f95', '-fbounds-check'])

    return config

if __name__ == "__main__":
    from numpy.distutils.core import setup
    setup(configuration=configuration)


GENDER_CRITERIA = 0.06
JWT_EXPIRE_MINUTES = 60*24
week_days = ("пн", "вт", "ср", "чт", "пт", "сб", "вс")
LOG_FILE = 'datamart.log'
DEBUG_LOG_FILE = 'datamart_debug.log'
CONFIG_PATH = '/path/kernel.conf'
TEMP_PHOTO_DIR = "/tmp/"
EMPLOYEES_PHOTO_DIR = "/opt/projects/images/employees"
DEBUG = False
EXCEL_DIR = '.'


from more_itertools import split_at

with open("calories.txt") as f:
    calories = [line for line in f]
    sublist = [
    sum([int(item.rstrip()) for item in sublist])
    for sublist in split_at(calories, lambda i: i == '\n')
    if sublist ]
    sublist.sort()
    print("largest element is:", sublist[-1])
    print("sum of three largest elements ",sum(sublist[-1:-4:-1]))


from main.data_base.query_to_parsing import select_filtered_results
import json

# ├── filteredQuery: {}
# │   ├──  domain_list: []
# │   │   │   ├──  domain_id: str
# │   ├──  start_date: str
# │   ├──  end_date: str
# │   ├──  only_last: bool
# │   ├──  get_checked: bool
# │   ├──  get_unchecked: bool
# пока не отправляется
# │   ├──  get_with_out_price: bool

def get_results_list(filteredQuery):
    results_list = []
    # print('Создаем запрос')
    start_date = filteredQuery['start_date'].split("-")
    filteredQuery['start_date'] = "/".join([start_date[2], start_date[1], start_date[0]])
    end_date = filteredQuery['end_date'].split("-")
    filteredQuery['end_date'] = "/".join([end_date[2], end_date[1], end_date[0]])

    if len(filteredQuery['domain_list']) == 0:
        del filteredQuery['domain_list']
        # print(filteredQuery)

    res_list = select_filtered_results(filteredQuery)
    # print(res_list)
    for el in res_list:
        # print(el.id)
        results_list.append(el.id)

    return json.dumps(results_list)


def incrementOperands(opList):
	"""
	Cycle through list of operations and increment to the next operand in the self.operands list
	"""

	operands = ["+", "-", "/", "*", "**"] # fully editable

	# Start from the end of the list and work backwards
	i = len(opList) - 1

	# If the last operation is the same as the last operand in operands list, increment to the first
	# operand in operands list (repeat over all cases)
	while opList[i] == operands[-1]:
		opList[i] = operands[0]
		i -= 1
	# After setting all the last operands to the first operand, set next operation to the next operand
	opList[i] = operands[operands.index(opList[i])+1]

def incrementNumbers(numList):
	nums = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

	i = len(numList) - 1
	while numList[i] == nums[-1]:
		numList[i] = nums[0]
		i -= 1

	numList[i] = nums[nums.index(numList[i])+1]

def evaluate(numList, opList):
	evalString = ""
	for i in range(0,3):
		evalString += str(numList[i]) + str(opList[i])
	evalString += str(numList[-1])

	stopCases = ["4**", "5**", "6**", "7**", "8**", "9**"]
	for case in stopCases:
		if case in evalString:
			return(0)

	try:
		return(eval(evalString))
	except ZeroDivisionError:
		return(0)

def getEvalString(numList, opList):
	evalString = ""
	for i in range(0,3):
		evalString += str(numList[i]) + str(opList[i])
	evalString += str(numList[-1])
	return(evalString)

numbers = [0, 0, 0, 0]
operands = ["+", "+", "+"]
hits = 0
idDict = {} # keys are number list ID, values are solutions (if any)

for i in range(0,10000):
	tmpList = numbers[:]
	tmpList.sort()
	numId = ''.join(str(x) for x in tmpList)
	if numId in idDict.keys() and idDict[numId] != "":
		incrementNumbers(numbers)
		continue
	else:
		idDict[numId] = ""

		print(numbers)
		operands = ["+", "+", "+"]

		for j in range(0,5**3-1):
			if evaluate(numbers, operands) == 13:
				print(str(numbers) + " is solvable.")
				print(getEvalString(numbers, operands))
				idDict[numId] = getEvalString(numbers, operands)
				hits += 1
				break
			incrementOperands(operands)
		incrementNumbers(numbers)

ans = float(hits)/len(idDict.keys())
print("Percentage of all license plates that are solvable: " + str(ans))


import os
import uuid
import yt_dlp
import shutil
import instaloader
from modules.logger import logger
from aiogram import types
from modules.mongo import save_error, save_stats
from modules.utils import reply_video, reply_voice, reply_photo, reply_text
from modules.config import (
    SUPPORTED_URLS,
    DOWNLOAD_STARTED,
    EX_VALID_LINK,
    MAX_SIZE_IN_MBYTES,
    EX_MAX_DURATION,
)
import copy

async def process_message(message: types.Message):
    url = message.text

    if not any(url.startswith(prefix) for prefix in SUPPORTED_URLS):
        await message.reply(EX_VALID_LINK)
        return

    try:
        file = None
        media_title = None
        await message.reply(DOWNLOAD_STARTED)
        if "instagram.com/" in url:
            url = url.replace("instagram.com", "ddinstagram.com")
            url = url.replace("www.", "")
            await reply_text(message, url)
        else:
            file = await download_media(url)
            #await reply_text(message, file)
            if file:
                temp_file = file[0]
                media_title = file[1]
                is_audio = file[2]
                if is_audio:
                    await reply_voice(message, temp_file, media_title)
                else:
                    await reply_video(message, temp_file)
    except Exception as e:
        await message.reply(f"[v1] Sorry, some error. {str(e)}")
        save_error(message.from_user.id, url, str(e))
    finally:
        media = media_title if file else ""
        save_stats(message.from_user.id, url, media)

async def download_media(url: str):
    temp_dir = "temp"
    temp_file_name = str(uuid.uuid4())
    temp_file = f"{temp_dir}/{temp_file_name}.mp4"
    os.makedirs(temp_dir, exist_ok=True)
    video_title = "untitled"

    ydl_opts = {
        "outtmpl": temp_file,
        "format": "best[filesize<=50M]/w[ext=mp4]",
    }

    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info = ydl.extract_info(url, download=False)  # Only fetch metadata
        duration = info.get("duration", 0)
        video_title = info.get("title", "untitled")

        if duration / 60 < 15:
            ydl.download([url])
        else:
            return await download_audio(url)

    return [temp_file, video_title, False]

async def download_audio(url: str):
    is_audio = True
    temp_dir = "temp"
    media_title = "untitled"
    temp_file_name = str(uuid.uuid4())
    temp_file = f"{temp_dir}/{temp_file_name}.mp3"

    os.makedirs(temp_dir, exist_ok=True)

    ydl_opts = {
        "outtmpl": temp_file,
        "format": "wa",
    }

    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
      info = ydl.extract_info(url, download=False)  # Only fetch metadata
      media_title = info.get("title", "untitled")
      filesize = info.get("filesize", 0) / 1024 / 1024  # MBytes

      if filesize <= MAX_SIZE_IN_MBYTES:
          ydl.download([url])
      else:
          raise Exception(EX_MAX_DURATION.format(MAX_SIZE_IN_MBYTES))

    return [temp_file, media_title, is_audio]


def quick_sort(arr, start, end):

    if start >= end:
        return
    low = start
    high = end
    mid = arr[low]

    while low < high:
        while low < high and arr[high] > mid:
            high -= 1
        arr[low] = arr[high]

        while low < high and arr[low] < mid:
            low += 1
        arr[high] = arr[low]

    arr[low] = mid

    quick_sort(arr, start, low-1)
    quick_sort(arr, low+1, end)

if __name__ == '__main__':
    arr = [123, 312, 12, 3, 122, 313, 2, 21313, 00, 566, 435, 23]
    quick_sort(arr, 0, len(arr)-1 )
    print(arr)


import copy

import numpy as np
from matplotlib import pyplot as plt, patches

from lib.aux import nam, moving_average
from lib import reg, aux, plot

def traj_1group(xy, c, unit='mm', title=None, single_color=False, **kwargs):
    ids = xy.index.unique('AgentID').values
    color = c.color if single_color else None
    scale = 1000 if unit == 'mm' else 1
    P = plot.AutoBasePlot(name=f'trajectories', **kwargs)
    ax = P.axs[0]
    tank = aux.get_tank_polygon(c, return_polygon=False) * scale
    for id in ids:
        xy0 = xy.xs(id, level="AgentID").values * scale
        ax.plot(xy0[:, 0], xy0[:, 1], color=color)
    ax.fill(tank[:, 0], tank[:, 1], fill=True, color='lightgrey', edgecolor='black', linewidth=4)
    for sid, sdic in c.env_params.food_params.source_units.items():
        px, py = sdic['pos']
        circle = plt.Circle((px * scale, py * scale), sdic['radius'] * scale, color=sdic['default_color'])
        ax.add_patch(circle)
    P.conf_ax(xMaxN=3, yMaxN=3, title=title, titlefontsize=25, xlab=f'X ({unit})', ylab=f'Y ({unit})',
              equal_aspect=True)
    return P.get()

def get_traj(d, mode='default'):
    if mode=='default':
        return d.load_traj(mode)
    elif mode == 'origin':
        try:
            ss=d.load_traj(mode)
            return ss[['x', 'y']]
        except:
            s = d.load_step(h5_ks=['contour', 'midline'])
            ss=reg.funcs.preprocessing["transposition"](s, c=d.config, store=True, replace=False, transposition='origin')
            return ss[['x', 'y']]

@reg.funcs.graph('trajectories')
def traj_grouped(unit='mm', name=None, subfolder='trajectories',
                 range=None, mode='default', single_color=False, **kwargs):
    if name is None:
        name = f'comparative_trajectories_{mode}'

    P = plot.AutoPlot(name=name, subfolder=subfolder,  # subplot_kw=dict(projection='polar'),
                 build_kws={'Nrows': 1, 'Ncols': 'Ndatasets', 'wh': 5, 'mode': 'both'}, **kwargs)
    for ii, (l, d) in enumerate(P.data_dict.items()):
        xy = get_traj(d, mode)
        c = d.config
        if range is not None:
            t0, t1 = range
            tick0, tick1 = int(t0 / c.dt), int(t1 / c.dt)
            xy = xy.loc[tick0:tick1]

        _ = traj_1group(xy, c, unit=unit, fig=P.fig, axs=P.axs[ii], title=l, single_color=single_color, save_to=None)
        if ii != 0:
            P.axs[ii].yaxis.set_visible(False)
    P.adjust((0.1, 0.9), (0.2, 0.9), 0.1, 0.01)
    return P.get()

def ax_conf_kws(kws, trange, Ndatasets,Nrows, i=0, ylab=None, ylim=None, xlim=None):
    conf_kws = {
        'ylab': kws.ylab if ylab is None else ylab,
        'ylim': kws.ylim if ylim is None else ylim,
        'xlim': (0, trange[-1]) if xlim is None else xlim,
        'xlab': r'time $(sec)$',
        'xvis' : True if i == Nrows - 1 else False,
    }

    leg_kws = {
        'leg_loc': "upper right",
        'leg_handles': [patches.Patch(color=col, label=l) for l, col in zip(kws.labels, kws.chunk_cols)],
        'leg_labels': kws.labels,
        'legfontsize': 15,
    }

    return {**conf_kws, **leg_kws}

def epoch_func(**kwargs):
    from lib.process.annotation import detect_turns, detect_strides, detect_pauses, process_epochs

    def stride_epochs(a, trange, show_extrema=True, a2plot=None, dt=0.1, **kwargs):

        if show_extrema and a2plot is None:
            i_min, i_max, strides, runs, run_counts = detect_strides(a=a, dt=dt, return_extrema=True)

            def func(ax):
                ax.plot(trange[i_max], a[i_max], linestyle='None', lw=10, color='green', marker='v')
                ax.plot(trange[i_min], a[i_min], linestyle='None', lw=10, color='red', marker='^')
        else:
            strides, runs, run_counts = detect_strides(a=a, dt=dt, return_extrema=False)

            def func(ax):
                pass
        pauses = detect_pauses(a, dt, runs=runs)
        # epochs1, epochs2, epochs0 = runs, pauses, strides
        epochs = [runs, pauses]
        epochs0 = strides
        return epochs, epochs0, func

    def turn_epochs(a, trange, min_amp=None, dt=0.1, **kwargs):
        def func(ax):
            ax.axhline(0, color='black', alpha=1, linestyle='dashed', linewidth=1)

        Lturns, Rturns = detect_turns(a, dt)
        if min_amp is not None:
            Lturns1, Ldurs, Lturn_slices, Lamps, Lturn_idx, Lmaxs = process_epochs(a, Lturns, dt)
            Rturns1, Rdurs, Rturn_slices, Ramps, Rturn_idx, Rmaxs = process_epochs(a, Rturns, dt)
            Lturns = Lturns[np.abs(Lamps) > min_amp]
            Rturns = Rturns[np.abs(Ramps) > min_amp]

        epochs = [Lturns, Rturns]
        epochs0 = Lturns.tolist() + Rturns.tolist()
        return epochs, epochs0, func

    epoch_dict = aux.AttrDict({
        'stride': {
            'ylab': "velocity (1/sec)",
            'ylim': (0.0,1.0),
            'labels': ['runs', 'pauses'],
            'chunk_cols': ["lightblue", "grey"],
            'func': stride_epochs,
            'k': 'sv'

        },
        'turn': {
            'ylab': "angular velocity (deg/sec)",
            'ylim': (-100.0,100.0),
            'labels': ['L turns', 'R turns'],
            'chunk_cols': ["lightgreen", "orange"],
            'func': turn_epochs,
            'k': 'fov'
        }
    })

    def epoch_f(epoch):
        kws = epoch_dict[epoch]

        def ss_f(ss, moving_average_interval=None, a2plot=None, dt=0.1, **kwargs):
            a = ss[reg.getPar(kws.k)].values
            trange = np.arange(0, a.shape[0] * dt, dt)
            if a2plot is not None:
                aa2plot = a2plot
            else:
                if moving_average_interval:
                    a = moving_average(a, n=int(moving_average_interval / dt))
                aa2plot = a

            epochs, epochs0, ax_func = kws.func(a, trange, a2plot=a2plot, dt=dt, **kwargs)

            def ax0_f(ax, epoch_boundaries=True):
                ax_func(ax)
                ax.plot(trange, aa2plot)

                if epoch_boundaries:
                    for s0, s1 in epochs0:
                        for s01 in [s0, s1]:
                            ax.axvline(trange[s01], color=f'{0.4 * (0 + 1)}', alpha=0.3, linestyle='dashed',
                                       linewidth=1)

                for color, epoch in zip(kws.chunk_cols, epochs):
                    for s0, s1 in epoch:
                        ax.axvspan(trange[s0], trange[s1], color=color, alpha=1.0)

            # return ax0_f

            def ax_conf0(P, **kwargs):
                ax_conf = ax_conf_kws(kws=kws, trange=trange, Ndatasets=P.Ndatasets,Nrows=P.Nrows, **kwargs)
                return ax_conf

            def P0_f(P, i=0,title=None,**kwargs):
                ax0_f(ax=P.axs[i], **kwargs)
                ax_conf = ax_conf0(P=P, i=i, **kwargs)
                P.conf_ax(i, title=title, **ax_conf)
                # for i,ax in enumerate(P.axs) :

            return P0_f

        return ss_f

    return epoch_f

def track_annotated(epoch='stride', a=None, dt=0.1, a2plot=None, ylab=None, ylim=None, xlim=None,
                    slice=None, agent_idx=0, agent_id=None,
                    subfolder='tracks', moving_average_interval=None, epoch_boundaries=True, show_extrema=True,
                    min_amp=None, **kwargs):
    from lib.process.annotation import detect_turns, detect_strides, detect_pauses, process_epochs
    temp = f'track_{slice[0]}-{slice[1]}' if slice is not None else f'track'
    name = f'{temp}_{agent_id}' if agent_id is not None else f'{temp}_{agent_idx}'
    P = plot.AutoPlot(name=name, subfolder=subfolder,
                 build_kws={'Nrows': 'Ndatasets', 'Ncols': 1, 'w': 20, 'h': 5, 'mode': 'both'}, ** kwargs)

    trange = np.arange(0, a.shape[0] * dt, dt)

    ax = P.axs[0]

    def stride_epochs(a, dt, ax):
        if show_extrema and a2plot is None:
            i_min, i_max, strides, runs, run_counts = detect_strides(a=a, dt=dt, return_extrema=True)
            ax.plot(trange[i_max], a[i_max], linestyle='None', lw=10, color='green', marker='v')
            ax.plot(trange[i_min], a[i_min], linestyle='None', lw=10, color='red', marker='^')
        else:
            strides, runs, run_counts = detect_strides(a=a, dt=dt, return_extrema=False)
        pauses = detect_pauses(a, dt, runs=runs)
        # epochs1, epochs2, epochs0 = runs, pauses, strides
        epochs = [runs, pauses]
        epochs0 = strides
        return epochs, epochs0

    def turn_epochs(a, dt, ax):
        ax.axhline(0, color='black', alpha=1, linestyle='dashed', linewidth=1)
        Lturns, Rturns = detect_turns(a, dt)
        if min_amp is not None:
            Ldurs, Lamps, Lmaxs = process_epochs(a, Lturns, dt, return_idx=False)
            Rdurs, Ramps, Rmaxs = process_epochs(a, Rturns, dt, return_idx=False)
            Lturns = Lturns[np.abs(Lamps) > min_amp]
            Rturns = Rturns[np.abs(Ramps) > min_amp]

        epochs = [Lturns, Rturns]
        epochs0 = Lturns.tolist() + Rturns.tolist()
        return epochs, epochs0

    epoch_dict = aux.AttrDict({
        'stride': {
            'ylab': "velocity (1/sec)",
            'labels': ['runs', 'pauses'],
            'chunk_cols': ["lightblue", "grey"],
            'func': stride_epochs,

        },
        'turn': {
            'ylab': "angular velocity (deg/sec)",
            'labels': ['L turns', 'R turns'],
            'chunk_cols': ["lightgreen", "orange"],
            'func': turn_epochs,
        }
    })

    kws = epoch_dict[epoch]

    epochs, epochs0 = kws.func(a, dt, ax=ax)

    conf_kws = {
        'ylab': kws.ylab if ylab is None else ylab,
        'ylim': ylim,
        'xlim': (0, trange[-1]) if xlim is None else xlim,
        'xlab': r'time $(sec)$' if 0 == P.Ndatasets - 1 else None,
    }

    if a2plot is not None:
        aa2plot = a2plot
    else:
        if moving_average_interval:
            a = aux.moving_average(a, n=int(moving_average_interval / dt))
        aa2plot = a

    ax.plot(trange, aa2plot)

    if epoch_boundaries:
        for s0, s1 in epochs0:
            for s01 in [s0, s1]:
                ax.axvline(trange[s01], color=f'{0.4 * (0 + 1)}', alpha=0.3, linestyle='dashed', linewidth=1)

    for color, epoch in zip(kws.chunk_cols, epochs):
        for s0, s1 in epoch:
            ax.axvspan(trange[s0], trange[s1], color=color, alpha=1.0)

    leg_kws = {
        'leg_loc': "upper right",
        'leg_handles': [patches.Patch(color=col, label=l) for l, col in zip(kws.labels, kws.chunk_cols)],
        'leg_labels': kws.labels,
        'legfontsize': 15,
    }

    P.conf_ax(0, **conf_kws, **leg_kws)
    return P.get()

def annotated_strideplot(**kwargs):
    return track_annotated(epoch='stride', **kwargs)

def annotated_turnplot(**kwargs):
    return track_annotated(epoch='turn', **kwargs)

def track_annotated_data(name=None, subfolder='tracks',
                         epoch='stride', a2plot_k=None, agent_idx=[3, 4, 5, 6, 7], dur=1, **kwargs):
    if name is None:
        name = f'annotated_{epoch}plot'
    Nidx = len(agent_idx)

    P = plot.AutoPlot(name=name, subfolder=subfolder,
                 build_kws={'Nrows': 'Ndatasets', 'Nrows_coef': Nidx, 'Ncols': 1, 'w': 15, 'h': 3, 'mode': 'both'},
                 **kwargs)
    epoch_kdic = {
        'stride': 'sv',
        'turn': 'fov'
    }
    apar = reg.getPar(epoch_kdic[epoch])

    def get_a(ss):
        return ss[apar].values

    if a2plot_k is not None:

        par, lab = reg.getPar(a2plot_k, to_return=['d', 'symbol'])
    else:
        par, lab = None, None

    def get_a2plot(ss):
        return ss[par].values if par is not None else None

    def get_title(idx,c,e, l):
        id = c.agent_ids[idx]
        ee = e.loc[id]

        length = np.round(ee['length'] * 1000, 2)
        cum_sd = np.round(ee[reg.getPar('cum_sd')], 2)
        run_tr = int(ee[reg.getPar('run_tr')] * 100)
        title = f'{l}  # {idx} track, l : {length} mm, pathlength {cum_sd}xl , {run_tr}% time crawling'
        return title

    for jj, (l, d) in enumerate(P.data_dict.items()):

        s, e, c = d.step_data, d.endpoint_data, d.config
        Nticks = int(dur * 60 / c.dt)
        kws0 = aux.AttrDict({
            'datasets': [d],
            'labels': [l],
            # 'agent_idx': idx,
            'slice': (0, dur * 60),
            'dt': c.dt,
            'fig': P.fig,
            'show': False,
            # 'a': a_dict[epoch],
            'epoch': epoch,
            'ylab': lab,
            # 'a2plot': ss[par].values if par is not None else None,
        })

        for i, idx in enumerate(agent_idx):
            ii = Nidx * jj + i
            id = c.agent_ids[idx]
            ss = s.xs(id, level='AgentID', drop_level=True).loc[:Nticks]
            # ee = e.loc[id]

            # length = np.round(ee['length'] * 1000, 2)
            # cum_sd = np.round(ee[preg.getPar('cum_sd')], 2)
            # run_tr = int(ee[preg.getPar('run_tr')] * 100)
            # title = f'{l}  # {idx} track, l : {length} mm, pathlength {cum_sd}xl , {run_tr}% time crawling'
            title=get_title(idx,c,e,l)
            kws1 = aux.AttrDict({
                'agent_idx': idx,
                'a': get_a(ss),
                'axs': P.axs[ii],
                'a2plot': get_a2plot(ss),
                **kws0
            })

            track_annotated(**kws1)
            P.conf_ax(ii, xvis=True if ii == P.Nrows - 1 else False, ylab=lab, title=title)
    P.adjust((0.1, 0.98), (0.05, 0.95), 0.001, 0.2)
    P.fig.align_ylabels(P.axs[:])
    return P.get()

@reg.funcs.graph('stride track')
def annotated_strideplot_data(**kwargs):
    return track_annotated_data(epoch='stride', **kwargs)

@reg.funcs.graph('turn track')
def annotated_turnplot_data(**kwargs):
    return track_annotated_data(epoch='turn', **kwargs)

@reg.funcs.graph('marked strides')
def plot_marked_strides(agent_idx=0, agent_id=None, slice=[20, 40], subfolder='individuals', **kwargs):
    temp = f'marked_strides_{slice[0]}-{slice[1]}' if slice is not None else f'marked_strides'
    name = f'{temp}_{agent_id}' if agent_id is not None else f'{temp}_{agent_idx}'
    P = plot.Plot(name=name, subfolder=subfolder, **kwargs)
    Nds = P.Ndatasets

    chunks = ['stride', 'pause']
    chunk_cols = ['lightblue', 'grey']
    p, ylab = reg.getPar('sv', to_return=['d', 'l'])
    figx = 15 * 6 * 3 if slice is None else int((slice[1] - slice[0]) / 3)
    figy = 5

    P.build(Nds, 1, figsize=(figx, figy * Nds), sharey=True, sharex=True)
    handles = [patches.Patch(color=col, label=n) for n, col in zip(['stride', 'pause'], chunk_cols)]

    for ii, (d, l) in enumerate(zip(P.datasets, P.labels)):
        ax = P.axs[ii]
        P.conf_ax(ii, xlab=r'time $(sec)$' if ii == Nds - 1 else None, ylab=ylab, ylim=[0, 1.0], xlim=slice,
                  leg_loc='upper right', leg_handles=handles)
        temp_id = d.agent_ids[agent_idx] if agent_id is None else agent_id
        s = copy.deepcopy(d.read('step').xs(temp_id, level='AgentID', drop_level=True))
        s.set_index(s.index * d.dt, inplace=True)
        ax.plot(s[p], color='blue')
        for i, (c, col) in enumerate(zip(chunks, chunk_cols)):
            s0s = s.index[s[nam.start(c)] == True]
            s1s = s.index[s[nam.stop(c)] == True]
            for s0, s1 in zip(s0s, s1s):
                ax.axvspan(s0, s1, color=col, alpha=1.0)
                ax.axvline(s0, color=f'{0.4 * (i + 1)}', alpha=0.6, linestyle='dashed', linewidth=1)
                ax.axvline(s1, color=f'{0.4 * (i + 1)}', alpha=0.6, linestyle='dashed', linewidth=1)

        ax.plot(s[p].loc[s[nam.max(p)] == True], linestyle='None', lw=10, color='green', marker='v')
        ax.plot(s[p].loc[s[nam.min(p)] == True], linestyle='None', lw=10, color='red', marker='^')
    P.adjust((0.08, 0.95), (0.15, 0.95), H=0.1)
    return P.get()

@reg.funcs.graph('sample tracks')
def plot_sample_tracks(mode=['strides', 'turns'], agent_idx=0, agent_id=None, slice=[20, 40], subfolder='individuals',
                       **kwargs):
    Nrows = len(mode)
    if Nrows == 2:
        suf = 'stridesVSturns'
    else:
        suf = mode[0]
    t0, t1 = slice
    temp = f'sample_marked_{suf}_{t0}-{t1}'
    name = f'{temp}_{agent_id}' if agent_id is not None else f'{temp}_{agent_idx}'
    P = plot.Plot(name=name, subfolder=subfolder, **kwargs)
    Nds = P.Ndatasets

    figx = 15 * 6 * 3 if slice is None else int((t1 - t0) / 3)
    figy = 5

    P.build(Nrows, Nds, figsize=(figx * Nds, figy * Nrows), sharey=False, sharex=True)
    for ii, (l, d) in enumerate(P.data_dict.items()):
        for jj, key in enumerate(mode):
            kk = ii + Nrows * jj
            ax = P.axs[kk]
            if key == 'strides':
                chunks = ['stride', 'pause']
                chunk_cols = ['lightblue', 'grey']

                p, ylab, ylim = reg.getPar('sv', to_return=['d', 'l', 'lim'])
                ylim = (0.0, 1.0)
            elif key == 'turns':
                chunks = ['Rturn', 'Lturn']
                chunk_cols = ['lightgreen', 'orange']

                b = 'bend'
                bv = nam.vel(b)
                ho = nam.orient('front')
                hov = nam.vel(ho)
                p, ylab, ylim = reg.getPar('fov', to_return=['d', 'l', 'lim'])

            handles = [patches.Patch(color=col, label=n) for n, col in zip(chunks, chunk_cols)]
            P.conf_ax(kk, xlab=r'time $(sec)$' if jj == Nrows - 1 else None, ylab=ylab, ylim=ylim, xlim=slice,
                      leg_loc='upper right', leg_handles=handles)

            temp_id = d.agent_ids[agent_idx] if agent_id is None else agent_id
            s = copy.deepcopy(d.read('step').xs(temp_id, level='AgentID', drop_level=True))
            s.set_index(s.index * d.dt, inplace=True)
            ax.plot(s[p], color='blue')
            for i, (c, col) in enumerate(zip(chunks, chunk_cols)):
                s0s = s.index[s[nam.start(c)] == True]
                s1s = s.index[s[nam.stop(c)] == True]
                for s0, s1 in zip(s0s, s1s):
                    ax.axvspan(s0, s1, color=col, alpha=1.0)
                    ax.axvline(s0, color=f'{0.4 * (i + 1)}', alpha=0.6, linestyle='dashed', linewidth=1)
                    ax.axvline(s1, color=f'{0.4 * (i + 1)}', alpha=0.6, linestyle='dashed', linewidth=1)

            ax.plot(s[p].loc[s[nam.max(p)] == True], linestyle='None', lw=10, color='green', marker='v')
            ax.plot(s[p].loc[s[nam.min(p)] == True], linestyle='None', lw=10, color='red', marker='^')
    P.adjust((0.08, 0.95), (0.12, 0.95), H=0.2)
    return P.get()

if __name__ == '__main__':

    refID = 'exploration.dish'
    d = reg.loadRef(refID)
    d.load()
    s, e,c = d.step_data,d.endpoint_data, d.config
    def get_title(idx):
        l=d.id

        id = c.agent_ids[idx]
        ss = s.xs(id, level='AgentID', drop_level=True)
        ee = e.loc[id]

        length = np.round(ee['length'] * 1000, 2)
        cum_sd = np.round(ee[reg.getPar('cum_sd')], 2)
        run_tr = int(ee[reg.getPar('run_tr')] * 100)
        title = f'{l}  # {idx} track, l : {length} mm, pathlength {cum_sd}xl , {run_tr}% time crawling'
        return title,ss

    P = plot.AutoPlot(name='name', subfolder=None, show=True,
                 build_kws={'Nrows': 5, 'Ncols': 2, 'w': 20, 'h':6, 'mode':'box'}, datasets=[d])

    # epoch = 'turn'

    f0 = epoch_func()
    # print(0, f0, type(f0))

    # f1 = f0(epoch=epoch)

    # print(1, f1, type(f1))

    for ii in range(P.Nrows):
        title, ss = get_title(idx=ii)
        for jj,ep in enumerate(['stride', 'turn']):
            f1 = f0(epoch=ep)
            f2 = f1(ss=ss, min_amp=30)
            f3 = f2(P=P, i=2*ii+jj, title=title)

     # print(3, f3, type(f3))

    P.get()


from PIL import Image
import os
import sys

def main():
    print(
        "+---------------------------------+\n" +
        "| Welcome to ImageOverlayMaker v1 |\n" +
        "+---------------------------------+"
    )
    print('First thing\'s first: what directory are we working in?')
    correct_dir = False
    old_dir = os.getcwd()
    while not correct_dir:
        dir_ = input('Enter directory: ')
        os.chdir(dir_)
        print(f'Set working directory to {os.getcwd()}')
        _correct_dir = input('Is this correct [Y/n]? ')
        if not _correct_dir or 'y' in _correct_dir:
            correct_dir = True
        else:
            os.chdir(old_dir)
            print('Whoops! Sorry about that, let\'s try again...')
    print('Great! Next, you\'ll need to select the image you want to put over your other images.')
    mask = None
    while not mask:
        mask = input('Enter image file name: ')
        print(f'Set image to {mask}')
        _correct_mask = input('Is this correct [Y/n]? ')
        if not _correct_mask or 'y' in _correct_mask:
            pass
        else:
            mask = None
            print('Whoops! Sorry about that, let\'s try again...')
    fore = Image.open(mask).convert('RGBA')
    print('Now, we\'re going to need the file(s) you plan to put this image on.')
    _images = input('Input image names: ')
    images = _images.split(',')
    for image in images:
        bg = Image.open(image).convert('RGBA')
        bg.paste(fore, (0, 0), fore)
        bg.save(image)
    print('All done!')
    input('Press any key to continue...')

def fast():
    dir_ = input('Enter directory: ')
    os.chdir(dir_)
    fore = Image.open('mask.png').convert('RGBA')
    for image in [f for f in os.listdir() if not f == 'mask.png']:
        bg = Image.open(image).convert('RGBA')
        bg.paste(fore, (0, 0), fore)
        bg.save(image)

if __name__ == '__main__':
    if 'fast' in sys.argv:
        fast()
    else:
        main()


import argparse
import glob
import os

import numpy as np
import statsmodels.stats.proportion as smp
from stable_baselines3.common.utils import set_random_seed

from indago.avf.config import (
    AVF_TEST_POLICIES,
    AVF_TRAIN_POLICIES,
    CLASSIFIER_LAYERS,
    CROSSOVER_RATE,
    DNN_SAMPLING_POLICIES,
    POPULATION_SIZE,
)
from indago.config import DONKEY_ENV_NAME, ENV_IDS, ENV_NAMES, HUMANOID_ENV_NAME, PARK_ENV_NAME
from indago.exp_configurator import ExpConfigurator
from indago.utils.env_utils import ALGOS
from log import Log, close_loggers

parser = argparse.ArgumentParser()
parser.add_argument("-f", "--folder", help="Log folder", type=str, default="logs")
parser.add_argument("--algo", help="RL Algorithm", default="sac", type=str, required=False, choices=list(ALGOS.keys()))
parser.add_argument("--seed", help="Random generator seed", type=int, default=-1)
parser.add_argument("--exp-id", help="Experiment ID (0: latest)", default=0, type=int)
parser.add_argument("--env-name", help="Env name", type=str, choices=ENV_NAMES, default="park")
parser.add_argument("--env-id", help="Env id", type=str, choices=ENV_IDS, default="parking-v0")

parser.add_argument("--num-episodes", help="Num episodes (i.e. num of env configurations) to run", type=int, default=1)
parser.add_argument("--failure-prob-dist", action="store_true", default=False)
parser.add_argument(
    "--num-runs-each-env-config",
    help="Num runs for each env configuration (valid when failure probability dist = True)",
    type=int,
    default=30,
)

# AVF params
parser.add_argument("--avf-test-policy", help="Avf policy testing", type=str, choices=AVF_TEST_POLICIES, default=None)
parser.add_argument("--avf-train-policy", help="Avf train policy", type=str, choices=AVF_TRAIN_POLICIES, default="mlp")
parser.add_argument("--exp-name", help="Experiment name suffix", type=str, default=None)
parser.add_argument(
    "--dnn-sampling",
    help="Sampling policy when using DNN predictors",
    type=str,
    choices=DNN_SAMPLING_POLICIES,
    default="original",
)
parser.add_argument("--budget", help="Timeout in seconds for the failure search technique", type=int, default=-1)
parser.add_argument(
    "--sampling-size", help="Sampling size when using original dnn sampling technique", type=int, default=100000
)
parser.add_argument("--neighborhood-size", help="Neighborhood size when using the hc techniques", type=int, default=50)
parser.add_argument("--hc-counter", help="Stopping counter for the hc techniques", type=int, default=100)
parser.add_argument(
    "--stagnation-tolerance", help="Tolerance when comparing max predictions in the hc techniques", type=float, default=0.005
)
parser.add_argument("--population-size", help="Population size for ga techniques", type=int, default=POPULATION_SIZE)
parser.add_argument("--crossover-rate", help="Mutation rate for ga techniques", type=float, default=CROSSOVER_RATE)
parser.add_argument("--training-progress-filter", help="Percentage of training to filter", type=int, default=10)
parser.add_argument(
    "--oversample",
    help="Percentage of oversampling of the minority class for the classification problem",
    type=float,
    default=0.0,
)
parser.add_argument("--layers", help="Num layers architecture", type=int, choices=CLASSIFIER_LAYERS, default=1)
parser.add_argument("--regression", action="store_true", default=False)
parser.add_argument(
    "--model-checkpoint",
    help="Model checkpoint to load (valid when estimate failure probability = True)",
    type=int,
    default=-1,
)
parser.add_argument(
    "--exp-file",
    help="Experiment file produced with a previous run that shows the failures the agent experienced at test time",
    type=str,
    default=None,
)
parser.add_argument(
    "--parallelize",
    help="Parallelize experiments for failure_probability_estimation (not valid for random and prioritized_replay avf_test_policies)",
    action="store_true",
    default=False,
)
parser.add_argument(
    "--build-denoised-failure-set",
    help="Execute all training failures and build a dataset discarding those in which the trained agent does not fail",
    action="store_true",
    default=False,
)

# DonkeyCar parameters
parser.add_argument("-vae", "--vae-path", help="Path to saved VAE", type=str, default=None)
parser.add_argument(
    "--add-to-port", help="Adding to default port 9091 in order to execute more simulators in parallel", type=int, default=-1
)
parser.add_argument("--simulation-mul", help="Speed up DonkeyCar simulation by at most 5x", type=int, default=1)
parser.add_argument(
    "--z-size", help="Latent space size. Needs to match the latent space of the trained VAE", type=int, default=64
)
parser.add_argument("--exe-path", help="DonkeyCar simulator execution path", type=str, default=None)

parser.add_argument("--num-runs-experiments", help="Number of times to run the experiments", type=int, default=1)

parser.add_argument(
    "--resume-dir",
    help="Directory where the model retrained with test failures is. "
         "By default the best model from the directory will be loaded.",
    type=str,
    default=None,
)

args = parser.parse_args()

def run_experiments(
    env_name: str,
    env_id: str,
    seed: int,
    regression: bool,
    num_envs: int,
    folder: str,
    algo: str,
    exp_id: int,
    avf_train_policy: str,
    avf_test_policy: str,
    dnn_sampling: str,
    sampling_size: int,
    neighborhood_size: int,
    hc_counter: int,
    stagnation_tolerance: float,
    model_checkpoint: int,
    training_progress_filter: int,
    layers: int,
    oversample_minority_class_percentage: float,
    failure_prob_dist: bool,
    num_episodes: int,
    num_runs_each_env_config: int,
    exp_name: str,
    vae_path: str,
    add_to_port: int,
    simulation_mul: int,
    z_size: int,
    exe_path: str,
    exp_file: str,
    parallelize: bool,
    budget: int,
    population_size: int,
    crossover_rate: float,
    resume_dir: str
) -> None:
    logger = Log(logger_prefix="experiments")
    logger.info("Args: {}".format(args))

    if seed == -1:
        seed = np.random.randint(2 ** 32 - 1)

    # also set when instantiating the algorithm
    set_random_seed(seed=seed)

    exp_configurator = ExpConfigurator(
        env_name=env_name,
        env_id=env_id,
        seed=seed,
        regression=regression,
        num_envs=num_envs,
        folder=folder,
        algo=algo,
        exp_id=exp_id,
        avf_train_policy=avf_train_policy,
        avf_test_policy=avf_test_policy,
        dnn_sampling=dnn_sampling,
        sampling_size=sampling_size,
        neighborhood_size=neighborhood_size,
        stagnation_tolerance=stagnation_tolerance,
        hc_counter=hc_counter,
        model_checkpoint=model_checkpoint,
        training_progress_filter=training_progress_filter,
        layers=layers,
        oversample_minority_class_percentage=oversample_minority_class_percentage,
        failure_prob_dist=failure_prob_dist,
        num_episodes=num_episodes,
        num_runs_each_env_config=num_runs_each_env_config,
        exp_name=exp_name,
        vae_path=vae_path,
        add_to_port=add_to_port,
        simulation_mul=simulation_mul,
        z_size=z_size,
        exe_path=exe_path,
        exp_file=exp_file,
        parallelize_fp_dist=parallelize,
        budget=budget,
        population_size=population_size,
        crossover_rate=crossover_rate,
        resume_dir=resume_dir
    )
    # It will change when exp_file is not None
    num_episodes = exp_configurator.num_episodes

    if failure_prob_dist:
        num_experiments = 0
        num_failures = 0
        previous_env_config = None
        map_env_config_failure_prob = dict()
        num_trials = 0
        episode_num = 0
        while num_experiments < num_runs_each_env_config * num_episodes:
            if num_experiments % num_runs_each_env_config == 0 and num_experiments != 0:
                map_env_config_failure_prob[previous_env_config.get_str()] = (
                    num_failures / num_runs_each_env_config,
                    smp.proportion_confint(count=num_failures, nobs=num_runs_each_env_config, method="wilson"),
                )
                logger.info(
                    "Failure probability for env config {}: {}".format(
                        previous_env_config.get_str(), map_env_config_failure_prob[previous_env_config.get_str()]
                    )
                )
                num_failures = 0
                episode_num = 0
                num_trials += 1
            failure, env_config = exp_configurator.test_single_episode(episode_num=episode_num, num_trials=num_trials)
            logger.debug("Env config: {}".format(env_config.get_str()))
            previous_env_config = env_config
            if failure:
                num_failures += 1
            num_experiments += 1
            episode_num += 1
            logger.debug("Num experiments: {}/{}".format(num_experiments, num_runs_each_env_config * num_episodes))

        if len(map_env_config_failure_prob) > 0:
            map_env_config_failure_prob[previous_env_config.get_str()] = (
                num_failures / num_runs_each_env_config,
                smp.proportion_confint(count=num_failures, nobs=num_runs_each_env_config, method="wilson"),
            )
        values = []

        num_failures = 0
        for key, value in map_env_config_failure_prob.items():
            if value[0] > 0.5:
                num_failures += 1
                logger.info("FAIL - Failure probability for env config {}: {}".format(key, value))
            else:
                logger.info("Failure probability for env config {}: {}".format(key, value))
            values.append(value[0])

        if len(values) > 0:
            logger.info(
                "Failure probabilities: {}, Mean: {:.2f}, Std: {:.2f}, Min: {}, Max: {}".format(
                    values, np.mean(values), np.std(values), np.min(values), np.max(values)
                )
            )

        logger.info("{}".format(map_env_config_failure_prob.keys()))

    logger.info(
        "Number of evaluation predictions (i.e. number of times a model was used to make predictions): {}".format(
            exp_configurator.get_num_evaluation_predictions()
        )
    )
    exp_configurator.close_env()

def sort_filenames(f: str) -> int:
    return int(f.split("-")[-2])

if __name__ == "__main__":

    num_envs = 1
    if args.num_runs_experiments == 1:
        if args.exp_file is not None and 'trial' in args.exp_file and args.avf_test_policy == 'replay_test_failure':
            if args.avf_train_policy not in args.exp_file:
                name = args.exp_file.split('-')[1]
                exp_filenames = glob.glob(os.path.join(args.folder, args.algo, args.env_id + '_' + str(args.exp_id), "testing-{}-*-trial.txt".format(name)))
            else:
                name = args.exp_file.split('-')[2]
                exp_filenames = glob.glob(os.path.join(args.folder, args.algo, args.env_id + '_' + str(args.exp_id), "testing-{}-{}-*-trial.txt".format(args.avf_train_policy, name)))
            exp_filenames = sorted(exp_filenames, key=sort_filenames)
            assert len(exp_filenames) > 0, 'No matches for {}'.format(args.exp_file)
            for exp_filename in exp_filenames:
                run_experiments(
                    env_name=args.env_name,
                    env_id=args.env_id,
                    seed=args.seed,
                    regression=args.regression,
                    num_envs=num_envs,
                    folder=args.folder,
                    algo=args.algo,
                    exp_id=args.exp_id,
                    avf_train_policy=args.avf_train_policy,
                    avf_test_policy=args.avf_test_policy,
                    dnn_sampling=args.dnn_sampling,
                    sampling_size=args.sampling_size,
                    neighborhood_size=args.neighborhood_size,
                    stagnation_tolerance=args.stagnation_tolerance,
                    hc_counter=args.hc_counter,
                    model_checkpoint=args.model_checkpoint,
                    training_progress_filter=args.training_progress_filter,
                    layers=args.layers,
                    oversample_minority_class_percentage=args.oversample,
                    failure_prob_dist=args.failure_prob_dist,
                    num_episodes=args.num_episodes,
                    num_runs_each_env_config=args.num_runs_each_env_config,
                    exp_name=args.exp_name,
                    vae_path=args.vae_path,
                    add_to_port=args.add_to_port,
                    simulation_mul=args.simulation_mul,
                    z_size=args.z_size,
                    exe_path=args.exe_path,
                    exp_file=exp_filename[exp_filename.rindex('/') + 1:],
                    parallelize=args.parallelize,
                    budget=args.budget,
                    population_size=args.population_size,
                    crossover_rate=args.crossover_rate,
                    resume_dir=args.resume_dir
                )
                close_loggers()
        else:
            run_experiments(
                env_name=args.env_name,
                env_id=args.env_id,
                seed=args.seed,
                regression=args.regression,
                num_envs=num_envs,
                folder=args.folder,
                algo=args.algo,
                exp_id=args.exp_id,
                avf_train_policy=args.avf_train_policy,
                avf_test_policy=args.avf_test_policy,
                dnn_sampling=args.dnn_sampling,
                sampling_size=args.sampling_size,
                neighborhood_size=args.neighborhood_size,
                stagnation_tolerance=args.stagnation_tolerance,
                hc_counter=args.hc_counter,
                model_checkpoint=args.model_checkpoint,
                training_progress_filter=args.training_progress_filter,
                layers=args.layers,
                oversample_minority_class_percentage=args.oversample,
                failure_prob_dist=args.failure_prob_dist,
                num_episodes=args.num_episodes,
                num_runs_each_env_config=args.num_runs_each_env_config,
                exp_name=args.exp_name,
                vae_path=args.vae_path,
                add_to_port=args.add_to_port,
                simulation_mul=args.simulation_mul,
                z_size=args.z_size,
                exe_path=args.exe_path,
                exp_file=args.exp_file,
                parallelize=args.parallelize,
                budget=args.budget,
                population_size=args.population_size,
                crossover_rate=args.crossover_rate,
                resume_dir=args.resume_dir
            )
    else:
        for i in range(args.num_runs_experiments):
            run_experiments(
                env_name=args.env_name,
                env_id=args.env_id,
                seed=args.seed,
                regression=args.regression,
                num_envs=num_envs,
                folder=args.folder,
                algo=args.algo,
                exp_id=args.exp_id,
                avf_train_policy=args.avf_train_policy,
                avf_test_policy=args.avf_test_policy,
                dnn_sampling=args.dnn_sampling,
                sampling_size=args.sampling_size,
                neighborhood_size=args.neighborhood_size,
                stagnation_tolerance=args.stagnation_tolerance,
                hc_counter=args.hc_counter,
                model_checkpoint=args.model_checkpoint,
                training_progress_filter=args.training_progress_filter,
                layers=args.layers,
                oversample_minority_class_percentage=args.oversample,
                failure_prob_dist=args.failure_prob_dist,
                num_episodes=args.num_episodes,
                num_runs_each_env_config=args.num_runs_each_env_config,
                exp_name="{}-trial".format(i),
                vae_path=args.vae_path,
                add_to_port=args.add_to_port,
                simulation_mul=args.simulation_mul,
                z_size=args.z_size,
                exe_path=args.exe_path,
                exp_file=args.exp_file,
                parallelize=args.parallelize,
                budget=args.budget,
                population_size=args.population_size,
                crossover_rate=args.crossover_rate,
                resume_dir=args.resume_dir
            )
            close_loggers()


import torch
import numpy as np
import wandb

# def log(loss, accuracy, epoch):

def train_epoch(epoch, model, train_loader, optimizer, args_config, criterion, device):
    model.cuda()
    wandb.watch(model, criterion, log="all")
    for batch_index, (input, labels) in enumerate(train_loader):
        input = input.to(device)
        labels = labels.to(device)

        # Forward pass
        outputs = model(input)
        labels = labels.long()
        loss = criterion(outputs, labels)

        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        wandb.log({"epoch": epoch, "loss": loss.item()})

    print(f' Loss: {loss.item():.4f}')

    return model, optimizer

def check_accuracy( model, validation_loader, optimizer, args_config, criterion, device):
    model.cuda()
    with torch.no_grad():
        n_correct = 0
        n_samples = 0
        n_class_correct = [0 for i in range(4)]
        n_class_samples = [0 for i in range(4)]
        class_acc = [0 for i in range(4)]
        count = 0
        for batch_index, (input, labels) in enumerate(validation_loader):
            labels = labels.to(device)
            input = input.to(device)
            outputs = model(input)

            _, predicted = torch.max(outputs.data, 1)

            n_samples += labels.size(0)
            n_correct += (predicted == labels).sum()
            acc = n_correct / n_samples
            wandb.log({"acc":acc})
                # print(f'accurcy :{acc} %')
            count += 1
            # for i in range(args_config.batch_size):
            #     pred = predicted[i]
            #     if (labels[i] == pred):
            #         n_correct += 1

            for j in range(args_config.batch_size):
                label = labels[j]
                pred = predicted[j]
                if (label == pred):
                    n_class_correct[int(label)] += 1
                n_class_samples[int(label)] += 1

    # try:
    for k in range(args_config.num_classes):
        if not n_class_samples[k] == 0:
            class_acc[k] = 100.0 * n_class_correct[k] / n_class_samples[k]
            print(f'Accuracy of {k}: {class_acc[k]} %')
    # except:
    #     print("no match")

    print(f'accurcy :{acc} %')
    return acc


import typing as t
from datetime import datetime, timedelta, timezone
from operator import itemgetter

import pytest

from valuable import load
from valuable.utils import compose

def dataclass_or_skip(*args, **kwargs):
    """create a dataclass or skip the test"""
    dataclasses = pytest.importorskip('dataclasses')
    return dataclasses.make_dataclass(*args, **kwargs)

def test_list_loader():
    assert load.list_loader((int, ), [4, '3', '-1']) == [4, 3, -1]

def test_set_loader():
    assert load.set_loader((int, ), [4, '3', '-1']) == {4, 3, -1}

def test_lookup_default():
    getter = load.lookup_defaults(itemgetter('foo'), 'bla')
    assert getter({}) == 'bla'
    assert getter({'foo': 4}) == 4

class TestPrimitiveRegistry:

    def test_found(self):
        registry = load.PrimitiveRegistry({int: round})
        loader = registry(int)
        assert loader(3.4) == 3

    def test_not_found(self):
        registry = load.PrimitiveRegistry({int: round})
        with pytest.raises(load.UnsupportedType):
            registry(str)

class TestGenericRegistry:

    def test_ok(self):

        Tag = t.NamedTuple('Tag', [('name', str)])

        registry = load.GenericRegistry({
            t.List: load.list_loader,
        }) | load.PrimitiveRegistry({
            Tag: compose(Tag, '<{}>'.format),
        })

        # simple case
        loader = registry(t.List[Tag])
        assert loader(['hello', 5, 'there']) == [
            Tag('<hello>'), Tag('<5>'), Tag('<there>')
        ]

        # recursive case
        loader = registry(t.List[t.List[Tag]])
        assert loader([
            ['hello', 9, 'there'],
            [],
            ['another', 'list']
        ]) == [
            [Tag('<hello>'), Tag('<9>'), Tag('<there>')],
            [],
            [Tag('<another>'), Tag('<list>')]
        ]

    def test_unsupported_type(self):

        Tag = t.NamedTuple('Tag', [('name', str)])

        registry = load.GenericRegistry({
            t.List: load.list_loader,
        }) | load.PrimitiveRegistry({
            Tag: compose(Tag, '<{}>'.format),
        })
        with pytest.raises(load.UnsupportedType):
            registry(t.Set[Tag])

        with pytest.raises(load.UnsupportedType):
            registry(t.List[str])

        with pytest.raises(load.UnsupportedType):
            registry(object)

class TestGetOptionalLoader:

    def test_ok(self):

        Tag = t.NamedTuple('Tag', [('name', str)])

        registry = load.PrimitiveRegistry({
            Tag: compose(Tag, '<{}>'.format)
        }) | load.get_optional_loader

        loader = registry(t.Optional[Tag])
        assert loader(None) is None
        assert loader(5) == Tag('<5>')

class TestParseIso8601:

    def test_with_timezone(self):
        parsed = load.parse_iso8601('2012-02-27T13:08:00+0100')
        assert parsed == datetime(
            2012, 2, 27, 13, 8,
            tzinfo=timezone(timedelta(hours=1)))

    def test_no_timezone(self):
        parsed = load.parse_iso8601('2014-06-10T17:25:29Z')
        assert parsed == datetime(2014, 6, 10, 17, 25, 29)

def test_dataclass_registry(registry):

    Post = dataclass_or_skip('Post', [
        ('title', str),
        ('posted_at', datetime),
        ('author_id', int)
    ])

    data = {
        'Title':     'hello',
        'date':      '2017-10-18T14:13:05Z',
        'author_id': 12,
    }

    registry |= load.DataclassRegistry({
        Post: {
            'title':     itemgetter('Title'),
            'posted_at': itemgetter('date'),
            'author_id': itemgetter('author_id'),
        }
    })

    loader = registry(Post)

    assert loader(data) == Post(
        'hello',
        datetime(2017, 10, 18, 14, 13, 5),
        author_id=12)

    Foo = dataclass_or_skip('Foo', [])

    with pytest.raises(load.UnsupportedType):
        registry(Foo)

def test_create_dataclass_loader(registry):
    User = dataclass_or_skip('User', [
        ('id', int),
        ('name', str),
        ('hobbies', t.Optional[t.List[str]]),
        ('nicknames', t.Set[str]),
        ('height', float),
        ('likes_pizza', t.Optional[bool], True),
    ])

    dloader = load.create_dataclass_loader(User, registry, {
        'id':          itemgetter('id'),
        'name':        itemgetter('username'),
        'likes_pizza': itemgetter('pizza'),
        'height':      itemgetter('height'),
        'nicknames':   itemgetter('nicknames'),
        'hobbies':     itemgetter('hobbies'),
    })
    assert dloader({
        'id': 98,
        'username': 'wilma',
        'hobbies': ['tennis', 'diving'],
        'pizza': 'true',
        'height': 1.64,
        'nicknames': []
    }) == User(98, 'wilma', hobbies=['tennis', 'diving'],
               likes_pizza=True, height=1.64, nicknames=set())

class TestAutoDataclassRegistry:

    def test_ok(self, registry):

        Post = dataclass_or_skip('Post', [
            ('title', str),
            ('posted_at', datetime),
        ])

        registry |= load.AutoDataclassRegistry()
        loader = registry(Post)

        data = {
            'title':     'hello',
            'posted_at': '2017-10-18T14:13:05Z'
        }
        assert loader(data) == Post(
            'hello',
            datetime(2017, 10, 18, 14, 13, 5))

    def test_not_supported(self, registry):

        class MyClass():
            pass

        registry |= load.AutoDataclassRegistry()

        with pytest.raises(load.UnsupportedType):
            registry(MyClass)

def test_simple_registry():

    User = dataclass_or_skip('User', [
        ('name', str),
        ('id', int),
        ('nickname', t.Optional[str]),
    ])

    Post = dataclass_or_skip('Post', [
        ('title', str),
        ('user', User),
        ('comments', t.List[str])
    ])

    loader = load.simple_registry(Post)

    loaded = loader({
        'title': 'hello',
        'comments': ['first!', 'another comment', 5],
        'user': {
            'name': 'bob',
            'nickname': 'bobby',
            'id': '543',
            'extra data': '...',
        }
    })
    assert loaded == Post(
        title='hello',
        comments=['first!', 'another comment', '5'],
        user=User(
            name='bob',
            id=543,
            nickname='bobby'
        )
    )


import traceback
from typing import Union, Tuple
from neo4j import GraphDatabase, Session, Driver, Result
from neo4j.exceptions import ServiceUnavailable

class Neo4jBaseConnector:
    """
    Base class to connect to a neo4j database with utilities to check
    connectivity and use a class in a with... statement

    Attributes
    ----------
    driver : neo4j.Driver
        database driver

    session : neo4j.Session
        database session
    """
    driver: Driver
    session: Session

    def __init__(
        self, uri: str, auth: Union[Tuple[str, str], None] = None, **kwargs
    ):
        """
        Parameters
        ----------
        uri : str
            database uri

        auth : Tuple[str, str], Optional
            Optional database credentials int the form of (user, password).
            If database is not secured pass None.

        data : ResourceData (NamedTuple)
            Data to dump into the database

        verbose : bool, Optional, default True
            Whether function should be verbose

        reset_at_start : bool, Optional, default False
            If True the database will be emptied before data is added

        Raises
        ------
        ConnectionError
            If database is not reachable with the given parameters
        """
        self.driver = GraphDatabase.driver(uri, auth=auth, **kwargs)
        self.session = self.driver.session()
        if not self.active_connection():
            if auth is None:
                add = "If your database requires login credentials please " \
                      "add them via the 'auth' parameter"
            else:
                add = "Please make sure you parsed the correct login " \
                      "credentials"
            raise ConnectionError(
                f"Connection to {uri} failed. {add}. If you have not "
                "installed mantra's neo4j database locally or if you are "
                "running the local database in the pymantra-db-api docker "
                "please use the 'APINetworkGenerator' class for all graph "
                "queries."
            )

    def valid_port(self) -> bool:
        try:
            # this is just a simple query to check the driver connection
            self.session.run(
                "MATCH (n:metabolite) RETURN count(n) as count")
            return True
        except ServiceUnavailable:
            return False

    def active_connection(self) -> bool:
        """
        Verifying the connection to the database

        Returns
        -------
        bool
            True if connection is established, else False
        """
        try:
            # this is just a simple query to check the driver connection
            self.session.run(
                "MATCH (n:metabolite) RETURN count(n) as count")
            return True
        except ServiceUnavailable:
            return False

    def close(self):
        """
        Closing the database connection. Should always be called
        once the object instance is not needed anymore
        """
        if hasattr(self, "session"):
            # if constructor fails at driver call session is not set
            self.session.close()
        if hasattr(self, "driver"):
            self.driver.close()

    def __del__(self):
        self.close()

    # for usage with with(NetworkGenerator()) as ...
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, tb):
        self.close()
        if exc_type is not None:
            traceback.print_exception(exc_type, exc_value, tb)

    def run(self, query: str, *parameters, **kwargs) -> Result:
        """Run a neo4j query

        Parameters
        ----------
        query: str
            Query string
        parameters: Dict[str, any], optional
           Query parameters to insert into query
        kwargs
           Query parameters to insert into query, Preceding over those in
           `parameters`

        Returns
        -------
        neo4j.Result
            Query result
        """
        return self.session.run(query, *parameters, **kwargs)


'''
@author: kalarita
@Date: 2020-05-04 08:16:11
@LastEditors: kalarita
@LastEditTime: 2020-05-06 09:23:45
@note:部署完成之后我才想起来服务器时间还得注意,现在只是暂时调了一下,不清楚重启之后定时任务的情况
'''
import json
import sqlite3
import uuid
import os

from flask import Flask, make_response, request,Response
import hashlib
from flask.json import jsonify

app = Flask(__name__)
app.config["JSON_AS_ASCII"] = False

#接下来三个函数是在flask运行的时候直接把当天需要获取的数据类型直接放到内存中,并且常驻,在每次需要用的时候可以直接获取,减少硬盘IO
def get_his():
    hisresult = []
    with open('today.txt','r')as f:
        lines = f.readlines()
        for line in lines:
            hisresult.append(line)
    return hisresult

def get_img():
    with open('today.jpg','rb')as f:
        image = f.read()
    return image

def get_idea():
    idearesult = []
    with open('Idea.txt','r')as f:
        lines = f.readlines()
        for line in lines:
            idearesult.append(line)
    return idearesult

#初始化三个常驻变量
todayinhis = get_his()
image = get_img()
idea = get_idea()

#创建数据库,返回一个打开的数据库连接
def createDb():
    conn = sqlite3.connect('usr.db')
    cursor = conn.cursor()
    sql = '''CREATE TABLE usr (
        name varchar,
        id varchar,
        mail varchar,
        pwd varchar,
        primary key(mail)
    )
                '''
    cursor.execute(sql)
    conn.commit()
    return conn

#返回一个打开的数据库连接
def  connectDb(dbName):
    if os.path.exists('usr.db'):
        conn = sqlite3.connect(dbName)
    else:
        print("could not connect to database!")
        conn = createDb()
    return  conn

#注册接口,返回success为注册成功,fail为注册失败
@app.route('/register/<Secretkey>',methods = ['get'])
def register(Secretkey):                #通过传入的账户名,邮箱和密码,生成id,同时对密码进行加盐加密
    conn = connectDb('usr.db')
    cursor = conn.cursor()
    sqlall = 'select *from usr'
    beforetotal = len(cursor.execute(sqlall).fetchall())
    lst = Secretkey.split('+')
    print(lst)
    id = generateId(lst)
    print('id='+id+str(type(id)))
    shapwd = hashlib.sha1((lst[2]+lst[1][2:6]).encode('utf-8')).hexdigest() #盐值为账户的2-6位
    name = lst[0]
    mail = lst[1]
    sql = 'insert into usr values (?,?,?,?)'
    try:
        cursor.execute(sql,(name,id,mail,shapwd))
    except:
         resp=make_response(jsonify(ststus='alreadyregisted',account=mail))
    aftertotal = len(cursor.execute(sqlall).fetchall())
    if aftertotal-beforetotal ==1:
        resp=make_response(jsonify(ststus='registed',account=mail))
    else:
        resp=make_response(jsonify(ststus='alreadyregisted',account=mail))
    # print(conn.execute(sqlall).fetchall())
    conn.commit()
    conn.close()
    return resp
#使用uuid根据用户的邮箱产生不同的id
def generateId(lst):
    id = uuid.uuid3(uuid.NAMESPACE_DNS,lst[1])
    id = str(id)
    return id

#登录接口
@app.route('/login/<account>')
def varify(account):
    result = 'fail'
    #假设传进来的account值为邮箱+密码,例如1078950499@qq.com+bcs@12345
    accountlst = account.split('+')
    print(accountlst)
    conn = connectDb('usr.db')
    cursor = conn.cursor()
    sql = 'select pwd from usr where mail=\''+str(accountlst[0])+'\''
    print(sql)
    print(cursor.execute(sql).fetchall())
    if len(cursor.execute(sql).fetchall()[0])>0:
        fetchedpwd = cursor.execute(sql).fetchall()[0][0]
    else:
        fetchedpwd = ''
    total = len(cursor.execute(sql).fetchall())
    if total>0:
        print(fetchedpwd)
        shapwd = hashlib.sha1((accountlst[1]+accountlst[0][2:6]).encode('utf-8')).hexdigest()
        print(shapwd)
        if fetchedpwd == shapwd:
            resp=make_response(jsonify(login='True',account = accountlst[0],msg = 'login success!'))
        else:
            resp=make_response(jsonify(login='False',account = accountlst[0],msg = 'please check your account or your password!'))
    else:
        resp=make_response(jsonify(login='False',account = accountlst[0],msg = 'please check your account or your password!'))
    return resp

#获取历史上的今天信息
@app.route('/todayinhis',methods=['get'])
def respond():
    resp = make_response(jsonify(todayinhis))
    return resp

#获取今日图片
@app.route('/todaypic',methods=['get'])
def respondpic():
    resp = Response(image,mimetype='image/jepg')
    return resp

#获取今日idea
@app.route('/todayidea',methods=['get'])
def respondidea():
    resp = make_response(jsonify(idea))
    return resp


from datetime import datetime, timedelta
from textwrap import dedent

# The DAG object; we'll need this to instantiate a DAG
from airflow import DAG

# Operators; we need this to operate!
from airflow.operators.bash import BashOperator
with DAG(
    'themapark_etl',
    # These args will get passed on to each operator
    # You can override them on a per-task basis during operator initialization
    default_args={
        'depends_on_past': False,
        'email': ['qkqhdksi0@gmail.com'],
        'email_on_failure': False,
        'email_on_retry': False,
        'retries': 4,  # 재시도 횟수
        'retry_delay': timedelta(minutes=3),  # 재시도 딜레이 - 3분
    },
    description='Themapark ETL Project',
    #schedule=timedelta(days=1), # 반복날짜 - 1일마다
    start_date=datetime(2022, 10, 29, 0, 0),  # 시작날짜
    schedule_interval = '0 20 * * *', # 반복 - 매일 미국시간 20시에
    catchup=False,
    tags=['themapark_etl'],
) as dag:

    # 태스크들 추가
    # ㅡㅡㅡㅡㅡㅡ EXTRACT ㅡㅡㅡㅡㅡㅡ
    t1 = BashOperator(
        task_id='extract_today_weather',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py extract today_weather',
    )

    t2 = BashOperator(
        task_id='extract_today_dust',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py extract today_dust',
    )

    t3 = BashOperator(
        task_id='extract_pre_weather',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py extract pre_weather',
    )

    t4 = BashOperator(
        task_id='extract_pre_dust',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py extract pre_dust',
    )

    t5 = BashOperator(
        task_id='extract_navi_search',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py extract navi_search',
    )

    t6 = BashOperator(
        task_id='extract_event_childpark',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py extract event_childpark',
    )

    t7 = BashOperator(
        task_id='extract_event_seoulpark',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py extract event_seoulpark',
    )

    t8 = BashOperator(
        task_id='extract_everland_info',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py extract everland_info',
    )

    t9 = BashOperator(
        task_id='extract_lotteworld_info',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py extract lotteworld_info',
    )

    t18 = BashOperator(
        task_id='extract_holiday',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py extract holiday',
    )

    t20 = BashOperator(
        task_id='extract_subway_inout',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py extract subway_inout',
    )

    # ㅡㅡㅡㅡㅡㅡ TRANSFORM ㅡㅡㅡㅡㅡㅡ
    t10 = BashOperator(
        task_id='transform_today_weather',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py transform today_weather',
    )

    t11 = BashOperator(
        task_id='transform_today_dust',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py transform today_dust',
    )

    t12 = BashOperator(
        task_id='transform_themepark_event',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py transform themepark_event',
    )

    t13 = BashOperator(
        task_id='transform_navi_search',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py transform navi_search',
    )

    t21 = BashOperator(
        task_id='transform_subway_inout',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py transform subway_inout',
    )

    # ㅡㅡㅡㅡㅡㅡ DATAMART ㅡㅡㅡㅡㅡㅡ
    t14 = BashOperator(
        task_id='datamart_pre_air_weather',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py transform pre_air_weather',
    )

    t15 = BashOperator(
        task_id='datamart_pre_themepark_event',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py datamart pre_themepark_event',
    )

    t19 = BashOperator(
        task_id='datamart_holiday',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py datamart holiday',
    )

    # ㅡㅡㅡㅡㅡㅡ OPERATION ㅡㅡㅡㅡㅡㅡ
    t16 = BashOperator(
        task_id='operation_themepark_time',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py operation themepark_time',
    )

    t17 = BashOperator(
        task_id='operation_themepark_hol_fac',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py operation themepark_hol_fac',
    )

    t22 = BashOperator(
        task_id='operation_pre_air_weather',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py operation pre_air_weather',
    )

    t25 = BashOperator(
        task_id='operation_pre_event',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py operation pre_event',
    )

    # ㅡㅡㅡㅡㅡㅡ MODELING ㅡㅡㅡㅡㅡㅡ
    t23 = BashOperator(
        task_id='modeling_childpark',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py modeling childpark',
    )

    t24 = BashOperator(
        task_id='modeling_other_themepark',
        cwd='/home/big/pj/ETL',
        bash_command='python3 main.py modeling other_themepark',
    )

    t1.doc_md = dedent(
        """\
    #### Task Documentation
    You can document your task using the attributes `doc_md` (markdown),
    `doc` (plain text), `doc_rst`, `doc_json`, `doc_yaml` which gets
    rendered in the UI's Task Instance Details page.
    ![img](http://montcs.bloomu.edu/~bobmon/Semesters/2012-01/491/import%20soul.png)
    **Image Credit:** Randall Munroe, [XKCD](https://xkcd.com/license.html)
    """
    )

    dag.doc_md = __doc__  # providing that you have a docstring at the beginning of the DAG; OR
    dag.doc_md = """
    This is a documentation placed anywhere
    """  # otherwise, type it like this
    templated_command = dedent(
        """
    {% for i in range(5) %}
        echo "{{ ds }}"
        echo "{{ macros.ds_add(ds, 7)}}"
    {% endfor %}
    """
    )

    # 태스크 우선순위 설정
    # extract는 병렬로, transform과 datamart는 직렬로

    t1 >> t10  # 오늘 날씨
    t2 >> t11  # 오늘 미세먼지
    [t8, t9] >> t16 >> t17  # 미래 운영시간, 운휴시설 정보
    t3 >> t4 >> t14 >> t22  # 예보 날씨
    t5 >> t13  # 2일전 네비게이션 검색건수
    t18 >> t19  # 미래 공휴일정보
    [t6, t7] >> t12 >> t15 >> t25  # 오늘, 미래 행사정보
    t20 >> t21  # 3일전 지하철 승하차 정보
    [t22, t13, t19, t15, t21] >> t23 >> t24


import logging
from warp_drive.utils.data_feed import DataFeed

class CUDAEnvironmentContext:
    """
    Environment Context class to manage APIs for the communication
    between EnvWrapper class and the Environment class
    """

    def __init__(self):
        self.cuda_data_manager = None
        self.cuda_function_manager = None
        self.cuda_step = None
        self.cuda_step_function_feed = None

    def initialize_step_function_context(
        self,
        cuda_data_manager,
        cuda_function_manager,
        cuda_step_function_feed,
        step_function_name,
    ):
        try:
            self.cuda_data_manager = cuda_data_manager
            self.cuda_function_manager = cuda_function_manager
            self.cuda_function_manager.initialize_functions([step_function_name])
            self.cuda_step = self.cuda_function_manager.get_function(step_function_name)
            self.cuda_step_function_feed = cuda_step_function_feed
            return True
        except Exception as err:
            logging.error(err)
            return False

    def get_data_dictionary(self):
        data_dict = DataFeed()
        return data_dict

    def get_tensor_dictionary(self):
        tensor_dict = DataFeed()
        return tensor_dict

    def get_reset_pool_dictionary(self):
        reset_pool_dict = DataFeed()
        return reset_pool_dict


import logging
import math
import chardet
import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
from torch.nn import init
from torch.nn.modules.normalization import LayerNorm

from third_party import transformer
import modules
import utils

class LSTM(nn.Module):
    def __init__(self, config):
        super(LSTM, self).__init__()
        self.config = config

        self.vocab_size = config["vocab_size"]
        self.hidden_size = config["hidden_size"]
        self.num_layers = config["num_layers"]
        self.dropout_rate = config["dropout_rate"]
        self.emb = nn.Embedding(self.vocab_size, self.hidden_size)
        self.rnn = nn.LSTM(self.hidden_size, self.hidden_size, num_layers=self.num_layers, dropout=self.dropout_rate, batch_first=True)
        self.dropout1 = nn.Dropout(self.dropout_rate)
        self.dropout2 = nn.Dropout(self.dropout_rate)
        self.output_affine = nn.Linear(self.hidden_size, self.vocab_size, bias=False)
        self.emb.weight = self.output_affine.weight

    def forward(self, ids, lengths=None):
        outputs = self.emb(ids)
        outputs = self.dropout1(outputs)
        outputs, (h, c) = self.rnn(outputs)
        outputs = self.dropout2(outputs)
        outputs = self.output_affine(outputs)
        return outputs

    def reset_parameters(self):
        for name, param in self.named_parameters():
            if 'weight_ih' in name:
                init.xavier_uniform_(param.data)
            elif 'weight_hh' in name:
                init.orthogonal_(param.data)
            elif 'bias' in name:
                param.data.fill_(0)
        init.uniform_(self.emb.weight, a=-0.01, b=0.01)

class TransformerLM(nn.Module):
    def __init__(self, config):
        super(TransformerLM, self).__init__()
        self.config = config

        self.vocab_size = config["vocab_size"]
        self.d_model = config["d_model"]
        self.nhead = config["nhead"]
        self.num_layers = config["num_layers"]
        self.dim_feedforward = config["dim_feedforward"]
        self.activation = config["activation"]
        self.dropout_rate = config["dropout_rate"]

        self.dropout = nn.Dropout(self.dropout_rate)
        self.scale = self.d_model ** 0.5
        self.pe = modules.PositionalEncoding(self.d_model)
        self.emb = nn.Embedding(self.vocab_size, self.d_model)
        encoder_layer = transformer.TransformerEncoderLayer(d_model=self.d_model,
                nhead=self.nhead, dim_feedforward=self.dim_feedforward,
                dropout=self.dropout_rate, activation=self.activation)
        self.transformer_encoder = transformer.TransformerEncoder(encoder_layer, self.num_layers)
        self.output_affine = nn.Linear(self.d_model, self.vocab_size, bias=False)
        self.emb.weight = self.output_affine.weight

    def forward(self, ids, lengths, return_atten=False):
        B, T = ids.shape

        key_padding_mask = utils.get_transformer_padding_byte_masks(
            B, T, lengths).to(ids.device)
        casual_masks = utils.get_transformer_casual_masks(T).to(ids.device)

        outputs = self.emb(ids) * self.scale
        outputs = self.pe(outputs)
        outputs = self.dropout(outputs)
        outputs = outputs.permute(1, 0, 2)

        outputs, self_atten_list = self.transformer_encoder(outputs,
                mask=casual_masks,
                src_key_padding_mask=key_padding_mask,
                return_atten=True)
        outputs = self.output_affine(outputs)
        if return_atten:
            return outputs, self_atten_list
        return outputs