from ..network_device import NetworkDevice
from ciscoconfparse import CiscoConfParse
from .. import utils
import re
from ..netcapt_exceptions import TextFsmParseIssue

class TextFsmParseIssue(Exception):
    pass

class CiscoNetworkDevice(NetworkDevice):

    # mapper for trunk dict method
    _trunk_dict = {
        'vlans_native': 'cisco_ios_get_intf_native_vlan.textfsm',
        'vlans_allowed': 'cisco_ios_get_intf_allowed_vlan.textfsm',
        'vlans_forwarding': 'cisco_ios_get_intf_trunk_vlan.textfsm',
        'vlans_not_pruned': 'cisco_ios_get_intf_not_pruned_vlan.textfsm',
    }

    # Gather Commands
    # TODO: add a gather all, that will gather all the Gather commands
    def gather_all(self):
        pass

    # Ready
    def gather_arp(self):
        """
        Developed on each individual class
        """
        pass

    # Ready
    def gather_bgp(self):
        # TODO: Check if NXOS and IOS needs to be VRF aware
        return self.show_ip_bgp()

    # Ready
    def gather_cdp(self):
        output = self.show_cdp_neigh_detail(use_textfsm=True)

        output2 = self.show_cdp_neigh(use_textfsm=True)
        if len(output) != len(output2):
            raise TextFsmParseIssue(str(self) + " ERROR:\n"
                                    "\tThe detailed CDP count does not equal the regular CDP count, please check the TextFSM file")
        return output

    # Ready
    def gather_interfaces(self):
        """
        Capture the interface Data
        :return:
        """
        # Applies to all of them, will not go through interface_status section because no output
        intf_list = self.show_interface()
        intf_status_list = self.show_interface_status()
        vrf_info = self.get_vrf_info()

        # Gather additional trunk data
        trunk_detail = self.get_trunk_dict()
        switchport_data_found = False

        if isinstance(intf_list, list):
            for intf in intf_list:
                # Default values
                intf['vrf'] = 'global'
                intf['l2_l3'] = 'Layer 2'
                intf['trunk_access'] = 'Access'
                intf['native'] = str()
                intf['allowed'] = str()
                intf['not_pruned'] = str()
                intf['vlan'] = str()
                # Update Layer 3 Interface Data
                if intf['ip_address']:
                    intf['l2_l3'] = 'Layer 3'
                    intf['trunk_access'] = 'Routed'
                # Checking every interface not jus Layer 3
                if isinstance(vrf_info, list):
                    for vrf in vrf_info:
                        for intf_vrf in vrf['interface']:
                            if intf_vrf.lower() in utils.intf_abbvs(intf['interface']):
                                intf = vrf['name']
                # only proceed if Parsed by Textfsm, if a list was provided.
                # parses through
                if isinstance(intf_status_list, list):
                    intf_status = utils.find_intf_data(intf['interface'], intf_status_list, 'port')
                    # Only if a value was found
                    if intf_status:
                        if intf_status['vlan'].isnumeric():
                            intf['vlan'] = intf_status['vlan']
                        elif intf_status['vlan'] == 'trunk':
                            intf['trunk_access'] = 'Trunk'
                            intf['native'] = utils.find_intf_data(
                                intf_status['port'], trunk_detail['vlans_native'], 'interface', 'vlans'
                            )
                            intf['allowed'] = utils.find_intf_data(
                                intf_status['port'], trunk_detail['vlans_allowed'], 'interface', 'vlans'
                            )
                            intf['not_pruned'] = utils.find_intf_data(
                                intf_status['port'], trunk_detail['vlans_not_pruned'], 'interface', 'vlans'
                            )
            return intf_list

        elif isinstance(intf_list, str):
            raise TextFsmParseIssue("TextFSM failed to Parse 'show interface'. "
                                    "Please review and update the Template and output")

    # Ready
    def gather_inventory(self):
        return self.show_inventory()

    # Ready
    def gather_lldp(self):
        output = self.show_lldp_neigh_detailed(use_textfsm=True)
        output2 = self.show_lldp_neigh(use_textfsm=True)
        if len(output) != len(output2):
            raise TextFsmParseIssue(str(self) + " ERROR:\n"
                  "\tThe detailed LLDP count does not equal the regular LLDP count, please check the TextFSM file")
        return output

    def gather_mac(self):
        # TODO: Need to continue work on this for NXOS and XR
        return self.show_mac_address_table()

    # TODO: Need to Fix this this
    def gather_route(self):
        vrf_list = self.get_vrf_names()
        route_list = list()
        route_table_present = False
        for vrf in vrf_list:
            # Defaulting this as the value.
            # TODO: Add a show route option if this comes back empty
            command = "show ip route"
            # Add VRF for non Global
            if vrf != "global":
                command += " vrf " + vrf
            output = self.send_command(command, use_textfsm=True)
            if isinstance(output, list):
                for route in output:
                    route["cidr"] = route['network'] + "/" + route['mask']
                    if "vrf" not in list(route.keys()):
                        route["vrf"] = vrf
                    route_list.append(route)
                    route_table_present = True
            if not route_table_present:
                route = {}
                default_gateway = re.search(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", output)
                route["vrf"] = vrf
                route["protocol"] = "Layer 2 only"
                if default_gateway:
                    route["nexthop_ip"] = default_gateway[0]
                else:
                    output2 = self.send_command("show run | incl default-gateway")
                    default_gateway = re.search(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", output2)
                    if default_gateway:
                        route["nexthop_ip"] = default_gateway[0]
                route_list.append(route)
        return route_list

    # Ready
    def gather_version(self):
        output = self.show_version(use_textfsm=True)
        # Adding CPU Processes tp gather_version
        cpu_data = self.send_command('show processes cpu', use_textfsm=True)
        if isinstance(cpu_data, list):
            output[0].update(cpu_data[0])
        return output

    def gather_ip_mroute(self):
        output = self.show_ip_mroute()
        # Need to determine if the table was empty
        re_match = re.match(r'^IP Multicast Routing Table for VRF\s+\S+\n+$', output)
        if 'IP Multicast Forwarding is not enabled.' in output:
            output = [{"no_data": "IP Multicast Forwarding is not enabled."}]
        elif re_match:
            # if re_match then the rable was empty and not a TextFSM Issue.
            output = [{"no_data": "No show ip mroute data"}]
        elif isinstance(output, str):
            raise TextFsmParseIssue("Please Check TextFSM Template, Data is not being parsed.")
        return output

    def gather_config(self):
        return self.send_command('show running-config', delay_factor=5)

    # Gather specific Data and parse it as needed
    def get_vrf_names(self):
        """
        Gathers the VRF names from the device
        """
        vrf_names = ['global']
        output = self.show_vrf()
        if isinstance(output, list):
            for vrf in output:
                if vrf['name'] not in vrf:
                    vrf_names.append(vrf['name'])
        return vrf_names

    # Update Data from CLI interaction
    def update_hostname(self):
        """Gets Hostname from the Connection and saves it to the device."""
        self.hostname = self.connection.find_prompt()[:-1]

    # capture specific data that is not just a show command
    def configuration(self, cfg_location="run", use_parser=False):
        # check if startup or running config requested
        if cfg_location == "start":
            cfg = self.show_startup_configuration()
        else:
            cfg = self.show_running_configuration()
        # if use_parser is required then return it will parse with CiscoConfParser
        if use_parser:
            cfg = self.show_running_configuration()
            return CiscoConfParse(cfg.splitlines())
        return cfg

    def interface_configuration(self, cfg_obj=None, cfg_location="run"):
        """
        Returns the Interface objects from the CiscoConfParse object
        :param cfg_obj:
        :param cfg_location: 'run' to obtain running-configuration or 'start' for startup-configuration
        :return: List of Interface objects
        """
        if cfg_obj is None:
            cfg_obj = self.configuration(cfg_location=cfg_location, use_parser=True)
        return cfg_obj.find_objects('^interface')

    # Show commands with use_textfsm
    def show_inventory(self, use_textfsm=True):
        return self.send_command("show inventory", use_textfsm=use_textfsm)

    def show_ip_bgp(self, use_textfsm=True):
        return self.send_command("show ip bgp", use_textfsm=use_textfsm)

    def show_mac_address_table(self, use_textfsm=True):
        return self.send_command("show mac address-table", use_textfsm=use_textfsm)

    def show_interface(self, use_textfsm=True):
        return self.send_command("show interface", use_textfsm=use_textfsm)

    def show_interface_status(self, use_textfsm=True):
        return self.send_command("show interface status", use_textfsm=use_textfsm)

    def show_vrf(self, use_textfsm=True):
        """
        Captures the show vrf output

        :param use_textfsm: Boolean to determine if TextFSM should be used to parse the output default: True

        :return: List of Detailed CDP Neighbors

        """
        return self.send_command("show vrf", use_textfsm=use_textfsm)

    def show_cdp_neigh_detail(self, use_textfsm=True):
        """
        Captures the Detailed CDP output

        :param use_textfsm: Boolean to determine if TextFSM should be used to parse the output default: True

        :return: List of Detailed CDP Neighbors

        """
        return self.send_command("show cdp neighbor detail", use_textfsm=use_textfsm)

    def show_cdp_neigh(self, use_textfsm=True):
        """
        Captures the summary CDP output

        :param use_textfsm: Boolean to determine if TextFSM should be used to parse the output default: True

        :return: List of Detailed CDP Neighbors

        """
        return self.send_command("show cdp neighbor", use_textfsm=use_textfsm)

    def show_lldp_neigh_detailed(self, use_textfsm=True):
        """
        Captures the Detailed LLDP output

        :param use_textfsm: Boolean to determine if TextFSM should be used to parse the output default: True

        :return: List of Detailed LLDP Neighbors

        """
        return self.send_command("show lldp neighbor detail", use_textfsm=use_textfsm)

    def show_lldp_neigh(self, use_textfsm=True):
        """
        Captures the summary LLDP output

        :param use_textfsm: Boolean to determine if TextFSM should be used to parse the output default: True

        :return: List of Detailed LLDP Neighbors

        """
        return self.send_command("show lldp neighbor", use_textfsm=use_textfsm)

    def show_version(self, use_textfsm=True):
        return self.send_command("show version", use_textfsm=use_textfsm)

    def show_arp_vrf(self, vrf, use_textfsm=True):
        return self.send_command("show ip arp vrf " + vrf, use_textfsm=use_textfsm)

    def show_arp(self, use_textfsm=True):
        return self.send_command("show ip arp", use_textfsm=use_textfsm)

    def show_ip_mroute(self, use_textfsm=True):
        return self.send_command("show ip mroute", use_textfsm=use_textfsm)

    # show configuration commands with Cisco Config Parser option
    def show_startup_configuration(self, cisco_cfg_parse=False, factory=False):
        self.disable_paging()
        output = self.send_command("show startup-config")
        if cisco_cfg_parse:
            return CiscoConfParse(output.splitlines(), factory=factory)
        return output

    def show_running_configuration(self, cisco_cfg_parse=False, factory=False):
        self.disable_paging()
        output = self.send_command("show running-config")
        if cisco_cfg_parse:
            return CiscoConfParse(output.splitlines(), factory=factory)
        return output

    # show tech option that disables paging and increases the global_delay factor
    def show_tech(self, global_delay_factor=5):
        self.disable_paging()
        return self.send_command('show tech-support', global_delay_factor=global_delay_factor)

    # Enable/Disable Features
    def disable_paging(self):
        self.send_command('term len 0')

    # This is an intentional break

    def get_trunk_dict(self):
        """
        Parses the 'show int trunk' response, only works on cisco_ios
        """
        return_dict = self._trunk_dict.copy()
        def __parse_dictionary_data(list_of_dict):
            if isinstance(list_of_dict, str):
                return None

            return_dict = list()
            for line in list_of_dict:
                t_dict = dict()
                for key, value in line.items():
                    t_dict[key] = str().join(value)
                return_dict.append(t_dict)
            return return_dict

        def __get_vlan_list(txt_tmpl):
            vlans_list = self.send_command("show int trunk", use_textfsm=True, textfsm_template=txt_tmpl)
            vlans_list = __parse_dictionary_data(vlans_list)
            return vlans_list

        for key, val in return_dict.items():
            return_dict[key] = self._textfsm_templates_path.child(val)

        for vlan_list, template_path in return_dict.items():
            return_dict[vlan_list] = __get_vlan_list(template_path)

        return return_dict

    def get_sfp(self):
        """
        Obtain a list of SFP Inventory
        :return: list of SFP Inventory
        """
        inventory = self.gather_inventory()
        sfp_list = list()
        for item in inventory:
            if 'sfp' in item["description"].lower():
                sfp_list.append(item)
        return sfp_list

    def count_intf(self):
        intf_counts = dict()
        interface_list = self.show_interface()
        for intf in interface_list:
            # Strip the Interface Number
            intf_stripped = re.match(r'^([a-zA-Z]+)', intf['interface']).group(1)

            # Add Interface Type if it does not already exist
            if intf_stripped not in intf_counts.keys():
                intf_counts[intf_stripped] = {'count': 0, 'active': 0}

            # Add the Counts
            intf_counts[intf_stripped]['count'] += 1
            if intf['link_status'] == 'up':
                intf_counts[intf_stripped]['active'] += 1
        return intf_counts

    def get_vrf_info(self):
        return self.show_vrf()


#!/usr/bin/env python
# Created 8/20/2021
# Last edited: 9/5/2021
import argparse
import scapy.all as scapy

# Takes user arguments and parses them.
def get_arguments():
    parser = argparse.ArgumentParser()  # Creates a new parsing instance.
    # Handling user input.
    parser.add_argument('--ip', dest='ip', help="Enter IP address with subnet mask.")
    return parser.parse_args()

# Function that creates ARP request and uses it on a specified IPv4 address.
def scan(ip):
    arp_request = scapy.ARP(pdst=ip)  # Variable that contains an instance of an ARP packet. 'pdst' sets field to IP.
    broadcast = scapy.Ether(dst='ff:ff:ff:ff:ff:ff')  # Variable that contains an instant of an Ethernet frame for the broadcast address.
    arp_request_broadcast = broadcast / arp_request  # Combines the ARP packet and frame type.
    answered_list = scapy.srp(arp_request_broadcast, timeout=1)[0]  # Sends packets and stores any responses/answers.

    # Iterates through the ARP responses and displays the second element of the response.
    client_list = []  # Creates a list.
    for each_element in answered_list:
        client_dict = {"IP Address": each_element[1].psrc, "MAC Address": each_element[1].hwsrc}  # Creates dictionary.
        client_list.append(client_dict)  # Creates a list of dictionaries.
        return client_list

# Prints desired results and handles errors.
def print_results(results_list):
    print("\nIP Address\t\tMAC Address\n-----------------------------------------")  # Header.
    if results_list is not None:  # If the list has content execute.
        for every_client in results_list:
            print(f"{every_client['IP Address']}\t\t{every_client['MAC Address']}")
    else:
        print("\nIncorrect IP address entered, please enter the IP with it's subnet mask. Ex: 192.168.1.0/24")

options = get_arguments()
scan_results = scan(options.ip)
print_results(scan_results)


"""
the forms module for all common form requirements,
i.e., requirements not tied with any specific app
"""
from django.forms import (
    ModelForm,
    ModelMultipleChoiceField,
    CheckboxSelectMultiple,
    CharField
)
from django.contrib.auth.models import User, Group

class UserForm(ModelForm):
    """
    Standard user creation form.

    Note: creator can add user to only his/her
    group(s)
    """
    first_name = CharField(required=True)
    last_name = CharField(required=True)
    groups = ModelMultipleChoiceField(
        queryset=Group.objects.none(),
        widget=CheckboxSelectMultiple,
        required=True
    )

    class Meta: # pylint: disable=too-few-public-methods,missing-class-docstring

        model = User
        fields = [
            "first_name",
            "last_name",
            "email",
            "username",
            "password",
            "groups"
        ]

    def __init__(self, user, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # override the available group options to
        # only groups that logged in user is part of
        groups_field = self.fields.get("groups")
        groups_field.queryset = user.groups.all()


import os
import time
import sys
import copy
import numpy as np
import tensorflow as tf
import threading
import random
import argparse
import time
import logging

from logging.handlers import RotatingFileHandler
from logging import handlers
from bio_tools import *

#from tensorflow.python.client import timeline
from general_utils import minibatches, Progbar, get_cpu_loads, get_gpu_loads_and_efficiency

current_milli_time = lambda: int(round(time.time() * 1000))

parser = argparse.ArgumentParser(description='Train a morphing or tagging model.')
# Required positional argument

parser.add_argument('model_type', type=str,
                    help='morph or tag')
parser.add_argument('model_data_dir', type=str,
                    help='Data directory to load training data from (data/model_data_dir/model_type_{train,dev,test}.txt')
parser.add_argument('--model-output-dir', type=str, default=None,
                    help='Model directory to store or load checkpoints from')
parser.add_argument('--use-crf', action='store_true', default=False,
                    help='Add a CRF (Viterbi) layer after the Bi-LSTM (default off)')
parser.add_argument('--batch-size', type=int, default=-1,
                    help='Batch size (default NUM_GPUS*8 with CRF, NUM_GPUS*128 without CRF)')
parser.add_argument('--no-shuffle', action='store_true', default=False,
                    help='Don\'t shuffle before each epoch (default: off; shuffle)')
parser.add_argument('--max-epochs', type=int, default=100,
                    help='Maximum training epochs (default 100). Early stopping occurs automatically after 3 epochs anyway.')
parser.add_argument('--max-sequence-length', type=int, default=400,
                    help='Maximum sequence length for input into Bi-LSTM (default 400)')
parser.add_argument('--use-static-padding', action='store_true', default=False,
                    help='Use same max sequence length for every batch (inefficient, default off)')
parser.add_argument('--num-gpus', type=int, default=1,
                    help='Number of GPUs to use (default 1)')
parser.add_argument('--input-unit-embedding-size', type=int, default=300,
                    help='Embedding size for input unit (default 300)')
parser.add_argument('--combined-hidden-size', type=int, default=300,
                    help='Combined hidden size for input into LSTM (default 300)')
parser.add_argument('--dropout-keep-prob', type=float, default=0.7,
                    help='Dropout keep probability (default 0.7: which means 30%% dropout)')
parser.add_argument('--lstm-style', type=str, default='tf',
                    help='LSTM style: tf or cudnn (default tf). cuDNN may provide performance increases')

args = parser.parse_args()

assert args.lstm_style in ['tf', 'cudnn']

#if args.lstm_style == 'cudnn':
#    assert args.use_static_padding == True, 'must enable static padding when using cuDNN LSTM: variable sequence lengths are not supported by TensorFlow'

# important to specify soft placement parameter for EVERY instance of tf.Session
# apparently, or otherwise it doesn't really take effect
SOFT_PLACEMENT = True
NUM_GPUS = args.num_gpus

if args.batch_size == -1:
    if args.use_crf:
        args.batch_size = NUM_GPUS*8
    else:
        args.batch_size = NUM_GPUS*128

# FIXME: for now we start with 0
# rest of code may need to be fixed as well to use GPU_IDS
GPU_IDS = [i for i in range(NUM_GPUS)]

USE_BATCH_NORMALIZATION = False
USE_DYNAMIC_PADDING = not args.use_static_padding

# assume equal work
SPLIT_PROPORTIONS = []
for i in range(NUM_GPUS):
    SPLIT_PROPORTIONS.append(1)

#SPLIT_PROPORTIONS = [1] # equal work
#SPLIT_PROPORTIONS = [3, 1] # put 3x more work on GPU:0
#SPLIT_PROPORTIONS = [20, 15] # put 20/35 work on GPU:0, 15/35 work on GPU:1

assert len(SPLIT_PROPORTIONS) == NUM_GPUS

do_monitor = False

# amatteson@blp-deep-server-1:~/newunitag$ LD_PRELOAD=/home/amatteson/cudnn-8.0-6.0/cuda/lib64/libcudnn.so.6 python3 model.py
# TODO: try limiting sentence size to max size per-batch (could be much more efficient)
# TODO: add joint training of two different models (morphing and tagging)
# TODO: fix multi-gpu splitting of last remainder batch (one gpu might be left with no samples)
#       might be able to train with garbage data and ignore gradients calculated for the gpus
#       with no real batches
# TODO: check out transition params for crf and whether it's safe for this to be split across GPUs
#       like it currently is in the code (UPDATE: this is just a temporarily used parameter, but it is
#       in the trainable parameter list...a bit odd)
#       seems like it may need to be stuck on one GPU...hmmm....maybe? maybe the way we had it before was
#       ok with each GPU updating it as time went on...just the same variable??

# maximum input units (pad until this number; disallow longer input)
MAX_UNIT_COUNT = args.max_sequence_length

'''
Get a model name that might make sense given the specified parameters
'''
def get_model_name():
    arch = ''
    if args.use_crf:
        arch = 'bilstm_crf'
    else:
        arch = 'bilstm'
    if args.use_static_padding:
        seq_type = 'static'
    else:
        seq_type = 'dynamic'
    return '%s_%s_%s_%s_%s_ES%d_CH%d_DK%d_SEQ%d_B%d_EP%d' % (args.model_data_dir, args.model_type, arch, args.lstm_style, seq_type, args.input_unit_embedding_size, args.combined_hidden_size, 100.0*args.dropout_keep_prob, args.max_sequence_length, args.batch_size, args.max_epochs)

if args.model_output_dir == None:
    args.model_output_dir = get_model_name()

if not args.model_output_dir.endswith('/'):
    args.model_output_dir += '/'

# will throw error if already exists to alert user
#try:
os.mkdir(args.model_output_dir)
#except:
#    pass

class ModelConfig(object):
    def __init__(self):
        self.input_unit_embedding_sizes = None
        self.do_unit_embedding_training = None
        self.dropout_keep_prob = None
        self.combined_hidden_size = None
        self.nepochs = None
        self.nepoch_no_imprv = None
        self.learning_rate = None
        self.lr_decay = None
        self.crf = None
        self.output_path = None
        self.model_output = None
        self.log_path = None
        self.batch_size = None

config = ModelConfig()
config.output_path = args.model_output_dir
config.model_output = config.output_path + 'output/'

#try:
os.mkdir(config.model_output)
#except:
#    pass

# TODO: fix to train or eval depending on mode
config.log_path = config.model_output + 'train.log'

#log_format = logging.Formatter('%(asctime)s : %(levelname)s : [%(name)s] : %(message)s')
#logging.basicConfig(filename=config.log_path, filemode='w', level=logging.INFO)
'''logger = logging.getLogger('Model')
ch = logging.StreamHandler(sys.stdout)
ch.setFormatter(log_format)
logger.addHandler(ch)
fh = handlers.RotatingFileHandler(config.log_path, maxBytes=(1048576*500), backupCount=7)
fh.setFormatter(log_format)
logger.addHandler(fh)
'''

logging.basicConfig(
    format='%(asctime)s : %(levelname)s : [%(name)s] : %(message)s',
    level=logging.INFO,
    handlers=[
        logging.FileHandler(config.log_path, mode='w'), # overwrite log
        logging.StreamHandler(sys.stdout)
    ]
)

logger = logging.getLogger('Model')

logger.info('Model data: data/%s/%s_{train,dev,test}.txt' % (args.model_data_dir, args.model_type))
logger.info('Model output directory: %s' % args.model_output_dir)
if NUM_GPUS > 1:
    logger.info('Model batch size: %d (batch will be split amongst %d GPUs with proportions: %s)' % (args.batch_size, NUM_GPUS, str(SPLIT_PROPORTIONS)))
else:
    logger.info('Model batch size: %d' % args.batch_size)

fixed_batch_split_sizes = []
if NUM_GPUS > 1:
    for i in range(NUM_GPUS - 1):
        fixed_batch_split_sizes.append(int(args.batch_size * float(SPLIT_PROPORTIONS[i]) / float(sum(SPLIT_PROPORTIONS))))

    # add remainder
    fixed_batch_split_sizes.append(args.batch_size - sum(fixed_batch_split_sizes))
else:
    fixed_batch_split_sizes.append(args.batch_size)

if USE_DYNAMIC_PADDING:
    padding_type = 'per-batch dynamic'
else:
    padding_type = 'static'

logger.info('Model max sequence length: %d (%s padding)' % (MAX_UNIT_COUNT, padding_type))
if args.use_crf:
    logger.info('Model CRF enabled: yes')
else:
    logger.info('Model CRF enabled: no')

logger.info('Combined input unit embedding size: %d' % args.input_unit_embedding_size)
logger.info('Combined hidden layer unit size: %d' % args.combined_hidden_size)
logger.info('Dropout keep probability: %f' % args.dropout_keep_prob)
logger.info('Model max epochs: %d' % args.max_epochs)

def average_gradients(tower_grads):
    '''Calculate the average gradient for each shared variable across all towers.
    Note that this function provides a synchronization point across all towers.
    Args:
        tower_grads: List of lists of (gradient, variable) tuples. The outer list
            is over individual gradients. The inner list is over the gradient
            calculation for each tower.
    Returns:
         List of pairs of (gradient, variable) where the gradient has been averaged
         across all towers.
    '''
    average_grads = []
    for grad_and_vars in zip(*tower_grads):
        # Note that each grad_and_vars looks like the following:
        #     ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
        grads = []
        for g, _ in grad_and_vars:
            # Add 0 dimension to the gradients to represent the tower.
            expanded_g = tf.expand_dims(g, 0)

            # Append on a 'tower' dimension which we will average over below.
            grads.append(expanded_g)

        # Average over the 'tower' dimension.
        grad = tf.concat(axis=0, values=grads)
        grad = tf.reduce_mean(grad, 0)

        # Keep in mind that the Variables are redundant because they are shared
        # across towers. So .. we will just return the first tower's pointer to
        # the Variable.
        v = grad_and_vars[0][1]
        grad_and_var = (grad, v)
        average_grads.append(grad_and_var)
    return average_grads

# Small epsilon value for the BN transform
epsilon = 1e-3

def batch_norm_wrapper(inputs, is_training, decay = 0.999):
    scale = tf.get_variable('batch_norm_scale',
                    dtype=tf.float32, initializer=tf.ones([inputs.get_shape()[-1]]))
    beta = tf.get_variable('batch_norm_beta',
                    dtype=tf.float32, initializer=tf.zeros([inputs.get_shape()[-1]]))
    pop_mean = tf.get_variable('batch_norm_pop_mean',
                    dtype=tf.float32, initializer=tf.zeros([inputs.get_shape()[-1]]), trainable=False)
    pop_var = tf.get_variable('batch_norm_pop_var',
                    dtype=tf.float32, initializer=tf.ones([inputs.get_shape()[-1]]), trainable=False)

    if is_training:
        batch_mean, batch_var = tf.nn.moments(inputs,[0])
        train_mean = tf.assign(pop_mean,
                               pop_mean * decay + batch_mean * (1 - decay))
        train_var = tf.assign(pop_var,
                              pop_var * decay + batch_var * (1 - decay))
        with tf.control_dependencies([train_mean, train_var]):
            return tf.nn.batch_normalization(inputs,
                batch_mean, batch_var, beta, scale, epsilon)
    else:
        return tf.nn.batch_normalization(inputs,
            pop_mean, pop_var, beta, scale, epsilon)

# TODO: joint model

'''
class JointModel(object):
    def __init__(self):
        self.models = []

    def add_submodel(self, submodel):
        self.models.append(submodel)

    def add_joint_loss_op(self):
        self.joint_loss = tf.sum(self.
'''

class MorphingTaggingModel(object):
    def __init__(self, lexicon, config):
        self.logger = logging.getLogger('MorphingTaggingModel')
        self.lexicon = lexicon
        self.config = config
        self.idx_to_unit0, self.unit0_to_idx = self.lexicon.get_input_lexicon_for_training(0)
        self.idx_to_label, self.label_to_idx = self.lexicon.get_label_lexicon_for_training()

        self.add_placeholders()
        self.add_embeddings_op()

        for g in range(NUM_GPUS):
            # this way, parameters will be shared amongst training and inference
            with tf.variable_scope('tower%d_arch' % g):
                self.add_logits_ops(gpu_num=g, is_training=True)
                tf.get_variable_scope().reuse_variables()
                self.add_logits_ops(gpu_num=g, is_training=False)

        self.add_pred_ops()
        self.add_loss_ops()
        self.add_train_op()
        self.add_init_op()

        self.logger.info('Trainable variables')

        with tf.Session(config=tf.ConfigProto(allow_soft_placement=SOFT_PLACEMENT)) as sess:
            sess.run(self.init)
            #options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
            #run_metadata = tf.RunMetadata()
            tvars = tf.trainable_variables()
            tvars_vals = sess.run(tvars)
                     #options=options,
                     #run_metadata=run_metadata)

            #fetched_timeline = timeline.Timeline(run_metadata.step_stats)
            #chrome_trace = fetched_timeline.generate_chrome_trace_format()
            #with open('timeline_debug.json', 'w') as f:
            #    f.write(chrome_trace)

            for var, val in zip(tvars, tvars_vals):
                self.logger.info('... %s' % var.name) # , val)  # Prints the name of the variable alongside its value.

    def add_placeholders(self):
        '''
        Adds placeholders to self
        '''
        # words or morphemes
        # shape = (batch size, max length of sentence in batch)
        self.unit0_ids = tf.placeholder(tf.int32, shape=[self.config.batch_size, None],
                        name='unit0_ids')

        # number of "units" (words or morphemes)
        # shape = (batch size)
        self.sequence_lengths = tf.placeholder(tf.int32, shape=[self.config.batch_size],
                        name='sequence_lengths')

        # shape = (batch size, max length of sentence in batch)
        self.labels = tf.placeholder(tf.int32, shape=[self.config.batch_size, None],
                        name='labels')

        # hyper parameters
        # dropout keep probability (1.0 means keep all data)
        self.dropout_keep_prob = tf.placeholder(dtype=tf.float32, shape=[],
                        name='dropout_keep_prob')

        # learning rate
        self.learning_rate = tf.placeholder(dtype=tf.float32, shape=[],
                        name='learning_rate')

        # split up batch for multi-gpu
        self.batch_split_sizes = tf.placeholder(tf.int32, shape=[NUM_GPUS], name='batch_split_sizes')
        self.split_unit0_ids = tf.split(self.unit0_ids, self.batch_split_sizes, axis=0)
        self.split_sequence_lengths = tf.split(self.sequence_lengths, self.batch_split_sizes, axis=0)
        self.split_labels = tf.split(self.labels, self.batch_split_sizes, axis=0)

        # used at inference time
        self.joined_labels = tf.concat(self.split_labels, axis=0)

        self.transition_params = []
        self.logits_train = []
        self.logits_infer = []
        self.loss = []
        self.labels_pred = []
        self.init_h = []
        self.init_c = []

        for _ in range(NUM_GPUS):
            self.transition_params.append(None)
            self.logits_train.append(None)
            self.logits_infer.append(None)
            self.loss.append(None)
            self.labels_pred.append(None)
            self.init_h.append(None)
            self.init_c.append(None)

        self.standard_trainable_variables = []

    def add_embeddings_op(self):
        '''
        Adds embeddings to self (for first unit)
        '''

        combined_embeddings = None

        with tf.variable_scope('unit0'):
            # initialize unit 0's vocabulary embeddings as random normal distribution
            unit0_initial_value = tf.random_normal([len(self.idx_to_unit0),
                                                   self.config.input_unit_embedding_sizes[0]],
                stddev=1.0 / (self.config.input_unit_embedding_sizes[0]**.5), \
                seed=0)

            _unit0_embeddings = tf.get_variable(initializer=unit0_initial_value,
                                           name='_unit0_embeddings', dtype=tf.float32,
                                           trainable=self.config.do_unit_embedding_training[0])

            if self.config.do_unit_embedding_training[0]:
                self.standard_trainable_variables.append(_unit0_embeddings)

            combined_embeddings = tf.nn.embedding_lookup(_unit0_embeddings,
                                                         self.unit0_ids,
                                                         name='unit0_embeddings')

            self.logger.info('... Adding unit 0 embeddings with shape: ' + \
                  str(combined_embeddings.get_shape()))

        self.logger.info('Concatenated embedding shape: ' + \
                  str(combined_embeddings.get_shape()))

        # dropout created in static graph: takes effect if necessary when
        # referring to self.combined_embeddings
        self.combined_embeddings = tf.nn.dropout(combined_embeddings,
                                       keep_prob=self.dropout_keep_prob)

        # TODO: evaluate if there's a better way to do this.
        # is it even necessary to split up unit0_ids when we have this var?
        self.split_combined_embeddings = tf.split(self.combined_embeddings,
                                       self.batch_split_sizes, axis=0)

    def add_pred_ops(self):
        '''
        Adds labels_pred to self
        '''
        for gpu_num in range(NUM_GPUS):
            with tf.device('/gpu:%d' % gpu_num):
                if not self.config.crf:
                    self.labels_pred[gpu_num] = tf.cast(tf.argmax(self.logits_infer[gpu_num], axis=-1), tf.int32)

    def add_loss_ops(self):
        '''
        Adds loss to self
        '''
        for gpu_num in range(NUM_GPUS):
            with tf.device('/gpu:%d' % gpu_num):
                if self.config.crf:
                    with tf.variable_scope('tower%d_crf' % gpu_num):
                        log_likelihood, self.transition_params[gpu_num] = tf.contrib.crf.crf_log_likelihood(
                            self.logits_train[gpu_num], self.split_labels[gpu_num], self.split_sequence_lengths[gpu_num])
                    self.loss[gpu_num] = tf.reduce_mean(-log_likelihood)
                else:
                    losses = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=self.logits_train[gpu_num], labels=self.split_labels[gpu_num])
                    mask = tf.sequence_mask(self.split_sequence_lengths[gpu_num])
                    losses = tf.boolean_mask(losses, mask)
                    self.loss[gpu_num] = tf.reduce_mean(losses)

                #tf.add_to_collection('losses', self.loss)

                # for tensorboard
                tf.summary.scalar('loss_%d' % gpu_num, self.loss[gpu_num])

    def add_logits_ops(self, gpu_num, is_training):
        '''
        Adds logits to self
        '''

        assert not USE_BATCH_NORMALIZATION, 'USE_BATCH_NORMALIZATION needs fix below for is_training'

        # number of output labels
        nlabels = len(self.idx_to_label)

        #if is_training:
        #    # reuse can't be false...so???
        #    do_reuse = tf.AUTO_REUSE
        #else:
        #    do_reuse = True

        with tf.device('/gpu:%d' % gpu_num):
            if args.lstm_style == 'tf':
                # self.split_combined_embeddings[gpu_num]: input shape:[64 400 300]
                lstm_cell = tf.contrib.rnn.LSTMCell(num_units=self.config.combined_hidden_size, state_is_tuple=True)
                (output_fw, output_bw), _ = tf.nn.bidirectional_dynamic_rnn(lstm_cell,
                    lstm_cell, self.split_combined_embeddings[gpu_num], sequence_length=self.split_sequence_lengths[gpu_num], dtype=tf.float32)

                output = tf.concat([output_fw, output_bw], axis=-1)

                #output = tf.Print(output, [tf.shape(output)], message='output shape:')
                # print('output.get_shape()', output.get_shape().as_list())
                # [64 400 600]
            elif args.lstm_style == 'cudnn':
                combined_embeddings_input = self.split_combined_embeddings[gpu_num]

                self.logger.debug('combined_embeddings_input shape: %s' % str(combined_embeddings_input.get_shape().as_list()))

                #combined_embeddings_input = tf.reshape(combined_embeddings_input, [-1, MAX_UNIT_COUNT, sum(self.config.input_unit_embedding_sizes)])
                combined_embeddings_input = tf.reshape(combined_embeddings_input, [fixed_batch_split_sizes[gpu_num], -1, sum(self.config.input_unit_embedding_sizes)])
                combined_embeddings_input = tf.transpose(combined_embeddings_input, [1, 0, 2])

                self.logger.debug('combined_embeddings_input transpose: %s' % str(combined_embeddings_input.get_shape().as_list()))

                num_lstm_layers = 1

                cudnn_cell = tf.contrib.cudnn_rnn.CudnnLSTM(
                                                          num_layers=num_lstm_layers,
                                                          num_units=self.config.combined_hidden_size,
                                                          direction='bidirectional',
                                                          input_mode='linear_input',
                                                          input_size=self.config.combined_hidden_size,
                                                          dtype=tf.float32
                                                          # TODO: dropout here?
                                                          )

                # bizarre: init_h,init_c require first dim to be 2 because it gets divided by dir_count(2 when bidirectional)
                # model_shapes->num_layers = (*input_h)->dim_size(0) / model_shapes->dir_count;
                # otherwise first dim ends up being 0.5 (0 as int)

                # don't change batch_split_sizes later
                self.logger.debug('fixed_batch_split_sizes[%d]: %d' % (gpu_num, fixed_batch_split_sizes[gpu_num]))

                self.init_h[gpu_num] = tf.get_variable('h', dtype=tf.float32, initializer=tf.zeros([2*num_lstm_layers, fixed_batch_split_sizes[gpu_num], self.config.combined_hidden_size]))
                self.init_c[gpu_num] = tf.get_variable('c', dtype=tf.float32, initializer=tf.zeros([2*num_lstm_layers, fixed_batch_split_sizes[gpu_num], self.config.combined_hidden_size]))

                # FIXME: tweak init_scale
                params_size_t = cudnn_cell.params_size()

                self.logger.debug('params_size_t: %s' % str(params_size_t))

                cudnn_params = tf.get_variable('lstm_params',
                         initializer=tf.random_uniform([params_size_t], -0.04, 0.04), validate_shape=False)

                outputs, h, c = cudnn_cell(
                    combined_embeddings_input,
                    input_h=self.init_h[gpu_num],
                    input_c=self.init_c[gpu_num],
                    params=cudnn_params,
                    is_training=is_training
                )

                self.logger.debug('outputs.get_shape(): %s' % str(outputs.get_shape().as_list()))
                self.logger.debug('h.get_shape(): %s' % str(h.get_shape().as_list()))
                self.logger.debug('c.get_shape(): %s' % str(c.get_shape().as_list()))

                output = tf.transpose(outputs, [1, 0, 2])

                self.logger.debug('after transpose output.get_shape(): %s' % str(output.get_shape().as_list()))

            output = tf.nn.dropout(output, self.dropout_keep_prob)

        with tf.device('/gpu:%d' % gpu_num):
            #with tf.variable_scope('tower%d_proj' % gpu_num, reuse=tf.AUTO_REUSE):
            W = tf.get_variable('tower%d_W' % gpu_num, shape=[2*self.config.combined_hidden_size, nlabels],
                dtype=tf.float32, initializer=tf.random_normal_initializer( \
                    stddev=1.0 / (2*self.config.combined_hidden_size)**.5, \
                    seed=0))

            b = tf.get_variable('tower%d_b' % gpu_num, shape=[nlabels], dtype=tf.float32,
                initializer=tf.zeros_initializer())

            ntime_steps = tf.shape(output)[1]
            output = tf.reshape(output, [-1, 2*self.config.combined_hidden_size])
            #output = tf.Print(output, [tf.shape(output)], 'after proj reshape:')
            #print('after proj reshape:', output.get_shape().as_list())
            # after proj reshape:[25600 600]
            pred = tf.matmul(output, W) + b

            #if USE_BATCH_NORMALIZATION:
            #    pred_batch_norm = batch_norm_wrapper(pred, is_training)
            #else:
            pred_batch_norm = pred

            logits = tf.reshape(pred_batch_norm, [-1, ntime_steps, nlabels])

            if is_training:
                #logits = tf.Print(logits, [logits], 'logits_train:')
                self.logits_train[gpu_num] = logits
            else:
                #logits = tf.Print(logits, [logits], 'logits_infer:')
                self.logits_infer[gpu_num] = logits

    def add_train_op(self):
        '''
        Add train_op to self
        '''
        root_scope = tf.get_default_graph().get_name_scope()
        all_trainable = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, root_scope)

        self.logger.debug('Root scope: %s' % str(root_scope))

        optimizer = tf.train.AdamOptimizer(self.learning_rate)

        self.global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0), trainable=False)

        tower_grads = []
        for i in range(NUM_GPUS):
            with tf.device('/gpu:%d' % i):
                logits = self.logits_train[i]
                loss_op = self.loss[i]

                trainable = []

                for elem in all_trainable:
                    if elem.name.startswith('%s/tower' % root_scope):
                        if elem.name.startswith('%s/tower%d_' % (root_scope, i)):
                            trainable.append(elem)
                    else:
                        trainable.append(elem)

                self.logger.info('Trainable variables [gpu %d]' % i)

                for elem in trainable:
                    self.logger.info('... %s' % elem.name)

                #tf.get_variable_scope().reuse_variables()
                #summaries = tf.get_collection(tf.GraphKeys.SUMMARIES, scope)
                #print('\n'.join('{}: {}'.format(*k) for k in enumerate(summaries)))

                grads = optimizer.compute_gradients(loss_op, trainable)
                #print('\n'.join('{}: {}'.format(*k) for k in enumerate(grads)))
                tower_grads.append(grads)

                # Create a dummy optimization operation to create variables needed for optimization.
                with tf.variable_scope('tower%d_adam_opt' % i, reuse=tf.AUTO_REUSE):
                    _ = optimizer.minimize(loss_op)

                tf.get_variable_scope().reuse_variables()

        # We must calculate the mean of each gradient. Note that this is the
        # synchronization point across all towers.
        grads = average_gradients(tower_grads)

        # Apply the gradients to adjust the shared variables.
        apply_gradient_op = optimizer.apply_gradients(grads, global_step=self.global_step)

        self.train_op = apply_gradient_op

    def add_init_op(self):
        ## TODO: check that already initialize variables don't get
        # reinitialized (so that models don't collide with each other)
        self.init = tf.global_variables_initializer()

    def add_summary(self, sess):
        # tensorboard stuff
        self.merged = tf.summary.merge_all()
        self.file_writer = tf.summary.FileWriter(self.config.output_path, sess.graph)

    '''
    Prepare a feed dict for the specified model based on the given input batch
    Fills tf variables and placeholders with necessary values for sentence batch

    bio_data_sentence_list: list of BIODataSentence objects
    dropout_keep_prob:      inverse dropout rate: 1.0 means keep all data
    learning_rate:          optimizer learning rate
    '''
    def prepare_feed_dict(self, bio_data_sentence_batch, dropout_keep_prob=None, learning_rate=None):
        input_batch = []
        labels_batch = []
        sequence_lengths = []

        if USE_DYNAMIC_PADDING:
            max_length = max([len(s.labels) for s in bio_data_sentence_batch])
            self.logger.debug('max_length (per-batch dynamic): %d' % (max_length))
        else:
            max_length = MAX_UNIT_COUNT
            self.logger.debug('max_length (per-session static): %d' % (max_length))

        assert max_length <= MAX_UNIT_COUNT

        for sidx, sent in enumerate(bio_data_sentence_batch): # sent:BIODataSentence
            ## for now, only using first input
            input_batch.append(sent.get_inputs_padded(
                               input_idx=0,
                               max_length=max_length,
                               padding_token=BIOLexicon.PAD_TOK,
                               embed_id_dict=self.unit0_to_idx,
                               embed_oov_token=BIOLexicon.UNK_TOK))
            labels_batch.append(sent.get_labels_padded(
                               max_length=max_length,
                               padding_token=BIOLexicon.PAD_TOK,
                               embed_id_dict=self.label_to_idx,
                               embed_oov_token=BIOLexicon.UNK_TOK))
            sequence_lengths.append(sent.sentence_length)

            self.logger.debug('input_batch[%d] (len=%d): %s' % (sidx, len(input_batch[sidx]), str(input_batch[sidx])))
            self.logger.debug('labels_batch[%d] (len=%d): %s' % (sidx, len(labels_batch[sidx]), str(labels_batch[sidx])))
            self.logger.debug('sequence_length[%d]: %d' % (sidx, sequence_lengths[sidx]))

            assert len(input_batch[sidx]) == len(labels_batch[sidx])

        feed = {}
        feed[self.unit0_ids] = input_batch
        feed[self.sequence_lengths] = sequence_lengths
        feed[self.labels] = labels_batch

        #actual_batch_size = len(bio_data_sentence_batch)

        ## FIXME: during inference, we do a concat and the concat should work
        # as we expect. always fill batch to may size for safety
        #actual_batch_size = self.config.batch_size

        assert len(bio_data_sentence_batch) == self.config.batch_size

        feed[self.batch_split_sizes] = fixed_batch_split_sizes

        if self.dropout_keep_prob != None:
            feed[self.dropout_keep_prob] = dropout_keep_prob

        if self.learning_rate != None:
            feed[self.learning_rate] = learning_rate

        return feed, sequence_lengths

    def run_epoch(self, sess, train, dev, epoch):
        '''
        Performs one complete pass over the train set and evaluates on dev

        Args:
            sess: tensorflow session
            train: large BIODataSentence list (training set)
            dev: large BIODataSentence list (dev set)
            epoch: (int) number of the epoch
        '''
        nbatches = len(train) // self.config.batch_size
        if len(train) % self.config.batch_size != 0:
            nbatches += 1

        # hmm..couldn't this be slightly inefficient? seems like dynamic graph
        # almost

        prog = Progbar(target=nbatches)
        for i, sent_batch in enumerate(minibatches(train, self.config.batch_size, always_fill=True)):
            #if i % 100 == 0:

            fd, _ = self.prepare_feed_dict( \
                        bio_data_sentence_batch=sent_batch, \
                        dropout_keep_prob=self.config.dropout_keep_prob, \
                        learning_rate=self.config.learning_rate)

            #_, train_loss, summary = sess.run([self.train_op, self.loss, self.merged], feed_dict=fd)
            #_, train_loss = sess.run([self.train_op, self.loss], feed_dict=fd)

            #prog.update(i + 1, [('train loss', train_loss)])
            loss_ops_to_add = [self.loss[z] for z in range(NUM_GPUS)]
            _, losses = sess.run([self.train_op, loss_ops_to_add], feed_dict=fd)

            mean_loss = np.mean(losses)

            #prog.update(i + 1, [('train loss', train_loss)])
            prog.update(i + 1, [('train loss', mean_loss)]) # test

            # tensorboard
            #if i % 10 == 0:
                #self.file_writer.add_summary(summary, epoch*nbatches + i)

        acc, f1, mod_p = self.run_evaluate(sess, dev)
        self.logger.info('- dev acc {:04.2f} - f1 {:04.2f} - mod prec {:04.2f}'.format(100*acc, 100*f1, 100*mod_p))
        return acc, f1

    def predict_batch(self, sess, sents):
        '''
        Args:
            sess: a tensorflow session
            sents: list of BIODataSentence objects (batch)
                   (labels can be filled with PAD or other reserved value by
                    default, or may be passed in as gold value for convenience)
        Returns:
            pred_sents: list of new BIODataSentence objects filled with
                        predicted label data
        '''

        # ?? enforce?
        # assert len(sents) == self.config.batch_size

        cumulative_sequence_length = 0
        total_inference_time = 0
        total_viterbi_time = 0
        num_sentences = 0

        # dropout_keep_prob forced to 1.0 at inference time
        fd, sequence_lengths = self.prepare_feed_dict( \
                                   bio_data_sentence_batch=sents, \
                                   dropout_keep_prob=1.0)
        assert len(sequence_lengths) == len(sents)

        pred_sents = copy.deepcopy(sents)

        for sidx, s in enumerate(pred_sents):
        #    s.labels = []
            assert sequence_lengths[sidx] == len(s.inputs)
            cumulative_sequence_length += sequence_lengths[sidx]
        #    for i in range(sequence_lengths[sidx]):
        #        s.labels.append(BIOLexicon.PAD_TOK)

        if self.print_time_details:
            if args.use_static_padding:
                self.logger.info('... Fixed batch sequence length: %d' % MAX_UNIT_COUNT)
            else:
                self.logger.info('... Max batch sequence length: %d' % max(sequence_lengths))

        ## TODO: detect batch extension len(sents) % self.config.batch_size  remove this extra remainder from the results
        if self.config.crf:
            logits_exec = []
            transition_params_exec = []
            for i in range(NUM_GPUS):
                with tf.device('/gpu:%d' % i):
                    logits = self.logits_infer[i]
                    logits_exec.append(logits)
                    pred = self.labels_pred[i]
                    transition_params_exec.append(self.transition_params[i])
                    tf.get_variable_scope().reuse_variables()

            viterbi_sequences = []

            total_ff_time = current_milli_time()
            all_logits, all_transition_params = sess.run([logits_exec, transition_params_exec],
                    feed_dict=fd)
            total_ff_time = current_milli_time() - total_ff_time
            total_inference_time = total_ff_time

            if self.print_time_details:
                self.logger.info('...... Feed forward time: %.4fms/unit (batch total %.4fms)' % (total_ff_time / cumulative_sequence_length, total_ff_time))

            sidx = 0
            for gpu_idx, (this_gpu_logits, this_gpu_transition_params) in enumerate(zip(all_logits, all_transition_params)):
                # logits may be longer due to padding
                this_gpu_sequence_lengths = sequence_lengths[sidx:sidx+len(this_gpu_logits)]

                # iterate over the sentences
                for logit, sequence_length in zip(this_gpu_logits, this_gpu_sequence_lengths):
                    if sidx >= len(sents):
                        self.logger.info('... Breaking before padding during inference')
                        break # ignore extra padding

                    # keep only the valid time steps
                    logit = logit[:sequence_length]
                    viterbi_time = current_milli_time()
                    viterbi_sequence, viterbi_score = tf.contrib.crf.viterbi_decode(
                                    logit, this_gpu_transition_params)
                    viterbi_time = current_milli_time() - viterbi_time
                    total_viterbi_time += viterbi_time

                    viterbi_sequences.append(viterbi_sequence)
                    sidx += 1
                    num_sentences += 1

            total_inference_time += total_viterbi_time

            if self.print_time_details:
                self.logger.info('...... Viterbi time: %.4fms/unit (batch total %.4fms)' % (total_viterbi_time / cumulative_sequence_length, total_viterbi_time))
                self.logger.info('... Total inference time: %.4fms/unit (batch total %.4fms)' % (total_inference_time / cumulative_sequence_length, total_inference_time))

            for sidx, s in enumerate(pred_sents):
                s.labels = []
                #print('sequence_lengths[sidx]', sequence_lengths[sidx])
                #print('viterbi_sequences[sidx]', viterbi_sequences[sidx])
                for i in range(sequence_lengths[sidx]):
                    #print('i', i)
                    #print('viterbi_sequences[sidx][i]', viterbi_sequences[sidx][i])

                    ## FIXME: is it possible for the NN to return a non-sensible value here?
                    s.labels.append(self.idx_to_label[viterbi_sequences[sidx][i]])
                #print('s.labels', s.labels)

            #return viterbi_sequences, sequence_lengths
        else:
            ## TODO

            labels_pred_exec = []
            for i in range(NUM_GPUS):
                with tf.device('/gpu:%d' % i):
                    logits = self.logits_infer[i]
                    pred = self.labels_pred[i]
                    labels_pred_exec.append(pred)
                    tf.get_variable_scope().reuse_variables()

            labels_pred_concat = tf.concat(labels_pred_exec, axis=0)

            total_ff_time = current_milli_time()
            labels_pred = sess.run(labels_pred_concat, feed_dict=fd)
            total_ff_time = current_milli_time() - total_ff_time
            total_inference_time = total_ff_time

            if self.print_time_details:
                self.logger.info('... Total inference time: %.4fms/unit (batch total %.4fms)' % (total_inference_time / cumulative_sequence_length, total_inference_time))

            ## TODO: check sequence_lengths against input lengths

            for sidx, s in enumerate(pred_sents):
                #print('crf=false, sidx=', sidx)
                if sidx >= len(sents):
                    self.logger.info('... Breaking before padding during inference')
                    break # ignore extra padding

                s.labels = []
                for i in range(sequence_lengths[sidx]):
                    ## FIXME: is it possible for the NN to return a non-sensible value here?
                    s.labels.append(self.idx_to_label[labels_pred[sidx][i]])
                #print('s.labels', s.labels)

            #return labels_pred, sequence_lengths

        return pred_sents

    def run_evaluate(self, sess, test):
        '''
        Evaluates performance on specified test/dev set

        Args:
            sess: tensorflow session
            test: large BIODataSentence list (dev/test set)
        '''

        nbatches = len(test) // self.config.batch_size
        if len(test) % self.config.batch_size != 0:
            nbatches += 1

        correct_preds = 0
        total_preds = 0
        total_correct = 0
        correct_mod = 0
        total_mod = 0
        accs = []

        #prog = Progbar(target=nbatches)
        # always fill to match batch size by wrapping around if necessary
        for i, gold_sent_batch in enumerate(minibatches(test, self.config.batch_size, always_fill=True)):
            if i % 50 == 0:
                self.logger.info('Evaluate: batch %d/%d' % (i+1, nbatches))
                self.print_time_details = True
            else:
                self.print_time_details = False

            pred_sent_batch = self.predict_batch(sess, gold_sent_batch)

            assert len(gold_sent_batch) == len(pred_sent_batch)

            for sidx in range(len(gold_sent_batch)):
                gold_chunks = gold_sent_batch[sidx].get_label_chunks()
                pred_chunks = pred_sent_batch[sidx].get_label_chunks()
                correct_chunks = gold_chunks & pred_chunks

                self.logger.debug('gold_chunks: ' + str(sorted(gold_chunks)))
                self.logger.debug('pred_chunks: ' + str(sorted(pred_chunks)))

                for (chunk_idx, chunk_label) in gold_chunks:
                    if chunk_label.startswith('MOD') or chunk_label.startswith('B-MOD') or chunk_label.startswith('I-MOD'):
                        total_mod += 1

                for (chunk_idx, chunk_label) in correct_chunks:
                    if chunk_label.startswith('MOD') or chunk_label.startswith('B-MOD') or chunk_label.startswith('I-MOD'):
                        correct_mod += 1

                correct_preds += len(correct_chunks)
                total_preds += len(pred_chunks)
                total_correct += len(gold_chunks)
                accs += map(lambda items: items[0] == items[1], list(zip(gold_sent_batch[sidx].labels, pred_sent_batch[sidx].labels)))

        self.logger.info('correct_preds: ' + str(correct_preds))
        self.logger.info('total_mod: ' + str(total_mod))
        self.logger.info('total_preds: ' + str(total_preds))
        self.logger.info('total_correct: ' + str(total_correct))

        p = correct_preds / total_preds if correct_preds > 0 else 0
        r = correct_preds / total_correct if correct_preds > 0 else 0
        mod_p = correct_mod / total_mod if correct_mod > 0 else 0
        f1 = 2 * p * r / (p + r) if correct_preds > 0 else 0
        acc = np.mean(accs)
        return acc, f1, mod_p

    def train(self, sent_train, sent_dev):
        best_score = 0
        saver = tf.train.Saver()
        # for early stopping
        nepoch_no_imprv = 0

        #config = tf.ConfigProto()
        #config.gpu_options.allow_growth = True

        #with tf.Session(config=config) as sess:
        with tf.Session(config=tf.ConfigProto(allow_soft_placement=SOFT_PLACEMENT)) as sess:
            sess.run(self.init)
            # tensorboard
            self.add_summary(sess)
            for epoch in range(self.config.nepochs):
                self.logger.info('Epoch {:} out of {:}'.format(epoch + 1, self.config.nepochs))

                if not args.no_shuffle:
                    self.logger.debug('Shuffling training set...')
                    # do shuffle of sent_train
                    random.shuffle(sent_train)

                acc, f1 = self.run_epoch(sess, sent_train, sent_dev, epoch)

                # decay learning rate
                self.config.learning_rate *= self.config.lr_decay

                # early stopping and saving best parameters
                if f1 >= best_score:
                    nepoch_no_imprv = 0
                    if not os.path.exists(self.config.model_output):
                        os.makedirs(self.config.model_output)
                    saver.save(sess, self.config.model_output)
                    best_score = f1
                    self.logger.info('- new best score!')

                else:
                    nepoch_no_imprv += 1
                    if nepoch_no_imprv >= self.config.nepoch_no_imprv:
                        self.logger.info('- early stopping {} epochs without improvement'.format(
                                        nepoch_no_imprv))
                        break

    def evaluate(self, sent_test):
        saver = tf.train.Saver()
        with tf.Session(config=tf.ConfigProto(allow_soft_placement=SOFT_PLACEMENT)) as sess:
            self.logger.info('Testing model over test set')
            saver.restore(sess, self.config.model_output)
            acc, f1, mod_p = self.run_evaluate(sess, sent_test)
            self.logger.info('- test acc {:04.2f} - f1 {:04.2f} - mod prec {:04.2f}'.format(100*acc, 100*f1, 100*mod_p))

