from PyQt5 import QtWidgets
from PyQt5 import QtCore
from PyQt5.QtCore import pyqtSlot
from PyQt5.QtWidgets import QHeaderView
from db.models import *
from gui.widgets.custom_widgets import DialogWithDisablingOptions

class MainWidget(QtWidgets.QWidget):
    def __init__(self, parent, model):
        super().__init__(parent)
        self.layout = QtWidgets.QHBoxLayout(self)

        self.league_list = QtWidgets.QTableView(self)
        self.league_list.setModel(model)
        self.league_list.setSelectionBehavior(QtWidgets.QAbstractItemView.SelectRows)
        self.league_list.clearSelection()
        self.league_list.resizeColumnsToContents()
        self.league_list.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.layout.addWidget(self.league_list, 2)

        self.button_layout = QtWidgets.QVBoxLayout(self)

        self.new_league_button = QtWidgets.QPushButton("New league", self)
        self.button_layout.addWidget(self.new_league_button)
        self.new_league_button.clicked.connect(self.on_new_league)

        self.new_round_button = QtWidgets.QPushButton("New round", self)
        self.button_layout.addWidget(self.new_round_button)
        self.new_round_button.clicked.connect(self.on_new_round)

        self.del_button = QtWidgets.QPushButton("Delete league", self)
        self.button_layout.addWidget(self.del_button)
        self.del_button.clicked.connect(self.on_delete)

        self.results_overview_button = QtWidgets.QPushButton("Results overview", self)
        self.button_layout.addWidget(self.results_overview_button)
        self.results_overview_button.clicked.connect(self.on_load)

        self.quit_button = QtWidgets.QPushButton("Quit", self)
        self.button_layout.addWidget(self.quit_button)
        self.quit_button.clicked.connect(self.on_quit)

        self.button_layout.addStretch()
        self.jsolutions_label = QtWidgets.QLabel("JSolutions")
        self.jsolutions_label.setAlignment(QtCore.Qt.AlignCenter)
        self.button_layout.addWidget(self.jsolutions_label)
        self.button_layout.setAlignment(QtCore.Qt.AlignTop)
        self.button_layout.setAlignment(self.jsolutions_label, QtCore.Qt.AlignBottom)

        self.layout.addLayout(self.button_layout, 1)
        self.setLayout(self.layout)

    def get_selected_league(self):
        selected_indexes = self.league_list.selectedIndexes()
        if len(selected_indexes) != 2 or (selected_indexes[0].row() != selected_indexes[1].row()):
            return None
        else:
            return self.league_list.model().get_league(selected_indexes[0])

    @pyqtSlot()
    def on_quit(self):
        QtCore.QCoreApplication.instance().quit()

    @pyqtSlot()
    def on_load(self):
        league = self.get_selected_league()
        if league is None:
            return
        if league.max_round == 0 or league.max_round is None:
            QtWidgets.QMessageBox.warning(self, "Zero rounds error", "Zero rounds have been played in this league. "
                                                                     "Unable to show the results overview.")
        else:
            from gui.windows import LeagueOverviewWindow
            win = LeagueOverviewWindow(league, parent=self)
            win.show()

    @pyqtSlot()
    def on_delete(self):
        league = self.get_selected_league()
        if league is None:
            return
        league = League.get_by_name(league.name)
        reply = QtWidgets.QMessageBox.question(self, "Message", f"Are you sure you want to delete {league.name} league?",
                                                      QtWidgets.QMessageBox.Yes, QtWidgets.QMessageBox.No)
        if reply == QtWidgets.QMessageBox.Yes:
            Result.delete_all_from_league(league)
            league.delete_instance()
            self.league_list.model().refresh()
        else:
            return

    @pyqtSlot()
    def on_new_league(self):
        league_names = list(map(lambda x: x.name, League.get_all()))
        dialog = DialogWithDisablingOptions("New league", "Please enter valid league name:", league_names)
        if dialog.exec_():
            league = League.create(name=dialog.ret_str)
            league.save()
            self.league_list.model().refresh()

    @pyqtSlot()
    def on_new_round(self):
        league = self.get_selected_league()
        if league is None:
            return
        from gui.windows import InputWindow
        new_round_win = InputWindow(self, league)
        new_round_win.show()

    def refresh_leagues_overview(self):
        self.league_list.model().refresh()


# Import necessary libraries
import openai
import subprocess
import sys
import json
import html
import re
import ssl
import os
import math
import glob
import pprint
import nltk
import pdb
import requests
import time
import random
from PIL import Image, ImageDraw, ImageFont
from PIL import UnidentifiedImageError
if not nltk.data.find('tokenizers/punkt'):
    nltk.download('punkt', quiet=True)

sitelist = [
  { "subdomain": "alamo", "site_id": 29 },
  { "subdomain": "burlingame", "site_id": 30 },
  { "subdomain": "campbell", "site_id": 7 },
  { "subdomain": "castrovalley", "site_id": 25 },
  { "subdomain": "concord", "site_id": 31 },
  { "subdomain": "danville", "site_id": 9 },
  { "subdomain": "dublin", "site_id": 8 },
  { "subdomain": "hillsborough", "site_id": 12 },
  { "subdomain": "lafayette", "site_id": 13 },
  { "subdomain": "livermore", "site_id": 14 },
  { "subdomain": "orinda", "site_id": 34 },
  { "subdomain": "pittsburg", "site_id": 28 },
  { "subdomain": "pleasanthill", "site_id": 35 },
  { "subdomain": "sanramon", "site_id": 33 },
  { "subdomain": "walnutcreek", "site_id": 32 }
]

def get_site_id(subdomain):
  for site in sitelist:
    if site["subdomain"] == subdomain:
      return site["site_id"]
  return None

# Get the first command line argument
location = sys.argv[1]
sku = sys.argv[2]

# Initialize an empty dictionary for credentials
credentials = {}

# Define the file path to the credentials file
creds_file_path = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),  # Get the directory of the current file
    "../creds2.txt"  # Append the relative path to the credentials file
)
if os.path.exists('product.json'):
    os.remove('product.json')

class Location:
    def __init__(self, website, user, city, phone, consumer_key, consumer_secret, api_key):
        self.website = website
        self.user = user
        self.city = city
        self.phone = phone
        self.consumer_key = consumer_key
        self.consumer_secret = consumer_secret
        self.api_key = api_key  # Here's the new attribute

def scp_file_to_remote(local_file, remote_file):
    try:
        # Run SCP command
        subprocess.Popen(["scp", local_file, remote_file])
        print("File transfer initiated.")

    except Exception as e:
        print("Error while copying the file:", e)

def download_image(url, filename):
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()
        with open(filename, "wb") as file:
            for chunk in response.iter_content(chunk_size=8192):
                file.write(chunk)
        print(f"Image downloaded successfully: {filename}")
    except requests.exceptions.RequestException as e:
        print(f"Error downloading image: {str(e)}")

def add_watermark_and_save(image_path, watermark_text, output_path):
    try:
        # Open the image
        image = Image.open(image_path).convert("RGBA")

        # Define the watermark text and font style
        font = ImageFont.truetype("font.ttf", 40)

        # Create a transparent overlay and draw the watermark text
        overlay = Image.new("RGBA", image.size, (0, 0, 0, 0))
        draw = ImageDraw.Draw(overlay)
        text_width, text_height = draw.textbbox((0, 0), watermark_text, font=font)[:2]
        # position = ((image.width - text_width) // 2, (image.height - text_height) // 2)
        position = (image.width - text_width - 10, image.height - text_height - 10) # Position the watermark in the lower right corner

        draw.text(position, watermark_text, font=font, fill=(128, 128, 128, 128))

        # Composite the image and watermark overlay
        watermarked = Image.alpha_composite(image, overlay)

        # Save the watermarked image with the specified output path
        watermarked.save(output_path)
        print(f"Watermarked image saved as {output_path}")
    except Exception as e:
        print(f"Error: {str(e)}")

def makeunique(new_unique_product_name):
    ai_response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": "You are a helpful budtender who knows all about the cannabis industry.",
            },
            {
                "role": "user",
                "content": f"Use this product name '{new_unique_product_name}'. Use this phrase to come up with a slightly different name that means the same thing."
            f"Come up with a new name that is max 70 chars long and will rank well with regard to SEO. If there is a mention of price. Change it to some other descriptive language instead."
            },
        ]
    )

def generate_new_product_name(sku):
    ai_response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": "You are a helpful budtender who knows all about the cannabis industry.",
            },
            {
                "role": "user",
                "content": f"Use this product slug '{product['slug']}' to rewrite the product title. The slug contains words separated by a -."
            f"Use them to come up with a new name that is max 70 chars long and will rank well with regard to SEO. If there is a mention of price. Change it to some other descriptive language. Dont put spaces in the names. Use underscores to separate words."
            },
        ]
    )

    new_product_name = ai_response['choices'][0]['message']['content'].strip()
    new_product_name = html.unescape(re.sub('<.*?>', '', new_product_name))

    return new_product_name

def generate_new_image_name(image_name):
    ai_response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": "You are a creative AI assistant and California Budtender for a delivery service.",
            },
            {
                "role": "user",
                "content": f"I have an image with the name '{image_name}'. Please suggest a new name for the image that does not use dates or times in the name. Limit the name to 70 characters. Dont put spaces in the names. Use underscores to separate words."
            },
        ]
    )

    new_image_name = ai_response['choices'][0]['message']['content'].strip()
    new_image_name = html.unescape(re.sub('<.*?>', '', new_image_name))

    return new_image_name

def remove_keys(images_data):
    keys_to_remove = [
    'date_created',
    'date_created_gmt',
    'date_modified',
    'date_modified_gmt',
    'id',
    'alt'
    ]
    new_images_data = []
    for index, image_data in enumerate(images_data):
        if index < 4:
            new_image_data = {key: value for key, value in image_data.items() if key not in keys_to_remove}
        else:
            new_image_data = {}
        new_images_data.append(new_image_data)
    return new_images_data

def generate(new_pics_prompt):
    res = openai.Image.create(
        prompt=new_pics_prompt,
        n=1,
        size="256x256",
    )
    return res["data"][0]["url"]

locations = []

# Open the credentials file
with open(creds_file_path) as f:
    # Initialize variables for parsing the file
    website = None
    user = None
    city = None
    phone = None
    consumer_key = None
    consumer_secret = None
    openai.api_key = None
    for line in f:
        line = line.strip()  # Remove trailing and leading whitespace
        if line.startswith("[") and line.endswith("]"):
            if website and user and city and phone and consumer_key and consumer_secret and openai.api_key:
                locations.append(Location(website, user, city, phone, consumer_key, consumer_secret, openai.api_key))
            website = line[1:-1].lstrip()  # Remove the brackets and any leading whitespace
            user = None
            city = None
            phone = None
            consumer_key = None
            consumer_secret = None
            openai.api_key = None
        elif website and " = " in line:
            key, value = line.split(" = ")
            if key == "user":
                user = value
            elif key == "city":
                city = value
            elif key == "phone":
                phone = value
            elif key.lower().endswith("_consumer_key"):
                consumer_key = value
            elif key.lower().endswith("_consumer_secret"):
                consumer_secret = value
            elif key == "openai.api_key":
                openai.api_key = value
                aikey = value
            elif key == "website":
                website = value

    locations.append(
        Location(website, user, city, phone, consumer_key,
                 consumer_secret, openai.api_key)
    )

