import argparse
import os
import readline
import sys
from optparse import make_option
from time import sleep, time

import etcd3
import grpc
import requests
from cmd2 import Cmd, options
from consul import Consul
from google.protobuf.empty_pb2 import Empty
from simplejson import dumps

from cli.device import DeviceCli
from cli.xpon import XponCli
from cli.omci import OmciCli
from cli.alarm_filters import AlarmFiltersCli
from cli.logical_device import LogicalDeviceCli
from cli.table import print_pb_list_as_table
from voltha.protos import third_party
from voltha.protos import voltha_pb2, voltha_pb2_grpc, health_pb2_grpc
from voltha.protos.openflow_13_pb2 import FlowTableUpdate, FlowGroupTableUpdate

_ = third_party
from cli.utils import pb2dict

defs = dict(
    # config=os.environ.get('CONFIG', './cli.yml'),
    etcd=os.environ.get('ETCD', 'etcd-cluster.default.svc.cluster.local:2379'),
    consul=os.environ.get('CONSUL', 'localhost:8500'),
    voltha_grpc_endpoint=os.environ.get('VOLTHA_GRPC_ENDPOINT',
                                        'localhost:50055'),
    voltha_sim_rest_endpoint=os.environ.get('VOLTHA_SIM_REST_ENDPOINT',
                                            'localhost:18880'),
    global_request=os.environ.get('GLOBAL_REQUEST', False)
)

banner = """\
         _ _   _            ___ _    ___
__ _____| | |_| |_  __ _   / __| |  |_ _|
\ V / _ \ |  _| ' \/ _` | | (__| |__ | |
 \_/\___/_|\__|_||_\__,_|  \___|____|___|
(to exit type quit or exit or hit Ctrl-D)
"""

class VolthaCli(Cmd):
    prompt = 'voltha'
    history_file_name = '.voltha_cli_history'

    # Settable CLI parameters
    voltha_grpc = 'localhost:50055'
    voltha_sim_rest = 'localhost:18880'
    global_request = False
    max_history_lines = 500
    default_device_id = None
    default_logical_device_id = None

    Cmd.settable.update(dict(
        voltha_grpc='Voltha GRPC endpoint in form of <host>:<port>',
        voltha_sim_rest='Voltha simulation back door for testing in form '
                        'of <host>:<port>',
        max_history_lines='Maximum number of history lines stored across '
                          'sessions',
        default_device_id='Device id used when no device id is specified',
        default_logical_device_id='Logical device id used when no device id '
                                  'is specified',
    ))

    # cleanup of superfluous commands from cmd2
    del Cmd.do_cmdenvironment
    del Cmd.do_load
    del Cmd.do__relative_load

    def __init__(self, voltha_grpc, voltha_sim_rest, etcd, global_request=False):
        VolthaCli.voltha_grpc = voltha_grpc
        VolthaCli.voltha_sim_rest = voltha_sim_rest
        VolthaCli.global_request = global_request
        VolthaCli.etcd = etcd
        Cmd.__init__(self)
        self.prompt = '(' + self.colorize(
            self.colorize(self.prompt, 'blue'), 'bold') + ') '
        self.channel = None
        self.stub = None
        self.device_ids_cache = None
        self.device_ids_cache_ts = time()
        self.logical_device_ids_cache = None
        self.logical_device_ids_cache_ts = time()

    # we override cmd2's method to avoid its optparse conflicting with our
    # command line parsing
    def cmdloop(self):
        self._cmdloop()

    def load_history(self):
        """Load saved command history from local history file"""
        try:
            with file(self.history_file_name, 'r') as f:
                for line in f.readlines():
                    stripped_line = line.strip()
                    self.history.append(stripped_line)
                    readline.add_history(stripped_line)
        except IOError:
            pass  # ignore if file cannot be read

    def save_history(self):
        try:
            with open(self.history_file_name, 'w') as f:
                f.write('\n'.join(self.history[-self.max_history_lines:]))
        except IOError as e:
            self.perror('Could not save history in {}: {}'.format(
                self.history_file_name, e))
        else:
            self.poutput('History saved as {}'.format(
                self.history_file_name))

    def perror(self, errmsg, statement=None):
        # Touch it up to make sure error is prefixed and colored
        Cmd.perror(self, self.colorize('***ERROR: ', 'red') + errmsg,
                   statement)

    def get_channel(self):
        if self.channel is None:
            self.channel = grpc.insecure_channel(self.voltha_grpc)
        return self.channel

    def get_stub(self):
        if self.stub is None:
            self.stub = \
                voltha_pb2_grpc.VolthaGlobalServiceStub(self.get_channel()) \
                    if self.global_request else \
                        voltha_pb2_grpc.VolthaLocalServiceStub(self.get_channel())
        return self.stub

    # ~~~~~~~~~~~~~~~~~ ACTUAL COMMAND IMPLEMENTATIONS ~~~~~~~~~~~~~~~~~~~~~~~~

    def do_reset_history(self, line):
        """Reset CLI history"""
        while self.history:
            self.history.pop()

    def do_exit(self,line):
        """exit from CLI"""
        quit()

    def do_launch(self, line):
        """If Voltha is not running yet, launch it"""
        raise NotImplementedError('not implemented yet')

    def do_restart(self, line):
        """Launch Voltha, but if it is already running, terminate it first"""
        pass

    def do_adapters(self, line):
        """List loaded adapter"""
        stub = self.get_stub()
        res = stub.ListAdapters(Empty())
        omit_fields = {'config.log_level', 'logical_device_ids'}
        print_pb_list_as_table('Adapters:', res.items, omit_fields, self.poutput)

    def get_devices(self):
        stub = self.get_stub()
        res = stub.ListDevices(Empty())
        return res.items

    def get_logical_devices(self):
        stub = self.get_stub()
        res = stub.ListLogicalDevices(Empty())
        return res.items

    def do_devices(self, line):
        """List devices registered in Voltha"""
        devices = self.get_devices()
        omit_fields = {
            'adapter',
            'vendor',
            'model',
            'hardware_version',
            'images',
            'firmware_version',
            'vendor_id'
        }
        print_pb_list_as_table('Devices:', devices, omit_fields, self.poutput)

    def do_logical_devices(self, line):
        """List logical devices in Voltha"""
        stub = self.get_stub()
        res = stub.ListLogicalDevices(Empty())
        omit_fields = {
            'desc.mfr_desc',
            'desc.hw_desc',
            'desc.sw_desc',
            'desc.dp_desc',
            'desc.serial_number',
            'switch_features.capabilities'
        }
        presfns = {
            'datapath_id': lambda x: "{0:0{1}x}".format(int(x), 16)
        }
        print_pb_list_as_table('Logical devices:', res.items, omit_fields,
                               self.poutput, presfns=presfns)

    def do_device(self, line):
        """Enter device level command mode"""
        device_id = line.strip() or self.default_device_id
        if not device_id:
            raise Exception('<device-id> parameter needed')
        if device_id not in self.device_ids():
            self.poutput( self.colorize('Error: ', 'red') +
                            'There is no such device')
            raise Exception('<device-id> is not a valid one')
        sub = DeviceCli(device_id, self.get_stub, self.etcd)
        sub.cmdloop()

    def do_logical_device(self, line):
        """Enter logical device level command mode"""
        logical_device_id = line.strip() or self.default_logical_device_id
        if not logical_device_id:
            raise Exception('<logical-device-id> parameter needed')
        if logical_device_id not in self.logical_device_ids():
            self.poutput( self.colorize('Error: ', 'red') +
                            'There is no such device')
            raise Exception('<logical-device-id> is not a valid one')
        sub = LogicalDeviceCli(logical_device_id, self.get_stub)
        sub.cmdloop()

    def device_ids(self, force_refresh=False):
        if force_refresh or self.device_ids is None or \
                        (time() - self.device_ids_cache_ts) > 1:
            self.device_ids_cache = [d.id for d in self.get_devices()]
            self.device_ids_cache_ts = time()
        return self.device_ids_cache

    def logical_device_ids(self, force_refresh=False):
        if force_refresh or self.logical_device_ids is None or \
                        (time() - self.logical_device_ids_cache_ts) > 1:
            self.logical_device_ids_cache = [d.id for d
                                             in self.get_logical_devices()]
            self.logical_device_ids_cache_ts = time()
        return self.logical_device_ids_cache

    def complete_device(self, text, line, begidx, endidx):
        if not text:
            completions = self.device_ids()[:]
        else:
            completions = [d for d in self.device_ids() if d.startswith(text)]
        return completions

    def complete_logical_device(self, text, line, begidx, endidx):
        if not text:
            completions = self.logical_device_ids()[:]
        else:
            completions = [d for d in self.logical_device_ids()
                           if d.startswith(text)]
        return completions

    def do_xpon(self, line):
        """xpon <optional> [device_ID] - Enter xpon level command mode"""
        device_id = line.strip()
        if device_id:
            stub = self.get_stub()
            try:
                res = stub.GetDevice(voltha_pb2.ID(id=device_id))
            except Exception:
                self.poutput(
                    self.colorize('Error: ', 'red') + 'No device id ' +
                    self.colorize(device_id, 'blue') + ' is found')
                return
        sub = XponCli(self.get_channel, device_id)
        sub.cmdloop()

    def do_omci(self, line):
        """omci <device_ID> - Enter OMCI level command mode"""

        device_id = line.strip() or self.default_device_id
        if not device_id:
            raise Exception('<device-id> parameter needed')
        sub = OmciCli(device_id, self.get_stub)
        sub.cmdloop()

    def do_pdb(self, line):
        """Launch PDB debug prompt in CLI (for CLI development)"""
        from pdb import set_trace
        set_trace()

    def do_version(self, line):
        """Show the VOLTHA core version"""
        stub = self.get_stub()
        voltha = stub.GetVoltha(Empty())
        self.poutput('{}'.format(voltha.version))

    def do_health(self, line):
        """Show connectivity status to Voltha status"""
        stub = health_pb2_grpc.HealthServiceStub(self.get_channel())
        res = stub.GetHealthStatus(Empty())
        self.poutput(dumps(pb2dict(res), indent=4))

    @options([
        make_option('-t', '--device-type', action="store", dest='device_type',
                    help="Device type", default='simulated_olt'),
        make_option('-m', '--mac-address', action='store', dest='mac_address',
                    default='00:0c:e2:31:40:00'),
        make_option('-i', '--ip-address', action='store', dest='ip_address'),
        make_option('-H', '--host_and_port', action='store',
                    dest='host_and_port'),
    ])
    def do_preprovision_olt(self, line, opts):
        """Preprovision a new OLT with given device type"""
        stub = self.get_stub()
        kw = dict(type=opts.device_type)
        if opts.host_and_port:
            kw['host_and_port'] = opts.host_and_port
        elif opts.ip_address:
            kw['ipv4_address'] = opts.ip_address
        elif opts.mac_address:
            kw['mac_address'] = opts.mac_address.lower()
        else:
            raise Exception('Either IP address or Mac Address is needed')
        # Pass any extra arguments past '--' to the device as custom arguments
        kw['extra_args'] = line

        device = voltha_pb2.Device(**kw)
        device = stub.CreateDevice(device)
        self.poutput('success (device id = {})'.format(device.id))
        self.default_device_id = device.id

    def do_enable(self, line):
        """
        Enable a device. If the <id> is not provided, it will be on the last
        pre-provisioned device.
        """
        device_id = line or self.default_device_id
        if device_id not in self.device_ids():
            self.poutput('Error: There is no such preprovisioned device')
            return

        try:
            stub = self.get_stub()
            device = stub.GetDevice(voltha_pb2.ID(id=device_id))
            if device.admin_state == voltha_pb2.AdminState.ENABLED:
                if device.oper_status != voltha_pb2.OperStatus.ACTIVATING:
                    self.poutput('Error: Device is already enabled')
                    return
            else:
                stub.EnableDevice(voltha_pb2.ID(id=device_id))
                self.poutput('enabling {}'.format(device_id))

            while True:
                device = stub.GetDevice(voltha_pb2.ID(id=device_id))
                # If this is an OLT then acquire logical device id
                if device.oper_status == voltha_pb2.OperStatus.ACTIVE:
                    if device.type.endswith('_olt'):
                        assert device.parent_id
                        self.default_logical_device_id = device.parent_id
                        self.poutput('success (logical device id = {})'.format(
                            self.default_logical_device_id))
                    else:
                        self.poutput('success (device id = {})'.format(device.id))
                    break
                self.poutput('waiting for device to be enabled...')
                sleep(.5)
        except Exception as e:
            self.poutput('Error enabling {}.  Error:{}'.format(device_id, e))

    complete_activate_olt = complete_device

    def do_reboot(self, line):
        """
        Rebooting a device. ID of the device needs to be provided
        """
        device_id = line or self.default_device_id
        self.poutput('rebooting {}'.format(device_id))
        try:
            stub = self.get_stub()
            stub.RebootDevice(voltha_pb2.ID(id=device_id))
            self.poutput('rebooted {}'.format(device_id))
        except Exception as e:
            self.poutput('Error rebooting {}.  Error:{}'.format(device_id, e))

    def do_self_test(self, line):
        """
        Self Test a device. ID of the device needs to be provided
        """
        device_id = line or self.default_device_id
        self.poutput('Self Testing {}'.format(device_id))
        try:
            stub = self.get_stub()
            res = stub.SelfTest(voltha_pb2.ID(id=device_id))
            self.poutput('Self Tested {}'.format(device_id))
            self.poutput(dumps(pb2dict(res), indent=4))
        except Exception as e:
            self.poutput('Error in self test {}.  Error:{}'.format(device_id, e))

    def do_delete(self, line):
        """
        Deleting a device. ID of the device needs to be provided
        """
        device_id = line or self.default_device_id
        self.poutput('deleting {}'.format(device_id))
        try:
            stub = self.get_stub()
            stub.DeleteDevice(voltha_pb2.ID(id=device_id))
            self.poutput('deleted {}'.format(device_id))
        except Exception as e:
            self.poutput('Error deleting {}.  Error:{}'.format(device_id, e))

    def do_disable(self, line):
        """
        Disable a device. ID of the device needs to be provided
        """
        device_id = line
        if device_id not in self.device_ids():
            self.poutput('Error: There is no such device')
            return
        try:
            stub = self.get_stub()
            device = stub.GetDevice(voltha_pb2.ID(id=device_id))
            if device.admin_state == voltha_pb2.AdminState.DISABLED:
                self.poutput('Error: Device is already disabled')
                return
            stub.DisableDevice(voltha_pb2.ID(id=device_id))
            self.poutput('disabling {}'.format(device_id))

            # Do device query and verify that the device admin status is
            # DISABLED and Operational Status is unknown
            device = stub.GetDevice(voltha_pb2.ID(id=device_id))
            if device.admin_state == voltha_pb2.AdminState.DISABLED:
                self.poutput('disabled successfully {}'.format(device_id))
            else:
                self.poutput('disabling failed {}.  Admin State:{} '
                             'Operation State: {}'.format(device_id,
                                                          device.admin_state,
                                                          device.oper_status))
        except Exception as e:
            self.poutput('Error disabling {}.  Error:{}'.format(device_id, e))

    def do_test(self, line):
        """Enter test mode, which makes a bunch on new commands available"""
        sub = TestCli(self.history, self.voltha_grpc,
                      self.get_stub, self.voltha_sim_rest)
        sub.cmdloop()

    def do_alarm_filters(self, line):
        sub = AlarmFiltersCli(self.get_stub)
        sub.cmdloop()