train = BIODataInput('data/%s/%s_train.txt' % (args.model_data_dir, args.model_type))
test = BIODataInput('data/%s/%s_test.txt' % (args.model_data_dir, args.model_type))
dev = BIODataInput('data/%s/%s_dev.txt' % (args.model_data_dir, args.model_type))
lexicon = BIOLexicon()
lexicon.add_lexicon_data(train.get_lexicon())
#lexicon.add_lexicon_data(test.get_lexicon())
#lexicon.add_lexicon_data(dev.get_lexicon())

## FIXME: make sure to delete existing dir first
lexicon.save('%slexicon' % args.model_output_dir)

config.input_unit_embedding_sizes = [args.input_unit_embedding_size]
config.do_unit_embedding_training = [True]
config.dropout_keep_prob = args.dropout_keep_prob
config.combined_hidden_size = args.combined_hidden_size
config.nepochs = args.max_epochs
config.nepoch_no_imprv = 3
config.learning_rate = 0.001
config.lr_decay = 0.9

# especially for CRF...output action count is important for memory
# max sentence length can also be a factor
# if you run out of memory, check current_train_vars
# here are some reasonable values (471 morphing output actions):
#
#current_train_vars [<tf.Variable 'morphing/unit0/_unit0_embeddings:0' shape=(1902, 300) dtype=float32_ref>, <tf.Variable 'morphing/bi-lstm/bidirectional_rnn/fw/lstm_cell/kernel:0' shape=(600, 1200) dtype=float32_ref>, <tf.Variable 'morphing/bi-lstm/bidirectional_rnn/fw/lstm_cell/bias:0' shape=(1200,) dtype=float32_ref>, <tf.Variable 'morphing/proj/W:0' shape=(600, 471) dtype=float32_ref>, <tf.Variable 'morphing/proj/b:0' shape=(471,) dtype=float32_ref>, <tf.Variable 'morphing/transitions:0' shape=(471, 471) dtype=float32_ref>]
#

