from django.shortcuts import render
from django.http import HttpResponse
from django.http import JsonResponse
from django.views.decorators.csrf import csrf_exempt
import random
import http.client
import urllib.request
import urllib.parse
import json
import hmac
import hashlib
import base64

accesskey="<accesskey>"
secretkey="<secretkey>"
environment="test"

remote_script="https://sandbox-payments.open.money/layer"
sample_data = dict()
sample_data["amount"] = "12.00"
sample_data["currency"] = "INR"
sample_data["name"] = "John Doe"
sample_data["email_id"]="john.doe@dummydomain.com"
sample_data["contact_number"]= "9831111111"
sample_data["mtx"]= ""
sample_data["empty"]=""

BASE_URL_SANDBOX = "sandbox-icp-api.bankopen.co";
BASE_URL_UAT = "icp-api.bankopen.co";

# Create your views here.
@csrf_exempt
def index(request):
	global accesskey,secretkey,environment,remote_script,sample_data
	error=""
	layer_payment_token_data=dict()
	payment_token_data = dict()
	token_id=""
	hash = ""
	layer_params=""
	sample_data["mtx"] = random.randint(1,200)

	layer_payment_token_data = create_payment_token(sample_data,accesskey,secretkey,environment)

	if layer_payment_token_data:
		for k in layer_payment_token_data.keys():
			if k == "error":
				error = layer_payment_token_data[k]

	if len(error) == 0 and len(layer_payment_token_data["id"]) < 1:
		error="E55 Payment error. Token data empty."

	if len(error) == 0 and len(layer_payment_token_data["id"]) > 0:
		payment_token_data = get_payment_token(layer_payment_token_data["id"],accesskey,secretkey,environment)

	if payment_token_data:
		for k in payment_token_data.keys():
			if k == "error":
				error = payment_token_data[k]

	if len(error) == 0 and len(payment_token_data["id"]) < 1:
		error="Payment error. Layer token ID cannot be empty."

	if len(error) == 0 and len(payment_token_data["id"]) > 0 and payment_token_data["status"]=="paid":
		error="Layer: this order has already been paid."

	if len(error) == 0 and str(payment_token_data["amount"]) != str(sample_data["amount"]):
		error="Layer: an amount mismatch occurred."

	if error == "":
		gen = dict()
		gen["amount"]=payment_token_data["amount"]
		gen["id"]=payment_token_data["id"]
		gen["mtx"]=sample_data["mtx"]
		hash=create_hash(gen,accesskey,secretkey)
		layer_params = "{payment_token_id:"+payment_token_data["id"]+",accesskey:"+accesskey+"}"
		token_id=payment_token_data["id"]

	return render(request,
	'layerpayment/checkout.html',
	{'txnid':str(sample_data["mtx"]),
	'fullname':sample_data["name"],
	'email':sample_data["email_id"],
	'mobile':sample_data["contact_number"],
	'amount':str(sample_data["amount"]),
	'currency':sample_data["currency"],
	'remote_script':remote_script,
	'token_id':token_id,
	'hash':hash,
	'accesskey':accesskey,
	'layer_params':layer_params,
	'error':error})

@csrf_exempt
def callback(request):
	global accesskey,secretkey,environment
	error=""
	status=""
	payment_data=dict()

	response = request.POST
	if len(response["layer_payment_id"]) == 0:
		error = "Invalid payment id"
	if len(error)==0:
		vhash=dict()
		vhash["amount"] =response["layer_order_amount"]
		vhash["id"]=response["layer_pay_token_id"]
		vhash["mtx"]=response["tranid"]
		if not verify_hash(vhash,response["hash"],accesskey,secretkey):
			error="Invalid payment response...Hash mismatch"
	if len(error) == 0:
		payment_data = get_payment_details(response["layer_payment_id"],accesskey,secretkey,environment)

	if payment_data:
		for k in payment_data.keys():
			if k == "error":
				error = payment_data[k]
	if len(error) == 0 and payment_data["payment_token"]["id"] != response["layer_pay_token_id"]:
		error = "Layer: received layer_pay_token_id and collected layer_pay_token_id doesnt match"
	if len(error) == 0 and payment_data["amount"] != response["layer_order_amount"]:
		error = "Layer: received amount and collected amount doesnt match"
	if len(error) == 0 and payment_data["payment_token"]["status"] != "paid":
		status = "Transaction failed..."+payment_data["payment_error_description"]
	elif len(error) == 0:
		status = "Transaction Successful"

	return render(request,
	'layerpayment/response.html',
	{'errorstring':error,
	 'status':status})

def create_payment_token(data,accesskey,secretkey,environment):
	response=dict()

	try:
		emptykeys=[]
		for k in data.keys():
			if len(str(data[k]))<1:
				emptykeys.append(k)
		for i in emptykeys:
			del data[i]
		response = http_post(data,"payment_token",accesskey,secretkey,environment)
	except Exception as ex:
		response["error"]=ex

	return response

def get_payment_token(payment_token_id,accesskey,secretkey,environment):
	response=dict()
	try:
		if len(payment_token_id)==0 or payment_token_id.isspace():
			response["error"]="payment_token_id cannot be empty"
		else:
			response = http_get("payment_token/" + payment_token_id,accesskey,secretkey,environment)
	except Exception as ex:
		response["error"] = ex

	return response

def get_payment_details(payment_id,accesskey,secretkey,environment):
	response=dict()
	try:
		if len(payment_id)==0 or payment_id.isspace():
			response["error"]="pyment_id cannot be empty"
		else:
			response=http_get("payment/"+payment_id,accesskey,secretkey,environment)
	except Exception as ex:
		response["error"] = ex

	return response

def http_post(data,route,accesskey,secretkey,environment):
	response = ""
	url = BASE_URL_SANDBOX
	if environment == "live":
		url = BASE_URL_UAT

	resource = "/api/"+route

	try:
		conn = http.client.HTTPSConnection(url,timeout=10)
		headers = {'Content-type': 'application/json',"Authorization":"Bearer "+accesskey+":"+secretkey}
		jdata = json.dumps(data)
		conn.request('POST', resource, jdata, headers)
		resp = conn.getresponse()
		rdata = resp.read().decode('utf-8')
		conn.close()
		response = json.loads(rdata)
	except Exception as ex:
		print(ex)

	return response

def http_get(route,accesskey,secretkey,environment):
	response = ""
	url = BASE_URL_SANDBOX
	if environment == "live":
		url = BASE_URL_UAT
	resource = "/api/"+route

	try:
		conn = http.client.HTTPSConnection(url,timeout=10)
		headers = {'Content-type': 'application/json',"Authorization":"Bearer "+accesskey+":"+secretkey}
		conn.request("GET", resource,"",headers)
		resp = conn.getresponse()
		rdata = resp.read().decode('utf-8')
		conn.close()
		response = json.loads(rdata)
	except Exception as ex:
		print(ex)

	return response

def create_hash(data,accesskey,secretkey):
	hash=""
	try:
		pipeSeperatedString=accesskey+"|"+str(data["amount"])+"|"+data["id"]+"|"+str(data["mtx"])
		signature = hmac.new(
			bytes(secretkey , 'latin-1'),
			msg = bytes(pipeSeperatedString , 'latin-1'),
			digestmod = hashlib.sha256).hexdigest().upper()

		base64_bytes = base64.b64encode(signature.encode('ascii'))
		hash = base64_bytes.decode('ascii')

	except Exception as ex:
		hash = ex

	return hash

def verify_hash(data,rec_hash,accesskey,secretkey):
	gen_hash = create_hash(data,accesskey,secretkey)
	if gen_hash == rec_hash:
		return True
	else:
		return False


import math

import torch

def seq_to_graph(x, stride=None):
    seq_len = x.size(0)
    if stride is None:
        stride = max(math.ceil(math.log(seq_len)), 1)

    src = []
    tgt = []
    for i in range(seq_len):
        num_neighbors = min(stride, seq_len - i)
        src += [i] * num_neighbors
        tgt += [i + k for k in range(num_neighbors)]

    edge_index = torch.tensor([src, tgt], device=x.device, dtype=torch.long)
    return edge_index


# Autor: Hector Fernandes

import requests as r
import sys as s
from datetime import datetime

emojis = {
	'trânsito': '🚚',
	'entregue': '✅',
	'Correios': '🛬',
	'aduaneira': '🚨',
	'exportação': '🛫',
	'erro': '❌',
	'postado': '📦',
	'saiu': '🚀',
	'confirmado' : '💸',
	'pagamento': '💰',
	'indo' : '➡'
}

def get_status(codigo):
	response = r.get('https://proxyapp.correios.com.br/v1/sro-rastro/' + codigo)
	if response.status_code == 200:
		obj = response.json()['objetos'][0]
		if 'mensagem' in obj:
			if "SRO-019: Objeto inválido" in obj['mensagem']:
				return None
		return obj
	else:
		return None

def create_status_string(item):
	unidade = item['unidade']
	if 'nome' in unidade:
		local = unidade['nome']
	else:
		local = unidade['endereco']['cidade'] + '-' + unidade['endereco']['uf']
	if 'unidadeDestino' in item:
		unidadeDestino = item['unidadeDestino']
		if 'nome' in unidadeDestino:
			dest = unidadeDestino['nome']
		else:
			dest = unidadeDestino['endereco']['cidade'] + '-' + unidadeDestino['endereco']['uf']
		return (f'{local} {emojis["indo"]}  {dest}')
	else:
		return (local)

def create_event_string(item):
	date = item['dtHrCriado']
	hour = date[11:16]
	date = date[8:10] + '/' + date[5:7] + '/' + date[0:4]
	desc = item['descricao']
	string = date + ' - ' + hour + ' - ' + desc
	for emoji in emojis:
		if emoji in desc:
			string = emojis[emoji] + ' ' + string
	status_string = create_status_string(item)
	string = string + " | " + status_string
	return string

def print_status(status):
	events = status['eventos']
	for item in events:
		print(create_event_string(item))

def convert_date(date):
	date = date[8:10] + '/' + date[5:7] + '/' + date[0:4] + ' ' + date[11:16]
	date = datetime.strptime(date, '%d/%m/%Y %H:%M')
	return date

def print_all(status):
	last_date = convert_date(status['eventos'][0]['dtHrCriado'])
	status_date = datetime.now() - last_date
	status_date = status_date.days

	if status_date < 2 and status_date != 0:
		status_date = f'Última atualização há {status_date} dia'
		print(f'\n{emojis["postado"]} Cód: {status["codObjeto"]} - {status_date}')
	else:
		status_date = f'Última atualização há {status_date} dias'
		print(f'\n{emojis["postado"]} Cód: {status["codObjeto"]} - {status_date}')

	print_status(status)

def main():

	if len(s.argv) < 2:
		print('Uso: rastreador "<codigo>" ...')
		return

	for i in range(1, len(s.argv)):
		code = s.argv[i]
		status = get_status(code)

		if status is None:
			print(f'\n{emojis["erro"]} Código {code} inválido')
		else:
			print_all(status)

main()


from scipy import stats
import copy
import time
import pandas as pd
import numpy as np
from scipy import stats

