from django import forms
from django.forms import ModelForm
from django.utils.translation import ugettext_lazy as _
from django.contrib.auth.models import User
from django.contrib.auth.forms import UserCreationForm
from restaurant.models import Restaurant
from localflavor.us.forms import USStateSelect, USZipCodeField, USStateField

class ContactForm(forms.Form):
    from_email = forms.EmailField(required=True, label='email')
    name = forms.CharField(required=True, label='name')
    message = forms.CharField(widget=forms.Textarea(), required=True, label='message')

    def __init__(self, *args, **kwargs):
        super(ContactForm, self).__init__(*args, **kwargs)
        self.fields['from_email'].widget.attrs.update({'placeholder': 'E-Mail'})
        self.fields['name'].widget.attrs.update({'placeholder': 'Name'})
        self.fields['message'].widget.attrs.update({'placeholder': 'Enter your message'})

class RestaurantForm(forms.ModelForm):
    name = forms.CharField(required=True)
    address1 = forms.CharField(required=True)
    city = forms.CharField(required=True)
    state = forms.CharField(widget=USStateSelect, required=True)

    class Meta:
        model = Restaurant
        fields = ('name', 'address1', 'address2', 'city', 'state', 'zipcode')

    def __init__(self, *args, **kwargs):
        super(RestaurantForm, self).__init__(*args, **kwargs)
        self.fields['name'].widget.attrs.update({'placeholder': 'Restaurant Name'})
        self.fields['address1'].widget.attrs.update({'placeholder': 'Address Line 1'})
        self.fields['address2'].widget.attrs.update({'placeholder': 'Address Line 2'})
        self.fields['city'].widget.attrs.update({'placeholder': 'City'})
        self.fields['zipcode'].widget.attrs.update({'placeholder': 'Zipcode'})

        self.fields['name'].error_messages = {'required': 'Please enter a restaurant name'}
        self.fields['address1'].error_messages = {'required': 'Please enter a address'}
        self.fields['address2'].error_messages = {'required': 'Please enter a address'}
        self.fields['city'].error_messages = {'required': 'Please enter a city'}
        self.fields['zipcode'].error_messages = {'required': 'Please enter a valid zip code'}

        for fieldname in ['name', 'address1', 'address2', 'city', 'state', 'zipcode']:
            self.fields[fieldname].label = ''


"""Традиционный подход к созданию таблиц, классов и отображений в SQLALchemy"""

# Иллюстрирует классический пример использования SQLAlchemy,
# в которой очень ценится разделение задач.

from sqlalchemy import create_engine, Table, Column, Integer, String, MetaData
from sqlalchemy.orm import mapper, sessionmaker

# Создаем объект подключения
engine = create_engine('sqlite:///test.sqlite3')

# Создаем экземпляр каталога MetaData
metadata = MetaData()

# Описываем таблицу
vendors_table = Table('vendors', metadata,
    Column('id', Integer, primary_key=True),
    Column('name', String),
    Column('phone', String),
    Column('address', String)
)

# Создаем структуру таблиц
metadata.create_all(engine)

# Создаем класс-шаблон записи БД
class Vendor:

    def __init__(self, name, phone, address):
        self.name = name
        self.phone = phone
        self.address = address

    def __repr__(self):
        return f"<Vendor({self.name}, {self.phone}, {self.address})>"

# Связываем шаблон записи с таблицей
mapper(Vendor, vendors_table)

# Создаем экземпляр записи с реальными параметрами
vendor = Vendor("ООО ‘Компани’ ", "8(495)77-77-77", "г. Москва, ул. Бажова, д. 9")

print(vendor)
# None. Как так? мы же создали объект записи и передали значения???
print(vendor.id)

# Оказывается мы еще ничего не сохранили. это нужно делать через сессии

#-------------------Cоздание сессии--------------------#
# С помощью конструктора sessionmaker создаем класс-сессия
SESSION = sessionmaker(bind=engine)

# создаем объект сессии
SESS_OBJ = SESSION()

# вот теперь все хорошо
SESS_OBJ.add(vendor)
SESS_OBJ.commit()
print(vendor.id)


#! /usr/bin/python
'''
Created on Feb 5, 2012

@author: charles
'''
import optparse
import os
from SelectionStrategyEvaluator import ExperimentResult
from functools import partial
from Datasets.non_textual import named_data_sets as non_textual_data_sets
from Datasets.textual import named_data_sets as textual_data_sets
from itertools import chain
from glob import glob
from utils import my_import

if __name__ == '__main__':
    parser = optparse.OptionParser("usage: %prog [options] data_info_name input_results_file output_file")
    parser.add_option('--cv', help='Cross Validation to include', dest='cv',
                      default=0, action='store', type='int')
    parser.add_option('--generateall', help='boolean option', dest='generate_all',
                      default=False, action='store_true')
    parser.add_option('--nocolour', help='boolean option forces greyscale graphs', dest='colour',
                      default=True, action='store_false')
    parser.add_option('--experiment', help='Use the experiment to figure out what files to get', dest='experiment',
                      default=None, action='store')
    (options, args) = parser.parse_args()

    data_info_name, input_results_file, output_file = args
    cv_no = options.cv
    input_results_file = os.path.expanduser(input_results_file)
    output_file = os.path.expanduser(output_file)

    data_sets_dict = dict(chain(non_textual_data_sets, textual_data_sets))

    data_info = data_sets_dict[data_info_name]()

    if options.experiment:
        experiment_obj = my_import(options.experiment).experiment
        var_names = [name for (name, variation)
                     in experiment_obj.named_experiment_variations]
        variation_files = [os.path.join(input_results_file, var_name+".tar.gz")
                           for var_name in var_names]
    elif os.path.isdir(input_results_file):
        variation_files = glob(os.path.join(input_results_file, ".tar.gz"))
    else:
        variation_files = [input_results_file]

    get_var_name = lambda f: os.path.splitext(os.path.splitext(os.path.basename(f))[0])[0]
    vn_to_fsg_list = [(get_var_name(f), partial(open, f, 'r'))
                      for f in variation_files]

    exp_result = ExperimentResult()
    exp_result.load_from_csvs(vn_to_fsg_list)

    def my_stream_getter(var_name, cv):
        # There's only going to be one variation, so only concerned with cv
        if cv == cv_no:
            if output_file.endswith('.pdf'):
                f = output_file
            else:
                if not os.path.exists(output_file):
                    os.makedirs(output_file)
                f = os.path.join(output_file, var_name + '.pdf')
            return open(f, 'wb')
        else:
            return None

    exp_result.write_to_selection_graphs(my_stream_getter,
                                         data_info,
                                         options.generate_all,
                                         colour=options.colour)


import unittest
import os
import json
from biothings_explorer.query_graph_handler.index import TRAPIQueryHandler

class TestTRAPIQueryHandler(unittest.TestCase):
    example_folder1 = os.path.abspath(
        os.path.join(os.path.dirname(__file__), 'data', 'increased_urinary_glycerol_affects_glycerol.json'))

    example_folder2 = os.path.abspath(
        os.path.join(os.path.dirname(__file__), 'data', 'FDFM_caused_by_ACDY5.json'))

    example_folder3 = os.path.abspath(
        os.path.join(os.path.dirname(__file__), 'data', 'chemicals_targeting_IL1_Signaling_Pathway.json'))

    # skip until we figure out why it returns no results
    @unittest.skip
    def test_when_looking_for_chemicals_affected_by_phenotype_increased_urinary_glycerol_glycerol_should_pop_up(self):
        query_handler = TRAPIQueryHandler({}, None, None, True)
        with open(self.example_folder1) as f1:
            query = json.load(f1)
            query_handler.set_query_graph(query['message']['query_graph'])
            query_handler.query()
            res = query_handler.get_response()
            self.assertIn('CHEBI:17754', res['message']['knowledge_graph']['nodes'])

    def test_when_looking_for_genes_related_to_disease_dyskinesia_familial_with_facial_myokymia_acdy5_should_pop_up(self):
        query_handler = TRAPIQueryHandler({}, None, None, True)
        with open(self.example_folder2) as f1:
            query = json.load(f1)
            query_handler.set_query_graph(query['message']['query_graph'])
            query_handler.query()
            res = query_handler.get_response()
            self.assertIn('NCBIGene:111', res['message']['knowledge_graph']['nodes'])

    # Can fail sometimes during testing, also slow to finish
    # Can throw requests.exceptions.SSLError
    @unittest.skip
    def test_when_looking_for_chemicals_targeting_il1_signaling_patway_curcumin_should_pop_up(self):
        query_handler = TRAPIQueryHandler({}, None, None, True)
        with open(self.example_folder3) as f1:
            query = json.load(f1)
            query_handler.set_query_graph(query['message']['query_graph'])
            query_handler.query()
            res = query_handler.get_response()
            self.assertIn('CHEBI:3962', res['message']['knowledge_graph']['nodes'])


from django.shortcuts import render
from rest_framework.decorators import api_view
from rest_framework.response import Response
from rest_framework import status
from .yahoo_finance import yahoo_finance_company_information, yahoo_finance_historical_data

def index(request):
    return render(request, "main/index.html")

@api_view(["GET"])
def get_yahoo_finance_company_information(request, company_name: str):
    information = yahoo_finance_company_information(company_name)

    return Response(
        {
            "information": information if information else {}
        },
        status.HTTP_200_OK if information else status.HTTP_404_NOT_FOUND
    )

@api_view(["GET"])
def get_yahoo_finance_historical_data(request, company_name: str):
    historical_data = yahoo_finance_historical_data(company_name)

    return Response(
        {
            "historical_data": historical_data if historical_data else []
        },
        status.HTTP_200_OK if historical_data else status.HTTP_404_NOT_FOUND
    )


from rest_framework import viewsets
from rest_framework import status
from rest_framework.decorators import action
from rest_framework.response import Response
from .serializers import MessageSerializer, TrainingDataSerializer
from .models import Message
import requests

class ChatbotViewSet(viewsets.ModelViewSet):
    """
    A viewset for interacting with the chatbot.
    """
    # Specify the queryset and serializer_class attributes for your viewset
    queryset = Message.objects.all()
    serializer_class = MessageSerializer

    @action(detail=False, methods=['post'])
    def message(self, request):
        """
        Endpoint for sending messages to the chatbot and receiving responses.
        """
        # Get an instance of the serializer with the request data
        serializer = self.get_serializer(data=request.data)

        # Validate the request data using the serializer
        serializer.is_valid(raise_exception=True)

        # Get the user's message from the validated data
        message = serializer.data['message']

        # Send the message to the Rasa server and get the response
        response = requests.post('http://0.0.0.0:5005/webhooks/rest/webhook', json={'message': message})
        response_data = response.json()

        # Return the chatbot's response
        return Response(response_data)

    @action(detail=False, methods=['post'])
    def train(self, request):
        """
        Endpoint for training the chatbot's machine learning models using training data.
        """
        # Specify a different serializer class for this action
        serializer = TrainingDataSerializer(data=request.data)

        # Validate the request data using the serializer
        serializer.is_valid(raise_exception=True)

        # Get the training data from the validated data
        training_data = serializer.data

        # Send the training data to the Rasa server and get the response
        response = requests.post('http://0.0.0.0:5005/model/train', json=training_data)
        response_data = response.json()

        # Return information about the training process
        return Response(response_data)


"""
Write a function that has a single argument (env name) and returns the "combined" dictionary
that merge 2 dictionaries together with the environment specific settings overriding any common settings already defined
"""
import json
from contextlib import ExitStack
from collections import ChainMap

# Declare setting files' name
SETTINGS_FILES = ['common.json', 'dev.json', 'prod.json']

# Read file
with ExitStack() as stack:
    SETTINGS = {
        f_name.split('.')[0]: json.load(stack.enter_context(open(f_name)))
        for f_name in SETTINGS_FILES
    }