#fetches the first product dataset to be edited and pushed to the other sites.
for locationa in locations[:1]:
    base_url = "https://" + locationa.website + "/wp-json/wc/v3/products"
    consumer_key = locationa.website + "_consumer_key:" + locationa.consumer_key
    consumer_secret = locationa.website + "_consumer_secret:" + locationa.consumer_secret
    city = locationa.city
    phone = locationa.phone
    website = locationa.website
    aikey = openai.api_key

    auth = (
        locationa.consumer_key,
        locationa.consumer_secret,
    )

    response = requests.get(f'{base_url}', auth=auth, params={'sku': sku})
    response.raise_for_status()

    product = response.json()[0]
    source_product = product
    source_product['images'] = remove_keys(source_product['images'])
    source_images = source_product['images'][:4]
    imagecounter = 0
    for item in source_images:
        imagecounter = imagecounter + 1
        print("Image:",imagecounter)
        #source_product_name = product['name'].strip()
        item['src'] = item['src'].replace("/29/","/30/")
        item['src'] = item['src'].replace("alamo","burlingame")
        #imgcnt = 0
        #pprint.pprint(source_images)
        #source_image_url = item['src']
        # for item in source_images:
        # source_product_name = product['name'].strip()
        # print("Source Product\n",source_product_name)
        # print(website, aikey)
        # print("Source Images")
        # imgcnt = 0
        # pprint.pprint(source_images)
        # source_image_url = item['src']

# new_product_name = generate_new_product_name(sku)
# print("New name suggestion:", new_product_name)

seq = 0
#fetches all but the first product and applies the updated first site product details.
print("Destination Products\n")
for locationb in locations[1:]:
    seq = seq + 1
    base_url = "https://" + locationb.website + "/wp-json/wc/v3/products"
    consumer_key = locationb.website + "_consumer_key:" + locationb.consumer_key
    consumer_secret = locationb.website + "_consumer_secret:" + locationb.consumer_secret
    city = locationb.city
    city = city.replace('"', '')
    phone = locationb.phone
    phone = phone.replace(' ', '').replace('-', '').replace('"', '').replace('(', '').replace(')', '')

    website = locationb.website
    aikey = openai.api_key

    auth = (
        locationb.consumer_key,
        locationb.consumer_secret,
    )

    response = requests.get(f'{base_url}', auth=auth, params={'sku': sku})
    response.raise_for_status()

    product = response.json()[0]
    #source_product = product
    source_product['images'] = remove_keys(source_product['images'])
    product['images'] = source_product['images']
    msgg = "#" + str(seq) + " " + str(sku)
    print(msgg)
    subdomain = website.split('.')[0]
    print("Domain: ", subdomain)
    site_id = get_site_id(subdomain)
    print("Site ID:", site_id)
    print(city, "Doap")
    print(city, " Ca ", phone)
    print("Sku: ", sku)
   # First AI call: generate new product name
    product['name'] = generate_new_product_name(sku).replace('"','').replace('"','').replace("'","").replace(" ","_").replace("(","").replace(")","").replace(",","").replace("$","")
    print("New dest product name: ", product['name'])
    print("New Images")
    imgcnt = 0
    for item in source_images:
        imgcnt = imgcnt + 1
        itemname = item['name'].replace('-',' ').capitalize()
        print("Image #", imgcnt)
        itemname = item['name'].replace('-',' ').capitalize()
        # print("Image #", imgcnt)
        new_unique_product_name = generate_new_image_name(product['name']).replace('"','').replace('"','').replace("'","").replace("!","").replace("(","").replace(")","").replace(",","").replace("→","")
        new_unique_file_name = new_unique_product_name
        item['name'] = new_unique_product_name
        # print(item['name'], " : ", item['src'])
        source_image_url = item['src']
        source_image_filename = os.path.basename(source_image_url)
        new_unique_file_name = new_unique_file_name + ".png"
        download_image(source_image_url, source_image_filename)
        print("Source image url: ", source_image_url)
        replaced_url = source_image_url.replace("https://alamo.", "/var/www/")
        stripped_path = "/".join(replaced_url.split("/")[:-1])
        print("Orig file path: ", stripped_path)

        new_path = stripped_path.split("/")
        new_path[7] = str(site_id)
        new_path = "/".join(new_path)

        print("New remote file path: ", new_path)
        #item['src'] = "https://" + subdomain + ".doap.com/" + stripped_path + "/" + new_unique_file_name
        item['src'] = "https://" + subdomain + ".doap.com/" + stripped_path + "/" + new_unique_file_name
        item['src'] = item['src'].replace("/var/www/doap.com/","")
        watermark_text = city + " Doap " + phone
        add_watermark_and_save(source_image_filename, watermark_text, new_unique_file_name)
        local_file = '/Users/dmenache/Nextcloud/Projects/doap-api/ai_product_updater/' + new_unique_file_name
        remote_server = 'dmenache@debian.doap.com'
        testpath = stripped_path.replace("https://burlingame.","/var/www/")
        remote_file = f'{remote_server}:{testpath}/{new_unique_file_name}'
        scp_file_to_remote(local_file, remote_file)
        pdb.set_trace()
        #pprint.pprint(item)
        #pprint.pprint(source_images)
        product['images'] = source_images
        #pprint.pprint(product)
        # pprint.pprint(product)
        for image in product['images']:
            image['src'] = image['src'].replace('https://burlingame.doap.com/https://burlingame.doap.com/', 'https://burlingame.doap.com/')
        print("product[images]",product['images'])
        print("source_images",source_images)
        print("product[images]",product['images'])
    break
pprint.pprint(product)
pdb.set_trace()
update_url = f'{base_url}/{product["id"]}'
update_response = requests.put(update_url, json=product, auth=auth)
update_response.raise_for_status()


# -*- coding: utf-8 -*-

#importando as bibliotecas

from matplotlib.pyplot import text
import yfinance as yf
import pandas as pd
import numpy as np
import os.path
import telegram

pd.options.mode.chained_assignment = None

#escolher uma ação
wege = yf.Ticker('WEGE3.SA')

#escolher inteervalo de dados
wege_dia = wege.history(period='id', interval='5m')

#pegar preço de fechamento
wege_dia = wege_dia.Close

#transformando em dataframe
df_wege_dia = pd.DataFrame(wege_dia)

#reset index
df_wege_dia.reset_index(inplace=True)

#pegar o ultimo valor negociado
wege_dia_ultimo_preco = df_wege_dia.tail(1)

#renomear as colunas
wege_dia_ultimo_preco.rename(columns={'Datetime':'data_pregao', 'Close':'preco_fechamento'}, inplace=True)

#Ajustar a data
wege_dia_ultimo_preco['data_pregao']=pd.to_datetime(wege_dia_ultimo_preco['data_pregao'], format='%Y-%m-%d')

#Usar o data frame historico e pegar apenas o preço de fechamento e data pregão

if os.path.isfile('wege.csv'):
    df_wege = pd.read_csv('wege.csv', delimiter=';')
else:
    df = pd.read_csv('all_bovesta.csv', delimiter=';') #colocar aqui o seu arquivo do bovespa
    df_wege = df[df['silga_acao']=='WEGE3']
    df_wege = df_wege[['data_pregao', 'preco_fechamento']]

#Ajustar a data
df_wege['data_pregao']=pd.to_datetime(df_wege['data_pregao'], format='%Y-%m-%d')

#Retirar a ultima data que queremos calcular
df_remove = df_wege.loc[(df_wege['data_pregao'] == pd.to_datetime('today').normalize())]

df_wege = df_wege.drop(df_wege.index)

#append data atual
df_wege_total = df_wege.append(wege_dia_ultimo_preco)

#Ajuste data atual

df_wege_total['data_pregao']=pd.to_datetime(df_wege_total['data_pregao'], utc=True).dt.date

df_wege_total.to_csv('wege.csv', sep=';', index=False)

#Calcular MACD

rapidaMME=df_wege_total.preco_fechamento.ewm(span=12).mean()

lentaMME = df_wege_total.preco_fechamento.ewm(span=26).mean()

MACD= rapidaMME - lentaMME

sinal=MACD.ewm(span=9).mean()

df_wege_total['MACD'] = MACD
df_wege_total['sinal'] = sinal

#Ajuste de indx e retirar o campo data pregão

df_wege_total = df_wege_total.set_index(pd.DatetimeIndex(df_wege_total['data_pregao'].values))
df_wege_total = df_wege_total.drop('data_pregao',1)

# Criar codigo para verificar a compra ou a venda

df_wege_total['flag']=''
df_wege_total['preco_compra']=np.nan
df_wege_total['preco_venda']=np.nan

for i in range(1, len(df_wege_total.sinal)):
    if df_wege_total['MACD'][i] > df_wege_total['sinal'][i]:
        if df_wege_total['flag'][i-1] == 'c':
            df_wege_total['flag'][i]='C'
        else:
            df_wege_total['flag'][i]='C'
            df_wege_total['preco_compra'][i] = df_wege_total['preco_fechamento'][i]

    elif df_wege_total['MACD'][i] < df_wege_total['sinal'][i]:
        if df_wege_total['flag'][i-1] =='V':
            df_wege_total['flag'][i]='V'
        else:
            df_wege_total['flag'][i]='V'
            df_wege_total['preco_venda'][i] = df_wege_total['preco_fechamento'][i]

#Verifica os 2 ultimos dias
hoje = df_wege_total.flag[-1]
ontem = df_wege_total.flag[-2]

flag= hoje
preco_fechamento = round(df_wege_total.preco_fechamento.tail(1)[-1],2)

print(flag, preco_fechamento)

my_token = '1840232813:AAHxoVmcDWHK3jAxiTWsMqTsiw9vTHaICpY'
chat_id = '-476980685'

def envia_mensagem(msg, chat_id, token=my_token):
    bot=telegram.Bot(token = token)
    bot.SendMessage(chat_id = chat_id, text=msg)

msg = f'WEGE3 (WEGE), {flag} preço de fechamento: {preco_fechamento}'

if ontem == hoje:
    envia_mensagem(msg, chat_id, my_token)


from functools import cmp_to_key

