text:
import datetime
import numpy as np
import pandas as pd
import ctypes
from enum import Enum

class TimePrecision(Enum):
    SECONDS = 0
    MILLISECONDS = 3
    MICROSECONDS = 6
    NANOSECONDS = 9

from types import MethodType
from ._mutators import set_tdf, get_tdf
from .ipc import load_buffer, shmdt

def seconds_to_time(seconds):
    """Convert seconds since midnight to a datetime.time"""
    m, s = divmod(seconds, 60)
    h, m = divmod(m, 60)
    return datetime.time(h, m, s)

def time_to_seconds(time):
    """Convert a datetime.time to seconds since midnight"""
    if time is None:
        return None
    return 3600 * time.hour + 60 * time.minute + time.second

def datetime_to_seconds(arr, precision):
    """Convert an array of datetime64[ns] to seconds since the UNIX epoch"""

    p = TimePrecision(precision)
    if arr.dtype != np.dtype('datetime64[ns]'):
        if arr.dtype == 'int64':
            # The user has passed a unix timestamp already
            return arr

        if not (
            arr.dtype == 'object'
            or str(arr.dtype).startswith('datetime64[ns,')
        ):
            raise TypeError(
                f"Invalid dtype '{arr.dtype}', expected one of: "
                "datetime64[ns], int64 (UNIX epoch), "
                "or object (string)"
            )

        # Convert to datetime64[ns] from string
        # Or from datetime with timezone information
        # Return timestamp in 'UTC'
        arr = pd.to_datetime(arr, utc=True)
        return arr.view('i8') // 10 ** 9  # ns -> s since epoch
    else:
        if p == TimePrecision.SECONDS:
            return arr.view('i8') // 10 ** 9
        elif p == TimePrecision.MILLISECONDS:
            return arr.view('i8') // 10 ** 6
        elif p == TimePrecision.MICROSECONDS:
            return arr.view('i8') // 10 ** 3
        elif p == TimePrecision.NANOSECONDS:
            return arr.view('i8')

def datetime_in_precisions(epoch, precision):
    """Convert epoch time value into s, ms, us, ns"""
    p = TimePrecision(precision)
    if p == TimePrecision.SECONDS:
        return np.datetime64(epoch, 's').astype(datetime.datetime)
    elif p == TimePrecision.MILLISECONDS:
        return np.datetime64(epoch, 'ms')
    elif p == TimePrecision.MICROSECONDS:
        return np.datetime64(epoch, 'us')
    elif p == TimePrecision.NANOSECONDS:
        return np.datetime64(epoch, 'ns')
    else:
        raise TypeError("Invalid timestamp precision: {}".format(precision))

def date_to_seconds(arr):
    """Converts date into seconds"""

    return arr.apply(lambda x: np.datetime64(x, "s").astype(int))

def _parse_tdf_gpu(tdf):
    """
    Parse the results of a select ipc_gpu into a GpuDataFrame

    Parameters
    ----------
    tdf : TDataFrame

    Returns
    -------
    gdf : GpuDataFrame
    """

    import pyarrow as pa
    from cudf.comm.gpuarrow import GpuArrowReader
    from cudf.core.dataframe import DataFrame
    from pyarrow._cuda import Context, IpcMemHandle
    from numba import cuda

    ipc_handle = IpcMemHandle.from_buffer(pa.py_buffer(tdf.df_handle))
    ctx = Context()
    ipc_buf = ctx.open_ipc_buffer(ipc_handle)
    ipc_buf.context.synchronize()

    schema_buffer, shm_ptr = load_buffer(tdf.sm_handle, tdf.sm_size)

    buffer = pa.BufferReader(schema_buffer)
    schema = pa.ipc.read_schema(buffer)

    # Dictionary Memo functionality used to
    # deserialize on the C++ side is not
    # exposed on the pyarrow side, so we need to
    # handle this on our own.
    dict_memo = {}

    try:
        dict_batch_reader = pa.RecordBatchStreamReader(buffer)
        updated_fields = []

        for f in schema:
            if pa.types.is_dictionary(f.type):
                msg = dict_batch_reader.read_next_batch()
                dict_memo[f.name] = msg.column(0)
                updated_fields.append(pa.field(f.name, f.type.index_type))
            else:
                updated_fields.append(pa.field(f.name, f.type))

        schema = pa.schema(updated_fields)
    except pa.ArrowInvalid:
        # This message does not have any dictionary encoded
        # columns
        pass

    dtype = np.dtype(np.byte)
    darr = cuda.devicearray.DeviceNDArray(
        shape=ipc_buf.size,
        strides=dtype.itemsize,
        dtype=dtype,
        gpu_data=ipc_buf.to_numba(),
    )

    reader = GpuArrowReader(schema, darr)
    df = DataFrame()
    df.set_tdf = MethodType(set_tdf, df)
    df.get_tdf = MethodType(get_tdf, df)

    for k, v in reader.to_dict().items():
        if k in dict_memo:
            df[k] = pa.DictionaryArray.from_arrays(v.to_arrow(), dict_memo[k])
        else:
            df[k] = v

    df.set_tdf(tdf)

    # free shared memory from Python
    # https://github.com/omnisci/pymapd/issues/46
    # https://github.com/omnisci/pymapd/issues/31
    free_sm = shmdt(ctypes.cast(shm_ptr, ctypes.c_void_p))  # noqa

    return df

mapd_to_slot = {
    'BOOL': 'int_col',
    'BOOLEAN': 'int_col',
    'SMALLINT': 'int_col',
    'INT': 'int_col',
    'INTEGER': 'int_col',
    'BIGINT': 'int_col',
    'FLOAT': 'real_col',
    'DECIMAL': 'int_col',
    'DOUBLE': 'real_col',
    'TIMESTAMP': 'int_col',
    'DATE': 'int_col',
    'TIME': 'int_col',
    'STR': 'str_col',
    'POINT': 'str_col',
    'LINESTRING': 'str_col',
    'POLYGON': 'str_col',
    'MULTIPOLYGON': 'str_col',
    'TINYINT': 'int_col',
    'GEOMETRY': 'str_col',
    'GEOGRAPHY': 'str_col',
}

mapd_to_na = {
    'BOOL': -128,
    'BOOLEAN': -128,
    'SMALLINT': -32768,
    'INT': -2147483648,
    'INTEGER': -2147483648,
    'BIGINT': -9223372036854775808,
    'FLOAT': 0,
    'DECIMAL': 0,
    'DOUBLE': 0,
    'TIMESTAMP': -9223372036854775808,
    'DATE': -9223372036854775808,
    'TIME': -9223372036854775808,
    'STR': '',
    'POINT': '',
    'LINESTRING': '',
    'POLYGON': '',
    'MULTIPOLYGON': '',
    'TINYINT': -128,
    'GEOMETRY': '',
    'GEOGRAPHY': '',
}

reason: import numpy,from numba


text:
#!/usr/bin/env python

import math
import numpy as np
import scipy.optimize as opt
import scipy.special as spc
import scipy.fftpack as fft

def Z(zeta):
    return 1j * np.sqrt(np.pi) * np.exp(-zeta**2) * (1+spc.erf(1j*zeta))
def derZ(zeta):
    return -2*(1 + zeta*Z(zeta))
def der2Z(zeta):
    return -2*(Z(zeta) + zeta*derZ(zeta))

def WeibelDispRel(omega, k, u, vth, c=1, omega_pe=1):
    zeta = omega/k/np.sqrt(2*vth**2)
    return 1.0 - omega_pe**2/(c**2*k**2)*(zeta*Z(zeta)*(1 + u**2/vth**2) + u**2/vth**2) - 2*vth**2/c**2 * zeta**2
def derWeibelDispRel(omega, k, u, vth, c=1, omega_pe=1):
    zeta = omega/k/np.sqrt(2*vth**2)
    return - omega_pe**2/(c**2*k**2) * (1 + u**2/vth**2) * (Z(zeta) + zeta*derZ(zeta))/(k*vth*np.sqrt(2)) - 4*zeta*vth/(c**2*k*np.sqrt(2))

def WeibelDispRelT(omega, k, vthx, vthy, c=1, omega_pe=1):
    zeta = omega/k/np.sqrt(2*vthx**2)
    return 1.0 - omega_pe**2/(c**2*k**2)*(vthy**2/vthx**2*(1+zeta*Z(zeta))-1) - 2*vthx**2/c**2*zeta**2

vt = math.sqrt(0.1)
kv = [0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]

for i in range(len(kv)):
    k = kv[i]
    theor = opt.newton(WeibelDispRel, 0.04j, fprime=derWeibelDispRel, args=(k, 1.5*vt, vt), maxiter=10000)
    print(k, theor.imag)

ext: py

max_stars_repo_name: ammarhakim/ammar-simjournal

max_stars_repo_licenses: MIT Unlicense

max_issues_repo_name: ammarhakim/ammar-simjournal

max_issues_repo_licenses: MIT Unlicense

max_forks_repo_name: ammarhakim/ammar-simjournal

max_forks_repo_licenses: MIT Unlicense

reason: import numpy,import scipy


text:
from strategies.botstrategy import BotStrategy
from botposition import BotPosition
import numpy as np
import warnings
with warnings.catch_warnings():
    warnings.filterwarnings('ignore',category=FutureWarning)
    from keras import models

# Strategy based on LSTM neural network which predicts the next closing price
class SLSTM(BotStrategy):
    def __init__(self):
        super(SLSTM,self).__init__()
        self.model = models.load_model('regression_model.h5')
        self.timestep = 60 # it MUST be the timestep used for lstm training
        self.wait = 0

    # Open a short or long position depending on the model prediction, waiting for 10 candelstick between two positions.
    def open_position(self,candlestick):
        if len(self.prices) > self.timestep and self.wait > 50:
            self.wait = 0
            array = np.array(self.prices[-self.timestep:])
            x_input = np.reshape(array, (1,array.shape[0],1))
            predicted = self.model.predict(x_input)
            predicted = predicted[0][0]
            if predicted > self.prices[-1]:
                return BotPosition(self.current_price,candlestick.start_time,short = False)
            else:
                return BotPosition(self.current_price,candlestick.start_time,short = True)
        self.wait += 1
        return None

    # We trust the model so when the next candlestick comes we close the position
    def try_close_position(self,position, candlestick):
        position.close(self.current_price,candlestick.start_time)

reason: import numpy


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Nov 11 21:38:43 2018