# Choose which setting is the base
BASE_SETTINGS = 'common'

def chain_recursive_fred(overrided, base):
    chain = ChainMap(overrided, base)
    for k, v in overrided.items():
        if isinstance(v, dict) and k in base:
            chain[k] = chain_recursive_fred(v, base[k])
    return chain

def settings_for_env(env):
    result = chain_recursive_fred(SETTINGS[env], SETTINGS[BASE_SETTINGS])
    return result

dev = settings_for_env('dev')
print(dev)


import logging
from spinn_utilities.log import FormatAdapter
from spinn_utilities.config_holder import get_config_bool
from spinn_utilities.progress_bar import ProgressBar
from spinn_front_end_common.utilities.system_control_logic import (
    run_system_application)
from spinnman.model.enums import ExecutableType
from spinnman.model import ExecutableTargets
from spinnman.model.enums import CPUState
from spynnaker.pyNN.data import SpynnakerDataView
from spynnaker.pyNN.models.abstract_models import (
    AbstractNeuronExpandable, NEURON_EXPANDER_APLX)

logger = FormatAdapter(logging.getLogger(__name__))

def neuron_expander():
    """
    Run the neuron expander.

    .. note::
        Needs to be done after data has been loaded.
    """
    # Find the places where the neuron expander should run
    expander_cores, expanded_pop_vertices = _plan_expansion()

    progress = ProgressBar(expander_cores.total_processors,
                           "Expanding Neuron Data")
    expander_app_id = SpynnakerDataView.get_new_id()
    run_system_application(
        expander_cores, expander_app_id,
        get_config_bool("Reports", "write_expander_iobuf"),
        None, [CPUState.FINISHED], False,
        "neuron_expander_on_{}_{}_{}.txt", progress_bar=progress,
        logger=logger)
    progress.end()

    _fill_in_initial_data(expanded_pop_vertices)

def _plan_expansion():
    """
    Plan the expansion of neurons and set up the regions using USER1.

    :rtype: tuple(ExecutableTargets, list(tuple(MachineVertex, Placement)))
    """
    neuron_bin = SpynnakerDataView.get_executable_path(NEURON_EXPANDER_APLX)

    expander_cores = ExecutableTargets()
    expanded_pop_vertices = list()

    progress = ProgressBar(
        SpynnakerDataView.get_n_placements(),
        "Preparing to Expand Neuron Data")
    for placement in progress.over(SpynnakerDataView.iterate_placemements()):
        # Add all machine vertices of the population vertex to ones
        # that need synapse expansion
        vertex = placement.vertex
        if isinstance(vertex, AbstractNeuronExpandable):
            if vertex.gen_neurons_on_machine():
                expanded_pop_vertices.append((vertex, placement))
                expander_cores.add_processor(
                    neuron_bin, placement.x, placement.y, placement.p,
                    executable_type=ExecutableType.SYSTEM)
                # Write the region to USER1, as that is the best we can do
                txrx = SpynnakerDataView.get_transceiver()
                txrx.write_user(placement.x, placement.y, placement.p, 1,
                                vertex.neuron_generator_region)

    return expander_cores, expanded_pop_vertices

def _fill_in_initial_data(expanded_pop_vertices):
    """
    Once expander has run, fill in the connection data.

    :param list(MachineVertex, Placement) expanded_pop_vertices:
        List of machine vertices to read data from
    :param ~spinnman.transceiver.Transceiver transceiver:
        How to talk to the machine
    """
    progress = ProgressBar(
        len(expanded_pop_vertices), "Getting initial values")
    for vertex, placement in progress.over(expanded_pop_vertices):
        vertex.read_generated_initial_values(placement)


import buy as b
import sell as s
import BTC_price as p

import time
import requests
from binance.client import Client
import config
from binance.enums import *

global currency
global tha_dough   #the amount of USDT we wanna buy or sell
global mark_price
global BREAK_PRICE
global orderId
global diff
global TELE_TOKEN
global chat_id
global ping_pong
global old
global client
global ping_pong_count
global side
global ans
global boo
global SECONDS
global end_price
global step_size
global initial_BP
global stop_limit

boo = False

end_price = 0

TELE_TOKEN = "6034277983:AAHKdYppG-cw3FBII7wLisDoSbWfw_GwbHo"
TELE_TOKEN2 = "5903882275:AAGApm67fCPcuMHynovZGmS2yMjmQjnQhlM"
chat_id = "5789191868"
chat_id2 = "566644123"
currency = ""
ping_pong_count = 100000000
BREAK_PRICE = 1000000000
tha_dough = 0
diff = 100000
SECONDS = time.time()

def main(BREAK,AMOUN,PING,currenc,stop_lm):
    #try:
    global mark_price
    make_values(BREAK,AMOUN,PING)
    global currency
    global side
    global tha_dough   #the amount of USDT we wanna buy or sell
    global BREAK_PRICE
    global orderId
    global diff
    global TELE_TOKEN
    global chat_id
    global ping_pong
    global old
    global client
    global ping_pong_count
    global ans
    global SECONDS
    global boo
    global end_price
    global step_price
    global stop_limit
    global initial_BP

    initial_BP = BREAK_PRICE

    stop_limit = stop_lm

    client = Client(config.API_KEY,config.API_SECRET)

    old = time.time()

    ping_pong = 0

    currency = currenc

    message = "None"

    mark_price = round(float(p.price(currency)),1)

    diff = stop_lm

    coo = "llllllll"

    while True:
        mark_price = round(float(p.price(currency)),1)
        print(mark_price)
        if(mark_price > BREAK_PRICE + diff and mark_price < BREAK_PRICE + diff + 10):
            yu = b.buy(client,tha_dough , currency)
            coo = "long"
            break
        elif(mark_price < BREAK_PRICE - diff and mark_price > BREAK_PRICE - diff -10):
            yu = s.sell(client,tha_dough , currency)
            coo = "short"
            break
        else:
            pass

    print(yu)
    orderId = yu["orderId"]
    mo = coo + " INTITIATED OF " + str(tha_dough) + " IN CURRENCY  " + currency + " AT " + str(mark_price) + " BOT"

    message = coo + " INTITIATED OF " + str(tha_dough) + " IN CURRENCY  " + currency + " AT " + str(mark_price) + " BOT"
    url = f"https://api.telegram.org/bot{TELE_TOKEN}/sendMessage?chat_id={chat_id}&text={message}"
    url2 = f"https://api.telegram.org/bot{TELE_TOKEN2}/sendMessage?chat_id={chat_id2}&text={message}"
    requests.get(url).json()
    requests.get(url2).json()
    ans = "WHERES YOUR BUGGATI TOURNE DANS LA VIDE ANSI BA LA VIDA TRISTAN IS THE BEST"
    while True:
        print(f"{ans} IN CURRENCY {currency} AND BREAK PRICE IS {BREAK_PRICE}, the total amount {tha_dough} (in BTC/ETH) and the difference is {diff}")
        SECONDS = time.time()

        '''if(boo == True):
            if(old - SECONDS > 300):
                boo = False'''

        #lg.logg(mark_price)

        if(ping_pong == ping_pong_count and boo == False):
                message = f"PING PONG COUNT IS AT {ping_pong}"
                url = f"https://api.telegram.org/bot{TELE_TOKEN}/sendMessage?chat_id={chat_id}&text={message}"
                url2 = f"https://api.telegram.org/bot{TELE_TOKEN2}/sendMessage?chat_id={chat_id2}&text={message}"
                requests.get(url).json()
                requests.get(url2).json()
                ping_pong = 0
                #boo = True
                old = time.time()

        '''YEAR        = datetime.date.today().year     # the current year
        MONTH       = datetime.date.today().month    # the current month
        DATE        = datetime.date.today().day      # the current day
        HOUR        = datetime.datetime.now().hour   # the current hour
        MINUTE      = datetime.datetime.now().minute # the current minute
        #SECONDS     = datetime.datetime.now().second # the current second'''

        mark_price = round(float(p.price(currency)),1)
        print(mark_price)

        ans = choose(mark_price,side)

        if(ans == "long"):
            yu = b.buy(client,tha_dough *2 , currency)
            copa = "INITIATED"
        elif(ans == "short"):
            yu = s.sell(client,tha_dough *2 , currency)
            copa = "INITIATED"
        else:
            copa = "n"

        if(copa != "n"):
            print(yu)
            orderId = yu["orderId"]
            mo = ans + " INTITIATED OF " + str(tha_dough) + " IN CURRENCY  " + currency + " AT " + str(mark_price) + " BOT"

            message = "long INTITIATED OF " + str(tha_dough) + " IN CURRENCY  " + currency + " AT " + str(mark_price) + " BOT"
            url = f"https://api.telegram.org/bot{TELE_TOKEN}/sendMessage?chat_id={chat_id}&text={message}"
            url2 = f"https://api.telegram.org/bot{TELE_TOKEN2}/sendMessage?chat_id={chat_id2}&text={message}"
            requests.get(url).json()
            requests.get(url2).json()

    '''except Exception as e:
        print(e)
        message = "bot has stopped cunt"
        url = f"https://api.telegram.org/bot{TELE_TOKEN}/sendMessage?chat_id={chat_id}&text={message}"
        url2 = f"https://api.telegram.org/bot{TELE_TOKEN2}/sendMessage?chat_id={chat_id2}&text={message}"
        requests.get(url).json()
        requests.get(url2).json()
        sys.exit()'''

def make_values(value1, value2, value4=0):
    global BREAK_PRICE
    global mark_price
    global ping_pong_count
    global tha_dough
    global currency

    mark_price = 0

    BREAK_PRICE = round(float(value1),1)

    tha_dough = round(float(value2),3)

    currency = value4

    ping_pong_count = int(value4)

    print("BREAK PRICE: " + str(BREAK_PRICE) + " THE AMOUNT INVESTED(BTC/ETH): " + str(tha_dough) )
    #lg.logg(f"BREAK PRICE: {BREAK_PRICE}      THE AMOUNT INVESTED(BTC/ETH): {tha_dough}       THE DIFFERENCE FOR STOP LOSS IS: {diff}")

def choose(mark_price,side):
    global end_price
    global diff
    global stop_limit
    global initial_BP
    global BREAK_PRICE

    if(side == "l"):
        if(mark_price > (BREAK_PRICE + diff)):
            #BREAK_PRICE += diff
            return "long"
        else:
            return "blah"
    elif(side == "s"):
        if(mark_price < (BREAK_PRICE - diff)):
            #BREAK_PRICE -= diff
            return "short"
        else:
            return "blah"
    else:
        return "CHANGE THE BLAH"

def gui():
    #BREAK,AMOUN,PING,currenc,stop_lm
    a5 = str(input("ENTER THE CURRENCY PAIR: ")).upper()
    a2 = round(float(input("ENTER THE TOTAL AMOUNT YOU WANNA INVEST IN BTC: ")),2)
    a7 = int(input("ENTER THE STOP LIMIT DIFFERENCE(in USDT): "))
    a1 = int(input("ENTER THE START PRICE OF THE TRADE: "))
    a3 = int(input("ENTER THE THE NUMBER OF PING-PONGS YOU WANT TO BE REMINDED OF: "))
    main(a1,a2,a3,a5,a7)

gui()


import argparse
import os
from corescore.masks import CoreImageProcessor

def process_images(image_dir, labels):
    """Create masks for labelled images.
    For now, merge rock fragment labels from both core boxes"""
    coreProcessor = CoreImageProcessor("Images",
                                        labels=labels,
                                        merge_fragment_labels=True)
    for image in coreProcessor.core_types:
        mask_file = coreProcessor.processImage(image)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--images', help="Directory of images to make masks for")
    parser.add_argument('--labels', help="JSON file with labels for images")
    args = parser.parse_args()

    images = args.images
    if not images:
        images = 'Images/train'
    labels = args.labels
    if not labels:
        labels = 'Core_labels.json'

    process_images(images, labels)


