from database import ActiveDatabase
from MatchClass import Match
from UserClass import User
from SensitiveConfig import *
from config import *
import requests
import time
import logging

logger = logging.getLogger(__name__)

def initLogger():
    logger.setLevel(logging.INFO)

    #Creating file handler
    handler = logging.FileHandler('{}/MatchFinder_{}.log'.format(LOG_DIRECTORY,time.strftime('%m%d_%H%M%S')))
    handler.setLevel(logging.INFO)

    #Creating logging format
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    handler.setFormatter(formatter)

    #add handlers to the logger
    logger.addHandler(handler)

def fetchActiveUsers():
    ''' fetches all active users
        Args: None

        Return:
            userList = dictionary of active users
            success = boolean on whether it grabbed users
    '''

    logger.info("Fetching active users")

    #accessing active database
    dbConnection = ActiveDatabase().getConn()
    userConnection= dbConnection.collection(USERS)
    try:
        userList = userConnection.stream()
        logger.info("Query was successful.")
        success = True
    except Exception as e:
        logger.warning("Exception was thrown!!! --- {}".format(e))
        userList = None
        success = False

    logger.info("Returning -- Success: {}".format(success))
    logger.info("          -- UserList: {}".format(userList))
    return userList,success

def RateLimitFromHeader(limitCount):
    ''' Detects if script gets close to rate limit for api, then slows it down
        Args:
            limitCount = response.header["X-App-Rate-Limit-Count"]
        Return:
            None
    '''
    list = limitCount.split(',')
    current,limit = list[1].split(':',1)
    ratio = float(current)/float(limit)
    sleepTime = 0

    logger.info("Checking rate limit. current: {} limit: {} ratio: {}".format(current,limit,ratio))

    #Will need to be changed with producer key. Built for 120 rate limit every 2 minutes
    if ratio == .75:
        sleepTime = 2

    #Keeping this open for change, more distict sleep times may be added
    if sleepTime:
        time.sleep(sleepTime)

def fetchMatchEntries(user):
    ''' fetches n number of matches from RiotGames api
    Args: dictionary
        region: "NA1"
        accountId: "6UoWroRfGgxqnI0Oj-KahipCC2TcQWiKkAar-BWyYGoxAw"
        summonerId: "afternoonview'
    Returns:

    '''
    logger.info("Fetching match entries for user: {}".format(user["summonerId"]))

    try:
        url = RIOT_FETCH_MATCHES_URL.format(user["region"], user["accountId"], RIOT_API_KEY, MAX_MATCH_FETCH)
        logger.info("GET method: {}".format(url))
        response = requests.get(url)
        if response.status_code == 200:

            if "X-App-Rate-Limit-Count" in response.headers:
                RateLimitFromHeader(response.headers["X-App-Rate-Limit-Count"])
            matchList = response.json()['matches']
            success = True
        else:
            raise Exception("Invalid response code: {}".format(response.status_code))
    except Exception as e:
        logger.warning("Exception was thrown!!! --- {}".format(e))
        matchList = None
        success = False

    logger.info("Returning -- Success: {}".format(success))
    logger.info("          -- matchList: {}".format(matchList))
    return matchList, success

def fetchUsersInMatch(region,matchId):
    '''Attempts to retrieve match details
    Args:   region = NA1
            matchId = 38853692

    Returns: Dictionary of match details

    '''

    logger.info("Fetching summoners from match: {}-{}".format(region,matchId))
    url = RIOT_FETCH_MATCH_DETAILS_URL.format(region,matchId,RIOT_API_KEY)
    logger.info("GET Method: {}".format(url))

    response = requests.get(url)
    if "X-App-Rate-Limit-Count" in response.headers:
        RateLimitFromHeader(response.headers["X-App-Rate-Limit-Count"])
    users = response.json()['participantIdentities']

    logger.info("Returns -- users: {}".format(str(users)))
    return users

def insertSingleMatch(match,account_id):
    '''Attempts to insert a match
    Args: dictionary of a Match class

    Returns: a response message
    '''
    logger.info("Inserting single match into db")
    try:
        dbConnection = ActiveDatabase()
        matchConnection = dbConnection.getConn().collection(MATCHES)
        inMatchConnection = dbConnection.getConn().collection(IN_MATCH)
        logger.info("Checking if match exists")
        doesExist = inMatchConnection.where(u'match_id',u'==',match['matchId']).where(u'personal_summoner_id',u'==',account_id).stream()
        for doc in doesExist:
            logger.info("Match already exists")
            logger.info(doc.to_dict())

            logger.info("Returns -- Exists")
            return 'Exists'
        else:
            matchConnection.document(match['matchId']).set(match)
            logger.info("Match not found!!!")
            logger.info("Attempting to insert match")

            logger.info("Returns -- Success")
            return 'Success'
    except Exception as e:
        logger.warning("An error has occured: {}".format(e))

        logger.info("Returns -- Error")
        return 'Error'

def insertParticipant(participant,matchId,userAccountId):
    ''' Inserts a document into the IN_MATCH database
        Args:
            participant = the participant's accountId
            matchId = current matchId
            userAccountId = the active user's personal accountId

    '''
    logger.info("Inserting participant: participant: {} matchId: {} userAccountId {}".format(participant,matchId,userAccountId))
    try:
        insertDict = dict()
        insertDict['match_id'] = matchId
        insertDict['summoner_id'] = participant
        insertDict['personal_summoner_id'] = userAccountId
        dbConnection = ActiveDatabase()
        inMatchConnection = dbConnection.getConn().collection(IN_MATCH)
        inMatchConnection.add(insertDict)

        logger.info("Successfully added participant")
    except Exception as e:
        logger.warning("An error has occured: {}".format(e))

def main():
    logger.info("... Initializing MatchFinder ...")

    userList, userListSuccess = fetchActiveUsers()
    if userListSuccess:
        logger.info("Found users...")
        for user in userList:
            #Changing connection into dictionary
            user = user.to_dict()

            logger.info("Initializing run for user: {}".format(user))

            #Creating User Class out of data received
            user = User(user['region'],user['summoner_id'],user['account_id'])

            logger.info("Converted user into User class: {}".format(user.toDict()))

            matchList, matchListSuccess = fetchMatchEntries(user.toDict())
            if matchListSuccess:
                logger.info("Found matches...")
                for match in matchList:
                    match = Match(match['gameId'],match['timestamp'],user.getRegion())

                    logger.info("Initializing run for match: {}".format(match.toDict()))

                    insertMatchMessage = insertSingleMatch(match.toDict(),user.getAccountId())
                    if insertMatchMessage is 'Success':
                        usersInMatch = fetchUsersInMatch(match.getRegion(),match.getMatchId())
                        for participant in usersInMatch:
                            logger.info("Initializing participant: {}".format(participant['player']['accountId']))

                            insertParticipant(participant['player']['accountId'],match.getMatchId(),user.getAccountId())
                    elif insertMatchMessage is 'Exists':
                        logger.info("Match found that currently exists. Skipping user...")
                        break
                    elif insertMatchMessage is 'Error':
                        logger.info("Error has occured when inserting match. Skipping singular match...")
            else:
                logger.info("Failed to retrieve any active runs")
    else:
        logger.info("Failed to retrieve any active users")

if __name__ == "__main__":
    initLogger()
    main()


from retrying import retry
from topology_common.utils.threading import ThreadGroup
from time import sleep
from pytest import (
    mark,
    fixture
)
from topology_common.ops.vxlan.ft_vxlan_lib import (
    TrafficType,
    verify_mac_table
)

from topology_common.ops.vlans.vlan import configure_switch_vlan

from topology_common.ops.l3_common.ft_l3_lib import (
    AddressType
)

from fixtures import (
    env_setup_hs_sw_hs_hs,
    l2_basic_phy,
    l3_basic_phy
)

__doc__ = """
Test Case Header:
=================
:Author: Anvita Mishra anvita.mishra@hpe.com
:TestId: 00000
:Release: 10_04
:TestName: test_ft_VxLAN_ECMPConfigBehavior
:Objective: Verify VXLAN encapsulation and decapsulation across AVPs and NVPs
			with multiple ECMP Configurations

:Requirements:
    - 1 switch
    - 3 hosts
:TestDescription: This test-case verifies VXLAN encapsulation is working for
    different traffic types on AVP-->All direction and decapsulation for
    NVP-->All direction with multiple ECMP Configurations. This test-case also
    verifies ECMP load balancing behavior for Known Unicast Traffic

    1.  Setup the topology as shown and enable all interfaces.
    2.  Configure VXLAN on the switch
    3.  Set up NVP1 as first next hop.
    4.  Send encapsulatedpacket from NVP1 to all other hsots. Verify that
        NVP1 MAC is learned.
    5.  Test ECMP Load Balancing behavior by sending a list of packets where we
        permute the source MAC from AVP1 to all other hosts. Verify
        NVP1 (the designated next hop) receives all the traffic.
    6.  Add NVP2 as the second next hop.
    7.  Send list of packet permutations from AVP1 to all other hosts.
        Verify load-balancing for known unicast traffic and verify
        NVP1 (the designated next hop) receives all the traffic for all
        other traffic types.
    8.  Remove NVP1 as a next hop. Verify that NVP1 MAC is retained in the
        MAC table.
    9.  Send list of packet permutations from AVP1 to all other hosts.
        Verify load-balancing for known unicast traffic and verify
        NVP1 (the designated next hop) Verify
        NVP2 (the designated next hop) receives all the traffic.
:PlanPriority: 3 - High
:TestPassCriteria: -
    1. Setup of the topology successfully created.
    2. Packet successfully transmitted from NVP1 to all other hosts.
    3. Source MAC from packet sucessfully learned
    4. Packet permutations successfully transmitted from AVP1 to NVP1.
    5. NVP2 is added as a next hop
    6. ECMP load balancing behavior occurs as expected for each traffic type
    7. NVP2 is removed as a next hop
    8. Packet permutations successfully transmitted from AVP1 to NVP2.
:PlatformIndependent: Y
:SupportedPlatforms: 8325
Topology:
=========
.. ditaa::

+-----------+   +-----------+
|           |   |           |
|   h2      |   |   h3      |
|   NVP1    |   |   NVP2    |
|           |   |           |
+----+------+   +------+----+
     |                 |
+----+-----------------+-----+
|                            |
|           s1               |
|                            |
+-------------+--------------+
              |
         +----+------+
         |           |
         |   h1      |
         |   AVP1    |
         |           |
         +-----------+

"""

TOPOLOGY = """
# Nodes
[type=halon_0 name="OpenSwitch 1" target="true"] ops1
[type=host name="Host 1" image="ubuntuscapy_2.4:latest"] hs1
[type=host name="Host 2" image="ubuntuscapy_2.4:latest"] hs2
[type=host name="Host 3" image="ubuntuscapy_2.4:latest"] hs3

# Links
ops1:if01 -- hs1:eth1
ops1:if02 -- hs2:eth1
ops1:if03 -- hs3:eth1
"""

# Global Variables
AVP_MAC = ""
NVP_MAC = ""
VNI = 1760
UDP_DPORT = 4789
UDP_SPORT = 1234
SRC_IP = ""
DST_IP = ""
OUTER_SRC_MAC = ""
OUTER_DST_MAC = ""
PKT_TYPE = ''

# Global fixed variables
STP_DEST_MAC = "01:80:c2:00:00:00"
BROADCAST_MAC = "ff:ff:ff:ff:ff:ff"
MULTICAST_MAC = "01:00:00:00:00:bb"
LLDP_MAC = "01:80:c2:00:00:0e"
AVP1_MAC = "00:00:00:00:00:11"
NVP1_MAC = "00:00:00:00:00:aa"
SW_MAC = ""

# NVP IPs and IPs
# IPs  and MACs only needed for NVPs h2 and h3
HOST2_MAC = ""
HOST3_MAC = ""
H2_IP = '30.0.0.100'
H3_IP = '40.0.0.100'

VTEP_PEER_IP = '20.0.0.2'

# Switch SVI IPs
S1_SVI_IP_1 = '30.0.0.1'
S1_SVI_IP_2 = '40.0.0.1'

# Switch Loopback IP
# Used for VxLAN tunnel Source IP
S1_LO_IP = '1.1.1.1'

MASK = '24'
LO_MASK = '24'
DEST_MASK = '32'

VLAN1 = '10'
VLAN2 = '30'
VLAN3 = '40'
VNI_HEX = '0x6e0'
TUN_NUM = '1'

PKT_COUNT = 5
SNIFF_PKT_COUNT = 10