config.crf = args.use_crf
config.batch_size = args.batch_size

SHOW_EFFICIENCY_EVERY_N = 1000

def efficiency_show(gpu_efficiency_rating_over_time, cpu_efficiency_rating_over_time):
    assert len(gpu_efficiency_rating_over_time) >= SHOW_EFFICIENCY_EVERY_N
    assert len(cpu_efficiency_rating_over_time) >= SHOW_EFFICIENCY_EVERY_N

    # show last N results
    #last_N_times, last_N_ratings = zip(*gpu_efficiency_rating_over_time[-SHOW_EFFICIENCY_EVERY_N:])

    # show avg of all time
    gpu_last_N_times, gpu_last_N_ratings = zip(*gpu_efficiency_rating_over_time)
    gpu_avg_eff = sum(gpu_last_N_ratings) / float(len(gpu_last_N_ratings))
    logger.info('Average GPU efficiency(0.0~1.0): %.2f' % gpu_avg_eff)

    cpu_last_N_times, cpu_last_N_ratings = zip(*cpu_efficiency_rating_over_time)
    cpu_avg_eff = sum(cpu_last_N_ratings) / float(len(cpu_last_N_ratings))
    logger.info('Average CPU efficiency(0.0~1.0): %.2f' % cpu_avg_eff)