'''
Final Project - Web Startup
userController.py
- JaeYoung Chang (jchang5)
- Maggie Farrell (mfarre22)
'''

import cherrypy
import re, json
from user_library import _user_database

class UserController(object):
    def __init__(self, udb=None):
        if udb is None:
            self.udb = _user_database()
        else:
            self.udb = udb

    def POST_USER(self, username):
        output = {'result': 'success'}

        data = json.loads(cherrypy.request.body.read().decode('utf-8'))

        self.udb.post_user(username, data)

        return json.dumps(output)

    def DELETE_USER(self, username):
        output = {'result': 'success'}

        self.udb.delete_user(username)
        return json.dumps(output)

    def GET_USERS(self):
        output = {'result' : 'success'}

        try:
            result = self.udb.get_users()
            if result is not None:
                output.update(result)
            else:
                output['result'] = 'error'
                output['message'] = 'users not found'
        except Exception as ex:
            output['result'] = 'error'
            output['message'] = str(ex)
        return json.dumps(output)


import numpy as np
from numpy import random
import cv2
def random_crop(image, boxes = None, tags = None,poly_fixed_points=True, crop_size=(640, 640), max_tries = 50):
    crop_h, crop_w = crop_size
    ori_h, ori_w, c = image.shape
    start_x_max = np.clip(ori_w - crop_w, 0, ori_w)
    start_y_max = np.clip(ori_h - crop_h, 0, ori_h)
    new_w = crop_w if start_x_max == 0 else ori_w
    new_h = crop_h if start_y_max == 0 else ori_h
    new_image = np.zeros((new_h, new_w, c))
    new_image[0:ori_h, 0:ori_w, :] = image
    if start_x_max==0 and start_y_max==0:
        return new_image, boxes, tags
    #ensure the croped area is not in box
    h_array = np.zeros((new_h), dtype=np.int32)
    w_array = np.zeros((new_w), dtype=np.int32)
    for box in boxes:
        box = np.round(box, decimals=0).astype(np.int32)
        minx = np.min(box[:, 0])
        maxx = np.max(box[:, 0])
        w_array[minx:maxx] = 1
        miny = np.min(box[:, 1])
        maxy = np.max(box[:, 1])
        h_array[miny:maxy] = 1

    # ensure the cropped area not across a text
    h_axis = np.where(h_array == 0)[0]
    w_axis = np.where(w_array == 0)[0]
    if len(h_axis) == 0 or len(w_axis) == 0:
        h , w, c = new_image.shape
        new_image = cv2.resize(new_image, (crop_w, crop_h))
        if poly_fixed_points==False:
            boxes = [box*[crop_w*1.0/w,crop_h*1.0/h] for box in boxes]
        else:
            boxes[:,:,0] *= crop_w*1.0/w
            boxes[:,:,1] *= crop_h*1.0/h
        return new_image, boxes, tags
    for i in range(max_tries):
        rand_start_x = random.randint(0, start_x_max+1)
        rand_end_x = rand_start_x + crop_w
        rand_start_y = random.randint(0, start_y_max+1)
        rand_end_y = rand_start_y + crop_h
        x_valid = rand_start_x in w_axis and rand_end_x in w_axis
        y_valid = rand_start_y in h_axis and rand_end_y in h_axis
        if not (x_valid and y_valid):
            continue
        if len(boxes) != 0:
            if poly_fixed_points==False:
                selected_boxes = []
                for idx, box in enumerate(boxes):
                    box_axis_in_area = (box[:, 0] >= rand_start_x) & (box[:, 0] <= rand_end_x) \
                                & (box[:, 1] >= rand_start_y) & (box[:, 1] <= rand_end_y)
                    if np.sum(box_axis_in_area) == box.shape[0]:
                        selected_boxes.append(idx)
            else:
                box_axis_in_area = (boxes[:, :, 0] >= rand_start_x) & (boxes[:, :, 0] <= rand_end_x) \
                                & (boxes[:, :, 1] >= rand_start_y) & (boxes[:, :, 1] <= rand_end_y)
                selected_boxes = np.where(np.sum(box_axis_in_area, axis=1) == boxes.shape[1])[0]
            if len(selected_boxes) > 0:
                break
        else:
            selected_boxes = []
            continue
    if i == max_tries-1:
        h , w, c = new_image.shape
        new_image = cv2.resize(new_image, (crop_w, crop_h))
        if poly_fixed_points==False:
            boxes = [box*[crop_w*1.0/w,crop_h*1.0/h] for box in boxes]
        else:
            boxes[:,:,0] *= crop_w*1.0/w
            boxes[:,:,1] *= crop_h*1.0/h
        return new_image, boxes, tags
    cropped_image = new_image[rand_start_y:rand_end_y, rand_start_x:rand_end_x,:]
    # boxes = boxes[selected_boxes]
    # tags = tags[selected_boxes]

    if poly_fixed_points==False:
        boxes = [boxes[idx] for idx in selected_boxes]
        tags = [tags[idx] for idx in selected_boxes]
        boxes = [box-[rand_start_x,rand_start_y] for box in boxes]
    else:
        boxes = boxes[selected_boxes]
        tags = tags[selected_boxes]
        boxes[:,:,0] -= rand_start_x
        boxes[:,:,1] -= rand_start_y
    return cropped_image, boxes, tags

def random_ratio_scale(image, boxes = None, tags = None,poly_fixed_points=True, ratios = np.arange(0.8,1.3,0.1)):
    rand_index = random.randint(0, len(ratios)-1)
    ratio = ratios[rand_index]
    ori_h, ori_w, c = image.shape
    new_h = int(ori_w * ratio)
    ratio_h = new_h*1.0 / ori_h
    new_image = cv2.resize(image, (ori_w, new_h))
    if poly_fixed_points==False:
        boxes = [box*[1,ratio_h] for box in boxes]
    else:
        boxes[:,:,1]*= ratio_h
    return new_image, boxes, tags
def random_resize(image, boxes = None, tags = None,poly_fixed_points=True, longer_sides=np.arange(640,2592, 32)):
    rand_index = random.randint(0, len(longer_sides)-1)
    longer_side = longer_sides[rand_index]
    ori_h, ori_w, c = image.shape
    if ori_h > ori_w:
        ratio_h = longer_side*1.0/ori_h
        # boxes[:,:,1]*= ratio_h
        if poly_fixed_points==False:
            boxes = [box*[1,ratio_h] for box in boxes]
        else:
            boxes[:,:,1]*= ratio_h
        new_image = cv2.resize(image, (ori_w, longer_side))
    else:
        ratio_w = longer_side*1.0/ori_w
        if poly_fixed_points==False:
            boxes = [box*[ratio_w,1] for box in boxes]
        else:
            # print(boxes)
            boxes[:,:,0]*= ratio_w
        # boxes[:,:,0]*= ratio_w
        new_image = cv2.resize(image, (longer_side, ori_h))
    return new_image, boxes, tags
def random_rotate(image, boxes = None, tags = None,poly_fixed_points=True,rotate_angles = np.arange(-10,10,1)):
    rand_index = random.randint(0, len(rotate_angles)-1)
    angle = rotate_angles[rand_index]
    ori_h, ori_w, _ = image.shape
    cX, cY = ori_w//2, ori_h//2
    matrix = cv2.getRotationMatrix2D((cX, cY), angle, 1.0)
    cos, sin = matrix[0,0], -matrix[0,1]
    abs_cos, abs_sin = np.abs(cos), np.abs(sin)
    nW = int((ori_h * abs_sin) + (ori_w * abs_cos))
    nH = int((ori_h * abs_cos) + (ori_w * abs_sin))
    matrix[0, 2] += (nW / 2) - cX
    matrix[1, 2] += (nH / 2) - cY
    new_image = cv2.warpAffine(image, matrix, (nW, nH))

    if poly_fixed_points==False:
        boxes = [box-[cX,cY] for box in boxes]
        new_boxes = []
        for box in boxes:
            temp_box = np.zeros_like(box)
            temp_box[:,0] = (box[:,0]*cos - box[:,1]*sin) + nW/2
            temp_box[:,1] = (box[:,1]*cos + box[:,0]*sin) + nH/2
            new_boxes.append(temp_box)
    else:
        boxes -= np.array([cX, cY])
        new_boxes = np.zeros_like(boxes)
        new_boxes[:,:,0] = (boxes[:,:,0]*cos - boxes[:,:,1]*sin) + nW/2
        new_boxes[:,:,1] = (boxes[:,:,1]*cos + boxes[:,:,0]*sin) + nH/2
    return new_image, new_boxes, tags


from utils_graph_learning import multi_class_accuracy
import torch.nn as nn
import numpy as np

def process_arguments(args):

    # ###### choose the substructures: usually loaded from networkx,
    # ###### except for 'all_simple_graphs' where they need to be precomputed,
    # ###### or when a custom edge list is provided in the input by the user

    # define if degree is going to be used as a feature and when (for each layer or only at initialization)

    # args['degree_as_tag'] = [args['degree_as_tag']] + [False for _ in range(args['num_layers']-1)]

    # define if existing features are going to be retained when the degree is used as a feature
    # args['retain_features'] = [args['retain_features']] + [True for _ in range(args['num_layers']-1)]

    # replicate d_out dimensions if the rest are not defined (msg function, mlp hidden dimension, encoders, etc.)
    # and repeat hyperparams for every layer
    if args['d_msg'] == -1:
        args['d_msg'] = [None for _ in range(args['num_layers'])]
    elif args['d_msg'] is None:
        args['d_msg'] = [args['d_out'] for _ in range(args['num_layers'])]
    else:
        args['d_msg'] = [args['d_msg'] for _ in range(args['num_layers'])]

    if args['d_h'] is None:
        args['d_h'] = [[args['d_out']] * (args['num_mlp_layers'] - 1) for _ in range(args['num_layers'])]
    else:
        args['d_h'] = [[args['d_h']] * (args['num_mlp_layers'] - 1) for _ in range(args['num_layers'])]

    args['d_out_edge_encoder'] = [args['d_out'] for _ in range(args['num_layers'])]

    args['d_out_node_encoder'] = args['d_out']

    args['d_out_id_embedding'] = args['d_out']

    args['d_out_degree_embedding'] = args['d_out']

    # repeat hyperparams for every layer
    args['d_out'] = [args['d_out'] for _ in range(args['num_layers'])]

    args['train_eps'] = [False for _ in range(args['num_layers'])]

    if len(args['final_projection']) == 1:
        args['final_projection'] = [args['final_projection'][0] for _ in range(args['num_layers'])] + [True]

    args['bn'] = [args['bn'] for _ in range(args['num_layers'])]
    args['dropout_features'] = [args['dropout_features'] for _ in range(args['num_layers'])] + [args['dropout_features']]

    # loss function & metrics
    if args['loss_fn'] == 'CrossEntropyLoss':
        assert args['regression'] is False, "Can't use Cross-Entropy loss in regression."
        loss_fn = nn.CrossEntropyLoss()
    elif args['loss_fn'] == 'BCEWithLogitsLoss':
        assert args['regression'] is False, "Can't use binary Cross-Entropy loss in regression."
        loss_fn = nn.BCEWithLogitsLoss()
    elif args['loss_fn'] == 'MSELoss':
        loss_fn = nn.MSELoss()
    elif args['loss_fn'] == 'L1Loss':
        loss_fn = nn.L1Loss()
    else:
        raise NotImplementedError

    if args['prediction_fn'] == 'multi_class_accuracy':
        assert args['regression'] is False, "Can't use Classification Accuracy metric in regression."
        prediction_fn = multi_class_accuracy
    elif args['prediction_fn'] == 'MSELoss':
        prediction_fn = nn.MSELoss(reduction='sum')
    elif args['prediction_fn'] == 'L1Loss':
        prediction_fn = nn.L1Loss(reduction='sum')
    elif args['prediction_fn'] == 'None':
        prediction_fn = None
    else:
        raise NotImplementedError

    if args['regression']:
        perf_opt = np.argmin
    else:
        perf_opt = np.argmax

    return args, loss_fn, prediction_fn, perf_opt