class TestCli(VolthaCli):
    def __init__(self, history, voltha_grpc, get_stub, voltha_sim_rest):
        VolthaCli.__init__(self, voltha_grpc, voltha_sim_rest)
        self.history = history
        self.get_stub = get_stub
        self.prompt = '(' + self.colorize(self.colorize('test', 'cyan'),
                                          'bold') + ') '

    def get_device(self, device_id, depth=0):
        stub = self.get_stub()
        res = stub.GetDevice(voltha_pb2.ID(id=device_id),
                             metadata=(('get-depth', str(depth)),))
        return res

    def do_arrive_onus(self, line):
        """
        Simulate the arrival of ONUs (available only on simulated_olt)
        """
        device_id = line or self.default_device_id

        # verify that device is of type simulated_olt
        device = self.get_device(device_id)
        assert device.type == 'simulated_olt', (
            'Cannot use it on this device type (only on simulated_olt type)')

        requests.get('http://{}/devices/{}/detect_onus'.format(
            self.voltha_sim_rest, device_id
        ))

    complete_arrive_onus = VolthaCli.complete_device

    def get_logical_ports(self, logical_device_id):
        """
        Return the NNI port number and the first usable UNI port of logical
        device, and the vlan associated with the latter.
        """
        stub = self.get_stub()
        ports = stub.ListLogicalDevicePorts(
            voltha_pb2.ID(id=logical_device_id)).items
        nni = None
        unis = []
        for port in ports:
            if port.root_port:
                assert nni is None, "There shall be only one root port"
                nni = port.ofp_port.port_no
            else:
                uni = port.ofp_port.port_no
                uni_device = self.get_device(port.device_id)
                vlan = uni_device.vlan
                unis.append((uni, vlan))

        assert nni is not None, "No NNI port found"
        assert unis, "Not a single UNI?"

        return nni, unis

    def do_install_eapol_flow(self, line):
        """
        Install an EAPOL flow on the given logical device. If device is not
        given, it will be applied to logical device of the last pre-provisioned
        OLT device.
        """
        logical_device_id = line or self.default_logical_device_id

        # gather NNI and UNI port IDs
        nni_port_no, unis = self.get_logical_ports(logical_device_id)

        # construct and push flow rule
        stub = self.get_stub()
        for uni_port_no, _ in unis:
            update = FlowTableUpdate(
                id=logical_device_id,
                flow_mod=mk_simple_flow_mod(
                    priority=2000,
                    match_fields=[in_port(uni_port_no), eth_type(0x888e)],
                    actions=[
                        # push_vlan(0x8100),
                        # set_field(vlan_vid(4096 + 4000)),
                        output(ofp.OFPP_CONTROLLER)
                    ]
                )
            )
            res = stub.UpdateLogicalDeviceFlowTable(update)
            self.poutput('success for uni {} ({})'.format(uni_port_no, res))

    complete_install_eapol_flow = VolthaCli.complete_logical_device

    def do_install_all_controller_bound_flows(self, line):
        """
        Install all flow rules for controller bound flows, including EAPOL,
        IGMP and DHCP. If device is not given, it will be applied to logical
        device of the last pre-provisioned OLT device.
        """
        logical_device_id = line or self.default_logical_device_id

        # gather NNI and UNI port IDs
        nni_port_no, unis = self.get_logical_ports(logical_device_id)

        # construct and push flow rules
        stub = self.get_stub()

        for uni_port_no, _ in unis:
            stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
                id=logical_device_id,
                flow_mod=mk_simple_flow_mod(
                    priority=2000,
                    match_fields=[
                        in_port(uni_port_no),
                        eth_type(0x888e)
                    ],
                    actions=[output(ofp.OFPP_CONTROLLER)]
                )
            ))
            stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
                id=logical_device_id,
                flow_mod=mk_simple_flow_mod(
                    priority=1000,
                    match_fields=[
                        in_port(uni_port_no),
                        eth_type(0x800),
                        ip_proto(2)
                    ],
                    actions=[output(ofp.OFPP_CONTROLLER)]
                )
            ))
            stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
                id=logical_device_id,
                flow_mod=mk_simple_flow_mod(
                    priority=1000,
                    match_fields=[
                        in_port(uni_port_no),
                        eth_type(0x800),
                        ip_proto(17),
                        udp_dst(67)
                    ],
                    actions=[output(ofp.OFPP_CONTROLLER)]
                )
            ))
        self.poutput('success')

    complete_install_all_controller_bound_flows = \
        VolthaCli.complete_logical_device

    def do_install_all_sample_flows(self, line):
        """
        Install all flows that are representative of the virtualized access
        scenario in a PON network.
        """
        logical_device_id = line or self.default_logical_device_id

        # gather NNI and UNI port IDs
        nni_port_no, unis = self.get_logical_ports(logical_device_id)

        # construct and push flow rules
        stub = self.get_stub()

        for uni_port_no, c_vid in unis:
            # Controller-bound flows
            stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
                id=logical_device_id,
                flow_mod=mk_simple_flow_mod(
                    priority=2000,
                    match_fields=[in_port(uni_port_no), eth_type(0x888e)],
                    actions=[
                        # push_vlan(0x8100),
                        # set_field(vlan_vid(4096 + 4000)),
                        output(ofp.OFPP_CONTROLLER)
                    ]
                )
            ))
            stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
                id=logical_device_id,
                flow_mod=mk_simple_flow_mod(
                    priority=1000,
                    match_fields=[eth_type(0x800), ip_proto(2)],
                    actions=[output(ofp.OFPP_CONTROLLER)]
                )
            ))
            stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
                id=logical_device_id,
                flow_mod=mk_simple_flow_mod(
                    priority=1000,
                    match_fields=[eth_type(0x800), ip_proto(17), udp_dst(67)],
                    actions=[output(ofp.OFPP_CONTROLLER)]
                )
            ))

            # Unicast flows:
            # Downstream flow 1
            stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
                id=logical_device_id,
                flow_mod=mk_simple_flow_mod(
                    priority=500,
                    match_fields=[
                        in_port(nni_port_no),
                        vlan_vid(4096 + 1000),
                        metadata(c_vid)  # here to mimic an ONOS artifact
                    ],
                    actions=[pop_vlan()],
                    next_table_id=1
                )
            ))
            # Downstream flow 2
            stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
                id=logical_device_id,
                flow_mod=mk_simple_flow_mod(
                    priority=500,
                    table_id=1,
                    match_fields=[in_port(nni_port_no), vlan_vid(4096 + c_vid)],
                    actions=[set_field(vlan_vid(4096 + 0)), output(uni_port_no)]
                )
            ))
            # Upstream flow 1 for 0-tagged case
            stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
                id=logical_device_id,
                flow_mod=mk_simple_flow_mod(
                    priority=500,
                    match_fields=[in_port(uni_port_no), vlan_vid(4096 + 0)],
                    actions=[set_field(vlan_vid(4096 + c_vid))],
                    next_table_id=1
                )
            ))
            # Upstream flow 1 for untagged case
            stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
                id=logical_device_id,
                flow_mod=mk_simple_flow_mod(
                    priority=500,
                    match_fields=[in_port(uni_port_no), vlan_vid(0)],
                    actions=[push_vlan(0x8100), set_field(vlan_vid(4096 + c_vid))],
                    next_table_id=1
                )
            ))
            # Upstream flow 2 for s-tag
            stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
                id=logical_device_id,
                flow_mod=mk_simple_flow_mod(
                    priority=500,
                    table_id=1,
                    match_fields=[in_port(uni_port_no), vlan_vid(4096 + c_vid)],
                    actions=[
                        push_vlan(0x8100),
                        set_field(vlan_vid(4096 + 1000)),
                        output(nni_port_no)
                    ]
                )
            ))

        # Push a few multicast flows
        # 1st with one bucket for our uni 0
        stub.UpdateLogicalDeviceFlowGroupTable(FlowGroupTableUpdate(
            id=logical_device_id,
            group_mod=mk_multicast_group_mod(
                group_id=1,
                buckets=[
                    ofp.ofp_bucket(actions=[
                        pop_vlan(),
                        output(unis[0][0])
                    ])
                ]
            )
        ))
        stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
            id=logical_device_id,
            flow_mod=mk_simple_flow_mod(
                priority=1000,
                match_fields=[
                    in_port(nni_port_no),
                    eth_type(0x800),
                    vlan_vid(4096 + 140),
                    ipv4_dst(0xe4010101)
                ],
                actions=[group(1)]
            )
        ))

        # 2nd with one bucket for uni 0 and 1
        stub.UpdateLogicalDeviceFlowGroupTable(FlowGroupTableUpdate(
            id=logical_device_id,
            group_mod=mk_multicast_group_mod(
                group_id=2,
                buckets=[
                    ofp.ofp_bucket(actions=[pop_vlan(), output(unis[0][0])])
                    #                    ofp.ofp_bucket(actions=[pop_vlan(), output(unis[1][0])])
                ]
            )
        ))
        stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
            id=logical_device_id,
            flow_mod=mk_simple_flow_mod(
                priority=1000,
                match_fields=[
                    in_port(nni_port_no),
                    eth_type(0x800),
                    vlan_vid(4096 + 140),
                    ipv4_dst(0xe4020202)
                ],
                actions=[group(2)]
            )
        ))

        # 3rd with empty bucket
        stub.UpdateLogicalDeviceFlowGroupTable(FlowGroupTableUpdate(
            id=logical_device_id,
            group_mod=mk_multicast_group_mod(
                group_id=3,
                buckets=[]
            )
        ))
        stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
            id=logical_device_id,
            flow_mod=mk_simple_flow_mod(
                priority=1000,
                match_fields=[
                    in_port(nni_port_no),
                    eth_type(0x800),
                    vlan_vid(4096 + 140),
                    ipv4_dst(0xe4030303)
                ],
                actions=[group(3)]
            )
        ))

        self.poutput('success')

    complete_install_all_sample_flows = VolthaCli.complete_logical_device

    def do_install_dhcp_flows(self, line):
        """
        Install all dhcp flows that are representative of the virtualized access
        scenario in a PON network.
        """
        logical_device_id = line or self.default_logical_device_id

        # gather NNI and UNI port IDs
        nni_port_no, unis = self.get_logical_ports(logical_device_id)

        # construct and push flow rules
        stub = self.get_stub()

        # Controller-bound flows
        for uni_port_no, _ in unis:
            stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
                id=logical_device_id,
                flow_mod=mk_simple_flow_mod(
                    priority=1000,
                    match_fields=[
                        in_port(uni_port_no),
                        eth_type(0x800),
                        ip_proto(17),
                        udp_dst(67)
                    ],
                    actions=[output(ofp.OFPP_CONTROLLER)]
                )
            ))

        self.poutput('success')

    complete_install_dhcp_flows = VolthaCli.complete_logical_device

    def do_delete_all_flows(self, line):
        """
        Remove all flows and flow groups from given logical device
        """
        logical_device_id = line or self.default_logical_device_id
        stub = self.get_stub()
        stub.UpdateLogicalDeviceFlowTable(FlowTableUpdate(
            id=logical_device_id,
            flow_mod=ofp.ofp_flow_mod(
                command=ofp.OFPFC_DELETE,
                table_id=ofp.OFPTT_ALL,
                cookie_mask=0,
                out_port=ofp.OFPP_ANY,
                out_group=ofp.OFPG_ANY
            )
        ))
        stub.UpdateLogicalDeviceFlowGroupTable(FlowGroupTableUpdate(
            id=logical_device_id,
            group_mod=ofp.ofp_group_mod(
                command=ofp.OFPGC_DELETE,
                group_id=ofp.OFPG_ALL
            )
        ))
        self.poutput('success')

    complete_delete_all_flows = VolthaCli.complete_logical_device

    def do_send_simulated_upstream_eapol(self, line):
        """
        Send an EAPOL upstream from a simulated OLT
        """
        device_id = line or self.default_device_id
        requests.get('http://{}/devices/{}/test_eapol_in'.format(
            self.voltha_sim_rest, device_id
        ))

    complete_send_simulated_upstream_eapol = VolthaCli.complete_device

    def do_inject_eapol_start(self, line):
        """
        Send out an an EAPOL start message into the given Unix interface
        """
        pass