'''
Info:
    - Python slicing
        Example: a[start:stop:step] # start through not past stop, by step (https://stackoverflow.com/questions/509211/understanding-slicing)
    - Label indexing (with hatchet node as zero index)
        hatchet_node = hatchet.node.Node({'name': '<program root>', 'type': 'function'})
        print(hatchet_node)
        print(type(hatchet_node))
        print(df[hatchet_node, 0]) # <- Does not work :(
    - location indexing
        print(df.iloc[0::step_size]) # Get every node's rank 0
        print(df.iloc[1::step_size]) # Get every node's rank 1
        print(df.index.levels[0]) # node
        print(df.index.levels[1]) # rank
'''

def slice_statistics(a_slice, nobs, N_TRAVERSALS):
    '''Computes statistics on a slice (pandas.core.series.Series). Should be a numeric slice so sum can be computed.

    Args:
        a_slice: Slice of pandas dataframe.
        nobs: Number of observations.

    Returns:
        ?
    '''
    num_rows = len(a_slice)
    deep_slice = copy.deepcopy(a_slice)
    t_arr = []
    dt_arr = []
    bt_arr = []
    bdt_arr = []

    # traversals
    for i in range(nobs):
        start = time.time_ns()
        for index, value in a_slice.iteritems():
            value
        end = (time.time_ns() - start)
        adj_metric = (end*N_TRAVERSALS)/(num_rows)
        t_arr.append(adj_metric)

        start = time.time_ns()
        for index, value in deep_slice.iteritems():
            value
        end = (time.time_ns() - start)
        adj_metric = (end*N_TRAVERSALS)/(num_rows)
        dt_arr.append(adj_metric)

    # built-in
    for i in range(nobs):
        start = time.time_ns()
        s_sum = a_slice.sum()
        end = time.time_ns() - start
        adj_metric = (end*N_TRAVERSALS)/(num_rows)
        bt_arr.append(adj_metric)

        start = time.time_ns()
        ds_sum = deep_slice.sum()
        end = time.time_ns() - start
        adj_metric = (end*N_TRAVERSALS)/(num_rows)
        bdt_arr.append(adj_metric)

        assert(s_sum==ds_sum)

    traversal_stats = stats.describe(t_arr)
    dtraversal_stats = stats.describe(dt_arr)
    built_stats = stats.describe(bt_arr)
    dbuilt_stats = stats.describe(bdt_arr)

    return traversal_stats, dtraversal_stats, built_stats, dbuilt_stats

def slicing_tests(df, nobs, X, N_TRAVERSALS, EXTRA_OUT, DEBUG):
    '''
    Args:
        df: A 2-level multi-indexed pandas dataframe.
        nobs: Number of observations.
    '''
    step_size = df.index.levels[1].size # number of elements in the second level (level 0).

    rank_zero_times = df.iloc[0::step_size]['time']
    rzt_stats, drzt_stats, rzt_built_stats, drzt_built_stats = slice_statistics(rank_zero_times, nobs, N_TRAVERSALS)

    rank_one_times = df.iloc[1::step_size]['time']
    rot_stats, drot_stats, rot_built_stats, drot_built_stats = slice_statistics(rank_one_times, nobs, N_TRAVERSALS)

    rank_ten_times = df.iloc[10::step_size]['time']
    rtt_stats, drtt_stats, rtt_built_stats, drtt_built_stats = slice_statistics(rank_ten_times, nobs, N_TRAVERSALS)

    rank_zero_ten_times = df.iloc[0::10]['time']
    rztt_stats, drztt_stats, rztt_built_stats, drztt_built_stats = slice_statistics(rank_ten_times, nobs, N_TRAVERSALS)

    if EXTRA_OUT:
        slices_dic = {"Rank Zero Times": rank_zero_times, "Rank One Times": rank_one_times, "Rank Ten Times": rank_ten_times, "Rank Zero through Ten Times": rank_zero_ten_times}
        for k in slices_dic:
            print(f"--- {k} ---\n{slices_dic[k]}\n")
        if DEBUG: # See the raw statistics.
            stats_list = [rzt_stats, drzt_stats, rzt_built_stats, drzt_built_stats,
                rot_stats, drot_stats, rot_built_stats, drot_built_stats,
                rtt_stats, drtt_stats, rtt_built_stats, drtt_built_stats,
                rztt_stats, drztt_stats, rztt_built_stats, drztt_built_stats]
            for i in stats_list:
                print(i)

    print(f"*Entries are relative time measurements per-{N_TRAVERSALS} element accesses. X={X}. nobs={nobs}.")
    eval_list=[f'Nodes=all, Rank=0 ({len(rank_zero_times)}x1)',
        f'Nodes=all, Rank=1 ({len(rank_one_times)}x1)',
        f'Nodes=all, Rank=10 ({len(rank_ten_times)}x1)',
        f'Nodes=all, Rank=10x ({len(rank_zero_ten_times)}x1)']
    time_data = {"Slice": [rzt_stats, rot_stats, rtt_stats, rztt_stats], "Deep Copy": [drzt_stats, drot_stats, drtt_stats, drztt_stats]}
    built_data = {"Slice": [rzt_built_stats, rot_built_stats, rtt_built_stats, rztt_built_stats], "Deep Copy": [drzt_built_stats, drot_built_stats, drtt_built_stats, drztt_built_stats]}
    X_lambda_func = lambda x: f"{round(x.mean/X)}*X" if (not np.isnan(x.mean)) else x.mean
    df_time = pd.DataFrame(data=time_data, index=eval_list)
    df_time = df_time.applymap(X_lambda_func)
    print(f"Traversal Times\n{df_time}")
    df_built = pd.DataFrame(data=built_data, index=eval_list)
    df_built = df_built.applymap(X_lambda_func)
    print(f"Built-in (Sum)\n{df_built}")


#### Create Histograms using Numpy library
#import the necessary libraries and load data files

from astropy.table import Table, Column, vstack, MaskedColumn

import matplotlib
import matplotlib.ticker as mtick
import matplotlib.pyplot as plt
import numpy as np

all_tbl = Table.read('./data/all_table.csv', format='csv')

all_tbl['KT_high'].fill_value = -99.9
all_tbl['KT_low'].fill_value = -99.9

all_tbl = all_tbl.filled()

msk_kt = all_tbl['KT_high'] < 0
lo_tbl = all_tbl[msk_kt]
go_tbl = all_tbl[~msk_kt]

#create KT histogram

hist, edges = np.histogram(go_tbl['KT'],bins=7,normed=True)
center = np.array([edges+(12)])
height = hist*21.5
print(type(edges))
print(edges)
print(center)
print(height)
center_2 = np.array([20.57, 42.19, 63.80, 85.42, 128.66, 150.28, 171.9])

x = edges[:-1]

plt.bar(edges[:-1], hist.astype(np.float32)/hist.sum(), width=(edges[1]-edges[0]),
            #align='center',
            bottom=None,
            color='none',alpha=1,edgecolor='black',linewidth=2,linestyle='dashed',
            label='good')
#
plt.show()


from typing import Optional
from unittest import TestCase

import pytest
from flexmock import flexmock

from ogr import GithubService
from ogr.abstract import AuthMethod
from ogr.exceptions import GithubAPIException
from ogr.services.github.auth_providers.token import TokenAuthentication
from ogr.services.github.auth_providers.tokman import Tokman
from ogr.services.github.check_run import (
    GithubCheckRunOutput,
    create_github_check_run_output,
)
from ogr.services.github.project import GithubProject
from ogr.services.github.pull_request import GithubPullRequest

@pytest.fixture()
def github_project(mock_github_repo):
    github_project = GithubProject(
        repo="test_repo",
        service="test_service",
        namespace="fork_username",
    )
    parent_github_project = GithubProject(
        repo="test_parent_repo",
        service="test_service",
        namespace="test_parent_namespace",
    )
    flexmock(github_project)
    flexmock(parent_github_project)
    flexmock(GithubPullRequest)

    github_project.should_receive("github_repo").and_return(mock_github_repo())
    parent_github_project.should_receive("github_repo").and_return(mock_github_repo())
    github_project.should_receive("parent").and_return(parent_github_project)
    return github_project

@pytest.fixture()
def mock_pull_request():
    def mock_pull_request_factory(id):
        return flexmock(id=id)

    return mock_pull_request_factory

@pytest.fixture()
def mock_github_repo(mock_pull_request):
    def mock_github_repo_factory():
        return flexmock(create_pull=mock_pull_request(42))

    return mock_github_repo_factory

class TestGithubProject:
    @pytest.mark.parametrize(
        "fork_username",
        [
            pytest.param("fork_username", id="fork_username_set"),
            pytest.param(None, id="fork_username_None"),
        ],
    )
    def test_pr_create_is_not_fork(self, github_project, fork_username):
        github_project.should_receive("is_fork").and_return(False)
        GithubPullRequest.should_receive("__init__").and_return()

        head = ":".join(filter(None, [fork_username, "master"]))

        github_project.github_repo.should_call("create_pull").with_args(
            title="test_title",
            body="test_content",
            base="master",
            head=head,
        )
        github_project.parent.github_repo.should_call("create_pull").never()
        github_project.github_repo.should_call("create_pull").once()

        github_project.create_pr(
            title="test_title",
            body="test_content",
            target_branch="master",
            source_branch="master",
            fork_username=fork_username,
        )

    @pytest.mark.parametrize(
        "fork_username",
        [pytest.param("fork_username", id="fork_username_set")],
    )
    def test_pr_create_is_fork(self, github_project, fork_username):
        github_project.should_receive("is_fork").and_return(True)
        GithubPullRequest.should_receive("__init__").and_return()

        github_project.parent.github_repo.should_call("create_pull").with_args(
            title="test_title",
            body="test_content",
            base="master",
            head=f"{github_project}:master",
            fork_username=fork_username,
        )
        github_project.parent.github_repo.should_call("create_pull").never()
        github_project.github_repo.should_call("create_pull").once()

        github_project.create_pr(
            title="test_title",
            body="test_content",
            target_branch="master",
            source_branch="master",
            fork_username=fork_username,
        )

class TestGitHubService(TestCase):
    def test_hostname(self):
        assert GithubService().hostname == "github.com"

@pytest.mark.parametrize(
    ("title", "summary", "text", "expected"),
    [
        (
            "test",
            "test summary",
            None,
            {
                "title": "test",
                "summary": "test summary",
            },
        ),
        (
            "bigger output",
            "no summary",
            "# Random title\n\n- [ ] TODO list\n---\n_italics_",
            {
                "title": "bigger output",
                "summary": "no summary",
                "text": "# Random title\n\n- [ ] TODO list\n---\n_italics_",
            },
        ),
    ],
)
def test_create_github_check_run_output(
    title: str,
    summary: str,
    text: Optional[str],
    expected: GithubCheckRunOutput,
) -> None:
    assert create_github_check_run_output(title, summary, text) == expected

@pytest.fixture()
def github_service_with_multiple_auth_methods():
    return GithubService(
        token="abcdef",
        github_app_id="123",
        github_app_private_key="id_rsa",
        github_app_private_key_path="/path",
        tokman_instance_url="http://tokman:8080",
        github_authentication=None,
    )