# Number of MACs to send for ECMP Loading
TOTAL_MACS = 50

# This timeout parameter is for sending 5 packets
TIMEOUT = 5

# This timeout parameter is for sending 50 packets
LONG_TIMEOUT = 15

# Global dict to store packet count
# Updated in parse_packet function
RECEIVED_PKT_DICT = {}

@fixture(scope='function')
def config(env_setup_hs_sw_hs_hs, request, testlog):
    '''
    Fixture to get the global fixture for this topology and run the teardown

    :param env_setup_hs_sw_hs: global fixture for the topology
    :param request: request param from pytest fixture
    :param testlog: testlog fixture
    '''
    yield env_setup_hs_sw_hs_hs

    testlog.log_step("Teardown")
    env_setup_hs_sw_hs_hs.teardown()

@fixture(scope='function')
def get_switch_mac(topology):
    '''
    Helper function to get mac address from the switch
    '''
    global SW_MAC
    s = topology.get('ops1')
    result = s.libs.vtysh.show_system()
    SW_MAC = result['base_mac']
    yield

@fixture(scope='function')
def get_host_macs(topology):
    '''
    Helper function to get mac addresses for NVP hosts h2 and h3
    '''
    global HOST2_MAC, HOST3_MAC

    h2 = topology.get('hs2')
    h3 = topology.get('hs3')

    h2_port = h2.ports['eth1']
    h3_port = h3.ports['eth1']

    ifconfig = h2.libs.ip.get_ifconfig(h2_port)
    HOST2_MAC = ifconfig[h2_port]['hwaddr']

    ifconfig = h3.libs.ip.get_ifconfig(h3_port)
    HOST3_MAC = ifconfig[h3_port]['hwaddr']
    yield

@fixture(scope='function',
         ids=['l2_basic', 'stp', 'broadcast', 'multicast', 'lldp'],
         params=[TrafficType.l2_basic, TrafficType.stp,
                 TrafficType.broadcast, TrafficType.multicast,
                 TrafficType.lldp])
def traffic_permutation(request):
    """
    Set the global variables for traffic
    """
    global AVP_MAC, NVP_MAC, UDP_SPORT, SRC_IP, DST_IP, OUTER_SRC_MAC,\
        OUTER_DST_MAC, PKT_TYPE

    # Values common to all packet types
    SRC_IP = VTEP_PEER_IP
    DST_IP = S1_LO_IP
    OUTER_SRC_MAC = HOST2_MAC
    OUTER_DST_MAC = SW_MAC
    AVP_MAC = AVP1_MAC

    # NVP_MAC updated based on pkt_type
    NVP_MAC_dict = {TrafficType.l2_basic: NVP1_MAC,
                    TrafficType.stp: STP_DEST_MAC,
                    TrafficType.broadcast: BROADCAST_MAC,
                    TrafficType.multicast: MULTICAST_MAC,
                    TrafficType.lldp: LLDP_MAC}

    PKT_TYPE_dict = {TrafficType.l2_basic: 'l2_basic',
                     TrafficType.stp: 'stp',
                     TrafficType.broadcast: 'broadcast',
                     TrafficType.multicast: 'multicast',
                     TrafficType.lldp: 'lldp'}

    NVP_MAC = NVP_MAC_dict[request.param]
    PKT_TYPE = PKT_TYPE_dict[request.param]

    yield request.param

def generate_params(sw):
    """
    Return arguments required for every single permutation

    :param sw: Topology switch object
    :param permutation: Main pytest fixture for this test
    """

    # List of vlan ids to use for this permutation
    vlan_ids = []
    # Physical ports  required for this permutation per L3 interface
    phy_ports = []
    # L3 interfaces to be created
    l3_interfaces = 0
    # List of ip address for every host
    ip_address_hs = []
    # VxLAN interfaces to be created
    vxlan_ids = []
    # VNIs to be created
    vnis = {}
    # VTEP Peers to be created
    vtep_peers = []

    vlan_ids = [VLAN1, VLAN2, VLAN3]
    vxlan_ids = [TUN_NUM]
    vnis = {VNI: {'vlan': [VLAN1], 'vtep_peer': [VTEP_PEER_IP]}}
    l3_interfaces = 1
    phy_ports = [sw.vtysh_ports['if01'], sw.vtysh_ports['if02'],
                 sw.vtysh_ports['if03']]
    ip_address_hs = [H2_IP, H3_IP]

    return {'vlan_ids': vlan_ids,
            'vxlan_ids': vxlan_ids,
            'vnis': vnis,
            'l3_interfaces': l3_interfaces,
            'phy_ports': phy_ports,
            'ip_address_hs': ip_address_hs}

def configure_host_ips(h2, h3, ip_address_hs):
    """
    Configure host IPs for the NVPs (h2 and h3)
    AVPs do not need IPs assigned for test

    :param h2: A Topology host object
    :param h3: A Topology host object
    :param ip_address_hs: A list of ip addresses to configure h2 and h3
    """

    h2.libs.ip.flush_ip('eth1')
    h2.libs.ip.interface('eth1', up=False)

    h3.libs.ip.flush_ip('eth1')
    h3.libs.ip.interface('eth1', up=False)

    h2.libs.ip.interface(portlbl='eth1', addr="{}/{}".format(
        ip_address_hs[0], MASK), up=True)
    h3.libs.ip.interface(portlbl='eth1', addr="{}/{}".format(
        ip_address_hs[1], MASK), up=True)

def start_stop_scapy_on_hosts(hs, action='start'):
    """
    Function to start and stop scapy daemon on hosts

    :param hs: A Topology host object
    :param action: By default, the value is 'start'
            If set to any other value, scapy will exit
    """
    if action == 'start':
        hs.libs.scapy.start_scapy()
    else:
        hs.libs.scapy.exit_scapy()

def scapy_configuration_avp(hs, pkt_type=None, macsrc=None, macdst=None):
    """
    Function to configure scapy for packet in AVP->All direction

    :param hs: A Topology host object
    :param pkt_type: l2_basic, stp, broadcast, multicast, or lldp
    :param macsrc: Source MAC address
    :param macdst: Destination MAC address
    """
    # L2 header
    eth_header = hs.libs.scapy.ether()

    eth_header['src'] = '{}'.format(macsrc)
    eth_header['dst'] = '{}'.format(macdst)

    if isinstance(macsrc, list):
        eth_header['src'] = macsrc
    packet_list = []
    packet_struct = ""

    # If pkt_type is stp, need additional STP and LLC headers
    if pkt_type == 'stp':
        # Raw data
        stp_header = hs.libs.scapy.stp()
        llc_header = hs.libs.scapy.llc()
        packet_list = [eth_header, llc_header, stp_header]
        packet_struct = "Eth/LLC()/STP()"

    # For all other packet types, packet structure is same
    else:
        # Update raw_data with pkt_type
        raw_data = {'data': "{} packet from AVP end".format(pkt_type)}
        packet_list = [eth_header, raw_data]
        packet_struct = "Eth/data"

    packet = [packet_struct, packet_list]
    return packet

def scapy_configuration_nvp(hs, pkt_type=None,
                            inner_macsrc=None, inner_macdst=None, vni=None,
                            src_port=None, dst_port=None, src_IP=None,
                            dst_IP=None, outer_macsrc=None, outer_macdst=None):
    """
    Function to configure scapy for packet in NVP->All direction

    :param hs: A Topology host object
    :param pkt_type: l2_basic, stp, broadcast, multicast, or lldp
    :param inner_macsrc: Encapsulated packet source MAC address
    :param inner_macdst: Encapsulated packet destination MAC address
    :param vni: VNI value for VxLAN header
    :param src_port: UDP source port
    :param dst_port: UDP destination port
    :param src_IP: Source VTEP's IP
    :param dst_IP: Destination VTEP's IP
    :param inner_macsrc: Source VTEP's MAC address
    :param inner_macdst: Next hop MAC address
    """
    # Inner L2 header
    inner_eth_header = hs.libs.scapy.ether()
    inner_eth_header['prot'] = 'ENCAP_ETH'
    inner_eth_header['src'] = '{}'.format(inner_macsrc)
    inner_eth_header['dst'] = '{}'.format(inner_macdst)

    # VxLAN header
    vxlan_header = hs.libs.scapy.vxlan()
    vxlan_header['vni'] = vni
    vxlan_header['flags'] = 'Instance'
    # UDP header
    udp_header = hs.libs.scapy.udp()
    udp_header['sport'] = src_port
    udp_header['dport'] = dst_port
    # IP header
    ip_header = hs.libs.scapy.ip()
    ip_header['src'] = '{}'.format(src_IP)
    ip_header['dst'] = '{}'.format(dst_IP)
    # Outer L2 header
    outer_eth_header = hs.libs.scapy.ether()
    outer_eth_header['src'] = '{}'.format(outer_macsrc)
    outer_eth_header['dst'] = '{}'.format(outer_macdst)
    outer_eth_header['type'] = 0x0800
    # Inner L2 header
    inner_eth_header['src'] = '{}'.format(inner_macsrc)
    inner_eth_header['dst'] = '{}'.format(inner_macdst)

    # If pkt_type is stp, need additional STP and LLC headers
    if pkt_type == 'stp':
        stp_header = hs.libs.scapy.stp()
        llc_header = hs.libs.scapy.llc()
        packet_list = [outer_eth_header, ip_header, udp_header,
                       vxlan_header, inner_eth_header, llc_header, stp_header]
        packet_struct = "Eth/IP/UDP/VXLAN/ENCAP_ETH/LLC()/STP()"

    # For all other packet types, packet structure is same
    else:
        # Update raw_data with pkt_type
        raw_data = {'data': "{} packet from NVP end".format(pkt_type)}
        packet_list = [outer_eth_header, ip_header, udp_header,
                       vxlan_header, inner_eth_header, raw_data]
        packet_struct = "Eth/IP/UDP/VXLAN/ENCAP_ETH/data"

    packet = [packet_struct, packet_list]
    return packet

def start_traffic(hs, packet, count):
    """
    Start Traffic
    :param hs: Host hs.
    :param packet: Custom header to send.
    :param count: Number of packets to send
    """

    # Waiting 2 sec to give some time to sniffer to be ready
    sleep(2)
    iface = hs.ports['eth1']
    # Start Traffic
    hs.libs.scapy.sendp(packet[0], packet[1],
                        "iface='{}', count = {}, inter=0.2".format(iface, count))

def check_packet(packet, recipient_type, pkt_type, exp_src, exp_dst):
    """
    Helper function to determine how many packets are correctly received.
    Return packet_count

    :param packet: Dictionary of sniffed packets.
    :param recipient_type: String describing if host is an AVP or NVP
    :param pkt_type: l2_basic, stp, broadcast, multicast, or lldp
    :param exp_src: String representing expected source MAC of packet
    :param exp_dst: String representing expected destination MAC of packet
    """
    packet_count = 0

    # If recipient is AVP, check for simple packet
    if ("AVP" in recipient_type):
        for pkt in packet.values():
            # Boolean to store if packet is correct. Set to False initially
            correctly_received = False

            # Check 1: Check src and dst MACs match expected
            if ((pkt['Ethernet']['src'] == exp_src) and
                    (pkt['Ethernet']['dst'] == exp_dst)):

                # Check 2: If packet type is LLDP, check raw string for "lldp"
                if (pkt_type == "lldp"):
                    # If "lldp" is present, set correctly_received to True
                    if ("lldp" in pkt['Ethernet']['Raw']['load']):
                        correctly_received = True

                # Packet type is not LLDP, so only Check 1 matters
                # Set correctly_received to True
                else:
                    correctly_received = True

            # Update total packet count with value of correctly_received
            packet_count += correctly_received

    # If recipient is NVP, then check for VXLAN encapsulated packet
    elif ("NVP" in recipient_type):
        for pkt in packet.values():
            # Boolean to store if packet is correct. Set to False initially
            correctly_received = False

            # Check 1: Check if VxLAN fields are correct
            if ((pkt['Ethernet']['IP']['src'] == S1_LO_IP) and
                (pkt['Ethernet']['IP']['dst'] == VTEP_PEER_IP) and
                (pkt['Ethernet']['IP']['UDP']['dport'] == str(UDP_DPORT)) and
                (pkt['Ethernet']['IP']['UDP']['VXLAN']['vni'] == VNI_HEX) and
                (pkt['Ethernet']['IP']['UDP']['VXLAN']['Ethernet']['src'] in exp_src) and
                    (pkt['Ethernet']['IP']['UDP']['VXLAN']['Ethernet']['dst'] == exp_dst)):

                # Check 2: If packet type is LLDP, check raw string for "lldp"
                if (pkt_type == "lldp"):
                    # If "lldp" is present, set correctly_received to True
                    if ("lldp" in pkt['Ethernet']['IP']['UDP']['VXLAN']['Ethernet']['Raw']['load']):
                        correctly_received = True
                # Packet type is not LLDP, so only Check 1 matters
                # Set correctly_received to True
                else:
                    correctly_received = True

            # Update total packet count with value of correctly_received
            packet_count += correctly_received

    return packet_count