def custom_split(s):
    if s == '': return []
    cnt = 0
    res = []
    last_comma = -1
    brackets = 0
    while (cnt < len(s)):
        if s[cnt] == '[': brackets += 1
        if s[cnt] == ']': brackets -= 1
        if s[cnt] == ',':
            if brackets == 0:
                res.append(s[last_comma+1:cnt])
                last_comma = cnt
        cnt += 1
    res.append(s[last_comma+1:])
    return res

def compare(l1, l2):
    len1 = len(l1)
    len2 = len(l2)
    cnt = 0
    while cnt < len1 and cnt < len2:
        len1 = len(l1)
        len2 = len(l2)
        if l1[cnt][0] != '[' and l2[cnt][0] != '[': # both values are integers
            if int(l1[cnt]) < int(l2[cnt]): return 1
            if int(l1[cnt]) > int(l2[cnt]): return -1
        else:
            if l1[cnt][0] == '[':
                l1_new = custom_split(l1[cnt][1:-1])
            else:
                temp = l1[cnt]
                l1_new = []
                l1_new.append(temp)
            if l2[cnt][0] == '[':
                l2_new = custom_split(l2[cnt][1:-1])
            else:
                temp = l2[cnt]
                l2_new = []
                l2_new.append(temp)
            compare_res = compare(l1_new,l2_new)
            if compare_res == 1: return 1
            if compare_res == -1: return -1
        cnt += 1
    if not cnt < len1 and cnt < len2: return 1
    if cnt < len1 and not cnt < len2: return -1
    return 0

def custom_sort(a,b):
    return True if compare(a,b) == 1 else False

packets = []

DIVIDER_PACKET_1 = ["[[2]]"]
DIVIDER_PACKET_2 = ["[[6]]"]

with open('input.txt') as f:
    lines = [line.rstrip('\n') for line in f]
    cnt = 1
    sum = 0
    while 3*(cnt-1) < len(lines):
        packet_one = lines[3*(cnt-1)]
        packet_two = lines[3*(cnt-1)+1]
        list_one = custom_split(packet_one[1:-1])
        list_two = custom_split(packet_two[1:-1])
        packets.append(list_one)
        packets.append(list_two)
        cnt += 1
    packets.append(DIVIDER_PACKET_1)
    packets.append(DIVIDER_PACKET_2)

packets = sorted(packets, key=cmp_to_key(compare), reverse=True)
cnt = 1
decoder_key = 1
for packet in packets:
    print(packet)
    if packet == DIVIDER_PACKET_1: decoder_key *= cnt
    if packet == DIVIDER_PACKET_2: decoder_key *= cnt
    cnt += 1
print("decoder_key: ", decoder_key)


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision

import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt

import numpy as np

import math
import yaml
import pickle
import pprint
import os
import logging
import sys

import data_loaders
import nets
import losses
import utils
import setup

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

'''
#SECTION VAE Helpers
'''

def sample_n_frames(init_frames, ts, dt, ae, mu, sigma, n_generate=64):

    # function to generate the next n frames conditioned on the input
    with torch.no_grad():

        # Get the latent variables
        q_mu, q_sigma, det = ae.encode(init_frames)
        _, z, _ = ae.get_increments(q_mu, q_sigma)

        z_samples = torch.zeros((n_generate, z.shape[1]))
        z = z[-1,:].unsqueeze(0)

        # sample in z according to the learned SDE
        for i in range(n_generate):
            z_n = ae.get_next_z(z, ts[-1].unsqueeze(0) + i*dt, dt, mu, sigma)
            z_samples[i,:] = z_n.clone()
            z = z_n

        global savepath

        plots_list = [z_samples.detach().cpu().numpy()]
        plot_titles = ['Latent Traj']
        utils.plot_subplots(plots_list, plot_titles, os.path.join(savepath, 'latent_traj.png'), plot_type='plot', axis=True)

        conditional_frame = init_frames[0].unsqueeze(0).repeat(z_samples.size(0),1,1,1)

        if det is not None:
            in_z = torch.cat((z_samples.to(device), det[-1].repeat(z_samples.shape[0],1)), dim = 1)
        else:
            in_z = z_samples.to(device)

        frames = ae.decode(in_z, x=conditional_frame)

        return z_samples, frames

def plot_images(ae, mu, sigma, frames, frames_hat, dt, ts, l2_small):
    # plot latent trajectories
    # plot next frame reconstructions
    z_samples, sampled_frames = sample_n_frames(frames[:2],
            ts[:2],
            dt,
            ae.eval(),
            mu,
            sigma)
    _, sampled_frames2 = sample_n_frames(frames[:2],
            ts[:2],
            dt,
            ae.eval(),
            mu,
            sigma*2)

    # create the image grids
    im_grid_hat_single = torchvision.utils.make_grid(frames_hat[:64].detach().cpu(), pad_value=1, normalize=True)
    im_grid_hat = torchvision.utils.make_grid(sampled_frames[:64].detach().cpu(), pad_value=1, normalize=True)
    im_grid     = torchvision.utils.make_grid(frames[:64].detach().cpu(), pad_value=1, normalize=True)

    odd_rows = []

    for row in range(4):
        odd_rows.append(frames[row*8:(row+1)*8])
        odd_rows.append(sampled_frames[row*8:(row+1)*8])

    comp_grid = torchvision.utils.make_grid(torch.cat(odd_rows), pad_value=1, normalize=True)
    plots_list = [comp_grid.cpu().numpy().transpose((1,2,0))]
    plot_titles = ['Comparison']
    utils.plot_subplots(plots_list, plot_titles, os.path.join(savepath, 'train_comparison.png'))

    plots_list = [im_grid_hat.numpy().transpose((1,2,0))]
    plot_titles = ['Sampled (trajectory)']
    utils.plot_subplots(plots_list, plot_titles, os.path.join(savepath, 'train_sample_traj.png'))

    # save the images
    plots_list = [im_grid.numpy().transpose((1,2,0)),im_grid_hat_single.numpy().transpose((1,2,0))]
    plot_titles = ['Original','Sampled (single)']

    if l2_small:
        utils.plot_subplots(plots_list, plot_titles, os.path.join(savepath, 'train_sample_best.png'))
    else:
        utils.plot_subplots(plots_list, plot_titles, os.path.join(savepath, 'train_sample.png'))

    # save the movies
    utils.save_gif(sampled_frames.detach().cpu(), os.path.join(savepath, 'movies/train_sample_traj.gif'))
    utils.save_gif(sampled_frames2.detach().cpu(), os.path.join(savepath, 'movies/train_sample_traj2.gif'))

def save_nets(ae, mu, sigma, suffix):
    '''
    Routine to save the current state of our network
    ae : autoencoder network
    mu : latent drift network
    sigma : latent diffusion network
    suffix : str that defines how we want to save the network
    '''

    # save all the networks
    torch.save(ae.state_dict(), os.path.join(savepath,'saved_nets/ae_{}.pth'.format(suffix)))

    if type(mu) == nets.MLP or type(mu) == nets.Lin or type(mu) == nets.Well:
        torch.save(mu.state_dict(), os.path.join(savepath,'saved_nets/mu_{}.pth'.format(suffix)))
    else:
        with open(os.path.join(savepath,'mu_{}.pkl'.format(suffix)),'wb') as f:
            pickle.dump(mu, f)

    if type(sigma) == nets.MLP:
        torch.save(sigma.state_dict(), os.path.join(savepath,'saved_nets/sigma_{}.pth'.format(suffix)))
    else:
        with open(os.path.join(savepath,'sigma_{}.pkl'.format(suffix)),'wb') as f:
            pickle.dump(sigma, f)