def test_multiple_auth_methods_default_is_tokman(
    github_service_with_multiple_auth_methods,
):
    service = github_service_with_multiple_auth_methods
    assert isinstance(service.authentication, Tokman)

def test_set_reset_customized_auth_method(github_service_with_multiple_auth_methods):
    service = github_service_with_multiple_auth_methods
    assert isinstance(service.authentication, Tokman)
    service.set_auth_method(AuthMethod.token)
    assert isinstance(service.authentication, TokenAuthentication)
    service.reset_auth_method()
    assert isinstance(service.authentication, Tokman)

@pytest.fixture()
def github_service_with_one_auth_method():
    return GithubService(
        tokman_instance_url="http://tokman:8080",
        github_authentication=None,
    )

def test_no_multiple_auth_methods_default_is_tokman(
    github_service_with_one_auth_method,
):
    service = github_service_with_one_auth_method
    assert isinstance(service.authentication, Tokman)

def test_no_set_reset_customized_auth_method(github_service_with_one_auth_method):
    service = github_service_with_one_auth_method
    assert isinstance(service.authentication, Tokman)
    with pytest.raises(GithubAPIException):
        service.set_auth_method(AuthMethod.github_app)
    assert isinstance(service.authentication, Tokman)


import os
from contextlib import contextmanager
from unittest import mock

import pytest

from detect_secrets_server.actions import add_repo
from detect_secrets_server.actions import initialize
from detect_secrets_server.core.usage.parser import ServerParserBuilder
from detect_secrets_server.storage.base import BaseStorage
from testing.factories import metadata_factory
from testing.factories import single_repo_config_factory
from testing.mocks import mock_git_calls
from testing.mocks import SubprocessMock
from testing.util import cache_buster

class TestInitialize:
    def teardown(self):
        cache_buster()

    @staticmethod
    def parse_args(argument_string='', has_s3=False):
        base_argument = (
            'add will_be_mocked --config '
        )
        if has_s3:
            base_argument += '--s3-config examples/s3.yaml '

        with mock.patch(
            'detect_secrets_server.core.usage.s3.should_enable_s3_options',
            return_value=has_s3,
        ):
            return ServerParserBuilder().parse_args(
                (base_argument + argument_string).split()
            )

    def test_no_tracked_repos(self):
        with mock_repos_config({
            'tracked': [],
        }):
            args = self.parse_args()

        assert not initialize(args)

    def test_simple_success(self, mock_rootdir):
        with mock_repos_config({
            'tracked': [
                single_repo_config_factory(
                    'git@github.com:yelp/detect-secrets',
                ),
            ]
        }), mock_repo_class(
            'BaseTrackedRepo'
        ) as repo_class:
            args = self.parse_args(
                '--root-dir {}'.format(mock_rootdir)
            )
            initialize(args)

            kwargs = repo_class.call_args[1]
            assert kwargs['repo'] == 'git@github.com:yelp/detect-secrets'
            assert kwargs['sha'] == ''
            assert kwargs['crontab'] == '0 0 * * *'
            assert kwargs['rootdir'] == mock_rootdir

    @pytest.mark.parametrize(
        'data,expected_repo_class',
        [
            (
                {
                    'is_local_repo': True,
                    'repo': 'examples',
                },
                'LocalTrackedRepo',
            ),
            (
                {
                    'storage': 's3',
                    'repo': 'git@github.com:yelp/detect-secrets',
                },
                'S3TrackedRepo',
            ),
            (
                {
                    'is_local_repo': True,
                    'repo': 'examples',
                    'storage': 's3',
                },
                'S3LocalTrackedRepo',
            ),
        ]
    )
    def test_flags_set_tracked_repo_classes(self, data, expected_repo_class):
        with mock_repos_config({
            'tracked': [
                single_repo_config_factory(
                    **data
                ),
            ]
        }):
            args = self.parse_args(has_s3=data.get('storage') == 's3')

        with mock_repo_class(expected_repo_class) as repo_class:
            initialize(args)
            assert repo_class.called

    def test_repo_config_overrides_defaults(self, mock_rootdir):
        with mock_repos_config({
            'tracked': [
                single_repo_config_factory(
                    'git@github.com:yelp/detect-secrets',
                    plugins={
                        # This checks that CLI overrides config file
                        'HexHighEntropyString': {
                            'hex_limit': 5,
                        },

                        # This checks it overrides default values
                        'Base64HighEntropyString': {
                            'base64_limit': 2,
                        },

                        # This checks for disabling functionality
                        'PrivateKeyDetector': False,
                    },

                    # This checks it overrides CLI (non-plugin)
                    baseline_filename='will_be_overriden',

                    # This checks it overrides default value (non-plugin)
                    exclude_regex='something_here',
                    crontab='* * 4 * *',
                )
            ],
        }):
            args = self.parse_args(
                '--hex-limit 4 '
                '--baseline baseline.file '
                '--root-dir {}'.format(mock_rootdir)
            )

        with mock_repo_class('BaseTrackedRepo') as repo_class:
            initialize(args)

            kwargs = repo_class.call_args[1]
            assert kwargs['repo'] == 'git@github.com:yelp/detect-secrets'
            assert kwargs['sha'] == ''
            assert kwargs['crontab'] == '* * 4 * *'
            # NOTE: This is disabled, since it's `False` above.
            assert 'PrivateKeyDetector' not in kwargs['plugins']
            assert kwargs['plugins']['Base64HighEntropyString']['base64_limit'] == 2.0
            assert kwargs['plugins']['HexHighEntropyString']['hex_limit'] == 4.0
            assert kwargs['rootdir'] == mock_rootdir
            assert kwargs['baseline_filename'] == 'baseline.file'
            assert kwargs['exclude_regex'] == 'something_here'

class TestAddRepo:
    @staticmethod
    def parse_args(argument_string='', has_s3=False):
        with mock.patch(
            'detect_secrets_server.core.usage.s3.should_enable_s3_options',
            return_value=has_s3,
        ):
            return ServerParserBuilder().parse_args(
                argument_string.split()
            )

    def teardown(self):
        cache_buster()

    def test_add_non_local_repo(self, mock_file_operations, mock_rootdir):
        self.add_non_local_repo(mock_rootdir)
        mock_file_operations.write.assert_called_with(
            metadata_factory(
                repo='git@github.com:yelp/detect-secrets',
                sha='mocked_sha',
                json=True,
            ),
        )

    def test_override_meta_tracking_if_already_exists(
        self,
        mock_file_operations,
        mock_rootdir,
    ):
        with mock.patch(
            'detect_secrets_server.storage.file.FileStorage.get_tracked_file_location',

            # This doesn't matter what it is, just that it exists.
            return_value='examples/config.yaml',
        ):
            self.add_non_local_repo(mock_rootdir)

        assert mock_file_operations.write.called

    def add_non_local_repo(self, mock_rootdir):
        repo = 'git@github.com:yelp/detect-secrets'
        directory = '{}/repos/{}'.format(
            mock_rootdir,
            BaseStorage.hash_filename('yelp/detect-secrets'),
        )

        git_calls = [
            SubprocessMock(
                expected_input='git clone {} {} --bare'.format(repo, directory),
            ),
            SubprocessMock(
                expected_input='git rev-parse HEAD',
                mocked_output='mocked_sha',
            ),
        ]

        with mock_git_calls(*git_calls):
            args = self.parse_args('add {} --root-dir {}'.format(repo, mock_rootdir))
            add_repo(args)

    def test_add_local_repo(self, mock_file_operations, mock_rootdir):
        # This just needs to exist; no actual operations will be done to this.
        repo = 'examples'

        git_calls = [
            # repo.update
            SubprocessMock(
                expected_input='git rev-parse HEAD',
                mocked_output='mocked_sha',
            ),
        ]

        with mock_git_calls(*git_calls):
            args = self.parse_args(
                'add {} --baseline .secrets.baseline --local --root-dir {}'.format(
                    repo,
                    mock_rootdir,
                )
            )

            add_repo(args)

        mock_file_operations.write.assert_called_with(
            metadata_factory(
                sha='mocked_sha',
                repo=os.path.abspath(
                    os.path.join(
                        os.path.dirname(__file__),
                        '../../examples',
                    ),
                ),
                baseline_filename='.secrets.baseline',
                json=True,
            ),
        )

    def test_add_s3_backend_repo(self, mock_file_operations, mocked_boto):
        args = self.parse_args(
            'add {} '
            '--local '
            '--storage s3 '
            '--s3-credentials-file examples/aws_credentials.json '
            '--s3-bucket pail'.format('examples'),
            has_s3=True,
        )

        git_calls = [
            # repo.update
            SubprocessMock(
                expected_input='git rev-parse HEAD',
                mocked_output='mocked_sha',
            ),
        ]

        with mock_git_calls(
            *git_calls
        ):
            mocked_boto.list_objects_v2.return_value = {}
            add_repo(args)

@contextmanager
def mock_repos_config(data):
    """Unfortunately, mocking this means that we can't test more than
    one config file at a time. However, all consolidation tests with
    --config-file should have been done in usage_test, so we should
    be OK.
    """
    with mock.patch(
        'detect_secrets_server.core.usage.add.config_file',
        return_value=data,
    ):
        yield

@contextmanager
def mock_repo_class(classname):
    """
    :type classname: str
    """
    with mock.patch(
        'detect_secrets_server.repos.factory.{}'.format(classname),
    ) as repo_class:
        yield repo_class


import torch
import torch.optim as optim
import pytest
import helpers
import poptorch

