# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#*** nmeta - Network Metadata - Abstractions of Switches for OpenFlow Calls

"""
This module is part of the nmeta suite running on top of Ryu SDN controller.

It provides classes that abstract the details of OpenFlow switches
"""

#*** General Imports:
import sys
import struct

#*** For timestamps:
import datetime

#*** Ryu Imports:
from ryu.lib import addrconv
from ryu.ofproto import ofproto_v1_3
from ryu.lib.packet import packet
from ryu.lib.packet import ipv4, ipv6
from ryu.lib.packet import tcp
from ryu.lib.packet import udp

#*** For logging configuration:
from baseclass import BaseClass

#*** mongodb Database Import:
import pymongo
from pymongo import MongoClient

#*** Constant to use for a port not found value:
PORT_NOT_FOUND = 999999999

#*** Supports OpenFlow version 1.3:
OFP_VERSIONS = [ofproto_v1_3.OFP_VERSION]

class Switches(BaseClass):
    """
    This class provides an abstraction for a set of OpenFlow
    Switches.

    It stores instances of the Switch class in a dictionary keyed
    by DPID. The switch instances are accessible externally.

    A standard (not capped) MongoDB database collection is used to
    record switch details so that they can be accessed via the
    external API.
    """
    def __init__(self, config):
        #*** Required for BaseClass:
        self.config = config
        #*** Set up Logging with inherited base class method:
        self.configure_logging(__name__, "switches_logging_level_s",
                                       "switches_logging_level_c")

        #*** Set up database collections:
        #*** Get parameters from config:
        mongo_addr = config.get_value("mongo_addr")
        mongo_port = config.get_value("mongo_port")
        mongo_dbname = config.get_value("mongo_dbname")

        #*** Start mongodb:
        self.logger.info("Connecting to MongoDB database...")
        mongo_client = MongoClient(mongo_addr, mongo_port)

        #*** Connect to MongoDB nmeta database:
        db_nmeta = mongo_client[mongo_dbname]

        #*** Delete (drop) previous switches collection if it exists:
        self.logger.debug("Deleting previous switches MongoDB collection...")
        db_nmeta.switches_col.drop()

        #*** Create the switches collection:
        self.switches_col = db_nmeta.create_collection('switches_col')

        #*** Index dpid key to improve look-up performance:
        self.switches_col.create_index([('dpid', pymongo.TEXT)], unique=False)

        #*** Get max bytes of new flow packets to send to controller from
        #*** config file:
        self.miss_send_len = config.get_value("miss_send_len")
        if self.miss_send_len < 1500:
            self.logger.info("Be aware that setting "
                             "miss_send_len to less than a full size packet "
                             "may result in errors due to truncation. "
                             "Configured value is %s bytes",
                             self.miss_send_len)
        #*** Tell switch how to handle fragments (see OpenFlow spec):
        self.ofpc_frag = config.get_value("ofpc_frag")

        #*** Flow mod cookie value offset indicates flow session direction:
        self.offset = config.get_value("flow_mod_cookie_reverse_offset")

        #*** Dictionary of the instances of the Switch class,
        #***  key is the switch DPID which is assumed to be unique:
        self.switches = {}

    def add(self, datapath):
        """
        Add a switch to the Switches class
        """
        dpid = datapath.id
        self.logger.info("Adding switch dpid=%s", dpid)
        switch = Switch(self.config, datapath, self.offset)
        switch.dpid = dpid
        (ip_address, port) = datapath.address
        #*** Record class instance into dictionary to make it accessible:
        self.switches[datapath.id] = switch
        #*** Record switch in database collection:
        self.switches_col.update_one({'dpid': dpid},
                                {
                                "$set":
                                    {
                                    'dpid': dpid,
                                    'time_connected': datetime.datetime.now(),
                                    'ip_address': ip_address,
                                    'port': port
                                }
                        }, upsert=True)
        #*** Set the switch up for operation:
        switch.set_switch_config(self.ofpc_frag, self.miss_send_len)
        switch.request_switch_desc()
        switch.set_switch_table_miss(self.miss_send_len)
        return 1

    def stats_reply(self, msg):
        """
        Read in a switch stats reply
        """
        body = msg.body
        dpid = msg.datapath.id
        #*** Look up the switch:
        if dpid in self.switches:
            self.logger.info('event=DescStats Switch dpid=%s is mfr_desc="%s" '
                      'hw_desc="%s" sw_desc="%s" serial_num="%s" dp_desc="%s"',
                      dpid, body.mfr_desc, body.hw_desc, body.sw_desc,
                      body.serial_num, body.dp_desc)
            switch = self.switches[dpid]
            switch.mfr_desc = body.mfr_desc
            switch.hw_desc = body.hw_desc
            switch.sw_desc = body.sw_desc
            switch.serial_num = body.serial_num
            switch.dp_desc = body.dp_desc

            #*** Update switch details in database collection:
            self.switches_col.update_one({'dpid': dpid},
                                {
                                "$set":
                                    {
                                    'dpid': dpid,
                                    'mfr_desc': switch.mfr_desc,
                                    'hw_desc': switch.hw_desc,
                                    'sw_desc': switch.sw_desc,
                                    'serial_num': switch.serial_num,
                                    'dp_desc': switch.dp_desc
                                }
                        }, upsert=True)
        else:
            self.logger.warning("Ignoring DescStats reply from unknown switch"
                                                              " dpid=%s", dpid)
            return 0

    def __getitem__(self, key):
        """
        Passed a dpid key and return corresponding switch
        object, or 0 if it doesn't exist.
        Example:
            switch = switches[dpid]
        """
        if key in self.switches:
            return self.switches[key]
        else:
            return 0

    def delete(self, datapath):
        """
        Delete a switch from the Switches class
        """
        dpid = datapath.id
        self.logger.info("Deleting switch dpid=%s", dpid)
        #*** Get relevant instance of switch class:
        if not dpid in self.switches:
            return 0
        #*** Delete from dictionary of switches:
        del self.switches[dpid]
        #*** Delete switch from database collection:
        self.switches_col.delete_one({'dpid': dpid})
        return 1

class Switch(BaseClass):
    """
    This class provides an abstraction for an OpenFlow
    Switch
    """
    def __init__(self, config, datapath, offset):
        #*** Required for BaseClass:
        self.config = config
        #*** Set up Logging with inherited base class method:
        self.configure_logging(__name__, "switches_logging_level_s",
                                       "switches_logging_level_c")
        #*** Initialise switch variables:
        self.config = config
        self.datapath = datapath
        self.switch_hash = ""
        self.dpid = 0
        self.ip_address = ""
        self.time = ""
        self.cxn_status = ""
        self.cxn_ver = ""
        self.mfr_desc = ""
        self.hw_desc = ""
        self.sw_desc = ""
        self.serial_num = ""
        self.dp_desc = ""
        #*** Instantiate a class that represents flow tables:
        self.flowtables = FlowTables(config, datapath, offset)

    def dbdict(self):
        """
        Return a dictionary object of switch
        parameters for storing in the database
        """
        return self.__dict__

    def request_switch_desc(self):
        """
        Send an OpenFlow request to the switch asking it to
        send us it's description data
        """
        parser = self.datapath.ofproto_parser
        req = parser.OFPDescStatsRequest(self.datapath, 0)
        self.logger.debug("Sending description request to dpid=%s",
                            self.datapath.id)
        self.datapath.send_msg(req)

    def set_switch_config(self, config_flags, miss_send_len):
        """
        Set config on a switch including config flags that
        instruct fragment handling behaviour and miss_send_len
        which controls the number of bytes sent to the controller
        when the output port is specified as the controller.
        """
        parser = self.datapath.ofproto_parser
        self.logger.info("Setting config on switch "
                         "dpid=%s to config_flags flag=%s and "
                         "miss_send_len=%s bytes",
                          self.dpid, config_flags, miss_send_len)
        try:
            self.datapath.send_msg(parser.OFPSetConfig(
                                     self.datapath,
                                     config_flags,
                                     miss_send_len))
        except:
            #*** Log the error and return 0:
            exc_type, exc_value, exc_traceback = sys.exc_info()
            self.logger.error("Failed to set switch config. "
                   "Exception %s, %s, %s",
                    exc_type, exc_value, exc_traceback)
            return 0
        return 1

    def packet_out(self, data, in_port, out_port, out_queue, no_queue=0):
        """
        Sends a supplied packet out switch port(s) in specific queue.

        Set no_queue=1 if want no queueing specified (i.e. for a flooded
        packet). Also use for Zodiac FX compatibility.

        Does not use Buffer IDs as they are unreliable resource.
        """
        ofproto = self.datapath.ofproto
        parser = self.datapath.ofproto_parser
        dpid = self.datapath.id
        #*** First build OF version specific list of actions:
        if no_queue:
            #*** Packet out with no queue:
            actions = [self.datapath.ofproto_parser.OFPActionOutput \
                             (out_port, 0)]

        else:
            #*** Note: out_port must come last!
            actions = [
                    parser.OFPActionSetQueue(out_queue),
                    parser.OFPActionOutput(out_port, 0)]

        #*** Now have we have actions, build the packet out message:
        out = parser.OFPPacketOut(
                    datapath=self.datapath, buffer_id=ofproto.OFP_NO_BUFFER,
                    in_port=in_port, actions=actions, data=data)

        self.logger.debug("Sending Packet-Out message dpid=%s port=%s",
                                    dpid, out_port)
        #*** Tell the switch to send the packet:
        self.datapath.send_msg(out)

    def set_switch_table_miss(self, miss_send_len):
        """
        Set a table miss rule on table 0 to send packets to
        the controller. This is required for OF versions higher
        than v1.0
        """
        ofproto = self.datapath.ofproto
        parser = self.datapath.ofproto_parser
        dpid = self.datapath.id
        self.logger.info("Setting table-miss flow entry on switch dpid=%s with"
                                       "miss_send_len=%s", dpid, miss_send_len)
        match = parser.OFPMatch()
        actions = [parser.OFPActionOutput(ofproto.OFPP_CONTROLLER,
                                                                miss_send_len)]
        inst = [parser.OFPInstructionActions(ofproto.OFPIT_APPLY_ACTIONS,
                                                 actions)]
        mod = parser.OFPFlowMod(datapath=self.datapath, priority=0,
                                                match=match, instructions=inst)
        self.datapath.send_msg(mod)