def train(ae,
        mu,
        sigma,
        dt,
        train_data,
        val_data,
        optimizer,
        scheduler,
        n_epochs,
        data_params,
        **kwargs):

    '''
    The main training routine:

    ae : neural network (torch.Module subclass) that represents our autoencoder
    mu : network or parameter that describes the latent drift
    sigma : network or parameter that describes the latent diffusion
    dt : time step
    train_data : dataloader with the training data
    val_data : dataloader with validation data
    optimizer : optimization algorithm torch.optim
    scheduler : lr decay schedule
    n_epochs  : number of epochs to run
    data_params : parameters associated with the dataset

    returns statistics with respect to training
    '''

    global savepath
    global loss_type

    train_dataset = train_data.dataset.dataset
    val_dataset   = val_data.dataset.dataset
    try:
        inner_num = data_params['inner_iter']
    except:
        inner_num = 1

    if n_epochs > 1000:
        reserve_epoch = 499
    else:
        reserve_epoch = 49

    # plotting parameters
    l2_small   = True
    l2_small_valid = True

    losses_train = []
    losses_valid = []

    try:
        plot_freq = data_params['plot_freq']
    except KeyError:
        plot_freq = 50

    try:
        plot_train = data_params['plot_train']
    except KeyError:
        plot_train = True

    # setup the stats dict
    stats = {'kl': np.Inf,
            'l2' : np.Inf,
            'l2_valid': np.Inf,
            'kl_valid': np.Inf,
            'mu_mse': 0,
            'mu_mse_valid': 0,
            'mu_rel': 0,
            'mu_rel_valid': 0,
            'sde_mse': 0,
            'sde_mse_valid': 0,
            'sde_rel': 0,
            'sde_rel_valid': 0,
            'val_cond_met': False}

    for epoch in range(n_epochs):

        ae.train()
        mu.train()
        #sigma.train()

        for idx, (frames, ts) in enumerate(train_data):

            # save a gif of the data
            if len(frames.shape) > 2:
                if idx == 0 and epoch == 0:
                    utils.save_gif(frames.detach().cpu(), os.path.join(savepath, 'orig_data.gif'))

            # transfer the data to the device
            # the rest is boilerplate
            frames = frames.float().to(device)
            ts     = ts.float().to(device)

            for _ in range(inner_num):

                optimizer.zero_grad()

                kl_loss, l2_loss,\
                        frames_hat, mu_hat, q_mu, sigma_hat_full, q_sigma_full, inc, z = ae.step(frames, ts, dt, mu, sigma)

                kl_loss1, l2_loss1,\
                        _, _, _, _, _, _, _ = ae.step(frames, ts, dt, mu, sigma, plus_one=True)

                sigma.data = sigma / sigma.norm(2) * torch.ones(z.shape[1]).norm(2)

                loss = kl_loss + kl_loss1 + l2_loss + l2_loss1 + 20*sigma.norm(1)

                losses_train.append((kl_loss.item(), l2_loss.item()))

                loss.backward()
                optimizer.step()

                # And that's the end of the train routine

            '''
            PLOT SECTION
            This is still quite messy and needs to be refactored,
            but this is all visualization calls
            '''

            if kl_loss < stats['kl']:

                stats['kl'] = kl_loss.item()
                stats['mu'] = mu_hat.mean().item()

            if plot_train and (epoch % plot_freq) == 0 and idx == 0:

                if l2_loss < stats['l2']:
                    l2_small = True
                    stats['l2'] = l2_loss.item()
                else:
                    l2_small = False

                if len(frames.shape) > 2:
                    plot_images(ae, mu, sigma, frames, frames_hat, dt, ts, l2_small)

                # plot mu hat
                mu_hat_samples, hat_domain = utils.plot_mu_hat(mu, sigma, q_mu, ts, os.path.join(savepath, 'mu_hat_plot.png'))

                if len(frames.shape) < 3:
                    plots = [frames.cpu(), frames_hat.detach().cpu()]
                    names = ['Original', 'Sampled']
                    utils.plot_subplots(plots, names, os.path.join(savepath, 'train_recon.png'), plot_type='plot', axis=True)
                    _, sampled_frames = sample_n_frames(frames[:2], ts[:2], dt, ae, mu, sigma, n_generate=1000)
                    plots = [frames.cpu(), sampled_frames.detach().cpu()]
                    names = ['Original', 'Sampled']
                    utils.plot_subplots(plots, names, os.path.join(savepath, 'train_sampled.png'), plot_type='plot', axis=True)

                    if frames.shape[1] == 1:

                        with torch.no_grad():

                            inx = torch.linspace(frames.min().item(), frames.max().item()).unsqueeze(1)
                            oned_enc = ae.encode(inx.cuda(0))[0].detach().data.clone().cpu()

                            enc_scale = ( inx.log() / oned_enc ).mean()
                            enc_shift = (inx.log() - enc_scale * oned_enc).mean()

                            plt.plot(inx.detach().cpu(), enc_scale * oned_enc.cpu(), label='encoder')
                            plt.plot(inx.detach().cpu(), inx.log().detach().cpu(),label='log')
                            plt.legend()

                            plt.savefig(os.path.join(savepath, 'encoder_plot.pdf'))

                            plt.close('all')

                '''
                AFFINE TRANSFORM SECTION
                '''

                # calculate the affine map between xt and z
                current_run = train_dataset.xt_orig[idx*z.shape[0]:(idx+1)*z.shape[0]]
                scale = (train_dataset.xt_orig.max() - train_dataset.xt_orig.min())

                q_mu = q_mu[:, :train_dataset.xt_orig.shape[1]]
                z    = z[:, :train_dataset.xt_orig.shape[1]]

                if not 'stocks' in savepath:
                    # if this is the stocks dataset, don't compute the scaling since there is none

                    if data_params['affine']:

                        transformed_xt, Q, b, sde_mse, sde_rel = utils.calc_affine(
                                current_run,
                                z.detach().cpu().numpy(),
                                savepath,
                                affine=data_params['affine'])

                        if z.shape[1] == mu_hat.shape[1]:

                            mu_residuals, mu_rel, mu_crlb = utils.compare_mu2(
                                    mu,
                                    q_mu,
                                    ts,
                                    Q,
                                    b,
                                    dt,
                                    train_dataset,
                                    os.path.join(savepath,'mu_comp_scaled.png'),
                                    affine=data_params['affine'],
                                    loss_type=loss_type)
                        else:
                            mu_residuals = torch.Tensor([np.NaN]).numpy()
                            mu_crlb = torch.Tensor([np.NaN]).numpy()
                            mu_rel  = torch.Tensor([np.NaN]).numpy()

                    else:

                        q_max = q_mu.max()
                        q_min = q_mu.min()

                        if loss_type == 'exact':
                            q_scaled = ((q_mu - q_min ) / (q_max - q_min) * (scale) ).detach().cpu().numpy()
                            #q_scaled = q_mu.detach().cpu().numpy()  / np.sqrt(scale)
                        else:
                            q_scaled = q_mu.detach().cpu().numpy() / scale

                        transformed_xt, Q, b, sde_mse, sde_rel = utils.calc_affine(
                                current_run,
                                q_scaled,
                                #z.detach().cpu().numpy() / scale,
                                savepath,
                                affine=data_params['affine'])
                        if z.shape[1] == mu_hat.shape[1]:
                            mu_residuals, mu_rel, mu_crlb = utils.compare_mu2(
                                    mu,
                                    q_mu,
                                    ts,
                                    Q,
                                    b,
                                    dt,
                                    train_dataset,
                                    os.path.join(savepath,'mu_comp_scaled.png'),
                                    affine=data_params['affine'],
                                    loss_type=loss_type)
                        else:
                            mu_residuals = torch.Tensor([np.NaN]).numpy()
                            mu_crlb = torch.Tensor([np.NaN]).numpy()
                            mu_rel  = torch.Tensor([np.NaN]).numpy()
                    stats['sde_mse'] = sde_mse.copy()
                    stats['sde_rel'] = sde_rel.copy()

                    # compare the estimated mu to the true mu with the affine map q
                    stats['mu_mse']  = mu_residuals.copy()
                    stats['mu_rel']  = mu_rel.copy()
                    stats['mu_crlb'] = mu_crlb.copy()
                else:
                    mu_residuals = torch.Tensor([np.NaN]).numpy()
                    mu_crlb = torch.Tensor([np.NaN]).numpy()
                    mu_rel  = torch.Tensor([np.NaN]).numpy()

                    stats['sde_mse'] = torch.Tensor([np.NaN]).numpy()
                    stats['sde_rel'] = torch.Tensor([np.NaN]).numpy()

                    # compare the estimated mu to the true mu with the affine map Q
                    stats['mu_mse']  = torch.Tensor([np.NaN]).numpy()
                    stats['mu_rel']  = torch.Tensor([np.NaN]).numpy()
                    stats['mu_crlb'] = torch.Tensor([np.NaN]).numpy()

            # plot and print
            print('Epoch {} iter {}'.format(epoch, idx))
            print('L2 loss {}'.format(l2_loss.item()))
            print('KL loss {}'.format(kl_loss.item()))

            plots_list = [(q_mu[1:]-q_mu[:-1]).detach().cpu().numpy(), mu_hat.detach().cpu().numpy()]
            plot_titles = ['q_mu', 'mu_hat']
            utils.plot_subplots(plots_list, plot_titles, os.path.join(savepath, 'mu_comp.png'), plot_type='plot', axis=True)

        if scheduler:
            if type(scheduler) == torch.optim.lr_scheduler.ReduceLROnPlateau:
                scheduler.step(l2_loss)
            else:
                scheduler.step()

        if (epoch % plot_freq) == 0:
            # save all the networks
            #if len(frames.shape) < 3:
            #    utils.plot_mu_hat(mu, None, z, ts, os.path.join(savepath, 'mu_hat_est.pdf'))
            save_nets(ae, mu, sigma, 'latest')
            with open(os.path.join(savepath, 'latent.pkl'), 'wb') as f:
                #lat_d = {'q_mu' : q_mu.detach().cpu().numpy(), 'ts' : ts, 'xt_orig' : dataset.xt_orig}
                lat_d = {'q_mu' : transformed_xt, 'ts' : ts, 'xt_orig' : train_dataset.xt_orig}
                pickle.dump(lat_d, f)

            if type(sigma) == nn.Parameter:
                print('Update sigma_hat')
                print(sigma)
                stats['sigma_hat'] = (sigma.sort(descending=True)[0]).detach().cpu().numpy()

        if (epoch % plot_freq) == 0 and plot_train:

            '''
            EVAL
            '''

            # with our validataion data, see how well we're predicting
            with torch.no_grad():
                ae.eval()

                # first, compute how well we predict the next step on the validation data
                for idxt, (frames_test, ts_test) in enumerate(val_data):

                    frames_test = frames_test.float().to(device)
                    ts_test     = ts_test.float().to(device)

                    kl_loss_test, l2_loss_test,\
                            frames_hat_test, mu_hat_test, q_mu_test, sigma_hat_full, q_sigma_full,  \
                            inc_test, z_test = ae.step(frames_test,
                                    ts_test,
                                    dt,
                                    mu,
                                    sigma)

                    losses_valid.append((kl_loss_test.item(), l2_loss_test.item()))

                q_mu_test = q_mu_test[:,:train_dataset.xt_orig.shape[1]]
                z_test    = z_test[:,:train_dataset.xt_orig.shape[1]]

                if len(frames_hat_test.shape) < 3 and l2_loss_test < stats['l2_valid']:
                    stats['l2_valid'] = l2_loss_test.item()
                    stats['kl_valid'] = kl_loss_test.item()
                    l2_small_valid = True
                    stats['val_cond_met'] = True
                    save_nets(ae, mu, sigma, 'best_val')
                    plots = [frames_test.cpu(), frames_hat_test.detach().cpu()]
                    names = ['Original', 'Sampled']
                    utils.plot_subplots(plots, names, os.path.join(savepath, 'valid_recon.png'), plot_type='plot')

                # if the l2 and kl are sufficiently small, save these as our current best networks
                if ((l2_loss_test < stats['l2_valid'] and epoch > reserve_epoch) or ('dna' in savepath)) and ('z={}'.format(train_dataset.xt_orig.shape[1]) in savepath):

                    stats['val_cond_met'] = True

                    #stats['l2_valid'] = kl_loss_test.item()*l2_loss_test.item()
                    stats['l2_valid'] = l2_loss_test.item()
                    stats['kl_valid'] = kl_loss_test.item()
                    l2_small_valid = True

                    save_nets(ae, mu, sigma, 'best_val')

                    # Compute the mapping over the training data since we want to see the fit within the whole time series
                    frames = torch.Tensor(train_dataset.frames).float().to(device)[:z.shape[0]]
                    ts     = torch.Tensor(train_dataset.ts).float().to(device)[:z.shape[0]]

                    kl_loss, l2_loss,\
                            frames_hat, mu_hat, q_mu, sigma_hat_full, q_sigma_full, inc, z = ae.step(frames,
                                    ts,
                                    dt,
                                    mu,
                                    sigma)

                    if 'gbm' in savepath:
                        gbm = True
                    else:
                        gbm = False

                    # compare the estimated mu to the true mu with the affine map Q
                    current_run = train_dataset.xt_orig[:z.shape[0]]
                    scale = train_dataset.xt_orig.max() - train_dataset.xt_orig.min()

                    if gbm:
                        scale = (np.log(train_dataset.xt_orig[:]).max() - np.log(train_dataset.xt_orig[:]).min())

                    if len(frames.shape) < 3:
                        plots = [frames.cpu(), frames_hat.detach().cpu()]
                        names = ['Original', 'Sampled']
                        utils.plot_subplots(plots, names, os.path.join(savepath, 'valid_recon.png'), plot_type='plot')
                        continue

                    if data_params['affine']:
                        transformed_xt, Q, b, sde_mse, sde_rel = utils.calc_affine(current_run,
                                z.detach().cpu().numpy(),
                                savepath,
                                affine=data_params['affine'])

                        mu_mse, mu_rel, mu_crlb = utils.compare_mu2(mu,
                                q_mu,
                                ts,
                                Q,
                                b,
                                dt,
                                train_dataset,
                                os.path.join(savepath,'mu_comp_best_val.png'),
                                affine=data_params['affine'],
                                loss_type=loss_type)
                    else:
                        q_max = q_mu.max()
                        q_min = q_mu.min()

                        if loss_type == 'exact':
                            q_scaled = ((q_mu - q_min ) / (q_max - q_min) * (scale) ).detach().cpu().numpy()
                        else:
                            q_scaled = q_mu.detach().cpu().numpy() / scale

                        transformed_xt, Q, b, sde_mse, sde_rel = utils.calc_affine(
                                current_run,
                                q_scaled,
                                #z.detach().cpu().numpy() / scale,
                                savepath,
                                affine=data_params['affine'],
                                gbm=gbm)

                        mu_mse, mu_rel, mu_crlb = utils.compare_mu2(mu,
                                q_mu,
                                ts,
                                Q,
                                b,
                                dt,
                                train_dataset,
                                os.path.join(savepath,'mu_comp_best_val.png'),
                                affine=data_params['affine'],
                                loss_type=loss_type)

                    stats['mu_mse_val']    = mu_mse.copy()
                    stats['mu_rel_val']    = mu_rel.copy()
                    stats['mu_crlb_val']   = mu_crlb.copy()
                    stats['sde_mse_valid'] = sde_mse.copy()
                    stats['sde_rel_valid'] = sde_rel.copy()

                else:

                    l2_small_valid = False

                    plt.plot(torch.arange(sigma.shape[0]).detach().cpu().numpy(), (sigma.sort(descending=True)[0]).detach().cpu().numpy())
                    plt.savefig(os.path.join(savepath, 'sigma_hat.pdf'))
                    plt.close('all')

                    if 'dna' in savepath or 'balls' in savepath:
                        stats['val_cond_met'] = True
                        save_nets(ae, mu, sigma, 'best_val')

                im_grid_test            = torchvision.utils.make_grid(frames_test[:64].detach().cpu(), pad_value=1, normalize=True)
                im_grid_hat_single_test = torchvision.utils.make_grid(frames_hat_test[:64].detach().cpu(), pad_value=1, normalize=True)

                # sample the frames for the next n images
                _, sampled_frames_test  = sample_n_frames(frames_test[:2], ts_test[:2], dt, ae.eval(), mu, sigma)
                _, sampled_frames_test2 = sample_n_frames(frames_test[:2], ts_test[:2], dt, ae.eval(), mu, sigma*2)
                im_grid_hat_test        = torchvision.utils.make_grid(sampled_frames_test[:64].detach().cpu(), pad_value=1, normalize=True)

                odd_rows = []

                for row in range(4):
                    odd_rows.append(frames_test[row*8:(row+1)*8])
                    odd_rows.append(sampled_frames_test[row*8:(row+1)*8])

                comp_grid = torchvision.utils.make_grid(torch.cat(odd_rows), pad_value=1, normalize=True)
                plots_list = [comp_grid.cpu().numpy().transpose((1,2,0))]
                plot_titles = ['Comparison']
                utils.plot_subplots(plots_list,
                        plot_titles,
                        os.path.join(savepath, 'valid_comparison.png'))

                if val_dataset.xt.shape[1] < 10:

                    utils.calc_affine(
                            val_dataset.xt[:z_test.shape[0]],
                            np.sqrt(dt)*z_test.detach().cpu().numpy(),
                            savepath, suffix='test')

            plots_list = [im_grid_test.numpy().transpose((1,2,0)), im_grid_hat_single_test.numpy().transpose((1,2,0))]
            plot_titles = ['Original','Sampled (single)']
            if l2_small_valid:
                utils.plot_subplots(plots_list, plot_titles, os.path.join(savepath, 'valid_sample_best.png'))
            else:
                utils.plot_subplots(plots_list, plot_titles, os.path.join(savepath, 'valid_sample.png'))

            plots_list = [im_grid_hat_test.numpy().transpose((1,2,0))]
            plot_titles = ['Sampled (trajectory)']
            if l2_small_valid:
                utils.plot_subplots(plots_list, plot_titles, os.path.join(savepath, 'valid_sample_traj_best.png'))
            else:
                utils.plot_subplots(plots_list, plot_titles, os.path.join(savepath, 'valid_sample_traj.png'))

            if len(sampled_frames_test.shape) > 2:
                utils.save_gif(sampled_frames_test.detach().cpu(), os.path.join(savepath, 'movies/valid_sample_traj.gif'))
                utils.save_gif(sampled_frames_test2.detach().cpu(), os.path.join(savepath, 'movies/valid_sample_traj_2.gif'))

            plt.figure(figsize=(10,5))
            plt.subplot(1,2,1)
            plt.title('NLL')
            plt.plot([kp[0] for kp in losses_train])
            plt.subplot(1,2,2)
            plt.title('l2')
            plt.yscale('log')
            plt.plot([kp[1] for kp in losses_train])
            plt.savefig(os.path.join(savepath, 'losses_train.png'))
            plt.close('all')

            plt.figure(figsize=(10,5))
            plt.subplot(1,2,1)
            plt.title('NLL')
            plt.plot([kp[0] for kp in losses_valid])
            plt.subplot(1,2,2)
            plt.title('l2')
            plt.yscale('log')
            plt.plot([kp[1] for kp in losses_valid])
            plt.savefig(os.path.join(savepath, 'losses_valid.png'))
            plt.close('all')

    return stats