#  Test the reductions work as expected
@pytest.mark.parametrize("reduction", ["none", "mean", "sum"])
def test_non_final_loss_reductions(reduction):
    torch.manual_seed(42)

    base_model = torch.nn.Linear(10, 10)

    class CustomLoss(torch.nn.Module):
        # Mean squared error scaled.
        def forward(self, x, target):
            partial_loss = poptorch.identity_loss(x - target,
                                                  reduction=reduction)
            loss = partial_loss * partial_loss * 5
            return partial_loss, poptorch.identity_loss(loss, reduction="mean")

    loss_fn = CustomLoss()

    class ModelWithLoss(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.base_model = base_model

        def forward(self, data, target):
            out = base_model(data)
            loss = loss_fn(out, target)
            return out, loss

    model = ModelWithLoss()
    poptorch_model = poptorch.trainingModel(model)

    target = torch.randn(10)
    input = torch.randn(10)

    # Capture what the loss function will see before the loss changes.
    x, _ = model(input, target)
    _, (partial_loss, _) = poptorch_model(input, target)

    # Check we have actually reduced the loss
    if reduction != "none":
        assert torch.numel(partial_loss) == 1

    if reduction == "mean":
        simulated_loss = torch.mean(x - target)
    elif reduction == "sum":
        simulated_loss = torch.sum(x - target)
    elif reduction == "none":
        simulated_loss = x - target

    helpers.assert_allclose(expected=simulated_loss.reshape_as(partial_loss),
                            actual=partial_loss,
                            rtol=1e-02,
                            atol=1e-02)

# Test custom loss by training to targets
def run_custom_loss_test(loss_fn,
                         base_model=None,
                         input=None,
                         target=None,
                         test_output_vs_target=True):
    torch.manual_seed(42)

    if base_model is None:
        base_model = torch.nn.Linear(10, 10)
    if input is None:
        input = torch.randn(1, 10)
    if target is None:
        target = torch.randint(0, 10, [1])

    class ModelWithLoss(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.base_model = base_model
            self.loss_fn = loss_fn

        def forward(self, data, target):
            out = base_model(data)
            loss = self.loss_fn(out, target)
            return out, loss

    model = ModelWithLoss()
    poptorch_model = poptorch.trainingModel(model)

    # Pytorch native.
    native_out, loss = model(input, target)

    #Make sure the first run doesn't already pass the test.
    original, original_loss = poptorch_model(input, target)

    assert original_loss > 0.1

    if test_output_vs_target:
        assert not torch.allclose(native_out, target, rtol=1e-02, atol=1e-02)
        assert not torch.allclose(original, target, rtol=1e-02, atol=1e-02)

    for _ in range(0, 2500):
        out, loss = poptorch_model(input, target)

    # Check we have trained the "model"
    assert loss < 0.1

    if test_output_vs_target:
        helpers.assert_allclose(actual=out,
                                expected=target,
                                rtol=1e-02,
                                atol=1e-02)

        # Check that the pytorch native model is also returning the trained
        # value that was trained on IPU.
        out, _ = model(input, target)
        helpers.assert_allclose(actual=out,
                                expected=target,
                                rtol=1e-02,
                                atol=1e-02)

    return poptorch_model

def test_custom_loss():
    torch.manual_seed(42)

    class CustomLoss(torch.nn.Module):
        # Mean squared error scaled.
        def forward(self, x, target):
            loss = poptorch.identity_loss(x - target, reduction="none")
            loss = loss * loss * 5.0
            return poptorch.identity_loss(loss, reduction="mean")

    run_custom_loss_test(loss_fn=CustomLoss(),
                         input=torch.randn(10),
                         target=torch.randn(10))

def test_custom_loss_l1():
    torch.manual_seed(42)

    class CustomLoss(torch.nn.Module):
        # Mean squared error scaled.
        def forward(self, x, target):
            loss = torch.nn.functional.l1_loss(x, target)
            loss = loss * loss * 5.0
            return poptorch.identity_loss(loss, reduction="mean")

    run_custom_loss_test(loss_fn=CustomLoss(),
                         input=torch.randn(10),
                         target=torch.randn(10))

def test_custom_loss_nll():
    torch.manual_seed(42)

    class CustomLoss(torch.nn.Module):
        # Mean squared error scaled.
        def forward(self, x, target):
            loss = torch.nn.functional.nll_loss(x, target)
            loss = loss * 5.0
            return poptorch.identity_loss(loss, reduction="mean")

    base_model = torch.nn.Sequential(torch.nn.Linear(10, 10),
                                     torch.nn.LogSoftmax(dim=1))

    input = torch.randn(1, 10)
    target = torch.randint(0, 10, [1])

    out = base_model(input)

    model = run_custom_loss_test(loss_fn=CustomLoss(),
                                 base_model=base_model,
                                 input=input,
                                 target=target,
                                 test_output_vs_target=False)
    model.copyWeightsToHost()

    # Check that the pytorch native model is also returning the trained
    # value that was trained on IPU.
    out = base_model(input)

    assert torch.argmax(out, dim=1) == target

def test_two_custom_losses():
    torch.manual_seed(42)

    base_model = torch.nn.Sequential(torch.nn.Linear(10, 10),
                                     torch.nn.LogSoftmax(dim=1))

    class CustomLoss(torch.nn.Module):
        # Mean squared error scaled.
        def forward(self, x, target):
            loss = torch.nn.functional.nll_loss(x, target)
            loss2 = torch.nn.functional.nll_loss(x, target) * 5.0
            return loss + loss2

    error_msg = ("Multiple independent losses found in graph. "
                 "Graph must have one final loss. "
                 "Wrap final graph loss in poptorch.identity_loss.")
    with pytest.raises(poptorch.Error, match=error_msg):
        run_custom_loss_test(loss_fn=CustomLoss(), base_model=base_model)

def test_two_custom_losses_with_id_wrapper():
    torch.manual_seed(42)

    base_model = torch.nn.Sequential(torch.nn.Linear(10, 10),
                                     torch.nn.LogSoftmax(dim=1))

    class CustomLoss(torch.nn.Module):
        # Mean squared error scaled.
        def forward(self, x, target):
            loss = torch.nn.functional.nll_loss(x, target)
            loss2 = torch.nn.functional.nll_loss(x, target) * 5.0
            return poptorch.identity_loss(loss + loss2, reduction="mean")

    run_custom_loss_test(loss_fn=CustomLoss(),
                         base_model=base_model,
                         test_output_vs_target=False)

def test_no_loss():
    torch.manual_seed(42)

    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.model = torch.nn.Sequential(torch.nn.Linear(10, 10),
                                             torch.nn.LogSoftmax(dim=1))

        # Mean squared error scaled.
        def forward(self, x, target):
            fwd = self.model(x)
            loss = fwd * 12
            loss2 = target + 1
            a = loss + loss2
            return fwd, a, loss

    model = Model()
    optimizer = optim.SGD(model.parameters(), lr=0.01)

    poptorch_model = poptorch.trainingModel(model, optimizer=optimizer)

    label = torch.randint(0, 10, [1])
    input = torch.randn(1, 10)

    error_msg = "Couldn't find a loss in graph"
    with pytest.raises(poptorch.Error, match=error_msg):
        _ = poptorch_model(input, label)


import copy
import sys

import sqlalchemy as sa

from dnrm.db.sqlalchemy import models
from dnrm.exceptions import db as exceptions
from dnrm.openstack.common.db.sqlalchemy import session as db_session

def get_backend():
    """The backend is this module itself."""
    return sys.modules[__name__]

def db_create():
    models.create_db()

def db_drop():
    models.drop_db()

def db_cleanup():
    db_session.cleanup()

def model_query(model, session=None):
    session = session or db_session.get_session()
    query = session.query(model)
    return query

def falsy(value):
    return bool(value) and (not isinstance(value, (str, unicode)) or
                            value.lower() != 'false')

def filters_to_condition(model, filter_fields, filter_values):
    filter_values = copy.deepcopy(filter_values)
    if 'class' in filter_values:
        filter_values['klass'] = filter_values.pop('class')
    and_list = []
    if 'unused' in filter_values:
        if falsy(filter_values.pop('unused')):
            and_list.append(model.pool == None)
        else:
            and_list.append(model.pool != None)
    for key in filter_fields:
        column = getattr(model, key)
        if key not in filter_values:
            continue
        value = filter_values.pop(key)
        if isinstance(column.property.columns[0].type, sa.Boolean):
            value = falsy(value)
        if isinstance(value, (list, tuple, set)):
            expr = column.in_(set(value))
        else:
            expr = (column == value)
        and_list.append(expr)
    if and_list:
        return sa.and_(*and_list)
    else:
        return None

###############################################################################
# Resources

def _resource_to_dict(resource):
    resource = dict(resource)
    resource['class'] = resource.pop('klass')
    data = resource.pop('data', {})
    resource.update(data)
    resource['unused'] = resource['pool'] is None
    return resource

def _update_resource(resource, values):
    values = copy.deepcopy(values)
    for key in ('id', 'unused'):
        if key in values:
            del values[key]
    if 'class' in values:
        values['klass'] = values.pop('class')
    validated_values = {}
    for key in models.Resource.FILTER_FIELDS:
        try:
            validated_values[key] = values.pop(key)
        except KeyError:
            pass
    if values:
        data = copy.deepcopy(resource['data']) or {}
        data.update(values)
        validated_values['data'] = data
    resource.update(validated_values)

def resource_create(driver_name, values):
    resource = models.Resource()
    _update_resource(resource, values)
    resource['type'] = driver_name
    resource.save()
    return _resource_to_dict(resource)

def _resource_get_by_id(id, session=None):
    task = (model_query(models.Resource, session=session)
            .filter_by(id=id)
            .first())
    if not task:
        raise exceptions.ResourceNotFound(id=id)
    return task

def resource_get_by_id(id):
    return _resource_to_dict(_resource_get_by_id(id))

def resource_update(id, values):
    session = db_session.get_session()
    with session.begin():
        resource = _resource_get_by_id(id, session=session)
        _update_resource(resource, values)
        return _resource_to_dict(resource)

def resource_delete(id):
    count = (model_query(models.Resource)
             .filter_by(id=id)
             .delete())
    if not count:
        raise exceptions.ResourceNotFound(id=id)

def make_query(model, search_opts, session=None):
    search_opts = copy.deepcopy(search_opts)

    filters = search_opts.pop('filters', {})
    limit = search_opts.pop('limit', None)
    offset = search_opts.pop('offset', None)

    if search_opts:
        raise ValueError(_('Unexpected search options: %(options)s'),
                         options=', '.join(search_opts.keys()))

    query = model_query(models.Resource, session=session)

    condition = filters_to_condition(model, model.FILTER_FIELDS, filters)

    if condition is not None:
        query = query.filter(condition)

    if offset is not None:
        query = query.offset(limit)

    if limit is not None:
        query = query.limit(limit)

    return query

def resource_find(search_opts):
    query = make_query(models.Resource, search_opts)
    return [_resource_to_dict(resource) for resource in query.all()]

def resource_count(search_opts):
    query = make_query(models.Resource, search_opts)
    return query.count()

def resource_compare_update(id, filters, values):
    session = db_session.get_session()
    with session.begin():
        filters = copy.deepcopy(filters)
        query = make_query(models.Resource, {'filters': filters}, session)
        query = query.filter(models.Resource.id == id)
        resource = query.first()
        if resource:
            _update_resource(resource, values)
            return _resource_to_dict(resource)
        else:
            return None


from django.shortcuts import render
from django.views.generic.base import TemplateView
from django.http import HttpResponse
from django.http import HttpResponseRedirect
from django.urls import reverse
from .models import Reservation
from .models import Table

from datetime import datetime
from time import strptime
import re

# Create your views here.
class IndexView(TemplateView):
    template_name = "reservations/index.html"
    confirmation_page = "reservations/confirmation.html"

    def __init__(self):
        self.reservation_concept = Reservation()
        self.reservation_checker = ReservationChecker()
        self.reservation_table = None

    def get_context_data(self, **kwargs):
        context = super().get_context_data(**kwargs)
        # context['variable'] = self.value
        print('context data called')
        return context

    def validate_email(self, email):
         if re.match("(^[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+$)", email) != None:
             return True
         return False

    def save_reservation(self):
        self.reservation_concept.save()

    def post(self,request,*args,**kwargs):
        #check availability
        context = super().get_context_data(**kwargs)

        if "reservation_check" in request.POST:
            print('this is a check only')
            response_for_customer_options = [ "Our restaurant is closed on this date.",
                                        "Our restaurant has no free tables for this number of guests at the selected time.",
                                        "The selected date and time are available! Please confirm your booking."]

            page_actions = ["request_another_date",
                            "request_another_time",
                            "request_booking_confirmation"]
            self.collect_reservation_data_from_post_request(request)
            self.format_reservation(request)

            datetime_valid = self.future_datetime(self.reservation_concept.start_date_time)

            email_valid = self.validate_email(self.email)
            #print('dt',datetime(self.reservation_concept.start_date_time))

            reservation_table = self.reservation_checker.check_reservation_possible(self.selected_num_people,
                                                                                self.reservation_concept.start_date_time,
                                                                                self.reservation_concept.end_date_time)
            if reservation_table:
                response_for_customer = response_for_customer_options[2]
                page_action = page_actions[2]
            else:
                response_for_customer = response_for_customer_options[1]
                page_action = page_actions[1]

            variable_content = {"response":response_for_customer,
                                "page_action":page_action,
                                "selected_date":self.selected_date,
                                "selected_time_hours":self.selected_time_hours,
                                "selected_time_minutes":self.selected_time_minutes,
                                "selected_num_people":self.selected_num_people,
                                "email":self.email,
                                "email_valid":email_valid,
                                "phone":self.phone,
                                "selected_num_people":self.selected_num_people,
                                "first_name":self.first_name,
                                "last_name":self.last_name,
                                "remark":self.remark}

            return render(request, IndexView.template_name, variable_content)

        elif "reservation_confirmation" in request.POST:
            # check again if all is available
            self.collect_reservation_data_from_post_request(request)
            self.format_reservation(request)
            self.save_reservation()
            return render(request, self.confirmation_page)
        else:
            return render(request, self.template_name)

    def collect_reservation_data_from_post_request(self,request):
        self.selected_date = request.POST.get("date","")
        self.selected_time_hours = request.POST.get("time-hour","12")
        self.selected_time_minutes = request.POST.get("time-minutes","00")
        self.selected_num_people = request.POST.get("people", "2")
        self.email = request.POST.get("email","")
        self.phone = request.POST.get("phone","")
        self.first_name = request.POST.get("first_name","")
        self.last_name = request.POST.get("last_name","")
        self.remark = request.POST.get("remark","")

    def future_datetime(self, rdatetime):
        return rdatetime >= datetime.now()

    def format_reservation(self,request):
        reservation_name = self.first_name + " " + self.last_name
        reservation_day = int(self.selected_date[:2])
        reservation_month = strptime(self.selected_date[3:6],'%b').tm_mon
        reservation_year = int(self.selected_date[7:11])
        reservation_hour = int(self.selected_time_hours)
        reservation_minutes = int(self.selected_time_minutes)

        reservation_start_datetime = datetime(reservation_year,reservation_month,reservation_day, reservation_hour, reservation_minutes)
        reservation_end_datetime = datetime(reservation_year,reservation_month,reservation_day, reservation_hour+3, reservation_minutes)

        reservation_guests = int(self.selected_num_people)
        reservation_special_request = self.remark
        reservation_table = self.reservation_checker.check_reservation_possible(reservation_guests,
                                                                                reservation_start_datetime,
                                                                                reservation_end_datetime)
        print('table in concept:',reservation_table)

        self.reservation_concept = Reservation( name=reservation_name,
                                                num_guests=reservation_guests,
                                                start_date_time=reservation_start_datetime,
                                                end_date_time=reservation_end_datetime,
                                                special_request=reservation_special_request,
                                                table=reservation_table)

        self.reservation_table = reservation_table
        print(self.reservation_concept)

class ReservationChecker:
    def __init__(self):
        pass

    def get_reservations_for_day(self, start_datetime, end_datetime):
        # smarter would be to check for the start time instantly
        # filter range : endtime >=
        try:
            # any reservations that overlap should remain, as their table numbers have to be checked.
            # overlap means:
            # their start datetime falls between start_datetime and end_datetime OR
            # their end datetime falls between start_time and end_time
            # for OR use Q - objects
            reservations1 = Reservation.objects.filter(start_date_time__range=(start_datetime,end_datetime))
        except Reservation.DoesNotExist:
            reservations1 = None

        try:
            reservations2 = Reservation.objects.filter(end_date_time__range=(start_datetime,end_datetime))
        except Reservation.DoesNotExist:
            reservations2 = None

        print('reservations that start in this range:',reservations1)
        print('reservations that end in this range:',reservations2)

        reservations = reservations1 | reservations2

        return reservations

    def filter_time_span(self, reservations, start_time, end_time):
        reservations.filter(start_date_time)
        filtered_reservations = None
        return filtered_reservations

    def check_reservation_possible(self, num_guests, start_date_time, end_date_time):
        # retrieve all reservations for date, so that these tables can be excluded
        #print('checking',start_date_time,end_date_time)
        reservations = self.get_reservations_for_day(start_date_time, end_date_time)
        print('reservations for ',start_date_time,":",reservations)

        try:
            available_tables = set(Table.objects.all())
        except Table.DoesNotExist:
            available_tables = None

        print('all tables:')
        for t in available_tables:
            print(t)

        if reservations and available_tables:
            for r in reservations:
                if r.table:
                    if t in available_tables:
                        available_tables.remove(t)
        else:
            # no reservations yet
            pass

        print('still available:')
        for t in available_tables:
            print(t)

        #with available tables, calculate what's possible
        resulting_table = self.find_exact_table_match(num_guests, available_tables)
        if not resulting_table:
            resulting_table = self.find_larger_table_match(num_guests, available_tables)

        print('result:',resulting_table)

        return resulting_table

    def find_exact_table_match(self, persons, tables):
        # tables = list
        print('guests:',persons,"; received table options:",tables)
        for t in tables:
            if int(t.capacity) == int(persons):
                return t

    def find_larger_table_match(self, persons, tables):
        #print('finding a larger match')
        #print('received these options for iteration: ',tables)
        table_list = [t for t in tables]
        table_list.sort(key=self.sort_by_capacity)
        #print('sorted tables by capacity:',table_list)
        for t in tables:
            #print('iterating... cap:',t.capacity)
            if int(t.capacity) > int(persons):
                return t

    def sort_by_capacity(self,x):
        return x.capacity

class ConfirmationView(TemplateView):
    print('confirmation')


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

# Define your item pipelines here
#
# Don't forget to add your pipeline to the ITEM_PIPELINES setting
# See: https://docs.scrapy.org/en/latest/topics/item-pipeline.html

import re
from bs4 import BeautifulSoup
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.exc import IntegrityError
from scrapy.exceptions import DropItem

from models import Base, Product, BranchProduct

class ProductPipeline(object):
    def __init__(self):
        self.engine = create_engine('sqlite:///db.sqlite')
        Session = sessionmaker(bind=self.engine)
        self.session = Session()

    def process_item(self, item, spider):

        if item.get('image_urls') and len(list(item.get('image_urls'))) > 0:
            item['image_urls'] = ','.join([','.join([x['large']['url'], x['small']['url']]) for x in item.get('image_urls')])

        if item.get('barcodes') and len(list(item.get('barcodes'))) > 0:
            item['barcodes'] = self.arrayToString(item.get('barcodes'))

        if item.get('categories') and len(list(item.get('categories'))) > 0:
            item['categories'] = ','.join([ ','.join(list(map(lambda x: x['seo']['text'] , x['hierarchy']))) for x in item.get('categories')])

        item['description'] = self.removeHTML(item['description']).strip().capitalize()
        item['name'] = item['name'].capitalize()
        item['brand'] = item['brand'].capitalize()
        item['package'] = item['package'] if len(re.findall(r"\d", item['package'])) > 0 else 'n/a'

        return self.save(item)

    def save(self, item):
        try:
            product = Product(
                store='Walmart',
                barcodes=item.get('barcodes'),
                sku=str(item.get('sku')),
                brand=item.get('brand'),
                name=item.get('name'),
                description=item.get('description'),
                package=item.get('package'),
                categories=item.get('categories'),
                image_urls=item.get('image_urls'))

            branch = BranchProduct(
                branch=item.get('store'),
                product=product,
                stock=item['stock'],
                price=float(item.get('price')),
            )

            self.session.add(product)
            self.session.commit()
            return item
        except IntegrityError:
            self.session.rollback()
            return None

    def removeHTML(self, text):
        soup = BeautifulSoup(text)
        return soup.get_text()

    def arrayToString(self, item):
        return ','.join([x for x in item])


"""pylandstats setup script."""

import platform
from pathlib import Path

from setuptools import setup

# pythran imports must go AFTER setuptools imports
# See: https://github.com/pypa/setuptools/issues/309 and https://bit.ly/300HKtK
from transonic.dist import init_transonic_extensions, make_backend_files

here = Path(__file__).parent.absolute()

if platform.system() == "Windows":
    backend = "numba"
else:
    backend = "pythran"

paths = ["pylandstats/landscape.py"]
make_backend_files([here / path for path in paths], backend=backend)

if platform.system() == "Linux":
    compile_args = ("-O3", "-DUSE_XSIMD")
else:
    compile_args = ("-O3",)

extensions = init_transonic_extensions(
    "pylandstats", compile_args=compile_args, backend=backend
)

setup(
    ext_modules=extensions,
)


import argparse
from bs4 import BeautifulSoup
import logging
import os
import pandas as pd
from pathlib import Path
import random
import requests
import sys
import time
import tweepy

class Bot:
	def __init__(self):
		auth = tweepy.OAuthHandler(os.environ.get('TWITTER_API_KEY'),
			os.environ.get('TWITTER_API_SECRET_KEY'))

		auth.set_access_token(os.environ.get('TWITTER_API_TOKEN'),
			os.environ.get('TWITTER_API_SECRET_TOKEN'))

		self.api = tweepy.API(auth)

		self.bird_data = pd.read_csv("bird_data/bird_urls.csv")

	def run(self):
		self.verify_credentials()

		while True:
			self.reply_to_mentions()
			logging.info('waiting to check mentions')
			time.sleep(60)

	def reply_to_mentions(self):
		logging.info('reading most recent mention id from file')
		file = open('src/data.txt', 'r')
		since_id = int(file.read())
		file.close()

		logging.info('retrieving mentions')
		new_since_id = since_id

		for tweet in tweepy.Cursor(self.api.mentions_timeline, since_id=since_id).items():
			new_since_id = max(tweet.id, new_since_id)
			if tweet.in_reply_to_status_id is not None:
				continue

			else:
				logging.info(f'replying to mention by {tweet.user.name}')

				if not tweet.user.following:
					tweet.user.follow()

				URL = self.get_random()
				self.send_bird(URL, tweet.id)

		file = open('src/data.txt', 'w')
		file.write(str(new_since_id))
		file.close()

	def send_bird(self, URL, tweet_id):
		page = requests.get(URL)
		soup = BeautifulSoup(page.content, 'html.parser')

		image_sent = False

		try:
			result = soup.find('div',class_="bird-guide-image")
			img_url = result.find('img')['src']
			img = requests.get(img_url)
			file = open('bird_data/bird.jpg', 'wb')
			file.write(img.content)
			file.close()

			image = 'bird_data/bird.jpg'

			tweet_status = self.api.update_with_media(image, status=URL,
				in_reply_to_status_id=tweet_id, auto_populate_reply_metadata=True)

			image_sent = True
			logging.info(f'{Path(URL).stem} image sent')

			os.remove("bird_data/bird.jpg")

		except Exception as e:
			logging.error(e, exc_info=True)
			logging.info(URL + ', problem with bird image')
			image_sent = False

			if (os.path.exists("bird_data/bird.jpg")):
				os.remove("bird_data/bird.jpg")

		if not image_sent:
			return False

		try:
			result = soup.find('div', class_="hide-for-tiny hide-for-small hide-for-medium")
			text = result.text.strip('\t \n')

			if len(text) > 275:
				text = self.trim_text(text)

			self.api.update_status(status=text, in_reply_to_status_id=tweet_status.id)

			logging.info(f'{Path(URL).stem} text sent')

		except Exception as e:
			logging.error(e, exc_info=True)
			logging.info(URL + ', problem with bird text')

	def trim_text(self, text):
		logging.info(f'text too long: {len(text)}')
		text = text.split(".")
		text.pop()
		while len('.'.join(text) + '.') > 275:
			text.pop()
		text = '.'.join(text)
		text = text + '.'

		logging.info(f'text trimmed to {len(text)}')

		return text

	def get_random(self):
		logging.info('picking a random bird')

		number = random.randint(0, self.bird_data.shape[0])

		URL = self.bird_data['BIRD URLs'][number]
		logging.info(f'{Path(URL).stem} picked')
		return URL

	def verify_credentials(self):
		try:
			self.api.verify_credentials()
			logging.info("All clear")

		except Exception as e:
			logging.error(e, exc_info=True)
			logging.info("error during authentication")

def prep_log(debug,console):
	log_format = '%(asctime)s %(funcName)s()_%(lineno)s %(levelname)s: %(message)s'
	date_format = '%m/%d/%Y %I:%M:%S %p'

	logger = logging.getLogger()
	logger.setLevel(logging.DEBUG)
	formatter = logging.Formatter(log_format, datefmt=date_format)
	if console:
		console_handler = logging.StreamHandler(sys.stdout)
		console_handler.setLevel(logging.DEBUG if debug else logging.INFO)
		console_handler.setFormatter(formatter)
		logger.addHandler(console_handler)

	file_handler = logging.FileHandler('logs/log.log')
	file_handler.setLevel(logging.INFO)
	file_handler.setFormatter(formatter)
	logger.addHandler(file_handler)

def setup():
	parser = argparse.ArgumentParser()
	parser.add_argument("-d", "--debug", help="DEBUG MODE", action="store_true")
	parser.add_argument("-c", "--console", help="LOG TO CONSOLE", action="store_true")
	argv = parser.parse_args()

	prep_log(argv.debug, argv.console)

if __name__ == "__main__":
	setup()
	bird = Bot()
	bird.run()


import math
import copy
from torch import nn
import torch.nn.functional as F
import torch
from inspect import isfunction
from functools import partial

from torch.utils import data
from pathlib import Path
from torch.optim import Adam
from torchvision import transforms, utils

from einops import rearrange

from PIL import Image
from torch import linalg as LA

try:
    from apex import amp

    APEX_AVAILABLE = True
except:
    APEX_AVAILABLE = False

def exists(x):
    return x is not None

def default(val, d):
    if exists(val):
        return val
    return d() if isfunction(d) else d

def cycle(dl):
    while True:
        for data in dl:
            yield data

import numpy as np

def spiral_cw(A):
    A = np.array(A)
    out = []
    while(A.size):
        out.append(A[0])        # take first row
        A = A[1:].T[::-1]       # cut off first row and rotate counterclockwise
    return np.concatenate(out)

def spiral_ccw(A):
    A = np.array(A)
    out = []
    while(A.size):
        out.append(A[0][::-1])    # first row reversed
        A = A[1:][::-1].T         # cut off first row and rotate clockwise
    return np.concatenate(out)

def base_spiral(nrow, ncol):
    return spiral_ccw(np.arange(nrow*ncol).reshape(nrow,ncol))[::-1]

def to_spiral(A):
    A = np.array(A)
    B = np.empty_like(A)
    B.flat[base_spiral(*A.shape)] = A.flat
    return B

def from_spiral(A):
    A = np.array(A)
    return A.flat[base_spiral(*A.shape)].reshape(A.shape)

def to_spiral(A):
    A = np.array(A)
    B = np.empty_like(A)
    B.flat[base_spiral(*A.shape)] = A.flat
    return B

def from_spiral(A):
    A = np.array(A)
    return A.flat[base_spiral(*A.shape)].reshape(A.shape)

def loss_backwards(fp16, loss, optimizer, **kwargs):
    if fp16:
        with amp.scale_loss(loss, optimizer) as scaled_loss:
            scaled_loss.backward(**kwargs)
    else:
        loss.backward(**kwargs)

class EMA():
    def __init__(self, beta):
        super().__init__()
        self.beta = beta

    def update_model_average(self, ma_model, current_model):
        for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
            old_weight, up_weight = ma_params.data, current_params.data
            ma_params.data = self.update_average(old_weight, up_weight)

    def update_average(self, old, new):
        if old is None:
            return new
        return old * self.beta + (1 - self.beta) * new

class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn

    def forward(self, x, *args, **kwargs):
        return self.fn(x, *args, **kwargs) + x

class SinusoidalPosEmb(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, x):
        device = x.device
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = x[:, None] * emb[None, :]
        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
        return emb

def Upsample(dim):
    return nn.ConvTranspose2d(dim, dim, 4, 2, 1)

def Downsample(dim):
    return nn.Conv2d(dim, dim, 4, 2, 1)

class LayerNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
        self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))

    def forward(self, x):
        var = torch.var(x, dim=1, unbiased=False, keepdim=True)
        mean = torch.mean(x, dim=1, keepdim=True)
        return (x - mean) / (var + self.eps).sqrt() * self.g + self.b