def efficiency_monitor_start(data):
    gpu_efficiency_rating_over_time = []
    cpu_efficiency_rating_over_time = []
    while do_monitor:
        # monitor efficiency
        # TODO: also add to tensorboard
        cpu_loads = get_cpu_loads()
        gpu_loads, gpu_efficency_rating = get_gpu_loads_and_efficiency(GPU_IDS)

        # if rating is 0, gpu is probably just not in use (could be invalid data?? or bottlenecked on other area of app...might be important to know anyway)

        curtime = current_milli_time()

        cpu_efficiency = sum(cpu_loads) / float(len(cpu_loads))
        cpu_efficiency_rating_over_time.append((curtime, cpu_efficiency))

        #if efficency_rating != 0.0 and 0.0 not in gpu_loads:
        gpu_efficiency_rating_over_time.append((curtime, gpu_efficency_rating))
        if len(gpu_efficiency_rating_over_time) > 0 and len(gpu_efficiency_rating_over_time) % SHOW_EFFICIENCY_EVERY_N == 0:
            efficiency_show(gpu_efficiency_rating_over_time, cpu_efficiency_rating_over_time)

        # we don't want the measurement to get stuck on a certain point in the epoch
        # so sleep every random amount of seconds
        time.sleep(0.5+random.random())