def parse_packet(packet, recipient_type, pkt_type, exp_src, exp_dst, step):
    """
    Function to parse received packets and check if the packet values
    match expected values

    :param packet: Dictionary of sniffed packets.
    :param recipient_type: String describing if host is an AVP or NVP
    :param pkt_type: l2_basic, stp, broadcast, multicast, or lldp
    :param expected_pkts: Boolean stating if we expect packets
    :param exp_src: Expected source MAC of packet
    :param exp_dst: Expected destination MAC of packet
    :param step: Step fixture
    """

    packet_count = 0
    expected_packet_count = PKT_COUNT

    packet_count = check_packet(
        packet, recipient_type, pkt_type, exp_src, exp_dst)

    # Print packet count at each host
    print("Packet count at {} was {}".format(recipient_type, packet_count))

    # Store packet_counts in global dict RECEIVED_PKT_DICT
    # Dict needed for ECMP check
    RECEIVED_PKT_DICT[recipient_type] = packet_count

def sniff_traffic(hs, count, timeout, recipient_type, pkt_type,
                  exp_src, exp_dst, testlog):
    """
    Sniff traffic on given interface
    :param hs: Host hs.
    :param count: Count of packets to be sniffed
    :param timeout: timeout value
    :param recipient_type: AVP or NVP
    :param pkt_type: l2_basic, stp, broadcast, multicast, or lldp
    :param step: Step fixture
    """
    iface = hs.ports['eth1']

    # If host is NVP, sniff using a filter that checks for UDP packets
    if ("NVP" in recipient_type):
        packet = hs.libs.scapy.sniff2("iface='{}', count={}, timeout={}, "
                                      " filter='port 4789 and (!icmp or !ip6)', "
                                      " prn=lambda x:x.show()".format(
                                          iface, count, timeout), True)
        parse_packet(packet, recipient_type, pkt_type,
                     exp_src, exp_dst, testlog)
    # If host is AVP, sniff using a filter that checks for Ethernet packets
    elif ("AVP" in recipient_type):
        packet = hs.libs.scapy.sniff2("iface='{}', count={}, "
                                      " timeout={}, filter='!icmp or !ip6', "
                                      " prn=lambda x:x.show()".format(
                                          iface, count, timeout), True)
        parse_packet(packet, recipient_type, pkt_type,
                     exp_src, exp_dst, testlog)

def print_transmission_results(pkts_expected_dict,
                               expected_pkt_count, testlog):
    """
    Helper function to display results of packet transmission

    :param pkts_expected_dict: Dictionary of Booleans representing if packets
          are expected at each host (Keys: Host Names, Values: True/False)
    :param expected_pkt_count: Number of packets expected
    :param testlog: Testlog fixture
    """

    # Global Dictionary representing packet count received at each host
    # (Keys: Host Names, Values: Packet Count)
    global RECEIVED_PKT_DICT

    testlog.log_subheader("Results:")

    # If packets are expected at AVP1 and packets were received
    if (pkts_expected_dict["AVP1"]) and ("AVP1 (h1)" in RECEIVED_PKT_DICT.keys()):
        # Check received packet count matches expected packet count
        assert ((RECEIVED_PKT_DICT["AVP1 (h1)"] == expected_pkt_count)),\
            "Packets expected at AVP1, but incorrectly received. " + \
            "Check packet count at the AVP1"
        testlog.log_info("All packets correctly received at AVP1")

    # If packets are expected at ECMP group
    if (pkts_expected_dict["ECMP group (NVP1/NVP2)"]):
        check_load_balancing(testlog)

    # If packets are expected at NVP1 and packets were received
    if (pkts_expected_dict["NVP1"]) and ("NVP1 (h2)" in RECEIVED_PKT_DICT.keys()):
        # Check received packet count matches expected packet count
        assert ((RECEIVED_PKT_DICT["NVP1 (h2)"] == expected_pkt_count)),\
            "Packets expected at NVP1, but incorrectly received. " + \
            "Check packet count at the NVP1"
        testlog.log_info("All packets correctly received at NVP1")

    # If packets are expected at NVP1 and packets were received
    if (pkts_expected_dict["NVP2"]) and ("NVP2 (h3)" in RECEIVED_PKT_DICT.keys()):
        # Check received packet count matches expected packet count
        assert ((RECEIVED_PKT_DICT["NVP2 (h3)"] == expected_pkt_count)),\
            "Packets expected at NVP2, but incorrectly received. " + \
            "Check packet count at the NVP2"
        testlog.log_info("All packets correctly received at NVP2")

def check_load_balancing(testlog):
    """
    Helper function to parse RECEIVED_PKT_DICT and report on
    if NVP1 or NVP2 is chosen

    :param testlog: testlog fixture
    """

    # Global Dictionary representing packet count received at each host
    # (Keys: Host Names, Values: Packet Count)
    global RECEIVED_PKT_DICT

    NVP1_count = RECEIVED_PKT_DICT["NVP1 (h2)"]
    NVP2_count = RECEIVED_PKT_DICT["NVP2 (h3)"]
    NVP1_percent = (NVP1_count/TOTAL_MACS)*100
    NVP2_percent = (NVP2_count/TOTAL_MACS)*100

    print("NVP1 was chosen {} out of {} times, which is {:.2f}% of the time\n".format(
        NVP1_count, TOTAL_MACS, NVP1_percent))
    print("NVP2 was chosen {} out of {} times, which is {:.2f}% of the time\n".format(
        NVP2_count, TOTAL_MACS, NVP2_percent))

    # Check ECMP Load Balancing if expected
    if (PKT_TYPE == "l2_basic"):
        # Check if at least 40% of total packets were received at one host
        # Check that total NVP1_count and NVP2_count equals the number of packets
        assert((NVP1_percent >= 0.4) or (NVP2_percent >= 0.4) and
               ((NVP1_count + NVP2_count) == TOTAL_MACS)),\
            "ECMP Load Balancing Failed for Known Unicast Traffic"

        step_desc = "All packets correctly received for Known Unicast traffic. " +\
            "Packets were load balanced betweeen both Next Hops NVP1 and NVP2"
        testlog.log_info(step_desc)
    else:
        # Check that all of the packets are only received at one NVP and
        # zero packets are received at the other NVP
        assert(((NVP1_count == 0) and (NVP2_count == TOTAL_MACS)) or
               (NVP2_count == 0) and (NVP1_count == TOTAL_MACS)), \
            "Load Balancing was not expected, but packets were not" +\
            " received correctly for {} traffic".format(PKT_TYPE)
        chosen_next_hop = "#1 (NVP1)" if (
            NVP1_count >= NVP2_count) else "#2 (NVP2)"
        step_desc = "All packets correctly received for {} traffic. ".format(PKT_TYPE) + \
                    "All traffic was directed to the Next Hop {}".format(
                        chosen_next_hop)
        testlog.log_info(step_desc)