import json
import os
import datetime

from model import db
import dao
from utils import *

from flask import Flask, request

app = Flask(__name__)
db_filename = "backend-template.db"

app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///%s" % db_filename
app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
app.config["SQLALCHEMY_ECHO"] = True

db.init_app(app)
with app.app_context():
    db.create_all()

@app.route("/")
def welcome():
    return "Welcome to Backend Template! :)"

# USER ROUTES ----------------------------------------------
@app.route("/users/", methods=["GET"])
def get_users():
    return success_response([user.serialize() for user in dao.get_users()])

@app.route("/users/", methods=["POST"])
def create_user():
    body = json.loads(request.data)
    name = body.get("name")
    email = body.get("email")
    username = body.get("username")

    optional_user, err = dao.create_user(name, email, username)
    if not optional_user:
        return failure_response(err)
    return success_response(optional_user.serialize(), 201)

@app.route("/users/<int:user_id>/", methods=["GET"])
def get_user(user_id):
    user, err = dao.get_user_by_id(user_id)
    if not user:
        return failure_response(err)
    return success_response(user.serialize())

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 5000))
    app.run(host="0.0.0.0", port=port)


"""
# Sync Water API with USGS Sites

[USGS Site Web Service](https://waterservices.usgs.gov/rest/Site-Service.html)

[Understanding Output](https://waterservices.usgs.gov/rest/Site-Service.html#Understanding)

URL

- https://waterservices.usgs.gov/nwis/site/?format=rdb&stateCd={state}&period=P52W&siteType=LK,ST&siteStatus=all&hasDataTypeCd=iv,aw

## USGS Fields

| Field                 | Desc                              |
| --------------------- | --------------------------------- |
| agency_cd             | Agency                            |
| site_no               | Site identification number        |
| station_nm            | Site name                         |
| site_tp_cd            | Site type                         |
| dec_lat_va            | Decimal latitude                  |
| dec_long_va           | Decimal longitude                 |
| coord_acy_cd          | Latitude-longitude accuracy       |
| dec_coord_datum_cd    | Decimal Latitude-longitude datum  |
| alt_va                | Altitude of Gage/land surface     |
| alt_acy_va            | Altitude accuracy                 |
| alt_datum_cd          | Altitude datum                    |
| huc_cd                | Hydrologic unit code              |

"""
import csv
from datetime import datetime, timedelta
from io import StringIO
from textwrap import dedent

import helpers.water as water
import requests
from airflow import DAG
from airflow.operators.python import PythonOperator

# Default arguments
default_args = {
    "owner": "airflow",
    "depends_on_past": False,
    "start_date": (datetime.utcnow() - timedelta(days=2)).replace(minute=0, second=0),
    "catchup_by_default": False,
    "email_on_failure": False,
    "email_on_retry": False,
    "retries": 1,
    "retry_delay": timedelta(minutes=10),
    # 'max_active_runs':1,
    # 'concurrency':4,
}

with DAG(
    default_args=default_args,
    dag_id="a2w_sync_usgs_sites",
    tags=["a2w", "usgs"],
    schedule="@daily",
    doc_md=dedent(__doc__),
) as dag:

    ########################################
    def fetch_and_write(state_abbrev):

        url = f"https://waterservices.usgs.gov/nwis/site/?format=rdb&stateCd={state_abbrev}&period=P52W&siteType=LK,ST&siteStatus=all&hasDataTypeCd=iv,aw"
        # print(url)
        r = requests.get(url)
        # print(r.text)
        buff = StringIO(r.text)
        reader = csv.reader(buff, delimiter="\t")

        prepped_data = []
        keys = []
        result = []

        horizontal_datum = {"NAD83": 4269}
        vertical_datum = {"UNKNOWN": 0, "COE1912": 1, "NGVD29": 2, "NAVD88": 3}

        # Cleanup data before parse fields
        # Store header and actual data rows in new variable
        for idx, line in enumerate(reader):

            # only look at the header and data lines
            if line[0].strip() in ["agency_cd", "USGS"]:
                prepped_data.append(line)

        for idx, line in enumerate(prepped_data):

            if idx == 0:
                # this is the header
                keys = line
            else:
                # Build each line (object) by setting the keys and values
                _line = {}
                for i, k in enumerate(keys):
                    _line[k] = line[i].strip()
                result.append(_line)

        state_sites = []

        # Each line (after cleanup) in the results from the USGS API call represents a site
        for line in result:

            site = {}
            site["provider"] = "usgs"
            site["datatype"] = "usgs-site"
            site["code"] = line["site_no"].strip()
            site["state"] = state_abbrev

            # Geometry
            # --------
            geom = {}
            geom["type"] = "Point"
            try:
                geom["coordinates"] = [
                    float(line["dec_long_va"].strip()),
                    float(line["dec_lat_va"].strip()),
                ]
            except:
                geom["coordinates"] = [0, 0]
            site["geometry"] = geom

            # Attributes
            # ---------
            attributes = {}
            attributes["station_name"] = line["station_nm"].strip()
            attributes["site_type"] = line["site_tp_cd"].strip()
            site["attributes"] = attributes

            # site["site_number"] = line["site_no"].strip()
            # site["name"] = line["station_nm"].replace("'", "").strip()
            # site["state_abbrev"] = state_abbrev
            # site["elevation"] = (
            #     float(line["alt_va"].strip()) if line["alt_va"].strip() != "" else None
            # )
            # try:
            #     site["horizontal_datum_id"] = horizontal_datum[
            #         line["dec_coord_datum_cd"]
            #     ]
            # except:
            #     site["horizontal_datum_id"] = 4269
            # site["huc"] = (
            #     f"{line['huc_cd'].strip()}" if line["huc_cd"].strip() != "" else None
            # )
            # try:
            #     site["vertical_datum_id"] = vertical_datum[line["alt_datum_cd"]]
            # except:
            #     site["vertical_datum_id"] = vertical_datum["UNKNOWN"]

            state_sites.append(site)

        # POST to the water API
        # water.sync_usgs_sites(state_sites)
        water_hook = water.WaterHook(method="POST")
        resp = water_hook.request(
            endpoint="/providers/usgs/locations",
            json=state_sites,
        )

    states = water.get_states()

    for state in states:
        fetch_task_id = f"fetch_and_write_{state}"
        get_usgs_sites_task = PythonOperator(
            task_id=fetch_task_id,
            python_callable=fetch_and_write,
            op_kwargs={
                "state_abbrev": state,
            },
        )


import torch
from torch import nn, Tensor
import torch.nn.functional as F

from .utils import check_images, check_targets

class RCNN(nn.Module):
    """
    Abstract RCNN
    """
    def __init__(self, transform, backbone, rpn, rcnn_heads):
        super(RCNN, self).__init__()
        self.transform = transform
        self.backbone = backbone
        self.rpn = rpn
        self.rcnn_heads = rcnn_heads

    def forward(self, images, targets=None):
        """
        :param images: list[Tensor], raw images
        :param targets: list[dict[Tensor]], ground truth boxes and corresponding labels
                dict[Tensor]: targets['boxes']=tensor((N,4)), targets['labels']=tensor((N,))
        :return: detection results and losses
        """

        check_images(images)
        if self.training and not check_targets(targets):
            raise ValueError("Invalid targets, but I'm too lazy to tell the details.")

        # step 1: image and targets transform
        images, targets = self.transform(images, targets)

        # step 2: feed into convolutional backbone, generating feature map
        features = self.backbone(images.tensors)

        # step 3: feed features into rpn to generate proposals
        proposals, rpn_losses = self.rpn(images, targets, features)

        # step 4: feed proposals into rcnn head to get final prediction results
        detections, rcnn_losses = self.rcnn_heads(features, proposals, images.image_sizes, targets)

        # step 5: post process, e.g. transform box and image scales
        detections = self.transform.postprocess(detections, images.image_sizes, images.original_image_sizes)

        # losses in training
        losses = {}
        losses.update(rpn_losses)
        losses.update(rcnn_losses)

        return detections, losses

class RCNNHeadMLP(nn.Module):
    """
    RCNN fully connected layers
    """
    def __init__(self, in_channels, out_channels):
        super(RCNNHeadMLP, self).__init__()
        self.fc6 = nn.Linear(in_channels, out_channels)
        self.fc7 = nn.Linear(out_channels, out_channels)

    def forward(self, x):
        x = x.flatten(start_dim=1)
        x = F.relu(self.fc6(x))
        x = F.relu(self.fc7(x))
        return x

class RCNNPredictor(nn.Module):
    """
    RCNN predict head
    """
    def __init__(self, in_channels, num_classes):
        super(RCNNPredictor, self).__init__()
        self.cls_score = nn.Linear(in_channels, num_classes)
        self.bbox_reg = nn.Linear(in_channels, num_classes * 4)

    def forward(self, x):
        x = x.flatten(start_dim=1)
        cls = self.cls_score(x)
        reg = self.bbox_reg(x)

        return cls, reg


"""
TODO: copy readme?
      remove haplotypes stuff
      remove blast mode
      remove db stuff?

This is where the magic happens!

It is meant to be able to run locally and in pbs_runner.py there is also specific support for pbs cluster systems.

___Overview___
The pipeline is divided into 4 parts each having it's own .py file.
I   -  Preperation (data_preperation.py)
II  -  Processing (processing.py)
III -  Aggregation (aggregation.py)
IV  -  Visual Summary (summarize.py)

I - Data Prepeation
Input: directory containing fastq/gz files or a directory containing such directories.
Output: fastq files in sizes ready for efficient processing.

Depending on currently availble RAM and CPUs or given values in -mm / --max_memory and -cc / --cpu_count the script
will divide the files into an efficient number of files to later run the processing script on.
If given an merge_opposing it will also merge the forward and backward reads in corresponding fastq files.

II - Processing
Input: fastq files from part I and a reference fasta file.
Output: called_bases - every base called and its attributes.
        ignored_bases - every base ignored, its attributes and why it was ignored.
        suspicious_reads - every suspicious read and why it is suspicious.
        ignored_reads - every read ignored and why.
        read_counter - every read and a number describing how many times it was aligned by blast.
        blast files - alignment output files.

This is where the main logic of the pipeline happens and the runner runs this on each of the fastq files in parallel.
Basically it first aligns the reads to the reference with blast and then goes over every read and nucleotide and decides
whether to filter them out or leave them in.

III - Aggregation
Input: directory containing output of basecalling of stage II
Output: concatenations - concatenated outputs of stage II
        read_id_prefix_file - a json file containing a dictionary of read prefixes to save memory in the other files.
        mutation_read_list - a file describing in which reads each mutation appeared.
        freqs - a frequencies file describing the different alleles and their frequencies for each position.
        consensus_aligned_to_ref - a fasta file of the majority frequency derived from the freqs file aligned to the original reference
        consensus - a fasta file of the majority frequency derived from the freqs file
        read_counter - a file counting how many alignments were called for each read.

IV - Haplotype Inference
Input: directory containing output of stage III
Output: linked_mutations - file containing pairs of mutatations, their pvalues and their frequencies.
        stretches - file containing the aggregations of the stretches by their frequencies.
"""
import argparse
import getpass
import json
import os
import shutil
import concurrent.futures
import Bio
import multiprocessing as mp
import pandas as pd
from Bio import pairwise2, SeqIO
from datetime import datetime
from data_preparation import prepare_data
from summarize import graph_summary, create_stats_file
from processing import process_fastq
from aggregation import aggregate_processed_output, create_freqs_file, create_mutation_read_list_file
from logger import pipeline_logger
from utils import get_files_in_dir, get_sequence_from_fasta, get_mp_results_and_report, create_consensus_file, \
    get_files_by_extension, concatenate_files_by_extension, get_config, md5_dir, md5_file