g1 = tf.Graph()
with g1.as_default() as g:
    with tf.variable_scope(args.model_type):
        model = MorphingTaggingModel(lexicon=lexicon, config=config)
        # begin efficiency monitoring

        do_monitor = True
        t = threading.Thread(target=efficiency_monitor_start, args=('',))
        t.start()

        model.train(sent_train=train.sentences, sent_dev=dev.sentences)

        do_monitor = False # stop monitoring

        model.evaluate(sent_test=test.sentences)


from pathlib import Path

import typer

from peppyproject import PyProjectConfiguration

app = typer.Typer(add_completion=False)

@app.command()
def main(
    directory: Path = typer.Argument(
        None, help="directory from which to read configuration"
    ),
    output_filename: Path = typer.Option(
        None, "-o", "--output", help="path to which to write TOML"
    ),
):
    """
    read a Python project configuration and output a PEP621-compliant `pyproject.toml`
    """

    if directory is None:
        directory = Path.cwd()

    configuration = PyProjectConfiguration.from_directory(directory=directory)
    toml_string = configuration.configuration
    if output_filename is None:
        print(toml_string)
    else:
        if not output_filename.parent.exists():
            output_filename.mkdir(parents=True, exist_ok=True)
        with open(output_filename, "w") as toml_file:
            toml_file.write(toml_string)