@author: zhao
"""

import argparse
import sys
import os
import numpy as np
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.nn.functional as F
from torch.autograd import Variable
from torchvision import datasets, transforms
import torch.optim as optim

from collections import OrderedDict
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.abspath(os.path.join(BASE_DIR, '../../models')))
sys.path.append(os.path.abspath(os.path.join(BASE_DIR, '../../dataloaders')))
import shapenet_part_loader

from pointcapsnet_ae import PointCapsNet
from open3d import *
import matplotlib.pyplot as plt

USE_CUDA = True
parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', type=int, default=8, help='input batch size')
parser.add_argument('--num_points', type=int, default=2048, help='number of poins')
parser.add_argument('--model', type=str, default='', help='model path')
parser.add_argument('--prim_caps_size', type=int, default=1024, help='number of prim_caps')
parser.add_argument('--prim_vec_size', type=int, default=16, help='scale of prim_vec')
parser.add_argument('--latent_caps_size', type=int, default=64, help='number of latent_caps')
parser.add_argument('--latent_vec_size', type=int, default=64, help='scale of latent_vec')

opt = parser.parse_args()
print(opt)

capsule_net = PointCapsNet(opt.prim_caps_size, opt.prim_vec_size, opt.latent_caps_size, opt.latent_vec_size, opt.num_points)
if opt.model != '':
    capsule_net.load_state_dict(torch.load(opt.model))
else:
    print ('pls set the model path')

if USE_CUDA:
    capsule_net = capsule_net.cuda()
capsule_net=capsule_net.eval()

pcd_list=[]
for i in range(opt.latent_caps_size):
    pcd_ = PointCloud()
    pcd_list.append(pcd_)

#random selected capsules to show their reconstruction with color
hight_light_caps=[np.random.randint(0, opt.latent_caps_size) for r in range(10)]
colors = plt.cm.tab20((np.arange(20)).astype(int))

test_dataset = shapenet_part_loader.PartDataset(classification=True, class_choice="Airplane", npoints=opt.num_points, split='test')
test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=opt.batch_size,
                                         shuffle=False, num_workers=4)

for batch_id, data in enumerate(test_dataloader):
    points, _,= data
    if(points.size(0)<opt.batch_size):
        break

    points = Variable(points)
    points = points.transpose(2, 1)
    if USE_CUDA:
        points = points.cuda()
    latent_capsules, reconstructions= capsule_net(points)

    for pointset_id in range(opt.batch_size):

        prc_r_all=reconstructions[pointset_id].transpose(1, 0).contiguous().data.cpu()
        prc_r_all_point=PointCloud()
        prc_r_all_point.points = Vector3dVector(prc_r_all)

        colored_re_pointcloud= PointCloud()
        jc=0
        for j in range(opt.latent_caps_size):
            current_patch=torch.zeros(int(opt.num_points/opt.latent_caps_size),3)
            for m in range(int(opt.num_points/opt.latent_caps_size)):
                current_patch[m,]=prc_r_all[opt.latent_caps_size*m+j,] # the reconstructed patch of the capsule m is not saved continuesly in the output reconstruction.
            pcd_list[j].points = Vector3dVector(current_patch)
            if (j in hight_light_caps):
                pcd_list[j].paint_uniform_color([colors[jc,0], colors[jc,1], colors[jc,2]])
                jc+=1
            else:
                pcd_list[j].paint_uniform_color([0.8,0.8,0.8])
            colored_re_pointcloud+=pcd_list[j]

        draw_geometries([colored_re_pointcloud])


text:
#### MASTER LIBRARY FUNCTIONS ################################
# USAGE: from cgx_function_library import *                 #
# from cgx_function_library import *                        #
#############################################################

##### EVENTS #####

#/----------------------
#| cgx_event_code_to_description - Gets an alarm_code as STR e.g. 'SITE_CONNECTIVITY_DEGRADED' and returns a STR of the long description for the code. Returns "unknown" on error
CGX_EVENT_CODES = {
            'DEVICEHW_POWER_LOST' : { 'type': 'alarm', 'category': 'device', 'display_name': 'Power Loss', 'severity': 'major', 'description': 'Power Supply Unit (PSU) on a device is indicating fault. Power to that PSU may be interrupted or the PSU may have failed.'},
            'DEVICEHW_INTERFACE_DOWN' : { 'type': 'alarm', 'category': 'device', 'display_name': 'Interface Down', 'severity': 'major', 'description': 'A configured admin-up interface is either not receiving a signal or has an error that is causing lack of data flow through that interface.'},
            'DEVICEHW_DISKENC_SYSTEM' : { 'type': 'alarm', 'category': 'device', 'display_name': 'Disk Encryption Upgrade failure', 'severity': 'critical', 'description': 'One of the disk partitions failed to convert into an encrypted partition during device upgrade.'},
            'DEVICEHW_DISKUTIL_PARTITIONSPACE' : { 'type': 'alarm', 'category': 'device', 'display_name': 'High Disk Capacity Utilization', 'severity': 'major', 'description': 'Disk Storage Utilization on a device has reached 85% of capacity. Non-critical functions including logging and statistics export might be impacted.'},
            'DEVICEHW_MEMUTIL_SWAPSPACE' : { 'type': 'alarm', 'category': 'device', 'display_name': 'High Memory Utilization', 'severity': 'critical', 'description': 'Memory Utilization on a device has reached maximum capacity forcing using of disk based swap space. Sub-optimal performance might be impacting device functions.'},
            'DEVICESW_DHCPSERVER_ERRORS' : { 'type': 'alarm', 'category': 'device', 'display_name': 'DHCP server failed to start', 'severity': 'critical', 'description': 'DHCP server listening on physical interfaces failed to start due to possible reasons: (a) DHCP server configuration error. (b) Lack of active element interface with static IP configuration. (c) Internal errors on element.'},
            'DEVICESW_DHCPSERVER_RESTART' : { 'type': 'alert', 'category': 'device', 'display_name': 'DHCP server restarted', 'severity': 'minor', 'description': 'DHCP server listening on physical interfaces has restarted and recovered from an error.'},
            'DEVICESW_DHCPRELAY_RESTART' : { 'type': 'alert', 'category': 'device', 'display_name': 'DHCP relay agent restarted', 'severity': 'minor', 'description': 'DHCP relay agent on device has restarted and recovered from an error.'},
            'DEVICESW_DISCONNECTED_FROM_CONTROLLER' : { 'type': 'alarm', 'category': 'device', 'display_name': 'Device disconnected from Controller', 'severity': 'major', 'description': 'Device has remained disconnected from the Controller for a prolonged duration.'},
            'DEVICESW_INITIATED_CONNECTION_ON_EXCLUDED_PATH' : { 'type': 'alarm', 'category': 'device', 'display_name': 'Device Initiated Connection on excluded path', 'severity': 'major', 'description': 'Device Initiated Connection on excluded interface.'},
            'DEVICESW_NTP_NO_SYNC' : { 'type': 'alarm', 'category': 'device', 'display_name': 'NTP synchronization failed', 'severity': 'major', 'description': 'Unable to sync up with all configured NTP servers for more than 24 hours.'},
            'DEVICESW_GENERAL_PROCESSRESTART' : { 'type': 'alert', 'category': 'device', 'display_name': 'Process Restart', 'severity': 'minor', 'description': 'A software process on the device has restarted either due to an error or as a self-recovery method. Process restart as a self-recovery does not impact long-term functions on the device but can cause short term sub-optimal functions and errors.'},
            'DEVICESW_GENERAL_PROCESSSTOP' : { 'type': 'alarm', 'category': 'device', 'display_name': 'Process Stopped', 'severity': 'major', 'description': 'A software process on the device has stopped due to an error and inability to recover with a self-restart. Functionality is likely impacted.'},
            'DEVICESW_MONITOR_DISABLED' : { 'type': 'alarm', 'category': 'device', 'display_name': 'System Monitoring Disabled', 'severity': 'major', 'description': 'A software process that monitors the health of device and its Hardware/Software components has been disabled. Though not currently impacting, operation of the device and the software on the device while this fault exists will impact recovery and handling of other Hardware/Software faults that need recovery action.'},
            'DEVICESW_SNMP_AGENT_FAILED_TO_START' : { 'type': 'alarm', 'category': 'device', 'display_name': 'SNMP Agent Failed to Start', 'severity': 'major', 'description': 'SNMP agent failed to start due to these possible reasons: (a) Invalid configuration, or (b) Decryption failure.'},
            'DEVICESW_SNMP_AGENT_RESTART' : { 'type': 'alert', 'category': 'device', 'display_name': 'SNMP agent restarted', 'severity': 'minor', 'description': 'SNMP agent on device has restarted and recovered from an error.'},
            'DEVICESW_SYSTEM_BOOT' : { 'type': 'alert', 'category': 'device', 'display_name': 'Device Reboot', 'severity': 'critical', 'description': 'Device rebooted either due to recovery on a fault condition or as part of normal operational procedures including user initiated reboots and software upgrades. Reboots due to fault conditions can cause sub-optimal or significantly reduced functionality on the device.'},
            'DEVICEIF_ADDRESS_DUPLICATE' : { 'type': 'alert', 'category': 'device', 'display_name': 'Duplicate IP Address', 'severity': 'major', 'description': 'Another device in the local network is using an IP address assigned to this device.'},
            'DEVICEHW_INTERFACE_ERRORS' : { 'type': 'alert', 'category': 'device', 'display_name': 'High rate of errors on the interface', 'severity': 'major', 'description': 'Number of transmission and/or reception errors seen on an interface over the last one hour period has exceeded the threshold (0.5% of received or transmitted packet counts in the same one hour period).'},
            'DEVICEHW_INTERFACE_HALFDUPLEX' : { 'type': 'alarm', 'category': 'device', 'display_name': 'Interface running in half-duplex mode', 'severity': 'major', 'description': 'A interface has negotiated half duplex, although it is allowed to run in full duplex, which is preferred.'},
            'FLAP_RATE_EXCEEDED' : { 'type': 'alarm', 'category': 'policy', 'display_name': 'Flap Rate Exceeded', 'severity': 'major', 'description': 'An alarm has repeatedly raised and cleared triggering a flapping condition per the event policy rule.'},
            'NETWORK_ANYNETLINK_DEGRADED' : { 'type': 'alarm', 'category': 'network', 'display_name': 'Secure Fabric Link Degraded', 'displayCode': 'NETWORK_SECUREFABRICLINK_DEGRADED', 'severity': 'minor', 'description': 'Secure Fabric Link is degraded with atleast 1 VPNlink UP from the active spoke and 1 or more VPNlinks DOWN from the active spoke.'},
            'NETWORK_ANYNETLINK_DOWN' : { 'type': 'alarm', 'category': 'network', 'display_name': 'Secure Fabric Link Down', 'displayCode': 'NETWORK_SECUREFABRICLINK_DOWN', 'severity': 'major', 'description': 'Secure Fabric Link is down with all VPNLinks DOWN from the active spoke.'},
            'NETWORK_DIRECTPRIVATE_DOWN' : { 'type': 'alarm', 'category': 'network', 'display_name': 'Private Wan Reachability Down', 'severity': 'major', 'description': 'For remote office (branch) sites, all data center sites with ion 7000x deployed have been declared unreachable on Private WAN. If there are no alternate paths in the application policy, the fault is traffic impacting and should be attended to immediately.'},
            'NETWORK_DIRECTINTERNET_DOWN' : { 'type': 'alarm', 'category': 'network', 'display_name': 'Direct Internet Reachability Down', 'severity': 'major', 'description': 'For remote office (branch) sites, reachability on an Internet circuit has been declared to be down. If there are no alternate paths in the application policy, the fault is traffic impacting and should be attended to immediately.'},
            'NETWORK_PRIVATEWAN_UNREACHABLE' : { 'type': 'alarm', 'category': 'network', 'display_name': 'Private Wan Unreachable', 'severity': 'major', 'description': 'For data center sites, one or more remote offices have been declared unreachable over Private WAN based on routing updates received from the network. If this fault occurred due to WAN edge peering failure PEERING_EDGE_DOWN fault(s) will also be raised.'},
            'NETWORK_PRIVATEWAN_DEGRADED' : { 'type': 'alarm', 'category': 'network', 'display_name': 'Private Wan Degraded', 'severity': 'major', 'description': 'For data center sites, a subset of IP prefixes from one or more remote sites have been determined to be unreachable over Private WAN based on routing updates received from the network.'},
            'OPERATOR_SIGNUP_TOKEN_DISABLED' : { 'type': 'alert', 'category': 'aaa', 'display_name': 'User Signup Disabled', 'severity': 'minor', 'description': 'A new user that was issued a signup token to self-complete the signup steps failed that process multiple times by using wrong combination of the signup token and unique ID supplied by his/her administrator. The same fault can happen if an existing user forgot his/her password, was required to finish the password reset steps but failed that process multiple times.'},
            'PEERING_EDGE_DOWN' : { 'type': 'alarm', 'category': 'network', 'display_name': 'WAN Edge Peer Down', 'severity': 'critical', 'description': 'Routing peer session with a configured WAN edge device is down. If alternate paths are available traffic is not affected; else the fault is critical.'},
            'PEERING_CORE_DOWN' : { 'type': 'alarm', 'category': 'network', 'display_name': 'Core Peer Down', 'severity': 'critical', 'description': 'Routing peer session with a configured Core device is down. If alternate paths are available traffic is not affected; else the fault is critical.'},
            'PEERING_BGP_DOWN' : { 'type': 'alarm', 'category': 'network', 'display_name': 'BGP Peer Down', 'severity': 'critical', 'description': 'Routing peer session is down. If alternate paths are available traffic is not affected; else the fault is critical.'},
            'APPLICATION_CUSTOM_RULE_CONFLICT' : { 'type': 'alarm', 'category': 'application', 'display_name': 'Application Custom Rule Conflict', 'severity': 'minor', 'description': 'An application rule conflict has been detected.'},
            'APPLICATION_PROBE_DISABLED' : { 'type': 'alarm', 'category': 'device', 'display_name': 'Application Probe Disabled', 'severity': 'major', 'description': 'Application probes are disabled either due to incomplete configuration or invalid state. Element will no longer issue application probe to detect application reachability unless the issue is resolved. Consequently, if application probes are disabled then application will no longer switch to alternative paths in case it fails on its current path.'},
            'SITE_CIRCUIT_ABSENT_FOR_POLICY' : { 'type': 'alarm', 'category': 'policy', 'display_name': 'Site missing circuit(s) for policy', 'severity': 'major', 'description': 'Site is missing all circuit definitions specified in the Policy Set assigned to the site. Applications at the site will be affected since there are no circuits to forward the traffic.'},
            'SITE_NETWORK_SERVICE_ABSENT_FOR_POLICY' : { 'type': 'alarm', 'category': 'network', 'display_name': 'Policy DC Group Missing Service Endpoint', 'severity': 'major', 'description': 'One or more DC groups used in the policy has not been assigned a valid service endpoint for the domain bound to the identified site.'},
            'DEVICESW_IMAGE_UNSUPPORTED' : { 'type': 'alarm', 'category': 'device', 'display_name': 'Unsupported Software Image', 'severity': 'critical', 'description': 'Device\'s software image is not recognized by the Controller. The software version may not be allowed in the network or may no longer exist.'},
            'DEVICESW_FPS_LIMIT_EXCEEDED' : { 'type': 'alarm', 'category': 'device', 'display_name': 'System Rate Limit Exceeded', 'severity': 'major', 'description': 'A consistently high flow rate above system limits has been detected.'},
            'DEVICESW_CONCURRENT_FLOWLIMIT_EXCEEDED' : { 'type': 'alarm', 'category': 'device', 'display_name': 'System Tracking Limit Exceeded', 'severity': 'critical', 'description': 'An abnormal and significantly high sessions/flow trackers have been detected.'},
            'DEVICESW_LICENSE_VERIFICATION_FAILED' : { 'type': 'alarm', 'category': 'device', 'display_name': 'Virtual ION license verification failed', 'severity': 'major', 'description': 'The License is no longer valid. The maximum ION deployment limit reached.'},
            'DEVICESW_TOKEN_VERIFICATION_FAILED' : { 'type': 'alert', 'category': 'device', 'display_name': 'Virtual ION token validation failed', 'severity': 'critical', 'description': 'The token is no longer valid. It\'s either utilized, expired or revoked.'},
            'DEVICESW_CRITICAL_PROCESSSTOP' : { 'type': 'alarm', 'category': 'device', 'display_name': 'Critical Process Stopped', 'severity': 'critical', 'description': 'A critical software process on the device has stopped due to an error and inability to recover with a self-restart. Data forwarding functionality is likely impacted.'},
            'DEVICESW_CRITICAL_PROCESSRESTART' : { 'type': 'alert', 'category': 'device', 'display_name': 'Critical Process Restart', 'severity': 'critical', 'description': 'A critical software process on the device has restarted either due to an error or as a self-recovery method. Process restart as a self-recovery does not impact long-term functions on the device but can cause short term sub-optimal data plane functions and errors.'},
            'DEVICESW_CONNTRACK_FLOWLIMIT_EXCEEDED' : { 'type': 'alarm', 'category': 'device', 'display_name': 'Conntrack table flow count exceeded threshold.', 'severity': 'critical', 'description': 'Number of flows in the connection tracking table used for features such as NAT and device management policy has exceeded 90% threshold.'},
            'NETWORK_POLICY_RULE_CONFLICT' : { 'type': 'alarm', 'category': 'policy', 'display_name': 'Network policy rule conflict', 'severity': 'minor', 'description': 'Two or more policy rules in a network policy set conflict, potentially resulting in incorrect policy being applied to some flows.'},
            'NETWORK_POLICY_RULE_DROPPED' : { 'type': 'alarm', 'category': 'policy', 'display_name': 'Network policy rule dropped', 'severity': 'major', 'description': 'Network policy configuration contains rules with too many permutations such that the required resources exceed the operational limits. Some rules have been dropped from the policy so as to not exceed the limit and therefore the desired policy actions may not be performed in some cases.'},
            'NAT_POLICY_STATIC_NATPOOL_OVERRUN' : { 'type': 'alarm', 'category': 'device', 'severity': 'minor', 'display_name': 'Static NAT pool range is overrun by selector prefix', 'description': 'Configured Nat pool range cannot map 1:1 with matching traffic selector prefix.'},
            'PRIORITY_POLICY_RULE_CONFLICT' : { 'type': 'alarm', 'category': 'policy', 'display_name': 'Priority policy rule conflict', 'severity': 'minor', 'description': 'Two or more policy rules in a priority policy set conflict, potentially resulting in incorrect policy being applied to some flows.'},
            'PRIORITY_POLICY_RULE_DROPPED' : { 'type': 'alarm', 'category': 'policy', 'display_name': 'Priority policy rule dropped', 'severity': 'major', 'description': 'Priority policy configuration contains rules with too many permutations such that the required resources exceed the operational limits. Some rules have been dropped from the policy so as to not exceed the limit and therefore the desired policy actions may not be performed in some cases.'},
            'NAT_POLICY_LEGACY_ALG_CONFIG_OVERRIDE' : { 'type': 'alert', 'category': 'policy', 'display_name': 'NAT policy ALG action overridden by legacy configuration', 'severity': 'major', 'description': 'ALG action configured in the NAT policy has been overridden by legacy configuration present on the device.'},
            'SECURITY_POLICY_RULE_INCOMPLETE' : { 'type': 'alarm', 'category': 'policy', 'display_name': 'Security rule configuration incomplete', 'severity': 'critical', 'description': 'The security policy rule configuration is incomplete. In this case the security policy rule is skipped.'},
            'SITE_CONNECTIVITY_DOWN' : { 'type': 'alarm', 'category': 'network', 'display_name': 'Site Connectivity down', 'severity': 'critical', 'description': 'All site WAN connectivity is down.'},
            'SITE_CONNECTIVITY_DEGRADED' : { 'type': 'alarm', 'category': 'network', 'display_name': 'Site Connectivity degraded', 'severity': 'major', 'description': 'Multiple issues are present impacting site WAN connectivity.'},
            'SPOKEHA_STATE_UPDATE' : { 'type': 'alert', 'category': 'spokeha', 'display_name': 'Device state in the spoke cluster changed', 'severity': 'major', 'description': 'Device changed its state from active to backup or backup to active. If the device changed its state to backup, and there is no other device eligible to become active, then network connectivity at the site will be affected.'},
            'SPOKEHA_MULTIPLE_ACTIVE_DEVICES' : { 'type': 'alarm', 'category': 'spokeha', 'display_name': 'Split Brain', 'severity': 'critical', 'description': 'A critical alarm will be raised on the spoke HA cluster resource by the controller when both elements declare themselves to be "active" (split brain).'},
            'SPOKEHA_CLUSTER_DEGRADED' : { 'type': 'alarm', 'category': 'spokeha', 'display_name': 'Cluster Degraded', 'severity': 'major', 'description': 'One of the element in the SpokeCluster has effective priority 0.'},
            'SPOKEHA_CLUSTER_DOWN' : { 'type': 'alarm', 'category': 'spokeha', 'display_name': 'Cluster Down', 'severity': 'critical', 'description': 'All elements in the SpokeCluster have effective priority 0.'},
            'CLAIMCERT_AUTO_RENEWAL_DISABLED' : { 'type': 'alert', 'category': 'device', 'display_name': 'Auto Renewal of Claim Certificate is Disabled', 'severity': 'major', 'description': 'The claim certificate shall expire in a short period of time. This may happen when either previous attempts to renew the certificate have failed or the certificate expiry has been set to a short timeframe.'},
            'CLAIMCERT_EXPIRY_WARNING' : { 'type': 'alert', 'category': 'device', 'display_name': 'Claim Certificate Expiration Date Approaching', 'severity': 'major', 'description': 'The claim certificate shall expire in a short period of time. This may happen when either previous attempts to renew the certificate have failed or the certificate expiry has been set to a short timeframe.'},
            'CLAIMCERT_RENEWAL_RETRY_LIMIT_EXCEEDED' : { 'type': 'alarm', 'category': 'device', 'display_name': 'Claim Certificate Renewal Attempts Exceeded', 'severity': 'critical', 'description': 'There were errors observed during the process of Claim Certificate renewal. Repeated attempts to renew the Claim Certificate exceeded three consecutive retries. Auto renewal is therefore disabled and a renewal must be triggered from the Controller. However, this event indicates a problem that is external to the process and it must be attended to immediately.'},
            'CLAIMCERT_RENEWAL_FAILED' : { 'type': 'alert', 'category': 'device', 'display_name': 'Renewal of Claim Certificate Failed', 'severity': 'major', 'description': 'The process of renewing the Claim Certificate encountered problems. These may be related to external events such as failures reported by CA or incorrect/invalid certificate being issued. Other reasons can be internal failures such as problems arising from generating a CSR request or receiving CSR details from controller.'},
            'CLAIMCERT_RENEWALS_TOO_FREQUENT' : { 'type': 'alert', 'category': 'device', 'display_name': 'Claim Certificate Renewed Too Frequently', 'severity': 'major', 'description': 'A condition is reached where a renewed Claim Certificate is determined to fall within a renewal window and subsequent renewals are attempted once again. This condition may occur when the Certificate issued by the CA has a shorter expiry time than the renewal window configured on the Controller.'},
            'DEVICESW_IPFIX_COLLECTORS_DOWN' : { 'type': 'alarm', 'category': 'device', 'display_name': 'IPFIX collectors are down', 'severity': 'major', 'description': 'The software process responsible to export IPFIX records has observed that there are no active connections to the IPFIX collectors. The process will continue to monitor the connection status and resume exporting of the IPFIX records once the connections are re-established.'},
            'DEVICEHW_ION9000X722FW_OUTOFDATE' : { 'type': 'alarm', 'category': 'device', 'display_name': 'ION 9000 Ports 9-12 Firmware Update Required', 'severity': 'major', 'description': 'A very important firmware update is required for stable operation of ports 9 through 12 on this device.'}
            }
def cgx_event_code_to_description(alarm_code):
    global CGX_EVENT_CODES
    try:
        return CGX_EVENT_CODES[alarm_code]['description']
    except:
        return "Unknown" ### event code not found
#\----------------------

#/----------------------
#| cgx_get_last_alarms - Gets last Alarms with or without filters. By default, pulls the last 100 active alarms. Set resolved to True to include resolved to view the past 24 hours of alarms.
def cgx_get_last_alarms(sdk, count=100, acknowledged=False, resolved=False, hours_interval=24, suppressed=False):
    true = True
    false = False
    acknowledged = True if str(acknowledged).lower() == "true" else False
    suppressed = True if str(suppressed).lower() == "true" else False
    resolved = True if str(resolved).lower() == "true" else False
    event_request_query = {"limit":{"count":count,"sort_on":"time","sort_order":"descending"},"view":{"summary":false},"severity":[],"priority":[],"acknowledged":acknowledged,"suppressed":false,"query":{"site":[],"category":[],"code":[],"correlation_id":[],"type":["alarm"]}}
    if resolved: ## to view resolved, you must indicate a start time and end time.
        (start_time, end_time) = cgx_generate_timestamps(hours_interval=hours_interval)
        event_request_query = {"limit":{"count":count,"sort_on":"time","sort_order":"descending"},"view":{"summary":false},"end_time":end_time,"start_time":start_time,"severity":[],"priority":[],"query":{"site":[],"category":[],"code":[],"correlation_id":[],"type":["alarm"]}}
    try:
        result = sdk.post.events_query(event_request_query)
        return result.cgx_content.get("items", None)
    except:
        return None
#\----------------------

##### SITES #####

#/----------------------
#| cgx_site_name_to_id - Attempts to find the best matching site id given a site name
def cgx_site_name_to_id(sdk, search_site):
    from fuzzywuzzy import fuzz
    search_ratio = 0
    site_dict = {}
    resp = sdk.get.sites()
    if resp.cgx_status:
        site_list = resp.cgx_content.get("items", None)    #site_list contains an list of all returned sites
        for site in site_list:                            #Loop through each site in the site_list
            check_ratio = fuzz.ratio(search_site.lower(),site['name'].lower())
            if (check_ratio > search_ratio ): ###Find the "best" matching site name
                search_ratio = check_ratio
                site_dict = site
    else:
        print("ERROR: API Call failure when enumerating SITES in tenant! Exiting!")
    try:
        return site_dict['id']
    except:
        return None
#\----------------------

#/----------------------
#| cgx_site_id_to_name - Attempts to find the site name given a site id
def cgx_site_id_to_name(sdk, search_id):
    resp = sdk.get.sites(site_id=search_id)
    if resp.cgx_status:
        return resp.cgx_content['name']
    else:
        print("ERROR: API Call failure when enumerating SITE ID in tenant! Exiting!")
        return None
#\----------------------

#/----------------------
#| cgx_site_name_to_object - Attempts to find the best matching site dict object given a site name
def cgx_site_name_to_dict(sdk, search_site):
    from fuzzywuzzy import fuzz
    search_ratio = 0
    site_dict = {}
    resp = sdk.get.sites()
    if resp.cgx_status:
        site_list = resp.cgx_content.get("items", None)    #site_list contains an list of all returned sites
        for site in site_list:                            #Loop through each site in the site_list
            check_ratio = fuzz.ratio(search_site.lower(),site['name'].lower())
            if (check_ratio > search_ratio ): ###Find the "best" matching site name
                search_ratio = check_ratio
                site_dict = site
    else:
        print("ERROR: API Call failure when enumerating SITES in tenant! Exiting!")
    try:
        return site_dict
    except:
        return None
#\----------------------

#/----------------------
#| cgx_destroy_all_sites - Attempts to iterate through all sites and
def cgx_destroy_all_sites(sdk, unclaim_device=True):
    from cloudgenix_config import do
    from cloudgenix_config import pull
    result = sdk.get.sites().cgx_content.get("items")
    site_list = []
    for i in result:
        site_list.append(i['id'])
    for site_id in site_list:
        try:
            elements_list = cgx_get_elements_in_site(sdk, site_id)
            result = pull.pull_config_sites(site_id, None, return_result=True, passed_sdk=sdk) # Pull the site config before destroying
            do.do_site(result, destroy=True, passed_sdk=sdk, ) # Destroy the site using the previously pulled config
            if unclaim_device:
                for element in elements_list:
                    cgx_unclaim_element(sdk,element) ## Unclaim each ION at the site
        except:
            print("Problem with destroying site",site_id)
#\----------------------

#/----------------------
#| cgx_destroy_site - Attempts to destroy a specific site by ID by destroying the site
def cgx_destroy_site(sdk, site_id):
    from cloudgenix_config import do
    from cloudgenix_config import pull
    try:
        result = pull.pull_config_sites(site_id, None, return_result=True, passed_sdk=sdk) # Pull the site config before destroying
        do.do_site(result, destroy=True, passed_sdk=sdk, ) # Destroy the site using the previously pulled config
    except:
        print("Problem with destroying site",site_id)
#\----------------------

##### ELEMENTS #####

#/----------------------
#| cgx_get_element_ids_in_site - returns a list of element ID's bound to that site
def cgx_get_element_ids_in_site(sdk, site_id):
    response = []
    elements_list = sdk.get.elements().cgx_content.get("items")
    for element in elements_list:
        if element['site_id'] == site_id:
            response.append(element['id'])
    return response
#\----------------------

#/----------------------
#| cgx_get_element_dict_in_site - returns a list of element DICT's bound to that site
def cgx_get_element_dict_in_site(sdk, site_id):
    response = []
    elements_list = sdk.get.elements().cgx_content.get("items")
    for element in elements_list:
        if element['site_id'] == site_id:
            response.append(element)
    return response
#\----------------------

#/----------------------
#| cgx_element_name_to_id - Attempts to find the best matching element id given an element name
def cgx_element_name_to_id(sdk, search_ion):
    from fuzzywuzzy import fuzz
    search_ratio = 0
    element_dict = {}
    resp = sdk.get.elements()
    if resp.cgx_status:
        element_list = resp.cgx_content.get("items", None)    #element_list contains an list of all returned elements
        for element in element_list:                            #Loop through each site in the element_list
            check_ratio = fuzz.ratio(search_ion.lower(),element['name'].lower())
            if (check_ratio > search_ratio ): ###Find the "best" matching site name
                search_ratio = check_ratio
                element_dict = element
    else:
        print("ERROR: API Call failure when enumerating SITES in tenant! Exiting!")
    try:
        return element_dict['id']
    except:
        return None
#\----------------------

#/----------------------
#| cgx_element_name_to_dict - Attempts to find the best matching element id given an element name
def cgx_element_name_to_dict(sdk, search_ion):
    from fuzzywuzzy import fuzz
    search_ratio = 0
    element_dict = {}
    resp = sdk.get.elements()
    if resp.cgx_status:
        element_list = resp.cgx_content.get("items", None)    #element_list contains an list of all returned elements
        for element in element_list:                            #Loop through each site in the element_list
            check_ratio = fuzz.ratio(search_ion.lower(),element['name'].lower())
            if (check_ratio > search_ratio ): ###Find the "best" matching site name
                search_ratio = check_ratio
                element_dict = element
    else:
        print("ERROR: API Call failure when enumerating SITES in tenant! Exiting!")
    try:
        return element_dict
    except:
        return None
#\----------------------

#/----------------------
#| cgx_unclaim_element - unclaims a claimed element by id
def cgx_unclaim_element(sdk, element_id):
    declaim_post = {"action":"declaim","parameters":None}
    sdk.post.operations_e(element_id, declaim_post)
#\----------------------

##### APP ID #####

#/----------------------
#| cgx_app_id_to_name_cached - Caching function which takes in an APP_ID and returns the NAME
app_id_name_cache = {}
app_def_cache = None
def cgx_app_id_to_name_cached(sdk, search_app_id):
    global app_id_name_cache, app_def_cache
    search_app_id = str(search_app_id).strip().lower()
    if search_app_id in app_id_name_cache.keys():
        return app_id_name_cache[search_app_id]['display_name']
    if not app_def_cache:
        app_def_cache = sdk.get.appdefs().cgx_content.get("items",None)
    for app in app_def_cache:
        if search_app_id == app['id']:
            app_id_name_cache[search_app_id] = app
            return app['display_name']
    app_id_name_cache[search_app_id] = None
    return None
#\----------------------

#/----------------------
#| cgx_app_id_to_name - Non-caching function which takes in an APP_ID and returns the NAME
def cgx_app_id_to_dict(sdk, search_app_id):
    try:
        result = sdk.get.appdefs(appdef_id=str(search_app_id)).cgx_content
        return result['display_name']
    except:
        return "Unknown"
#\----------------------

#/----------------------
#| cgx_app_id_to_name_cached - Caching function which takes in an APP Name and returns the ID
app_id_name_cache = {}
app_def_cache = None
def cgx_app_name_to_id_cached(sdk, search_app_name):
    global app_id_name_cache, app_def_cache
    from fuzzywuzzy import fuzz
    search_ratio = 0
    app_dict = {}
    try:
        if not app_def_cache:
            app_def_cache = sdk.get.appdefs().cgx_content.get("items",None)
        for app in app_def_cache:
            check_ratio = fuzz.ratio(search_app_name.lower(),app['display_name'].lower())
            if (check_ratio > search_ratio ): ###Find the "best" matching site name
                search_ratio = check_ratio
                app_dict = app
        return app_dict['id']
    except:
        return "N/A"
#\----------------------

#/----------------------
#| cgx_app_id_to_dict_cached - Caching function which takes in an APP Name and returns the app dict
app_id_name_cache = {}
app_def_cache = None
def cgx_app_id_to_dict_cached(sdk, search_app_name):
    global app_id_name_cache, app_def_cache
    from fuzzywuzzy import fuzz
    search_ratio = 0
    app_dict = {}
    try:
        if not app_def_cache:
            app_def_cache = sdk.get.appdefs().cgx_content.get("items",None)
        for app in app_def_cache:
            check_ratio = fuzz.ratio(search_app_name.lower(),app['display_name'].lower())
            if (check_ratio > search_ratio ): ###Find the "best" matching site name
                search_ratio = check_ratio
                app_dict = app
        return app_dict
    except:
        return "unknown"
#\----------------------

##### ANALYTICS #####

#/----------------------
#| cgx_get_last_flows - Gets last flows based on a Site_ID, start time, end time, and optional list of app-id's to filter by
def cgx_get_last_flows(sdk, site_id, start_time, end_time, filter_list_of_apps=[]):
    true = True
    false = False
    (start_time, end_time) = cgx_generate_timestamps(hours_interval=1)
    flow_request_query = {"start_time":start_time,"end_time":end_time,"filter":{"app":filter_list_of_apps,"site":[site_id]},"debug_level":"all"}
    try:
        result = sdk.post.flows_monitor(flow_request_query)
        if result.cgx_status:
            return result.cgx_content.get("flows").get("items")
        else:
            print("Error getting flows")
            return result
    except:
        return None
#\----------------------

#/----------------------
#| cgx_get_bw_consumption - Gets bandwidth consumption average for a give time period. If no site_id is given, Aggregate BW consumption for the tenant is provided.
def cgx_get_bw_consumption(sdk, start_time, end_time, site_id=None):
    true = True
    false = False
    post_request = {"start_time":start_time, "end_time":end_time, "interval":"5min","metrics":[{"name":"BandwidthUsage","statistics":["average"],"unit":"Mbps"}],"view":{},"filter":{"site":[]}}
    if site_id: post_request['filter']['site'].append(str(site_id))
    result = sdk.post.monitor_metrics(post_request)
    metrics = result.cgx_content
    series = metrics.get("metrics",[{}])[0].get("series",[])[0]
    return(cgx_average_series(series))
#\----------------------

##### Topology #####

#/----------------------
#| cgx_mesh_two_sites - Takes two sites and forms VPN tunnels between the two across all capable WAN connections
def cgx_mesh_two_sites(sdk, site1_id, site2_id ):
    site1_wans = sdk.get.waninterfaces(site1_id).cgx_content.get("items", None)
    site2_wans = sdk.get.waninterfaces(site2_id).cgx_content.get("items", None)
    for left_wan in site1_wans:
        for right_wan in site2_wans:
            if left_wan['type'] == right_wan['type']: ##WANS must be of same type. I.E. PrivateWAN must connect to Private WAN's
                cgx_add_anynet_link(sdk, site1_id, left_wan['id'], site2_id, right_wan['id'])
#\----------------------

#/----------------------
#| cgx_add_anynet_link - Adds a VPN tunnel between two sites on a specific WAN circuit
def cgx_add_anynet_link(sdk, site1, wan1, site2, wan2 ):
    ### This function adds a SecureFabric/Anynet link between two sites
    post_data = {'name': None, 'description': None, 'tags': None,
                        'ep1_site_id': str(site1),
                            'ep1_wan_if_id': str(wan1),
                        'ep2_site_id': str(site2),
                            'ep2_wan_if_id': wan2,
                    'admin_up': True, 'forced': True, 'type': None}
    result = sdk.post.tenant_anynetlinks(post_data)
    if result.cgx_status == False:
        print("FAILURE:",result.cgx_errors)
    else:
        print("SUCCESS",site1,result.cgx_status)
#\----------------------

##### MISC #####

#/----------------------
#| cgx_generate_timestamps - Generates CGX start and end Timestamps used in events, alarms, reporting for specified interval in hours (default 1 hour). The time may be offset backwards (I.E Set offset to 8 to get a start/end time 8 hours ago)
def cgx_generate_timestamps(hours_interval=1, offset_hours=0):
    from datetime import timedelta
    from datetime import datetime
    now = (datetime.utcnow() - timedelta(hours=(offset_hours)))
    end_time = now.isoformat()
    start_time = (now - timedelta(hours=(hours_interval))).isoformat()
    return (str(start_time)+"Z", str(end_time)+"Z") ## return (start_time, end_time) I.E. call with "(start_time, end_time) = cgx_generate_timestamps(hours_interval=8) for the last 8 hours"
#\----------------------

#/----------------------
#| cgx_generate_timestamps_days - Generates CGX start and end Timestamps used in events, alarms, reporting for specified interval in days (default 1 day or 24 hours)
def cgx_generate_timestamps_days(days_interval=1, offset_days=0):
    from datetime import timedelta
    from datetime import datetime
    now = (datetime.utcnow() - timedelta(days=(offset_days)))
    end_time = now.isoformat()
    start_time = (now - timedelta(days=(days_interval))).isoformat()
    return (str(start_time)+"Z", str(end_time)+"Z") ## return (start_time, end_time) I.E. call with "(start_time, end_time) = cgx_generate_timestamps(days_interval=1) for the last 24 hours"
#\----------------------

#/----------------------
#| cgx_mililseconds_to_timestamp - Takes a millisecond value and converts it to a datetime stamp. Friendly_Output makes it nicer when printing output
def cgx_millseconds_to_timestamp(epoc_milliseconds, friendly_output=False):
    from datetime import datetime
    try:
        dt_object = datetime.fromtimestamp(int(epoc_milliseconds)/1000) #divide by 1000 to convert to epoch
        if friendly_output: return str(dt_object.isoformat()).split(".")[0]
        return(str(dt_object.isoformat())+"Z")
    except:
        print("Error converting time",epoc_milliseconds)
        return(epoc_milliseconds)
#\----------------------

#/----------------------
#| cgx_convert_protocol_to_name - Translates a protocol number to a friendly name. I.E. Proto 6 would return TCP
def cgx_convert_protocol_to_name(protocol_number):
    value = str(protocol_number).lower().strip()
    protocols_dict = {
        "1"   : "ICMP",
        "2"   : "IGMP",
        "6"   : "TCP",
        "4"   : "IPv4 Encap",
        "17"  : "UDP",
        "41"  : "IPv6",
        "47"  : "GRE",
        "51"  : "AH",
        "115" : "L2TP",
        "136" : "UDP-Lite",
        "143" : "Ethernet",
    }
    if value in protocols_dict.keys(): return protocols_dict[value]
    return value
#\----------------------

#/----------------------
#| safe_resolve - SAFELY finds nested keys in a DICT object with error correction in a function.
#                Usage - Instead of doing "street_address = sites['items'][0]['address']['street']"" which could have key errors and break, do this:
#                        street_address = safe_resolve(sites,'items',0,'address','street'). Any key error will be cleanly handled and "None" will be returned on error
def safe_resolve(nested_object, *keys):
   for key in keys:
       try:
           nested_object = nested_object[key]
       except:
           return None
   return nested_object
#\----------------------

#/----------------------
#| safe_sum - SAFELY sums values passed regardless of original type. E.G. safe_sum(1,2,3,"a","b",None,False,True,True) returns 6 ignoring non numerical values
def safe_sum(*values):
    total = 0
    for key in values:
        try:
            if type(key) is not bool:
                total += float(key)
        except:
            total += 0
    try:
        if int(total) == total:
            return int(total)
        else:
            return total
    except:
        return total
#\----------------------

#/----------------------
#| cgx_average_series - takes a metrics series structure (input['metrics']['series']['data']['datapoints'][**list**]['value']) and averages it to the decimal places (default:2)
#|                  A typical call would look like:
#|                      metrics = cgx_get_bw_consumption(sdk,start,end) ## Calls sdk.post.monitor_metrics and returns Metrics with Series in them
#|                      for series in metrics.get("metrics",[{}])[0].get("series",[]):
#|                          average = cgx_average_series(series)
def cgx_average_series(metrics_series_structure, decimal_places=2):
    count = 0
    sum = 0
    for datapoints in metrics_series_structure.get("data",[{}])[0].get("datapoints",[{}]):
        if (datapoints.get("value",None) is not None):
            count += 1
            sum += datapoints.get("value",0)
    if count != 0:
        if decimal_places == 0:
            return round((sum/count))
        return round((sum/count),decimal_places)
    return 0
#\----------------------

#/----------------------
#| cgx_sum_series - takes a series structure (input['series']['data']['datapoints'][**list**]['value']) and sums it to the decimal places (default:2)
#|                  A typical call would look like:
#|                      metrics = cgx_get_bw_consumption(sdk,start,end) ## Calls sdk.post.monitor_metrics and returns Metrics with Series in them
#|                      for series in metrics.get("metrics",[{}])[0].get("series",[]):
#|                          sum = cgx_sum_series(series)
def cgx_sum_series(metrics_series_structure, decimal_places=2):
    count = 0
    sum = 0
    for datapoints in metrics_series_structure.get("data",[{}])[0].get("datapoints",[{}]):
        if (datapoints.get("value",None) is not None):
            count += 1
            sum += datapoints.get("value",0)
    if count != 0:
        if decimal_places == 0:
            return round((sum))
        return round((sum),decimal_places)
    return 0
#\----------------------

#/----------------------
#| validate_2d_array - Validates an array is 2-dimensional for use in CSV Export
def validate_2d_array(test_list):
    import numpy as np
    np_list = np.array(test_list)
    if len(np_list.shape) == 2:
        return True
    return False
#\----------------------

#/----------------------
#| write_2d_list_to_csv - Writes a 2-Dimensional list to a CSV file
def write_2d_list_to_csv(csv_file, list_2d, write_mode="w"):
    import csv
    try:
        file = open(csv_file, write_mode)
        with file:
            write = csv.writer(file)
            write.writerows(list_2d)
            return True
        return False
    except:
        return False

#\----------------------

#######################################Examples#########################################################################
def print_alarms(sdk,alarms):
    import re
    counter=0
    for alarm in alarms:
        site_name = cgx_site_id_to_name(sdk,alarm['site_id'])
        event_desc = cgx_event_code_to_description(alarm['code'] )
        counter+=1
        #print(counter,site_name,alarm['severity'],,alarm['code'],)
        print("")
        print("Alarm",counter," - Site: ",site_name)
        print("  - Timestamp :",re.sub("\..*Z","",str(alarm['time']))   )
        print("  - Site      :",site_name   )
        print("  - Severity  :",alarm['severity']   )
        print("  - Code      :",alarm['code']   )
        print("  - Code Descr:",event_desc   )

def print_flows(sdk,flows):
    import re
    counter=0
    for flow in flows:
        app_id = cgx_app_id_to_name_cached(sdk,flow['app_id'])
        counter+=1
        direction = "LAN to WAN" if safe_resolve(flow,'lan_to_wan') else "WAN to LAN"
        print("")
        print("Flow",counter," - App: ",app_id, "- Direction:",direction)
        print("  - Flow Start  :",re.sub("\..*Z","",str(cgx_millseconds_to_timestamp(flow['flow_start_time_ms']))   ))
        print("  - Source      :",flow['source_ip'], "(" + str(flow['source_port']) + ")"  )
        print("  - Destination :",flow['destination_ip'], "(" + str(flow['destination_port']) + ")"  )
        print("  - Appliction  :",app_id , "(" + str(flow['app_id']) + ")")
        print("  - Protocol    :",cgx_convert_protocol_to_name(flow['protocol']   ))
        print("  - INIT Success:",flow['init_success']   )
        print("  - TXN Success :",flow['success_transactions']   )
        print("  - TXN Fail    :",flow['incomplete_trans']   )
        print("  - TXN Fail    :",safe_resolve(flow,'incomplete_trans')   )
        print("  - OOO Packets :", " (" + str(safe_resolve(flow,'ooo_pkts_c2s')) + "/" + str(safe_resolve(flow,'ooo_pkts_s2c') ) + ") " )
        print("  - Packet Count:", " (" + str(safe_resolve(flow,'packets_c2s')) + "/" + str(safe_resolve(flow,'packets_s2c') ) + ") " )
        print("  - Retrans Pkts:", " (" + str(safe_resolve(flow,'retransmit_pkts_c2s')) + "/" + str(safe_resolve(flow,'retransmit_pkts_s2c') ) + ") " )
        print("  - Domain      :",safe_resolve(flow,'unknown_domain') )
        print("  - PathChng rsn:",safe_resolve(flow,'wan_path_change_reason') )

max_stars_repo_name: karl-sc/CGX-Library-Functions

max_issues_repo_name: karl-sc/CGX-Library-Functions

max_forks_repo_name: karl-sc/CGX-Library-Functions


text:
"""
Rank constraints via convex iteration (Dattorro's Convex Optimization and Euclidean Distance Geometry textbook).

"""
import numpy as np
import cvxpy as cp
import networkx as nx
from timeit import default_timer
from progress.bar import ShadyBar as Bar
from liegroups.numpy import SE3

from graphik.solvers.sdp_formulations import SdpSolverParams
from graphik.solvers.sdp_snl import (
    distance_range_constraints,
    solve_linear_cost_sdp,
    distance_constraints_graph,
    extract_full_sdp_solution,
    extract_solution,
    chordal_sparsity_overlap_constraints
)
from graphik.solvers.constraints import get_full_revolute_nearest_point
from graphik.utils.roboturdf import load_ur10
from graphik.utils.constants import *
from graphik.graphs.graph_base import ProblemGraph
from graphik.graphs.graph_revolute import ProblemGraphRevolute

def random_psd_matrix(N: int, d: int = None, normalize: bool = True) -> np.ndarray:
    """
    Return a random PSD matrix.

    :param N: size of the matrix (NxN)
    :param d: rank of the matrix (full rank if not specified)
    """
    d = N if d is None else d
    P = np.random.rand(N, d)
    W = P @ P.T
    if normalize:
        W = W / np.linalg.norm(W, ord="fro")
    return W

def solve_fantope_closed_form(G: np.ndarray, d:int):
    """

    :param G:
    :param d:
    """
    start = default_timer()
    _, Q = np.linalg.eigh(G)
    Q = np.flip(Q, 1)
    U = Q[:, d:]
    return U@U.T, default_timer() - start

def solve_fantope_sdp(G: np.ndarray, d: int):
    n = G.shape[0]
    Z, cons = fantope_constraints(n, d)
    solve_fantope_iterate(G, Z, cons)
    return Z.value

def fantope_constraints(n: int, rank: int, sparsity_pattern: set = None):
    assert rank < n, "Needs a desired rank less than the problem dimension."
    Z = cp.Variable((n, n), PSD=True)
    constraints = [cp.trace(Z) == float(n - rank), np.eye(Z.shape[0]) - Z >> 0]

    if sparsity_pattern is not None:
        for idx in range(n):
            for jdx in range(idx, n):
                if (idx, jdx) not in sparsity_pattern:
                    constraints += [Z[idx, jdx] == 0.]
                    if idx != jdx:
                        constraints += [Z[jdx, idx] == 0.]

    return Z, constraints

def solve_fantope_iterate(
    G: np.ndarray, Z: cp.Variable, constraints: list, verbose=False, solver_params=None
):
    # TODO: templatize for speed? Ask Filip about that feature, I forget what it's called
    prob = cp.Problem(cp.Minimize(cp.trace(G @ Z)), constraints)
    if solver_params is None:
        solver_params = SdpSolverParams()
    prob.solve(verbose=verbose, solver="MOSEK", mosek_params=solver_params.mosek_params)
    return prob

def solve_fantope_sparse(sdp_variable_map: dict, d: int):
    C_mapping = {}
    t_fantope = 0.0
    for clique in sdp_variable_map:
        G_clique = sdp_variable_map[clique].value  # Assumes it's been solved by cvxpy
        C_mapping[clique], t_clique = solve_fantope_closed_form(G_clique, d)
        t_fantope += t_clique
    return C_mapping, t_fantope

def solve_fantope_sdp_sparse(constraint_clique_dict: dict, sdp_variable_map: dict, d: int, verbose=False,
                             solver_params=None):

    # Make cvxpy variables and constraints for each Fantope
    fantope_sdp_variable_map = {}
    constraints = []
    for clique in sdp_variable_map:
        n_clique = sdp_variable_map[clique].shape[0]
        Z_clique = cp.Variable(sdp_variable_map[clique].shape, PSD=True)
        fantope_sdp_variable_map[clique] = Z_clique
        constraints += [cp.trace(Z_clique) == float(n_clique - d), np.eye(Z_clique.shape[0]) - Z_clique >> 0]

    # Get the overlap constraints that link each Fantope's overlapping variables
    constraints += chordal_sparsity_overlap_constraints(constraint_clique_dict, fantope_sdp_variable_map, d)

    # Solve the sparse Fantope SDP
    if solver_params is None:
        solver_params = SdpSolverParams()
    cost = 0.
    for clique in sdp_variable_map:
        Z_clique = fantope_sdp_variable_map[clique]
        G_clique = sdp_variable_map[clique].value
        cost += cp.trace(G_clique @ Z_clique)
    prob = cp.Problem(cp.Minimize(cost), constraints)
    prob.solve(verbose=verbose, solver="MOSEK", mosek_params=solver_params.mosek_params)

    # Return the desired cost function matrices
    C_mapping = {}
    for clique in sdp_variable_map:
        C_mapping[clique] = fantope_sdp_variable_map[clique].value
    return C_mapping, prob.solver_stats.solve_time

def sparse_eigenvalue_sum(sdp_variable_map: dict, d: int):
    running_eigenvalue_sum = 0.
    for clique in sdp_variable_map:
        Z_clique = sdp_variable_map[clique].value
        running_eigenvalue_sum += np.sum(np.linalg.eigvalsh(Z_clique)[:-d])
    return running_eigenvalue_sum

def get_sparsity_pattern(G, canonical_point_order: list) -> set:
    sparsity_pattern = set()
    G = G.copy()
    # remove the edges that don't have distances defined
    edges = []
    for u, v, data in G.edges(data=True):
        if not data.get(DIST, False):
            edges += [(u, v)]
    G.remove_edges_from(edges)
    undirected = G.to_undirected()
    for idx, point_idx in enumerate(canonical_point_order):
        for jdx in range(idx, len(canonical_point_order)):
            point_jdx = canonical_point_order[jdx]
            if undirected.has_edge(point_idx, point_jdx):
                sparsity_pattern.add((idx, jdx))

    return sparsity_pattern

def convex_iterate_sdp_snl_graph(
    graph: ProblemGraph,
    anchors: dict = {},
    ranges=False,
    max_iters=10,
    sparse=False,
    verbose=False,
    closed_form=True,
    W_init=None,
    abs_eig_sum_tol=1e-6,
    rel_eig_sum_tol=1e-3,
    floor_mode=False,
    scs=False
):
    # get a copy of the current robot + environment graph
    G = nx.DiGraph(graph)
    # G = graph.directed.copy()

    # remove base nodes and all adjacent edges
    G.remove_node("x")
    G.remove_node("y")

    robot = graph.robot
    d = robot.dim
    eig_value_sum_vs_iterations = []

    # If a position is pre-defined for a node, set to anchor
    for node, data in G.nodes(data=True):
        if data.get(POS, None) is not None:
            if node not in ('p0', 'q0') or not floor_mode:
                anchors[node] = data[POS]
    planar_constraints = {'p0': (np.array([0., 0., 1.]), 0.),
                          'q0': (np.array([0., 0., 1.]), 1.)} if floor_mode else None

    # full_points = [node for node in G if node not in ["x", "y"]]
    canonical_point_order = [point for point in G if point not in anchors.keys()]
    constraint_clique_dict = distance_constraints_graph(
        G, anchors, sparse, ee_cost=False, angle_limits=ranges  # TODO: is a param other than ranges needed?
    )

    # Add inequalities (angluar limits, obstacles) if present
    inequality_map = distance_range_constraints(G, constraint_clique_dict, anchors) if ranges else None

    # Save runtimes
    primal_sdp_runtime = 0.
    fantope_solver_runtime = 0.

    # Return feasibility
    feasible = FEASIBLE

    # Track the cost for convergence
    last_cost = 1e6

    n = len(canonical_point_order)
    N = n + d
    C = np.eye(N) if W_init is None else W_init  # Identity satisfies any sparsity pattern by default
    prob = None
    sdp_variable_map = None
    for iter in range(max_iters):
        solution, prob, sdp_variable_map, _ = solve_linear_cost_sdp(
            robot,
            anchors,
            constraint_clique_dict,
            C,
            prob,
            sdp_variable_map,
            canonical_point_order,
            verbose=False,
            inequality_constraints_map=inequality_map,
            planar_constraints=planar_constraints,
            scs=scs,
            warm_start=True
        )
        # Handle infeasibility case
        if solution is INFEASIBLE:
            feasible = INFEASIBLE
            primal_sdp_runtime += prob.solver_stats.solve_time
            break
        elif solution is SOLVER_ERROR:
            feasible = SOLVER_ERROR
            primal_sdp_runtime += 0.0  # TODO: handle this in post-processing
            break
        primal_sdp_runtime += prob.solver_stats.solve_time

        if not sparse:
            G = extract_full_sdp_solution(constraint_clique_dict, canonical_point_order, sdp_variable_map, N, d)
            eigvals_G = np.linalg.eigvalsh(G)  # Returns in ascending order (according to docs)
            eig_value_sum_vs_iterations.append(np.sum(eigvals_G[0:n]))
            C, t_fantope = solve_fantope_closed_form(G, robot.dim)

        else:
            if closed_form:
                C, t_fantope = solve_fantope_sparse(sdp_variable_map, d)
            else:
                C, t_fantope = solve_fantope_sdp_sparse(constraint_clique_dict, sdp_variable_map, d)
            eig_value_sum_vs_iterations.append(sparse_eigenvalue_sum(sdp_variable_map, d))
        fantope_solver_runtime += t_fantope

        # Check for convergence
        eigval_sum_change = last_cost - prob.value
        rel_change = np.abs(eigval_sum_change)/np.abs(last_cost)
        if np.abs(eigval_sum_change) <= abs_eig_sum_tol or prob.value <= abs_eig_sum_tol or rel_change < rel_eig_sum_tol:
            break
        else:
            last_cost = prob.value

    return (
        C,
        constraint_clique_dict,
        sdp_variable_map,
        canonical_point_order,
        eig_value_sum_vs_iterations,
        prob,
        primal_sdp_runtime,
        fantope_solver_runtime,
        feasible
    )

def solve_with_cidgik(graph: ProblemGraphRevolute, T_goal: SE3) -> (dict, dict):
    robot = graph.robot
    n = robot.n

    # Set up the anchors needed as input to CIDGIK
    anchors = {
        "p0": graph.nodes["p0"][POS],
        "q0": graph.nodes["q0"][POS],
        f"p{n}": T_goal.trans,
        f"q{n}": T_goal.trans + T_goal.rot.as_matrix()[:, 2]
    }

    # Solve with CIDGIK
    _, constraint_clique_dict, sdp_variable_map, _, _, _, _, _, feasible = \
        convex_iterate_sdp_snl_graph(
            graph,
            anchors,
            ranges=True,
            sparse=False,
            closed_form=True,
            scs=False
        )

    # Extract the angular configuration
    if feasible is FEASIBLE:
        solution = extract_solution(constraint_clique_dict, sdp_variable_map, robot.dim)

        # Add the end-effector goal points to the solution
        solution[f"p{robot.n}"] = anchors[f"p{robot.n}"]
        solution[f"q{robot.n}"] = anchors[f"q{robot.n}"]

        # Add the base points to the solution
        base_nodes = ["p0", "x", "y", "q0"]
        for node in base_nodes:
            solution[node] = graph.nodes[node][POS]
        G_sol = graph.from_pos(solution)
        q_sol = graph.joint_variables(G_sol, {f"p{n}": T_goal})

        return q_sol, solution
    else:
        return None, None

if __name__ == "__main__":

    # TODO: use the graph convex iteration from above and deprecate the old one
    # UR10 Test
    robot, graph = load_ur10()
    d = robot.dim
    full_points = [f"p{idx}" for idx in range(0, graph.robot.n + 1)] + [
        f"q{idx}" for idx in range(0, graph.robot.n + 1)
    ]
    n_runs = 1000

    # Store results
    final_eigvalue_sum_list = []
    primal_sdp_runtime = []
    fantope_runtime = []
    feasibility_list = []

    final_eigvalue_sum_list_sparse_naive = []
    primal_sdp_runtime_sparse_naive = []
    fantope_runtime_sparse_naive = []
    feasibility_list_sparse_naive = []
    final_eigvalue_sum_list_sparse_sdp = []
    primal_sdp_runtime_sparse_sdp = []
    fantope_runtime_sparse_sdp = []
    feasibility_list_sparse_sdp = []

    bar = Bar("", max=n_runs, check_tty=False, hide_cursor=False)
    for idx in range(n_runs):
        # Generate a random feasible target
        q = robot.random_configuration()
        input_vals = get_full_revolute_nearest_point(graph, q, full_points)

        # End-effectors don't include the base pair at this step, that's inside of convex_iterate_sdp_snl_graph()
        end_effectors = {
            key: input_vals[key] for key in [f"p{robot.n}", f"q{robot.n}"]
        }

        # Run 'dense' default solver
        (
            _,
            _,
            _,
            _,
            eig_value_sum_vs_iterations,
            _,
            t_primal,
            t_fantope,
            feasible
        ) = convex_iterate_sdp_snl_graph(graph, anchors=end_effectors, ranges=False, max_iters=10,
                                         sparse=False, verbose=False, closed_form=True)
        # solution = extract_solution(constraint_clique_dict, sdp_variable_map, d)
        if feasible is FEASIBLE:
            final_eigvalue_sum_list.append(eig_value_sum_vs_iterations[-1])
        else:
            final_eigvalue_sum_list.append(np.NaN)
        primal_sdp_runtime.append(t_primal)
        fantope_runtime.append(t_fantope)
        feasibility_list.append(feasible)

        # Run 'naive' sparse solver (Wang and Yu, 2018)
        (
            _,
            _,
            _,
            _,
            eig_value_sum_vs_iterations,
            _,
            t_primal,
            t_fantope,
            feasible
        ) = convex_iterate_sdp_snl_graph(graph, anchors=end_effectors, ranges=False, max_iters=10,
                                         sparse=True, verbose=False, closed_form=True)
        primal_sdp_runtime_sparse_naive.append(t_primal)
        fantope_runtime_sparse_naive.append(t_fantope)
        if feasible is FEASIBLE:
            final_eigvalue_sum_list_sparse_naive.append(eig_value_sum_vs_iterations[-1])
        else:
            final_eigvalue_sum_list_sparse_naive.append(np.NaN)

        # Run sparse solver where Fantope program has a sparsity pattern matched to the primal SDP's
        (
            _,
            _,
            _,
            _,
            eig_value_sum_vs_iterations,
            _,
            t_primal,
            t_fantope,
            feasible
        ) = convex_iterate_sdp_snl_graph(graph, anchors=end_effectors, ranges=False, max_iters=10,
                                         sparse=True, verbose=False, closed_form=False)
        primal_sdp_runtime_sparse_sdp.append(t_primal)
        fantope_runtime_sparse_sdp.append(t_fantope)
        if feasible is FEASIBLE:
            final_eigvalue_sum_list_sparse_sdp.append(eig_value_sum_vs_iterations[-1])
        else:
            final_eigvalue_sum_list_sparse_sdp.append(np.NaN)

        bar.next()
    bar.finish()

    # Visualize results
    from matplotlib import pyplot as plt
    from matplotlib import rc
    rc("font", **{"family": "serif", "serif": ["Computer Modern"]})
    rc("text", usetex=True)

    # Plot parameters
    n_bins = 100
    colors = ['r', 'g', 'b']
    # Runtime histograms
    primal_sdp_runtime_full = np.vstack([primal_sdp_runtime, primal_sdp_runtime_sparse_naive, primal_sdp_runtime_sparse_sdp]).T
    primal_sdp_runtime_min = np.min(primal_sdp_runtime_full)
    primal_sdp_runtime_max = np.max(primal_sdp_runtime_full)
    primal_sdp_runtime_bins = np.linspace(primal_sdp_runtime_min, primal_sdp_runtime_max, n_bins)
    plt.figure()
    plt.hist(primal_sdp_runtime_full, bins=primal_sdp_runtime_bins, color=colors)
    plt.legend(('Dense', 'Sparse (Naive)', 'Sparse (SDP)'))
    plt.title('Distribution of Primal SDP Runtimes')
    plt.xlabel('Runtime (s)')
    plt.ylabel('Count')
    plt.grid()
    plt.show()

    fantope_runtime_full = np.vstack([fantope_runtime, fantope_runtime_sparse_naive, fantope_runtime_sparse_sdp]).T
    fantope_runtime_min = np.min(fantope_runtime_full)
    fantope_runtime_max = np.max(fantope_runtime_full)
    fantope_runtime_bins = np.linspace(fantope_runtime_min, fantope_runtime_max, n_bins)
    plt.figure()
    plt.hist(fantope_runtime_full, bins=fantope_runtime_bins, color=colors)
    plt.legend(('Dense', 'Sparse (Naive)', 'Sparse (SDP)'))
    plt.title('Distribution of Fantope Solver Runtimes')
    plt.xlabel('Runtime (s)')
    plt.ylabel('Count')
    plt.grid()
    plt.show()

    # Eigenvalue sum histograms
    eigenvalue_sum_full = np.vstack([final_eigvalue_sum_list, final_eigvalue_sum_list_sparse_naive, final_eigvalue_sum_list_sparse_sdp]).T
    eigenvalue_sum_min = max(np.min(eigenvalue_sum_full), 1e-12)
    eigenvalue_sum_max = np.max(eigenvalue_sum_full)
    # eigenvalue_sum_bins = np.log10(np.linspace(eigenvalue_sum_min, eigenvalue_sum_max, n_bins))
    eigenvalue_sum_bins = np.log10(np.logspace(np.log10(eigenvalue_sum_min), np.log10(eigenvalue_sum_max), n_bins))
    plt.figure()
    plt.hist(np.log10(eigenvalue_sum_full), bins=eigenvalue_sum_bins, color=colors)
    plt.legend(('Dense', 'Sparse (Naive)', 'Sparse (SDP)'))
    plt.title('Distribution of Excess Rank')
    plt.xlabel('log$_{10}$ Excess Rank')
    plt.ylabel('Count')
    plt.grid()
    plt.show()

max_stars_repo_name: utiasSTARS/GraphIK

max_issues_repo_name: utiasSTARS/GraphIK

max_forks_repo_name: utiasSTARS/GraphIK

reason: import numpy,import networkx,import cvxpy


text:
# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the LICENSE file in
# the root directory of this source tree. An additional grant of patent rights
# can be found in the PATENTS file in the same directory.abs

from fairseq import logging
from pathlib import Path
from typing import List, Dict, Optional, Any
from dataclasses import dataclass

import numpy as np
import random
import re
import logging
import torch

from examples.s2s_trans.data.data_cfg import S2STDataConfig
from fairseq.data.audio.speech_to_text_dataset import (
    SpeechToTextDataset, SpeechToTextDatasetCreator,
    _collate_frames, get_features_or_waveform
)
from fairseq.data import Dictionary, data_utils as fairseq_data_utils
from fairseq.data.audio.feature_transforms import CompositeAudioFeatureTransform

logger = logging.getLogger(__name__)

@dataclass
class TextToSpeechDatasetItem(object):
    index: int
    src_speech: torch.Tensor
    src_text: torch.Tensor
    tgt_speech: torch.Tensor
    tgt_text: torch.Tensor
    src_orig: Optional[torch.Tensor] = None  # for hubert
    tgt_text_orig: Optional[str] = None  # for kd encoder
    speaker_id: Optional[int] = None
    duration: Optional[torch.Tensor] = None
    pitch: Optional[torch.Tensor] = None
    energy: Optional[torch.Tensor] = None

class S2STDataset(SpeechToTextDataset):
    def __init__(
            self,
            split: str,
            is_train_split: bool,
            cfg: S2STDataConfig,
            src_audio_paths: List[str],
            src_orig_paths: List[str],
            tgt_audio_paths: List[str],
            src_n_frames: List[int],
            tgt_n_frames: List[int],
            src_texts: Optional[List[str]] = None,
            tgt_texts: Optional[List[str]] = None,
            tgt_text_orig: Optional[List[str]] = None,
            speakers: Optional[List[str]] = None,
            src_langs: Optional[List[str]] = None,
            tgt_langs: Optional[List[str]] = None,
            ids: Optional[List[str]] = None,
            src_dict: Optional[Dictionary] = None,
            tgt_dict: Optional[Dictionary] = None,
            pre_tokenizer=None,
            bpe_tokenizer=None,
            n_frames_per_step=1,
            speaker_to_id=None,
            max_sample_size = 9600000,
            durations: Optional[List[List[int]]] = None,
            pitches: Optional[List[str]] = None,
            energies: Optional[List[str]] = None,
            random_crop: Optional[bool] = False,
            pad_audio: Optional[bool] = True,
            normalize: Optional[bool] = False,
    ):

        super(S2STDataset, self).__init__(
            split, is_train_split, cfg, src_audio_paths, src_n_frames,
            src_texts=src_texts, tgt_texts=tgt_texts, speakers=speakers,
            src_langs=src_langs, tgt_langs=tgt_langs, ids=ids,
            tgt_dict=tgt_dict, pre_tokenizer=pre_tokenizer,
            bpe_tokenizer=bpe_tokenizer, n_frames_per_step=n_frames_per_step,
            speaker_to_id=speaker_to_id
        )
        self.check_tgt_lang_tag()
        self.cfg = cfg
        self.shuffle = cfg.shuffle if is_train_split else False

        self.feature_transforms_src = CompositeAudioFeatureTransform.from_config_dict_for_src(
            self.cfg.get_feature_transforms_for_src(split, is_train_split)
        )
        self.feature_transforms_tgt = CompositeAudioFeatureTransform.from_config_dict_for_tgt(
            self.cfg.get_feature_transforms_for_tgt(split, is_train_split)
        )
        self.src_audio_paths = src_audio_paths
        self.src_orig_paths = src_orig_paths
        self.tgt_audio_paths = tgt_audio_paths
        self.tgt_text_orig = tgt_text_orig
        self.src_n_frames = src_n_frames
        self.tgt_n_frames = tgt_n_frames
        self.src_dict = src_dict
        #print(self.src_dict.indices)
        self.tgt_dict = tgt_dict
        self.durations = durations
        self.pitches = pitches
        self.energies = energies

        #for hubert
        self.pad_audio = pad_audio
        self.normalize = normalize
        self.random_crop = random_crop
        self.max_sample_size = max_sample_size
        logger.info(
            f"pad_audio={pad_audio}, random_crop={random_crop}, "
            f"normalize={normalize}, max_sample_size={self.max_sample_size}"
        )

        self.tgt_lens = self.get_tgt_lens_and_check_oov()
        self.src_lens = self.get_src_lens_and_check_oov()
    def __repr__(self):
        return (
            self.__class__.__name__
            + f'(split="{self.split}", n_samples={self.n_samples:_}, '
            f"prepend_tgt_lang_tag={self.cfg.prepend_tgt_lang_tag}, "
            f"shuffle={self.shuffle}, transforms={self.feature_transforms}, "
            f"n_frames_per_step={self.n_frames_per_step}"
        )
    def get_tgt_lens_and_check_oov(self):
        if self.tgt_texts is None:
            return [0 for _ in range(self.n_samples)]
        tgt_lens = []
        n_tokens, n_oov_tokens = 0, 0
        for i in range(self.n_samples):
            tokenized = self.get_tokenized_tgt_text(i).split(" ")
            oov_tokens = [
                t
                for t in tokenized
                if self.tgt_dict.index(t) == self.tgt_dict.unk_index
            ]
            n_tokens += len(tokenized)
            n_oov_tokens += len(oov_tokens)
            tgt_lens.append(len(tokenized))
        logger.info(f"'{self.split}' has {n_oov_tokens / n_tokens * 100:.2f}% OOV in target texts")
        return tgt_lens
    def get_src_lens_and_check_oov(self):
        if self.src_texts is None:
            return [0 for _ in range(self.n_samples)]
        src_lens = []
        n_tokens, n_oov_tokens = 0, 0
        for i in range(self.n_samples):
            tokenized = self.get_tokenized_src_text(i).split(" ")
            oov_tokens = [
                t
                for t in tokenized
                if self.src_dict.index(t) == self.src_dict.unk_index
            ]
            n_tokens += len(tokenized)
            n_oov_tokens += len(oov_tokens)
            src_lens.append(len(tokenized))
        logger.info(f"'{self.split}' has {n_oov_tokens / n_tokens * 100:.2f}% OOV in source texts")
        return src_lens

    def __getitem__(self, index: int) -> TextToSpeechDatasetItem:

        if self.cfg.use_hubert:
            src_orig = self.get_audio(self.src_orig_paths[index])
        else:
            src_orig = None
        #print(self.src_audio_paths[index])
        src_speech = get_features_or_waveform(
            self.src_audio_paths[index],
            need_waveform=self.cfg.use_audio_input,
            use_sample_rate=self.cfg.use_sample_rate,
        )
        # this feature transform may be used to specaugment
        if self.feature_transforms_src is not None:
            assert not self.cfg.use_audio_input
            src_speech = self.feature_transforms_src(src_speech)
        src_speech = torch.from_numpy(src_speech).float()

        #print(self.tgt_audio_paths[index])
        tgt_speech = get_features_or_waveform(
            self.tgt_audio_paths[index],
            need_waveform=self.cfg.use_audio_input,
            use_sample_rate=self.cfg.use_sample_rate,
        )
        if self.feature_transforms_tgt is not None:
            assert not self.cfg.use_audio_input
            tgt_speech = self.feature_transforms_tgt(tgt_speech)

        tgt_speech = torch.from_numpy(tgt_speech).float()
        # stack fbank for target speech of s2st task
        tgt_speech = self.pack_frames(tgt_speech)

        tgt_text = None
        if self.tgt_texts is not None:
            tokenized = self.get_tokenized_tgt_text(index)
            tgt_text = self.tgt_dict.encode_line(
                tokenized, add_if_not_exist=False, append_eos=True
            ).long()
            if self.cfg.prepend_tgt_lang_tag:
                lang_tag_idx = self.get_lang_tag_idx(
                    self.tgt_langs[index], self.tgt_dict
                )
                tgt_text = torch.cat((torch.LongTensor([lang_tag_idx]), tgt_text), 0)

        src_text = None
        if self.src_texts is not None:
            tokenized = self.get_tokenized_src_text(index)
            src_text = self.src_dict.encode_line(
                tokenized, add_if_not_exist=False, append_eos=True
            ).long()
            # remove the eos token
            #src_text = src_text[:len(src_text)-1]
            if self.cfg.prepend_tgt_lang_tag:
                lang_tag_idx = self.get_lang_tag_idx(
                    self.src_langs[index], self.src_dict
                )
                src_text = torch.cat((torch.LongTensor([lang_tag_idx]), src_text), 0)

        speaker_id = None
        if self.speaker_to_id is not None:
            speaker_id = self.speaker_to_id[self.speakers[index]]

        duration, pitch, energy = None, None, None
        if self.durations is not None:
            duration = torch.tensor(
                self.durations[index] + [0], dtype=torch.long  # pad 0 for EOS
            )
        if self.pitches is not None:
            pitch = get_features_or_waveform(self.pitches[index])
            pitch = torch.from_numpy(
                np.concatenate((pitch, [0]))  # pad 0 for EOS
            ).float()
        if self.energies is not None:
            energy = get_features_or_waveform(self.energies[index])
            energy = torch.from_numpy(
                np.concatenate((energy, [0]))  # pad 0 for EOS
            ).float()

        tgt_text_orig = None
        if self.tgt_text_orig is not None:
            tgt_text_orig = self.tgt_text_orig[index]

        return TextToSpeechDatasetItem(
            index=index,
            src_speech=src_speech,
            src_orig= src_orig,
            src_text=src_text,
            tgt_text_orig = tgt_text_orig,
            tgt_speech=tgt_speech,
            tgt_text = tgt_text,
            speaker_id=speaker_id, duration=duration, pitch=pitch,
            energy=energy
        )
    def check_tgt_lang_tag(self):
        if self.cfg.prepend_tgt_lang_tag:
            assert self.tgt_langs is not None and self.tgt_dict is not None
            tgt_lang_tags = [
                self.LANG_TAG_TEMPLATE.format(t) for t in set(self.tgt_langs)
            ]
            assert all(t in self.tgt_dict for t in tgt_lang_tags)
    def get_tokenized_src_text(self, index: int):
        text = self.tokenize(self.pre_tokenizer,self.src_texts[index])
        text = self.tokenize(self.bpe_tokenizer, text)

        return text

    #starting point of fuction for hubert audio input
    def get_audio(self, wav_path):
        import soundfile as sf

        #wav_path = os.path.join(self.audio_root, self.audio_names[index])
        wav, cur_sample_rate = sf.read(wav_path)
        wav = torch.from_numpy(wav).float()
        wav = self.postprocess(wav, cur_sample_rate)
        return wav
    def crop_to_max_size(self, wav, target_size):
        size = len(wav)
        diff = size - target_size
        if diff <= 0:
            return wav, 0

        start, end = 0, target_size
        if self.random_crop:
            start = np.random.randint(0, diff + 1)
            end = size - diff + start
        return wav[start:end], start
    def collater_audio(self, audios, audio_size):
        collated_audios = audios[0].new_zeros(len(audios), audio_size)
        padding_mask = (
            torch.BoolTensor(collated_audios.shape).fill_(False)
            # if self.pad_audio else None
        )
        audio_starts = [0 for _ in audios]
        for i, audio in enumerate(audios):
            diff = len(audio) - audio_size
            if diff == 0:
                collated_audios[i] = audio
            elif diff < 0:
                assert self.pad_audio
                collated_audios[i] = torch.cat(
                    [audio, audio.new_full((-diff,), 0.0)]
                )
                padding_mask[i, diff:] = True
            else:
                collated_audios[i], audio_starts[i] = self.crop_to_max_size(
                    audio, audio_size
                )
        return collated_audios, padding_mask, audio_starts
    def postprocess(self, wav, cur_sample_rate):
        if wav.dim() == 2:
            wav = wav.mean(-1)
        assert wav.dim() == 1, wav.dim()

        # if cur_sample_rate != self.sample_rate:
        #     raise Exception(f"sr {cur_sample_rate} != {self.sample_rate}")

        if self.normalize:
            with torch.no_grad():
                wav = F.layer_norm(wav, wav.shape)
        return wav

    # stopping point
    def collater(self, samples: List[TextToSpeechDatasetItem]) -> Dict[str, Any]:
        if len(samples) == 0:
            return {}

        src_feat_lengths, order = torch.tensor(
            [s.src_speech.shape[0] for s in samples], dtype=torch.long
        ).sort(descending=True)
        id_ = torch.tensor([s.index for s in samples],
                           dtype=torch.long).index_select(0, order)

        src_feat = None
        if not self.cfg.use_hubert:
            src_feat = _collate_frames(
                [s.src_speech for s in samples], self.cfg.use_audio_input
            ).index_select(0, order)

        #for hubert
        if self.cfg.use_hubert:
            audios = [s.src_orig for s in samples]
            audio_sizes = [ s.size(0) for s in audios]
            if self.pad_audio:
                audio_size = min(max(audio_sizes), self.max_sample_size)
            else:
                audio_size = min(min(audio_sizes), self.max_sample_size)
            collated_audios, padding_mask, audio_starts = self.collater_audio(
                audios, audio_size
            )
            collated_audios = collated_audios.index_select(0, order)
            padding_mask = padding_mask.index_select(0, order)
        else:
            collated_audios = None
            padding_mask = None

        src_text = fairseq_data_utils.collate_tokens(
            [s.src_text for s in samples],
            self.src_dict.pad(),
            self.src_dict.eos(),
            left_pad=False,
            move_eos_to_beginning=False,
        ).index_select(0, order)
        src_text_len = torch.tensor([s.src_text.size(0) for s in samples], dtype=torch.long).index_select(0, order)

        tgt_feat_lengths = torch.tensor(
            [s.tgt_speech.shape[0] for s in samples], dtype=torch.long
        ).index_select(0, order)
        tgt_feat = _collate_frames(
            [s.tgt_speech for s in samples], self.cfg.use_audio_input
        ).index_select(0, order)
        tgt_text = fairseq_data_utils.collate_tokens(
            [s.tgt_text for s in samples],
            self.tgt_dict.pad(),
            self.tgt_dict.eos(),
            left_pad=False,
            move_eos_to_beginning=False,
        ).index_select(0, order)
        tgt_text_len = torch.tensor([s.tgt_text.size(0) for s in samples], dtype=torch.long).index_select(0, order)

        speaker = None
        if self.speaker_to_id is not None:
            speaker = torch.tensor(
                [s.speaker_id for s in samples], dtype=torch.long
            ).index_select(0, order).view(-1, 1)

        bsz, _, d = tgt_feat.size()
        prev_s2s_output_tokens = torch.cat(
            (tgt_feat.new_zeros((bsz, 1, d)), tgt_feat[:, :-1, :]), dim=1
        )
        prev_src_text_tokens = fairseq_data_utils.collate_tokens(
                [x.src_text for x in samples],
                self.src_dict.pad(),
                self.src_dict.eos(),
                left_pad=False,
                move_eos_to_beginning=True,
            ).index_select(0, order)
        prev_tgt_text_tokens = fairseq_data_utils.collate_tokens(
                [x.tgt_text for x in samples],
                self.tgt_dict.pad(),
                self.tgt_dict.eos(),
                left_pad=False,
                move_eos_to_beginning=True,
            ).index_select(0, order)

        durations, pitches, energies = None, None, None
        if self.durations is not None:
            durations = fairseq_data_utils.collate_tokens(
                [s.duration for s in samples], 0
            ).index_select(0, order)
            assert src_text.shape[1] == durations.shape[1]
        if self.pitches is not None:
            pitches = _collate_frames([s.pitch for s in samples], True)
            pitches = pitches.index_select(0, order)
            assert src_text.shape[1] == pitches.shape[1]
        if self.energies is not None:
            energies = _collate_frames([s.energy for s in samples], True)
            energies = energies.index_select(0, order)
            assert src_text.shape[1] == energies.shape[1]
        target_texts = [self.tgt_dict.string(samples[i].tgt_text) for i in order]
        tgt_text_orig = [ samples[i].tgt_text_orig for i in order ]
        return {
            "id": id_,
            "net_input": {
                "src_speech": src_feat,
                "src_speech_lens": src_feat_lengths,
                "prev_output_tokens": prev_s2s_output_tokens,
                "prev_src_text_tokens": prev_src_text_tokens,
                "prev_tgt_text_tokens": prev_tgt_text_tokens,
                "collated_audios_orig": collated_audios,
                "padding_mask": padding_mask,
                "speaker": speaker,
            },
            "speaker": speaker,
            "src_text": src_text,
            "src_text_len": src_text_len,
            "tgt_text": tgt_text,
            "tgt_text_len": tgt_text_len,
            "tgt_speech": tgt_feat,
            "target_lengths": tgt_feat_lengths,
            "durations": durations,
            "pitches": pitches,
            "energies": energies,
            "ntokens": sum(tgt_feat_lengths).item(),
            "src_txt_ntokens": sum(src_text_len).item(),
            "tgt_txt_ntokens": sum(tgt_text_len).item(),
            "nsentences": len(samples),
            "target_texts": target_texts,
            "tgt_text_orig": tgt_text_orig,
        }

class S2STDatasetCreator(SpeechToTextDatasetCreator):
    KEY_DURATION = "duration"
    KEY_PITCH = "pitch"
    KEY_ENERGY = "energy"
    # mandatory columns
    KEY_ID, KEY_SRC_AUDIO, KEY_SRC_ORIG, KEY_TGT_AUDIO = "id", "src_audio", "src_orig", "tgt_audio"
    KEY_SRC_N_FRAMES,KEY_TGT_N_FRAMES = "src_n_frames", "tgt_n_frames"
    TGT_ORIG_TXT = "tgt_text_orig"
    KEY_SRC_TEXT, KEY_TGT_TEXT = "src_text", "tgt_text"
    # optional columns
    KEY_SPEAKER = "speaker"
    KEY_SRC_LANG, KEY_TGT_LANG = "src_lang", "tgt_lang"
    # default values
    DEFAULT_SPEAKER = DEFAULT_SRC_TEXT = DEFAULT_LANG = ""
    @classmethod
    def _from_list(
        cls,
        split_name: str,
        is_train_split,
        samples: List[Dict],
        cfg: S2STDataConfig,
        src_dict,
        tgt_dict,
        pre_tokenizer,
        bpe_tokenizer,
        n_frames_per_step,
        speaker_to_id
    ) -> S2STDataset:
        audio_root = Path(cfg.audio_root)
        ids = [s[cls.KEY_ID] for s in samples]
        src_audio_paths = [(audio_root / s[cls.KEY_SRC_AUDIO]).as_posix() for s in samples]
        src_orig_paths = None
        tgt_orig_text = None
        if cfg.use_hubert:
            src_orig_paths = [(audio_root / s[cls.KEY_SRC_ORIG]).as_posix() for s in samples]
        if cfg.kd_encoder:
            tgt_orig_text = [s[cls.TGT_ORIG_TXT] for s in samples]
        tgt_audio_paths = [(audio_root / s[cls.KEY_TGT_AUDIO]).as_posix() for s in samples]
        src_n_frames = [int(s[cls.KEY_SRC_N_FRAMES]) for s in samples]
        tgt_n_frames = [int(s[cls.KEY_TGT_N_FRAMES]) for s in samples]
        tgt_texts = [s[cls.KEY_TGT_TEXT] for s in samples]
        src_texts = [s.get(cls.KEY_SRC_TEXT, cls.DEFAULT_SRC_TEXT) for s in samples]
        speakers = [s.get(cls.KEY_SPEAKER, cls.DEFAULT_SPEAKER) for s in samples]
        src_langs = [s.get(cls.KEY_SRC_LANG, cls.DEFAULT_LANG) for s in samples]
        tgt_langs = [s.get(cls.KEY_TGT_LANG, cls.DEFAULT_LANG) for s in samples]

        durations = [s.get(cls.KEY_DURATION, None) for s in samples]
        durations = [
            None if dd is None else [int(d) for d in dd.split(" ")]
            for dd in durations
        ]
        durations = None if any(dd is None for dd in durations) else durations

        pitches = [s.get(cls.KEY_PITCH, None) for s in samples]
        pitches = [
            None if pp is None else (audio_root / pp).as_posix()
            for pp in pitches
        ]
        pitches = None if any(pp is None for pp in pitches) else pitches

        energies = [s.get(cls.KEY_ENERGY, None) for s in samples]
        energies = [
            None if ee is None else (audio_root / ee).as_posix()
            for ee in energies]
        energies = None if any(ee is None for ee in energies) else energies

        return S2STDataset(
            split_name, is_train_split, cfg, src_audio_paths, src_orig_paths, tgt_audio_paths,
            src_n_frames, tgt_n_frames,
            src_texts, tgt_texts, tgt_orig_text,
            speakers, src_langs, tgt_langs, ids,
            src_dict, tgt_dict,
            pre_tokenizer, bpe_tokenizer, n_frames_per_step,
            speaker_to_id = speaker_to_id, durations=durations, pitches=pitches, energies=energies

        )
    @classmethod
    def _from_tsv(
        cls,
        root: str,
        cfg: S2STDataConfig,
        split: str,
        src_dict,
        tgt_dict,
        is_train_split: bool,
        pre_tokenizer,
        bpe_tokenizer,
        n_frames_per_step,
        speaker_to_id
    ) -> SpeechToTextDataset:
        samples = cls._load_samples_from_tsv(root, split)
        return cls._from_list(
            split, is_train_split, samples, cfg, src_dict,tgt_dict, pre_tokenizer,
            bpe_tokenizer, n_frames_per_step, speaker_to_id
        )
    @classmethod
    def from_tsv(
        cls,
        root: str,
        cfg: S2STDataConfig,
        splits: str,
        src_dict,
        tgt_dict,
        pre_tokenizer,
        bpe_tokenizer,
        is_train_split: bool,
        epoch: int,
        seed: int,
        n_frames_per_step: int = 1,
        speaker_to_id=None
    ) -> SpeechToTextDataset:
        datasets = [
            cls._from_tsv(
                root, cfg, split, src_dict, tgt_dict, is_train_split, pre_tokenizer,
                bpe_tokenizer, n_frames_per_step, speaker_to_id
            )
            for split in splits.split(",")
        ]

        if is_train_split and len(datasets) > 1 and cfg.sampling_alpha != 1.0:
            # temperature-based sampling
            size_ratios = cls.get_size_ratios(datasets, alpha=cfg.sampling_alpha)
            datasets = [
                ResamplingDataset(
                    d, size_ratio=r, seed=seed, epoch=epoch, replace=(r >= 1.0)
                )
                for r, d in zip(size_ratios, datasets)
            ]

        return ConcatDataset(datasets) if len(datasets) > 1 else datasets[0]