def parallel_process(processing_dir, fastq_files, reference_file, quality_threshold, task, evalue, dust, num_alignments,
                     soft_masking, perc_identity, mode, reads_overlap):
    with concurrent.futures.ProcessPoolExecutor() as executor:
        future_tasks = {executor.submit(process_fastq, fastq_file, reference=reference_file, output_dir=processing_dir,
                                        quality_threshold=quality_threshold, task=task, evalue=evalue, dust=dust,
                                        num_alignments=num_alignments, soft_masking=soft_masking,
                                        perc_identity=perc_identity, mode=mode, reads_overlap=reads_overlap): fastq_file
                        for fastq_file in fastq_files}

        for future in concurrent.futures.as_completed(future_tasks):
            future.result()

def set_filenames(output_dir):
    filenames = {"freqs_file_path": os.path.join(output_dir, 'freqs.tsv'),
                 "linked_mutations_path": os.path.join(output_dir, 'linked_mutations.tsv'),
                 "mutation_read_list_path": os.path.join(output_dir, 'mutation_read_list.tsv'),
                 "stretches": os.path.join(output_dir, 'stretches.tsv'),
                 "blast_file": os.path.join(output_dir, 'blast.tsv'),
                 "read_counter_file": os.path.join(output_dir, 'read_counter.tsv'),
                 "summary_graphs": os.path.join(output_dir, 'summary.png'),
                 "processing_dir": os.path.join(output_dir, "processing"),
                 'linked_mutations_dir': os.path.join(output_dir, "linked_mutations"),
                 'data_dir': os.path.join(output_dir, "data")}
    filenames['basecall_dir'] = os.path.join(filenames['processing_dir'], 'basecall')
    os.makedirs(filenames['data_dir'], exist_ok=True)
    os.makedirs(filenames["basecall_dir"], exist_ok=True)
    return filenames

def create_consensus_and_check_alignment_with_ref(reference_file, align_to_ref, min_coverage, iteration_data_dir, basecall_dir,
                                                  iteration_counter, min_frequency):
    reference = get_sequence_from_fasta(reference_file)
    freqs_file_path = os.path.join(iteration_data_dir, f"freqs_{iteration_counter}.tsv")
    called_bases_files = get_files_by_extension(basecall_dir, "called_bases")
    create_freqs_file(called_bases_files=called_bases_files, output_path=freqs_file_path, reference=reference)
    if align_to_ref == "Y":
        consensus_path = os.path.join(iteration_data_dir, f"consensus_aligned_to_ref_{iteration_counter}.fasta")
        align_to_ref = True
    else:
        consensus_path = os.path.join(iteration_data_dir, f"consensus_{iteration_counter}.fasta")
        align_to_ref = False
    create_consensus_file(freqs_file=freqs_file_path, min_frequency=min_frequency,
                          min_coverage=min_coverage, output_file=consensus_path, align_to_ref=align_to_ref)
    consensus = get_sequence_from_fasta(consensus_path)
    #TODO: get helpdesk to create the right environment and remove this crap..
    alignment = pairwise2.align.globalxx(consensus, reference)[0]
    if float(Bio.__version__) > 1.76:
        alignment_score = alignment.score / max(len(consensus), len(reference))
        alignment = alignment.seqA
    else:
        alignment_score = int(alignment[2]) / max(len(consensus), len(reference))
        alignment = alignment[0]
    return alignment_score, alignment

def get_consensus_path(basecall_iteration_counter, align_to_ref, iteration_data_dir):
    consensus = os.path.join(iteration_data_dir, "consensus")
    if align_to_ref == "Y":
        consensus += "_aligned_to_ref"
    consensus += f"_{basecall_iteration_counter}.fasta"
    return consensus

def update_meta_data(output_dir, status, db_path, params=None):
    # TODO: write running time in meta_data
    json_file = os.path.join(output_dir, 'meta_data.json')
    if params is not None:
        meta_data = params
        del meta_data['filenames']
        del meta_data['log']
        meta_data['username'] = getpass.getuser()
        meta_data['start_time'] = datetime.now().strftime('%Y-%m-%d-%H:%M')
    else:
        with open(json_file) as read_handle:
            meta_data = json.load(read_handle)
    meta_data['status'] = status
    with open(json_file, 'w') as write_handle:
        json.dump(meta_data, write_handle)
    build_db(db_path)

def build_db(db_path):
    # TODO: get user feedback on what db functionality should be
    os.makedirs(db_path, exist_ok=True)
    outputs = [f.path for f in os.scandir(db_path) if f.is_dir()]
    db_rows = []
    for directory in outputs:
        json_file = os.path.join(directory, "meta_data.json")
        if os.path.isfile(json_file):
            with open(json_file) as read_handle:
                meta_data = json.load(read_handle)
            db_rows.append(meta_data)
    db = pd.DataFrame.from_dict(db_rows)
    db.to_csv(os.path.join(db_path, 'db.tsv'), sep='\t')

def create_dir_also_if_exists(output_dir):
    try:
        os.makedirs(output_dir, exist_ok=False)
    except FileExistsError:
        i = 2
        while i < 1000:
            try:
                output_dir = f"{output_dir}_{i}"
                os.makedirs(output_dir, exist_ok=False)
            except FileExistsError:
                i += 1
    return output_dir

def assign_output_dir(db_path, alias=None):
    now = datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
    output_dir_name = now
    if alias:
        output_dir_name = alias + "_" + output_dir_name
    output_dir = os.path.join(db_path, output_dir_name)
    output_dir = create_dir_also_if_exists(output_dir)
    return output_dir

def process_data(align_to_ref, dust, evalue, fastq_files, log, max_basecall_iterations, min_frequency,
                 min_coverage, mode, num_alignments, overlapping_reads, output_dir, perc_identity, processing_dir,
                 quality_threshold, reference_file, soft_masking, task, basecall_dir):
    alignments = []
    for basecall_iteration_counter in range(1, max_basecall_iterations + 1):
        log.info(f"Processing fastq files iteration {basecall_iteration_counter}/{max_basecall_iterations}")
        parallel_process(processing_dir=processing_dir, fastq_files=fastq_files, reference_file=reference_file,
                         quality_threshold=quality_threshold, task=task, evalue=evalue, dust=dust, mode=mode,
                         num_alignments=num_alignments, soft_masking=soft_masking, perc_identity=perc_identity,
                         reads_overlap=overlapping_reads)
        iteration_data_dir = os.path.join(output_dir, 'iteration_data')
        os.makedirs(iteration_data_dir, exist_ok=True)
        alignment_score, alignment = create_consensus_and_check_alignment_with_ref(reference_file=reference_file,
                                                                                iteration_counter=basecall_iteration_counter,
                                                                                basecall_dir=basecall_dir,
                                                                                align_to_ref=align_to_ref,
                                                                                iteration_data_dir=iteration_data_dir,
                                                                                min_coverage=min_coverage,
                                                                                min_frequency=min_frequency)
        alignments.append(alignment)
        log.info(f'Iteration alignment score: {round(alignment_score, 4)}')
        if alignment_score == 1:
            break
        consensus_path = get_consensus_path(basecall_iteration_counter=basecall_iteration_counter,
                                            align_to_ref=align_to_ref, iteration_data_dir=iteration_data_dir)
        reference_file = consensus_path
    return reference_file, alignments, basecall_iteration_counter

def validate_input(output_dir, input_dir, reference_file, mode):
    if mode != 'RefToSeq' and mode != 'SeqToRef':
        raise Exception("blast mode must be either RefToSeq or SeqToRef! ")
    if os.path.exists(output_dir):
        list_dir = os.listdir(output_dir)
        if len(list_dir) > 0:
            if not ((len(list_dir) == 1) and (list_dir[0] == 'pbs_logs')):
                raise Exception("output_dir must be path to a new or empty directory!")
    if not os.path.isfile(reference_file):
        raise Exception("reference_file must exist!")
    else:
        with open(reference_file, "r") as handle:
            fasta = SeqIO.parse(handle, "fasta")
            if not any(fasta):
                raise Exception("reference_file must be of type fasta!")
    if not os.path.isdir(input_dir):
        raise Exception("Input_dir must exist!")
    files_fasta = get_files_by_extension(input_dir, "fastq")
    files_fastagz = get_files_by_extension(input_dir, "fastq.gz")
    if len(files_fasta) == 0 and len(files_fastagz) == 0:
        raise Exception("Could not find files ending with '.fastq' or 'fastq.gz' in input_dir !")

def try_to_rmtree(path_to_delete, retry_attempts=5):
    # this is a hack to avoid a bug rmtree throws about a non-empty dir sometimes...
    if retry_attempts:
        try:
            if os.path.isdir(path_to_delete):
                shutil.rmtree(path_to_delete)
        except:
            try_to_rmtree(path_to_delete, retry_attempts-1)

def remove_unnecessary_files(dirs_to_remove, files_to_remove):
    for file_path in files_to_remove:
        os.remove(file_path)
    for dir_path in dirs_to_remove:
        try_to_rmtree(dir_path)

def runner(input_dir, reference_file, output_dir, max_basecall_iterations, min_coverage, db_comment,
           quality_threshold, task, evalue, dust, num_alignments, soft_masking, perc_identity, mode,
           align_to_ref, cleanup, min_frequency, cpu_count, overlapping_reads,
           db_path, max_memory):
    if not db_path:
        db_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'db')
    if not output_dir:
        output_dir = assign_output_dir(db_path)
    validate_input(output_dir, input_dir, reference_file, mode)
    log = pipeline_logger(logger_name='AccuNGS-Runner', log_folder=output_dir)
    try:
        filenames = set_filenames(output_dir=output_dir)
        if not cpu_count:
            cpu_count = mp.cpu_count()
        input_dir_hash = md5_dir(input_dir)
        reference_file_hash = md5_file(reference_file)
        params = locals().copy()
        update_meta_data(params=params, output_dir=output_dir, status='Setting up...', db_path=db_path)
        log.debug(f"runner params: {params}")  # TODO: why does this contain status..?
        log.info("Preparing data")
        update_meta_data(output_dir=output_dir, status='Preparing data...', db_path=db_path)
        prepare_data(input_dir=input_dir, output_dir=filenames['data_dir'], overlapping_reads=overlapping_reads,
                     cpu_count=cpu_count, max_memory=max_memory)
        data_files = get_files_in_dir(filenames['data_dir'])
        fastq_files = [file_path for file_path in data_files if "fastq.part_" in os.path.basename(file_path)]
        log.info(f"Processing {len(fastq_files)} fastq files.")
        update_meta_data(output_dir=output_dir, status='Processing data...', db_path=db_path)
        reference_file, alignments, iterations = process_data(
            align_to_ref=align_to_ref, dust=dust, min_frequency=min_frequency,
            evalue=evalue, fastq_files=fastq_files, log=log, soft_masking=soft_masking,
            max_basecall_iterations=max_basecall_iterations, min_coverage=min_coverage,
            mode=mode, num_alignments=num_alignments, overlapping_reads=overlapping_reads,
            output_dir=output_dir, perc_identity=perc_identity, reference_file=reference_file,
            processing_dir=filenames['processing_dir'], quality_threshold=quality_threshold,
            task=task, basecall_dir=filenames['basecall_dir'])
        log.info("Aggregating processed fastq files outputs...")
        last_freqs = os.path.join(output_dir, 'iteration_data', f'freqs_{iterations}.tsv')
        shutil.copy(last_freqs, filenames['freqs_file_path'])
        aggregate_processed_output(input_dir=filenames['processing_dir'], output_dir=output_dir,
                                   min_coverage=min_coverage, min_frequency=min_frequency, cleanup=cleanup)
        create_stats_file(output_dir, filenames, alignments)
        log.info("Generating graphs...")
        graph_summary(freqs_file=filenames['freqs_file_path'], blast_file=filenames['blast_file'],
                      read_counter_file=filenames['read_counter_file'],
                      output_file=filenames['summary_graphs'], min_coverage=min_coverage)  # TODO: drop low quality mutations?
        log.info(f"Most outputs are ready in {output_dir} !")
        if cleanup == "Y":
            dirs_to_remove = [filenames['basecall_dir'], filenames['data_dir']]
            files_to_remove = [filenames['blast_file']]
            log.info(f"Removing intermediary files to save space...")
            remove_unnecessary_files(dirs_to_remove, files_to_remove)
        update_meta_data(output_dir=output_dir, status='Done', db_path=db_path)
        log.info(f"Done!")
    except Exception as e:
        log.exception(e)
        update_meta_data(output_dir=output_dir, status="Failed! see logs for details.", db_path=db_path)