if __name__ == "__main__":
    app()


"""
    Unit tests over parameter combinations of the library.

    TODO ADD MORE
"""
from minorminer import find_embedding as find_embedding_orig
from warnings import warn
import os
import sys
import time

# Given that this test is in the tests directory, the calibration data should be
# in a sub directory. Use the path of this source file to find the calibration
calibration_dir = os.path.join(os.path.dirname(
    os.path.abspath(__file__)), "calibration")

def find_embedding(Q, A, return_overlap=False, **args):
    args['verbose'] = 0
    args['tries'] = 1
    if return_overlap:
        emb, succ = find_embedding_orig(
            Q, A, return_overlap=return_overlap, **args)
        if not succ:
            return emb, succ
        elif check_embedding(Q, A, emb, **args):
            if check_embedding.warning:
                warn(check_embedding.warning, RuntimeWarning)
            return emb, succ
        else:
            raise RuntimeError(
                "bad embedding reported as success (%s)" % (check_embedding.errcode))
    else:
        emb = find_embedding_orig(Q, A, return_overlap=return_overlap, **args)
        if emb:
            if not check_embedding(Q, A, emb, **args):
                raise RuntimeError(
                    "bad embedding reported as success (%s)" % (check_embedding.errcode))
            elif check_embedding.warning:
                warn(check_embedding.warning, RuntimeWarning)
        return emb