class FlowTables(BaseClass):
    """
    This class provides an abstraction for the flow tables on
    an OpenFlow Switch
    """
    def __init__(self, config, datapath, offset):
        #*** Required for BaseClass:
        self.config = config
        #*** Set up Logging with inherited base class method:
        self.configure_logging(__name__, "switches_logging_level_s",
                                       "switches_logging_level_c")
        self.config = config
        self.datapath = datapath
        self.offset = offset
        self.dpid = datapath.id
        self.parser = datapath.ofproto_parser
        self.suppress_idle_timeout = config.get_value('suppress_idle_timeout')
        self.suppress_hard_timeout = config.get_value('suppress_hard_timeout')
        self.suppress_priority = config.get_value('suppress_priority')
        self.drop_idle_timeout = config.get_value('drop_idle_timeout')
        self.drop_hard_timeout = config.get_value('drop_hard_timeout')
        self.drop_priority = config.get_value('drop_priority')
        #*** Unique value counters for Flow Mod cookies:
        self.flow_mod_cookie_forward = 1
        self.flow_mod_cookie_reverse = offset

    def suppress_flow(self, msg, in_port, out_port, out_queue):
        """
        Add flow entries to a switch to suppress further packet-in
        events while the flow is active.

        Prefer to do fine-grained match where possible.
        Install reverse matches as well for TCP flows.

        Do not install suppression for these types of flow:
        - DNS (want to harvest identity)
        - ARP (want to harvest identity)
        - DHCP (want to harvest identity)
        - LLDP (want to harvest identity)
        """
        #*** Extract parameters:
        pkt = packet.Packet(msg.data)
        pkt_ip4 = pkt.get_protocol(ipv4.ipv4)
        pkt_ip6 = pkt.get_protocol(ipv6.ipv6)
        pkt_tcp = pkt.get_protocol(tcp.tcp)
        pkt_udp = pkt.get_protocol(udp.udp)
        idle_timeout = self.suppress_idle_timeout
        hard_timeout = self.suppress_hard_timeout
        priority = self.suppress_priority
        #*** Dict for results:
        result = {'match_type': 'ignore', 'forward_cookie': 0,
                 'forward_match': '', 'reverse_cookie': 0, 'reverse_match': '',
                 'client_ip': ''}
        self.logger.debug("event=add_flow out_queue=%s", out_queue)
        #*** Install flow entry(ies) based on type of flow:
        if pkt_tcp:
            #*** Do not suppress TCP DNS:
            if pkt_tcp.src_port == 53 or pkt_tcp.dst_port == 53:
                return result
            #*** Install two flow entries for TCP so that return traffic
            #*** is also suppressed:
            if pkt_ip4:
                forward_match = self.match_ipv4_tcp(pkt_ip4.src, pkt_ip4.dst,
                                            pkt_tcp.src_port, pkt_tcp.dst_port)
                reverse_match = self.match_ipv4_tcp(pkt_ip4.dst, pkt_ip4.src,
                                            pkt_tcp.dst_port, pkt_tcp.src_port)
            elif pkt_ip6:
                forward_match = self.match_ipv6_tcp(pkt_ip6.src, pkt_ip6.dst,
                                            pkt_tcp.src_port, pkt_tcp.dst_port)
                reverse_match = self.match_ipv6_tcp(pkt_ip6.dst, pkt_ip6.src,
                                            pkt_tcp.dst_port, pkt_tcp.src_port)
            else:
                #*** Unknown protocol so warn and exit:
                self.logger.warning("Unknown protocol, not installing flow "
                                    "suppression entries")
                return result
            #*** Actions:
            forward_actions = self.actions(out_port, out_queue)
            reverse_actions = self.actions(in_port, out_queue)
            #*** Cookies:
            forward_cookie = self.flow_mod_cookie_forward
            reverse_cookie = self.flow_mod_cookie_reverse
            #*** Now have matches and actions. Install to switch:
            self.add_flow(forward_match, forward_actions,
                                 priority=priority,
                                 idle_timeout=idle_timeout,
                                 hard_timeout=hard_timeout,
                                 cookie=forward_cookie)
            self.add_flow(reverse_match, reverse_actions,
                                 priority=priority,
                                 idle_timeout=idle_timeout,
                                 hard_timeout=hard_timeout,
                                 cookie=reverse_cookie)
            if pkt_ip4:
                #*** Convert IPv4 addrs back to dotted decimal for storing:
                forward_match['ipv4_src'] = pkt_ip4.src
                forward_match['ipv4_dst'] = pkt_ip4.dst
                reverse_match['ipv4_src'] = pkt_ip4.dst
                reverse_match['ipv4_dst'] = pkt_ip4.src
            result['match_type'] = 'dual'
            result['forward_cookie'] = forward_cookie
            result['forward_match'] = forward_match
            result['reverse_cookie'] = reverse_cookie
            result['reverse_match'] = reverse_match
            result['client_ip'] = pkt_ip4.src
            #*** Increment flow mod cookies ready for next use:
            if self.flow_mod_cookie_forward < self.offset:
                self.flow_mod_cookie_forward += 1
            else:
                self.logger.info("flow_mod_cookie_forward rolled")
                self.flow_mod_cookie_forward = 1
            self.flow_mod_cookie_reverse += 1
            return result
        else:
            if pkt_udp:
                #*** Do not suppress UDP DNS OR DHCP:
                if (pkt_udp.src_port == 53 or pkt_udp.dst_port == 53 or
                             pkt_udp.src_port == 67 or pkt_udp.dst_port == 67):
                    return result
            if pkt_ip4:
                #*** Match IPv4 packet
                match = self.match_ipv4(pkt_ip4.src, pkt_ip4.dst,
                                                                 pkt_ip4.proto)
            elif pkt_ip6:
                #*** Match IPv6 packet
                match = self.match_ipv6(pkt_ip6.src, pkt_ip6.dst)
            else:
                #*** Non-IP packet, ignore:
                return result
            #*** Actions:
            actions = self.actions(out_port, out_queue)
            #*** Cookie:
            cookie = self.flow_mod_cookie_forward
            #*** Now have matches and actions. Install to switch:
            self.add_flow(match, actions,
                                 priority=priority,
                                 idle_timeout=idle_timeout,
                                 hard_timeout=hard_timeout,
                                 cookie=cookie)
            if pkt_ip4:
                #*** Convert IPv4 addrs back to dotted decimal for storing:
                match['ipv4_src'] = pkt_ip4.src
                match['ipv4_dst'] = pkt_ip4.dst
            result['match_type'] = 'single'
            result['forward_cookie'] = cookie
            result['forward_match'] = match
            result['client_ip'] = pkt_ip4.src
            #*** Increment flow mod cookie ready for next use:
            self.flow_mod_cookie_forward += 1
            return result

    def drop_flow(self, msg):
        """
        Add flow entry to a switch to suppress further packet-in
        events for a particular flow.

        Prefer to do fine-grained match where possible.

        TCP or UDP source ports are not matched as ephemeral
        """
        #*** Extract parameters:
        pkt = packet.Packet(msg.data)
        pkt_ip4 = pkt.get_protocol(ipv4.ipv4)
        pkt_ip6 = pkt.get_protocol(ipv6.ipv6)
        pkt_tcp = pkt.get_protocol(tcp.tcp)
        pkt_udp = pkt.get_protocol(udp.udp)
        idle_timeout = self.drop_idle_timeout
        hard_timeout = self.drop_hard_timeout
        priority = self.drop_priority
        #*** Dict for results:
        result = {'match_type': 'ignore', 'forward_cookie': 0,
                 'forward_match': '', 'reverse_cookie': 0, 'reverse_match': '',
                 'client_ip': ''}
        self.logger.debug("event=drop_flow")
        #*** Drop action is the implicit in setting no actions:
        drop_action = 0
        #*** Install flow entry based on type of flow:
        if pkt_tcp:
            if pkt_ip4:
                drop_match = self.match_ipv4_tcp(pkt_ip4.src, pkt_ip4.dst,
                                            0, pkt_tcp.dst_port)
            elif pkt_ip6:
                drop_match = self.match_ipv6_tcp(pkt_ip6.src, pkt_ip6.dst,
                                            0, pkt_tcp.dst_port)
            else:
                #*** Unknown protocol so warn and exit:
                self.logger.warning("Unknown IP protocol, not installing flow "
                                    "drop entry for TCP")
                return result
        elif pkt_udp:
            if pkt_ip4:
                drop_match = self.match_ipv4_udp(pkt_ip4.src, pkt_ip4.dst,
                                            0, pkt_udp.dst_port)
            elif pkt_ip6:
                drop_match = self.match_ipv6_udp(pkt_ip6.src, pkt_ip6.dst,
                                            0, pkt_udp.dst_port)
            else:
                #*** Unknown protocol so warn and exit:
                self.logger.warning("Unknown IP protocol, not installing flow "
                                    "drop entry for UDP")
                return result
        elif pkt_ip4:
            #*** Match IPv4 packet
            drop_match = self.match_ipv4(pkt_ip4.src, pkt_ip4.dst,
                                                                 pkt_ip4.proto)
        elif pkt_ip6:
            #*** Match IPv6 packet
            drop_match = self.match_ipv6(pkt_ip6.src, pkt_ip6.dst)
        else:
            #*** Non-IP packet, ignore:
            self.logger.warning("Drop not installed as non-IP")
            return 0
        #*** Cookie:
        cookie = self.flow_mod_cookie_forward
        #*** Now have match and action. Install to switch:
        self.logger.debug("Installing drop rule to dpid=%s", self.dpid)
        self.add_flow(drop_match, drop_action, priority=priority,
                          idle_timeout=idle_timeout, hard_timeout=hard_timeout,
                          cookie=cookie)
        result['match_type'] = 'single'
        result['forward_cookie'] = cookie
        if pkt_ip4:
            #*** Convert IPv4 addrs back to dotted decimal for storing:
            drop_match['ipv4_src'] = pkt_ip4.src
            drop_match['ipv4_dst'] = pkt_ip4.dst
            result['client_ip'] = pkt_ip4.src
        result['forward_match'] = drop_match
        #*** Increment flow mod cookie ready for next use:
        self.flow_mod_cookie_forward += 1
        return result

    def add_flow(self, match_d, actions, priority, idle_timeout, hard_timeout,
                    cookie):
        """
        Add a flow entry to a switch
        """
        #*** Convert match dict to an OFPMatch object:
        match = self.parser.OFPMatch(**match_d)
        ofproto = self.datapath.ofproto
        parser = self.datapath.ofproto_parser
        if actions:
            inst = [parser.OFPInstructionActions(ofproto.OFPIT_APPLY_ACTIONS,
                                                                      actions)]
        else:
            inst = []
        mod = parser.OFPFlowMod(datapath=self.datapath,
                                cookie=cookie,
                                idle_timeout=idle_timeout,
                                hard_timeout=hard_timeout,
                                priority=priority,
                                flags=ofproto.OFPFF_SEND_FLOW_REM,
                                match=match,
                                instructions=inst)
        self.logger.debug("Installing Flow Entry to dpid=%s match=%s",
                                    self.dpid, match)
        self.datapath.send_msg(mod)

    def actions(self, out_port, out_queue, no_queue=0):
        """
        Create actions for a switch flow entry. Specify the out port
        and QoS queue, and set no_queue=1 if don't want QoS set.
        Returns a list of action objects
        """
        if no_queue:
            #*** Set flow entry action without queueing specified:
            return [self.datapath.ofproto_parser.OFPActionOutput(out_port, 0)]
        else:
            return [self.datapath.ofproto_parser.OFPActionSetQueue(out_queue),
                    self.datapath.ofproto_parser.OFPActionOutput(out_port, 0)]

    def match_ipv4_tcp(self, ipv4_src, ipv4_dst, tcp_src, tcp_dst):
        """
        Match an IPv4 TCP flow on a switch.
        Passed IPv4 and TCP parameters and return
        an OpenFlow match object for this flow
        """
        if tcp_src:
            return dict(eth_type=0x0800,
                    ipv4_src=_ipv4_t2i(str(ipv4_src)),
                    ipv4_dst=_ipv4_t2i(str(ipv4_dst)),
                    ip_proto=6,
                    tcp_src=tcp_src,
                    tcp_dst=tcp_dst)
        else:
            return dict(eth_type=0x0800,
                    ipv4_src=_ipv4_t2i(str(ipv4_src)),
                    ipv4_dst=_ipv4_t2i(str(ipv4_dst)),
                    ip_proto=6,
                    tcp_dst=tcp_dst)

    def match_ipv4_udp(self, ipv4_src, ipv4_dst, udp_src, udp_dst):
        """
        Match an IPv4 UDP flow on a switch.
        Passed IPv4 and UDP parameters and return
        an OpenFlow match object for this flow
        """
        if udp_src:
            return dict(eth_type=0x0800,
                    ipv4_src=_ipv4_t2i(str(ipv4_src)),
                    ipv4_dst=_ipv4_t2i(str(ipv4_dst)),
                    ip_proto=17,
                    udp_src=udp_src,
                    udp_dst=udp_dst)
        else:
            return dict(eth_type=0x0800,
                    ipv4_src=_ipv4_t2i(str(ipv4_src)),
                    ipv4_dst=_ipv4_t2i(str(ipv4_dst)),
                    ip_proto=17,
                    udp_dst=udp_dst)

    def match_ipv6_tcp(self, ipv6_src, ipv6_dst, tcp_src, tcp_dst):
        """
        Match an IPv6 TCP flow on a switch.
        Passed IPv6 and TCP parameters and return
        an OpenFlow match object for this flow
        """
        if tcp_src:
            return dict(eth_type=0x86DD,
                    ipv6_src=ipv6_src,
                    ipv6_dst=ipv6_dst,
                    ip_proto=6,
                    tcp_src=tcp_src,
                    tcp_dst=tcp_dst)
        else:
            return dict(eth_type=0x86DD,
                    ipv6_src=ipv6_src,
                    ipv6_dst=ipv6_dst,
                    ip_proto=6,
                    tcp_dst=tcp_dst)

    def match_ipv6_udp(self, ipv6_src, ipv6_dst, udp_src, udp_dst):
        """
        Match an IPv6 UDP flow on a switch.
        Passed IPv6 and UDP parameters and return
        an OpenFlow match object for this flow
        """
        if udp_src:
            return dict(eth_type=0x86DD,
                    ipv6_src=ipv6_src,
                    ipv6_dst=ipv6_dst,
                    ip_proto=17,
                    udp_src=udp_src,
                    udp_dst=udp_dst)
        else:
            return dict(eth_type=0x86DD,
                    ipv6_src=ipv6_src,
                    ipv6_dst=ipv6_dst,
                    ip_proto=17,
                    udp_dst=udp_dst)

    def match_ipv4(self, ipv4_src, ipv4_dst, ip_proto):
        """
        Match an IPv4 flow on a switch.
        Passed IPv4 parameters and return
        an OpenFlow match object for this flow
        """
        return dict(eth_type=0x0800,
                    ipv4_src=_ipv4_t2i(str(ipv4_src)),
                    ipv4_dst=_ipv4_t2i(str(ipv4_dst)),
                    ip_proto=ip_proto)

    def match_ipv6(self, ipv6_src, ipv6_dst):
        """
        Match an IPv6 flow on a switch.
        Passed IPv6 parameters and return
        an OpenFlow match object for this flow
        """
        return dict(eth_type=0x86DD,
                    ipv6_src=ipv6_src,
                    ipv6_dst=ipv6_dst)