def create_runner_parser():
    # TODO: dynamic defaults?
    parser = argparse.ArgumentParser(description="Note: Default parameter values are retrieved from config.ini in your "
                                                 "installation directory.")
    parser.add_argument("-i", "--input_dir", required=True,
                        help="Path to directory containing fastq/gz files or sub directories containg fastq/gz files.")
    parser.add_argument("-o", "--output_dir", help="A directory for output files. "
                                                   "If none is given will put it in the db")
    parser.add_argument("-r", "--reference_file", required=True, help="Full path to reference file (including "
                                                                      "extension) of type fasta to align against.")
    parser.add_argument("-m", "--max_basecall_iterations", type=int,
                        help="Number of times to rerun with previous consensus as the new reference before giving up.")
    parser.add_argument("-or", "--overlapping_reads",
                        help="Y/N/P, run pipeline with, without, or with partial overlapping reads. Y- merge opposing "
                             "reads in the same directory and drop non overlapping areas of the reads. "
                             "P - Merge opposing reads but keep non overlapping areas."
                             "N - No merge, assume reads are independent. "
                             "Y & P assume 2 fastq/gz files in each sub directory of the input_dir.")
    parser.add_argument("-bt", "--blast_task", help="blast's task parameter")
    parser.add_argument("-be", "--blast_evalue", help="blast's e value parameter", type=float)
    parser.add_argument("-bd", "--blast_dust", help="blast's dust parameter")
    parser.add_argument("-bn", "--blast_num_alignments", type=int, help="blast's num_alignments parameter")
    parser.add_argument("-bp", "--blast_perc_identity", type=int, help="blast's perc_identity parameter")
    parser.add_argument("-bs", "--blast_soft_masking", help="blast's soft_masking parameter")
    parser.add_argument("-bm", "--blast_mode", help="RefToSeq or SeqToRef")  # TODO: docs
    parser.add_argument("-qt", "--quality_threshold", type=int,
                        help="phred score must be higher than this to be included")
    parser.add_argument("-mc", "--min_coverage", type=int,
                        help="positions with less than this coverage will be substituted by Ns in the consensus")
    parser.add_argument("-mf", "--min_frequency", type=float,
                        help="positions with less than this frequency will be substituted by Ns in the consensus")
    parser.add_argument("-ar", "--align_to_ref", help="Y/N, generate consensus aligned to the original reference")
    parser.add_argument("-c", "--cleanup", help="Remove intermediary files when done in order to save space")
    parser.add_argument("-cc", "--cpu_count", help="max number of cpus to use (None means all)", type=int)
    parser.add_argument("-db", "--db_path", help='path to db directory')
    parser.add_argument("-dbc", "--db_comment", help='comment to store in db')
    parser.add_argument("-mm", "--max_memory", help='limit memory usage to this many megabytes '
                                                    '(None would use available memory when starting to run)')
    parser.add_argument("-ch", "--calculate_haplotypes", help='Y/N, Run pipeline including calculating haplotypes')
    return parser

if __name__ == "__main__":
    parser = create_runner_parser()
    parser_args = vars(parser.parse_args())
    args = dict(get_config()['runner_defaults'])
    args.update({key: value for key, value in parser_args.items() if value is not None})
    runner(input_dir=args['input_dir'], output_dir=args['output_dir'], reference_file=args['reference_file'],
           max_basecall_iterations=int(args['max_basecall_iterations']), overlapping_reads=args['overlapping_reads'],
           quality_threshold=int(args['quality_threshold']), task=args['blast_task'], max_memory=args['max_memory'],
           evalue=float(args['blast_evalue']), dust=args['blast_dust'],
           num_alignments=int(args['blast_num_alignments']),
           mode=args['blast_mode'], perc_identity=float(args['blast_perc_identity']), cpu_count=args['cpu_count'],
           min_coverage=int(args['min_coverage']), db_comment=args['db_comment'],
           soft_masking=args['blast_soft_masking'], min_frequency=float(args['min_frequency']),
           cleanup=args['cleanup'], align_to_ref=args['align_to_ref'], db_path=args['db_path'])


import collections
import torch
import numpy as np
import os
import h5py
import time
import math
import platform
import typing
from torch import Tensor

def count_net_params(net):
    n_param = 0
    for name, param in net.named_parameters():
        layer_param = torch.numel(param)
        print("::: Layer Name: ", name, "Layer Size", layer_param)
        n_param += layer_param
    return n_param

def count_rnn_params(net):
    n_param = 0
    for name, param in net.named_parameters():
        if 'rnn' in name:
            sizes = 1
            for el in param.size():
                sizes = sizes * el
            n_param += sizes
    return n_param

def get_dataset_path(args):
    # Dictionary for datasets
    if args.dataset_name == 'gscdv1':
        print("Preparing Google Speech Command Dataset...")
        dataset_name = 'speech_commands_v0.01'
        testset_name = 'speech_commands_test_set_v0.01'
    elif args.dataset_name == 'gscdv2':
        print("Preparing Google Speech Command Dataset...")
        dataset_name = 'speech_commands_v0.02'
        testset_name = 'speech_commands_test_set_v0.02'
    elif args.dataset_name == 'timit':
        print("Preparing TIMIT Dataset...")
        dataset_name = 'TIMIT/TIMIT'
        testset_name = None
    elif args.dataset_name == 'tidigits':
        print("Preparing TIDIGITS Dataset...")
        dataset_name = 'tidigits'
        testset_name = None
    else:
        raise RuntimeError('Dataset not supported.')

    data_dir = args.data_dir
    dataset_path = os.path.join(data_dir, dataset_name)
    testset_path = os.path.join(data_dir, testset_name)

    return dataset_path, testset_path

def load_h5py_data(f):
    dict_data = {}
    for k in f.keys():
        dict_data[k] = np.array(f[k])
    return dict_data

def log_lut(x, qi_in, qf_in, qi_out, qf_out):
    lut_in = quantize_array(x, qi_in, qf_in, 1)
    lut_out = np.log10(lut_in + 1)
    lut_out = quantize_array(lut_out, qi_out, qf_out, 1)
    return lut_out

def gen_description_file(augment_noise, target_snr):
    if augment_noise:
        description_file = 'description_snr' + str(target_snr) + '.csv'
    else:
        description_file = 'description.csv'
    return description_file

def piecewise_linear_log(x,
                         k0, b0,
                         k1, b1,
                         k2, b2,
                         k3, b3,
                         k4, b4,
                         k5, b5,
                         k6, b6,
                         k7, b7,
                         k8, b8,
                         k9, b9
                         ):
    return np.piecewise(x=x,
                        condlist=[x >= 0,
                                  x >= 6,
                                  x >= 12,
                                  x >= 18,
                                  x >= 24,
                                  x >= 108,
                                  x >= 192,
                                  x >= 276,
                                  x >= 1549,
                                  x >= 2792],
                        funclist=[lambda x: k0 * x + b0,
                                  lambda x: k1 * x + b1,
                                  lambda x: k2 * x + b2,
                                  lambda x: k3 * x + b3,
                                  lambda x: k4 * x + b4,
                                  lambda x: k5 * x + b5,
                                  lambda x: k6 * x + b6,
                                  lambda x: k7 * x + b7,
                                  lambda x: k8 * x + b8,
                                  lambda x: k9 * x + b9]
                        )

def create_folder(folder_list):
    for folder in folder_list:
        try:
            os.makedirs(folder)
        except:
            pass

def load_model(proj, net, model_path):
    if proj.use_cuda and torch.cuda.is_available():
        model_location = 'cuda:' + str(proj.gpu_device)
    else:
        model_location = 'cpu'
    pretrained_model = torch.load(model_path, map_location=model_location)
    dict_pretrained_model = dict(pretrained_model.items())
    list_pretrained_model = list(dict_pretrained_model.items())
    dict_converted_model = {}
    net_state_dict = net.state_dict()
    list_trained_model = list(net_state_dict.items())
    num_pretrain_keys = len(list_pretrained_model)
    num_train_keys = len(list_trained_model)

    # Convert Model
    print("num_pretrain_keys: ", num_pretrain_keys)
    print("num_train_keys:    ", num_train_keys)
    if num_pretrain_keys != num_train_keys:
        for k_model, v_model in dict_pretrained_model.items():
            if 'ih' in k_model:
                if 'weight' in k_model:
                    param_ih = v_model
                    continue
                if 'bias' in k_model:
                    param_ih = v_model
                    continue
            elif 'hh' in k_model:
                if 'weight' in k_model:
                    param_hh = v_model
                    param_cat = torch.cat((param_ih, param_hh), dim=-1)
                if 'bias' in k_model:
                    param_hh = v_model
                    param_cat = param_ih + param_hh
                dict_converted_model[k_model] = param_cat
            else:
                dict_converted_model[k_model] = v_model
        list_converted_model_params = list(dict_converted_model.items())
    else:
        list_converted_model_params = list_pretrained_model
    # Copy param of the pretrained model to the training network
    new_state_dict = collections.OrderedDict()
    count = 0
    for key, value in net.state_dict().items():
        layer_name, weights = list_converted_model_params[count]
        net_state_dict[key] = weights
        print("Pre-trained Layer: %s - Loaded into new layer: %s" % (layer_name, key))
        count += 1

    net.load_state_dict(net_state_dict)
    return net

# def load_model(net, model_path, gpu_device):
#     # Get adapted DeltaRNN dimensions
#     # adapted_input_size = int(math.ceil(input_size / float(num_array)) * num_array)
#     # adapted_hidden_size = int(math.ceil(hidden_size / float(num_array)) * num_array)
#     # Load pretrained model
#     pretrained_net = torch.load(model_path, map_location='cuda:' + str(gpu_device))
#     pretrained_net = list(pretrained_net.items())
#
#     # Iterate throughput retrain net parameters to load pretrained net parameters correctly
#     new_state_dict = collections.OrderedDict()
#     count = 0
#     num_param_key = len(pretrained_net)
#     for key, value in net.state_dict().items():
#         if count >= num_param_key:
#             break
#         layer_name, weights = pretrained_net[count]
#         new_state_dict[key] = weights
#         print("Pre-trained Layer: %s - Loaded into new layer: %s" % (layer_name, key))
#         count += 1
#     net.load_state_dict(new_state_dict)
#     return net

def timeSince(since):
    now = time.time()
    s = now - since
    m = math.floor(s / 60)
    s -= m * 60
    return '%dm %ds' % (m, s)