def check_embedding(Q, A, emb, **args):
    from networkx import Graph, is_connected
    check_embedding.warning = None
    Qg = Graph()
    Ag = Graph()
    Qg.add_edges_from(Q)
    Ag.add_edges_from(A)

    qubhits = 0
    footprint = set()
    var = {}
    for x in Qg:
        try:
            embx = emb[x]
        except KeyError:
            check_embedding.errcode = "missing chain"
            return False
        for q in embx:
            var[q] = x
        footprint.update(embx)
        qubhits += len(embx)
        if not is_connected(Ag.subgraph(embx)):
            check_embedding.errcode = "broken chain for %s: (%s)" % (x, embx)
            return False
    if len(footprint) != qubhits:
        check_embedding.errcode = "overlapped chains"
        return False

    Qv = Graph()
    for p, q in Ag.edges():
        try:
            Qv.add_edge(var[p], var[q])
        except KeyError:
            continue
    for x, y in Qg.edges():
        if not Qv.has_edge(x, y):
            check_embedding.errcode = "missing edge"
            return False

    for x, chain in args.get("fixed_chains", {}).items():
        if set(chain) != set(emb[x]):
            check_embedding.errcode = "fixed chain mismatch"
            return False
    for x, domain in args.get("restrict_chains", {}).items():
        if not set(domain) >= set(emb[x]):
            check_embedding.warning = "restrict chain mismatch"

    return True

def Path(n):
    return [(i, i + 1) for i in range(n - 1)]

def Grid(n):
    return [((x, y), (x + dx, y + dy)) for a in range(n) for b in range(n - 1) for (x, y, dx, dy) in [(a, b, 0, 1), (b, a, 1, 0)]]

def Clique(n):
    return [(u, v) for u in range(n) for v in range(u)]

def Biclique(n):
    return [(u, v) for u in range(n) for v in range(n, 2 * n)]

def Chimera(n, l=4):
    return [((x, y, u, k), (x + dx, y + dy, u, k))
            for a in range(n)
            for b in range(n - 1)
            for k in range(l)
            for x, y, u, dx, dy in [(b, a, 0, 1, 0), (a, b, 1, 0, 1)]
            ] + [((x, y, 0, k), (x, y, 1, kk)) for x in range(n) for y in range(n) for k in range(l) for kk in range(l)]

def NAE3SAT(n):
    import networkx
    from math import ceil
    from random import seed, randint
    seed(18293447845779813366)
    c = int(ceil(sum(randint(1, ceil(n * 4.2)) for _ in range(100)) / 100.))
    return networkx.generators.k_random_intersection_graph(c, n, 3).edges()

def ChordalCycle(p):
    import networkx
    G = networkx.generators.chordal_cycle_graph(p)
    G.remove_edges_from(list(G.selfloop_edges()))
    return G.edges()

def GeometricGraph(n, pos=None):
    import networkx
    G = networkx.generators.geometric.random_geometric_graph(
        n, n**-.333, dim=2, pos=pos)
    if pos is not None:
        for g in G:
            if len(list(G[g])) == 0:
                del pos[g]
    return G.edges()

def CartesianProduct(n):
    import networkx
    K = networkx.generators.complete_graph(n)
    return networkx.product.cartesian_product(K, K).edges()