class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.fn = fn
        self.norm = LayerNorm(dim)

    def forward(self, x):
        x = self.norm(x)
        return self.fn(x)

class ConvNextBlock(nn.Module):
    """ https://arxiv.org/abs/2201.03545 """

    def __init__(self, dim, dim_out, *, time_emb_dim=None, mult=2, norm=True):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.GELU(),
            nn.Linear(time_emb_dim, dim)
        ) if exists(time_emb_dim) else None

        self.ds_conv = nn.Conv2d(dim, dim, 7, padding=3, groups=dim)

        self.net = nn.Sequential(
            LayerNorm(dim) if norm else nn.Identity(),
            nn.Conv2d(dim, dim_out * mult, 3, padding=1),
            nn.GELU(),
            nn.Conv2d(dim_out * mult, dim_out, 3, padding=1)
        )

        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()

    def forward(self, x, time_emb=None):
        h = self.ds_conv(x)

        if exists(self.mlp):
            assert exists(time_emb), 'time emb must be passed in'
            condition = self.mlp(time_emb)
            h = h + rearrange(condition, 'b c -> b c 1 1')

        h = self.net(h)
        return h + self.res_conv(x)

class LinearAttention(nn.Module):
    def __init__(self, dim, heads=4, dim_head=32):
        super().__init__()
        self.scale = dim_head ** -0.5
        self.heads = heads
        hidden_dim = dim_head * heads
        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
        self.to_out = nn.Conv2d(hidden_dim, dim, 1)

    def forward(self, x):
        b, c, h, w = x.shape
        qkv = self.to_qkv(x).chunk(3, dim=1)
        q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> b h c (x y)', h=self.heads), qkv)
        q = q * self.scale

        k = k.softmax(dim=-1)
        context = torch.einsum('b h d n, b h e n -> b h d e', k, v)

        out = torch.einsum('b h d e, b h d n -> b h e n', context, q)
        out = rearrange(out, 'b h c (x y) -> b (h c) x y', h=self.heads, x=h, y=w)
        return self.to_out(out)