def create_context_window(tensor, window_size=2):
    """
    Create a new tensor with context window
    :param tensor: 3D tensor (T, N, F)
    :param window_size: Integer
    :return:
    """
    tensor = tensor.transpose(0, 1)
    dim_T = tensor.shape[0]
    dim_N = tensor.shape[1]
    dim_F = tensor.shape[2]
    dim_FC = int(dim_F / 3)
    dim_C = int(window_size * 3)
    pad_len = window_size - 1
    pad = torch.zeros((pad_len, dim_N, dim_F)).cuda()
    padded_dim_T = dim_T + pad_len
    feature = torch.cat((tensor, pad), dim=0)
    feature = torch.reshape(feature, (padded_dim_T, dim_N, 3, dim_FC))
    output = []
    for tstep in range(0, dim_T):
        temp = feature[tstep:tstep + window_size, :, :, :]
        # print(temp.shape)
        temp = temp.transpose(0, 1)
        temp = torch.reshape(temp, (dim_N, dim_C, dim_FC))
        output.append(temp)
    output = torch.stack(output)
    return output, dim_C, dim_FC

def write_hdf5_normal(filename, X, y):
    # fun
    feature_lens = np.asarray([sample.shape[1] for sample in X]).astype(np.float32)
    label_lens = np.asarray([len(target) for target in y]).astype(np.float32)
    features = np.concatenate(X, axis=1).T.astype(np.float32)
    labels = np.concatenate(y).astype(np.float32)

    with h5py.File(filename, 'w') as f:
        f.create_dataset('features', data=features)
        f.create_dataset('feature_lens', data=feature_lens)
        f.create_dataset('labels', data=labels)
        f.create_dataset('label_lens', data=label_lens)

def print_size_of_model(model):
    torch.save(model.state_dict(), "temp.p")
    print('Size (MB):', os.path.getsize("temp.p") / 1e6)
    os.remove('temp.p')

def to_sparse(x):
    """ converts dense tensor x to sparse format """
    x_typename = torch.typename(x).split('.')[-1]
    sparse_tensortype = getattr(torch.sparse, x_typename)

    indices = torch.nonzero(x)
    if indices.nelement() == 0:  # if all elements are zeros
        print("1", indices)
        return sparse_tensortype(*x.shape)
    else:
        print("2", indices)
        indices = indices.torch()
        values = x[tuple(indices[i] for i in range(indices.shape[0]))]
        return sparse_tensortype(indices, values, x.size())

def quantize_tensor(x: Tensor, qi: int, qf: int, en: int = 0, floor=False):
    """
    :param x: input tensor
    :param qi: number of integer bits before the decimal point
    :param qf: number of fraction bits after the decimal point
    :param en: Type of Quantziation (0 - none, 1 - round, 2 - floor)
    :param floor: Whether use floor() instead of round()
    :return: tensor quantized to fixed-point precision
    """
    if en == 0:
        return x
    else:
        power = float(2. ** qf)
        clip_val = float(2. ** (qi + qf - 1))
        if floor:
            value = torch.floor(x * power)
        else:
            value = torch.round(x * power)
        value = torch.clamp(value, -clip_val, clip_val - 1)  # saturation arithmetic
        value = value / power
        return value

def quantize_array(x, qi, qf, enable, unsigned=0, floor=False):
    """
    :param x: input numpy array
    :param qi: number of integer bits before the decimal point
    :param qf: number of fraction bits after the decimal point
    :param enable: whether enable the function
    :param unsigned: whether unsigned input
    :return: array quantized to fixed-point precision
    """
    if enable == 0:
        return x
    power = np.asarray(2. ** qf).astype(np.float64)
    if unsigned == 0:
        clip_val = float(2. ** (qi + qf - 1))
        if floor:
            value = np.floor(x * power)
        else:
            value = np.round(x * power)
        value = np.clip(value, -clip_val, clip_val - 1)  # saturation arithmetic
        result = value / power
    else:
        clip_val = float(2. ** (qi + qf))
        if floor:
            value = np.floor(x * power)
        else:
            value = np.round(x * power)
        value = np.clip(value, 0, clip_val - 1)  # saturation arithmetic
        result = value / power
    return result

# Auto determine minimal possible qi
def quantize_tensor_auto(x, bw, enable):
    if enable == 0:
        return x, 0, 0
    qi = torch.ceil(torch.log2(x)) + 1
    qf = bw - qi
    # if qi < 0 or qf < 0:
    #     print('Quantizing {} to Q({}, {})!'.format(x, qi, qf))
    power = float(2. ** qf)
    clip_val = float(2. ** (qi + qf - 1) - 1)
    value = torch.round(x * power)
    value = torch.clamp(value, -clip_val, clip_val)  # saturation arithmetic
    value = value / power
    return value, qi, qf

def quantize_rnn(net, qi, qf, enable):
    for name, param in net.named_parameters():
        if 'rnn' in name and 'swap' not in name:
            # print("Quantizing Parameter: ", name)
            param.features = quantize_tensor(param.features, qi, qf, enable)
    return net

def quantize_net(net, qi, qf, enable):
    for name, param in net.named_parameters():
        param.features = quantize_tensor(param.features, qi, qf, enable)
    return net

def hard_quantize_tensor(x, m, n, enable=0):
    """
    :param x: input tensor
    :param m: number of integer bits before the decimal point
    :param n: number of fraction bits after the decimal point
    :return: tensor quantized to fixed-point precision
    """

    if not enable:
        return x

    power = 2. ** n
    max_val = 2. ** (m + n - 1) - 1
    value = x * power
    value = torch.floor(value)  # rounding
    value = torch.clamp(value, -max_val, max_val)  # saturation arithmetic
    value = value / power
    return value

def get_temporal_sparsity(list_layer, seq_len, threshold):
    # Evaluate Sparsity
    num_zeros = 0
    num_elems = 0
    # print(seq_len.size())
    # Iterate through layers
    for layer in list_layer:
        all_delta_vec = layer.transpose(0, 1)
        all_delta_vec = torch.abs(all_delta_vec)  # Take absolute values of all delta vector elements
        for i, delta_vec in enumerate(all_delta_vec):
            seq = delta_vec[:seq_len[i], :]
            zero_mask = seq < threshold
            num_zeros += torch.sum(zero_mask)
            num_elems += torch.numel(zero_mask)
    sparsity = float(num_zeros) / float(num_elems)
    return sparsity

# def targeted_dropout(x, gamma, alpha, epoch):
#     """
#     :param x: input tensor
#     :param m: number of integer bits before the decimal point
#     :param n: number of fraction bits after the decimal point
#     :return: tensor quantized to fixed-point precision
#     """
#     import math
#     torch.manual_seed(epoch)
#     torch.cuda.manual_seed_all(epoch)
#     n_dims = len(x.size())
#     n_elements = x.numel()
#     drop_part = math.ceil(n_elements * gamma)
#     weight_vec = x.view(-1)
#     weight_vec_abs = torch.abs(weight_vec)
#     sorted, indices = torch.sort(weight_vec_abs)
#     # print(sorted)
#     drop_indices = indices[0:drop_part]
#     drop_rand_mask = torch.rand(drop_indices.size(0)).cuda()
#     drop_mask = torch.ones(drop_indices.size(0)).cuda()
#     drop_mask = drop_mask.masked_fill_(drop_rand_mask <= alpha, 0)
#     weight_vec[drop_indices] *= drop_mask
#     if n_dims == 1:
#         weight = weight_vec
#     elif n_dims > 1:
#         weight = torch.reshape(weight_vec, (x.size(0), x.size(1)))
#
#     return weight

def aligned_cbwdrop(x, gamma, alpha, num_pe):
    """
    :param x: input tensor (weight matrix)
    :param m: number of integer bits before the decimal point
    :param n: number of fraction bits after the decimal point
    :return: tensor quantized to fixed-point precision
    """
    import math

    n_dims = len(x.size())
    n_rows = x.size(0)

    if n_dims == 1:
        n_cols = 1
    elif n_dims > 1:
        n_cols = x.size(1)
    else:
        raise RuntimeError('Input tensor must have at least 1 dimension.')

    # Split and shuffle weight matrix
    # for i in tqdm(range(0, num_pe), total=num_pe, unit="Submat"):
    for i in range(0, num_pe):
        if n_dims == 1:
            pe_work = x[np.arange(i, n_rows, num_pe)]
            pe_work_abs = torch.abs(pe_work)
            drop_part = math.floor(pe_work_abs.shape[0] * gamma)
            _, indices = torch.sort(pe_work_abs, dim=0)
            drop_indices = indices[0:drop_part]
            drop_rand_mask = torch.rand(drop_indices.size(0)).cuda()
            pe_work[drop_indices] = pe_work[drop_indices].masked_fill_(drop_rand_mask <= alpha, 0)
            x[np.arange(i, n_rows, num_pe)] = pe_work
        elif n_dims > 1:
            pe_work = x[np.arange(i, n_rows, num_pe), :]
            pe_work_abs = torch.abs(pe_work)
            drop_part = math.floor(pe_work_abs.shape[0] * gamma)
            _, indices = torch.sort(pe_work_abs, dim=0)
            drop_indices = indices[0:drop_part, :]
            drop_rand_mask = torch.rand(drop_indices.size(0), drop_indices.size(1)).cuda()
            for j in range(0, n_cols):
                pe_work[drop_indices[:, j], j] = pe_work[drop_indices[:, j], j].masked_fill_(
                    drop_rand_mask[:, j] <= alpha, 0)
            x[np.arange(i, n_rows, num_pe), :] = pe_work
    return x

# def cbtd(x, gamma, alpha):
#     """
#     x: input tensor
#     gamma: target sparsity
#     alpha: drop probability
#     mode: 0 - normal; 1 -
#     """
#     import math
#     device = x.device
#     n_dims = len(x.size())
#     n_rows = x.size(0)
#
#     if n_dims == 1:
#         n_cols = 1
#     elif n_dims == 2:
#         n_cols = x.size(1)
#     else:
#         raise RuntimeError('Input tensor must have 1 or 2 dimensions.')
#
#     x_abs = torch.abs(x)
#     drop_part = math.floor(n_rows * gamma)
#     _, indices = torch.sort(x_abs, dim=0)
#     if n_dims == 1:
#         drop_indices = indices[0:drop_part]
#         drop_rand_mask = torch.rand(drop_indices.size(0)).cuda()
#         x[drop_indices] = x[drop_indices].masked_fill_(drop_rand_mask <= alpha, 0)
#     elif n_dims > 1:
#         drop_indices = indices[0:drop_part, :]
#         drop_rand_mask = torch.rand(drop_indices.size(0), drop_indices.size(1), device=device)
#         for j in range(0, n_cols):
#             x[drop_indices[:, j], j] = x[drop_indices[:, j], j].masked_fill_(
#                 drop_rand_mask[:, j] <= alpha, 0)
#     return x

def castToFixed(x, m, n):
    """
    :param x: input numpy array
    :param m: number of integer bits before the decimal point
    :param n: number of fraction bits after the decimal point
    :return: tensor quantized to fixed-point precision
    """
    power = 2. ** n
    max_val = 2. ** (m + n) - 1
    value = x * power
    value = np.round(value)  # rounding
    value = np.clip(value, -max_val - 1, max_val)  # saturation arithmetic
    return value

def interleave(x, width):
    """
    PyTorch Function to interleave a
    :param x:
    :param width:
    :return:
    """

    x_dim0 = x.size(0)
    num_dim = len(x.size())
    interleaved = []
    for i in range(0, width):
        idx_curr = torch.arange(i, x_dim0, width).long()
        if num_dim == 1:
            interleaved.append(x[idx_curr])
        elif num_dim == 2:
            interleaved.append(x[idx_curr, :])
    interleaved = torch.stack(interleaved)
    return interleaved