def GridChimeraEmbedding(n):
    emb = {}
    M = [[0, 2, 2, 0], [1, 3, 3, 1], [1, 3, 3, 1], [0, 2, 2, 0]]
    for x in range(n):
        for y in range(n):
            emb[x, y] = [(x // 2, y // 2, 0, M[x % 4][y % 4]),
                         (x // 2, y // 2, 1, M[y % 4][x % 4])]
    return emb

def mask_wxw(n, w=2, l=4):
    return {(X // w, Y // w): [(x, y, u, k) for x in range(X, X + w) for y in range(Y, Y + w) for u in (0, 1) for k in range(l)] for X in range(0, n, w) for Y in range(0, n, w)}

success_count_functions = []

def success_count(n, *a, **k):
    from functools import wraps
    from math import log

    def count_successes(f):
        global success_count_functions
        success_count_functions.append([f, n, a, k])
        if os.path.exists(os.path.join(calibration_dir, f.__name__)):
            S, N = load_success_count_calibration(f)
            N += (S == N)
            accept_prob = .01  # 1% false negative rate
            tts = int(log(accept_prob * S / N, 1 - S / N) + 1)
            false_prob = (S / N) * (1 - S / N)**tts

            @wraps(f)
            def test_run():
                for i in range(tts):
                    if f(*a, **k):
                        break
                else:
                    assert False, "took %d tries without success, this should only happen %.02f%% of the time" % (
                        tts, false_prob * 100)
        else:
            def test_run():
                raise RuntimeError(
                    "%s is not calibrated -- run calibrate_all() or calibrate_new()" % (f.__name__))

        test_run.original = f
        return test_run
    return count_successes

def calibrate_success_count(f, n, a, k, directory=calibration_dir, M=None):
    succ = 0
    if M is None:
        M = 10000
    N = M * n
    print("calibrating %s, %d trials " % (f.__name__, N))
    t0 = time.clock()
    for i in range(N):
        if i % (N / 10) == 0:
            print("%d " % (10 * i // N))
            sys.stdout.flush()
        succ += bool(f(*a, **k))
    print()
    dt = time.clock() - t0
    print("%s: %.04e per trial; success rate %.01f%% " %
          (f.__name__, dt / N, succ * 100. / N))
    if directory != calibration_dir and os.path.exists(os.path.join(calibration_dir, f.__name__)):
        olds, oldn = load_success_count_calibration(f)
        print("standard is %.01f%%" % (olds * 100. / oldn))
    else:
        print()
    with open(os.path.join(directory, f.__name__), "w") as cal_f:
        cal_f.write(repr((succ, float(N))))

def load_success_count_calibration(f, directory=calibration_dir):
    with open(os.path.join(directory, f.__name__)) as cal_f:
        return eval(cal_f.read())

def calibrate_all(directory=calibration_dir, M=None):
    global success_count_functions

    if not os.path.exists(directory):
        os.mkdir(directory)

    for f, n, a, k in success_count_functions:
        calibrate_success_count(f, n, a, k, directory=directory, M=M)
        print()

def calibrate_new(directory=calibration_dir, M=None):
    for f, n, a, k in success_count_functions:
        if os.path.exists(os.path.join(directory, f.__name__)):
            continue
        else:
            calibrate_success_count(f, n, a, k, directory=directory, M=M)

def success_perfect(n, *a, **k):
    from functools import wraps

    def is_perfect(f):
        @wraps(f)
        def test_run():
            for _ in range(n):
                assert bool(f(*a, **k)), "test fail"
        test_run.original = f
        return test_run
    return is_perfect

def success_bounce(n, *a, **k):
    from functools import wraps

    def is_perfect(f):
        @wraps(f)
        def test_run():
            succs = sum(bool(f(*a, **k)) for _ in range(n))
            assert False, "%d successes out of %d trials" % (succs, n)
        test_run.original = f
        return test_run
    return is_perfect

def check_args(prob, hard, initial_chains=None, fixed_chains=None, restrict_chains=None, skip_initialization=False):
    import networkx
    probg = networkx.Graph()
    probg.add_edges_from(prob)

    hardg = networkx.Graph()
    hardg.add_edges_from(hard)

    assert networkx.is_connected(hardg), "hardware graph not connected"
    assert networkx.is_connected(probg), "problem graph not connected"

    if fixed_chains is not None:
        for v, chain in fixed_chains.items():
            assert probg.has_node(
                v), "fixed_chains vars not contained in problem graph"
            for q in chain:
                assert hardg.has_node(
                    q), "fixed_chains chains not contained in hardware graph"
        if initial_chains is not None:
            for v in fixed_chains:
                assert v not in initial_chains, "fixed_chains chains overwrite initial chains"
        if restrict_chains is not None:
            for v in fixed_chains:
                assert v not in restrict_chains, "fixed_chains chains are restricted"

    if initial_chains is not None:
        for v, chain in initial_chains.items():
            assert probg.has_node(
                v), "initial vars not contained in problem graph"
            for q in chain:
                assert hardg.has_node(
                    q), "initial chains not contained in hardware graph"
        if skip_initialization:
            for u, v in probg.edges():
                edgelord = {z for q in initial_chains[v] for z in hardg.neighbors(
                    q)} | set(initial_chains[v])
                assert set(
                    initial_chains[u]) & edgelord, "%s and %s are connected as variables but not as initials" % (u, v)

    if restrict_chains is not None:
        fullset = set(hardg.nodes())
        for v, chain in restrict_chains.items():
            assert probg.has_node(
                v), "restricted vars not contained in problem graph"
            for q in chain:
                assert hardg.has_node(
                    q), "restricted chains not contained in hardware graph"
        for u, v in probg.edges():
            edgelord = {z for q in restrict_chains.get(v, fullset) for z in hardg.neighbors(
                q)} | set(restrict_chains.get(v, fullset))
            assert set(restrict_chains.get(
                u, fullset)) & edgelord, "%s and %s are connected as variables but not as domains" % (u, v)

@success_count(100, 5)
def test_path_label_00(n):
    p = Path(n)
    return find_embedding(p, p)

@success_count(100, 5)
def test_path_label_01(n):
    p = Path(n)
    L = [str(i) for i in range(n)]
    Lp = [(L[x], L[y]) for x, y in p]
    return find_embedding(p, Lp)

@success_count(100, 5)
def test_path_label_10(n):
    p = Path(n)
    L = [str(i) for i in range(n)]
    Lp = [(L[x], L[y]) for x, y in p]
    return find_embedding(Lp, p)

@success_count(100, 5)
def test_path_label_11(n):
    p = Path(n)
    L = [str(i) for i in range(n)]
    Lp = [(L[x], L[y]) for x, y in p]
    return find_embedding(Lp, Lp)

@success_count(30, 3)
def test_grid_init_restrict(n):
    from random import choice
    chim = Chimera(n, l=4)
    mask = mask_wxw(n, 1, l=4)
    grid = Grid(2 * n)

    init = {(x, y): [choice(mask[x // 2, y // 2])]
            for x in range(2 * n) for y in range(2 * n)}
    doms = {(x, y): mask[x // 2, y // 2]
            for x in range(2 * n) for y in range(2 * n)}

    return find_embedding(grid, chim, initial_chains=init, restrict_chains=doms, skip_initialization=False)

@success_count(30, 3)
def test_grid_init(n):
    from random import choice
    chim = Chimera(n, l=4)
    mask = mask_wxw(n, 1, l=2)
    grid = Grid(2 * n)

    init = {(x, y): mask[x // 2, y // 2]
            for x in range(2 * n) for y in range(2 * n)}

    return find_embedding(grid, chim, initial_chains=init, skip_initialization=False)

@success_count(30, 15, 7)
def test_nae3sat(n, m):
    from random import choice
    chim = Chimera(m)
    prob = NAE3SAT(n)

    return find_embedding(prob, chim)

@success_count(30, 79, 6)
def test_expander(p, m):
    prob = ChordalCycle(p)
    chim = Chimera(m)

    return find_embedding(prob, chim)

@success_count(30, 5)
def test_cartesian(n):
    prob = CartesianProduct(n)
    chim = Chimera(n, l=n)

    return find_embedding(prob, chim)

@success_count(30, 45, 6)
def test_geometric_nohint(n, m):
    prob = GeometricGraph(n)
    chim = Chimera(m)

    return find_embedding(prob, chim)

@success_count(30, 55, 6)
def test_geometric_hint(n, m):
    from random import randint
    pos = {}
    chains = {}
    for i in range(n):
        x = randint(0, m - 1)
        k1 = randint(0, 3)
        y = randint(0, m - 1)
        k2 = randint(0, 3)
        pos[i] = (4 * x + k2) / 4. / m, (4 * y + k1) / 4. / m
        chains[i] = (x, y, 0, k1), (x, y, 1, k2)
    prob = GeometricGraph(n, pos)
    chim = Chimera(m)

    return find_embedding(prob, chim, initial_chains={i: c for i, c in chains.items() if i in pos})

@success_count(30, 3)
def test_grid_restrict(n):
    chim = Chimera(n)
    mask = mask_wxw(n, 1)
    grid = Grid(2 * n)

    doms = {(x, y): mask[x // 2, y // 2]
            for x in range(2 * n) for y in range(2 * n)}

    check_args(grid, chim, restrict_chains=doms)

    return find_embedding(grid, chim, restrict_chains=doms)

@success_perfect(100, 4)
def test_grid_with_answer_fast(n):
    chim = Chimera(n)
    mask = mask_wxw(n, 1)
    grid = Grid(2 * n)

    init = GridChimeraEmbedding(2 * n)

    check_args(grid, chim, initial_chains=init, skip_initialization=True)

    return find_embedding(grid, chim, initial_chains=init, skip_initialization=True, chainlength_patience=0)

@success_perfect(100, 2)
def test_grid_with_answer_slow(n):
    chim = Chimera(n)
    mask = mask_wxw(n, 1)
    grid = Grid(2 * n)

    init = GridChimeraEmbedding(2 * n)

    check_args(grid, chim, initial_chains=init, skip_initialization=True)

    return find_embedding(grid, chim, initial_chains=init, skip_initialization=True, chainlength_patience=10)

@success_count(30, 5)
def test_grid_suspend(n):
    chim = Chimera(n)
    mask = mask_wxw(n, 1)
    grid = Grid(2 * n)

    suspg = [((x, y), (x // 2, y // 2, 0))
             for x in range(2 * n) for y in range(2 * n)]
    suspc = [((x, y, 0), m) for x in range(n)
             for y in range(n) for m in mask[x, y]]

    suspension = {(x, y, 0): [(x, y, 0)] for x in range(n) for y in range(n)}

    return find_embedding(grid + suspg, chim + suspc, fixed_chains=suspension, chainlength_patience=0)

@success_count(30, 5)
def test_grid_plant_suspend(n):
    chim = Chimera(n)
    mask = mask_wxw(n, 1)
    grid = Grid(2 * n)

    suspg = [((x, y), (x // 2, y // 2, 0))
             for x in range(2 * n) for y in range(2 * n)]
    suspc = [(m, (x, y, 0)) for x in range(n)
             for y in range(n) for m in mask[x, y]]

    suspension = {(x, y, 0): [(x, y, 0)] for x in range(n) for y in range(n)}
    init = {(x, y): mask[x // 2, y // 2]
            for x in range(2 * n) for y in range(2 * n)}

    return find_embedding(grid + suspg, chim + suspc, fixed_chains=suspension, initial_chains=init, chainlength_patience=0)

@success_count(30, 5)
def test_grid_suspend_domain(n):
    chim = Chimera(n)
    mask = mask_wxw(n, 1)
    grid = Grid(2 * n)

    suspg = [((x, y), (x // 2, y // 2, 0))
             for x in range(2 * n) for y in range(2 * n)]
    suspc = [((x, y, 0), m) for x in range(n)
             for y in range(n) for m in mask[x, y]]

    suspension = {(x, y, 0): [(x, y, 0)] for x in range(n) for y in range(n)}
    doms = {(x, y): mask[x // 2, y // 2]
            for x in range(2 * n) for y in range(2 * n)}

    check_args(grid + suspg, chim + suspc, fixed_chains=suspension,
               skip_initialization=False, restrict_chains=doms)

    return find_embedding(grid + suspg, chim + suspc, fixed_chains=suspension, restrict_chains=doms, chainlength_patience=0)

@success_count(30, 5)
def test_grid_cheat_domain(n):
    chim = Chimera(n)
    grid = Grid(2 * n)
    cheat = GridChimeraEmbedding(2 * n)

    return find_embedding(grid, chim, restrict_chains=cheat, chainlength_patience=0)

@success_count(30, 2)
def test_biclique_chimera(n):
    chim = Chimera(n)
    kliq = Biclique(4 * n)

    return find_embedding(kliq, chim, chainlength_patience=0)

@success_count(30, 5)
def test_path_cheat_domain(n):
    P = Path(n)
    cheat = {p: [p] for p in range(n)}

    return find_embedding(P, P, restrict_chains=cheat, chainlength_patience=0)

@success_count(30, 6, 25)
def test_clique(n, k):
    chim = Chimera(n)
    cliq = Clique(k)

    return find_embedding(cliq, chim, chainlength_patience=0)

@success_perfect(20, 25, 25)
def test_clique_clique(n, k):
    cliq = Clique(k)

    return find_embedding(cliq, cliq, chainlength_patience=0)

@success_perfect(3, 16)
def test_clique_large_nosegfault(n):
    chim = Chimera(n)
    cliq = Clique(4 * n + 2)

    return not find_embedding(cliq, chim, chainlength_patience=0, timeout=1)

@success_count(30, 6, 25)
def test_clique_parallel(n, k):
    chim = Chimera(n)
    cliq = Clique(k)

    return find_embedding(cliq, chim, chainlength_patience=0, threads=2)

@success_count(30, 3, 13)
def test_clique_term(n, k):
    chim = Chimera(n)
    cliq = Clique(k)
    cterm = [((n // 2, n // 2, 0, 0), k)]
    kterm = [(0, k)]
    fix = {k: [k]}
    return find_embedding(cliq + kterm, chim + cterm, fixed_chains=fix, chainlength_patience=0)

@success_count(30, 8)
def test_grid_heal_A(n):
    from random import randint
    grid = Grid(2 * n)
    chim = Chimera(n + 2)
    breaks = {(x, x, x % 2, randint(0, 3)) for x in range(1, 4)}
    chim = [e for e in chim if not breaks.intersection(e)]

    emb = GridChimeraEmbedding(2 * n)
    i_emb = {}
    for v, chain in emb.items():
        remainder = {(x + 1, y + 1, u, k)
                     for x, y, u, k in chain}.difference(breaks)
        if remainder:
            i_emb[v] = remainder

    return find_embedding(grid, chim, initial_chains=i_emb, chainlength_patience=0)

@success_count(30, 4)
def test_grid_heal_B(n):
    from random import randint
    grid = Grid(2 * n)
    chim = Chimera(n + 2)
    breaks = {(x, x, x % 2, randint(0, 3)) for x in range(1, 4)}

    chim = [e for e in chim]
    chimb = [(b, (b, None)) for b in breaks]
    gridb = [(b, (b, None)) for b in breaks]
    f_emb = {(b, None): [(b, None)] for b in breaks}

    emb = GridChimeraEmbedding(2 * n)

    return find_embedding(grid + gridb, chim + chimb, initial_chains=emb, fixed_chains=f_emb, chainlength_patience=0)

@success_perfect(1000, 3)
def test_fail_impossible(n):
    Kn = Clique(n)  # we're gonna try to embed this here clique
    Pn = Path(n)  # into this here path, and it ain't gonna work

    return not find_embedding(Kn, Pn)

@success_perfect(1, 16, .1)
def test_fail_timeout(n, t):
    Kn = Clique(4 * n + 1)  # we're gonna try to embed this here clique
    # into this here chimera, and it might work but we'll time out
    Cn = Chimera(n)

    return not find_embedding(Kn, Cn, tries=1e6, max_no_improvement=1e6, inner_rounds=1e6, timeout=t, threads=4)