class Unet(nn.Module):
    def __init__(
            self,
            dim,
            out_dim=None,
            dim_mults=(1, 2, 4, 8),
            channels=3,
            with_time_emb=True,
            residual=False
    ):
        super().__init__()
        self.channels = channels
        self.residual = residual
        print("Is Time embed used ? ", with_time_emb)

        dims = [channels, *map(lambda m: dim * m, dim_mults)]
        in_out = list(zip(dims[:-1], dims[1:]))

        if with_time_emb:
            time_dim = dim
            self.time_mlp = nn.Sequential(
                SinusoidalPosEmb(dim),
                nn.Linear(dim, dim * 4),
                nn.GELU(),
                nn.Linear(dim * 4, dim)
            )
        else:
            time_dim = None
            self.time_mlp = None

        self.downs = nn.ModuleList([])
        self.ups = nn.ModuleList([])
        num_resolutions = len(in_out)

        for ind, (dim_in, dim_out) in enumerate(in_out):
            is_last = ind >= (num_resolutions - 1)

            self.downs.append(nn.ModuleList([
                ConvNextBlock(dim_in, dim_out, time_emb_dim=time_dim, norm=ind != 0),
                ConvNextBlock(dim_out, dim_out, time_emb_dim=time_dim),
                Residual(PreNorm(dim_out, LinearAttention(dim_out))),
                Downsample(dim_out) if not is_last else nn.Identity()
            ]))

        mid_dim = dims[-1]
        self.mid_block1 = ConvNextBlock(mid_dim, mid_dim, time_emb_dim=time_dim)
        self.mid_attn = Residual(PreNorm(mid_dim, LinearAttention(mid_dim)))
        self.mid_block2 = ConvNextBlock(mid_dim, mid_dim, time_emb_dim=time_dim)

        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
            is_last = ind >= (num_resolutions - 1)

            self.ups.append(nn.ModuleList([
                ConvNextBlock(dim_out * 2, dim_in, time_emb_dim=time_dim),
                ConvNextBlock(dim_in, dim_in, time_emb_dim=time_dim),
                Residual(PreNorm(dim_in, LinearAttention(dim_in))),
                Upsample(dim_in) if not is_last else nn.Identity()
            ]))

        out_dim = default(out_dim, channels)
        self.final_conv = nn.Sequential(
            ConvNextBlock(dim, dim),
            nn.Conv2d(dim, out_dim, 1)
        )

    def forward(self, x, time):
        orig_x = x
        t = self.time_mlp(time) if exists(self.time_mlp) else None

        h = []

        for convnext, convnext2, attn, downsample in self.downs:
            x = convnext(x, t)
            x = convnext2(x, t)
            x = attn(x)
            h.append(x)
            x = downsample(x)

        x = self.mid_block1(x, t)
        x = self.mid_attn(x)
        x = self.mid_block2(x, t)

        for convnext, convnext2, attn, upsample in self.ups:
            x = torch.cat((x, h.pop()), dim=1)
            x = convnext(x, t)
            x = convnext2(x, t)
            x = attn(x)
            x = upsample(x)
        if self.residual:
            return self.final_conv(x) + orig_x

        return self.final_conv(x)