if __name__ == '__main__':

    parser = argparse.ArgumentParser()

    _help = '<hostname>:<port> to consul agent (default: %s)' % defs['consul']
    parser.add_argument(
        '-C', '--consul', action='store', default=defs['consul'], help=_help)

    _help = '<hostname>:<port> to etcd container (default: %s)' % defs['etcd']
    parser.add_argument(
        '-E', '--etcd', action='store', default=defs['etcd'], help=_help)

    _help = 'Lookup Voltha endpoints based on service entries in Consul'
    parser.add_argument(
        '-L', '--lookup', action='store_true', help=_help)

    _help = 'All requests to the Voltha gRPC service are global'
    parser.add_argument(
        '-G', '--global_request', action='store_true', help=_help)

    _help = '<hostname>:<port> of Voltha gRPC service (default={})'.format(
        defs['voltha_grpc_endpoint'])
    parser.add_argument('-g', '--grpc-endpoint', action='store',
                        default=defs['voltha_grpc_endpoint'], help=_help)

    _help = '<hostname>:<port> of Voltha simulated adapter backend for ' \
            'testing (default={})'.format(
        defs['voltha_sim_rest_endpoint'])
    parser.add_argument('-s', '--sim-rest-endpoint', action='store',
                        default=defs['voltha_sim_rest_endpoint'], help=_help)

    args = parser.parse_args()

    if args.lookup:
        host = args.consul.split(':')[0].strip()
        port = int(args.consul.split(':')[1].strip())
        consul = Consul(host=host, port=port)

        _, services = consul.catalog.service('voltha-grpc')
        if not services:
            print('No voltha-grpc service registered in consul; exiting')
            sys.exit(1)
        args.grpc_endpoint = '{}:{}'.format(services[0]['ServiceAddress'],
                                            services[0]['ServicePort'])

        _, services = consul.catalog.service('voltha-sim-rest')
        if not services:
            print('No voltha-sim-rest service registered in consul; exiting')
            sys.exit(1)
        args.sim_rest_endpoint = '{}:{}'.format(services[0]['ServiceAddress'],
                                                services[0]['ServicePort'])

    host = args.etcd.split(':')[0].strip()
    port = int(args.etcd.split(':')[1].strip())
    etcd = etcd3.client(host=host, port=port)
    c = VolthaCli(args.grpc_endpoint, args.sim_rest_endpoint, etcd,
                  args.global_request)
    c.poutput(banner)
    c.load_history()
    c.cmdloop()
    c.save_history()


import numpy as np
import sys
import csv
import matplotlib.pyplot as plt

def print_hi(name):
    # Use a breakpoint in the code line below to debug your script.
    print(f'Hi, {name}')  # Press Ctrl+F8 to toggle the breakpoint.

def printPic(x,y,name):
    plt.title('pic')
    colors=['red','yellow','green','blue','cyan','violet']
    for i in range(0,len(y)):
        print(name[i+1])
        plt.scatter(x , y[i],s=30 , marker='o',c=colors[i],label=name[i+1],alpha=0.6)

    plt.show()