def get_parser():
    """Get parser object."""
    from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter

    parser = ArgumentParser(
        description=__doc__, formatter_class=ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "-f",
        "--file",
        dest="filename",
        help="experiment definition file",
        metavar="FILE",
        required=True,
    )
    return parser

if __name__ == '__main__':
    import shutil

    args = get_parser().parse_args()

    yaml_filepath = args.filename

    with open(yaml_filepath, 'r') as f:
        cfg = yaml.load(f, yaml.SafeLoader)

    global savepath

    all_stats = {'config':cfg, 'runs':[]}

    try:
        n_runs = cfg['n_runs']
    except KeyError:
        n_runs = 5

    try:
        n_tries = cfg['n_tries']
    except KeyError:
        n_tries = 1
    print(n_tries)

    for run in range(n_runs):

        savepath = 'results/{}_d={}w={}z={}det={}lat={}loss={}sigma={}/run{}'.format(
                cfg['head'],
                cfg['dataset']['name'],
                cfg['ae']['net']['width'],
                cfg['ae']['net']['latent_dim'],
                cfg['ae']['net']['add_det'],
                cfg['sde']['type'],
                cfg['ae']['net']['loss'],
                cfg['ae']['net']['sigma_type'],
                run)

        global loss_type
        loss_type = cfg['ae']['net']['loss']

        #if os.path.isfile(os.path.join(savepath, 'data.pkl')):
        #    os.remove(os.path.join(savepath, 'data.pkl'))

        if not os.path.exists(savepath):
            os.makedirs(savepath)

        if not os.path.exists(os.path.join(savepath,'movies')):
            os.makedirs(os.path.join(savepath,'movies'))

        if not os.path.exists(os.path.join(savepath,'saved_nets')):
            os.makedirs(os.path.join(savepath,'saved_nets'))

        log_format = "%(asctime)s %(message)s"
        logging.basicConfig(
            stream=sys.stdout,
            level=logging.INFO,
            format=log_format,
            datefmt="%m/%d %I:%M:%S %p",
        )
        fh = logging.FileHandler(os.path.join(savepath, "log.txt"))
        fh.setFormatter(logging.Formatter(log_format))
        logging.getLogger().addHandler(fh)
        logging.info("config = %s", cfg)

        pp = pprint.PrettyPrinter(indent=4)
        pp.pprint(cfg)

        best_loss = np.Inf

        for t_num in range(n_tries):

            while True:
                initialized = setup.setup(cfg, savepath)
                stats = train(**initialized)
                val_cond_met = stats['val_cond_met']
                if val_cond_met or 'dna' in cfg['head'] or 'stocks' in cfg['head']:
                    break
            src_ae = os.path.join(savepath,'saved_nets/ae_best_val.pth')
            dst_ae = os.path.join(savepath,'saved_nets/ae_best_val_{}.pth'.format(t_num))
            src_mu = os.path.join(savepath,'saved_nets/mu_best_val.pth')
            dst_mu = os.path.join(savepath,'saved_nets/mu_best_val_{}.pth'.format(t_num))
            shutil.copyfile(src_ae, dst_ae)
            shutil.copyfile(src_mu, dst_mu)
            print('=========== End of Training ===========')
            print('Printing results for try {}'.format(t_num))
            print('STAT: L2 on Train: {}'.format(stats['l2']))
            print('STAT: KL on Train: {}'.format(stats['kl']))
            print('STAT: L2 on Validation: {}'.format(stats['l2_valid']))
            print('STAT: KL on Validation: {}'.format(stats['kl_valid']))
            print('STAT: mu mse on Validation: {}'.format(stats['mu_mse']))
            print('STAT: SDE mse on Validation: {}'.format(stats['sde_mse']))
            print('========== End of Results ============')

            if stats['kl_valid'] + stats['l2_valid'] < best_loss:

                best_loss = stats['kl_valid'] + stats['l2_valid']
                src_ae = os.path.join(savepath,'saved_nets/ae_best_val.pth')
                dst_ae = os.path.join(savepath,'saved_nets/ae_best_val_bt.pth')
                src_mu = os.path.join(savepath,'saved_nets/mu_best_val.pth')
                dst_mu = os.path.join(savepath,'saved_nets/mu_best_val_bt.pth')
                shutil.copyfile(src_ae, dst_ae)
                shutil.copyfile(src_mu, dst_mu)

        all_stats['runs'].append(stats)
        print(stats)

    with open(os.path.join(savepath,'saved_stats.pkl'), 'wb') as f:
        pickle.dump(all_stats, f)


import unittest
from selenium import webdriver
from data.constants import Constants
from helpers.keywords import Helpers
from pom.pages.login import Login
from pom.pages.project import Project
from pom.locators.base_loc import BaseLoc
from selenium.webdriver.chrome.options import Options
from webdriver_manager.chrome import ChromeDriverManager
from pom.pages.home import Home
from pom.locators.sidebar_loc import SideBarLoc