#=============== Private functions:

def _ipv4_t2i(ip_text):
    """
    Turns an IPv4 address in text format into an integer.
    Borrowed from rest_router.py code
    """
    if ip_text == 0:
        return ip_text
    assert isinstance(ip_text, str)
    return struct.unpack('!I', addrconv.ipv4.text_to_bin(ip_text))[0]


from sqlite3 import adapters
from numpy.core.shape_base import stack
import torch
import torch.nn.functional as F
from torch import nn
torch.backends.cudnn.benchmark = True
from einops.einops import rearrange
from einops.layers.torch import Rearrange
import numpy as np

from util import box_ops
from util.misc import (NestedTensor, nested_tensor_from_tensor_list,
                       accuracy, get_world_size, interpolate,
                       is_dist_avail_and_initialized)

from .convnext import build_convbackbone, LayerNorm
from .matcher import build_matcher

from .rpet_pose import (RPETpose, PostProcessPoseDR, PostProcessPoseHM, NMTNORMCritierion, filter_target_simdr)
from .transformer import build_transformer

class RPET(nn.Module):
    """ This is the Person Detection Network module that performs perrson detection """
    def __init__(self, backbone, transformer, num_classes, num_queries,
                    num_txrx=8, aux_loss=False, box_feature='x', signal_ratio=0.75):

        super().__init__()
        self.num_queries = num_queries
        self.transformer = transformer
        hidden_dim = transformer.d_model
        print("transformer.d_model = ", hidden_dim)

        self.class_embed = nn.Linear(hidden_dim, num_classes + 1)
        self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
        self.query_embed = nn.Embedding(num_queries, hidden_dim)

        self.signal_ratio=signal_ratio
        self.hidden_dim = hidden_dim
        self.input_proj_dim = hidden_dim if box_feature =='x' else int(hidden_dim * self.signal_ratio)#75)
        self.input_proj = nn.Sequential(
                LayerNorm(backbone.num_channels, eps=1e-6, data_format="channels_first"),
                nn.Conv1d(backbone.num_channels, self.input_proj_dim, kernel_size=1),
        ) if self.input_proj_dim > 0. else None

        self.backbone = backbone
        self.aux_loss = aux_loss
        self.stack_num = backbone.stack_num
        self.frame_skip = backbone.frame_skip

        self.num_txrx = num_txrx
        self.box_feature = box_feature

        if box_feature != 'x':
            print("use visual clue")
            vc_input_chan = hidden_dim - self.input_proj_dim #64
            self.vc_input_proj = nn.Sequential(
                nn.Conv2d(256, vc_input_chan, kernel_size=1),
                nn.BatchNorm2d(vc_input_chan, momentum=0.1),
            )

        dummy_input = torch.zeros((128, self.stack_num//self.frame_skip, num_txrx**2, 768))

        self.check_dim(dummy_input)

    def forward(self, samples, targets=None):
        """ The forward expects a NestedTensor, which consists of:
               - samples.tensor: batched images, of shape [batch_size x 3 x H x W]
               - samples.mask: a binary mask of shape [batch_size x H x W], containing 1 on padded pixels

            It returns a dict with the following elements:
               - "pred_logits": the classification logits (including no-object) for all queries.
                                Shape= [batch_size x num_queries x (num_classes + 1)]
               - "pred_boxes": The normalized boxes coordinates for all queries, represented as
                               (center_x, center_y, height, width). These values are normalized in [0, 1],
                               relative to the size of each individual image (disregarding possible padding).
                               See PostProcess for information on how to retrieve the unnormalized bounding box.
               - "aux_outputs": Optional, only returned when auxilary losses are activated. It is a list of
                                dictionnaries containing the two above keys for each decoder layer.
        """
        pos, vc_query = None, None
        backbone_out = self.backbone(samples)
        src, mask, pos, vc = backbone_out['src'], backbone_out['mask'], \
                        backbone_out['pos'], backbone_out['pred_feature']
        assert mask is not None

        src = src.flatten(2)
        if self.input_proj != None:
            input_proj = self.input_proj(src)
            input_proj = rearrange(input_proj, 'b n (t1 t2) -> b n t1 t2', t1=16)

        if vc is not None:
            vc_query = self.vc_input_proj(vc)
            input_proj = torch.cat((input_proj, vc_query), dim=1) if self.input_proj != None else vc_query

        hs = self.transformer(input_proj, mask, self.query_embed.weight, pos)[0]
        hs = hs[-1]

        outputs_class = self.class_embed(hs)
        outputs_coord = self.bbox_embed(hs).sigmoid()
        out = {'pred_logits': outputs_class, 'pred_boxes': outputs_coord}

        if self.box_feature != 'x':
            out['pred_feature'] = vc

        if self.aux_loss:
            out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord)

        return out

    def check_dim(self, samples): #samples: NestedTensor):
        print("____check dimension in RPET_____")
        print(f"input shape : {samples.shape}")

        #features= self.backbone(samples)
        pos, vc_query= None, None
        backbone_out = self.backbone(samples)
        src, mask, pos, vc = backbone_out['src'], backbone_out['mask'], \
                        backbone_out['pos'], backbone_out['pred_feature']
        print(f"backbone = src : {src.shape}, mask : {mask.shape}")
        if pos is not None:
            print("position embedding = ", pos.shape)
        if vc is not None:
            print("visual clue = ", vc.shape)
            vc_query = self.vc_input_proj(vc)
            print("visual clue to patch = ", vc_query.shape)
        assert mask is not None

        src = src.flatten(2)
        if self.input_proj != None:
            input_proj = self.input_proj(src)
            input_proj = rearrange(input_proj, 'b n (t1 t2) -> b n t1 t2', t1=16)
            if vc is not None:
                print(f"input_proj[{input_proj.shape}] + vc_query[{vc_query.shape}]" )
        else:
            print(f"input_proj[X] + vc_query[{vc_query.shape}]" )

        if vc is not None:
            input_proj = torch.cat((input_proj, vc_query), dim=1) if self.input_proj != None else vc_query

        print(f"transofmer input self.input_proj(src) : {input_proj.shape}, mask : {mask.shape}", \
                f"\n self.query_embed.weight {self.query_embed.weight.shape}")
        hs = self.transformer(input_proj, mask, self.query_embed.weight, pos)[0]
        print(f"hs(decoder) = {hs.shape}")

        hs = hs[-1]
        outputs_class = self.class_embed(hs)
        outputs_coord = self.bbox_embed(hs).sigmoid()
        print("after ffn ", outputs_class.shape, outputs_coord.shape)

    @torch.jit.unused
    def _set_aux_loss(self, outputs_class, outputs_coord):
        # this is a workaround to make torchscript happy, as torchscript
        # doesn't support dictionary with non-homogeneous values, such
        # as a dict having both a Tensor and a list.
        return [{'pred_logits': a, 'pred_boxes': b}
                for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]

def _expand(tensor, length: int):
    return tensor.unsqueeze(1).repeat(1, int(length), 1, 1, 1).flatten(0, 1)

class SetCriterion(nn.Module):
    """ This class computes the loss for DETR.
    The process happens in two steps:
        1) we compute hungarian assignment between ground truth boxes and the outputs of the model
        2) we supervise each pair of matched ground-truth / prediction (supervise class and box)
    """
    def __init__(self, num_classes, matcher, weight_dict, eos_coef, losses, device=None, pose_method=None):
        """ Create the criterion.
        Parameters:
            num_classes: number of object categories, omitting the special no-object category
            matcher: module able to compute a matching between targets and proposals
            weight_dict: dict containing as key the names of the losses and as values their relative weight.
            eos_coef: relative classification weight applied to the no-object category
            losses: list of all the losses to be applied. See get_loss for list of available losses.
        """
        super().__init__()
        self.num_classes = num_classes
        self.matcher = matcher
        self.weight_dict = weight_dict
        self.eos_coef = eos_coef
        self.losses = losses
        empty_weight = torch.ones(self.num_classes + 1)
        empty_weight[-1] = self.eos_coef
        self.register_buffer('empty_weight', empty_weight)
        if pose_method == 'hm' or pose_method =='hmdr':
            self.hm_criterion = nn.MSELoss(reduction='mean').to(device)
        if pose_method == 'simdr' or pose_method =='hmdr':
            self.cd_criterion = NMTNORMCritierion(label_smoothing=0.1)

        self.device = device
        if 'feature' in losses:
            self.feature_criterion = nn.MSELoss(reduction='mean').to(device)

        self.pose_method=pose_method

    def loss_labels(self, outputs, targets, indices, num_boxes, log=True):
        """Classification loss (NLL)
        targets dicts must contain the key "labels" containing a tensor of dim [nb_target_boxes]
        """
        assert 'pred_logits' in outputs
        src_logits = outputs['pred_logits']
        idx = self._get_src_permutation_idx(indices)
        target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)])
        target_classes = torch.full(src_logits.shape[:2], self.num_classes,
                                    dtype=torch.int64, device=src_logits.device)
        target_classes[idx] = target_classes_o
        loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight)
        losses = {'loss_ce': loss_ce}

        if log:
            losses['class_error'] = 100 - accuracy(src_logits[idx][..., :-1], target_classes_o)[0]
        return losses

    @torch.no_grad()
    def loss_cardinality(self, outputs, targets, indices, num_boxes):
        """ Compute the cardinality error, ie the absolute error in the number of predicted non-empty boxes
        This is not really a loss, it is intended for logging purposes only. It doesn't propagate gradients
        """
        pred_logits = outputs['pred_logits']
        device = pred_logits.device
        tgt_lengths = torch.as_tensor([len(v["labels"]) for v in targets], device=device)
        card_pred = (pred_logits.argmax(-1) != pred_logits.shape[-1] - 1).sum(1)
        card_err = F.l1_loss(card_pred.float(), tgt_lengths.float())
        losses = {'cardinality_error': card_err}
        return losses

    def loss_boxes(self, outputs, targets, indices, num_boxes):
        """Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss
           targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]
           The target boxes are expected in format (center_x, center_y, w, h), normalized by the image size.
        """
        assert 'pred_boxes' in outputs
        idx = self._get_src_permutation_idx(indices)
        src_boxes = outputs['pred_boxes'][idx]
        target_boxes = torch.cat([t['boxes'][i] for t, (_, i) in zip(targets, indices)], dim=0)

        loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction='none')

        losses = {}
        losses['loss_bbox'] = loss_bbox.sum() / num_boxes

        loss_giou = 1 - torch.diag(box_ops.generalized_box_iou(
            box_ops.box_cxcywh_to_xyxy(src_boxes),
            box_ops.box_cxcywh_to_xyxy(target_boxes)))

        losses['loss_giou'] = loss_giou.sum() / num_boxes
        return losses

    def loss_pose_dr(self, outputs, targets, indices, num_boxes):
        assert 'x_coord' in outputs

        idx = self._get_src_permutation_idx(indices)
        src_x = outputs['x_coord'][idx]
        src_y = outputs['y_coord'][idx]
        output_size = src_x.shape[2]
        target_pose = torch.cat([t["cd"][J] for t, (_, J) in zip(targets, indices)])
        target_coord = target_pose[:,:,:2]
        target_vis = target_pose[:,:,2]

        target_pose = torch.round(target_coord * output_size)
        target_weight, filtered_joints = filter_target_simdr(target_pose.clone(), target_vis.clone(), output_size)
        target_weight = target_weight.to(src_x)
        filtered_joints = filtered_joints.to(src_x).type(torch.int64)
        loss = self.cd_criterion(src_x, src_y, filtered_joints, target_weight)
        losses = {}
        losses['loss_simdr'] = loss
        return losses

    def loss_pose_hm(self, outputs, targets, indices, num_boxes):
        assert 'pred_hm' in outputs

        src_idx = self._get_src_permutation_idx(indices)
        tgt_idx = self._get_tgt_permutation_idx(indices)
        src_pose = outputs['pred_hm'][src_idx]  #(idx_num, 17, 256)

        tgt_pose = torch.cat([t["hm"][J] for t, (_, J) in zip(targets, indices)])
        tgt_pose = tgt_pose.to(src_pose)
        num_joints = src_pose.size(1)
        heatmaps_pred = src_pose.split(1,1)
        heatmaps_gt = tgt_pose.split(1,1)
        loss = 0

        for j in range(num_joints):
            heatmap_pred = heatmaps_pred[j].squeeze()
            heatmap_gt = heatmaps_gt[j].squeeze()
            tmp_loss = 0.5 * self.hm_criterion(heatmap_pred, heatmap_gt)
            loss += tmp_loss

        loss = loss

        losses = {'loss_hm': loss}
        return losses

    def loss_feature(self, outputs, targets, indices, num_boxes):
        assert 'pred_feature' in outputs

        src_feature = outputs['pred_feature']
        batch_size = src_feature.size(0)

        tgt_feature = torch.cat([t["features"] for t in targets])
        tgt_feature = tgt_feature.to(src_feature)
        losses = {}

        features_pred = src_feature.reshape((batch_size, -1))
        features_gt = tgt_feature.reshape((batch_size, -1))
        feature_loss = self.feature_criterion(features_pred, features_gt)

        losses['loss_feature'] = feature_loss #/ (x*y)
        if indices == None:
            losses['class_error'] = 0.

        return losses

    def _get_src_permutation_idx(self, indices):
        batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
        src_idx = torch.cat([src for (src, _) in indices])
        return batch_idx, src_idx

    def _get_tgt_permutation_idx(self, indices):
        batch_idx = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])
        tgt_idx = torch.cat([tgt for (_, tgt) in indices])
        return batch_idx, tgt_idx

    def get_loss(self, loss, outputs, targets, indices, num_boxes, **kwargs):
        loss_map = {
            'labels': self.loss_labels,
            'cardinality': self.loss_cardinality,
            'boxes': self.loss_boxes,
            'posedr' : self.loss_pose_dr,
            'posehm' : self.loss_pose_hm,
            'feature': self.loss_feature,
        }
        assert loss in loss_map, f'do you really want to compute {loss} loss?'
        return loss_map[loss](outputs, targets, indices, num_boxes, **kwargs)

    def forward(self, outputs, targets):
        """ This performs the loss computation.
        Parameters:
             outputs: dict of tensors, see the output specification of the model for the format
             targets: list of dicts, such that len(targets) == batch_size.
                      The expected keys in each dict depends on the losses applied, see each loss' doc
        """
        outputs_without_aux = {k: v for k, v in outputs.items() if k != 'aux_outputs'}

        if 'pred_logits' in outputs:
            indices = self.matcher(outputs_without_aux, targets)

            # Compute the average number of target boxes accross all nodes, for normalization purposes
            num_boxes = sum(len(t["labels"]) for t in targets)
            num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device)
            if is_dist_avail_and_initialized():
                torch.distributed.all_reduce(num_boxes)
            num_boxes = torch.clamp(num_boxes / get_world_size(), min=1).item()
        else:
            indices = None
            num_boxes = 0
        # Compute all the requested losses
        losses = {}
        for loss in self.losses:
            losses.update(self.get_loss(loss, outputs, targets, indices, num_boxes))

        # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
        if 'aux_outputs' in outputs:
            for i, aux_outputs in enumerate(outputs['aux_outputs']):
                indices = self.matcher(aux_outputs, targets)
                for loss in self.losses:
                    if loss == 'posehm' or loss =='posedr':
                        # Intermediate masks losses are too costly to compute, we ignore them.
                        continue
                    kwargs = {}
                    if loss == 'labels':
                        # Logging is enabled only for the last layer
                        kwargs = {'log': False}
                    l_dict = self.get_loss(loss, aux_outputs, targets, indices, num_boxes, **kwargs)
                    l_dict = {k + f'_{i}': v for k, v in l_dict.items()}
                    losses.update(l_dict)

        return losses