if __name__ == '__main__':
    dict_1 = {}

    Name = []
    with open("../scripts/temp.csv", "r") as csvfile:
        reader = csv.reader(csvfile)
        for line in reader:
            if dict_1.get(line[1]) is None:
                dict_1[line[1]] = float(line[7])
                Name = np.append(Name,line[1])
            else:
                dict_1[line[1]] = np.append(dict_1.get(line[1]),float(line[7]))

    Other = []
    for i in range(1,len(Name)):
        Other = np.append(Other,dict_1[Name[i]],axis=0)
    Other = np.reshape(Other,(len(Name) - 1,len(Other)//(len(Name)-1)))

    printPic(dict_1[Name[0]],Other,Name)


from pyspark.sql import SparkSession
from pyspark.sql import Row

spark = SparkSession.builder.appName('filter-example').getOrCreate()
sc = spark.sparkContext
sc.setLogLevel("ERROR")

def run():
    print()
    student = spark.createDataFrame([
        Row(name='Song', subject='English', term='Spring', score=80),
        Row(name='Song', subject='English', term='Fall', score=90),
        Row(name='Song', subject='Math', term='Spring', score=90),
        Row(name='Song', subject='Math', term='Fall', score=100),
        Row(name='Donald trump', subject='English', term='Spring', score=20),
        Row(name='Donald trump', subject='English', term='Fall', score=40),
        Row(name='Donald trump', subject='Math', term='Spring', score=30),
        Row(name='Donald trump', subject='Math', term='Fall', score=25),
    ])

    student.printSchema()
    print(f'No. of partitons - {student.rdd.getNumPartitions()}')
    student.show()

    print('repartition - 3')
    df = student.repartition(3)
    print(f'No. of partitons - {df.rdd.getNumPartitions()}')
    df.show()

    print('repartition - 4')
    df = df.repartition(4)
    print(f'No. of partitons - {df.rdd.getNumPartitions()}')
    df.show()

    print('coalesce - 1')
    df = df.coalesce(1)
    print(f'No. of partitons - {df.rdd.getNumPartitions()}')
    df.show()

    print('coalesce - 2')
    df = df.coalesce(2)
    print(f'No. of partitons - {df.rdd.getNumPartitions()}')
    print('Increase the No. of partitions happened with coalesce')
    df.show()

if __name__ == '__main__':
    run()


import numpy as np

import matplotlib.pyplot as plt

#data generation
X1, Y1 = np.random.multivariate_normal([4,4], [[1,0], [0, 1]], 100).T
X2, Y2 = np.random.multivariate_normal([-3,3], [[1.5,0], [0, .5]], 100).T
X3, Y3 = np.random.multivariate_normal([0,-3], [[1,0], [0, .75]], 100).T

X = np.append(X1, X2)
X = np.append(X, X3)
Y = np.append(Y1, Y2)
Y = np.append(Y, Y3)

data = np.zeros((300, 2))
#likely a more efficent way to do This
for i in range(300):
    data[i, 0] = X[i]
    data[i, 1] = Y[i]

#clustering
def clustering(X, iterations):
    centroids = np.random.rand(3, 2);

    assignments = np.zeros((300));

    dist = np.zeros((3));

    n1 = 0
    n2 = 0
    n3 = 0

    for j in range(iterations):

        for k in range(300):

            for i in range(3):

                dist[i] = np.linalg.norm(data[k, :] - centroids[i, :])

                assignments[k] = np.argmin(dist)

        for k in range(300):

            if assignments[k] == 0:

                plt.scatter(data[k, 0], data[k, 1], color = "red")

                centroids[0, 0] = data[k, 0] + centroids[0, 0]
                centroids[0, 1] = data[k, 1] + centroids[0, 1]

                n1 = n1 + 1

            elif assignments[k] == 1:

                plt.scatter(data[k, 0], data[k, 1], color = "blue")

                centroids[1, 0] = data[k, 0] + centroids[1, 0]
                centroids[1, 1] = data[k, 1] + centroids[1, 1]

                n2 = n2 + 1
            else:

                plt.scatter(data[k, 0], data[k, 1], color = "green")

                centroids[2, 0] = data[k, 0] + centroids[2, 0]
                centroids[2, 1] = data[k, 1] + centroids[2, 1]

                n3 = n3 + 1

        centroids[0,0] = centroids[0,0]/n1
        centroids[0,1] = centroids[0,1]/n1
        centroids[1,0] = centroids[1,0]/n2
        centroids[1,1] = centroids[1,1]/n2
        centroids[2,0] = centroids[2,0]/n3
        centroids[2,1] = centroids[2,1]/n3

    plt.show()

for i in range(3):
    clustering(data, i)


from datetime import datetime
import pykrige
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import mean_squared_error
import h5py
import xarray as xr
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing import MinMaxScaler
from sklearn.ensemble import GradientBoostingRegressor
from sklearn import tree
from sklearn.model_selection import GridSearchCV
from sklearn.model_selection import cross_val_score
from sklearn.model_selection import KFold
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeRegressor
from sklearn.metrics import accuracy_score
from sklearn.metrics import roc_auc_score
from pykrige import variogram_models, UniversalKriging, OrdinaryKriging
from scipy.interpolate import griddata

"""# Define functions"""

def getKNearest(fromX, fromY, measuringPtX, measuringPtY, k):
  # returns a list of k indices of pts from fromX and fromY (assumed to correspond to x and y coordinates of the same pts) that are the nearest to measuringPtX and measuringPtY
  distances = np.sqrt((fromX - measuringPtX)**2 + (fromY - measuringPtY)**2)
  # Get the indices of the k-nearest points
  k_nearest_indices = np.argsort(distances)[:k]
  return k_nearest_indices

def visualizePointsInBox(allTrackX, allTrackY, xCoordMin, xCoordMax, yCoordMin, yCoordMax):
  fig, ax = plt.subplots()
  #ax.scatter(allSurfX, allSurfY, label='Surface points', s=0.2)
  ax.scatter(allTrackX, allTrackY, label='Track bed points', s=0.2)
  # Plot the batches
  # Vertical lines
  ax.axvline(x=xCoordMin, color="black")
  ax.axvline(x=xCoordMax, color="black")
  ax.axhline(yCoordMin, color="black")
  ax.axhline(yCoordMax, color="black")

  ax.set_xlabel('X')
  ax.set_ylabel('Y')
  ax.set_title('Surface data vs track data: batches visualized')
  #ax.legend(loc = "lower right")

  plt.show()
# Store the direction (degrees) of each vector (surf_vx, surf_vy)
def getVectorDir(x, y):
    # Calculate the angle in radians
    #angle_rad = np.arctan2(np.abs(y), x)
    angle_rad = np.arctan2(y,x)
    # Convert the angle from radians to degrees
    angle_deg = np.degrees(angle_rad)

    # Ensure the angle is within the range [0, 360)
    angle_deg = (angle_deg + 360) % 360

    #return abs(angle_deg -90) # Subtract 90 b/c 0 degrees = up in the coordinate system used by pykrige. also it just wants the axis, so we constrain to abs
    return angle_deg -90

# Box settings
SQUAREBOXWIDTH = 50 # m
EXPANDSTEP = 10 # m
SHRINKSTEP = EXPANDSTEP
MAX_NUM_PTS = 80
MIN_NUM_PTS = 15

# Which variogram models to even try in order to find the best fit
METHODS = ['exponential', 'power', 'linear', 'gaussian', 'spherical']

# Pick average of K nearest pts' angles from velocity vectors
ANISOTROPY_ANGLE_METH = "KNN"
K = 1

# Define slice of dataset to run on
START_INDEX = 0
END_SLICE_INDEX = 40

TRACKBED_XIND = 7
TRACKBED_YIND = 8
TRACKBED_TARGETIND = 9

# Display plots of variograms?
PLOT = False
RUN_COLAB = False

VERBOSE_OUTPUT = False

# Extract slice of prediction points array from command line args
import sys

if (len(sys.argv) > 1):
  START_INDEX = int(sys.argv[1])
  END_SLICE_INDEX = int(sys.argv[2])
  print("recieved args: ", START_INDEX, " ", END_SLICE_INDEX)

TRACKBED_XIND = 7
TRACKBED_YIND = 8
TRACKBED_TARGETIND = 9

# Display plots of variograms?
PLOT = False
RUN_COLAB = False

if RUN_COLAB:
  pass
else:
  log_file_name = "output_log_" + "(" + str(START_INDEX) + "," + str(END_SLICE_INDEX) + ")" + ".txt"
  sys.stdout = open(log_file_name, "w")

"""# Import data"""

if RUN_COLAB:
  df_merged = pd.read_csv('/content/drive/MyDrive/Predicted and preprocessed data sets /training_test_combined_vmag_ky_230710.csv')
else:
  df_merged =  pd.read_csv('training_test_combined_vmag_ky_230710.csv')
df_merged.columns

"""# Data Preprocessing"""

# Remove duplicates by setting to mean
df2 = df_merged.groupby(['track_bed_x', 'track_bed_y']).agg({'surf_x':'mean', 'surf_y':'mean', 'surf_vx':'mean', 'surf_vy': 'mean', 'surf_elv': 'mean', 'surf_dhdt':'mean', 'surf_SMB': 'mean', 'track_bed_x':'first', 'track_bed_y':'first', 'track_bed_target':'mean', 'v_mag':'mean'})
print("num duplicates found: ", len(df_merged) - len(df2))
df_merged = df2.reset_index(drop=True)
df_merged.columns
print("Number of rows after duplicates removed: ", df_merged.shape)

"""# Execute kriging"""

trackDatTmp = np.array(df_merged)
trackDatTmp.shape
trackDat_x = trackDatTmp[:, TRACKBED_XIND]
trackDat_y = trackDatTmp[:, TRACKBED_YIND]
trackDat_angle = getVectorDir(trackDatTmp[:, 2], trackDatTmp[:, 3])
len(trackDat_angle)

trackDat = np.column_stack((trackDatTmp, trackDat_angle))
trackDat.shape

# For each point in tracking point dataset
  # Grab the track bed data points that are within  within box (except for that that point)
  # Print size of points grabbed in order to ensure it is reasonable
  # train kriging based on that dataset

boxWidth = SQUAREBOXWIDTH
outputCols = ['surf_x', 'surf_y', 'surf_vx', 'surf_vy','surf_elv', 'surf_dhdt', 'surf_SMB', 'track_bed_x', 'track_bed_y', 'v_mag', 'track_bed_target', 'residual', 'prediction', 'estPredictionError']
out = np.empty( (0, 14) ) # (track_bed_x, track_bed_y, target_prediction, target_actual, residual)

slice = df_merged[START_INDEX:END_SLICE_INDEX].copy()
print("\t".join(str(value) for value in outputCols))
Kriging_startime = datetime.now()
for index, row in slice.iterrows():
  #print("Pixel " + str(index + 1) + " out of " + str(df_merged.shape[0]) )
  # Hold the prediction and ground truth for track_bed_target at the current pixel
  predictionThisRow = None
  groundTruthThisRow = row['track_bed_target']

  trackRowsToTrainOn = None

  # Define a spatial batch centered on the point to make the prediction on
  xMin, xMax, yMin, yMax = (None, None, None, None)
  expandStep = EXPANDSTEP
  shrinkStep = SHRINKSTEP
  previousProblem = 0 # 1 = too big; -1 = too small
  while True:
    #print("Top of loop")

    xMin, xMax = (row['track_bed_x'] - boxWidth, row['track_bed_x'] + boxWidth)
    yMin, yMax = (row['track_bed_y'] - boxWidth, row['track_bed_y'] + boxWidth)

    # We are still fiting kriging on the entire dataset (not just the slice)
    x_selected_ind = np.where((trackDat_x >= xMin) & (trackDat_x <= xMax) )[0]
    y_selected_ind = np.where((trackDat_y >= yMin) & (trackDat_y <= yMax) )[0]
    #visualizePointsInBox(trackDat_x, trackDat_y, xMin, xMax, yMin, yMax)

    # Select indices that are in selected x and y range (inside the rectangle) to init object on
    indicesXY = list(set(x_selected_ind).intersection(y_selected_ind))
    trackRowsToTrainOn = np.array([trackDat[ind] for ind in indicesXY])

    # If none were selected, perform same code as if some but too few were selected
    if (len(trackRowsToTrainOn) < MIN_NUM_PTS):

      # Prevent overshooting by reducing step size
      if (previousProblem == 1): # Was too big last time
        expandStep /= 2
      else:
        pass#expandStep *= 1.5

      boxWidth += expandStep
      previousProblem = -1
      continue

    # Remove the row corresponding to the center point
    mask = np.logical_and(trackRowsToTrainOn[:, TRACKBED_XIND] == row['track_bed_x'], trackRowsToTrainOn[:, TRACKBED_YIND] == row['track_bed_y'])
    trackRowsToTrainOn = trackRowsToTrainOn[~mask] # Grab track rows that exclude it
    #print(trackRowsToTrainOn)

    # Print number of points
    numPts = len(trackRowsToTrainOn)
    #print(numPts)

    # Perform iterative resizing of box to ensure correct length
    if (numPts < MIN_NUM_PTS):

      # Prevent overshooting by reducing step size
      if (previousProblem == 1): # Was too big last time
        expandStep /= 2
      else:
        pass#expandStep *= 1.5

      boxWidth += expandStep
      previousProblem = -1
    elif (numPts > MAX_NUM_PTS):

      # Prevent overshooting by reducing step size
      if (previousProblem == -1): # Was too small last time
        shrinkStep /= 2
      else:
        pass#shrinkStep *= 1.5

      boxWidth -= shrinkStep
      previousProblem = 1
    else:
      break # We have chosen an appropriate number of trackRowsToTrainOn

  # Define settings for UK covariance model fitting
  nlags = len(trackRowsToTrainOn[:, 1])
  kNearest_PtsToTrainOn_indices = getKNearest(trackRowsToTrainOn[:, TRACKBED_XIND], trackRowsToTrainOn[:, TRACKBED_YIND], row['track_bed_x'], row['track_bed_y'], K)
  angles = trackDat_angle[kNearest_PtsToTrainOn_indices]
  anisotropy_angle = angles.mean()

  # Fit each possible variogram model
  linear = UniversalKriging(trackRowsToTrainOn[:, TRACKBED_XIND], trackRowsToTrainOn[:, TRACKBED_YIND], trackRowsToTrainOn[:, TRACKBED_TARGETIND], variogram_model="linear", anisotropy_scaling = 3, anisotropy_angle = anisotropy_angle, enable_plotting=PLOT, nlags=nlags)
  exponential = UniversalKriging(trackRowsToTrainOn[:, TRACKBED_XIND], trackRowsToTrainOn[:, TRACKBED_YIND], trackRowsToTrainOn[:, TRACKBED_TARGETIND], variogram_model="exponential",anisotropy_scaling = 3, anisotropy_angle = anisotropy_angle, enable_plotting=PLOT, nlags=nlags)
  power = UniversalKriging(trackRowsToTrainOn[:, TRACKBED_XIND], trackRowsToTrainOn[:, TRACKBED_YIND], trackRowsToTrainOn[:, TRACKBED_TARGETIND], variogram_model="power", anisotropy_scaling = 3, anisotropy_angle = anisotropy_angle, enable_plotting=PLOT, nlags=nlags)
  gaussian = UniversalKriging(trackRowsToTrainOn[:, TRACKBED_XIND], trackRowsToTrainOn[:, TRACKBED_YIND], trackRowsToTrainOn[:, TRACKBED_TARGETIND], variogram_model="gaussian",anisotropy_scaling = 3, anisotropy_angle = anisotropy_angle, enable_plotting=PLOT, nlags=nlags)
  spherical = UniversalKriging(trackRowsToTrainOn[:, TRACKBED_XIND], trackRowsToTrainOn[:, TRACKBED_YIND], trackRowsToTrainOn[:, TRACKBED_TARGETIND], variogram_model="spherical",anisotropy_scaling = 3, anisotropy_angle = anisotropy_angle, enable_plotting=PLOT, nlags=nlags)

  model_names = ['linear', 'exponential', 'power', 'gaussian', 'spherical']
  models = [linear, exponential, power,  gaussian, spherical]
  stats = []
  for i, model in enumerate(models):
    stats.append(model.get_statistics()[2])

  # Choose the best model as the one with the lowest capacity ratio
  ind_closest_to_one = np.abs(np.array(stats)).argmin()
  kriging = models[ind_closest_to_one]

  # Predict the ground truth pixel
  prediction, predictionVariance = kriging.execute("points", row['track_bed_x'], row['track_bed_y'])
  predictionThisRow = prediction[0]

  residualThisRow = groundTruthThisRow - predictionThisRow
  existingRowData = row[['surf_x', 'surf_y', 'surf_vx', 'surf_vy','surf_elv', 'surf_dhdt', 'surf_SMB', 'track_bed_x', 'track_bed_y', 'v_mag', 'track_bed_target']].copy() # Is this a copy...?
  existingRowData['prediction'] = predictionThisRow
  existingRowData['residual'] = row['track_bed_target'] - existingRowData['prediction']
  existingRowData['estPredictionError'] = predictionVariance
  #print(existingRowData)
  print("\t".join(str(value) for value in existingRowData.values.tolist()))
  out = np.vstack( (out, np.array(existingRowData)))

Kriging_stoptime = datetime.now()
print("---------------------------------------------------------")
df_krigout = pd.DataFrame(columns = outputCols, data=out)
df_krigout.to_csv( "./pixels (" + str(START_INDEX) + "," + str(END_SLICE_INDEX) + ") residuals.csv" )
sys.stdout.close()
sys.stdout = sys.__stdout__
print(f"Finished applying kriging to pixels: {START_INDEX}:{END_SLICE_INDEX}")
print("Time to run: ", Kriging_stoptime - Kriging_startime)


import unittest

from mock import patch, MagicMock
from rawes.elastic import Elastic
from requests.models import Response
from rawes.http_connection import HttpConnection

class TestConnectionPooling(unittest.TestCase):
    """Connection pooling was added on top of Rawes, it wasn't designed from
    the beggingin. We need some tests to ensure our expectations of the
    connection pooling are met.
    """

    def testBasicRoundRobin(self):
        """ Set up a client with three different hosts to connect to, make
        multiple calls and check that each call goes on a different host in a
        Round Robin fashion
        """
        hosts = ['http://someserver1:9200', 'http://someserver2:9200',
                 'http://someserver3:9200']
        es = Elastic(hosts, connection_pool_kwargs={'dead_timeout': 10})
        with patch('rawes.http_connection.requests.Session.request',
                MagicMock(return_value=None)) as request:
            request.return_value = Response()
            called = []
            for _ in xrange(len(hosts)):
                es.get()
                # Save a list of called hosts (and remove trailing /)
                called.append(request.call_args[0][1][:-1])
            # Check against original hosts list
            self.assertSetEqual(set(hosts), set(called),
                            'All hosts in coonnection pool should be used')
            called_again = []
            for _ in xrange(len(hosts)):
                es.get()
                # Call the same hosts again (don't forget about the trailing /)
                called_again.append(request.call_args[0][1][:-1])
            # Check they were called in the same order as before
            self.assertListEqual(called, called_again,
                                    'Round robin order wasn\'t preserved')


from sklearn.metrics import mean_absolute_error,mean_squared_error

def test_mean_absolute_error():

    y_true=[1,1,1,1,1,2,2,2,0,0]
    y_pred=[0,0,0,1,1,1,0,0,0,0]

    print("Mean Absolute Error:",mean_absolute_error(y_true,y_pred))
def test_mean_squared_error():

    y_true=[1,1,1,1,1,2,2,2,0,0]
    y_pred=[0,0,0,1,1,1,0,0,0,0]

    print("Mean Absolute Error:",mean_absolute_error(y_true,y_pred))
    print("Mean Square Error:",mean_squared_error(y_true,y_pred))

if __name__=="__main__":
    test_mean_absolute_error()
    test_mean_squared_error()

12. Model evaluation/12.7 regression_metrics.py


#!../.env/bin/python3
import argparse
import os
import requests
import http.server
import socketserver
import re
from pathlib import Path
from creds import TOKEN

def handler_from_argparse(PATH):
    class HandleAll(http.server.SimpleHTTPRequestHandler):
        def __init__(self, *args, **kwargs):
            os.chdir(PATH)
            super().__init__(*args, **kwargs)
    return HandleAll

def merge(PATH, PORT, TOKEN):
    all_images = check_local_dir(PATH)
    uploaded_images = check_Label_Studio_images(PORT, TOKEN)

    Handler = handler_from_argparse(PATH)

    i = 0
    with socketserver.TCPServer(('', PORT), Handler) as httpd:
        print('\nServing at port:', PORT)
        for img in all_images:
            if img not in uploaded_images and img not in ('Thumbs.db', 'desktop.ini'):
                json_data = {'image': f'http://0.0.0.0:{PORT}/{img}'}

                requests.post('http://localhost:8080/api/projects/1/import',
                              headers={'Authorization': TOKEN if 'Token' in TOKEN else f'Token {TOKEN}'},
                              json=[json_data]).json()
                i += 1
        info = f'Uploaded {i} image(s) ✅.' if i > 0 else 'No images uploaded.'
        print(info)

        try:
            httpd.serve_forever()
        except KeyboardInterrupt:
            httpd.shutdown()

def check_Label_Studio_images(PORT, TOKEN):
    """
    Checks if any objects were previously uploaded either from LS or via API. Use one PORT from the beginning.
    """
    response = requests.get('http://localhost:8080/api/projects/1/tasks/?page=1&page_size=2000',
                            headers={'Authorization': TOKEN if 'Token' in TOKEN else f'Token {TOKEN}'}).json()
    images_range = len(response)
    uploaded_images = set()
    for idx in range(images_range):
        strip_LS_id = re.search(f'(upload/1/[a-z0-9]+-|:{PORT}/)(.+)', response[idx]['data']['image'])
        if strip_LS_id:
            uploaded_images.add(strip_LS_id.group(2))
    return uploaded_images

def check_local_dir(PATH):
    data_path = Path(PATH)
    all_images = os.listdir(data_path)

    return all_images

def parse_options():
    parser = argparse.ArgumentParser()
    parser.add_argument('-path', '--PATH', type=str,
                        default="/run/user/1000/gvfs/smb-share:server=192.168.1.12,share=e/zzzProjekty/labels yolo/all",
                        help='path to the directory containing dataset')
    parser.add_argument('-port', '--PORT', type=int, default=8082, help='specified port (default - 8082),'
                                                                        ' use one to avoid duplicates')
    parser.add_argument('-token', '--TOKEN', type=str, default=TOKEN, help='only: a-z0-9')
    options = parser.parse_args()
    return options

if __name__ == '__main__':
    options = parse_options()
    print(merge(**vars(options)))


import argparse
import colorsys
import glob
import json
import os
import shutil
import sys
import tqdm

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import pycocotools.mask as rletools

from PIL import Image
from multiprocessing import Pool
# from mots_common.io import load_sequences, load_seqmap
from functools import partial
from subprocess import call

from tools.visualize import generate_colors, colormap
from tools.video_utils import frames2video

# =======================================================
# Global variables
# =======================================================
"""
known_tao_ids: set of tao ids that can be mapped exactly to coco ids.
neighbor_classes: tao classes that are similar to coco_classes.
unknown_tao_ids: all_tao_ids that exclude known_tao_ids and neighbor_classes..
"""

IOU_THRESHOLD = 0.5
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = '/'.join(ROOT_DIR.split('/')[:-1])
all_ids = set([i for i in range(1, 1231)])

# Category IDs in TAO that are known (appeared in COCO)
with open(ROOT_DIR + "/datasets/tao/coco_id2tao_id.json") as f:
    coco_id2tao_id = json.load(f)
known_tao_ids = set([v for k, v in coco_id2tao_id.items()])
# Category IDs in TAO that are unknown (comparing to COCO)
unknown_tao_ids = all_ids.difference(known_tao_ids)

# neighbor classes
with open(ROOT_DIR + "/datasets/tao/neighbor_classes.json") as f:
    coco2neighbor_classes = json.load(f)
# Gather tao_ids that can be categorized in the neighbor_classes
neighbor_classes = set()
for coco_id, neighbor_ids in coco2neighbor_classes.items():
    neighbor_classes = neighbor_classes.union(set(neighbor_ids))

# Exclude neighbor classes from unknown_tao_ids
unknown_tao_ids = unknown_tao_ids.difference(neighbor_classes)

# =======================================================
# =======================================================

class SegmentedObject:
    def __init__(self, bbox, mask, score, class_id, track_id):
        self.bbox = bbox
        self.mask = mask
        self.score = score
        self.class_id = class_id
        self.track_id = track_id

def box_IoU_xywh(boxA, boxB):
    """input box: [x,y,w,h]"""
    # convert [x,y,w,h] to [x1,y1,x2,y2]
    xA, yA, wA, hA = boxA
    boxA = [xA, yA, xA + wA, yA + hA]
    xB, yB, wB, hB = boxB
    boxB = [xB, yB, xB + wB, yB + hB]

    # determine the (x, y)-coordinates of the intersection rectangle
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])
    # compute the area of intersection rectangle
    interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)
    # compute the area of both the prediction and ground-truth
    # rectangles
    boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
    boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)
    # compute the intersection over union by taking the intersection
    # area and dividing it by the sum of prediction + ground-truth
    # areas - the interesection area
    iou = interArea / float(boxAArea + boxBArea - interArea)
    # return the intersection over union value
    return iou