class GaussianDiffusion(nn.Module):
    def __init__(
            self,
            defade_fn,
            *,
            image_size,
            device_of_kernel,
            channels=3,
            timesteps=1000,
            loss_type='l1',
            start_fade_factor=0.1,
            fade_routine='Incremental',
            train_routine='Final',
            sampling_routine='default'
    ):
        super().__init__()
        self.channels = channels
        self.image_size = image_size
        self.defade_fn = defade_fn
        self.device_of_kernel = device_of_kernel

        self.num_timesteps = int(timesteps)
        self.loss_type = loss_type
        self.start_fade_factor = start_fade_factor
        self.fade_routine = fade_routine

        self.fade_factors = self.get_fade_factors()
        self.train_routine = train_routine
        self.sampling_routine = sampling_routine

    def get_fade_factors(self):
        fade_factors = []
        for i in range(self.num_timesteps):
            if self.fade_routine == 'Incremental':
                fade_factors.append(1 - self.start_fade_factor * (i + 1))
            elif self.fade_routine == 'Constant':
                fade_factors.append(1 - self.start_fade_factor)
            elif self.fade_routine == 'Spiral':
                A = np.arange(32 * 32).reshape(32, 32)
                spiral = to_spiral(A)
                k = spiral > i
                k = torch.tensor(k).float()
                fade_factors.append(k.cuda())
            elif self.fade_routine == 'Spiral_2':
                A = np.arange(32 * 32).reshape(32, 32)
                spiral = to_spiral(A)
                k = spiral > i
                k = torch.tensor(k).float()
                fade_factors.append(k.cuda())

        return fade_factors

    @torch.no_grad()
    def sample(self, batch_size=16, img=None, t=None):

        if t is None:
            t = self.num_timesteps
        for i in range(t):
            with torch.no_grad():
                img = self.fade_factors[i] * img

        if self.fade_routine == 'Spiral_2':
            new_mean = torch.rand((img.shape[0], 3))
            new_mean = new_mean.unsqueeze(2).repeat(1, 1, img.shape[2])
            new_mean = new_mean.unsqueeze(3).repeat(1, 1, 1, img.shape[3]).cuda()

        for i in range(t):
            with torch.no_grad():
                if self.fade_routine == 'Spiral_2':
                    img = self.fade_factors[i] * img + new_mean * (torch.ones_like(self.fade_factors[i]) - self.fade_factors[i])
                else:
                    img = self.fade_factors[i] * img

        xt = img
        direct_recons = None
        while t:
            step = torch.full((batch_size,), t - 1, dtype=torch.long).cuda()
            x = self.defade_fn(img, step)

            if "Final" in self.train_routine:
                if direct_recons is None:
                    direct_recons = x

                if self.sampling_routine == 'default':
                    for i in range(t - 1):
                        with torch.no_grad():
                            x = self.fade_factors[i] * x

                elif self.sampling_routine == 'x0_step_down':
                    x_times = x
                    x_times_sub_1 = x

                    for i in range(t):
                        with torch.no_grad():
                            x_times_sub_1 = x_times
                            x_times = self.fade_factors[i] * x_times

                    x = img - x_times + x_times_sub_1

                elif self.sampling_routine == 'x0_step_down_spiral_2_fix':
                    x_times = x
                    x_times_sub_1 = x

                    for i in range(t):
                        with torch.no_grad():
                            x_times_sub_1 = x_times
                            x_times = self.fade_factors[i] * x_times + new_mean * (
                                        torch.ones_like(self.fade_factors[i]) - self.fade_factors[i])

                    x = img - x_times + x_times_sub_1

                elif self.sampling_routine == 'x0_step_down_spiral_2_rand':
                    x_times = x
                    x_times_sub_1 = x

                    for i in range(t):
                        new_mean = torch.rand((img.shape[0], 3))
                        new_mean = new_mean.unsqueeze(2).repeat(1, 1, img.shape[2])
                        new_mean = new_mean.unsqueeze(3).repeat(1, 1, 1, img.shape[3]).cuda()

                        with torch.no_grad():
                            x_times_sub_1 = x_times
                            x_times = self.fade_factors[i] * x_times + new_mean * (
                                        torch.ones_like(self.fade_factors[i]) - self.fade_factors[i])

                    x = img - x_times + x_times_sub_1

            elif self.train_routine == 'Step':
                if direct_recons is None:
                    direct_recons = x

            elif self.train_routine == 'Gradient_norm':
                if direct_recons is None:
                    direct_recons = img - x
                x = img - x

            img = x
            t = t - 1

        return xt, direct_recons, img

    @torch.no_grad()
    def all_sample(self, batch_size=16, img=None, t=None, times=None):

        if t is None:
            t = self.num_timesteps
        if times is None:
            times = t

        if self.fade_routine == 'Spiral_2':
            new_mean = torch.rand((img.shape[0], 3))
            new_mean = new_mean.unsqueeze(2).repeat(1, 1, img.shape[2])
            new_mean = new_mean.unsqueeze(3).repeat(1, 1, 1, img.shape[3]).cuda()

        for i in range(t):
            with torch.no_grad():
                if self.fade_routine == 'Spiral_2':
                    img = self.fade_factors[i] * img + new_mean * (
                                torch.ones_like(self.fade_factors[i]) - self.fade_factors[i])
                else:
                    img = self.fade_factors[i] * img

        x0_list = []
        xt_list = []

        while times:
            step = torch.full((batch_size,), times - 1, dtype=torch.long).cuda()
            x = self.defade_fn(img, step)
            x0_list.append(x)

            if "Final" in self.train_routine:
                if self.sampling_routine == 'default':
                    print("Normal")

                    x_times_sub_1 = x
                    for i in range(times - 1):
                        with torch.no_grad():
                            x_times_sub_1 = self.fade_factors[i] * x_times_sub_1

                    x = x_times_sub_1

                elif self.sampling_routine == 'x0_step_down':
                    print("x0_step_down")

                    x_times = x
                    for i in range(times):
                        with torch.no_grad():
                            x_times = self.fade_factors[i] * x_times

                    x_times_sub_1 = x
                    for i in range(times - 1):
                        with torch.no_grad():
                            x_times_sub_1 = self.fade_factors[i] * x_times_sub_1

                    x = img - x_times + x_times_sub_1

                elif self.sampling_routine == 'x0_step_down_spiral_2_fix':
                    x_times = x
                    x_times_sub_1 = x

                    for i in range(times):
                        with torch.no_grad():
                            x_times_sub_1 = x_times
                            x_times = self.fade_factors[i] * x_times + new_mean * (
                                        torch.ones_like(self.fade_factors[i]) - self.fade_factors[i])

                    x = img - x_times + x_times_sub_1

                elif self.sampling_routine == 'x0_step_down_spiral_2_rand':
                    x_times = x
                    x_times_sub_1 = x

                    for i in range(times):
                        new_mean = torch.rand((img.shape[0], 3))
                        new_mean = new_mean.unsqueeze(2).repeat(1, 1, img.shape[2])
                        new_mean = new_mean.unsqueeze(3).repeat(1, 1, 1, img.shape[3]).cuda()

                        with torch.no_grad():
                            x_times_sub_1 = x_times
                            x_times = self.fade_factors[i] * x_times + new_mean * (
                                        torch.ones_like(self.fade_factors[i]) - self.fade_factors[i])

                    x = img + x_times_sub_1 - img #- x_times

                elif self.sampling_routine == 'no_time_embed':
                    x = x
                    for i in range(100):
                        with torch.no_grad():
                            x = self.fade_factors[i] * x

            elif self.train_routine == 'Gradient_norm':
                x = img - 0.1 * x
                for i in range(10):
                    with torch.no_grad():
                        x = self.fade_factors[i] * x

            img = x
            xt_list.append(img)
            times = times - 1

        return x0_list, xt_list

    def q_sample(self, x_start, t):

        if self.fade_routine == 'Spiral':
            choose_fade = []
            for img_index in range(t.shape[0]):
                choose_fade.append(x_start[img_index,:] * self.fade_factors[t[img_index]] )

            choose_fade = torch.stack(choose_fade)
            return choose_fade

        elif self.fade_routine == 'Spiral_2':

            choose_fade = []
            for img_index in range(t.shape[0]):
                new_mean = torch.rand((1, 3))
                new_mean = new_mean.unsqueeze(2).repeat(1, 1, x_start.shape[2])
                new_mean = new_mean.unsqueeze(3).repeat(1, 1, 1, x_start.shape[3]).cuda()

                cf = x_start[img_index,:] * self.fade_factors[t[img_index]] + new_mean * (torch.ones_like(self.fade_factors[t[img_index]]) - self.fade_factors[t[img_index]])
                choose_fade.append(cf[0,:])

            choose_fade = torch.stack(choose_fade)
            return choose_fade

        else:
            max_iters = torch.max(t)
            all_fades = []
            x = x_start
            for i in range(max_iters + 1):
                with torch.no_grad():
                    x = self.fade_factors[i] * x
                    all_fades.append(x)

            all_fades = torch.stack(all_fades)

            choose_fade = []
            for step in range(t.shape[0]):
                if step != -1:
                    choose_fade.append(all_fades[t[step], step])
                else:
                    choose_fade.append(x_start[step])

            choose_fade = torch.stack(choose_fade)
            return choose_fade

    def p_losses(self, x_start, t):
        if self.train_routine == 'Final':
            x_fade = self.q_sample(x_start=x_start, t=t)
            x_recon = self.defade_fn(x_fade, t)

            if self.loss_type == 'l1':
                loss = (x_start - x_recon).abs().mean()
            elif self.loss_type == 'l2':
                loss = F.mse_loss(x_start, x_recon)
            else:
                raise NotImplementedError()

        elif self.train_routine == 'Gradient_norm':
            x_fade = self.q_sample(x_start=x_start, t=t)
            grad_pred = self.defade_fn(x_fade, t)
            gradient = (x_fade - x_start)
            norm = LA.norm(gradient, dim=(1, 2, 3), keepdim=True)
            gradient_norm = gradient / (norm + 1e-5)

            if self.loss_type == 'l1':
                loss = (gradient_norm - grad_pred).abs().mean()
            elif self.loss_type == 'l2':
                loss = F.mse_loss(gradient_norm, grad_pred)
            else:
                raise NotImplementedError()

        elif self.train_routine == 'Step':
            x_fade = self.q_sample(x_start=x_start, t=t)
            x_fade_sub = self.q_sample(x_start=x_start, t=t - 1)
            x_blur_sub_pred = self.defade_fn(x_fade, t)

            if self.loss_type == 'l1':
                loss = (x_fade_sub - x_blur_sub_pred).abs().mean()
            elif self.loss_type == 'l2':
                loss = F.mse_loss(x_fade_sub, x_blur_sub_pred)
            else:
                raise NotImplementedError()

        return loss

    def forward(self, x, *args, **kwargs):
        b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size
        assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
        t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
        return self.p_losses(x, t, *args, **kwargs)

class Dataset(data.Dataset):
    def __init__(self, folder, image_size, exts=['jpg', 'jpeg', 'png']):
        super().__init__()
        self.folder = folder
        self.image_size = image_size
        self.paths = [p for ext in exts for p in Path(f'{folder}').glob(f'**/*.{ext}')]

        self.transform = transforms.Compose([
            transforms.Resize(image_size),
            transforms.CenterCrop(image_size),
            transforms.ToTensor(),
            transforms.Lambda(lambda t: (t * 2) - 1)
        ])

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

    def __getitem__(self, index):
        path = self.paths[index]
        img = Image.open(path)
        return self.transform(img)