def soft_nms_pytorch(dets, box_scores, sigma=0.5, thresh=0., cuda=1):
    # Indexes concatenate boxes with the last column
    N = dets.shape[0]
    if cuda:
        #indexes = torch.arange(0, N, dtype=torch.float).cuda().view(N, 1)
        indexes = torch.arange(0, N, dtype=torch.float).view(N, 1).to(dets)
    else:
        indexes = torch.arange(0, N, dtype=torch.float).view(N, 1)
    dets = torch.cat((dets, indexes), dim=1)
    #print(dets)
    # The order of boxes coordinate is [y1,x1,y2,x2]
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = box_scores
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    for i in range(N):
        # intermediate parameters for later parameters exchange
        tscore = scores[i].clone()
        pos = i + 1

        if i != N - 1:
            maxscore, maxpos = torch.max(scores[pos:], dim=0)
            if tscore < maxscore:
                dets[i], dets[maxpos.item() + i + 1] = dets[maxpos.item() + i + 1].clone(), dets[i].clone()
                scores[i], scores[maxpos.item() + i + 1] = scores[maxpos.item() + i + 1].clone(), scores[i].clone()
                areas[i], areas[maxpos + i + 1] = areas[maxpos + i + 1].clone(), areas[i].clone()

        # IoU calculate
        yy1 = np.maximum(dets[i, 0].to("cpu").numpy(), dets[pos:, 0].to("cpu").numpy())
        xx1 = np.maximum(dets[i, 1].to("cpu").numpy(), dets[pos:, 1].to("cpu").numpy())
        yy2 = np.minimum(dets[i, 2].to("cpu").numpy(), dets[pos:, 2].to("cpu").numpy())
        xx2 = np.minimum(dets[i, 3].to("cpu").numpy(), dets[pos:, 3].to("cpu").numpy())

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = torch.tensor(w * h).to(dets) if cuda else torch.tensor(w * h)
        ovr = torch.div(inter, (areas[i] + areas[pos:] - inter))

        # Gaussian decay
        weight = torch.exp(-(ovr * ovr) / sigma)
        scores[pos:] = weight * scores[pos:]

    # select the boxes and keep the corresponding indexes
    #keep = dets[:, 4][scores > thresh].int()
    keep = dets[:, 4].int()

    return dets[:, :4], scores, keep

class PostProcess(nn.Module):
    """ This module converts the model's output into the format expected by the coco api"""
    @torch.no_grad()
    def forward(self, outputs, target_sizes, nms=False):
        """ Perform the computation
        Parameters:
            outputs: raw outputs of the model
            target_sizes: tensor of dimension [batch_size x 2] containing the size of each images of the batch
                          For evaluation, this must be the original image size (before any data augmentation)
                          For visualization, this should be the image size after data augment, but before padding
        """
        out_logits, out_bbox = outputs['pred_logits'], outputs['pred_boxes']

        assert len(out_logits) == len(target_sizes)
        assert target_sizes.shape[1] == 2

        #batch_size = out_bbox.shape[0]
        prob = F.softmax(out_logits, -1)
        #print(prob.shape)
        #print(prob[..., :-1].shape)
        scores, labels = prob[..., :-1].max(-1) # No Obejct는 제외함

        # convert to [x0, y0, x1, y1] format
        boxes = box_ops.box_cxcywh_to_xyxy(out_bbox)
        # and from relative [0, 1] to absolute [0, height] coordinates
        img_h, img_w = target_sizes.unbind(1)
        scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
        boxes = boxes * scale_fct[:, None, :]

        batch_size, N = boxes.shape[0], boxes.shape[1]
        indices = torch.zeros((batch_size, N))
        if nms:
            for b in range(batch_size):
                dets = boxes[b]
                score = scores[b]
                #print(boxes[b], scores[b])
                boxes[b], scores[b], keep = soft_nms_pytorch(dets, score)
                labels[b] = torch.index_select(labels[b], 0, keep)
                indices[b] = keep
        else:
            for b in range(batch_size):
                indices[b] = torch.arange(0, N, dtype=torch.float)
        #print(indices)
        results = [{'scores': s, 'labels': l, 'boxes': b, 'indices': idx} for s, l, b, idx in zip(scores, labels, boxes, indices)]

        return results

class MLP(nn.Module):
    """ Very simple multi-layer perceptron (also called FFN)"""

    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
        super().__init__()
        self.num_layers = num_layers
        h = [hidden_dim] * (num_layers - 1)
        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
        return x