@mark.timeout(1000)
def test_ft_VxLAN_ECMPConfigBehavior(topology,
                                     testlog,
                                     config,
                                     get_switch_mac,
                                     get_host_macs,
                                     traffic_permutation,
                                     l2_basic_phy,
                                     l3_basic_phy):
    """
    Test sends different types of traffic between 3 hosts connected
    through VxLAN tunnel.

    Build a topology of one switch and four hosts and connect
    the hosts to the switches.
    Setup a VxLAN between switches and VLANs between ports
    and switches.
    """

    """
    Step: Create the topology as shown
    Result: Topology created successfully
    """
    testlog.log_step("Set up the switch and hosts")
    s1 = topology.get('ops1')
    h1 = topology.get('hs1')
    h2 = topology.get('hs2')
    h3 = topology.get('hs3')

    argv = generate_params(s1)

    """
    Step: Configure the topology as shown
    Result: Topology has been configured successfully
    """
    # Configuring VLAN and four interfaces on switch 1
    testlog.log_subheader("Creating L2 Interfaces")
    config.create_l2_interface_type(s1,
                                    l2_basic_phy['l2_iface'],
                                    False,
                                    [argv['phy_ports'][0]],
                                    argv['vlan_ids'][0],
                                    None, testlog)

    testlog.log_subheader("Creating SVIs")
    # Creating SVI interfaces of if02 and if03
    config.create_l2_interface_type(s1,
                                    l2_basic_phy['l2_iface'],
                                    False,
                                    [argv['phy_ports'][1]],
                                    argv['vlan_ids'][1],
                                    None, testlog)

    config.create_l2_interface_type(s1,
                                    l2_basic_phy['l2_iface'],
                                    False,
                                    [argv['phy_ports'][2]],
                                    argv['vlan_ids'][2],
                                    None, testlog)

    configure_switch_vlan(s1,
                          argv['vlan_ids'][1],
                          S1_SVI_IP_1 + '/' + MASK,
                          [argv['phy_ports'][1]])

    configure_switch_vlan(s1,
                          argv['vlan_ids'][2],
                          S1_SVI_IP_2 + '/' + MASK,
                          [argv['phy_ports'][2]])

    testlog.log_subheader("Creating Switch Loopback Interface")
    # Creating a loopback interface on s1
    config.create_l3_loopback(s1,
                              '0',
                              S1_LO_IP + '/' + LO_MASK)

    testlog.log_subheader("Configure Host IPs")
    configure_host_ips(h2, h3, argv['ip_address_hs'])

    """
    Step: Configure VxLAN interface on switch 1
    Result: VxLAN interface has been configured successfully
    """
    testlog.log_step("Configure VXLAN")
    config.configure_vxlan(s1,
                           TUN_NUM,
                           S1_LO_IP,
                           argv['vnis'])
    testlog.log_subheader("Current Running Configuration")
    s1.libs.vtysh.show_running_config()

    """
  	Step:   Adding Next Hop #1 (NVP1)
    Result: IP Route from NVP1 to VTEP Peer is configured successfully
    """
    testlog.log_step("Adding Next Hop #1 (NVP1)")
    testlog.log_subheader(
        "Adding an IP Route from NVP1 to VTEP Peer (20.0.0.2)")
    config.switch_add_route(s1, VTEP_PEER_IP, DEST_MASK, H2_IP)
    testlog.log_subheader("Current Running Configuration")
    s1.libs.vtysh.show_running_config()

    testlog.log_step("Start Scapy on hosts")
    # Starting scapy
    start_stop_scapy_on_hosts(h1)
    start_stop_scapy_on_hosts(h2)
    start_stop_scapy_on_hosts(h3)

    # Update src MAC and dst MAC based on packet type
    dst_MAC_dict = {"l2_basic": AVP1_MAC, "stp": STP_DEST_MAC,
                    "broadcast": BROADCAST_MAC, "multicast": MULTICAST_MAC,
                    "lldp": LLDP_MAC}

    cur_src_MAC = NVP1_MAC
    cur_dst_MAC = dst_MAC_dict[PKT_TYPE]

    """
    Step:   Use scapy to configure the traffic for NVP to all hosts
    Result: Packet created
    """
    step_desc = "Configure traffic at NVP1"
    testlog.log_step(step_desc)

    # Configure Scapy on hs2 for NVP-->All traffic
    packet_nvp = scapy_configuration_nvp(h2, pkt_type=PKT_TYPE,
                                         inner_macsrc=cur_src_MAC,
                                         inner_macdst=cur_dst_MAC, vni=VNI,
                                         src_port=UDP_SPORT, dst_port=UDP_DPORT,
                                         src_IP=SRC_IP, dst_IP=DST_IP,
                                         outer_macsrc=OUTER_SRC_MAC,
                                         outer_macdst=OUTER_DST_MAC)

    """
    Step:   Transmit the traffic from NVP to all hosts and
            parse the captured packet to check
            if the packet is received as expected
    Result: Traffic transmitted successfully and
            Correct values received
    """
    # Updating step_desc based on traffic type
    if (PKT_TYPE == "l2_basic"):
        step_desc = "Transmitting from NVP1 (Packet Type: Unknown Unicast)\n" + \
                    "\t     Expected Result: NVP1 MAC is learned"
    else:
        step_desc = "Transmitting from NVP1 (Packet Type: {})\n".format(PKT_TYPE) + \
                    "\t     Expected Result:  NVP1 MAC is learned"
    testlog.log_step(step_desc)

    testlog.log_subheader("Sending Packet from NVP1")
    # Sending traffic from NVP1 h2 to all other hosts
    with ThreadGroup() as ctx:
        ctx.run(start_traffic, hs=h2, packet=packet_nvp, count=PKT_COUNT)

    # Verifying if NVP1_MAC was learned on interface vxlan 1
    testlog.log_subheader("Verifying NVP1 MAC was learned")
    verify_mac = retry(
        stop_max_attempt_number=3,
        wait_fixed=1000)(verify_mac_table)
    verify_mac(sw1=s1, mac=cur_src_MAC, exp_port="vxlan1",
               exp_vlan=VLAN1, vtep_peer=VTEP_PEER_IP)

    """
    Step:   Transmit the traffic from AVP to all hosts and
            parse the captured packet to check
            if the packet is received as expected
    Result: Traffic transmitted successfully and
            Correct values received
    """

    # Updating step_desc based on traffic type
    if (PKT_TYPE == "l2_basic"):
        step_desc = "Transmitting from AVP1. (Packet Type: Known Unicast)\n" + \
                    "\t     Permuting Source MAC to Test ECMP Load Balancing\n" + \
                    "\t     Expected Result: All Packets Transmited to Next Hop #1 (NVP1)"
    else:
        step_desc = "Transmitting from AVP1. (Packet Type: {})\n".format(PKT_TYPE) + \
                    "\t     Permuting Source MAC to Test ECMP Load Balancing\n" + \
                    "\t     Expected Result: All Packets Transmited to Next Hop #1 (NVP1)"
    testlog.log_step(step_desc)

    # Dictionary of Booleans representing if packets are expected at each host
    # (Keys: Host Names, Values: True/False)
    pkts_expected_dict = {"AVP1": False, "NVP1": True, "NVP2": False,
                          "ECMP group (NVP1/NVP2)": False}

    # Creating a list of potential source MACs
    lst_source_macs = ['00:00:00:00:00:{}'.format(
        i) for i in range(21, 21+TOTAL_MACS)]

    # Update src MAC and dst MAC based on packet type
    dst_MAC_dict = {"l2_basic": AVP1_MAC, "stp": STP_DEST_MAC,
                    "broadcast": BROADCAST_MAC, "multicast": MULTICAST_MAC,
                    "lldp": LLDP_MAC}

    cur_src_MAC = lst_source_macs
    cur_dst_MAC = NVP_MAC

    packet_avp_lst = scapy_configuration_avp(h1, pkt_type=PKT_TYPE,
                                             macsrc=cur_src_MAC, macdst=cur_dst_MAC)

    with ThreadGroup() as ctx:
        # Sniffing on hs2 (NVP1)
        ctx.run(sniff_traffic, hs=h2, count=TOTAL_MACS+10, timeout=LONG_TIMEOUT,
                recipient_type="NVP1 (h2)", pkt_type=PKT_TYPE,
                exp_src=cur_src_MAC, exp_dst=cur_dst_MAC, testlog=testlog)
        # Sniffing on hs3 (NVP2)
        ctx.run(sniff_traffic, hs=h3, count=TOTAL_MACS+10, timeout=LONG_TIMEOUT,
                recipient_type="NVP2 (h3)", pkt_type=PKT_TYPE,
                exp_src=cur_src_MAC, exp_dst=cur_dst_MAC, testlog=testlog)
        ctx.run(start_traffic, hs=h1, packet=packet_avp_lst, count=1)

    # Checking Transmission Results
    print_transmission_results(pkts_expected_dict, TOTAL_MACS, testlog)
    print(RECEIVED_PKT_DICT)

    RECEIVED_PKT_DICT.clear()

    """
    Step:   Adding Another Next Hop
    Result: IP Route from NVP2 to VTEP Peer is configured successfully
    """
    testlog.log_step("Adding Next Hop #2 (NVP2)")
    testlog.log_subheader(
        "Adding an IP Route from NVP2 to VTEP Peer (20.0.0.2)")
    config.switch_add_route(s1, VTEP_PEER_IP, DEST_MASK, H3_IP)
    testlog.log_subheader("Current Running Configuration")
    s1.libs.vtysh.show_running_config()

    # Updating step_desc based on traffic type
    if (PKT_TYPE == "l2_basic"):
        step_desc = "Transmitting from AVP1. (Packet Type: Known Unicast)\n" + \
                    "\t     Permuting Source MAC to Test ECMP Load Balancing\n" + \
                    "\t     Expected Result: Packets are Load Balanced between both Next Hops"
    else:
        step_desc = "Transmitting from AVP1. (Packet Type: {})\n".format(PKT_TYPE) + \
                    "\t     Permuting Source MAC to Test ECMP Load Balancing\n" + \
                    "\t     Expected Result: All Packets Transmited to Next Hop #1 (NVP1)"
    testlog.log_step(step_desc)

    # Dictionary of Booleans representing if packets are expected at each host
    # (Keys: Host Names, Values: True/False)
    pkts_expected_dict = {"AVP1": False, "NVP1": False, "NVP2": False,
                          "ECMP group (NVP1/NVP2)": True}

    with ThreadGroup() as ctx:
        # Sniffing on hs2 (NVP1)
        ctx.run(sniff_traffic, hs=h2, count=TOTAL_MACS+10, timeout=LONG_TIMEOUT,
                recipient_type="NVP1 (h2)", pkt_type=PKT_TYPE,
                exp_src=cur_src_MAC, exp_dst=cur_dst_MAC, testlog=testlog)
        # Sniffing on hs3 (NVP2)
        ctx.run(sniff_traffic, hs=h3, count=TOTAL_MACS+10, timeout=LONG_TIMEOUT,
                recipient_type="NVP2 (h3)", pkt_type=PKT_TYPE,
                exp_src=cur_src_MAC, exp_dst=cur_dst_MAC, testlog=testlog)
        ctx.run(start_traffic, hs=h1, packet=packet_avp_lst, count=1)

    # Checking Transmission Results
    print_transmission_results(pkts_expected_dict, TOTAL_MACS, testlog)
    print(RECEIVED_PKT_DICT)

    RECEIVED_PKT_DICT.clear()

    step_desc = "Removing Next Hop #1 (NVP1)"
    testlog.log_step(step_desc)
    testlog.log_subheader("Removing an IP Route from NVP1 to VTEP Peer")
    config.switch_remove_route(s1, VTEP_PEER_IP, DEST_MASK, H2_IP)
    testlog.log_subheader("Current Running Configuration")
    s1.libs.vtysh.show_running_config()

    testlog.log_subheader("Verifying NVP1 MAC is retained in MAC Table")
    verify_mac = retry(
        stop_max_attempt_number=10,
        wait_fixed=1000)(verify_mac_table)
    verify_mac(sw1=s1, mac=NVP1_MAC, exp_port="vxlan1",
               exp_vlan=VLAN1, vtep_peer=VTEP_PEER_IP)

    # Updating step_desc based on traffic type
    if (PKT_TYPE == "l2_basic"):
        step_desc = "Transmitting from AVP1. (Packet Type: Known Unicast)\n" + \
                    "\t     Permuting Source MAC to Test ECMP Load Balancing\n" + \
                    "\t     Expected Result: All Packets Transmited to Next Hop #2 (NVP2)"
    else:
        step_desc = "Transmitting from AVP1. (Packet Type: {})\n".format(PKT_TYPE) + \
                    "\t     Permuting Source MAC to Test ECMP Load Balancing\n" + \
                    "\t     Expected Result: All Packets Transmited to Next Hop #2 (NVP2)"
    testlog.log_step(step_desc)

    # Dictionary of Booleans representing if packets are expected at each host
    # (Keys: Host Names, Values: True/False)
    pkts_expected_dict = {"AVP1": False, "NVP1": False, "NVP2": True,
                          "ECMP group (NVP1/NVP2)": False}

    with ThreadGroup() as ctx:
        # Sniffing on hs2 (NVP1)
        ctx.run(sniff_traffic, hs=h2, count=TOTAL_MACS+10, timeout=LONG_TIMEOUT,
                recipient_type="NVP1 (h2)", pkt_type=PKT_TYPE,
                exp_src=cur_src_MAC, exp_dst=cur_dst_MAC, testlog=testlog)
        # Sniffing on hs3 (NVP2)
        ctx.run(sniff_traffic, hs=h3, count=TOTAL_MACS+10, timeout=LONG_TIMEOUT,
                recipient_type="NVP2 (h3)", pkt_type=PKT_TYPE,
                exp_src=cur_src_MAC, exp_dst=cur_dst_MAC, testlog=testlog)
        ctx.run(start_traffic, hs=h1, packet=packet_avp_lst, count=1)

    # Checking Transmission Results
    print_transmission_results(pkts_expected_dict, TOTAL_MACS, testlog)
    print(RECEIVED_PKT_DICT)

    RECEIVED_PKT_DICT.clear()

    testlog.log_step("Stop Scapy on hosts")
    # Stopping scapy
    start_stop_scapy_on_hosts(h1, action='stop')
    start_stop_scapy_on_hosts(h2, action='stop')
    start_stop_scapy_on_hosts(h3, action='stop')


#! /usr/bin/env python
# encoding: utf-8
# hhaim, 2014 (IL) base on WAF book

"""
call 'waf --targets=waf.pdf' or use 'waf list' to see the targets available
"""

VERSION='0.0.1'
APPNAME='wafdocs'

import os, re, shutil
import sys
import shlex
import subprocess
import json
import re
import tempfile
from waflib import Logs
import string

top = '.'
out = 'build'
sphinx_version = None

try:
    from HTMLParser import HTMLParser
except:
    from html.parser import HTMLParser

# HTML injector for TOC and DISQUS
from HTMLInjector import HTMLInjector

class CTocNode:
    def __init__ (self, link_fmt = None):
        self.name="root"
        self.level=1; # 1,2,3,4
        self.link=None;
        self.parent=None
        self.childs=[]; # link to CTocNode

        # by default, link_fmt will formatted as in-page links
        self.link_fmt = link_fmt if link_fmt else lambda x: '#' + x

    def get_link (self):
        if self.link==None:
            name=self.name
            l=name.split('.');
            l=l[-1].lower()
            s='';
            for c in l:
                if c.isalpha() or c.isspace():
                    s+=c

            link = '_'.join(s.lower().split());
        else:
            link = self.link

        return self.link_fmt(link)

    def add_new_child (self,name,level,link):
        n=CTocNode();
        n.name=name
        n.link=link
        n.level=level
        n.parent=self
        n.link_fmt=n.parent.link_fmt
        self.childs.append(n);
        return n

    def to_json_childs (self):
        l=[]
        for obj in self.childs:
            l.append(obj.to_json());
        return (l);

    def to_open (self):
        if self.level <3:
            return True
        else:
            return False

    def to_json (self):
        d={"text" : self.name,
           "link" : self.get_link(),
           "state"       : {
                "opened"    : self.to_open()
                }
          }
        if len(self.childs)>0 :
            d["children"]= self.to_json_childs()
        return d

class TocHTMLParser(HTMLParser):

    def __init__ (self, link_fmt = None):
        HTMLParser.__init__(self);
        self.state=0
        self.root = CTocNode(link_fmt)
        self.root.parent=self.root
        self.level=2;
        self.attrs=None
        self.d={};
        self.last_level=1
        self.set_level(1,self.root)

    def set_level (self,level,node):
        assert(node!=None);
        assert(isinstance(node,CTocNode)==True);
        self.d[str(level)]=node

        # in case we change from high to low level remove the higher level
        if level<self.last_level:
            for l in range(level+1,self.last_level+1):
                self.d.pop(str(l),None)

    def _get_level (self,level):
        k=str(level)
        if k in self.d:
            n=self.d[k]
            assert(n!=None);
            return n
        else:
            return None

    def get_level (self,level):
        for l in range(level,0,-1):
            n=self._get_level(l)
            if n != None:
                return n
        assert(0);

    def is_header (self,tag):
        if len(tag)==2 and tag[0]=='h' and tag[1].isdigit() and (int(tag[1])>1):
            return (True);

    def handle_starttag(self, tag, attrs):
        if self.is_header (tag):
            self.attrs=attrs
            self.state=True;
            self.level=int(tag[1]);

    def handle_endtag(self, tag):
        if self.is_header (tag):
            self.state=False;

    def get_id (self):
        if self.attrs:
            for obj in self.attrs:
                if obj[0]=='id':
                    return obj[1]
        else:
            return None

    def handle_data(self, data):
        if self.state:

           level=self.level

           cnode=self.get_level(level-1)

           n=cnode.add_new_child(data,level,self.get_id());
           assert(n!=None);
           self.set_level(level,n)
           self.last_level=level

    def dump_as_json (self):
        return json.dumps(self.root.to_json_childs(), sort_keys=False, indent=4)