class Dataset_Cifar10(data.Dataset):
    def __init__(self, folder, image_size, exts=['jpg', 'jpeg', 'png']):
        super().__init__()
        self.folder = folder
        self.image_size = image_size
        self.paths = [p for ext in exts for p in Path(f'{folder}').glob(f'**/*.{ext}')]

        self.transform = transforms.Compose([
            transforms.RandomCrop(image_size, padding=4),
            transforms.Resize(image_size),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Lambda(lambda t: (t * 2) - 1)
        ])

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

    def __getitem__(self, index):
        path = self.paths[index]
        img = Image.open(path)
        return self.transform(img)

class Trainer(object):
    def __init__(
            self,
            diffusion_model,
            folder,
            *,
            ema_decay=0.995,
            image_size=128,
            train_batch_size=32,
            train_lr=2e-5,
            train_num_steps=100000,
            gradient_accumulate_every=2,
            fp16=False,
            step_start_ema=2000,
            update_ema_every=10,
            save_and_sample_every=1000,
            results_folder='./results',
            load_path=None,
            dataset=None
    ):
        super().__init__()
        self.model = diffusion_model
        self.ema = EMA(ema_decay)
        self.ema_model = copy.deepcopy(self.model)
        self.update_ema_every = update_ema_every

        self.step_start_ema = step_start_ema
        self.save_and_sample_every = save_and_sample_every

        self.batch_size = train_batch_size
        self.image_size = diffusion_model.image_size
        self.gradient_accumulate_every = gradient_accumulate_every
        self.train_num_steps = train_num_steps

        if dataset == 'cifar10':
            self.ds = Dataset_Cifar10(folder, image_size)
        else:
            self.ds = Dataset(folder, image_size)
        self.dl = cycle(data.DataLoader(self.ds, batch_size=train_batch_size, shuffle=True, pin_memory=True))
        self.opt = Adam(diffusion_model.parameters(), lr=train_lr)

        self.step = 0

        assert not fp16 or fp16 and APEX_AVAILABLE, 'Apex must be installed in order for mixed precision training to be turned on'

        self.fp16 = fp16
        if fp16:
            (self.model, self.ema_model), self.opt = amp.initialize([self.model, self.ema_model], self.opt,
                                                                    opt_level='O1')

        self.results_folder = Path(results_folder)
        self.results_folder.mkdir(exist_ok=True)

        self.reset_parameters()

        if load_path is not None:
            self.load(load_path)

    def reset_parameters(self):
        self.ema_model.load_state_dict(self.model.state_dict())

    def step_ema(self):
        if self.step < self.step_start_ema:
            self.reset_parameters()
            return
        self.ema.update_model_average(self.ema_model, self.model)

    def save(self):
        data = {
            'step': self.step,
            'model': self.model.state_dict(),
            'ema': self.ema_model.state_dict()
        }
        torch.save(data, str(self.results_folder / f'model.pt'))

    def load(self, load_path):
        print("Loading : ", load_path)
        data = torch.load(load_path)

        self.step = data['step']
        self.model.load_state_dict(data['model'])
        self.ema_model.load_state_dict(data['ema'])

    def add_title(self, path, title):
        import cv2
        import numpy as np

        img1 = cv2.imread(path)
        black = [0, 0, 0]
        constant = cv2.copyMakeBorder(img1, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=black)
        height = 20
        violet = np.zeros((height, constant.shape[1], 3), np.uint8)
        violet[:] = (255, 0, 180)

        vcat = cv2.vconcat((violet, constant))

        font = cv2.FONT_HERSHEY_SIMPLEX

        cv2.putText(vcat, str(title), (violet.shape[1] // 2, height - 2), font, 0.5, (0, 0, 0), 1, 0)
        cv2.imwrite(path, vcat)

    def train(self):
        backwards = partial(loss_backwards, self.fp16)

        acc_loss = 0
        while self.step < self.train_num_steps:
            u_loss = 0
            for i in range(self.gradient_accumulate_every):
                data = next(self.dl).cuda()
                loss = self.model(data)
                print(f'{self.step}: {loss.item()}')
                u_loss += loss.item()
                backwards(loss / self.gradient_accumulate_every, self.opt)

            acc_loss = acc_loss + (u_loss / self.gradient_accumulate_every)

            self.opt.step()
            self.opt.zero_grad()

            if self.step % self.update_ema_every == 0:
                self.step_ema()

            if self.step != 0 and self.step % self.save_and_sample_every == 0:
                milestone = self.step // self.save_and_sample_every
                batches = self.batch_size
                og_img = next(self.dl).cuda()
                xt, direct_recons, all_images = self.ema_model.sample(batch_size=batches, faded_recon_sample=og_img)

                og_img = (og_img + 1) * 0.5
                utils.save_image(og_img, str(self.results_folder / f'sample-og-{milestone}.png'), nrow=6)

                all_images = (all_images + 1) * 0.5
                utils.save_image(all_images, str(self.results_folder / f'sample-recon-{milestone}.png'), nrow=6)

                direct_recons = (direct_recons + 1) * 0.5
                utils.save_image(direct_recons, str(self.results_folder / f'sample-direct_recons-{milestone}.png'),
                                 nrow=6)

                xt = (xt + 1) * 0.5
                utils.save_image(xt, str(self.results_folder / f'sample-xt-{milestone}.png'),
                                 nrow=6)

                acc_loss = acc_loss / (self.save_and_sample_every + 1)
                print(f'Mean of last {self.step}: {acc_loss}')
                acc_loss = 0

                self.save()

            self.step += 1

        print('training completed')

    def test_from_data(self, extra_path, s_times=None):
        batches = self.batch_size
        og_img = next(self.dl).cuda()
        x0_list, xt_list = self.ema_model.all_sample(batch_size=batches, faded_recon_sample=og_img, times=s_times)

        og_img = (og_img + 1) * 0.5
        utils.save_image(og_img, str(self.results_folder / f'og-{extra_path}.png'), nrow=6)

        import imageio
        frames_t = []
        frames_0 = []

        for i in range(len(x0_list)):
            print(i)

            x_0 = x0_list[i]
            x_0 = (x_0 + 1) * 0.5
            utils.save_image(x_0, str(self.results_folder / f'sample-{i}-{extra_path}-x0.png'), nrow=6)
            self.add_title(str(self.results_folder / f'sample-{i}-{extra_path}-x0.png'), str(i))
            frames_0.append(imageio.imread(str(self.results_folder / f'sample-{i}-{extra_path}-x0.png')))

            x_t = xt_list[i]
            all_images = (x_t + 1) * 0.5
            utils.save_image(all_images, str(self.results_folder / f'sample-{i}-{extra_path}-xt.png'), nrow=6)
            self.add_title(str(self.results_folder / f'sample-{i}-{extra_path}-xt.png'), str(i))
            frames_t.append(imageio.imread(str(self.results_folder / f'sample-{i}-{extra_path}-xt.png')))

        imageio.mimsave(str(self.results_folder / f'Gif-{extra_path}-x0.gif'), frames_0)
        imageio.mimsave(str(self.results_folder / f'Gif-{extra_path}-xt.gif'), frames_t)

    def test_with_mixup(self, extra_path):
        batches = self.batch_size
        og_img_1 = next(self.dl).cuda()
        og_img_2 = next(self.dl).cuda()
        og_img = (og_img_1 + og_img_2) / 2

        x0_list, xt_list = self.ema_model.all_sample(batch_size=batches, faded_recon_sample=og_img)

        og_img_1 = (og_img_1 + 1) * 0.5
        utils.save_image(og_img_1, str(self.results_folder / f'og1-{extra_path}.png'), nrow=6)

        og_img_2 = (og_img_2 + 1) * 0.5
        utils.save_image(og_img_2, str(self.results_folder / f'og2-{extra_path}.png'), nrow=6)

        og_img = (og_img + 1) * 0.5
        utils.save_image(og_img, str(self.results_folder / f'og-{extra_path}.png'), nrow=6)

        frames_t = []
        frames_0 = []

        for i in range(len(x0_list)):
            print(i)
            x_0 = x0_list[i]
            x_0 = (x_0 + 1) * 0.5
            utils.save_image(x_0, str(self.results_folder / f'sample-{i}-{extra_path}-x0.png'), nrow=6)
            self.add_title(str(self.results_folder / f'sample-{i}-{extra_path}-x0.png'), str(i))
            frames_0.append(Image.open(str(self.results_folder / f'sample-{i}-{extra_path}-x0.png')))

            x_t = xt_list[i]
            all_images = (x_t + 1) * 0.5
            utils.save_image(all_images, str(self.results_folder / f'sample-{i}-{extra_path}-xt.png'), nrow=6)
            self.add_title(str(self.results_folder / f'sample-{i}-{extra_path}-xt.png'), str(i))
            frames_t.append(Image.open(str(self.results_folder / f'sample-{i}-{extra_path}-xt.png')))

        frame_one = frames_0[0]
        frame_one.save(str(self.results_folder / f'Gif-{extra_path}-x0.gif'), format="GIF", append_images=frames_0,
                       save_all=True, duration=100, loop=0)

        frame_one = frames_t[0]
        frame_one.save(str(self.results_folder / f'Gif-{extra_path}-xt.gif'), format="GIF", append_images=frames_t,
                       save_all=True, duration=100, loop=0)

    def test_from_random(self, extra_path):
        batches = self.batch_size
        og_img = next(self.dl).cuda()
        og_img = og_img * 0.9  # torch.randn_like(og_img) + 0.1
        x0_list, xt_list = self.ema_model.all_sample(batch_size=batches, faded_recon_sample=og_img)

        og_img = (og_img + 1) * 0.5
        utils.save_image(og_img, str(self.results_folder / f'og-{extra_path}.png'), nrow=6)

        frames_t_names = []
        frames_0_names = []

        for i in range(len(x0_list)):
            print(i)

            x_0 = x0_list[i]
            x_0 = (x_0 + 1) * 0.5
            utils.save_image(x_0, str(self.results_folder / f'sample-{i}-{extra_path}-x0.png'), nrow=6)
            self.add_title(str(self.results_folder / f'sample-{i}-{extra_path}-x0.png'), str(i))
            frames_0_names.append(str(self.results_folder / f'sample-{i}-{extra_path}-x0.png'))

            x_t = xt_list[i]
            all_images = (x_t + 1) * 0.5
            utils.save_image(all_images, str(self.results_folder / f'sample-{i}-{extra_path}-xt.png'), nrow=6)
            self.add_title(str(self.results_folder / f'sample-{i}-{extra_path}-xt.png'), str(i))
            frames_t_names.append(str(self.results_folder / f'sample-{i}-{extra_path}-xt.png'))

        import imageio
        frames_0 = []
        frames_t = []
        for i in range(len(x0_list)):
            print(i)
            frames_0.append(imageio.imread(frames_0_names[i]))
            frames_t.append(imageio.imread(frames_t_names[i]))

        imageio.mimsave(str(self.results_folder / f'Gif-{extra_path}-x0.gif'), frames_0)
        imageio.mimsave(str(self.results_folder / f'Gif-{extra_path}-xt.gif'), frames_t)