class ProjectTest(unittest.TestCase):

    def setUp(self):
        options = Options()
        options.headless = True
        print("\n========== PROJECTS TESTS ==========")
        self.driver = webdriver.Chrome(ChromeDriverManager().install(), options=options)
        self.driver.maximize_window()
        self.driver.get(Constants.url["prod"])
        self.driver.find_element(*BaseLoc.sign_in_lnk).click()
        Login.login_form(self, Constants.credentials["users"]["real"]["user"], Constants.credentials["users"]["real"]["pass"])
        Helpers.click_visible_element(self, SideBarLoc.inbox_li_btn)

    def test_create_project(self):
        Project.create_projects(self, 1, Constants.project_data["name"], False, Constants.project_data["view"]["panel"])
        Home.validate_project(self)

    def test_create_project_fav(self):
        Project.create_projects(self, 1, Constants.project_data["name"], True, Constants.project_data["view"]["list"])
        Home.validate_project(self)

    def test_create_projects(self):
        Project.create_projects(self, 3, Constants.project_data["name"], True, Constants.project_data["view"]["list"])
        Home.validate_project(self)

    def test_delete_all_projects(self):
        Project.delete_all_projects(self)

    def tearDown(self):
        Helpers.wait_seconds(self, 3)
        self.driver.quit()

    if __name__ == "__main__":
        unittest.main()


from flask import Blueprint, jsonify, g, request
from wrappers.auth_required import auth_required, rate_limited
from models.jobs import TOPJob
from utils.json_helper import jsonify_payload

bp = Blueprint("management", __name__, url_prefix="/management")

@bp.route("/jobs", methods=["GET"])
@auth_required
def get_jobs():
    print('user info: ', g.user_id)

    jobs = TOPJob.find_by_user_id(g.user_id)

    print('jobs: ', [job.serialize() for job in jobs])

    return jsonify_payload({'jobs': [job.serialize() for job in jobs]})

@bp.route("/jobs", methods=["POST"])

@auth_required
@rate_limited
def create_job():
    payload = request.get_json()

    print('payload: ', payload)
    print('subscription data: ', g.subscription)

    job_name = payload.get('job_name')
    job_description = payload.get('job_description')
    job_id = payload.get('job_id')

    if not job_id:
        print('creating job')
        job = TOPJob(job_name, job_description, g.user_id)
        job.save()

    else:
        print('updating job ', job_id)
        job = TOPJob.find_by_id(job_id)
        job.update(job_name, job_description)

    return jsonify_payload({'job': job.serialize()})


import datetime
import sys
import uuid

import pandas as pd
import pytest
from dagster_gcp import (
    bigquery_resource,
    bq_create_dataset,
    bq_delete_dataset,
    bq_solid_for_queries,
    import_df_to_bq,
    import_gcs_paths_to_bq,
)
from dagster_pandas import DataFrame
from google.cloud import bigquery
from google.cloud.exceptions import NotFound

from dagster import (
    DagsterExecutionStepExecutionError,
    InputDefinition,
    List,
    ModeDefinition,
    Nothing,
    OutputDefinition,
    Path,
    execute_pipeline,
    pipeline,
    solid,
)
from dagster.config.validate import validate_config
from dagster.core.definitions import create_environment_type
from dagster.seven import mock

def dataset_exists(name):
    '''Check if dataset exists - ensures we have properly cleaned up after tests and haven't leaked
    any datasets'''
    client = bigquery.Client()
    dataset_ref = client.dataset(name)

    try:
        client.get_dataset(dataset_ref)
        return True
    except NotFound:
        return False

def get_dataset():
    '''Creates unique dataset names of the form: test_ds_83791a53
    '''
    return 'test_ds_' + str(uuid.uuid4()).replace('-', '_')

def bq_modes():
    return [ModeDefinition(resource_defs={'bigquery': bigquery_resource})]

def test_simple_queries():
    @pipeline(mode_defs=bq_modes())
    def bq_pipeline():
        bq_solid_for_queries(
            [
                # Toy example query
                'SELECT 1 AS field1, 2 AS field2;',
                # Test access of public BQ historical dataset (only processes ~2MB here)
                # pylint: disable=line-too-long
                '''SELECT *
            FROM `weathersource-com.pub_weather_data_samples.sample_weather_history_anomaly_us_zipcode_daily`
            ORDER BY postal_code ASC, date_valid_std ASC
            LIMIT 1''',
            ]
        ).alias('bq_query_solid')()

    res = execute_pipeline(bq_pipeline).result_for_solid('bq_query_solid')
    assert res.success

    values = res.output_value()
    for df in values:
        assert isinstance(df, pd.DataFrame)
    assert values[0].to_dict('list') == {'field1': [1], 'field2': [2]}
    assert values[1].to_dict('list') == {
        'postal_code': ['02101'],
        'country': ['US'],
        'date_valid_std': [datetime.date(2014, 1, 1)],
        'doy_std': [1],
        'avg_temperature_air_2m_f': [25.05],
        'avg_temperature_anomaly_air_2m_f': [-7.81],
        'tot_precipitation_in': [0.0],
        'tot_precipitation_anomaly_in': [-0.28],
        'tot_snowfall_in': [0.0],
        'tot_snowfall_anomaly_in': [-1.36],
        'avg_wind_speed_10m_mph': [7.91],
        'avg_wind_speed_10m_anomaly_mph': [-1.85],
    }

# pylint: disable=line-too-long
def test_bad_config():
    configs_and_expected_errors = [
        (
            # Create disposition must match enum values
            {'create_disposition': 'this is not a valid create disposition'},
            'Value not in enum type BQCreateDisposition',
        ),
        (
            # Dataset must be of form project_name.dataset_name
            {'default_dataset': 'this is not a valid dataset'},
            'Value at path root:solids:test:config:query_job_config:default_dataset is not valid. Expected "_Dataset"',
        ),
        (
            # Table must be of form project_name.dataset_name.table_name
            {'destination': 'this is not a valid table'},
            'Value at path root:solids:test:config:query_job_config:destination is not valid. Expected "_Table"',
        ),
        (
            # Priority must match enum values
            {'priority': 'this is not a valid priority'},
            'Value not in enum type BQPriority',
        ),
        (
            # Schema update options must be a list
            {'schema_update_options': 'this is not valid schema update options'},
            'Value at path root:solids:test:config:query_job_config:schema_update_options must be list. Expected: [BQSchemaUpdateOption]',
        ),
        (
            {'schema_update_options': ['this is not valid schema update options']},
            'Value not in enum type BQSchemaUpdateOption',
        ),
        (
            {'write_disposition': 'this is not a valid write disposition'},
            'Value not in enum type BQWriteDisposition',
        ),
    ]

    @pipeline(mode_defs=bq_modes())
    def test_config_pipeline():
        bq_solid_for_queries(['SELECT 1']).alias('test')()

    env_type = create_environment_type(test_config_pipeline)
    for config_fragment, error_message in configs_and_expected_errors:
        config = {'solids': {'test': {'config': {'query_job_config': config_fragment}}}}
        result = validate_config(env_type, config)
        assert result.errors[0].message == error_message

def test_create_delete_dataset():
    dataset = get_dataset()

    @pipeline(mode_defs=bq_modes())
    def create_pipeline():
        bq_create_dataset.alias('create_solid')()

    config = {'solids': {'create_solid': {'config': {'dataset': dataset, 'exists_ok': True}}}}

    assert execute_pipeline(create_pipeline, config).result_for_solid('create_solid').success

    config = {'solids': {'create_solid': {'config': {'dataset': dataset, 'exists_ok': False}}}}
    with pytest.raises(DagsterExecutionStepExecutionError) as exc_info:
        execute_pipeline(create_pipeline, config)
    assert 'Dataset "%s" already exists and exists_ok is false' % dataset in str(
        exc_info.value.user_exception
    )

    @pipeline(mode_defs=bq_modes())
    def delete_pipeline():
        bq_delete_dataset.alias('delete_solid')()

    # Delete should succeed
    config = {'solids': {'delete_solid': {'config': {'dataset': dataset}}}}
    assert execute_pipeline(delete_pipeline, config).result_for_solid('delete_solid').success

    # Delete non-existent with "not_found_ok" should succeed
    config = {'solids': {'delete_solid': {'config': {'dataset': dataset, 'not_found_ok': True}}}}
    assert execute_pipeline(delete_pipeline, config).result_for_solid('delete_solid').success

    # Delete non-existent with "not_found_ok" False should fail
    config = {'solids': {'delete_solid': {'config': {'dataset': dataset, 'not_found_ok': False}}}}
    with pytest.raises(DagsterExecutionStepExecutionError) as exc_info:
        execute_pipeline(delete_pipeline, config)
    assert 'Dataset "%s" does not exist and not_found_ok is false' % dataset in str(
        exc_info.value.user_exception
    )

    assert not dataset_exists(dataset)

# See: https://github.com/dagster-io/dagster/issues/1711
@pytest.mark.skip
def test_pd_df_load():
    dataset = get_dataset()
    table = '%s.%s' % (dataset, 'df')

    test_df = pd.DataFrame({'num1': [1, 3], 'num2': [2, 4]})

    create_solid = bq_create_dataset.alias('create_solid')
    load_solid = import_df_to_bq.alias('load_solid')
    query_solid = bq_solid_for_queries(['SELECT num1, num2 FROM %s' % table]).alias('query_solid')
    delete_solid = bq_delete_dataset.alias('delete_solid')

    @solid(
        input_defs=[InputDefinition('success', Nothing)], output_defs=[OutputDefinition(DataFrame)]
    )
    def return_df(_context):  # pylint: disable=unused-argument
        return test_df

    config = {
        'solids': {
            'create_solid': {'config': {'dataset': dataset, 'exists_ok': True}},
            'load_solid': {'config': {'destination': table}},
            'delete_solid': {'config': {'dataset': dataset, 'delete_contents': True}},
        }
    }

    @pipeline(mode_defs=bq_modes())
    def bq_pipeline():
        delete_solid(query_solid(load_solid(return_df(create_solid()))))

    result = execute_pipeline(bq_pipeline, config)
    assert result.success

    values = result.result_for_solid('query_solid').output_value()
    assert values[0].to_dict() == test_df.to_dict()

    # BQ loads should throw an exception if pyarrow and fastparquet aren't available
    with mock.patch.dict(sys.modules, {'pyarrow': None, 'fastparquet': None}):
        with pytest.raises(DagsterExecutionStepExecutionError) as exc_info:
            result = execute_pipeline(bq_pipeline, config)
        assert (
            'loading data to BigQuery from pandas DataFrames requires either pyarrow or fastparquet'
            ' to be installed' in str(exc_info.value.user_exception)
        )
        cleanup_config = {
            'solids': {'delete_solid': {'config': {'dataset': dataset, 'delete_contents': True}}}
        }

        @pipeline(mode_defs=bq_modes())
        def cleanup():
            delete_solid()

        assert execute_pipeline(cleanup, cleanup_config).success

    assert not dataset_exists(dataset)