def create_toc_json(input_file, output_file, link_fmt=None):
    with open(input_file, encoding="utf-8") as f:
        l = f.readlines()
    html_input = ''.join(l)
    parser = TocHTMLParser(link_fmt)
    parser.feed(html_input);
    with open(output_file, 'w') as f:
        f.write(parser.dump_as_json())

re_xi = re.compile('''^(include|image)::([^.]*.(asciidoc|\\{PIC\\}))\[''', re.M)
def ascii_doc_scan(self):
    p = self.inputs[0].parent
    node_lst = [self.inputs[0]]
    seen = []
    depnodes = []
    while node_lst:
        nd = node_lst.pop(0)
        if nd in seen: continue
        seen.append(nd)

        code = nd.read()
        for m in re_xi.finditer(code):
            name = m.group(2)
            if m.group(3) == '{PIC}':

                ext = '.eps'
                if self.generator.rule.rfind('A2X') > 0:
                    ext = '.png'

                k = p.find_resource(name.replace('{PIC}', ext))
                if k:
                    depnodes.append(k)
            else:
                k = p.find_resource(name)
                if k:
                    depnodes.append(k)
                    node_lst.append(k)
    return [depnodes, ()]

def scansize(self):
    name = 'image::%s\\{PIC\\}\\[.*,(width|height)=(\\d+)' % self.inputs[0].name[:-4]
    re_src = re.compile(name)
    lst = self.inputs[0].parent.get_src().ant_glob('*.txt')
    for x in lst:
        m = re_src.search(x.read())
        if m:
            val = str(int(1.6 * int(m.group(2))))
            if m.group(1) == 'width':
                w = val
                h = "800"
            else:
                w = "800"
                h = val

            ext = self.inputs[0].name[-3:]
            if ext == 'eps':
                code = '-geometry %sx%s' % (w, h)
            elif ext == 'dia':
                if m.group(1) == 'width':
                    h = ''
                else:
                    w = ''
                code = '--size %sx%s' % (w, h)
            else:
                code = '-Gsize="%s,%s"' % (w, h)
            break
    else:
        return ([], '')

    return ([], code)

def options(opt):
    opt.add_option('--exe', action='store_true', default=False, help='Execute the program after it is compiled')
    opt.add_option('--performance', action='store_true', help='Build a performance report based on google analytics')
    opt.add_option('--performance-detailed',action='store_true',help='print detailed test results (date,time, build id and results) to csv file named _detailed_table.csv.')
    opt.add_option('--ndr', action = 'store_true', help = 'Include build of NDR report.')

search_path_arr = ['~/.local/bin', '/usr/local/bin/', '/usr/bin', './extensions']
os.environ['PATH'] = os.environ['PATH'] + ':' + ':'.join(search_path_arr)

def configure(conf):
    search_path = ' '.join(search_path_arr)
    conf.find_program('asciidoc', path_list=search_path, var='ASCIIDOC')
    conf.find_program('sphinx-build', path_list=search_path, var='SPHINX')
    conf.find_program('source-highlight', path_list=search_path, var='SRC_HIGHLIGHT')
    conf.find_program('dblatex', path_list=search_path, var='DBLATEX')
    conf.find_program('a2x', path_list=search_path, var='A2X')

    # asciidoctor
    conf.find_program('asciidoctor', path_list=search_path, var='ASCIIDOCTOR')
    #conf.find_file('multipage-html5-converter.rb', path_list=search_path)

def convert_to_pdf(task):
    input_file = task.outputs[0].abspath()
    out_dir = task.outputs[0].parent.get_bld().abspath()
    return  os.system('a2x --no-xmllint %s -f pdf  -d  article %s -D %s ' %('-v' if Log.verbose else '', task.inputs[0].abspath(),out_dir ) )

def convert_to_html_toc_book (task, disqus=False, generator='asciidoc', docinfo = True):
    """
        generate HTML output with TOC

        disqus     - add disqus (only if document contains <disqus></disqus> at the apropriate place)
        generator  - can be either 'asciidoc' or 'asciidoctor'
    """

    tools = {'asciidoc':    task.env['ASCIIDOC'][0],
             'asciidoctor': task.env['ASCIIDOCTOR'][0]}

    input_file = task.inputs[0].abspath()

    json_out_file = os.path.splitext(task.outputs[0].abspath())[0]+'.json'
    tmp = os.path.splitext(task.outputs[0].abspath())[0]+'.tmp'
    json_out_file_short = os.path.splitext(task.outputs[0].name)[0]+'.json'
    cmd='{0} -a stylesheet={1} -a  icons=true {2} -d book  -o {3} {4}'.format(
            tools[generator],
            task.inputs[1].abspath(),
            '-a docinfo' if docinfo else '',
            tmp,
            task.inputs[0].abspath());

    res = os.system(cmd)
    if res != 0 :
        return (1)

    create_toc_json(tmp,json_out_file)

    in_file  = tmp
    out_file = task.outputs[0].abspath()

    # inject TOC and DISQUS if needed
    with HTMLInjector(in_file, out_file) as injector:
        injector.inject_toc(json_out_file_short)
        if disqus:
            injector.inject_disqus(os.path.split(out_file)[1])

    return os.system('rm {0}'.format(tmp));

def convert_to_html_toc_book_no_docinfo(task):
    return convert_to_html_toc_book(task, docinfo = False)

# strip hierarchy from a multipage doc
# used to generate equal hierarchy for generating
# the pages
# also a 'main' hook will be added
def multipage_strip_hierarchy (in_file, out_file):

    with open(in_file) as f:
        lines = f.readlines()

    for i, line in enumerate(lines):
        if re.match("==+ [^ ].*", line):
            prefix, name = line.split(' ', 1)
            lines[i] = '== {0}\n'.format(name)

    # add a hook
    lines.append('\n== main\n')

    with open(out_file, "w") as f:
        f.write(''.join(lines))

# generate an asciidoctor chunk book for a target
def convert_to_asciidoctor_chunk_book(task, title, css):

    in_file          = task.inputs[0].abspath()
    src_dir          = os.path.dirname(in_file)
    out_dir          = os.path.splitext(task.outputs[0].abspath())[0]
    target_name      = os.path.basename(out_dir)

    css = os.path.abspath(css)

    # build chunked with no hierarchy
    with tempfile.NamedTemporaryFile() as tmp_file:
        # strip the hierarchy to make sure all pages are generated
        multipage_strip_hierarchy(in_file, tmp_file.name)

        multipage_backend = os.path.join('./extensions', 'multipage-html5-converter.rb')
        cmd = '{0} -a stylesheet={1} -r {2} -b multipage_html5 -D {3} {4} -o {5}.html -B {6}'.format(
                task.env['ASCIIDOCTOR'][0],
                css,
                multipage_backend,
                out_dir,
                tmp_file.name,
                target_name,
                src_dir)

        res = os.system(cmd)
        if res != 0:
            return (1)

    # build single only to get the TOC
    with tempfile.NamedTemporaryFile() as tmp_file:
        cmd = '{0} -a toc=left -b html5 -o {2} {3}'.format(
                task.env['ASCIIDOCTOR'][0],
                multipage_backend,
                tmp_file.name,
                in_file)

        res = os.system( cmd )
        if res != 0:
            return (1)

        # create TOC JSON in the library
        create_toc_json(tmp_file.name, os.path.join(out_dir, 'toc.json'), link_fmt = lambda link : '{0}{1}.html'.format(target_name, link))

    # iterate over all files and inject DISQUS if the pattern matches
    for filename in [f for f in os.listdir(out_dir) if f.endswith('.html')]:
        with HTMLInjector(os.path.join(out_dir, filename)) as injector:
            injector.inject_disqus(page_id = filename)

    # add TOC and disqus to the main page
    main = '{0}_main.html'.format(target_name)

    with HTMLInjector(os.path.join(out_dir, main),
                      os.path.join(out_dir, 'index.html')) as injector:
        injector.inject_title(title)
        injector.clear_class('sect1')
        injector.inject_toc('toc.json', fmt = 'multi', image_path = '..')
        injector.inject_ga()

    os.remove(os.path.join(out_dir, main))

    return 0

def convert_to_pdf_book(task):
    input_file = task.outputs[0].abspath()
    out_dir = task.outputs[0].parent.get_bld().abspath()
    return os.system('a2x --no-xmllint %s -f pdf  -d book %s -D %s ' %('-v' if Logs.verbose else '', task.inputs[0].abspath(),out_dir ) )

def ensure_dir(f):
    if not os.path.exists(f):
        os.makedirs(f)

def my_copy(task):
    input_file = task.outputs[0].get_src()
    out_dir = task.outputs[0].parent.get_bld().abspath()
    ensure_dir(out_dir)
    shutil.copy2(str(input_file), str(task.outputs[0]))
    return (0)

def do_visio(bld):
    for x in bld.path.ant_glob('visio\\*.vsd'):
        tg = bld(rule='${VIS} -i ${SRC} -o ${TGT} ', source=x, target=x.change_ext('.png'))

def get_sphinx_version(sphinx_path):
    try:
        global sphinx_version
        if not sphinx_version:
            sphinx_version_regexp = '^Sphinx \(sphinx-build\) (\d+)\.(\d+)\.\d+$'
            cmd = '%s %s --version' % (sys.executable, sphinx_path)
            output = subprocess.check_output(shlex.split(cmd), universal_newlines = True)
            for line in output.splitlines():
                ver = re.match(sphinx_version_regexp, line)
                if ver:
                    sphinx_version = float('%s.%s' % (ver.group(1), ver.group(2)))
        return sphinx_version
    except Exception as e:
        print('Error getting Sphinx version: %s' % e)

def get_trex_core_git():
    trex_core_git_path = os.path.abspath(os.path.join(os.getcwd(), os.pardir))
    if not os.path.isdir(trex_core_git_path):
        trex_core_git_path = os.getenv('TREX_CORE_GIT', None)
    return trex_core_git_path

def parse_hlt_args(task):
    trex_core_git_path = get_trex_core_git()
    if not trex_core_git_path:
        return 1
    hltapi_path = os.path.abspath(os.path.join(trex_core_git_path, 'scripts', 'automation', 'trex_control_plane', 'interactive', 'trex', 'stl', 'trex_stl_hltapi.py'))
    header = ['[options="header",cols="<.^1,^.^1,9<.^e"]', '|=================', '^| Argument | Default ^| Comment']
    footer = ['|=================\n']
    hlt_asciidoc = []
    category_regexp = '^(\S+)_kwargs = {$'
    comment_line_regexp = '^\s*#\s*(.+)$'
    arg_line_regexp = "^\s*'([^']+)':\s*'?([^,']+)'?,\s*#?\s*(.+)?$"
    if not os.path.exists(hltapi_path):
        raise Exception('Could not find hltapi file: %s' % hltapi_path)
    with open(hltapi_path) as f:
        in_args = False
        for line in f.read().splitlines():
            if not in_args:
                if line.startswith('import'):
                    break
                category_line = re.match(category_regexp, line)
                if category_line:
                    hlt_asciidoc.append('\n===== %s\n' % category_line.group(1))
                    hlt_asciidoc += header
                    in_args = True
                continue
            comment_line = re.match(comment_line_regexp, line)
            if comment_line:
                hlt_asciidoc.append('3+^.^s| %s' % comment_line.group(1).replace('|', '\|'))
                continue
            arg_line = re.match(arg_line_regexp, line)
            if arg_line:
                arg, default, comment = arg_line.groups()
                hlt_asciidoc.append('| %s | %s | %s' % (arg, default, comment.replace('|', '\|') if comment else ''))
                continue
            if line == '}':
                hlt_asciidoc += footer
                in_args = False
    if not len(hlt_asciidoc):
        raise Exception('Parsing of hltapi args failed')
    with open('build/hlt_args.asciidoc', 'w') as f:
        f.write('\n'.join(hlt_asciidoc))
    return 0