def build_rpet(args):
    num_classes = 91
    device = torch.device(args.device)
    backbone = build_convbackbone(args)

    transformer = build_transformer(args)

    model = RPET(
        backbone,
        transformer,
        num_classes=num_classes,
        num_queries=args.num_queries,
        aux_loss=args.aux_loss,
        box_feature=args.box_feature,
        signal_ratio=args.signal_ratio
    )

    if args.pose is not None:
        model = RPETpose(model, freeze_rpet=(args.frozen_weights is not None),
                            method=args.pose, dr_size=args.dr_size)

    matcher = build_matcher(args)
    weight_dict = {'loss_ce': 1, 'loss_bbox': args.bbox_loss_coef}
    weight_dict['loss_giou'] = args.giou_loss_coef

    if args.pose == 'simdr':
        weight_dict["loss_simdr"] = args.pose_loss_coef
    elif args.pose =='hm':
        weight_dict["loss_hm"] = args.pose_loss_coef * 5
    elif args.pose =='hmdr':
        weight_dict["loss_simdr"] = args.pose_loss_coef
        weight_dict["loss_hm"] = args.pose_loss_coef

    if args.box_feature == '16' and args.pose == None:
        weight_dict['loss_feature'] = args.feature_loss_coef

    if args.aux_loss:
        aux_weight_dict = {}
        for i in range(args.dec_layers - 1):
            aux_weight_dict.update({k + f'_{i}': v for k, v in weight_dict.items()})
        weight_dict.update(aux_weight_dict)

    losses = ['labels', 'boxes', 'cardinality']

    if args.pose == 'simdr':
        losses += ['posedr']
    elif args.pose == 'hm':
        losses += ['posehm']
    elif args.pose =='hmdr':
        losses += ['posedr']
        losses += ['posehm']

    if args.box_feature == '16' and args.pose == None:
        losses += ['feature']

    criterion = SetCriterion(num_classes, matcher=matcher, weight_dict=weight_dict,
                            eos_coef=args.eos_coef, losses=losses, device=device, pose_method=args.pose)
    criterion.to(device)
    postprocessors = {'bbox': PostProcess()}
    if args.pose == 'simdr':
        postprocessors['posedr'] = PostProcessPoseDR()
    elif args.pose == 'hm':
        postprocessors['posehm'] = PostProcessPoseHM()
    elif args.pose == 'hmdr':
        postprocessors['posedr'] = PostProcessPoseDR()

    print(f"losses = {losses}")
    print(f"weight_dict = {weight_dict}")
    print(f"postprocessors = {postprocessors.keys()}")

    return model, criterion, postprocessors


from yt_env_setup import (YTEnvSetup, Restarter, NODES_SERVICE)

from yt_helpers import profiler_factory

from yt_commands import (
    ls, get, set, print_debug, authors, wait, run_test_vanilla, create_user,
    wait_breakpoint, with_breakpoint, release_breakpoint)

from yt.common import YtError, update_inplace
from yt.wrapper import YtClient

import pytest

import datetime
from hashlib import sha1
import os.path
import re
import shutil
import requests

##################################################################

class JobProxyHider(object):
    def __init__(self, env_setup):
        self.bin_path = env_setup.bin_path

    def _job_proxy_path(self):
        return os.path.join(self.bin_path, "ytserver-job-proxy")

    def __enter__(self):
        job_proxy_path_hidden = self._job_proxy_path() + ".hidden"
        print_debug("Hiding {} to {}".format(self._job_proxy_path(), job_proxy_path_hidden))
        shutil.move(self._job_proxy_path(), job_proxy_path_hidden)

    def __exit__(self, exc_type, exc_val, exc_tb):
        job_proxy_path_hidden = self._job_proxy_path() + ".hidden"
        print_debug("Unhiding {} to {}".format(job_proxy_path_hidden, self._job_proxy_path()))
        shutil.move(job_proxy_path_hidden, self._job_proxy_path())

class JobProxyTestBase(YTEnvSetup):
    NUM_MASTERS = 1
    NUM_NODES = 1
    NUM_SCHEDULERS = 1
    ENABLE_HTTP_PROXY = True
    NUM_HTTP_PROXIES = 1

    DELTA_CONTROLLER_AGENT_CONFIG = {
        "controller_agent": {
            "safe_online_node_count": 2,
        },
    }

class TestJobProxyBinary(JobProxyTestBase):
    DELTA_NODE_CONFIG = {
        "exec_node": {
            "job_controller": {
                "job_proxy_build_info_update_period": 300,
            }
        }
    }

    @authors("galtsev")
    def test_orchid_build_info(self):
        def check_iso8601_date(string):
            date = datetime.datetime.fromisoformat(string.rstrip("Z"))
            assert date > datetime.datetime.fromisoformat("2010-02-02")

        n = ls("//sys/cluster_nodes")[0]
        service_path = f"//sys/cluster_nodes/{n}/orchid/service/"

        assert get(service_path + "name") == "node"
        assert len(get(service_path + "build_host")) > 0
        assert re.match(r"^[0-9]+\.[0-9]", get(service_path + "version"))

        for time in ("build_time", "start_time"):
            check_iso8601_date(get(service_path + time))

    @authors("galtsev")
    def test_job_proxy_build_attribute(self):
        n = ls("//sys/cluster_nodes")[0]
        path = f"//sys/cluster_nodes/{n}"
        attribute_name = "job_proxy_build_version"
        attribute_path = f"{path}/@{attribute_name}"
        orchid_path = f"{path}/orchid/exec_node/job_controller/job_proxy_build/version"

        assert attribute_name in ls(f"{path}/@")

        assert get(attribute_path) == get(orchid_path)

        assert re.match(r"^[0-9]+\.[0-9]", get(attribute_path))

    @authors("max42")
    def test_job_proxy_build_info(self):
        n = ls("//sys/cluster_nodes")[0]
        orchid_path = "//sys/cluster_nodes/{}/orchid".format(n)

        expected_version = get(orchid_path + "/service/version")

        def check_build(build, expect_error):
            if expect_error:
                return "error" in build
            else:
                return "error" not in build and "version" in build and \
                       build["version"] == expected_version

        def check_direct(expect_error):
            job_proxy_build = get(orchid_path + "/exec_node/job_controller/job_proxy_build")
            return check_build(job_proxy_build, expect_error)

        def check_discover_versions(expect_error):
            url = "http://" + self.Env.get_proxy_address() + "/internal/discover_versions/v2"
            print_debug("HTTP GET", url)
            rsp = requests.get(url).json()
            print_debug(rsp)
            job_proxies = [instance for instance in rsp["details"] if instance["type"] == "job_proxy"]
            assert len(job_proxies) == 1
            return check_build(job_proxies[0], expect_error)

        # At the beginning job proxy version should be immediately visible.
        assert check_direct(False)
        assert check_discover_versions(False)

        with JobProxyHider(self):
            wait(lambda: check_direct(True))
            wait(lambda: check_discover_versions(True))

        wait(lambda: check_direct(False))
        wait(lambda: check_discover_versions(False))

    @authors("max42")
    def test_slot_disabling_on_unavailable_job_proxy(self):
        wait(lambda: get("//sys/scheduler/orchid/scheduler/cluster/total_node_count") == 1)
        wait(lambda: get("//sys/scheduler/orchid/scheduler/cluster/resource_limits/user_slots") == 1)

        with JobProxyHider(self):
            with Restarter(self.Env, NODES_SERVICE):
                wait(lambda: get("//sys/scheduler/orchid/scheduler/cluster/total_node_count") == 0)
                wait(lambda: get("//sys/scheduler/orchid/scheduler/cluster/resource_limits/user_slots") == 0)
            wait(lambda: get("//sys/scheduler/orchid/scheduler/cluster/total_node_count") == 1)
            wait(lambda: get("//sys/scheduler/orchid/scheduler/cluster/resource_limits/user_slots") == 0)

        wait(lambda: get("//sys/scheduler/orchid/scheduler/cluster/total_node_count") == 1)
        wait(lambda: get("//sys/scheduler/orchid/scheduler/cluster/resource_limits/user_slots") == 1)

    @authors("alex-shishkin")
    def test_rpc_proxy_socket_path_env_variable(self):
        op = run_test_vanilla(with_breakpoint("echo $YT_JOB_PROXY_SOCKET_PATH >&2; BREAKPOINT"))
        job_id = wait_breakpoint()[0]
        content = op.read_stderr(job_id).decode("ascii").strip()
        release_breakpoint()
        assert re.match(r"^.*/pipes/.*-job-proxy-[0-9]+$", content)

class TestRpcProxyInJobProxy(YTEnvSetup):
    NUM_MASTERS = 1
    NUM_NODES = 1
    NUM_SCHEDULERS = 1
    ENABLE_RPC_PROXY = True
    NUM_RPC_PROXIES = 1
    ENABLE_HTTP_PROXY = False

    DELTA_NODE_CONFIG = {
        "exec_node": {
            "job_proxy_authentication_manager": {
                "enable_authentication": True,
                "cypress_token_authenticator": {
                    "secure": True,
                },
            },
        }
    }

    def setup_method(self, method):
        super(TestRpcProxyInJobProxy, self).setup_method(method)
        create_user("tester_name")
        set("//sys/tokens/" + sha1(b"tester_token").hexdigest(), "tester_name")

    def create_client_from_uds(self, socket_file, config=None):
        default_config = {
            "backend": "rpc",
            "driver_config": {
                "connection_type": "rpc",
                "proxy_unix_domain_socket": socket_file,
            },
        }
        if config:
            update_inplace(default_config, config)
        return YtClient(proxy=None, config=default_config)

    def run_job_proxy(self, enable_rpc_proxy, rpc_proxy_thread_pool_size=None,  time_limit=2000):
        task_patch = {
            "enable_rpc_proxy_in_job_proxy": enable_rpc_proxy,
        }
        if rpc_proxy_thread_pool_size is not None:
            task_patch["rpc_proxy_worker_thread_pool_size"] = rpc_proxy_thread_pool_size
        op = run_test_vanilla(
            with_breakpoint("echo $YT_JOB_PROXY_SOCKET_PATH >&2; BREAKPOINT; sleep {}".format(time_limit)),
            task_patch=task_patch
        )
        job_id = wait_breakpoint()[0]
        socket_file = op.read_stderr(job_id).decode("ascii").strip()
        release_breakpoint()
        return socket_file

    @authors("alex-shishkin")
    def test_disabled_rpc_proxy(self):
        with pytest.raises(YtError):
            socket_file = self.run_job_proxy(enable_rpc_proxy=False)
            client = self.create_client_from_uds(socket_file, config={"token": "tester_token"})
            client.list("/")

    @authors("alex-shishkin")
    def test_rpc_proxy_simple_query(self):
        socket_file = self.run_job_proxy(enable_rpc_proxy=True)
        client = self.create_client_from_uds(socket_file, config={"token": "tester_token"})
        root_listing = client.list("/")
        assert root_listing == ["sys", "tmp"]

    @authors("alex-shishkin")
    def test_failed_rpc_proxy_auth(self):
        with pytest.raises(YtError):
            socket_file = self.run_job_proxy(enable_rpc_proxy=True)
            client = self.create_client_from_uds(socket_file, config={"token": "wrong_token"})
            client.list("/")

    @authors("alex-shishkin")
    def test_incorrect_thread_count(self):
        with pytest.raises(YtError):
            self.run_job_proxy(enable_rpc_proxy=True, rpc_proxy_thread_pool_size=0)

class TestUnavailableJobProxy(JobProxyTestBase):
    DELTA_NODE_CONFIG = {
        "exec_node": {
            "job_controller": {
                "job_proxy_build_info_update_period": 300,
            },
            "slot_manager": {
                "testing": {
                    "skip_job_proxy_unavailable_alert": True,
                }
            }
        }
    }

    @authors("max42")
    def test_job_abort_on_unavailable_job_proxy(self):
        # JobProxyUnavailable alert is racy by its nature, so we still must ensure
        # that whenever job is scheduled to a node that does not have ytserver-job-proxy now,
        # job is simply aborted instead of accounting as failed. We do so
        # by forcefully skipping aforementioned alert and checking that operation successfully
        # terminates despite job proxy being unavailable for some period.

        job_count = 4
        op = run_test_vanilla("sleep 0.6", job_count=job_count, spec={"max_failed_job_count": 0}, track=False)

        wait(lambda: op.get_job_count("completed") > 0)

        with JobProxyHider(self):
            wait(lambda: op.get_job_count("aborted") > 2)

        op.track()