# from https://github.com/matterport/Mask_RCNN/blob/master/mrcnn/visualize.py
def apply_mask(image, mask, color, alpha=0.5):
    """Apply the given mask to the image.
    """
    for c in range(3):
        image[:, :, c] = np.where(mask == 1,
                                  image[:, :, c] * (1 - alpha) + alpha * color[c],
                                  image[:, :, c])
    return image

def load_txt(path):
    objects_per_frame = {}
    track_ids_per_frame = {}  # To check that no frame contains two objects with same id
    combined_mask_per_frame = {}  # To check that no frame contains overlapping masks
    with open(path, "r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            fields = line.split(",")

            frame = int(fields[0])
            if frame not in objects_per_frame:
                objects_per_frame[frame] = []
            if frame not in track_ids_per_frame:
                track_ids_per_frame[frame] = set()
            if int(fields[1]) in track_ids_per_frame[frame]:
                assert False, "Multiple objects with track id " + fields[1] + " in frame " + fields[0]
            else:
                track_ids_per_frame[frame].add(int(fields[1]))

            # class_id = int(fields[2])
            class_id = 1
            if not (class_id == 1 or class_id == 2 or class_id == 10):
                assert False, "Unknown object class " + fields[2]

            fields[12] = fields[12].strip()
            score = float(fields[6])
            mask = {'size': [int(fields[10]), int(fields[11])], 'counts': fields[12].encode(encoding='UTF-8')}
            bbox = [float(fields[2]), float(fields[3]), float(fields[4]), float(fields[5])]  # [x, y, w, h]
            # if frame not in combined_mask_per_frame:
            #   combined_mask_per_frame[frame] = mask
            # elif rletools.area(rletools.merge([combined_mask_per_frame[frame], mask], intersect=True)) > 0.0:
            #   assert False, "Objects with overlapping masks in frame " + fields[0]
            # else:
            #   combined_mask_per_frame[frame] = rletools.merge([combined_mask_per_frame[frame], mask], intersect=False)

            objects_per_frame[frame].append(SegmentedObject(
                bbox,
                mask,
                score,
                class_id,
                int(fields[1])
            ))

    return objects_per_frame

def load_sequences(seq_paths):
    objects_per_frame_per_sequence = {}
    print("Loading Sequences")
    for seq_path_txt in tqdm.tqdm(seq_paths):
        seq = seq_path_txt.split("/")[-1].replace(".txt", "")
        if os.path.exists(seq_path_txt):
            objects_per_frame_per_sequence[seq] = load_txt(seq_path_txt)
            # Sort proposals in each frame by score
            for frame_id, props in objects_per_frame_per_sequence[seq].items():
                props.sort(key=lambda p: p.score, reverse=True)
        else:
            assert False, "Can't find data in directory " + seq_path_txt

    return objects_per_frame_per_sequence

def process_sequence(seq_fpaths, tracks_folder, img_folder, output_folder, max_frames, all_frames_dict,
                     annot_frames_dict,
                     topN_proposals, gt_frame2anns, tao_id2name, only_annotated, draw_boxes=True, create_video=True):
    folder_name = tracks_folder.split("/")[-1]
    # print("Processing sequence", seq_name)
    os.makedirs(output_folder, exist_ok=True)
    # tracks = load_sequences(seq_fpaths)
    for seq_fpath in seq_fpaths:
        tracks = load_sequences([seq_fpath])
        seq_id = seq_fpath.split('/')[-1].replace(".txt", "")
        max_frames_seq = max_frames[seq_id]
        all_frames = all_frames_dict[seq_id]
        if seq_id in annot_frames_dict.keys():
            annot_frames = annot_frames_dict[seq_id]

            visualize_sequences(seq_id, tracks, max_frames_seq, img_folder, all_frames, annot_frames, output_folder,
                                topN_proposals, gt_frame2anns, tao_id2name, 'unknown', only_annotated, draw_boxes,
                                create_video)
            visualize_sequences(seq_id, tracks, max_frames_seq, img_folder, all_frames, annot_frames, output_folder,
                                topN_proposals, gt_frame2anns, tao_id2name, 'known', only_annotated, draw_boxes,
                                create_video)
            visualize_sequences(seq_id, tracks, max_frames_seq, img_folder, all_frames, annot_frames, output_folder,
                                topN_proposals, gt_frame2anns, tao_id2name, 'neighbor', only_annotated, draw_boxes,
                                create_video)
        else:
            visualize_all_sequences(seq_id, tracks, max_frames_seq, all_frames, output_folder, topN_proposals, draw_boxes=False,
                                    create_video=True)

def process_sequence_coco(track_result_map, img_id2name, datasrc, img_folder, output_folder, max_frames,
                          annot_frames_dict,
                          topN_proposals, gt_frame2anns, tao_id2name, draw_boxes=True, create_video=True):
    """
    Args:
        track_result_map: Dict { img_id: List[detections] }. tracking result
        img_id2name: a Dict, that maps img_id to img_name (datasrc/video_name/frame_name)
        img_folder:
        output_folder:
        max_frames:
        annot_frames_dict:
        draw_boxes:
        create_video:

    Returns:

    """

    # print("Processing sequence", seq_name)
    os.makedirs(output_folder, exist_ok=True)

    split = img_folder.split('/')
    if "train" in split:
        gt_path = os.path.join(BASE_DIR, "data/TAO/annotations/train.json")
    elif "val" in split:
        gt_path = os.path.join(BASE_DIR, "data/TAO/annotations/validation.json")
    else:
        raise Exception("invalid img_folder")

    with open(gt_path, 'r') as f:
        gt_dict = json.load(f)
    gt_images = gt_dict["images"]
    img_name2id = dict()
    for img in gt_images:
        img_name = img["file_name"]
        if img["id"] not in img_id2name.keys():
            img_name2id[img_name] = -1
        img_name2id[img_name] = img["id"]

    # tracks: List of
    #   { video_name: {frame_id: List[SegmentedObject]}, sorted in the order of frame }
    tracks = dict()
    for video, frames in annot_frames_dict.items():
        frames.sort()
        if video not in tracks.keys():
            tracks[video] = dict()
        for idx, frame in enumerate(frames):
            frame_id = idx + 1
            tracks[video][frame_id] = list()
            frame_name = '/'.join(frame.split('/')[7:])
            img_id = img_name2id[frame_name]
            if img_id not in track_result_map.keys():
                # There are not detections in the current frame
                dets = list()
            else:
                dets = track_result_map[img_id]
            # Convert each dets to SegmentedObject
            dets.sort(key=lambda k: k['score'])
            for det in dets:
                tracks[video][frame_id].append(SegmentedObject(bbox=det["bbox"],
                                                               class_id=det["category_id"],
                                                               track_id=det["track_id"]))
            # only keep detections with top 10 highest scores.
            tracks[video][frame_id] = tracks[video][frame_id][:10]

    # Now we have tracks
    # Get the seq_id, which is just the video name
    # Get the max_frames

    for seq_id in annot_frames_dict.keys():
        annot_frames = annot_frames_dict[seq_id]
        max_frames_seq = len(annot_frames)
        visualize_sequences(seq_id, tracks, max_frames_seq, img_folder, annot_frames, output_folder,
                            topN_proposals, draw_boxes, create_video)

def visualize_all_sequences(seq_id, tracks, max_frames_seq, all_frames, output_folder, topN_proposals, draw_boxes=False,
                            create_video=True):
    colors = colormap()
    dpi = 100.0
    frames_with_annotations = [frame.split('/')[-1] for frame in all_frames]
    all_frame_names = [frame.split('/')[-1] for frame in all_frames]
    for t in range(max_frames_seq):
        print("Processing frame", all_frames[t])
        filename_t = all_frames[t]
        img = np.array(Image.open(filename_t), dtype="float32") / 255
        img_sizes = img.shape
        fig = plt.figure()
        fig.set_size_inches(img_sizes[1] / dpi, img_sizes[0] / dpi, forward=True)
        fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=None, hspace=None)
        ax = fig.subplots()
        ax.set_axis_off()

        if t + 1 in tracks[seq_id].keys():
            for obj in tracks[seq_id][t + 1][:topN_proposals]:
                color = colors[obj.track_id % len(colors)]
                if obj.class_id == 1:
                    category_name = ""
                elif obj.class_id == 2:
                    category_name = "Pedestrian"
                else:
                    category_name = "Ignore"
                    color = (0.7, 0.7, 0.7)

                if obj.class_id == 1 or obj.class_id == 2:  # Don't show boxes or ids for ignore regions
                    x, y, w, h = rletools.toBbox(obj.mask)
                    if draw_boxes:
                        import matplotlib.patches as patches
                        rect = patches.Rectangle((x, y), w, h, linewidth=1,
                                                 edgecolor=color, facecolor='none', alpha=1.0)
                        ax.add_patch(rect)
                    category_name += ":" + str(obj.track_id)
                    # ax.annotate(category_name, (x + 0.5 * w, y + 0.5 * h), color=color, weight='bold',
                    #             fontsize=7, ha='center', va='center', alpha=1.0)
                binary_mask = rletools.decode(obj.mask)
                apply_mask(img, binary_mask, color)

        ax.imshow(img)
        if not os.path.exists(os.path.join(output_folder + "/" + seq_id)):
            os.makedirs(os.path.join(output_folder + "/" + seq_id))
        fig.savefig(output_folder + "/" + seq_id + "/" + all_frame_names[t])
        plt.close(fig)
    if create_video:
        fps = 10
        frames2video(pathIn=output_folder + "/" + seq_id + '/',
                     pathOut=output_folder + "/" + seq_id + ".mp4", fps=fps)

        # Delete the frames
        shutil.rmtree(output_folder + "/" + seq_id)