def build(bld):
    bld(rule=my_copy, target='symbols.lang')

    for x in bld.path.ant_glob('images\\**\**.png'):
            x1 = os.path.relpath(str(x), str(bld.path))
            bld(rule=my_copy, target=x1)
            bld.add_group()

    for x in bld.path.ant_glob('images\\**\**.jpg'):
        x1 = os.path.relpath(str(x), str(bld.path))
        bld(rule=my_copy, target=x1)
        bld.add_group()

    bld.add_group() # separator, the documents may require any of the pictures from above

    bld(rule=convert_to_html_toc_book,
        source='trex_emu.asciidoc waf.css', target='trex_emu.html', scan=ascii_doc_scan)

class Env(object):
    @staticmethod
    def get_env(name) :
        s= os.environ.get(name);
        if s == None:
            print("You should define $",name)
            raise Exception("Env error");
        return (s);

    @staticmethod
    def get_release_path () :
        s= Env().get_env('TREX_LOCAL_PUBLISH_PATH');
        s +=get_build_num ()+"/"
        return  s;

    @staticmethod
    def get_remote_release_path () :
        s= Env().get_env('TREX_REMOTE_PUBLISH_PATH');
        return  s;

    @staticmethod
    def get_local_web_server () :
        s= Env().get_env('TREX_WEB_SERVER');
        return  s;

    # extral web
    @staticmethod
    def get_trex_ex_web_key() :
        s= Env().get_env('TREX_EX_WEB_KEY');
        return  s;

    @staticmethod
    def get_trex_ex_web_path() :
        s= Env().get_env('TREX_EX_WEB_PATH');
        return  s;

    @staticmethod
    def get_trex_ex_web_user() :
        s= Env().get_env('TREX_EX_WEB_USER');
        return  s;

    @staticmethod
    def get_trex_ex_web_srv() :
        s= Env().get_env('TREX_EX_WEB_SRV');
        return  s;

    @staticmethod
    def get_trex_core() :
        s= Env().get_env('TREX_CORE_GIT');
        return  s;

def release(bld):
    # copy all the files to our web server
    core_dir = Env().get_trex_core()
    release_dir = core_dir +"/scripts/doc/";
    os.system('mkdir -p '+release_dir)
    os.system('cp -rv build/release_notes.* '+ release_dir)

def rsync_int(bld, src, dst):
    cmd = 'rsync -av  --rsh=ssh build/{src} {host}:{dir}/{dst}'.format(
           src = src,
           host = Env().get_local_web_server(),
           dir = Env().get_remote_release_path() + '../doc',
           dst = dst)
    ret = os.system(cmd)
    if ret:
        bld.fatal("cmd '%s' exited with return status: %s" % (cmd, ret))

def rsync_ext(bld, src, dst):
    cmd = 'rsync -avz  -e "ssh -i {key}" --rsync-path=/usr/bin/rsync build/{src} {user}@{host}:{dir}/doc/{dst}'.format(
           key  = Env().get_trex_ex_web_key(),
           src  = src,
           user = Env().get_trex_ex_web_user(),
           host = Env().get_trex_ex_web_srv(),
           dir  = Env().get_trex_ex_web_path(),
           dst  = dst)
    ret = os.system(cmd)
    if ret:
        bld.fatal("cmd '%s' exited with return status: %s" % (cmd, ret))

def publish(bld):
    # copy all the files to internal web server
    rsync_int(bld, '', '')

def publish_ext(bld):
    # copy all the files to external web server
    rsync_ext(bld, '', '')

def publish_perf(bld):
    pass

def publish_test(bld):
    # copy all the files to our web server
    remote_dir = "%s:%s" % ( Env().get_local_web_server(), Env().get_remote_release_path ()+'../test/')
    os.system('rsync -av  --rsh=ssh build/ %s' % (remote_dir))

def publish_both(bld):
    publish(bld)
    publish_ext(bld)

def run (bld):
    import BaseHTTPServer, SimpleHTTPServer
    import ssl

    cwd = os.getcwd()
    try:
        os.chdir('./build')
        print('Starting local HTTPS server at {0}:{1}'.format(os.uname()[1], 8080))
        httpd = BaseHTTPServer.HTTPServer(('0.0.0.0', 8080), SimpleHTTPServer.SimpleHTTPRequestHandler)
        httpd.socket = ssl.wrap_socket (httpd.socket, certfile='../server.pem', server_side=True)
        httpd.serve_forever()

    finally:
        os.chdir(cwd)


##**********************
#### 0: Preparation ####
##**********************
# 0a Load modules

import pandas as pd
import os
from sas7bdat import SAS7BDAT

# 0b Define icd codes for stroke
# test and pick one type
icdstroke = ['43301', '43311', '43321', '43331', '43381', '43391','43401', '43411',
             '43491', '436', '436', '430' ,'431', '4310', '43100', '432', '4320','43200',
             '4321' ,'4329']
'''icdstroke = [43301, 43311, 43321, 43331, 43381, 43391,43401, 43411,
             43491, 436, 436, 430,431, 4310, 43100, 432, 4320,43200,
             4321 ,4329]'''

# 0c Define the columns we will keep for subsequent analysis
cols_ip = ['ADMDT', 'ADMHR', 'AGE', 'DOB', 'DXA', 'DX01', 'DX02', 'DX03', 'DX04', 'DX05',
        'DX06', 'DX07', 'DX08', 'DX09', 'DX10','DX11', 'DX12', 'DX13', 'DX14', 'DX15',
        'DX16', 'DX17', 'DX18', 'DX19', 'DX20','DX21', 'DX22', 'DX23', 'DX24', 'DX25',
	'ETHNIC', 'PR01', 'PR02', 'PR03', 'PR04', 'PR05', 'RACE', 'SEX',
	'SOURCE1', 'SOURCE2', 'SOURCE3', 'SOURCE4', 'SOURCE5', 'UPID', 'ZIP', 'TPADM']
# the outpatient datasets do not have TPADM colun
cols_op = cols_ip[0:(len(cols_ip)-1)]

##***********************************
#### 1: Function to Curate SPARCS Admissions ####
##***********************************

# 1a Define function to keep only stroke admissions
def curateStrokes(YY, PP):
    # 1b Move to directory with the data
    os.chdir("N:/CSID/SPARCS/datasets")
    # 1c Define our SAS7BDAT object based on the year and pp
    iter = pd.read_sas(('sparcs_' + PP + YY+ '.sas7bdat'), iterator = True)
    # 1d Set the size of the chunks
    Chunk = 10000
    # 1e Determine the number of chunks we will readin
    # we add one to account for the final bit we readin
    NChunk = int(1 + round(iter.row_count / Chunk))
    # 1f Create the empty dataframe we will fill
    df = iter.read(0)
    df = df[cols_op]
    # 1g Loop over the chuncks
    for i in range(NChunk):
        # 1g.i For all but the last chunk, readin the set number of rows
        if i < NChunk:
            NRow = NChunk
        # 1g.ii For the final chunk, readin the number of remaining rows
        if i == NChunk:
            NRow = iter.row_count - (NChunk-1) * Chunk
        # 1g.iii Readin the rows
        df2 = iter.read(NRow)
        # 1g.iv Keep only stroke cases
        df2 = df2.query('DXA in @icdstroke | DX01 in @icdstroke | DX02 in @icdstroke | DX03 in @icdstroke')
        # 1g.v Keep only the revelvant columns
        df2 = df2[cols_op]
        # 1g.vi Add this data to that dataframe
        df = df.append(df2)
    # 1h Change directories to the project directory
    os.chdir('H:/Temp_CVD_Analysis/DataPrep/Data/Raw_data/Outcome_Data')
    # 1i Save as csv
    with open(('raw_' + YY+ "_" + PP+ '_stroke.csv'), 'w') as csvfile:
            df.to_csv(csvfile, index=False)

##***********************************
#### 2: Curate SPARCS Admissions ####
##***********************************

# 2a Curate the data
'''curateStrokes('95', 'op')
curateStrokes('96', 'op')
curateStrokes('97', 'op')
curateStrokes('98', 'op')
curateStrokes('99', 'op')
curateStrokes('00', 'op')
curateStrokes('01', 'op')
curateStrokes('02', 'op')
curateStrokes('03', 'op')
curateStrokes('04', 'op')
curateStrokes('05', 'op')
curateStrokes('06', 'op')
curateStrokes('07', 'op')
curateStrokes('08', 'op')
curateStrokes('09', 'op')
curateStrokes('10', 'op')
curateStrokes('11', 'op')'''

'''curateStrokes('12', 'op')
curateStrokes('13', 'op')
curateStrokes('14', 'op')
curateStrokes('15', 'op')'''

curateStrokes('95', 'ip')
curateStrokes('96', 'ip')
curateStrokes('97', 'ip')
curateStrokes('98', 'ip')
curateStrokes('99', 'ip')
'''curateStrokes('00', 'ip')
curateStrokes('01', 'ip')
curateStrokes('02', 'ip')
curateStrokes('03', 'ip')
curateStrokes('04', 'ip')
curateStrokes('05', 'ip')
curateStrokes('06', 'ip')
curateStrokes('07', 'ip')
curateStrokes('08', 'ip')
curateStrokes('09', 'ip')
curateStrokes('10', 'ip')
curateStrokes('11', 'ip')
curateStrokes('12', 'ip')
curateStrokes('13', 'ip')
curateStrokes('14', 'ip')
curateStrokes('15', 'ip')'''


"""
Old-fashioned DQAS code on QAOA ansatz design, now deprecated.
"""

# pylint: disable=wildcard-import

import sys

sys.path.insert(0, "../")

from collections import namedtuple
from pickle import dump
from matplotlib import pyplot as plt
import numpy as np
import tensorflow as tf
import tensorcircuit as tc
from tensorcircuit.applications.dqas import *
from tensorcircuit.applications.vags import *
from tensorcircuit.applications.layers import *
from tensorcircuit.applications.graphdata import regular_graph_generator

# qaoa_block_vag_energy = partial(qaoa_block_vag, f=(_identity, _neg))

tc.set_backend("tensorflow")

def main_layerwise_encoding():
    p = 5
    c = 7

    def noise():
        n = np.random.normal(loc=0.0, scale=0.002, size=[p, c])
        return tf.constant(n, dtype=tf.float32)

    def penalty_gradient(stp, nnp, lbd=0.15, lbd2=0.01):
        c = stp.shape[1]
        p = stp.shape[0]
        cost = tf.constant(
            [1.0, 1.0, 1.0, 1.0, 27 / 2 * 2.0, 27 / 2 * 2.0, 15 / 2.0], dtype=tf.float32
        )
        with tf.GradientTape() as t:
            prob = tf.math.exp(stp) / tf.tile(
                tf.math.reduce_sum(tf.math.exp(stp), axis=1)[:, tf.newaxis], [1, c]
            )
            penalty = 0.0
            for i in range(p - 1):
                penalty += lbd * tf.tensordot(prob[i], prob[i + 1], 1)
                penalty += lbd2 * tf.tensordot(cost, prob[i], [0, 0])
            penalty += lbd2 * tf.tensordot(cost, prob[p - 1], [0, 0])
            penalty_g = t.gradient(penalty, stp)
        return penalty_g

    learning_rate_sch = tf.keras.optimizers.schedules.InverseTimeDecay(0.12, 1, 0.01)

    DQAS_search(
        qaoa_vag_energy,
        g=regular_graph_generator(n=8, d=3),
        p=p,
        batch=8,
        prethermal=0,
        prethermal_preset=[0, 6, 1, 6, 1],
        epochs=3,
        parallel_num=2,
        pertubation_func=noise,
        nnp_initial_value=np.random.normal(loc=0.23, scale=0.06, size=[p, c]),
        stp_regularization=penalty_gradient,
        network_opt=tf.keras.optimizers.Adam(learning_rate=0.03),
        prethermal_opt=tf.keras.optimizers.Adam(learning_rate=0.04),
        structure_opt=tf.keras.optimizers.SGD(learning_rate=learning_rate_sch),
    )

result = namedtuple("result", ["epoch", "cand", "loss"])