class TestJobProxyProfiling(YTEnvSetup):
    NUM_MASTERS = 1
    NUM_NODES = 1
    NUM_SCHEDULERS = 1
    ENABLE_RESOURCE_TRACKING = True

    @authors("prime")
    def test_sensors(self):
        op = run_test_vanilla("sleep 100", job_count=10, track=False)

        node = ls("//sys/cluster_nodes")[0]
        profiler = profiler_factory().at_job_proxy(node)

        thread_count = profiler.gauge("resource_tracker/thread_count")

        wait(lambda: thread_count.get() and thread_count.get() > 0)

        op.abort()

        wait(lambda: thread_count.get() is None)


from selenium import webdriver
import time
import math
from selenium.webdriver.chrome.options import Options

options = Options()
options.binary_location = "D:\\Program Files (x86)\\Google\\Chrome\\Application\\chrome.exe"

def calc(number):
    return str(math.log(abs(12 * math.sin(int(number)))))

try:
    link = "http://suninjuly.github.io/redirect_accept.html"
    browser = webdriver.Chrome(options=options)
    browser.get(link)
    time.sleep(4)

    button = browser.find_element_by_class_name("trollface")
    button.click()

    new_window = browser.window_handles[1]
    browser.switch_to.window(new_window)
    time.sleep(1)

    x_element = browser.find_element_by_id("input_value")
    x = x_element.text

    y = calc(x)

    input = browser.find_element_by_id("answer")
    input.send_keys(y)

    button1 = browser.find_element_by_class_name("btn")
    button1.click()

finally:
    time.sleep(5)
    browser.quit()


from dataclasses import dataclass, field
from typing import Optional
from nova.verbindung_tarifierbarkeits_request import VerbindungTarifierbarkeitsRequest

__NAMESPACE__ = "http://nova.voev.ch/services/v14/vertrieb"

@dataclass
class PruefenVerbindungTarifierbarkeits:
    class Meta:
        name = "pruefenVerbindungTarifierbarkeits"
        namespace = "http://nova.voev.ch/services/v14/vertrieb"

    verbindung_tarifierbarkeits_request: Optional[VerbindungTarifierbarkeitsRequest] = field(
        default=None,
        metadata={
            "name": "verbindungTarifierbarkeitsRequest",
            "type": "Element",
            "required": True,
        }
    )


import os, sys
import numpy as np
import paddle
# paddle.enable_static()
import paddle.fluid as fluid
from paddle import ParamAttr
import paddle.nn as nn
import paddle.nn.functional as F
from collections import OrderedDict
import torch

SEED = 666
INPUT_SIZE = 89
KERNEL_SIZE = 1
STRIDES = (1,1)
PADDING = 0

def get_para_bias_attr(l2_decay, k, name):
    import math
    regularizer = paddle.regularizer.L2Decay(l2_decay)
    stdv = 1.0 / math.sqrt(k * 1.0)
    initializer = nn.initializer.Uniform(-stdv, stdv)
    weight_attr = ParamAttr(
        regularizer=regularizer, initializer=initializer, name=name + "_w_attr")
    bias_attr = ParamAttr(
        regularizer=regularizer, initializer=initializer, name=name + "_b_attr")
    return [weight_attr, bias_attr]

def paddle_lstm():
    np.random.seed(SEED)
    x = np.random.rand(1, 80, 512).astype(np.float32)
    # np.save('org.npy', x)
    with fluid.dygraph.guard():
        lstm = nn.LSTM(512, 256, num_layers=2, direction='bidirectional')

        # sd = np.load('lstm.npy', allow_pickle=True).tolist()
        # lstm.set_state_dict(sd)

        state_dict = lstm.state_dict()
        sd = OrderedDict()
        for key, value in state_dict.items():
            v = value.numpy()
            print(key, value.shape, np.sum(v), np.mean(v), np.max(v), np.min(v))
            sd[key] = v

        np.save('lstm.npy', sd)

        inp = fluid.dygraph.to_variable(x)
        ret, _ = lstm(inp)
        print(len(ret))
    return ret.numpy()

def torch_lstm():
    np.random.seed(SEED)
    org = torch.Tensor(np.random.rand(1, 80, 512).astype(np.float32))
    lstm = torch.nn.LSTM(512, 256, num_layers=2, batch_first=True, bidirectional=True)
    sd = np.load('lstm.npy', allow_pickle=True)
    sd = sd.tolist()

    for key, value in lstm.state_dict().items():
        lstm.state_dict()[key].copy_(torch.Tensor(sd[key]))

    ret, _ = lstm(org)
    return ret.data.numpy()

def torch_lstm_m():
    class EncoderWithRNN(torch.nn.Module):
        def __init__(self, in_channels, hidden_size):
            super(EncoderWithRNN, self).__init__()
            self.out_channels = hidden_size * 2
            self.rnn1 = torch.nn.LSTM(in_channels, hidden_size, bidirectional=False, batch_first=True, num_layers=2)
            self.rnn2 = torch.nn.LSTM(in_channels, hidden_size, bidirectional=False, batch_first=True, num_layers=2)

        def forward(self, x):
            self.rnn1.flatten_parameters()
            self.rnn2.flatten_parameters()
            out1, h1 = self.rnn1(x)
            out2, h2 = self.rnn2(torch.flip(x, [1]))
            return torch.cat([out1, torch.flip(out2, [1])], 2)

    lstm = EncoderWithRNN(10, 3)
    for key, value in lstm.state_dict().items():
        print(key, value.shape)

if __name__ == '__main__':
    print('==========paddle=================')
    a = paddle_lstm()
    print(a.shape)
    print('a: ', np.sum(a), np.mean(a), np.max(a), np.min(a))
    print('===========pytorch================')
    b = torch_lstm()
    print(b.shape)
    print('b: ', np.sum(b), np.mean(b), np.max(b), np.min(b))
    # print('===========pytorch_m================')
    # torch_lstm_m()