def get_array_work(workload, num_array_pe, num_array):
    n_col = workload.size(0)
    hw_n_col = int(num_array_pe * np.ceil(float(n_col) / float(num_array_pe)))
    num_pad_col = int(hw_n_col - n_col)
    workload_pad_col = torch.zeros(num_pad_col)
    act_workload = torch.cat((workload, workload_pad_col), dim=0)

    array_col_work = torch.zeros(int(num_array), int(hw_n_col / num_array))
    array_work = torch.zeros(int(num_array))
    for i in range(0, int(num_array)):
        array_col_work[i] = act_workload[np.arange(i, hw_n_col, num_array)]
        array_work[i] = torch.sum(array_col_work[i])

    return array_work, array_col_work

def cast_fp_to_fxp(x, m, n):
    """
    Converts a floating point number into a fixed point number

    :param x: input numpy array
    :param m: number of integer bits before the decimal point
    :param n: number of fraction bits after the decimal point
    :return: tensor quantized to fixed-point precision
    """
    power = 2. ** n
    max_val = 2. ** (m + n)
    value = x * power
    value = np.round(value)  # rounding
    value = np.clip(value, -max_val, max_val - 1).astype(np.int32)  # saturation arithmetic
    return value

def cast_fxp_to_hex(x, data_type):
    x = np.asarray(x, dtype=data_type)
    return x

# def cast_signed_to_unsigned(x, bw):
#     temp = x
#     idx_below_zero = np.argwhere(temp<0)
#     abs_below_zero = np.abs(temp[idx_below_zero])
#     max_val = 2**bw
#     unsigned_temp = max_val - abs_below_zero
#     temp[idx_below_zero] = unsigned_temp
#     return temp

def cast_signed_to_unsigned(x, bw):
    temp = x
    if temp.ndim == 1:
        temp = np.expand_dims(temp, axis=-1)
    # print("Temp Size: ", temp.shape)
    for i in range(0, temp.shape[-1]):
        # print("0")
        idx_below_zero = np.argwhere(temp[:, i] < 0)
        # print("idx_below_zero", idx_below_zero.shape)
        abs_below_zero = np.abs(temp[idx_below_zero, i])
        # print("abs_below_zero", abs_below_zero.shape)
        max_val = 2 ** bw
        unsigned_temp = max_val - abs_below_zero
        temp[idx_below_zero, i] = unsigned_temp
    return temp

def gen_sim_file(file_path, pdict):
    """
    Generate Modelsim Simulation Files for NN parameters

    :param pdict:    Dictionary from which to save numpy arrays of NN parameters in fp
    :param qi:       number of integer bits before the decimal point in fxp
    :param qf:       number of integer bits after the decimal point in fxp
    :return:         mat_param_fxp - parameter in fixed-point numbers
    """

    # Iterate over params in dict
    pdict_out = {}
    for key, value in pdict.items():

        # Get Inputs
        param = value[0]
        qi = value[1]
        qf = value[2]
        idx = value[3]
        idx_bw = value[4]

        if param.ndim == 1:
            param = np.expand_dims(param, axis=-1)
        if idx is not False:
            if idx.ndim == 1:
                idx = np.expand_dims(idx, axis=-1)

        # Get unsigned variable type
        if idx is not False:
            bit_width = qi + qf + idx_bw
        else:
            bit_width = qi + qf

        if 0 < bit_width <= 8:
            var_type = 'uint8'
        elif 8 < bit_width <= 16:
            var_type = 'uint16'
        elif 16 < bit_width <= 32:
            var_type = 'uint32'
        else:
            raise ValueError('Bit width of parameter must not be larger than 32...')

        # Convert fp to fxp
        if idx is not False:
            # param = quantizeTensor(param, qi, qf, 1) # Remove if necessary
            # fxp_param = param
            fxp_param = cast_fp_to_fxp(param, qi, qf)
            fxp_param_unsigned = cast_signed_to_unsigned(fxp_param, qi + qf)
            fxp_param_shift = fxp_param_unsigned * (2 ** idx_bw)
            param = fxp_param_shift + idx
        else:
            param = cast_fp_to_fxp(param, qi, qf)

        # Get param dimension
        dim0 = param.shape[0]
        dim1 = param.shape[1]

        # Create file
        try:
            f = open(file_path + '/' + key + '.txt', 'w')
        except IOError as err:
            print("I/O error({0}): {1}".format(err.errno, err.strerror))

        # Write parameters column-wisely
        param_hex = cast_fxp_to_hex(param, var_type)
        for i in range(0, dim1):
            for j in range(0, dim0):
                if var_type == 'uint8':
                    f.write('%2.2X\n' % param_hex[j][i])
                elif var_type == 'uint16':
                    f.write('%4.4X\n' % param_hex[j][i])
                elif var_type == 'uint32':
                    f.write('%8.8X\n' % param_hex[j][i])
                else:
                    f.write('%X\n' % param_hex[j][i])
        f.close()

def gen_clib(file_path, file_name, pdict, num_layer=None, input_size=None, hid_size=None):
    """
    Generate C Library for NN parameters

    :param filename: name of the C Library
    :param pdict:    Dictionary from which to save numpy arrays of NN parameters in fp
    :param qdict:    Dictionary from which to save properties of corresponding variables in pdict in tuples
        Elements of each tuple in qdict in order:
        - quantize
        - qi:        number of integer bits before the decimal point in fxp
        - qf:        number of integer bits after the decimal point in fxp
        - var_class: 0 - Constant | 1 - General Tensors | 2 - NN Parameter |
        - print_row_size
    :return:         mat_param_fxp - parameter in fixed-point numbers
    """
    print_row_size = 20
    # Create file
    try:
        f_h = open(file_path + '/' + file_name + '.h', 'w')
    except IOError as err:
        print("I/O error({0}): {1}".format(err.errno, err.strerror))
    try:
        f = open(file_path + '/' + file_name + '.c', 'w')
    except IOError as err:
        print("I/O error({0}): {1}".format(err.errno, err.strerror))

    # Write macro definitions
    macro_name = file_name.upper() + '_H'
    f_h.write('#ifndef ' + macro_name + '\n' + '#define ' + macro_name + '\n')
    f.write('#include "' + file_name + '.h"\n')

    # Iterate over params in dict
    for key, value in pdict.items():
        # Get Inputs
        param = np.asarray(value[0])
        qi = value[1]
        qf = value[2]
        idx = value[3]
        idx_bw = value[4]
        param_type = value[5]

        # Get unsigned variable type
        if idx is not False:
            bit_width = qi + qf + idx_bw
        else:
            bit_width = qi + qf

        if 0 < bit_width <= 8:
            var_type = 'char'
            nn_var_type = 'unsigned char'
            np_var_type = 'uint8'
        elif 8 < bit_width <= 16:
            var_type = 'short'
            nn_var_type = 'unsigned short'
            np_var_type = 'uint16'
        elif 16 < bit_width <= 32:
            var_type = 'long'
            nn_var_type = 'unsigned long'
            np_var_type = 'uint32'
        else:
            raise ValueError('Bit width of parameter must not be larger than 32...')

        # Convert fp to fxp
        if param_type != 0:
            if idx is not False:
                # param = quantizeTensor(param, qi, qf, 1) # Remove if necessary
                fxp_param = cast_fp_to_fxp(param, qi, qf)
                fxp_param_unsigned = cast_signed_to_unsigned(fxp_param, qi + qf)
                fxp_param_shift = fxp_param_unsigned * (2 ** idx_bw)
                param = fxp_param_shift + idx
                param = cast_fxp_to_hex(param, np_var_type)
                param = param
            else:
                param = cast_fp_to_fxp(param, qi, qf)

        if param_type == 0:  # Constant
            f_h.write('#define ' + key.upper() + ' %d\n' % param)
        elif param_type == 1:  # General Tensor
            if param.ndim == 1:
                param = np.expand_dims(param, axis=0)

            # Get param dimension
            size = param.size
            dim0 = param.shape[0]
            dim1 = param.shape[1]

            # Write variable declaration
            f_h.write('/*\n'
                      ' * Var Type: General Matrix \n'
                      ' * Var Name: %s' % key + '\n'
                                                ' * Bit Width: %d' % bit_width + '\n'
                                                                                ' * Dimension: (%d, %d)' % (
                          dim0, dim1) + '\n'
                                        ' */\n'
                      )
            f_h.write('#define ' + key.upper() + '_NUM_ROWS %d\n' % dim0)
            f_h.write('#define ' + key.upper() + '_NUM_COLS %d\n' % dim1)
            f_h.write('#define ' + key.upper() + '_MAT_SIZE %d\n' % size)
            f_h.write('extern const %s %s[' % (var_type, key) + key.upper() + '_MAT_SIZE];\n')
            f.write('const %s %s[' % (var_type, key) + key.upper() + '_MAT_SIZE] = {\n')

            # Write parameters column-wisely
            for i in range(0, dim1):
                for j in range(0, dim0):
                    if var_type == 'float':
                        f.write('%.20f' % param[j][i])
                    else:
                        f.write('%d' % param[j][i])
                    if not (i >= dim1 - 1 and j >= dim0 - 1):
                        f.write(',')
                    if (i * dim0 + j) % print_row_size == print_row_size - 1:
                        f.write('\n')
            f.write('};\n\n')

        elif param_type == 2:  # NN Parameter
            if param.ndim == 1:
                param = np.expand_dims(param, axis=0)

            # Get param dimension
            size = param.size
            dim0 = param.shape[0]
            dim1 = param.shape[1]

            # Write variable declaration
            f_h.write('/*\n'
                      ' * Var Type: NN Parameter Matrix \n'
                      ' * Var Name: %s' % key + '\n'
                                                ' * Bit Width: %d' % bit_width + '\n'
                                                                                ' * Dimension: (%d, %d)' % (
                          dim0, dim1) + '\n'
                                        ' */\n'
                      )

            f_h.write('#define ' + key.upper() + '_NUM_ROWS %d\n' % dim0)
            f_h.write('#define ' + key.upper() + '_NUM_COLS %d\n' % dim1)
            f_h.write('#define ' + key.upper() + '_MAT_SIZE %d\n' % size)
            f_h.write('extern const %s %s[' % (
                nn_var_type, key) + key.upper() + '_MAT_SIZE] __attribute__ ((aligned (%d)));\n' % bit_width)
            f.write('const %s %s[' % (nn_var_type, key) + key.upper() + '_MAT_SIZE] = {\n')
            if num_layer is not None:
                f_h.write('#define ' + key.upper() + '_NUM_LAYERS %d\n' % num_layer)
            if input_size is not None:
                f_h.write('#define ' + key.upper() + '_INP_SIZE %d\n' % input_size)
            if hid_size is not None:
                f_h.write('#define ' + key.upper() + '_HID_SIZE %d\n' % hid_size)
            # if 'rnn' in key:
            #     f_h.write('extern const %s %s[RNN_PARAM_Hweight_tensor_SIZE] __attribute__ ((aligned (%d)));\n' % (var_type, key, bit_width))
            #     f.write('const %s %s[RNN_PARAM_Hweight_tensor_SIZE] = {\n' % (var_type, key))
            # else:
            #     f_h.write('extern const %s %s[%d];\n' % (var_type, key, size))
            #     f.write('const %s %s[%d] = {\n' % (var_type, key, size))

            # Write parameters column-wisely
            for i in range(0, dim1):
                for j in range(0, dim0):
                    f.write('%d' % param[j][i])
                    if not (i >= dim1 - 1 and j >= dim0 - 1):
                        f.write(',')
                    if (i + j) % print_row_size == print_row_size - 1:
                        f.write('\n')
            f.write('};\n\n')

    f_h.write('#endif')
    f.close()
    f_h.close()