def main_block_encoding():
    learning_rate_sch = tf.keras.optimizers.schedules.InverseTimeDecay(0.3, 1, 0.01)
    p = 6
    c = 8

    def record():
        return result(
            get_var("epoch"), get_var("cand_preset_repr"), get_var("avcost1").numpy()
        )

    def noise():
        # p = 6
        # c = 6
        n = np.random.normal(loc=0.0, scale=0.2, size=[2 * p, c])
        return tf.constant(n, dtype=tf.float32)

    stp, nnp, h = DQAS_search(
        qaoa_block_vag_energy,
        g=regular_graph_generator(n=8, d=3),
        batch=8,
        prethermal=0,
        prethermal_preset=[0, 6, 1, 6, 1],
        epochs=3,
        parallel_num=2,
        pertubation_func=noise,
        p=p,
        history_func=record,
        nnp_initial_value=np.random.normal(loc=0.23, scale=0.06, size=[2 * p, c]),
        network_opt=tf.keras.optimizers.Adam(learning_rate=0.04),
        prethermal_opt=tf.keras.optimizers.Adam(learning_rate=0.04),
        structure_opt=tf.keras.optimizers.SGD(learning_rate=learning_rate_sch),
    )

    with open("qaoa_block.result", "wb") as f:
        dump([stp.numpy(), nnp.numpy(), h], f)

    epochs = np.arange(len(h))
    data = np.array([r.loss for r in h])
    plt.plot(epochs, data)
    plt.xlabel("epoch")
    plt.ylabel("objective")
    plt.savefig("qaoa_block.pdf")

layer_pool = [Hlayer, rxlayer, rylayer, rzlayer, xxlayer, yylayer, zzlayer]
block_pool = [
    Hlayer,
    rx_zz_block,
    zz_ry_block,
    zz_rx_block,
    zz_rz_block,
    xx_rz_block,
    yy_rx_block,
    rx_rz_block,
]
set_op_pool(block_pool)

if __name__ == "__main__":
    # main_layerwise_encoding()
    main_block_encoding()


# define a list of valid operation signs
operations = ['+', '-', '*', '/']

# define a function to perform the calculator operations
def calculate(input_string):
    # initialize the operand variables
    operand1 = 0
    operand2 = 0
    operation = ''

    # iterate through the input string to extract the operands and operation
    for i, char in enumerate(input_string):
        # check if the character is a digit
        if char.isdigit():
            # check if the operation has been set
            if operation:
                # set the second operand
                operand2 = operand2 * 10 + int(char)
            else:
                # set the first operand
                operand1 = operand1 * 10 + int(char)
        # check if the character is an operation sign
        elif char in operations:
            # set the operation
            operation = char

    # perform the operation and return the result
    if operation == '+':
        return operand1 + operand2
    elif operation == '-':
        return operand1 - operand2
    elif operation == '*':
        return operand1 * operand2
    elif operation == '/':
        # check for division by zero
        if operand2 == 0:
            return "Cannot divide by zero"
        else:
            return operand1 / operand2
    else:
        return "Invalid operation"

# get the input from the user
input_string = input("Enter an operation (e.g. 2+3): ")

# calculate and print the result
result = calculate(input_string)
print(result)


import plum
from plum.types import Variable
import torch
import numpy as np

class GreedySearch:

    def __init__(self, vocab, max_steps=999):
        self.vocab = vocab
        self.max_steps = max_steps
        self.reset()

    def reset(self):
        self.is_finished = False
        self.steps = 0
        self._states = []
        self._outputs = []

    def init_state(self, batch_size, device, source_indices,
                   encoder_state=None):

        output = Variable(
            torch.LongTensor([self.vocab.start_index] * batch_size)\
                .view(1, -1),
            lengths=torch.LongTensor([1] * batch_size),
            length_dim=0, batch_dim=1, pad_value=self.vocab.pad_index)

        if str(device) != "cpu":
            output = output.cuda(device)

        mask = source_indices.mask.t()

        return {"target": output, "decoder_state": encoder_state,
                "inputs": output, 'pointer_mask': mask}

    def next_state(self, decoder, prev_state, context, active_items,
                   controls):

        # Get next state from the decoder.
        next_state = decoder.next_state(prev_state, context, controls=controls)

        next_pointer = next_state['log_probs']\
            .permute_as_sequence_batch_features().max(2)[1]\

        target = context['source_indices'].data.t().gather(
            1, next_pointer.data.t())
        target = prev_state['target'].new_with_meta(target.t())
        next_state['target'] = target

        pm = prev_state['pointer_mask']
        pm.scatter_(1, next_pointer.data.t(), 1)
        next_state['pointer_mask'] = pm

        # Mask outputs if we have already completed that batch item.
        next_state["target"].data.view(-1).masked_fill_(
            ~active_items, int(self.vocab.pad_index))

        next_inputs = torch.cat(
            [
                prev_state['inputs'].data,
                next_state['target'].data,
            ],
            0)
        next_inputs = Variable(
            next_inputs,
            prev_state['inputs'].lengths + 1,
            length_dim=0, batch_dim=1,
            pad_value=prev_state['inputs'].pad_value)
        next_state['inputs'] = next_inputs
#
        return next_state

    def check_termination(self, next_state, active_items):

        # Check for stop tokens and batch item inactive if so.
        nonstop_tokens = next_state["target"].data.view(-1).ne(
            int(self.vocab.stop_index))
        active_items = active_items.data.mul_(nonstop_tokens)

        return active_items

    def _collect_search_states(self, active_items):
        # TODO implement search states api.
        #search_state = self._states[0]
        #for next_state in self._states[1:]:
        #    search_state.append(next_state)
        #self._states = search_state
        self._outputs = torch.cat([o.data for o in self._outputs], dim=0)

    def __call__(self, decoder, encoder_state, controls=None):
        self.reset()
        batch_size = encoder_state["batch_size"]
        device = encoder_state["device"]

        search_state = self.init_state(batch_size, device,
                                       encoder_state['source_indices'],
                                       encoder_state.get("state", None))

        context = {
            "encoder_output": encoder_state["output"],
            "source_indices": encoder_state['source_indices'],
        }
        active_items = torch.BoolTensor(batch_size).fill_(1)
        output_length = torch.LongTensor(batch_size).fill_(0)
        if str(encoder_state["device"]) != "cpu":
            active_items = active_items.cuda(encoder_state["device"])
            output_length = output_length.cuda(encoder_state["device"])

        prev_decoder_inputs = search_state["target"]
        step_masks = []
        # Perform search until we either trigger a termination condition for
        # each batch item or we reach the maximum number of search steps.
        while self.steps < self.max_steps and not self.is_finished:

            inactive_items = ~active_items

            # Mask any inputs that are finished, so that greedy would
            # be identitcal to forward passes.
            search_state["target"].data.view(-1).masked_fill_(
                inactive_items, int(self.vocab.pad_index))

            step_masks.append(inactive_items)
            self.steps += 1
            search_state = self.next_state(
                decoder, search_state, context, active_items, controls)

            self._states.append(search_state)
            self._outputs.append(search_state["target"].clone()\
                .permute_as_sequence_batch_features())
            output_length = output_length + active_items.long()
            temp_outputs = torch.cat(
                    [prev_decoder_inputs.tensor] + \
                            [x.tensor for x in self._outputs], 0)
            temp_outputs = Variable(temp_outputs, lengths=output_length + 1,
                                    batch_dim=1, length_dim=0,
                                    pad_value=self.vocab.pad_index)
            search_state["prev_output"] = temp_outputs

            active_items = self.check_termination(search_state, active_items)
            self.is_finished = torch.all(~active_items)

        # Finish the search by collecting final sequences, and other
        # stats.
        self._collect_search_states(active_items)
        self._incomplete_items = active_items
        self._is_finished = True

        self._mask_T = torch.stack(step_masks)
        self._mask = self._mask_T.t().contiguous()
        return self

    def __getitem__(self, key):
        if key == "target":
            return self._outputs

    def output(self, as_indices=False):

        if as_indices:
            lengths = (~self._mask_T).sum(0)
            return Variable(self._outputs, lengths=lengths, length_dim=0,
                            batch_dim=1, pad_value=self.vocab.pad_index)\
                            .apply_sequence_mask()

        tokens = []
        for output in self._outputs.t():
            tokens_i = []
            for index in output:
                if index in [self.vocab.pad_index, self.vocab.stop_index]:
                    break
                tokens_i.append(self.vocab.token(index))
            tokens.append(tokens_i)

        return tokens

    def state_dict(self):
        skip_keys = ["_parameter_names", "_parameter_tags", "_submodule_names",
                     "_submodule_tags", "_backend", "_parameters", "_buffers",
                     "_backward_hooks", "_forward_hooks", "_forward_pre_hooks",
                     "_state_dict_hooks", "_load_state_dict_pre_hooks",
                     "_modules", "_hyperparameters", "_hyperparameter_tags",
                     "_name2hyperparameter"]

        hyperparameters = {}
        for name, hp in self._name2hyperparameter.items():
            hyperparameters[name] = _state_dict_helper(
                self._hyperparameters[hp])

        modules = {}
        parameters = {}
        if hasattr(self, "_modules"):
            submod_names = set(self._submodule_names.values())
            for name, module in self._modules.items():
                if name in submod_names:
                    hyperparameters[name] = _state_dict_helper(module)
                else:
                    modules[name] = _state_dict_helper(module)
                    for n, p in module.named_parameters():
                        parameters[tuple([name]) + tuple(n.split("."))] = p
        if hasattr(self, "_parameters"):
            for k, p in self._parameters.items():
                parameters[k] = p
        return {
            "__plum_type__": self.plum_id,
            "attributes": {},
            "hyperparameters": hyperparameters,
            "modules": modules,
            "parameters": parameters,
        }

def _state_dict_helper(obj):
    if isinstance(obj, (list, tuple)):
        return [_state_dict_helper(x) for x in obj]
    elif isinstance(obj, dict):
        return {k: _state_dict_helper(v) for k, v in obj.items()}
    elif hasattr(obj, "state_dict"):
        return obj.state_dict()
    else:
        return obj