# -*- coding: utf-8 -*-
"""
Created on Sat Oct  8 13:45:38 2022

@author: mzeigha
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import TwoSlopeNorm
import seaborn as sns
from textwrap import wrap
import matplotlib.colors as colors

production_no_tax = pd.read_parquet('Results/production_df_no_tax.parquet')
# production_SNP = pd.read_parquet('Results/production_df_SNP.parquet')

generators = pd.read_csv('Results/generators.csv',index_col=0 )
county_damage_per_generator = pd.read_csv('Results/county_damage_per_generator.csv', index_col=0).T
county_damage_per_generator = county_damage_per_generator[15:73]

population = pd.read_csv('Results/USA_Counties/white_population.csv')
population["non_white"] = (population['P1_001N']-population['P1_003N'])/population['P1_001N']

for i in range(len(population)):
    population.loc[i,'county'] = str(population.loc[i,'GEO_ID'])[-5:]
population.index = population['county']

egrid = pd.read_csv('C:/Users/mzeigha/Documents/GitHub/CAPOW_Emission/Generator/eGRID Plants.csv', dtype = str)
gen = pd.read_csv('C:/Users/mzeigha/Documents/GitHub/CAPOW_Emission/Generator/generator_final_DOE.csv', index_col=('name'))

for i in gen.index:
    a = egrid[egrid['ORISPL']==str(gen.loc[i,'DOE'])]

    try:
        gen.loc[i,'county2'] = a.iloc[0,16] + a.iloc[0,17]
    except Exception:
        pass

generators = pd.concat([generators,gen['county2']], axis = 1)

for i in generators.index:

    if generators.loc[i,'zone'] == 'PGE_valley':
        generators.loc[i,'gas'] = 4.66

    elif generators.loc[i,'zone']=='PGE_bay':
        generators.loc[i,'gas'] = 4.66

    else:
        generators.loc[i,'gas'] = 4.47

generators['no_tax_cost'] =generators['gas']*generators['seg2']+generators['var_om']

generators = generators.sort_values(by = 'no_tax_cost')

M = pd.DataFrame([])
for i in generators.index:
    for j in population.index:
        if generators.loc[i, 'county2'] == str(j):
            M.loc[i,j] = generators.loc[i,'NOXTax'] + generators.loc[i,'SO2Tax'] + generators.loc[i,'PMTax']
        else:
            M.loc[i,j] = 0

no_tax= production_no_tax.groupby(['Year'])[generators.index].sum()
# SNP= production_SNP.groupby(['Year'])[generators.index].sum()

portion = ((county_damage_per_generator.T).loc[generators.index]).T

no_tax_damage = {}
# SNP_damage = {}
for i in range(500):
    no_tax_damage[i] = (portion*no_tax.loc[i])
    # SNP_damage[i] = (county_damage_per_generator*SNP.loc[i]).T

no_tax_damage_avg = (portion*no_tax.sum()/500).T
# SNP_damage_avg = county_damage_per_generator*SNP.sum()/500

# M2 = pd.DataFrame([])
# for i in generators.index:
#     for j in population.index:
#         if generators.loc[i, 'county2'] == str(j):
#             M2.loc[i,j] = no_tax_damage_avg.loc[i,j]
#         # else:
#         #     M2.loc[i,j] = 0

#%%
amir = pd.DataFrame([])
for i in range (500):
    amir.loc[i, 'sum'] = no_tax_damage[i].sum().sum()

#%%

# M3 = pd.DataFrame([])
# for i in generators.index:
#     for j in population.index:
#         if generators.loc[i, 'county2'] == str(j):
#             M3.loc[i,j] = no_tax_damage[151].loc[i,j]
#         # else:
#         #     M3.loc[i,j] = 0

M2 = no_tax_damage_avg
M3 = no_tax_damage[151].T
M4 = M3-M2

xcsfont = {'fontname':'arial' , 'fontweight':'bold'}
csfont = {'fontname':'arial'}

fig = plt.figure( figsize = (15,10))

gs = fig.add_gridspec(4,1, width_ratios=[1], height_ratios = [1, 1, 1, 1] , wspace=0, hspace=0.3 )

############# Wind
cf_ax1 = fig.add_subplot(gs[0,0])
cf_ax2 = fig.add_subplot(gs[1,0])
cf_ax3 = fig.add_subplot(gs[2,0])
cf_ax4 = fig.add_subplot(gs[3,0])

cmap1 = 'Reds'

A = pd.concat([M.T , population['non_white']], axis = 1)
A = A.sort_values(by = 'non_white')

ylabel = [np.round(A.iloc[0,-1],decimals=2), np.round(A.iloc[10,-1],decimals=2), np.round(A.iloc[20,-1],decimals=2), np.round(A.iloc[30,-1],decimals=2), np.round(A.iloc[40,-1],decimals=2),np.round(A.iloc[50,-1],decimals=2)]
yticks = [0,10,20,30,40,50]

xlabel = [np.round(generators.iloc[0,-1],decimals=2), np.round(generators.iloc[50,-1],decimals=2), np.round(generators.iloc[100,-1],decimals=2), np.round(generators.iloc[150,-1],decimals=2), np.round(generators.iloc[200,-1],decimals=2)]
xticks = [0,50,100,150,200]

A.index = A['non_white']
A = A.drop(columns = ['non_white'])

norm=colors.LogNorm(vmin = 1 , vmax = 100 )
cbar = plt.cm.ScalarMappable(norm=norm, cmap=cmap1)

a = cf_ax1.matshow(A, norm=norm, cmap=cmap1)
cf_ax1.set_yticks( yticks, ylabel ,fontsize=14, **csfont)
cf_ax1.set_xticks( xticks, xlabel, fontsize=14, **csfont)
cf_ax1.xaxis.set_label_position('top')
cf_ax1.set_xlabel('a)',loc='left',labelpad = 6, fontsize=15, **xcsfont)
# cf_ax1.set_ylabel('SNP Tax',fontsize=16, **csfont)
cf_ax1.set_title('Power Plant (increasing Marginal Cost in $/MWh)',y = 1.2, fontsize=16, **csfont)
cf_ax1.axhline(32,color='black', linestyle='dotted')

cf_ax1.annotate('32 Majority white counties', (150,5),annotation_clip=False , fontsize=11,**csfont)
cf_ax1.annotate('26 Majority non-white counties', (140,37),annotation_clip=False , fontsize=11,**csfont)
cf_ax1.annotate('Non-White portion of County Population', (-32,210), rotation=90 , annotation_clip=False , fontsize=16,**csfont)

cax = plt.axes([0.72, 0.73, 0.02, 0.15])
cbar = plt.colorbar(a, cax= cax , orientation="vertical", ticks = [1,10,100] )#.set_ticklabels([3.13,8.58])
cbar.set_ticklabels([0, 10, 100])
cbar.ax.tick_params(labelsize=12)
cbar.set_label('\n'.join(wrap("Damage Rate ($/MWh)", 12)),labelpad=2,fontsize=14,**csfont)
# '\n'.join(wrap("Anomaly (C)", 12)),labelpad=15,fontsize=12,**csfont)

B = pd.concat([M2.T , population['non_white']], axis = 1)
B = B.sort_values(by = 'non_white')
B.index = B['non_white']
B = B.drop(columns = ['non_white'])
norm=colors.LogNorm(vmin=1, vmax=M3.max().max())
cbar = plt.cm.ScalarMappable(norm=norm, cmap=cmap1)

cf_ax2.matshow(B, norm = colors.LogNorm(), cmap=cmap1)
cf_ax2.set_yticks( yticks, ylabel ,fontsize=14, **csfont)
cf_ax2.set_xticks([])
# cf_ax2.set_ylabel('Average Year',fontsize=16, **csfont)
# cf_ax2.set_title('Net Cap',fontsize=16, **csfont)
cf_ax2.xaxis.set_label_position('top')
cf_ax2.set_xlabel('b)',loc='left',labelpad = 6, fontsize=15, **xcsfont)

C = pd.concat([M3.T , population['non_white']], axis = 1)
C = C.sort_values(by = 'non_white')
C.index = C['non_white']
C = C.drop(columns = ['non_white'])

cf_ax3.matshow(C, norm = colors.LogNorm(), cmap=cmap1)
cf_ax3.set_yticks( yticks, ylabel ,fontsize=14, **csfont)
cf_ax3.set_xticks([])
# cf_ax3.set_ylabel('Worst Year',fontsize=16, **csfont)
# cf_ax3.title('Net Cap')
cf_ax3.xaxis.set_label_position('top')
cf_ax3.set_xlabel('c)',loc='left',labelpad = 6, fontsize=15, **xcsfont)

cax2 = plt.axes([0.72, 0.33, 0.02, 0.34])
cbar = plt.colorbar(cbar, cax= cax2 , orientation="vertical")#, ticks = [1,100,200] )#.set_ticklabels([3.13,8.58])
# cbar.set_ticklabels([0, 10, 200])
cbar.ax.tick_params(labelsize=12)
cbar.set_label("Damage ($)",labelpad=2,fontsize=14,**csfont)

# norm= TwoSlopeNorm(vmin=M4.min().min(), vcenter=0, vmax=M4.max().max())

D = pd.concat([M4.T , population['non_white']], axis = 1)
D = D.sort_values(by = 'non_white')
D.index = D['non_white']
D = D.drop(columns = ['non_white'])

 # norm=colors.SymLogNorm
a = cf_ax4.matshow(D, norm = colors.LogNorm(), cmap=cmap1)
cf_ax4.set_yticks( yticks, ylabel ,fontsize=14, **csfont)
cf_ax4.set_xticks([])
# cf_ax4.set_ylabel('Disparity',fontsize=16, **csfont)
cf_ax4.xaxis.set_label_position('top')
cf_ax4.set_xlabel('d)',loc='left',labelpad = 6, fontsize=15, **xcsfont)

# ylabel = [-10000, -100, 0 ,100,10000,1000000,100000000]
# yticks = [-10000, -100, 0 ,100,10000,1000000,100000000]

cax3 = plt.axes([0.72, 0.127, 0.02, 0.15])
cbar = plt.colorbar(a, cax= cax3 , orientation="vertical")#, ticks = [ 0, 100] )#.set_ticklabels([3.13,8.58])
# cbar.set_ticklabels(['-10000','','','','', '0','','','','','' ,'100000000'])
cbar.ax.tick_params(labelsize=13)
cbar.set_label("Damage ($)",labelpad=13,fontsize=14,**csfont)

# plt.savefig('Plots/REVIEW_MP2_1.png' , bbox_inches='tight',dpi=250)

# #%%
# N2 = SNP_damage_avg.T
# N3 = SNP_damage[151]
# N4 = N3-N2

# csfont = {'fontname':'arial'}

# fig = plt.figure( figsize = (15,10))

# gs = fig.add_gridspec(4,1, width_ratios=[1], height_ratios = [1, 1, 1, 1] , wspace=0, hspace=0.1 )

# ############# Wind
# cf_ax1 = fig.add_subplot(gs[0,0])
# cf_ax2 = fig.add_subplot(gs[1,0])
# cf_ax3 = fig.add_subplot(gs[2,0])
# cf_ax4 = fig.add_subplot(gs[3,0])

# A = pd.concat([M.T , population['non_white']], axis = 1)
# A = A.sort_values(by = 'non_white')

# ylabel = [np.round(A.iloc[0,-1],decimals=2), np.round(A.iloc[10,-1],decimals=2), np.round(A.iloc[20,-1],decimals=2), np.round(A.iloc[30,-1],decimals=2), np.round(A.iloc[40,-1],decimals=2),np.round(A.iloc[50,-1],decimals=2)]
# yticks = [0,10,20,30,40,50]

# xlabel = [np.round(generators.iloc[0,2],decimals=2), np.round(generators.iloc[50,2],decimals=2), np.round(generators.iloc[100,2],decimals=2), np.round(generators.iloc[150,2],decimals=2), np.round(generators.iloc[200,2],decimals=2)]
# xticks = [0,50,100,150,200]

# A.index = A['non_white']
# A = A.drop(columns = ['non_white'])

# norm=colors.LogNorm(vmin = 1 , vmax = A.max().max() )
# cbar = plt.cm.ScalarMappable(norm=norm, cmap='RdBu_r')

# cf_ax1.matshow(A, norm=colors.LogNorm(), cmap='RdBu_r')
# cf_ax1.set_yticks( yticks, ylabel ,fontsize=14, **csfont)
# cf_ax1.set_xticks( xticks, xlabel ,fontsize=14, **csfont)
# cf_ax1.set_ylabel('SNP Tax',fontsize=16, **csfont)
# cf_ax1.set_title('Net Cap',fontsize=16, **csfont)
# cf_ax1.axhline(32,color='black', linestyle='dotted')

# cf_ax1.annotate('32 Majority white counties', (160,5),annotation_clip=False , fontsize=11,**csfont)
# cf_ax1.annotate('26 Majority non-white counties', (150,37),annotation_clip=False , fontsize=11,**csfont)

# cax = plt.axes([0.75, 0.705, 0.02, 0.17])
# cbar = plt.colorbar(cbar, cax= cax , orientation="vertical", ticks = [1,10,200] )#.set_ticklabels([3.13,8.58])
# cbar.set_ticklabels([0, 10, 200])
# cbar.ax.tick_params(labelsize=15)

# B = pd.concat([N2.T , population['non_white']], axis = 1)
# B = B.sort_values(by = 'non_white')
# B.index = B['non_white']
# B = B.drop(columns = ['non_white'])
# norm=colors.LogNorm(vmin=1, vmax=M3.max().max())
# cbar = plt.cm.ScalarMappable(norm=norm, cmap='RdBu_r')

# cf_ax2.matshow(B, norm = colors.LogNorm(), cmap='RdBu_r')
# cf_ax2.set_yticks( yticks, ylabel ,fontsize=14, **csfont)
# cf_ax2.set_xticks([])
# cf_ax2.set_ylabel('Average Year',fontsize=16, **csfont)
# # cf_ax2.title('Net Cap')

# C = pd.concat([N3.T , population['non_white']], axis = 1)
# C = C.sort_values(by = 'non_white')
# C.index = C['non_white']
# C = C.drop(columns = ['non_white'])

# cf_ax3.matshow(C, norm = colors.LogNorm(), cmap='RdBu_r')
# cf_ax3.set_yticks( yticks, ylabel ,fontsize=14, **csfont)
# cf_ax3.set_xticks([])
# cf_ax3.set_ylabel('Worst Year',fontsize=16, **csfont)
# # cf_ax3.title('Net Cap')

# cax2 = plt.axes([0.75, 0.33, 0.02, 0.35])
# cbar = plt.colorbar(cbar, cax= cax2 , orientation="vertical")#, ticks = [1,100,200] )#.set_ticklabels([3.13,8.58])
# # cbar.set_ticklabels([0, 10, 200])
# cbar.ax.tick_params(labelsize=15)

# norm= TwoSlopeNorm(vmin=N4.min().min(), vcenter=0, vmax=N4.max().max())
# cbar = plt.cm.ScalarMappable(norm=norm, cmap='RdBu_r')
# D = pd.concat([N4.T , population['non_white']], axis = 1)
# D = D.sort_values(by = 'non_white')
# D.index = D['non_white']
# D = D.drop(columns = ['non_white'])

# cf_ax4.matshow(D, norm = colors.LogNorm(), cmap='RdBu_r')
# cf_ax4.set_yticks( yticks, ylabel ,fontsize=14, **csfont)
# cf_ax4.set_xticks([])
# cf_ax4.set_ylabel('Worst Year',fontsize=16, **csfont)

# cax3 = plt.axes([0.75, 0.125, 0.02, 0.17])
# cbar = plt.colorbar(cbar, cax= cax3 , orientation="vertical")#, ticks = [1,10,200] )#.set_ticklabels([3.13,8.58])
# # cbar.set_ticklabels([0, 10, 200])
# cbar.ax.tick_params(labelsize=15)

# plt.savefig('Plots/REVIEW2_SNP.png' , bbox_inches='tight',dpi=250)


import logging
from flearn.client import Client
import math
import numpy as np
from flearn.fedave import Server as ServerBase

class Server(ServerBase):
    def __init__(self, data, model, parsed):
        logging.info('Using Federated avg to Train')
        self.data = data
        for key, val in parsed.items():
            logging.info("%s: %s" % (key, val))
            setattr(self, key, val)

        # initialize client
        self.clients = self.setup_clients(data, model)
        self.latest_model = None

    def update_local_model(self, client, learning_rates, betas):
        client.nesterov_grad_descent(self.num_epochs, self.batch_size, learning_rates, betas)

    def get_betas(self, rounds, num_epochs):
        # betas = [(rounds + i) / (rounds + i + 3) for i in range(num_epochs)]
        betas = [0.1] * num_epochs
        return betas

    def train(self, results_filename):
        # learning_rate = self.learning_rate
        if self.is_decay:
            learning_rates = self.get_epoch_learning_rates(0, self.learning_rate, self.lrconst, self.num_epochs)
        else:
            learning_rates = [self.learning_rate * self.lrconst] * self.num_epochs

        betas = self.get_betas(0, self.num_epochs)
        train_loss, test_loss = 0, 0
        final_iteration = 0
        is_break = False
        train_losss, test_losss, train_accs, test_accs = [], [], [], []
        num_selected_clients = max(int(self.partial_participation * len(self.clients)), 1)
        for i in range(self.num_rounds):
            final_iteration = i * self.num_epochs
            if i % self.eval_every == 0:
                train_loss, test_loss, train_acc, test_acc = self.evaluation()
                train_losss.append(train_loss); test_losss.append(test_loss)
                train_accs.append(train_acc); test_accs.append(test_acc)
                loginfostr = "Round {} Train_loss {}, Test_loss {}, " \
                             "Train_acc {}, Test_acc {}, lr {} beta {}".format(i,
                                                                               train_loss,
                                                                               test_loss,
                                                                               train_acc,
                                                                               test_acc,
                                                                               learning_rates[-1],
                                                                               betas[-1])
                logging.info(loginfostr)
                if train_loss < self.epsilon and self.model == 'linearregression':
                    logging.info("Final iterations {}".format(final_iteration))
                    is_break = True
                    break
                if train_acc > self.epsilon and self.model == 'logisticregression':
                    logging.info("Final iterations {}".format(final_iteration))
                    is_break = True
                    break
                if train_loss < self.epsilon and self.model == 'binarylogisticregression':
                    logging.info("Final iterations {}".format(final_iteration))
                    is_break = True
                    break
                if math.isnan(train_loss) or math.isinf(train_loss):
                    logging.info("Final iterations {}".format(1e8))
                    is_break = True
                    break

            if final_iteration > self.num_iterations:
                break

            active_clients = self.selected_clients(i,
                                                   num_selected_clients) if self.partial_participation != 1 else self.clients

            csolns = []  # buffer for receiving client solutions
            for client in active_clients:
                # client.set_params(self.latest_model)
                self.update_local_model(client, learning_rates, betas)
                csolns.append(client.get_params())
            self.latest_model = self.aggregate(csolns)
            self.sync_clients(self.latest_model)

            betas = self.get_betas(i, self.num_epochs)
            # warmup = 100
            if self.is_decay:
                learning_rates = self.get_epoch_learning_rates(i, self.learning_rate, self.lrconst, self.num_epochs)
                # learning_rate = 1/(i * self.learning_rate + self.lrconst)  # see ICLR'20 Xiang Li, C.5

            if (i * self.num_epochs) % 100 == 0:
                results = [train_losss, test_losss, train_accs, test_accs, final_iteration]
                np.save(results_filename, results)

        if is_break is False:
            logging.info("Final iterations {}".format(final_iteration))

        results = [train_losss, test_losss, train_accs, test_accs, final_iteration]
        np.save(results_filename, results)

        return results


from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QHBoxLayout, QLabel, QPushButton, QFrame, QMessageBox

from compoment.atx2agent.core.common.logs.log_uru import Logger
from compoment.atx2agent.core.task.task_manage import TaskType
from ui.test.execTaskThread import TaskManagerThread
from ui.test.model.model import Task, TaskStatus

logger = Logger().get_logger
class TestTaskRunModule(QFrame):
    def __init__(self, case_manager_controller=None,device_manager_controller=None,task_type_controller=None,task_manager_table_controller=None,parent=None):
        super(TestTaskRunModule, self).__init__(parent)
        self.initUI()
        self.case_manager_controller = case_manager_controller
        self.device_manager_controller =device_manager_controller
        self.task_type_controller = task_type_controller
        self.task_manager_table_controller = task_manager_table_controller

    def initUI(self):
        ################ 执行任务操作 ###############
        self.testTaskRunLayout = QHBoxLayout()
        self.testTaskRunLayout.setAlignment(Qt.AlignBottom)
        self.setLayout(self.testTaskRunLayout)

        self.run_case_label = QLabel("操作")
        self.testTaskRunLayout.addWidget(self.run_case_label)

        self.task_button = QPushButton('添加任务')
        self.testTaskRunLayout.addWidget(self.task_button)
        self.task_button.clicked.connect(self.on_add_task)

        self.run_task_button = QPushButton('执行任务')
        self.testTaskRunLayout.addWidget(self.run_task_button)
        self.run_task_button.clicked.connect(self.on_run_task)

        self.cancel_task_button = QPushButton('取消任务')
        self.testTaskRunLayout.addWidget(self.cancel_task_button)
        self.cancel_task_button.clicked.connect(self.on_cancel_task)

    def on_add_task(self):
        cases=[]
        selected_cases = self.case_manager_controller.get_selected_cases()
        selected_devices = self.device_manager_controller.get_selected_devices()
        selected_task_type =self.task_type_controller.get_task_type()
        if not selected_cases:
            QMessageBox.warning(self, '警告', '请选择测试用例')
            return
        if not selected_devices:
            QMessageBox.warning(self, '警告', '请选择一个设备')
            return
        if selected_task_type==TaskType.NONE:
            QMessageBox.warning(self, '警告', '请选择一种任务类型')
            return
        if self.task_manager_table_controller.is_device_occupied(selected_devices):
            QMessageBox.warning(self, '警告', f'设备[{selected_devices[0].serialsName}]已经被占用了')
            return
        # todo 根据任务类型的不同，需要选择对应的用例, 如果不同则弹出提示框
        try:
            if selected_task_type == TaskType.KW_UI_TEST:
                cases = self.case_manager_controller.get_test_case_list_view()
            elif selected_task_type == TaskType.PO_UI_TEST:
                cases = self.case_manager_controller.get_po_py_case_list()
                cases = sum(cases, [])
            if cases:
                # 根据设备的数量，创建多个任务
                for device in selected_devices:
                    task = Task(task_id=len(self.task_manager_table_controller.task_list), task_name='task1', test_cases=cases, device=device,
                                task_type=selected_task_type, status=TaskStatus.PENDING, result='',task_timeout=300)
                    self.task_manager_table_controller.add_task(task)
            else:
                QMessageBox.warning(self, '警告', '没有可执行的用例')
                return
        except Exception as e:
            logger.error("添加任务失败{}".format(e))
            QMessageBox.warning(self, '警告', '添加任务失败')

    def on_run_task(self):
        tasks =self.task_manager_table_controller.get_task_list()
        if tasks is None:
            QMessageBox.warning(self, '警告', '没有可以执行的任务')
            return
        if tasks:
            for task in tasks:   # 根据任务数量，创建多个线程
                # TODO 这里要判断一下有没有选中的任务
                if not self.task_manager_table_controller.no_selected_task_messagebox():
                    return
                # 如果任务已经在执行了，或者任务取消中，就不再执行
                # TODO 这里要改成通知 另一个专门做自动化用例执行的线程，去执行任务，而不是在这里直接执行，取消任务也是一样
                if task.status == TaskStatus.RUNNING or task.status == TaskStatus.CANCELLING:
                    QMessageBox.warning(self, '警告', '任务已经在执行了')
                    continue
                # 设置任务状态为执行中
                task.status = TaskStatus.RUNNING
                task_thread = TaskManagerThread(task)
                self.task_manager_table_controller.set_thread(task_thread)
                print("task_thread_id------------>:", id(task_thread))
                # task_thread.task_cancelled_signal.connect(self.task_manager.cancel_task)
                task_thread.update_status_signal.connect(self.task_manager_table_controller.update_task_status)
                task_thread.update_result_signal.connect(self.task_manager_table_controller.update_task_result)
                task_thread.start()

    def on_cancel_task(self):
        """
        取消任务
        @return:
        """
        # 获取选中的任务
        tasks = self.task_manager_table_controller.get_selected_task()
        print('取消任务: {}'.format(tasks))
        for task in tasks:
            print("任务类型：", type(task))
            self.task_manager_table_controller.cancel_task(task)


#!/usr/bin/env python
# coding: utf-8

# # model 5

# In[6]:

import polars as pl
import numpy as np
import pandas as pd

import tensorflow as tf

import matplotlib.pyplot as plt
import seaborn as sns

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OneHotEncoder

# import from the auxFunctions.py file
from auxFunctions import calculate_mae_and_mrrmse, mean_rowwise_rmse_loss, custom_mean_rowwise_rmse, create_model_checkpoint, plot_training_history

# ## pre-process data
# one-hot encode cell type
# <br>map sm_name to affinities_MPNN_CNN_BindingDB.csv

# In[7]:

de_train = pl.scan_parquet('de_train.parquet')
de_train_df = de_train.collect().to_pandas()

# test provided by kaggle --> upload predictions to kaggle to get the score
id_map = pd.read_csv('id_map.csv')

affinities = pd.read_csv('affinities_Transformer_CNN_BindingDB.csv', index_col=0)

# train dataset provided by kaggle
# - will be split into train/test/validation for internal testing before model is trained on the entire train and used to predict on the test in id_map

# In[8]:

def extractAffinities(sm_names, affinities):
    """
    Function to extract affinities from the affinities dataframe

    Parameters:
    - sm_names: List/Array of sm_names
    - affinities: Stored affinities predicted using DeepPurpose

    Returns:
    - Affinities as a numpy array
    """
    encoded_affinities = []
    for name in sm_names:
        filtered = affinities[affinities['sm_name'] == name]
        sm_affinities = filtered.iloc[:, 2:].values[0]
        encoded_affinities.append(sm_affinities)

    np_encoded_affinities = np.array(encoded_affinities)

    return np_encoded_affinities

# In[9]:

# one-hot encode cell_type
cell_type = de_train_df['cell_type'].to_numpy().reshape(-1, 1)
encoder = OneHotEncoder()
encoder.fit(cell_type)
# of type scipy.sparse._csr.csr_matrix
encoded_cell_type = encoder.transform(cell_type)

# map sm_name to affinities
sm_name = de_train_df['sm_name']

# has shape (614, 12766), of type numpy.ndarray
np_encoded_affinities = extractAffinities(sm_name, affinities)

# concatenate encoded_cell_type and np_encoded_affinities
# final shape (614, 12772)
encoded_features = np.hstack((encoded_cell_type.toarray(), np_encoded_affinities))

# wanted output
genes_lfc = de_train_df.drop(columns=['cell_type', 'sm_name', 'sm_lincs_id', 'SMILES', 'control'])

# In[10]:

# repeat for kaggle test set
kaggle_cell_type = id_map['cell_type'].to_numpy().reshape(-1, 1)
encoded_kaggle_cell_type = encoder.transform(kaggle_cell_type)

kaggle_sm_name = id_map['sm_name']
encoded_kaggle_affinities = extractAffinities(kaggle_sm_name, affinities)

# final shape (255, 12772)
encoded_kaggle_features = np.hstack((encoded_kaggle_cell_type.toarray(), encoded_kaggle_affinities))

# In[11]:

# Split the data into 70% training, 15% validation, and 15% testing
X_train, X_temp, y_train, y_temp = train_test_split(encoded_features, genes_lfc.values, test_size=0.3, shuffle=False)
X_val, X_test, y_val, y_test = train_test_split(X_temp, y_temp, test_size=0.5, shuffle=False)

# used for final training before predicting on kaggle_test
full_features = encoded_features
full_labels = genes_lfc.values

# ## model training

# In[12]:

from tensorflow.keras.layers import Dense, Dropout, BatchNormalization, Activation
from tensorflow.keras.models import Sequential

# In[ ]:

tf.random.set_seed(42)

model_2 = Sequential([
    Dense(3400, activation="tanh"),
    Dense(3000, activation="tanh"),
    Dense(2100, activation="tanh"),
    Dense(200, activation="tanh"),
    Dense(200, activation="tanh"),
    Dropout(0.2),
    Dense(18211, activation="linear")
])

model_2.compile(loss="mae",
                optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),
                metrics=[custom_mean_rowwise_rmse])

history_2 = model_2.fit(X_train, y_train,
                       epochs=10,
                       validation_data=(X_val,y_val),
                       batch_size=32,
                       callbacks=[create_model_checkpoint("model_2", monitor="val_custom_mean_rowwise_rmse")])

calculate_mae_and_mrrmse(model=model_2, data=X_test, y_true=y_test)

# In[21]:

plot_training_history(history_2, metrics=["custom_mean_rowwise_rmse"])

kaggle_prediction = model_2.predict(encoded_kaggle_features, batch_size=1)

sample_submission = pd.read_csv('./sample_submission.csv')
sample_submission.iloc[:,1:] = kaggle_predictions

sample_submission.to_csv('submission_model_3.csv', index=False)