# See: https://github.com/dagster-io/dagster/issues/1711
@pytest.mark.skip
def test_gcs_load():
    dataset = get_dataset()
    table = '%s.%s' % (dataset, 'df')

    create_solid = bq_create_dataset.alias('create_solid')
    query_solid = bq_solid_for_queries(
        [
            'SELECT string_field_0, string_field_1 FROM %s ORDER BY string_field_0 ASC LIMIT 1'
            % table
        ]
    ).alias('query_solid')
    delete_solid = bq_delete_dataset.alias('delete_solid')

    @solid(
        input_defs=[InputDefinition('success', Nothing)], output_defs=[OutputDefinition(List[Path])]
    )
    def return_gcs_uri(_context):  # pylint: disable=unused-argument
        return ["gs://cloud-samples-data/bigquery/us-states/us-states.csv"]

    config = {
        'solids': {
            'create_solid': {'config': {'dataset': dataset, 'exists_ok': True}},
            'import_gcs_paths_to_bq': {
                'config': {
                    'destination': table,
                    'load_job_config': {
                        'autodetect': True,
                        'skip_leading_rows': 1,
                        'source_format': 'CSV',
                        'write_disposition': 'WRITE_TRUNCATE',
                    },
                }
            },
            'delete_solid': {'config': {'dataset': dataset, 'delete_contents': True}},
        }
    }

    @pipeline(mode_defs=bq_modes())
    def bq_pipeline():
        delete_solid(query_solid(import_gcs_paths_to_bq(return_gcs_uri(create_solid()))))

    result = execute_pipeline(bq_pipeline, config)
    assert result.success

    values = result.result_for_solid('query_solid').output_value()
    assert values[0].to_dict() == {'string_field_0': {0: 'Alabama'}, 'string_field_1': {0: 'AL'}}

    assert not dataset_exists(dataset)


import matplotlib.pyplot as plt

from random_walk import RandomWalk

# cd Documents/python_work/data_visualization

while True:
    # Create instance of RandomWalk.
    rw = RandomWalk(5000)
    rw.fill_walk()

    # Set the size of the interactive window.
    plt.figure(dpi=128, figsize=(10, 5))

    # Plot random walk with gradient.
    point_numbers = list(range(rw.num_points))
    # plt.scatter(rw.x_values, rw.y_values, c=point_numbers, cmap=plt.cm.Blues,
    #     s=1)
    plt.plot(rw.x_values, rw.y_values, c='blue')

    # Start and end points (green and red, respectively).
    plt.scatter(0, 0, c='green', s=100)
    plt.scatter(rw.x_values[-1], rw.y_values[-1], c='red', s=100)

    # Remove axes.
    # plt.axes().get_xaxis().set_visible(False)
    # plt.axes().get_yaxis().set_visible(False)

    plt.show()

    keep_running = input("Make another walk? (y/n): ")
    if keep_running == 'n':
        break


import imaplib
import email
from time import sleep
from random import randint
import importlib
from src.Analyser import mark_email
from src.Email import Email
import numpy as np
from goto import with_goto
from src.save import save

ai = importlib.import_module("Neural_Network", package=None)

"""
Fonction qui efface tous les fichiers emailX.txt
De nouveaux fichiers seront créés, avec X commençant à 1
"""
def efface_old_mail():
    import os

    arret = False
    x = 1
    while not arret:
        if os.path.exists("email" + str(x) + ".txt"):
            os.remove("email" + str(x) + ".txt")
            x += 1
        else:
            print("Fichiers supprimés")
            arret = True

old_mes_nb = -1
x = 1

efface_old_mail()

@with_goto
def run():
    print('start')

    label.connexion

    ### Récupération de l'adresse mail à tester, et du mot de passe ###
    connexion_pos = False
    while not connexion_pos:
        adresse = input("Adresse mail: ")
        mdp = input("Mot de passe: ")

        if("@gmail.com" in adresse):
            connexion_pos = True
        else:
            print("Adresse mail non valide\n")
            continue

    label.start
    try:
        old_mes_nb = -1
        x = 1

        ### Connexion à la boite mail ###
        try :
            mail = imaplib.IMAP4_SSL('imap.gmail.com')

            #mail.login('yncrea.test.projet.M1@gmail.com', 'ujikolpm')
            mail.login(adresse, mdp)

        except Exception:
            ### Cas où la connexion échoue ###
            print("Echec connexion\n")
            goto.connexion

        while True:

            mail.list()
            mail.select('inbox')

            result, data = mail.uid('search', None, "ALL")

            i = len(data[0].split())
            new_mes_nb = i

            if (old_mes_nb == -1):
                old_mes_nb = new_mes_nb

            ### Un nouveau message arrive dans la boite mail ###
            if (new_mes_nb > old_mes_nb):
                print("\n---NOUVEAU MESSAGE : %i---" % x)
                latest_email_uid = data[0].split()[new_mes_nb - 1]
                result, email_data = mail.uid('fetch', latest_email_uid, '(RFC822)')

                raw_email = email_data[0][1]

                raw_email_string = raw_email.decode('utf-8')
                email_message = email.message_from_string(raw_email_string)

                ### Création d'un fichier texte contenant le message                ###
                ### Création d'un objet Email récupérant les infos du fichier texte ###
                for part in email_message.walk():
                    save_string = r"email" + str(x) + ".txt"
                    myfile = open(save_string, 'a')
                    myfile.write(str(part))
                    mailo = Email(save_string)
                    myfile.close()

                ### L'email est déplacé dans le dossier traitement ###
                cible_dossier = 'traitement'

                result_move, err_mes = mail.uid('move', latest_email_uid, cible_dossier)
                if (result_move == 'OK'):
                    print("Mail déplacé avec succès")
                else:
                    print(err_mes)

                mail.select(cible_dossier)

                result, data = mail.uid('search', None, "ALL")
                latest_email_uid = data[0].split()[- 1]

                ### Analyse du message et attribution de son niveau de dangerosit ###
                mark = mark_email(mailo)
                marks = np.array([mark])
                sortie_traitement = ai.analyse_mail(marks)[0][0]
                save(mailo, marks=mark, grade=sortie_traitement.item())

                print("Résultat traitement :", sortie_traitement)

                if (sortie_traitement >= 0.6): ### Cas d'un message sur ###
                    result_move, err_mes = mail.uid('move', latest_email_uid, "sur")
                    if (result_move == 'OK'):
                        print("Mail déplacé dans sur")
                    else:
                        print(err_mes)

                elif (sortie_traitement >= 0.4 and sortie_traitement < 0.6): ### Cas d'un message pour lequel l'IA a un doute ###
                    result_move, err_mes = mail.uid('move', latest_email_uid, "moyen")
                    if (result_move == 'OK'):
                        print("Mail déplacé dans moyen")
                    else:
                        print(err_mes)

                else: ### Cas d'un message dangereux ###
                    result_move, err_mes = mail.uid('move', latest_email_uid, "danger")
                    if (result_move == 'OK'):
                        print("Mail déplacé dans danger")
                    else:
                        print(err_mes)

                x += 1
                old_mes_nb = new_mes_nb
                print("Analyse effectuée")

            elif (new_mes_nb < old_mes_nb):### Cas où des messages ont été supprimés ###
                old_mes_nb = new_mes_nb
    except TimeoutError:### Timeout de la connexion avec la boite mail atteint, retour au label start pour rafraichir la connexion ###
        goto.start
    except KeyboardInterrupt:
        goto.end

    label.end
    mail.logout()
    print("Good bye")

run()


# coding: utf-8

"""
    Yapily API

    To access endpoints that require authentication, use your application key and secret created in the Dashboard (https://dashboard.yapily.com)  # noqa: E501

    The version of the OpenAPI document: 0.0.358
    Generated by: https://openapi-generator.tech
"""

from __future__ import absolute_import

import unittest
import datetime

import yapily
from yapily.models.bulk_user_delete_details import BulkUserDeleteDetails  # noqa: E501
from yapily.rest import ApiException

class TestBulkUserDeleteDetails(unittest.TestCase):
    """BulkUserDeleteDetails unit test stubs"""

    def setUp(self):
        pass

    def tearDown(self):
        pass

    def make_instance(self, include_optional):
        """Test BulkUserDeleteDetails
            include_option is a boolean, when False only required
            params are included, when True both required and
            optional params are included """
        # model = yapily.models.bulk_user_delete_details.BulkUserDeleteDetails()  # noqa: E501
        if include_optional :
            return BulkUserDeleteDetails(
                id = '0',
                invalid_application_user_ids = [
                    '0'
                    ],
                invalid_user_uuids = [
                    '0'
                    ],
                status = 'IN_PROGRESS',
                started_at = datetime.datetime.strptime('2013-10-20 19:20:30.00', '%Y-%m-%d %H:%M:%S.%f'),
                users = [
                    yapily.models.user_delete_response.UserDeleteResponse(
                        id = '0',
                        delete_status = 'SUCCESS',
                        creation_date = datetime.datetime.strptime('2013-10-20 19:20:30.00', '%Y-%m-%d %H:%M:%S.%f'),
                        user_consents = [
                            yapily.models.consent_delete_response.ConsentDeleteResponse(
                                id = '0',
                                delete_status = 'SUCCESS',
                                institution_id = '0',
                                institution_consent_id = '0',
                                creation_date = datetime.datetime.strptime('2013-10-20 19:20:30.00', '%Y-%m-%d %H:%M:%S.%f'), )
                            ], )
                    ],
                links = {
                    'key' : '0'
                    }
            )
        else :
            return BulkUserDeleteDetails(
        )

    def testBulkUserDeleteDetails(self):
        """Test BulkUserDeleteDetails"""
        inst_req_only = self.make_instance(include_optional=False)
        inst_req_and_optional = self.make_instance(include_optional=True)

if __name__ == '__main__':
    unittest.main()


from django.conf.urls import url, include
from rest_framework.urlpatterns import format_suffix_patterns
from devup.views import UpList, UpDetail, UpCreate, UpUpdate

app_name = 'devup'

urlpatterns = [
    url(r'^up_list$', UpList.as_view(), name='up_list'),
    url(r'^up_create$', UpCreate.as_view(), name='up_create'),
    url(r'^up_detail/(?P<pk>[-\w]+)/$', UpDetail.as_view(), name='up_detail'),
    url(r'^(?P<pk>[-\w]+)/$', UpDetail.as_view(), name='detail'),
    url(r'^(?P<pk>\d+)/update$', UpUpdate.as_view(), name='update'),

]


from utils import connector

async def declare_queue(queue_name, durable=False):
    conct = connector.Connector()
    channel = await conct.get_channel()
    await channel.queue_declare(
        queue=queue_name,
        durable=durable,
    )

async def bind_queue(queue_name, exchange_name, routing_key):
    conct = connector.Connector()
    channel = await conct.get_channel()
    await channel.queue_bind(
        queue=queue_name,
        exchange=exchange_name,
        routing_key=routing_key,
    )

async def declare_exchange(
        exchange_name,
        exchange_type="direct",
        durable=False,
):
    conct = connector.Connector()
    channel = await conct.get_channel()
    await channel.exchange_declare(
        exchange=exchange_name,
        exchange_type=exchange_type,
        durable=durable,

        passive=False,
        # if passive is True - it will raise exception if exchange
        # doesn't exist

        internal=False,
        # If set, the exchange may not be used directly by publishers,
        # but only when bound to other exchanges. Internal exchanges are
        # used to construct wiring that is not visible to applications.
        # Hint: could be used as "dead-letter-exchange" for queues
    )