def visualize_sequences(seq_id, tracks, max_frames_seq, img_folder, all_frames, annot_frames, output_folder,
                        topN_proposals, gt_frame2anns, tao_id2name, split, only_annotated=True, draw_boxes=True,
                        create_video=True):
    # colors = generate_colors(min(60, topN_proposals))
    colors = colormap()
    dpi = 100.0
    # frames_with_annotations = [frame for frame in tracks.keys() if len(tracks[frame]) > 0]
    # img_sizes = next(iter(tracks[frames_with_annotations[0]])).mask["size"]
    frames_with_annotations = [frame.split('/')[-1] for frame in annot_frames]
    all_frame_names = [frame.split('/')[-1] for frame in all_frames]
    # img_sizes = next(iter(tracks[seq_id])).bbox
    if only_annotated:
        frame_idx_with_annotations = list()
        for frame_name in frames_with_annotations:
            frame_idx = all_frame_names.index(frame_name)
            frame_idx_with_annotations.append(frame_idx)
        annot_frame_idx = 0
    for t in range(max_frames_seq):
        if only_annotated:
            if t != frame_idx_with_annotations[annot_frame_idx]:
                continue
            else:
                annot_frame_idx += 1
        print("Processing frame", all_frames[t])
        # filename_t = img_folder + "/" + seq_id + "/" + frames_with_annotations[t]
        filename_t = all_frames[t]
        img = np.array(Image.open(filename_t), dtype="float32") / 255
        img_sizes = img.shape
        fig = plt.figure()
        fig.set_size_inches(img_sizes[1] / dpi, img_sizes[0] / dpi, forward=True)
        fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=None, hspace=None)
        ax = fig.subplots()
        ax.set_axis_off()

        # GT bbox and masks
        valid_cat_ids = list()
        if split == 'known':
            valid_cat_ids = known_tao_ids
        elif split == "neighbor":
            valid_cat_ids = neighbor_classes
        elif split == "unknown":
            valid_cat_ids = unknown_tao_ids
        # Filter out gt-annotations that don't belong to current split
        gt_anns_filtered = list()
        k = '/'.join(filename_t.split('/')[-2:]).replace(".jpg", "")
        for ann in gt_frame2anns[k]:
            if ann["category_id"] in valid_cat_ids:
                gt_anns_filtered.append(ann)

        # Apply bbox and maks from Proposals to image
        if t + 1 in tracks[seq_id].keys():
            for obj in tracks[seq_id][t + 1][:topN_proposals]:
                # Compare obj.bbox with every GT-bbox
                # bbox from GT and proposal is in the form of [x,y,w,h]
                matched_gt = None
                max_IoU_score = 0
                for ann in gt_anns_filtered:
                    iou = box_IoU_xywh(obj.bbox, ann['bbox'])
                    if iou >= IOU_THRESHOLD and iou > max_IoU_score:
                        matched_gt = ann
                        max_IoU_score = max(max_IoU_score, iou)
                if not matched_gt:
                    continue

                color = colors[obj.track_id % len(colors)]
                if obj.class_id == 1:
                    category_name = ""
                elif obj.class_id == 2:
                    category_name = "Pedestrian"
                else:
                    category_name = "Ignore"
                    color = (0.7, 0.7, 0.7)

                if obj.class_id == 1 or obj.class_id == 2:  # Don't show boxes or ids for ignore regions
                    x, y, w, h = obj.bbox
                    if draw_boxes:
                        rect = patches.Rectangle((x, y), w, h, linewidth=1,
                                                 edgecolor=color, facecolor='none', alpha=1.0)
                        ax.add_patch(rect)
                    category_name += ":" + str(obj.track_id)
                    # ax.annotate(category_name, (x + 0.5 * w, y + 0.5 * h), color=color, weight='bold',
                    #             fontsize=7, ha='center', va='center', alpha=1.0)
                binary_mask = rletools.decode(obj.mask)
                apply_mask(img, binary_mask, color)
            # Draw GT bbox
            color_gt = (1.0, 0.0, 0.0)
            for ann in gt_anns_filtered:
                x, y, w, h = ann['bbox']
                rect = patches.Rectangle((x, y), w, h, linewidth=5, linestyle='-.',
                                         edgecolor=color_gt, facecolor='none', alpha=0.5)
                ax.add_patch(rect)
                category_name = tao_id2name[ann["category_id"]]
                ax.annotate(category_name, (x + 0.5 * w, y + 0.5 * h), color=color_gt, weight='bold',
                            fontsize=7, ha='center', va='center', alpha=1.0)

        ax.imshow(img)
        if not os.path.exists(os.path.join(output_folder + "/" + seq_id)):
            os.makedirs(os.path.join(output_folder + "/" + seq_id))
        fig.savefig(output_folder + "/" + seq_id + "/" + all_frame_names[t])
        plt.close(fig)

    if create_video:
        fps = 1 if only_annotated else 10
        frames2video(pathIn=output_folder + "/" + seq_id + '/',
                     pathOut=output_folder + '/' + split + "/" + seq_id + ".mp4", fps=fps)

        # Delete the frames
        shutil.rmtree(output_folder + "/" + seq_id)

        # os.chdir(output_folder + "/" + seq_id)
        # call(["ffmpeg", "-framerate", "10", "-y", "-i", "%06d.jpg", "-c:v", "libx264", "-profile:v", "high", "-crf",
        #       "20",
        #       "-pix_fmt", "yuv420p", "-vf", "pad=\'width=ceil(iw/2)*2:height=ceil(ih/2)*2\'", "output.mp4"])