class BeamSearch:
    def __init__(self, vocab, beam_size=8, max_steps=999):
        self.vocab = vocab
        self.beam_size = beam_size
        self.max_steps = max_steps
        self.reset()

    def reset(self):
        self.is_finished = False
        self.steps = 0
        self._states = []
        self._outputs = []

    def init_state(self, batch_size, device, source_indices, encoder_state):

        beam_state = {}

        if encoder_state:
            n, bs, ds = encoder_state.size()
            assert batch_size == bs
            beam_state["decoder_state"] = encoder_state.unsqueeze(2)\
                .repeat(1, 1, self.beam_size, 1)\
                .view(n, batch_size * self.beam_size, ds)

        beam_state["target"] = Variable(
            torch.LongTensor(
                [self.vocab.start_index] * batch_size * self.beam_size)\
                .view(1, -1),
            lengths=torch.LongTensor([1] * batch_size * self.beam_size),
            length_dim=0, batch_dim=1, pad_value=self.vocab.pad_index)

        beam_state["target"] = beam_state["target"].to(device)
        beam_state['inputs'] = beam_state['target']

        source_indices = source_indices.permute_as_batch_sequence_features()
        mask = source_indices.mask.unsqueeze(1).repeat(1, self.beam_size, 1)\
            .contiguous().view(-1, source_indices.size(1))
        beam_state['pointer_mask'] = mask

        # Start the first beam of each batch with 0 log prob, and all others
        # with -inf.
        beam_state["accum_log_prob"] = (
            self._init_accum_log_probs(batch_size, device)
        )

        # At the first time step no sequences have been terminated so this mask
        # is all 0s.
        beam_state["terminal_mask"] = (
            torch.BoolTensor(
                1, batch_size * self.beam_size, 1).fill_(0).to(device)
        )

        return beam_state

    def _init_accum_log_probs(self, batch_size, device):
        lp = torch.FloatTensor(1, batch_size, self.beam_size, 1)
        if "cuda" in str(device):
            lp = lp.cuda(device)
        lp.data.fill_(0)
        lp.data[:,:,1:].fill_(float("-inf"))
        return lp.view(1, batch_size * self.beam_size, 1)

    def init_context(self, encoder_state):
        n, bs, ds = encoder_state["output"].size()
        beam_encoder_output = encoder_state["output"].repeat_batch_dim(
            self.beam_size)
            #.repeat(1, 1, self.beam_size, 1)\
            #.view(n, bs * self.beam_size, ds)

        source_indices = encoder_state['source_indices']\
            .permute_as_batch_sequence_features()
        source_indices = source_indices.data.unsqueeze(1)\
            .repeat(1, self.beam_size, 1).view(-1, source_indices.size(1))

        return {"encoder_output": beam_encoder_output,
                'source_indices': source_indices}

    def __call__(self, decoder, encoder_state, controls=None):
        self.reset()

        batch_size = encoder_state["batch_size"]
        device = encoder_state['device']
        search_state = self.init_state(
            batch_size,
            device,
            encoder_state['source_indices'],
            encoder_state.get("state", None))

        search_context = self.init_context(encoder_state)
        active_items = torch.BoolTensor(batch_size).fill_(1).to(device)

        if controls is not None:
            controls = controls.repeat_batch_dim(self.beam_size)

        self._beam_scores = [list() for _ in range(batch_size)]
        self._num_complete = torch.LongTensor(batch_size).fill_(0).to(device)
        self._terminal_info = [list() for _ in range(batch_size)]

        # Perform search until we either trigger a termination condition for
        # each batch item or we reach the maximum number of search steps.
        while self.steps < self.max_steps and not self.is_finished:

            search_state = self.next_state(
                decoder, batch_size, search_state, search_context,
                active_items, controls)
            active_items = self.check_termination(search_state, active_items)
            self._is_finished = torch.all(~active_items)

            self._states.append(search_state)
            self.steps += 1

        # Finish the search by collecting final sequences, and other
        # stats.
        self._collect_search_states(active_items)
        self._incomplete_items = active_items
        self._is_finished = True

        return self

    def next_state(self, decoder, batch_size, prev_state, context,
                   active_items, controls):

        # Get next state from the decoder.
        next_state = decoder.next_state(prev_state, context, controls=controls)

        # Compute the top beam_size next outputs for each beam item.
        # topk_lps (1 x batch size x beam size x beam size)
        # candidate_outputs (1 x batch size x beam size x beam size)
        topk_lps, candidate_outputs = torch.topk(
            next_state["log_probs"].data \
                .view(1, batch_size, self.beam_size, -1),
            k=self.beam_size, dim=3)

        # If any sequence was completed last step, we should mask it's log
        # prob so that we don't generate from the terminal token.
        # slp (1 x batch_size x beam size x 1)
        slp = prev_state["accum_log_prob"] \
            .masked_fill(prev_state["terminal_mask"], float("-inf")) \
            .view(1, batch_size, self.beam_size, 1)

        # Combine next step log probs with the previous sequences cumulative
        # log probs, i.e.
        #     log P(y_t) = log P(y_<t) + log P(y_t)
        # candidate_log_probs (1 x batch size x beam size x beam size)
        candidate_log_probs = slp + topk_lps

        # Rerank and select the beam_size best options from the available
        # beam_size ** 2 candidates.
        # b_seq_lps (1 x (batch size * beam size) x 1)
        # b_scores (1 x (batch size * beam size) x 1)
        # b_pointer (1 x (batch size * beam size))
        # b_indices ((batch size * beam size))
        b_seq_lps, b_scores, b_pointer, b_indices = self._next_candidates(
            batch_size, candidate_log_probs, candidate_outputs)

        b_pointer_mask = (
            (b_pointer.data >= context['source_indices'].size(1))
            |
            (b_pointer.data < 0)
        )

        b_pointer = b_pointer.masked_fill(b_pointer_mask, 0)
#        print(context['source_indices'])
#        print(b_pointer)
        b_output = context['source_indices'].gather(1, b_pointer.data.t())
        b_output = b_pointer.new_with_meta(b_output.t())
#        print(b_output)

#        print(b_indices)
        pm = prev_state['pointer_mask']
#        print(pm.size())
#        print(pm.long())
        pm = pm.index_select(0, b_indices)
#        print(pm.long())
        pm.scatter_(1, b_pointer.data.t(), 1)

#        print(pm.long())
#        print(pm.size())

        # TODO re-implement this behavior
        #next_state.stage_indexing("batch", b_indices)

#        print("BEAM INDICES")
#        print(b_indices)
#        print(prev_state['inputs'])
#        print(prev_state['inputs'].index_select(1, b_indices))
#        print(type(prev_state['inputs'].index_select(1, b_indices)))
#        print(b_output)
#        print(type(b_output))
        next_inputs = torch.cat(
            [
                prev_state['inputs'].index_select(1, b_indices).data,
                b_output.data,
            ],
            0)
        next_inputs = Variable(
            next_inputs,
            prev_state['inputs'].lengths + 1,
            length_dim=0, batch_dim=1,
            pad_value=prev_state['inputs'].pad_value)
#        print(next_inputs)

        next_state = {
            "decoder_state": next_state["decoder_state"]\
                .index_select(1, b_indices),
            "target": b_output,
            "accum_log_prob": b_seq_lps,
            "beam_score": b_scores,
            "beam_indices": b_indices,
            "inputs": next_inputs,
            'pointer_mask': pm,
        }
        return next_state
        #exit()
        #next_state = {"decoder_state": next_state["decoder_state"]
        #print(next_state.keys())

        next_state["output"] = (b_output, ("batch", "sequence"))
        next_state["cumulative_log_probability"] = (
            b_seq_lps, ("sequence", "batch", "placeholder")
        )
        next_state["beam_score"] = (
            b_scores, ("sequence", "batch", "placeholder")
        )
        next_state["beam_indices"] = (b_indices, ("batch"))

        return next_state

    def _next_candidates(self, batch_size, log_probs, candidates):
        # TODO seq_lps should really be called cumulative log probs.

        # flat_beam_lps (batch size x (beam size ** 2))
        flat_beam_lps = log_probs.view(batch_size, -1)

        flat_beam_scores = flat_beam_lps / (self.steps + 1)

        # beam_seq_scores (batch size x beam size)
        # relative_indices (batch_size x beam size)
        beam_seq_scores, relative_indices = torch.topk(
            flat_beam_scores, k=self.beam_size, dim=1)
        relative_indices_mask = (
            (relative_indices >= flat_beam_scores.size(1))
            |
            (relative_indices < 0)
        )
        relative_indices = relative_indices.masked_fill(
            relative_indices_mask, 0)

        # beam_seq_lps (batch size x beam size)
        beam_seq_lps = flat_beam_lps.gather(1, relative_indices)

        # TODO make these ahead of time.
        offset1 = (
            torch.arange(batch_size, device=beam_seq_lps.device)
                * self.beam_size
        ).view(batch_size, 1)

        offset2 = offset1 * self.beam_size

        beam_indexing = ((relative_indices // self.beam_size) + offset1) \
            .view(-1)

        # beam_seq_lps (1 x (batch_size * beam_size) x 1)
        beam_seq_lps = beam_seq_lps \
            .view(1, batch_size * self.beam_size, 1)

        # beam_seq_scores (1 x (batch_size * beam_size) x 1)
        beam_seq_scores = beam_seq_scores \
            .view(1, batch_size * self.beam_size, 1)

        # next_output (1 x (batch size * beam size))
        next_candidate_indices = (relative_indices + offset2).view(-1)
        next_output = Variable(
            candidates.view(-1)[next_candidate_indices].view(1, -1),
            lengths=candidates.new().long().new(batch_size * self.beam_size)\
                .fill_(1),
            length_dim=0, batch_dim=1, pad_value=self.vocab.pad_index)

        return beam_seq_lps, beam_seq_scores, next_output, beam_indexing

    def check_termination(self, next_state, active_items):

        # view as batch size x beam size
        next_output = next_state["target"].data \
            .view(-1, self.beam_size)
        batch_size = next_output.size(0)

        is_complete = next_output.eq(int(self.vocab.stop_index))
        complete_indices = np.where(is_complete.cpu().data.numpy())

        for batch, beam in zip(*complete_indices):
            if self._num_complete[batch] == self.beam_size:
                continue
            else:
                self._num_complete[batch] += 1

                # Store step and beam that finished so we can retrace it
                # later and recover arbitrary search state item.
                self._terminal_info[batch].append(
                    (self.steps, beam + batch * self.beam_size))

                IDX = batch * self.beam_size + beam
                self._beam_scores[batch].append(
                    next_state["beam_score"][0, IDX, 0].view(1))

        next_state["terminal_mask"] = (
            is_complete.view(1, batch_size * self.beam_size, 1)
        )
        active_items = self._num_complete < self.beam_size

        return active_items

    def _collect_search_states(self, active_items):

        batch_size = active_items.size(0)

        last_state = self._states[-1]
        last_step = self.steps - 1
        for batch in range(batch_size):
            beam = 0
            while len(self._beam_scores[batch]) < self.beam_size:
                IDX = batch * self.beam_size + beam
                self._beam_scores[batch].append(
                    last_state["beam_score"][0, IDX, 0].view(1))
                self._terminal_info[batch].append(
                    (last_step, beam + batch * self.beam_size))
                beam += 1

        # TODO consider removing beam indices from state
        beam_indices = torch.stack([state["beam_indices"]
                                    for state in self._states])

        self._beam_scores = torch.stack([torch.cat(bs)
                                         for bs in self._beam_scores])

        lengths = self._states[0]["target"].new(
            [[step + 1 for step, beam in self._terminal_info[batch]]
             for batch in range(batch_size)])

        selector = self._states[0]["target"].new(
            batch_size, self.beam_size, lengths.max())
        mask = selector.new().bool().new(selector.size()).fill_(1)

        for batch in range(batch_size):
            for beam in range(self.beam_size):
                step, real_beam = self._terminal_info[batch][beam]
                mask[batch, beam,:step + 1].fill_(0)
                self._collect_beam(batch, real_beam, step,
                                   beam_indices,
                                   selector[batch, beam])
        selector = selector.view(batch_size * self.beam_size, -1)

        ## RESORTING HERE ##
        #if self.sort_by_score:
        # TODO make this an option again
        if True:
            self._beam_scores, I = torch.sort(self._beam_scores, dim=1,
                                              descending=True)
            offset1 = (
                torch.arange(batch_size, device=I.device) * self.beam_size
            ).view(batch_size, 1)
            II = I + offset1
            selector = selector[II.view(-1)]
            mask = mask.view(batch_size * self.beam_size,-1)[II]\
                .view(batch_size, self.beam_size, -1)
            lengths = lengths.gather(1, I)
        ##

        # TODO reimplement staged indexing
#        for step, sel_step in enumerate(selector.split(1, dim=1)):
#            self._states[step].stage_indexing("batch", sel_step.view(-1))

        self._output = []
        for step, sel_step in enumerate(selector.split(1, dim=1)):
            self._output.append(
                self._states[step]["target"].index_select(1, sel_step.view(-1))
            )
        #print(self._states[0]["output"].size())
        self._output = plum.cat([o.data for o in self._output], 0).t()\
            .view(batch_size, self.beam_size, -1)

        for i in range(batch_size):
            for j in range(self.beam_size):
                self._output[i,j,lengths[i,j]:].fill_(
                    int(self.vocab.pad_index))

        self._lengths = lengths

        return self

#        for batch_out in self._output:
#        #print(self._output.t().view(batch_size)
#        #for batch_out in self._output.t().view(batch_size, self.beam_size, -1):
#            for row in batch_out:
#                print(" ".join([self.vocab[t] for t in row if t != self.vocab.pad_index]))
#            print()
#        print(lengths)
#        print(lengths.size())
#        print(batch_size)
#        exit()
#
#        states = self._states[0]
#        for state in self._states[1:]:
#            states.append(state)
#
#        self._states = states
#        self._selector = selector
#        self._lengths = lengths
#        self._selector_mask = mask.view(self.batch_size * self.beam_size, -1)
#        self._selector_mask_T = self._selector_mask.t().contiguous()

    def _collect_beam(self, batch, beam, step, beam_indices,
                      selector_out):
        selection = [0] * beam_indices.size(0)
        selector_out[step + 1:].fill_(0)
        while step >= 0:
            selection[step] = beam
            selector_out[step].fill_(beam)
            next_beam = beam_indices[step, beam].item()
            beam = next_beam
            step -= 1

    def output(self, as_indices=False, n_best=-1):
        if n_best < 1:
            o = self._output[:,0]
            if as_indices:
                olen = (self._lengths[:,0])
                return Variable(o, lengths=olen, batch_dim=0, length_dim=1,
                                pad_value=int(self.vocab.pad_index))
            tokens = []
            for row in o:
                tokens.append([self.vocab.token(t) for t in row
                               if t != self.vocab.pad_index])
            return tokens

        elif n_best < self.beam_size:
            o = self._output[:,:n_best]
        else:
            o = self._output

        if as_indices:
            print(o)
            print(self._lengths)
            return o

        beams = []
        for beam in o:
            tokens = []
            for row in beam:
                tokens.append([self.vocab.token(t) for t in row
                               if t != self.vocab.pad_index])
            beams.append(tokens)
        return beams