from argparse import ArgumentParser

from gitrello import Gitrello
import github
import trello
import settings

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument('--pr_id', required=True)
    parser.add_argument('--repo', required=True)
    args = parser.parse_args()

    g = github.Github(settings.GITHUB_TOKEN).get_user()
    client = trello.TrelloClient(api_key=settings.API_KEY, token=settings.API_TOKEN)
    board = client.get_board(settings.BOARD_ID)
    repo = [x for x in g.get_repos() if x.name == args.repo][0]
    pull = repo.get_pull(int(args.pr_id))
    gitrello = Gitrello(pull, board)
    card = gitrello.create_card()


import json
import os
import subprocess
import sys
from pathlib import Path
import youtube_dl

ydl_opts_download = {
    "format": "bestaudio/best",
    "cachedir": False,
    "outtmpl": "%(id)s%(ext)s",
    "postprocessors": [
        {
            "key": "FFmpegExtractAudio",
            "preferredcodec": "mp3",
            "preferredquality": "192",
        }
    ],
}

def download_single_yt(url_list):
    skipped = []
    for i in range(len(url_list)):
        try:
            with youtube_dl.YoutubeDL(ydl_opts_download) as ydl:
                ydl.download([url_list[i]])
        except:
            skipped.append(url_list[i])
    if skipped == []:
        return 0
    else:
        download_single_yt(skipped)

def ytdownload(link):
    with youtube_dl.YoutubeDL(
        {
            "outtmpl": "%(id)s%(ext)s",
            "quiet": True,
        }
    ) as ydl:
        result = ydl.extract_info(link, download=False)
        if "entries" in result:
            # Can be a playlist or a list of videos
            video = result["entries"]
        playlist_urls = [
            result["entries"][i]["webpage_url"] for i, item in enumerate(video)
        ]

    download_single_yt(playlist_urls)
    print("-" * 15)

def download(title, link, out_folder, i):
    print("downloading ", title, " OTS")
    os.chdir("./" + out_folder)
    fname = ""
    if i < 10:
        fname = "0" + str(i) + " - " + title
    else:
        fname = str(i) + " - " + title
    os.mkdir(fname)
    os.chdir("./" + fname)
    if "spotify" in link.lower():
        subprocess.check_call(["spotdl", link, "--output-format wav"])
    elif "youtube" in link.lower():
        ytdownload(link)
    os.chdir("..")
    os.chdir("..")

def download_all(json_source, out_folder):
    print("open file...")
    file = open(json_source)
    movies = json.load(file)

    print("creating main folder...")
    ost = Path(out_folder)
    if not ost.exists():
        ost.mkdir()

    for i, movie in enumerate(movies):
        link = movie["link"].replace(" ", "_")
        title = movie["title"].replace(" ", "_")
        download(title, link, out_folder, i + 1)

    print("--- DONE ---")

# download_all("ESC.json", "ESC")


from pathlib import Path
import numpy as np
import torch
from torch.utils.data import Dataset
from torchvision import transforms
import torchvision.transforms.functional as TF
from PIL import Image

from src.draw_utils import save_img_with_kps
from src.readers.image_reader import ImageReader
from typing import Dict
from torch import Tensor
from typing import Tuple

class Phase0PointsDataset(Dataset):
    MEAN = [0.5, 0.5, 0.5]
    STD = [0.2, 0.2, 0.2]
    IMG_SIZE = 768
    TRANSFORMS = transforms.Compose(
        [
            transforms.Resize(IMG_SIZE),
            transforms.ToTensor(),
            transforms.Normalize(mean=MEAN, std=STD),
        ]
    )

    def __init__(self, reader: ImageReader, augment: bool = False):
        assert isinstance(reader, ImageReader)
        self.reader = reader
        self.augment = augment

    def __len__(self):
        return len(self.reader)

    def __getitem__(self, idx: int) -> Dict[str, Tensor]:
        sample = self.reader[idx]

        img = sample.phase_0_image
        kps = torch.tensor(sample.phase_0_points).to(dtype=torch.float32)

        if self.augment:
            if np.random.rand() < 0.5:
                img, kps = Phase0PointsDataset.color_augment(img, kps)

            if np.random.rand() < 0.5:
                img, kps = Phase0PointsDataset.rotate(img, kps)

            if np.random.rand() < 0.5:
                img, kps = Phase0PointsDataset.perspective_augment(img, kps)

            if np.random.rand() < 0.5:
                img, kps = Phase0PointsDataset.crop_augment(img, kps)

        kps = kps / torch.tensor([img.width, img.height])
        kps = kps.flatten()

        img_tensor = Phase0PointsDataset.TRANSFORMS(img)

        sample_t = {
            "img": img_tensor,
            "kps": kps,
        }

        return sample_t

    @staticmethod
    def color_augment(img: Image.Image, kps: Tensor) -> Tuple[Image.Image, Tensor]:
        img = TF.adjust_brightness(img, 0.7 + np.random.rand() * 1.5)
        img = TF.adjust_contrast(img, 0.5 + np.random.rand() * 1.5)
        img = TF.adjust_gamma(img, gamma=0.5 + np.random.rand(), gain = 0.5 + np.random.rand())
        img = TF.adjust_hue(img, -0.5 + np.random.rand())
        img = TF.adjust_saturation(img, np.random.rand() * 1.5)
        return img, kps

    @staticmethod
    def rotate(img: Image.Image, kps: Tensor) -> Tuple[Image.Image, Tensor]:
        rotation_angle_deg = np.random.rand() * 30 - 15
        rotation_angle_rad = np.deg2rad(rotation_angle_deg)
        rotation_matrix = np.array(
            [
                [np.cos(rotation_angle_rad), -np.sin(rotation_angle_rad)],
                [np.sin(rotation_angle_rad), np.cos(rotation_angle_rad)],
            ]
        )
        rot_torch = torch.from_numpy(rotation_matrix.astype(np.float32))
        img = TF.rotate(img, np.rad2deg(rotation_angle_rad))

        center = torch.tensor([img.width, img.height]) / 2
        kps = kps - center
        kps = torch.matmul(kps, rot_torch)
        kps = kps + center
        return img, kps

    @staticmethod
    def perspective_augment(img: Image.Image, kps: Tensor) -> Tuple[Image.Image, Tensor]:
        topleft = kps[0]
        topright = kps[1]
        bottomleft = kps[2]
        bottomright = kps[3]

        startpoints = [
            topleft.to(dtype=torch.int32).tolist(),
            topright.to(dtype=torch.int32).tolist(),
            bottomright.to(dtype=torch.int32).tolist(),
            bottomleft.to(dtype=torch.int32).tolist(),
        ]

        a = min([torch.linalg.norm(topleft - topright) * 0.1, torch.linalg.norm(topleft - bottomleft) * 0.1])
        new_topleft = topleft + (-a + np.random.rand() * 2*a)
        new_topleft = torch.clip(new_topleft, torch.tensor([0, 0]), torch.tensor([img.width, img.height]))
        new_topright = topright + (-a + np.random.rand() * 2*a)
        new_topright = torch.clip(new_topright, torch.tensor([0, 0]), torch.tensor([img.width, img.height]))
        new_bottomleft = bottomleft + (-a + np.random.rand() * 2*a)
        new_bottomleft = torch.clip(new_bottomleft, torch.tensor([0, 0]), torch.tensor([img.width, img.height]))
        new_bottomright = bottomright + (-a + np.random.rand() * 2*a)
        new_bottomright = torch.clip(new_bottomright, torch.tensor([0, 0]), torch.tensor([img.width, img.height]))

        endpoints = [
            new_topleft.to(dtype=torch.int32).tolist(),
            new_topright.to(dtype=torch.int32).tolist(),
            new_bottomright.to(dtype=torch.int32).tolist(),
            new_bottomleft.to(dtype=torch.int32).tolist(),
        ]
        img = transforms.functional.perspective(img, startpoints, endpoints)
        kps = torch.stack([new_topleft, new_topright, new_bottomleft, new_bottomright])

        return img, kps

    @staticmethod
    def crop_augment(img: Image.Image, kps: Tensor) -> Tuple[Image.Image, Tensor]:
        kps_x0 = kps[:, 0].min().item()
        kps_x1 = kps[:, 0].max().item()
        kps_y0 = kps[:, 1].min().item()
        kps_y1 = kps[:, 1].max().item()

        crop_x0 = int(kps_x0 * np.random.rand())
        crop_x1 = int(kps_x1 + np.random.rand() * (img.width - kps_x1))
        crop_y0 = int(kps_y0 * np.random.rand())
        crop_y1 = int(kps_y1 + np.random.rand() * (img.height - kps_y1))

        # make square
        crop_1 = max(crop_x1 - crop_x0, crop_y1 - crop_y0)
        crop_y1 = crop_y0 + crop_1
        crop_x1 = crop_x0 + crop_1

        img = img.crop((crop_x0, crop_y0, crop_x1, crop_y1))
        kps = kps - torch.tensor([crop_x0, crop_y0])

        return img, kps

    @staticmethod
    def img_from_tensor(img_tensor: Tensor) -> Image.Image:
        img: np.ndarray = img_tensor.permute(1, 2, 0).numpy()
        img = (
            img * np.array(Phase0PointsDataset.STD) + np.array(Phase0PointsDataset.MEAN)
        ) * 255
        img = img.astype(np.uint8)
        img = Image.fromarray(img)
        return img

    def show(self, idx: int, out_folder: Path, repeat_idx=0, verbose: bool = False):
        sample_t = self[idx]
        img_tensor = sample_t["img"]
        kps_tensor = sample_t["kps"]

        img = Phase0PointsDataset.img_from_tensor(img_tensor)
        kps = kps_tensor.reshape(-1, 2).numpy() * Phase0PointsDataset.IMG_SIZE
        filename = out_folder / f"sample_{idx}_{repeat_idx}.jpg"

        save_img_with_kps(img, kps, filename, circle_radius=10, verbose=verbose)


# django imports
from django import template

import itertools
import datetime
import pytz
import dateutil

register = template.Library()

@register.filter
def group_by_date(dates, timezone):
    tz = pytz.timezone(timezone)
    dates_parser = []
    for day in dates:
        try:
            new_date = pytz.utc.localize(dateutil.parser.parse(day))
        except ValueError:
            new_date = dateutil.parser.parse(day)
        dates_parser.append(new_date)
    days = [
        tz.normalize(day.replace(tzinfo=pytz.utc)) for day in dates_parser
    ]
    days2 = [
        list(group) for k, group in itertools.groupby(
            days, key=datetime.datetime.toordinal,
        )
    ]

    return [(day[0].date, day) for day in days2]