def load_and_preprocessing_gt(gt_path, datasrc):
    with open(gt_path, 'r') as f:
        gt = json.load(f)

    # Get map: image_id --> frame_name
    img_id2frame_name = dict()
    frame_name2ima_id = dict()
    for image in gt["images"]:
        img_id2frame_name[image["id"]] = image["file_name"]
        frame_name2ima_id[image["file_name"]] = image["id"]

    # Get map: frame_name --> annotations
    frame_name2anns = dict()
    for ann in gt["annotations"]:
        frame_path = img_id2frame_name[ann["image_id"]]
        curr_datasrc = frame_path.split('/')[1]
        if curr_datasrc == datasrc:
            video_name = frame_path.split('/')[-2]
            frame_name = frame_path.split('/')[-1].replace(".jpg", "")
            k = video_name + '/' + frame_name
            if k not in frame_name2anns.keys():
                frame_name2anns[k] = list()
            frame_name2anns[k].append(ann)

    tao_id2name = dict()
    for cat in gt["categories"]:
        tao_id2name[cat["id"]] = cat["name"]

    # There are soem frames although listed as annotated frames, but have no annotations
    with open(ROOT_DIR + '/datasets/tao/val_annotated_{}.txt'.format(datasrc), 'r') as f:
        content = f.readlines()
    for frame in content:
        frame = frame.strip()
        frame_name = '/'.join(frame.split('/')[2:]).replace(".jpg", "")
        if frame_name not in frame_name2anns.keys():
            # print(datasrc + '/' + frame_name)
            frame_name2anns[frame_name] = list()

    return frame_name2anns, tao_id2name

def main():
    # if len(sys.argv) != 5:
    #     print("Usage: python visualize_mots.py tracks_folder(gt or tracker results) img_folder output_folder seqmap")
    #     sys.exit(1)

    # tracks_folder = sys.argv[1]
    # img_folder = sys.argv[2]
    # output_folder = sys.argv[3]
    # seqmap_filename = sys.argv[4]

    parser = argparse.ArgumentParser(description='Visualization script for tracking result')
    parser.add_argument("--tracks_folder", type=str, default="/home/kloping/OpenSet_MOT/Tracking/unovost_RAFT_noReID/")
    parser.add_argument("--gt_path", type=str, default="/home/kloping/OpenSet_MOT/data/TAO/annotations/validation.json")
    parser.add_argument("--img_folder", type=str, default="/home/kloping/OpenSet_MOT/data/TAO/frames/val/")
    parser.add_argument("--datasrc", type=str, default='LaSOT')
    parser.add_argument("--phase", default="objectness", help="objectness, score or one_minus_bg_score", type=str)
    parser.add_argument("--only_annotated", action="store_true")
    parser.add_argument("--tao_subset", action="store_true", help="if only process fixed tao-subsets")
    parser.add_argument("--topN_proposals", default="1000",
                        help="for each frame, only display top N proposals (according to their scores)", type=int)
    parser.add_argument("--track_format", help="The file format of the tracking result", type=str, default='mot')
    args = parser.parse_args()

    if args.track_format == "mot":
        """
        The tracking result is stored in the mot-format.
            https://motchallenge.net/instructions/
        """
        curr_data_src = args.datasrc

        tracks_folder = args.tracks_folder + '/_' + args.phase + '/' + curr_data_src
        img_folder = args.img_folder + '/' + curr_data_src
        output_folder = args.tracks_folder + "/viz_" + args.phase + str(args.topN_proposals) + '/' + curr_data_src
        topN_proposals = args.topN_proposals
        print("For each frame, display top {} proposals".format(topN_proposals))

        if args.tao_subset:
            with open('../datasets/tao/tao_val_subset.txt', 'r') as f:
                content = f.readlines()
            content = [c.strip() for c in content]
            seqmap_filenames = [
                os.path.join(args.tracks_folder, '_' + args.phase, curr_data_src, c.split('/')[-1] + '.txt')
                for c in content if c.split('/')[0] == curr_data_src]
        else:
            seqmap_filenames = sorted(glob.glob(tracks_folder + '/*' + '.txt'))

        # Pre-processing GT
        gt_frame_name2anns, tao_id2name = load_and_preprocessing_gt(args.gt_path, curr_data_src)

        # Image path in all sequences
        all_frames = dict()  # {seq_name: List[frame_paths]}
        sequences = [fn.split('/')[-1] for fn in sorted(glob.glob(os.path.join(img_folder, '*')))]
        for video in sequences:
            fpath = os.path.join(img_folder, video)
            frames = sorted(glob.glob(fpath + '/*' + '.jpg'))
            all_frames[video] = frames

        annot_frames = dict()  # annotated frames for each sequence
        if args.only_annotated:
            # Get the annotated frames in the current sequence.
            txt_fname = ROOT_DIR + "/datasets/tao/val_annotated_{}.txt".format(curr_data_src)
            with open(txt_fname) as f:
                content = f.readlines()
            content = ['/'.join(c.split('/')[2:]) for c in content]
            annot_seq_paths = [os.path.join(img_folder, x.strip()) for x in content]

            for s in annot_seq_paths:
                seq_name = s.split('/')[-2]
                if seq_name not in annot_frames.keys():
                    annot_frames[seq_name] = []
                annot_frames[seq_name].append(s)

        max_frames = dict()
        for seq_name, frames in annot_frames.items():
            max_frames[seq_name] = len(frames)

        for seq_fpath in seqmap_filenames:
            seq_name = seq_fpath.split('/')[-1][:-4]
            max_frames[seq_name] = len(all_frames[seq_name])

        process_sequence_part = partial(process_sequence, max_frames=max_frames, all_frames_dict=all_frames,
                                        annot_frames_dict=annot_frames,
                                        tracks_folder=tracks_folder, img_folder=img_folder,
                                        gt_frame2anns=gt_frame_name2anns, tao_id2name=tao_id2name,
                                        only_annotated=args.only_annotated,
                                        output_folder=output_folder, topN_proposals=topN_proposals)
        process_sequence_part(seqmap_filenames)

    elif args.track_format == "coco":
        """
        The tracking result is stored in the coco-format.
            https://github.com/TAO-Dataset/tao/blob/master/docs/evaluation.md
        """
        max_frames = dict()  # dummy variable here that is not used.
        datasrc = "Charades"  # ["ArgoVerse", "BDD", "Charades", "LaSOT", "YFCC100M"]

        tracks_folder = "/home/kloping/OpenSet_MOT/Tracking/tao_track_results/SORT_val/sort_tao_val_results.json"
        img_folder = "/home/kloping/OpenSet_MOT/data/TAO/frames/val/"
        output_folder = "/home/kloping/OpenSet_MOT/Tracking/tao_track_results/viz/SORT_val/" + datasrc

        annot_frames = dict()  # annotated frames for each sequence
        # Get the annotated frames in the current sequence.
        txt_fname = "../data/tao/val_annotated_{}.txt".format(datasrc)
        with open(txt_fname) as f:
            content = f.readlines()
        content = ['/'.join(c.split('/')[1:]) for c in content]
        annot_seq_paths = [os.path.join(img_folder, datasrc, x.strip()) for x in content]

        for s in annot_seq_paths:
            seq_name = s.split('/')[-2]
            if seq_name not in annot_frames.keys():
                annot_frames[seq_name] = []
            annot_frames[seq_name].append(s)
            annot_frames[seq_name].sort()

        # Load tracking result json file
        with open(tracks_folder, 'r') as f:
            track_results = json.load(f)

        # Get the (train/val) GT annotation json from img_folder
        #   expected img_folder: /home/kloping/OpenSet_MOT/data/TAO/frames/val/
        split = img_folder.split('/')
        if "train" in split:
            gt_path = os.path.join(BASE_DIR, "data/TAO/annotations/train.json")
        elif "val" in split:
            gt_path = os.path.join(BASE_DIR, "data/TAO/annotations/validation.json")
        else:
            raise Exception("invalid img_folder")

        with open(gt_path, 'r') as f:
            gt_dict = json.load(f)
        gt_images = gt_dict["images"]
        img_id2name = dict()
        for img in gt_images:
            img_name = img["file_name"][4:-4]
            if img["id"] not in img_id2name.keys():
                img_id2name[img["id"]] = ''
            img_id2name[img["id"]] = img_name

        # Filter, let the images in the current datasrc remain.
        track_result_map = dict()  # {img_id: List[detections]}
        for det in track_results:
            img_name = img_id2name[det["image_id"]]
            video_name = '/'.join(img_name.split('/')[:2])  # datasrc/video_name
            curr_d = img_name.split('/')[0]
            if curr_d == datasrc:
                if det["image_id"] not in track_result_map.keys():
                    track_result_map[det["image_id"]] = list()
                track_result_map[det["image_id"]].append(det)

        process_sequence_part = partial(process_sequence_coco, max_frames=max_frames, datasrc=datasrc,
                                        annot_frames_dict=annot_frames,
                                        img_id2name=img_id2name, img_folder=img_folder, output_folder=output_folder)
        process_sequence_part(track_result_map)

    # with Pool(10) as pool:
    #     pool.map(process_sequence_part, seqmap)
    # for seq in seqmap:
    #  process_sequence_part(seq)

if __name__ == "__main__":
    main()


# -*- coding: utf-8 -*-
# @Time    : 2020/8/24 9:22
# @Author  : Bo Hu
# @Email   : hubosist@mail.ustc.edu.cn
# @File    : test_merge.py
# @Software: PyCharm

import os
import cupy as cp
import SimpleITK as sitk
from tqdm import tqdm
import numpy as np
import time
from skimage import morphology
from merge_func1 import merge_func

block_name = "z-3500_y-21034_x-24866_1"
root_folder = os.path.join("/braindat/lab/hubo/DATASET/FAFB_data/test_set/testset", block_name)
id_save_root_folder = os.path.join(root_folder, "Cupy_Merged")
seg_file = os.path.join(root_folder, "Test_Results", block_name + "_seg.nii.gz")
mask_file = os.path.join(root_folder, "Step3_result.nii.gz")

def check_dir(dir):
    if not os.path.exists(dir):
        os.mkdir(dir)

def load_nii(path):
    nii = sitk.ReadImage(path)
    nii = sitk.GetArrayFromImage(nii)
    print(path.split("\\")[-1], "loaded!")
    return nii

def save_nii(img,path):
    print(img.shape)
    img = sitk.GetImageFromArray(img)
    sitk.WriteImage(img, path)
    print(path.split("\\")[-1], "saving succeed!")

if __name__ == '__main__':
    check_dir(id_save_root_folder)

    seg = load_nii(seg_file)
    mask = load_nii(mask_file)

    merged = merge_func.main_merge(mask, seg)

    save_nii(merged, "/braindat/lab/hubo/TEMP/test_merged.nii.gz")


from fastapi.testclient import TestClient

from utils.auth_utils import ADMIN_HEADER, BASIC_HEADER, MANAGER_HEADER

category = {"name": "test", "description": "test"}

def test_put_category_as_admin(client: TestClient):
    response = client.put(
        "/categoria/1",
        json=category,
        headers=ADMIN_HEADER
    )
    assert response.status_code == 200
    verify = client.get("/categoria?category_id=1")
    assert verify.json()["data"]["name"] == "test"
    assert verify.json()["data"]["description"] == "test"

def test_put_category_invalid_name_as_admin(client: TestClient):
    category["name"] = None
    response = client.put(
        "/categoria/1",
        json=category,
        headers=ADMIN_HEADER
    )
    assert response.status_code == 422

def test_put_invalid_category_as_admin(client: TestClient):
    category["name"] = "test"
    response = client.put(
        "/categoria/50",
        json=category,
        headers=ADMIN_HEADER
    )
    assert response.status_code == 200
    assert response.json()["message"] == "Categoria não encontrada"

def test_put_category_as_manager(client: TestClient):
    response = client.put(
        "/categoria/2",
        json=category,
        headers=MANAGER_HEADER
    )
    assert response.status_code == 200
    verify = client.get("/categoria?category_id=2")
    assert verify.json()["data"]["name"] == "test"
    assert verify.json()["data"]["description"] == "test"

def test_put_category_as_basic(client: TestClient):
    response = client.put(
        "/categoria/3",
        json=category,
        headers=BASIC_HEADER
    )
    assert response.status_code == 401
    assert response.json()["message"] == "Acesso negado"


import pika
import json
from class_types import Endpoint
import requests

#connection = pika.BlockingConnection(pika.ConnectionParameters('localhost'))
#channel = connection.channel()
endpoints = {}
with open("endpoints.json") as endpoints_config:
    data = json.load(endpoints_config)
    for idx,ep in enumerate(data['services']):
        endpoints[ep['domain']] = Endpoint(ep['domain'],ep['ip'],ep['port'])

to = 'mpcs.uchicago.51205@gmail.com' # send it to itself
dummy_email = {'to':to,'subject':'Holiday Wish from Team 3', 'body':'Happy Thanksgiving !'}
headers = {'Content-Type':'application/json'}
#channel.queue_declare(queue='email_queue')
#channel.basic_publish(exchange='',
#                      routing_key='email_queue',
#                      body=json.dumps(dummy_email))
r= requests.post(endpoints['mediator'].get_prefix() + "send_email", data=json.dumps(dummy_email),headers=headers)
print(r.json(),r.status_code)
print(" [x] Sent 'dummy email!'")
#connection.close()


import numpy as np
import torch
from torch.optim.lr_scheduler import StepLR
from collections import defaultdict
from tqdm import tqdm
from topmost.utils import static_utils

# transform tensor list to numpy list
def to_nparray(tensor_list):
    return np.asarray([item.detach().cpu().numpy() for item in tensor_list], dtype=object)

class HierarchicalTrainer:
    def __init__(self, model, dataset_handler, epochs=200, learning_rate=0.002, batch_size=200, lr_scheduler=None, lr_step_size=125, log_interval=5):
        self.model = model
        self.dataset_handler = dataset_handler
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.batch_size = batch_size
        self.lr_scheduler = lr_scheduler
        self.lr_step_size = lr_step_size
        self.log_interval = log_interval

    def make_optimizer(self,):
        args_dict = {
            'params': self.model.parameters(),
            'lr': self.learning_rate,
        }

        optimizer = torch.optim.Adam(**args_dict)
        return optimizer

    def make_lr_scheduler(self, optimizer,):
        lr_scheduler = StepLR(optimizer, step_size=self.lr_step_size, gamma=0.5, verbose=True)
        return lr_scheduler

    def train(self):
        optimizer = self.make_optimizer()

        if self.lr_scheduler:
            print("===>Warning: use lr_scheduler")
            lr_scheduler = self.make_lr_scheduler(optimizer)

        data_size = len(self.dataset_handler.train_dataloader.dataset)

        for epoch in tqdm(range(1, self.epochs + 1), leave=False):
            self.model.train()
            loss_rst_dict = defaultdict(float)

            for batch_data in self.dataset_handler.train_dataloader:

                rst_dict = self.model(batch_data)
                batch_loss = rst_dict['loss']

                optimizer.zero_grad()
                batch_loss.backward()
                optimizer.step()

                for key in rst_dict:
                    loss_rst_dict[key] += rst_dict[key] * len(batch_data)

            if self.lr_scheduler:
                lr_scheduler.step()

            if epoch % self.log_interval == 0:
                output_log = f'Epoch: {epoch:03d}'
                for key in loss_rst_dict:
                    output_log += f' {key}: {loss_rst_dict[key] / data_size :.3f}'

                print(output_log)

    def test(self, bow):
        data_size = bow.shape[0]

        num_topics_list = self.model.num_topics_list
        theta_list = [list() for _ in range(len(num_topics_list))]
        all_idx = torch.split(torch.arange(data_size), self.batch_size)

        with torch.no_grad():
            self.model.eval()
            for idx in all_idx:
                batch_input = bow[idx]
                batch_theta_list = self.model.get_theta(batch_input)

                for layer_id in range(len(num_topics_list)):
                    theta_list[layer_id].extend(batch_theta_list[layer_id].cpu().numpy().tolist())

        theta = np.empty(len(num_topics_list), object)
        theta[:] = [np.asarray(item) for item in theta_list]
        return theta

    def export_phi(self):
        phi = to_nparray(self.model.get_phi_list())
        return phi

    def export_beta(self):
        beta_list = to_nparray(self.model.get_beta())
        return beta_list

    def export_top_words(self, num_top_words=15):
        beta = self.export_beta()
        top_words_list = list()
        for layer in range(beta.shape[0]):
            print(f"======= Layer: {layer} number of topics: {beta[layer].shape[0]} =======")
            top_words = static_utils.print_topic_words(beta[layer], vocab=self.dataset_handler.vocab, num_top_words=num_top_words)
            top_words_list.append(top_words)
        return top_words_list

    def export_theta(self):
        train_theta = self.test(self.dataset_handler.train_bow)
        test_theta = self.test(self.dataset_handler.test_bow)
        return train_theta, test_theta


from typing import Union, Dict, List, AnyStr, Any, TextIO, Literal

import tempfile
import xmltodict

from commons import *
from varvault import ResourceNotFoundError

vault_file_new = f"{DIR}/new-vault.xml"
existing_vault = f"{DIR}/existing-vault.xml"

class XmlResource(varvault.BaseResource):

    KEY = "VAULT"

    def __init__(self, path: str, mode: Literal["r", "w", "a", "r+", "w+", "a+"] = "r"):
        self.file_io = None
        super(XmlResource, self).__init__(path, mode)

    @property
    def resource(self) -> TextIO:
        return self.file_io

    def create(self) -> None:
        """Creates the resource self.file_io for this handler which we'll use to read and write to."""
        path = self.path
        assert path, "Path is not defined"
        dirname = os.path.dirname(path)

        create_dir = lambda: os.makedirs(dirname, exist_ok=True) if dirname else None
        write = lambda: self.do_write({})

        if self.mode_properties.create and not self.mode_properties.load:
            create_dir()
            write()

        elif self.mode_properties.load and self.mode_properties.create:
            if not self.exists():
                create_dir()
                write()

        elif self.mode_properties.load and not self.mode_properties.create:
            try:
                self.do_read()
            except Exception as e:
                if not self.mode_properties.live_update:
                    raise ResourceNotFoundError(f"Unable to read from resource at {path} (mode is {self.mode})", self) from e
                else:
                    return
        else:
            raise NotImplementedError(f"Mode {self.mode} is not valid ({self.mode_properties})")

        self.file_io = open(self.path, "r+")
        self.file_io.close()

    @property
    def state(self):
        import hashlib
        hash_md5 = hashlib.md5()
        with open(self.path, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hash_md5.update(chunk)
        return hash_md5.hexdigest()

    @property
    def path(self) -> Union[AnyStr, Any]:
        return self.raw_path

    def writable(self, obj: Dict) -> bool:
        obj = {self.KEY: obj}
        try:
            xmltodict.unparse(obj, pretty=True)
            return True
        except:
            return False

    def exists(self) -> bool:
        return os.path.exists(self.path)

    def do_write(self, vault: dict):
        # An XML file must have exactly one root. We assign the vault to the root self.KEY
        vault = {self.KEY: vault}
        xmltodict.unparse(vault, open(self.path, "w"), pretty=True)

    def do_read(self) -> Union[Dict, List]:
        vault = xmltodict.parse(open(self.path, "rb"))
        return vault[self.KEY]

class Keyring(varvault.Keyring):
    string_value = varvault.Key("string_value")
    int_value = varvault.Key("int_value")
    list_value = varvault.Key("list_value")

class TestXmlVault:
    @classmethod
    def setup_class(cls):
        tempfile.tempdir = "/tmp" if sys.platform == "darwin" or sys.platform == "linux" else tempfile.gettempdir()

    def setup_method(self):
        try:
            os.remove(vault_file_new)
        except:
            pass

    def test_create_xml_vault(self):
        vault = varvault.create(keyring=Keyring, resource=XmlResource(vault_file_new, mode="w"))

        @vault.manual(output=(Keyring.string_value, Keyring.int_value, Keyring.list_value))
        def _set():
            return "valid", 1, [1, 2, 3, 4, 5]

        _set()

        data_in_file = xmltodict.parse(open(vault_file_new, "rb"))[XmlResource.KEY]
        logger.info(data_in_file)
        assert Keyring.string_value in data_in_file and data_in_file[Keyring.string_value] == "valid"
        # XML data is actually just strings. But we still want to support using XML files, but it will take some manual labour to get it working correctly.
        assert Keyring.int_value in data_in_file and data_in_file[Keyring.int_value] == "1"
        assert Keyring.list_value in data_in_file

    def test_read_from_xml_vault(self):
        vault = varvault.create(keyring=Keyring, resource=XmlResource(existing_vault, "r"))

        @vault.manual(input=(Keyring.string_value, Keyring.int_value))
        def _get(string_value=None, int_value=None):
            assert string_value == "valid"
            assert int_value == "1"

        _get()


#Библиотеки
import pygame
import random
import sys
import time #Чтобы сделать задержку

#Размер экрана
MAX_X = 1366
MAX_Y = 768

#Количество снежинок
MAX_SNOW = 150

#Размер снежинки
SNOW_SIZE = 64 #В пикселях

#Класс снежинки
class Snow():
    def __init__(self, x, y): #x и y - координаты снежинки
        """Инициализация снежинки"""
        self.xx = x
        self.yy = y
        self.speed = random.randint(5, 10) #От 5 до 10
        self.img_num = random.randint(1, 4)
        self.image_filename = "snow" + str(self.img_num) + ".PNG"
        self.image = pygame.image.load(self.image_filename).convert_alpha() #Загружаем картинку
        self.image = pygame.transform.scale(self.image, (SNOW_SIZE, SNOW_SIZE)) #Уменьшаем картинку по x и y на SNOW_SIZE

    def move_snow(self):
        """Перемещение снежинки"""
        self.yy = self.yy + self.speed #Движение снежинки вниз
        if self.yy > MAX_Y: #Если снежинка дошла до самого нижней границы экрана, то она должна появиться сверху
            self.yy = (0 - SNOW_SIZE)

        i = random.randint(1, 3)
        if i == 1: #Движение направо
            self.xx = self.xx + 10
            if self.xx > MAX_X:
                self.xx = (0 - SNOW_SIZE)
        elif i == 2: #Движение налево
            self.xx = self.xx - 10
            if self.xx < (0 - SNOW_SIZE):
                self.xx = MAX_X

    def draw_snow(self):
        screen.blit(self.image, (self.xx, self.yy))
        pygame.display.flip()
"""Конец класса Snow()"""

def initialize_snow(MAX_SNOW, snowfall):
    """Функция создает снежинки на экране"""
    #Теперь собственно создание массива объектов, используя класс
    for i in range(0, MAX_SNOW):
        xxx = random.randint(0, MAX_X)
        yyy = random.randint(0, MAX_Y)
        snowfall.append(Snow(xxx, yyy)) #Берем класс Snow и перекидываем туда координаты снежинки, так как def __init__(self, x, y) | затем добавляем в массив

def check_for_exit():
    """Функция, которая осуществляет выход"""
    for event in pygame.event.get():
        if event.type == pygame.KEYDOWN:
            sys.exit()

#---------------------MAIN---------------------
#pygame.init
screen = pygame.display.set_mode((MAX_X, MAX_Y))
bg_color = (255, 255, 255) #0,0,0 - black | 255,255,255 - white
snowfall = [] #Для массива снежинок

initialize_snow(MAX_SNOW, snowfall)

while True:
    screen.fill(bg_color) #Закрашиваем старое изображение картинки перед ее отрисовкой
    check_for_exit()
    for i in snowfall:
        i.move_snow() #Передача методов
        i.draw_snow() #Передача методов
    time.sleep(0.10)
    pygame.display.flip()