from blending.src.connector.connect_simple import ConnectSimple
from blending.src.connector.connect_conflict_random import ConnectConflictRandom
from blending.src.connector.connect_conflict_viable import \
    ConnectConflictAllViable
from blending.src.connector.connect_conflict_interaction_information import \
    ConnectConflictInteractionInformation
from blending.util.blending_config import BlendConfig
from blending.util.blending_error import blending_status

__author__ = 'DongMin Kim'

class ConnectorFinder(object):
    """Provider class to make link connector instance.

    This provider will made the instance of link connector, and returns them to
    the blender.

    Attributes:
        a: An instance of AtomSpace.
        last_status: A last status of class.
        connectors: An available link connector list.
        :type a: opencog.atomspace.AtomSpace
        :type last_status: int
        :type connectors: dict[BaseConnector]
    """

    def __init__(self, a):
        self.a = a
        self.last_status = blending_status.UNKNOWN_ERROR

        self.connectors = {
            ConnectSimple.__name__: ConnectSimple,
            ConnectConflictRandom.__name__: ConnectConflictRandom,
            ConnectConflictAllViable.__name__: ConnectConflictAllViable,
            ConnectConflictInteractionInformation.__name__:
                ConnectConflictInteractionInformation
        }

    def get_connector(self, config_base):
        """Provider method for link connector.

        Args:
            config_base: A Node to save custom config.
            :param config_base: Atom
        Returns:
            The instance of link connector.
            :rtype : BaseChooser
        Raises:
            UserWarning: Can't find the link connector with given name.
        """
        self.last_status = blending_status.IN_PROCESS

        connector = self.connectors.get(
            BlendConfig().get_str(self.a, "link-connector", config_base)
        )
        if connector is not None:
            self.last_status = blending_status.SUCCESS
            return connector(self.a)
        else:
            self.last_status = blending_status.PARAMETER_ERROR
            raise UserWarning


import time
import socket
import struct

try:
    from urllib.parse import urlparse
except ImportError:
    from urlparse import urlparse

import utils
from xlog import getLogger
xlog = getLogger("x_tunnel")

from . import global_var as g
from . import proxy_session
from . import openai_handler

def netloc_to_host_port(netloc, default_port=80):
    if isinstance(netloc, str):
        netloc = netloc.encode("ascii")

    if b":" in netloc:
        host, _, port = netloc.rpartition(b':')
        port = int(port)
    else:
        host = netloc
        port = default_port
    return host, port

class Socks5Server():
    handle_num = 0

    def __init__(self, sock, client, args):
        self.connection = sock
        self.rfile = self.connection.makefile("rb", -1)
        self.wfile = self.connection.makefile("wb", 0)
        self.client_address = client
        self.read_buffer = b""
        self.buffer_start = 0
        self.args = args

    def handle(self):
        self.__class__.handle_num += 1
        try:
            socks_version = self.read_bytes(1)
            if not socks_version:
                return

            if socks_version == b"\x04":
                self.socks4_handler()
            elif socks_version == b"\x05":
                self.socks5_handler()
            elif socks_version == b"C":
                self.https_handler()
            elif socks_version in [b"G", b"P", b"D", b"O", b"H", b"T"]:
                self.http_handler(socks_version)
                return
            else:
                xlog.warn("socks version:%s not supported",  utils.str2hex(socks_version))
                return

        except socket.error as e:
            xlog.warn('proxy handler read error %r', e)
            self.connection.close()
        except Exception as e:
            xlog.exception("proxy handler err:%r", e)
            self.connection.close()

    def read_null_end_line(self):
        sock = self.connection
        sock.setblocking(0)
        try:
            while True:
                n1 = self.read_buffer.find(b"\x00", self.buffer_start)
                if n1 > -1:
                    line = self.read_buffer[self.buffer_start:n1]
                    self.buffer_start = n1 + 1
                    return line

                try:
                    data = sock.recv(8192)
                except socket.error as e:
                    # logging.exception("e:%r", e)
                    if e.errno in [2, 11, 10035]:
                        time.sleep(0.01)
                        continue
                    else:
                        raise e

                self.read_buffer += data
        finally:
            sock.setblocking(1)

    def read_crlf_line(self):
        sock = self.connection
        sock.setblocking(0)
        try:
            while True:
                n1 = self.read_buffer.find(b"\r\n", self.buffer_start)
                if n1 > -1:
                    line = self.read_buffer[self.buffer_start:n1]
                    self.buffer_start = n1 + 2
                    return line

                try:
                    data = sock.recv(8192)
                except socket.error as e:
                    # logging.exception("e:%r", e)
                    if e.errno in [2, 11, 10035]:
                        time.sleep(0.01)
                        continue
                    else:
                        raise e

                self.read_buffer += data
        finally:
            sock.setblocking(1)

    def read_headers(self):
        sock = self.connection
        sock.setblocking(0)
        try:
            while True:
                if self.read_buffer[self.buffer_start:] == b"\r\n":
                    self.buffer_start += 2
                    return b""

                n1 = self.read_buffer.find(b"\r\n\r\n", self.buffer_start)
                if n1 > -1:
                    block = self.read_buffer[self.buffer_start:n1]
                    self.buffer_start = n1 + 4
                    return block

                try:
                    data = sock.recv(8192)
                except socket.error as e:
                    # logging.exception("e:%r", e)
                    if e.errno in [2, 11, 10035]:
                        time.sleep(0.01)
                        continue
                    else:
                        raise e

                self.read_buffer += data
        finally:
            sock.setblocking(1)

    def read_bytes(self, size):
        sock = self.connection
        sock.setblocking(1)
        try:
            while True:
                left = len(self.read_buffer) - self.buffer_start
                if left >= size:
                    break

                need = size - left

                try:
                    data = sock.recv(need)
                except socket.error as e:
                    # logging.exception("e:%r", e)
                    if e.errno in [2, 11, 10035]:
                        time.sleep(0.01)
                        continue
                    else:
                        raise e

                if len(data):
                    self.read_buffer += data
                else:
                    raise socket.error("recv fail")
        finally:
            sock.setblocking(1)

        data = self.read_buffer[self.buffer_start:self.buffer_start + size]
        self.buffer_start += size
        return data

    def socks4_handler(self):
        # Socks4 or Socks4a
        sock = self.connection
        cmd = ord(self.read_bytes(1))
        if cmd != 1:
            xlog.warn("Socks4 cmd:%d not supported", cmd)
            return

        data = self.read_bytes(6)
        port = struct.unpack(">H", data[0:2])[0]
        addr_pack = data[2:6]
        if addr_pack[0:3] == b'\x00\x00\x00' and addr_pack[3] != b'\x00':
            domain_mode = True
        else:
            ip = socket.inet_ntoa(addr_pack)
            domain_mode = False

        user_id = self.read_null_end_line()
        if len(user_id):
            xlog.debug("Socks4 user_id:%s", user_id)

        if domain_mode:
            addr = self.read_null_end_line()
        else:
            addr = ip

        conn_id = proxy_session.create_conn(sock, addr, port)
        if not conn_id:
            xlog.warn("Socks4 connect fail, no conn_id")
            reply = b"\x00\x5b\x00" + addr_pack + struct.pack(">H", port)
            sock.send(reply)
            return

        xlog.info("Socks4:%r to %s:%d, conn:%d", self.client_address, addr, port, conn_id)
        reply = b"\x00\x5a" + addr_pack + struct.pack(">H", port)
        sock.send(reply)

        if len(self.read_buffer) - self.buffer_start:
            g.session.conn_list[conn_id].transfer_received_data(self.read_buffer[self.buffer_start:])

        g.session.conn_list[conn_id].start(block=True)

    def socks5_handler(self):
        sock = self.connection
        auth_mode_num = ord(self.read_bytes(1))
        data = self.read_bytes(auth_mode_num)

        sock.send(b"\x05\x00")  # socks version 5, no auth needed.
        try:
            data = self.read_bytes(4)
        except Exception as e:
            xlog.debug("socks5 auth num:%d, list:%s", auth_mode_num, utils.str2hex(data))
            xlog.warn("socks5 protocol error:%r", e)
            return

        socks_version = ord(data[0:1])
        if socks_version != 5:
            xlog.warn("request version:%d error", socks_version)
            return

        command = ord(data[1:2])
        if command != 1:  # 1. Tcp connect
            xlog.warn("request not supported command mode:%d", command)
            sock.send(b"\x05\x07\x00\x01")  # Command not supported
            return

        addrtype_pack = data[3:4]
        addrtype = ord(addrtype_pack)
        if addrtype == 1:  # IPv4
            addr_pack = self.read_bytes(4)
            addr = socket.inet_ntoa(addr_pack)
        elif addrtype == 3:  # Domain name
            domain_len_pack = self.read_bytes(1)[0:1]
            domain_len = ord(domain_len_pack)
            domain = self.read_bytes(domain_len)
            addr_pack = domain_len_pack + domain
            addr = domain
        elif addrtype == 4:  # IPv6
            addr_pack = self.read_bytes(16)
            addr = socket.inet_ntop(socket.AF_INET6, addr_pack)
        else:
            xlog.warn("request address type unknown:%d", addrtype)
            sock.send(b"\x05\x07\x00\x01")  # Command not supported
            return

        port = struct.unpack('>H', self.rfile.read(2))[0]

        conn_id = proxy_session.create_conn(sock, addr, port)
        if not conn_id:
            xlog.warn("create conn fail")
            reply = b"\x05\x01\x00" + addrtype_pack + addr_pack + struct.pack(">H", port)
            sock.send(reply)
            return

        xlog.info("socks5 %r connect to %s:%d conn:%d", self.client_address, addr, port, conn_id)
        reply = b"\x05\x00\x00" + addrtype_pack + addr_pack + struct.pack(">H", port)
        try:
            sock.send(reply)
        except Exception as e:
            if conn_id in g.session.conn_list:
                g.session.conn_list[conn_id].do_stop("close_on_Socks5_reply")
            xlog.warn("socks5 %r connect to %s:%d conn_id:%d closed:%r", self.client_address, addr, port, conn_id, e)
            return

        if len(self.read_buffer) - self.buffer_start:
            g.session.conn_list[conn_id].transfer_received_data(self.read_buffer[self.buffer_start:])

        g.session.conn_list[conn_id].start(block=True)

    def https_handler(self):
        line = self.read_crlf_line()
        line = line.decode('iso-8859-1')
        words = line.split()
        if len(words) == 3:
            command, path, version = words
        elif len(words) == 2:
            command, path = words
            version = b"HTTP/1.1"
        else:
            xlog.warn("https req line fail:%s", line)
            return

        if command != "ONNECT":
            xlog.warn("https req line fail:%s", line)
            return

        host, _, port = path.rpartition(':')
        host = host.encode()
        port = int(port)

        header_block = self.read_headers()

        sock = self.connection
        conn_id = proxy_session.create_conn(sock, host, port)
        if not conn_id:
            xlog.warn("create conn fail")
            sock.send(b'HTTP/1.1 500 Fail\r\n\r\n')
            return

        xlog.info("https %r connect to %s:%d conn:%d", self.client_address, host, port, conn_id)
        try:
            sock.send(b'HTTP/1.1 200 OK\r\n\r\n')
        except:
            xlog.warn("https %r connect to %s:%d conn:%d closed.", self.client_address, host, port, conn_id)

        if (len(self.read_buffer) - self.buffer_start) > 0:
            g.session.conn_list[conn_id].transfer_received_data(self.read_buffer[self.buffer_start:])

        g.session.conn_list[conn_id].start(block=True)

    def http_handler(self, first_char):
        req_line = self.read_crlf_line()
        words = req_line.split()
        if len(words) == 3:
            method, url, http_version = words
        elif len(words) == 2:
            method, url = words
            http_version = b"HTTP/1.1"
        else:
            xlog.warn("http req line fail:%s", req_line)
            return

        method = first_char + method
        # if method not in ["GET", "HEAD", "POST", "PUT", "DELETE", "OPTIONS", "TRACE", "PATCH"]:
        #    xlog.warn("https req method not known:%s", method)

        if url.startswith(b"http://") or url.startswith(b"HTTP://"):
            o = urlparse(url)
            host, port = netloc_to_host_port(o.netloc)

            p = url[7:].find(b"/")
            if p >= 0:
                path = url[7+p:]
            else:
                path = b"/"
        else:
            header_block = self.read_headers()
            lines = header_block.split(b"\r\n")
            path = url
            host = None
            headers = {}
            for line in lines:
                key, _, value = line.partition(b":")
                headers[key] = value
                if key.lower() == b"host":
                    host, port = netloc_to_host_port(value)
            if host is None:
                xlog.warn("http proxy host can't parsed. %s %s", req_line, header_block)
                self.connection.send(b'HTTP/1.1 500 Fail\r\n\r\n')
                return

            if url.startswith(b"/openai/"):
                content_length = int(headers.get(b"Content-Length", 0))
                req_body = self.read_bytes(content_length)
                return openai_handler.handle_openai(method, url, headers, req_body, self.connection)

        sock = self.connection
        conn_id = proxy_session.create_conn(sock, host, port)
        if not conn_id:
            xlog.warn("create conn fail")
            sock.send(b'HTTP/1.1 500 Fail\r\n\r\n')
            return

        xlog.info("http %r connect to %s:%d conn:%d", self.client_address, host, port, conn_id)

        new_req_line = b"%s %s %s" % (method, path, http_version)
        left_buf = new_req_line + self.read_buffer[(len(req_line) + 1):]
        g.session.conn_list[conn_id].transfer_received_data(left_buf)

        g.session.conn_list[conn_id].start(block=True)


#import csv
 #precision 2
#with open("mpg.csv") as csvfile:
   #mpg=list(csv.DictReader(csvfile))

#print(mpg)

#numpy notes

# x.dot(Y)
# eye(n)--> I matrix
#.arange(start, end, interval)
#.diag-->对角线
#.linspace(start, end, how many number)
#.ones()
#.zeros()
#.repeat([],number)--->依照指定个数复制指定向量中每个元素并放到新的阵列中
#.array([]*3)-->整个阵列重复3次
#.vstack()/.hstack()
#.atype()---转换资料型态
#.max()/.min()/.sum/.mean()/.std()/.argmax()--maxium position
#切割原本array后得到新的array，如果对新array做修改会影响到旧的---除非用.copy()
#iterate in two array-->.zip  : for i, j in zip(x,y)
#lambda:return function reference

#--------------------------------

#pandas:
#(1)series: list-->series，string=object, int=int
#nan-->not a number:数字List中的None转换成Series后自动转换成的资料(格式:float)，String中的None则被转成object
#.index:列出series中的index
#pd.series([],index:[])
#pd.iloc[]/pd.loc[]-->用index位子找item/index值找item，这两个不是Method 是 attributes(?)
#a.head()--??
#a method of computatiom: vactorization
#(%%)timeit n= 100 ---->计算循环次数与运算数度 (Default=1000/ 可以指定n)
#.set_value(label,value)
#.append() : 修改物件衍伸的资料却不会修改物件本身
#可以直接对Series运算, ex: data=data*2, data=data==2 (输出True or False, dtype: bool)
#数字运算：.sum()/.max()/.prod()/ .mean()/.median()/.std()/.nlargest(n)取前n大的数字/.nsmallest(n)
#字串运算：data.str.upper()/.lower()/.len()-->把字串内容转换为其长度/.cat(sep='')-->把字串串起来，中间用""连接/.contains('P')-->判断每个字串是否有此字串，输出布林/.replace("b","a")
#Series是一维资料--资料-索引

#(2)DataFrame
#.size: 数量/ .shape: 行列数
#cleaning tasks?
#.DataFrame(字典,index=[]):输入的配对为字段，可以用index自订索引
#print(,sep="\n")-->分隔
#.DataFrame([],index=[])--->与numpy的array类似
#可以用.loc(row index, column index)->根据索引，注意.iloc/.loc()输出的格式:输出类似字典格式(key:value)-->根据序列
#可以直接用值呼叫对应的栏: ex: data[column]
#取得特定栏或列--->变成Series,可用Series的运算
#data["new column"]-->直接建立新栏/ 或是用正是语法:pd.series()
#可以直接对栏为运算(对Series)后指派给新的字段： data["new"]=data["n"]*data['m']
#.T--->transpose
#chaining: .log[][]--->python have to coppy the data, may cost problems
#.drop()--->return a copy of data which drop the specific data(rows )
#和Series一样，要增加新的资料可以直接用Data[]
#可以直接对字段进行运算: df['Cost'] *= 0.8
#-->Series到DataFrame的index是怎么分布的??
#index_col(?)--设定读取档案或资料时转成DataFrame的Index的字段(将指定的栏转成index)，
# skiprows=n -->Line numbers to skip (0-indexed) or number of lines to skip (int) at the start of the file.
#Boolean Masking: a good querying  --用.where(条件) 来mask
#.dropna()-->去掉NAN资料 Default=0--row
#Boolean mask可以做多重逻辑条件--&,|
#df['Name'][df['Cost']>3]
#.set_index()/.reset_index()---?
#df['country'] = df.index---->为何能保存?
#df = df.set_index('Gold')
#df.head()
#

#2020/03/17 (week3)
#.merge
#pd.merge(dataframe1,dataframe2,how='inner/outer/right/left',left_index=/left_on=,right_index=/right_on=)
#  -->inner:交集，outer:联集,right/left-->将资料并入左边或右边的dataFrame
# -->left: use only keys from left frame, similar to a SQL left outer join; preserve key order.
#--->right: use only keys from right frame, similar to a SQL right outer join; preserve key order.
#  -->left/right_on: 用column 并入

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

#小结：
# 1.python中String是object, 且不可直接更改单一内容(测试)
# 2. Series是一维资料，但也有index(会被一起列出)，index和Series都可以用list表示或循环(iterate)，可以指定index(用list)
# 3. Dataframe是二维资料，用字典形式的方式指派，每一个Key:value是一个栏，value和key都可以用list/tuple(??)表示。
# 4. Dataframe也有Index 可以用各种方式指派index也能用巢状方式指派，可以将整栏指派为新的Index，然后用index取出资料

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

#idiom

import pandas as pd
purchase_1 = pd.Series({'Name': 'Chris',
                        'Item Purchased': 'Dog Food',
                        'Cost': 22.50})
purchase_2 = pd.Series({'Name': 'Kevyn',
                        'Item Purchased': 'Kitty Litter',
                        'Cost': 2.50})
purchase_3 = pd.Series({'Name': 'Vinod',
                        'Item Purchased': 'Bird Seed',
                        'Cost': 5.00})

df = pd.DataFrame([purchase_1, purchase_2, purchase_3], index=['Store 1', 'Store 1', 'Store 2'])

print(df)
print(df.drop(df[df["Item Purchased"=="Dog Food"].index]))

df = pd.read_csv('olympics.csv', index_col=0, skiprows=1)

for col in df.columns:
    if col[:2]=='01':
        df.rename(columns={col:'Gold'+col[4:]}, inplace=True)
    if col[:2]=='02':
        df.rename(columns={col:'Silver'+col[4:]}, inplace=True)
    if col[:2]=='03':
        df.rename(columns={col:'Bronze'+col[4:]}, inplace=True)
    if col[:1]=='№':
        df.rename(columns={col:'#'+col[1:]}, inplace=True)

names_ids = df.index.str.split('\s\(') # split the index by '('

df.index = names_ids.str[0] # the [0] element is the country name (new index)
df['ID'] = names_ids.str[1].str[:3] # the [1] element is the abbreviation or ID (take first 3 characters from that)

df = df.drop('Totals')
df.head()


#
#

import isce
import isceobj
import stdproc
from isceobj.Util.Poly2D import Poly2D
import logging
from isceobj.Util.decorators import use_api

import os
import numpy as np
import shelve

logger = logging.getLogger('isce.insar.runResampleSlc')

def runResampleSlc(self, kind='coarse'):
    '''
    Kind can either be coarse, refined or fine.
    '''

    if kind not in ['coarse', 'refined', 'fine']:
        raise Exception('Unknown operation type {0} in runResampleSlc'.format(kind))

    if kind == 'fine':
        if not (self.doRubbersheetingRange | self.doRubbersheetingAzimuth): # Modified by V. Brancato 10.10.2019
            print('Rubber sheeting not requested, skipping resampling ....')
            return

    logger.info("Resampling secondary SLC")

    secondaryFrame = self._insar.loadProduct( self._insar.secondarySlcCropProduct)
    referenceFrame = self._insar.loadProduct( self._insar.referenceSlcCropProduct)

    inimg = isceobj.createSlcImage()
    inimg.load(secondaryFrame.getImage().filename + '.xml')
    inimg.setAccessMode('READ')

    prf = secondaryFrame.PRF

    doppler = secondaryFrame._dopplerVsPixel
    coeffs = [2*np.pi*val/prf for val in doppler]

    dpoly = Poly2D()
    dpoly.initPoly(rangeOrder=len(coeffs)-1, azimuthOrder=0, coeffs=[coeffs])

    rObj = stdproc.createResamp_slc()
    rObj.slantRangePixelSpacing = secondaryFrame.getInstrument().getRangePixelSize()
    rObj.radarWavelength = secondaryFrame.getInstrument().getRadarWavelength()
    rObj.dopplerPoly = dpoly

    # for now let's start with None polynomial. Later this should change to
    # the misregistration polynomial

    misregFile = os.path.join(self.insar.misregDirname, self.insar.misregFilename)
    if ((kind in ['refined','fine']) and os.path.exists(misregFile+'_az.xml')):
        azpoly = self._insar.loadProduct(misregFile + '_az.xml')
        rgpoly = self._insar.loadProduct(misregFile + '_rg.xml')
    else:
        print(misregFile , " does not exist.")
        azpoly = None
        rgpoly = None

    rObj.azimuthOffsetsPoly = azpoly
    rObj.rangeOffsetsPoly = rgpoly
    rObj.imageIn = inimg

    #Since the app is based on geometry module we expect pixel-by-pixel offset
    #field
    offsetsDir = self.insar.offsetsDirname

    # Modified by V. Brancato 10.10.2019
    #rgname = os.path.join(offsetsDir, self.insar.rangeOffsetFilename)

    if kind in ['coarse', 'refined']:
        azname = os.path.join(offsetsDir, self.insar.azimuthOffsetFilename)
        rgname = os.path.join(offsetsDir, self.insar.rangeOffsetFilename)
        flatten = True
    else:
        azname = os.path.join(offsetsDir, self.insar.azimuthRubbersheetFilename)
        if self.doRubbersheetingRange:
           print('Rubbersheeting in range is turned on, taking the cross-correlation offsets')
           print('Setting Flattening to False')
           rgname = os.path.join(offsetsDir, self.insar.rangeRubbersheetFilename)
           flatten=False
        else:
           print('Rubbersheeting in range is turned off, taking range geometric offsets')
           rgname = os.path.join(offsetsDir, self.insar.rangeOffsetFilename)
           flatten=True

    rngImg = isceobj.createImage()
    rngImg.load(rgname + '.xml')
    rngImg.setAccessMode('READ')

    aziImg = isceobj.createImage()
    aziImg.load(azname + '.xml')
    aziImg.setAccessMode('READ')

    width = rngImg.getWidth()
    length = rngImg.getLength()

# Modified by V. Brancato 10.10.2019
    #flatten = True
    rObj.flatten = flatten
    rObj.outputWidth = width
    rObj.outputLines = length
    rObj.residualRangeImage = rngImg
    rObj.residualAzimuthImage = aziImg

    if referenceFrame is not None:
        rObj.startingRange = secondaryFrame.startingRange
        rObj.referenceStartingRange = referenceFrame.startingRange
        rObj.referenceSlantRangePixelSpacing = referenceFrame.getInstrument().getRangePixelSize()
        rObj.referenceWavelength = referenceFrame.getInstrument().getRadarWavelength()

    # preparing the output directory for coregistered secondary slc
    coregDir = self.insar.coregDirname

    os.makedirs(coregDir, exist_ok=True)

    # output file name of the coregistered secondary slc
    img = secondaryFrame.getImage()

    if kind  == 'coarse':
        coregFilename = os.path.join(coregDir , self._insar.coarseCoregFilename)
    elif kind == 'refined':
        coregFilename = os.path.join(coregDir, self._insar.refinedCoregFilename)
    elif kind == 'fine':
        coregFilename = os.path.join(coregDir, self._insar.fineCoregFilename)
    else:
        print('Exception: Should not have gotten to this stage')

    imgOut = isceobj.createSlcImage()
    imgOut.setWidth(width)
    imgOut.filename = coregFilename
    imgOut.setAccessMode('write')

    rObj.resamp_slc(imageOut=imgOut)

    imgOut.renderHdr()

    return


import io
import re

from setuptools import setup

# with io.open('README.md', 'rt', encoding='utf8') as f:
#     readme = f.read()

with io.open('src/__init__.py', 'rt', encoding='utf8') as f:
    version = re.search(r'__version__ = \'(.*?)\'', f.read()).group(1)

setup(
    name='astportal',
    version=version,

    license='MIT',
    description='Automation Platform powered by Flask',
    # long_description=readme,
    long_description_content_type='text/markdown',

    author='Maksim Avramenko',
    author_email='MAvramenko@humans.net',

    platforms=['any'],
    install_requires=[
        'flask >= 1.0.2',
        'flask_cors >= 3.0.6',
        'flask_sqlalchemy >= 2.3.2',
        'flask_jwt_extended >= 3.12.1',
        'python-dotenv >= 0.9.1',
        'rq >= 0.12.0',
        'cerberus >= 1.2',
        'rq-scheduler >= 0.8.3',
        'ldap3>=2.6.1',
        'PyMySQL>=0.9.3'
    ],
    setup_requires=[],
    tests_require=[],
    extras_require={},
    test_suite='tests',
    entry_points={
        'console_scripts': [
            'src = src.cli:cli'
        ]
    },
    packages=['kitchen'],
    include_package_data=True,
    zip_safe=False,
    classifiers=[
        'Development Status :: 5 - Production/Stable',
        'Environment :: Web Environment',
        'License :: OSI Approved :: MIT License',
        'Operating System :: OS Independent',
        'Programming Language :: Python :: 3',
        'Programming Language :: Python :: 3.6',
        'Programming Language :: Python :: 3.7',
        'Topic :: Software Development :: Libraries :: Application Frameworks',
        'Topic :: Software Development :: Libraries :: Python Modules',
    ]
)


# Configuration file for the Sphinx documentation builder.
#
# This file only contains a selection of the most common options. For a full
# list see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Path setup --------------------------------------------------------------

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.

import os
import subprocess
import sys
from pathlib import Path

from packaging import version as Version

PROJECT_ROOT_DIR = Path(__file__).parent.parent.resolve()

sys.path.insert(0, os.fspath(PROJECT_ROOT_DIR))

# -- Project information -----------------------------------------------------

project = "onETL"
copyright = "2023, DataOps.ETL"
author = "DataOps.ETL"

# The version info for the project you're documenting, acts as replacement for
# |version| and |release|, also used in various other places throughout the
# built documents.
#
# The short X.Y version.

ver = Version.parse(subprocess.check_output("python ../setup.py --version", shell=True, text=True).strip())
version = ver.base_version
# The full version, including alpha/beta/rc tags.
release = ver.public

# -- General configuration ---------------------------------------------------

# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
extensions = [
    "numpydoc",
    "sphinx_design",
    "sphinx_substitution_extensions",
    "sphinx_tabs.tabs",
    "sphinx_toolbox.more_autodoc.autoprotocol",
    "sphinx_toolbox.github",
    "sphinx_copybutton",
    "sphinx.ext.autodoc",
    "sphinx.ext.autosummary",
    "sphinxcontrib.autodoc_pydantic",
    "sphinxcontrib.towncrier",  # provides `towncrier-draft-entries` directive
    "sphinxcontrib.plantuml",
    "sphinx.ext.extlinks",
    "sphinx_favicon",
]
numpydoc_show_class_members = True
autodoc_pydantic_model_show_config = False
autodoc_pydantic_model_show_config_summary = False
autodoc_pydantic_model_show_config_member = False
autodoc_pydantic_model_show_json = False
autodoc_pydantic_model_show_validator_summary = False
autodoc_pydantic_model_show_validator_members = False
autodoc_pydantic_field_list_validators = False
sphinx_tabs_disable_tab_closing = True

towncrier_draft_autoversion_mode = "draft"
towncrier_draft_include_empty = False
towncrier_draft_working_directory = PROJECT_ROOT_DIR

github_username = "MobileTeleSystems"
github_repository = "onetl"

rst_prolog = f"""
.. |support_hooks| image:: https://img.shields.io/badge/%20-support%20hooks-blue
    :target: https://onetl.readthedocs.io/en/{ver}/hooks/index.html
"""

# Add any paths that contain templates here, relative to this directory.
templates_path = ["_templates"]

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This pattern also affects html_static_path and html_extra_path.
exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]

# -- Options for HTML output -------------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.

html_theme = "furo"

# Theme options are theme-specific and customize the look and feel of a theme
# further.  For a list of options available for each theme, see the
# documentation.
html_theme_options = {}

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ["_static"]
html_logo = "./_static/logo.svg"
favicons = [
    {"rel": "icon", "href": "icon.svg", "type": "image/svg+xml"},
]
# The master toctree document.
master_doc = "index"

# The language for content autogenerated by Sphinx. Refer to documentation
# for a list of supported languages.
#
# This is also used if you do content translation via gettext catalogs.
# Usually you set "language" from the command line for these cases.
language = "en"

# The name of the Pygments (syntax highlighting) style to use.
pygments_style = "sphinx"

# Create an alias for etl-entities lib in onetl documentation
extlinks = {
    "etl-entities": ("https://etl-entities.readthedocs.io/en/stable/%s", None),
}

# If true, `todo` and `todoList` produce output, else they produce nothing.
todo_include_todos = False

# -- Options for HTMLHelp output ------------------------------------------

# Output file base name for HTML help builder.
htmlhelp_basename = "onetl-doc"


from queue import deque
class Solution:
    def largestRectangleArea(self, heights: List[int]) -> int:
        n = len(heights)
        left_boundry, right_boundry = [0 for i in range(n)], [n-1 for i in range(n)]
        left_stack, right_stack = deque(), deque()
        i = 0
        while i < n:
            if left_stack:
                if heights[left_stack[-1]] >= heights[i]:
                    while left_stack and heights[left_stack[-1]] >= heights[i]:
                        left_stack.pop()
                    if left_stack:
                        left_boundry[i] = left_stack[-1]+1
                    left_stack.append(i)
                else:
                    left_stack.append(i)
                    left_boundry[i] = i
            else:
                left_stack.append(i)
            i += 1

        i = n-1
        while i >= 0:
            if right_stack:
                if heights[right_stack[-1]] >= heights[i]:
                    while right_stack and heights[right_stack[-1]] >= heights[i]:
                        right_stack.pop()
                    if right_stack:
                        right_boundry[i] = right_stack[-1]-1
                    right_stack.append(i)
                else:
                    right_stack.append(i)
                    right_boundry[i] = i
            else:
                right_stack.append(i)
            i -= 1

        max_area = 0
        for i in range(n):
            max_area = max(max_area, (right_boundry[i]-left_boundry[i]+1)*heights[i])
        return max_area


###########################################################################################
# Implementation of Geometric-Complete Perceptron layers
#
# Papers:
# (1) Geometry-Complete Perceptron Networks for 3D Molecular Graphs,
#     by A Morehead, J Cheng
# (2) Geometry-Complete Diffusion for 3D Molecule Generation,
#     by A Morehead, J Cheng
#
# Orginal repositories:
# (1) https://github.com/BioinfoMachineLearning/GCPNet
# (2) https://github.com/BioinfoMachineLearning/Bio-Diffusion
###########################################################################################

from copy import copy
from functools import partial
from typing import Any, Optional, Tuple, Union

import torch
import torch_scatter
from beartype import beartype
from graphein.protein.tensor.data import ProteinBatch
from jaxtyping import Bool, Float, Int64, jaxtyped
from omegaconf import DictConfig, OmegaConf
from torch import nn
from torch_geometric.data import Batch

from proteinworkshop.models.graph_encoders.components import radial
from proteinworkshop.models.graph_encoders.components.wrappers import (
    ScalarVector,
)
from proteinworkshop.models.utils import (
    get_activations,
    is_identity,
    safe_norm,
)

class VectorDropout(nn.Module):
    """
    Adapted from https://github.com/drorlab/gvp-pytorch
    """

    def __init__(self, drop_rate: float):
        super().__init__()
        self.drop_rate = drop_rate

    @jaxtyped
    @beartype
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        :param x: `torch.Tensor` corresponding to vector channels
        """
        device = x[0].device
        if not self.training:
            return x
        mask = torch.bernoulli(
            (1 - self.drop_rate) * torch.ones(x.shape[:-1], device=device)
        ).unsqueeze(-1)
        x = mask * x / (1 - self.drop_rate)
        return x

class GCPDropout(nn.Module):
    """
    Adapted from https://github.com/drorlab/gvp-pytorch
    """

    def __init__(self, drop_rate: float, use_gcp_dropout: bool = True):
        super().__init__()
        self.scalar_dropout = (
            nn.Dropout(drop_rate) if use_gcp_dropout else nn.Identity()
        )
        self.vector_dropout = (
            VectorDropout(drop_rate) if use_gcp_dropout else nn.Identity()
        )

    @jaxtyped
    @beartype
    def forward(
        self, x: Union[torch.Tensor, ScalarVector]
    ) -> Union[torch.Tensor, ScalarVector]:
        if isinstance(x, torch.Tensor) and x.shape[0] == 0:
            return x
        elif isinstance(x, ScalarVector) and (
            x.scalar.shape[0] == 0 or x.vector.shape[0] == 0
        ):
            return x
        elif isinstance(x, torch.Tensor):
            return self.scalar_dropout(x)
        return ScalarVector(
            self.scalar_dropout(x[0]), self.vector_dropout(x[1])
        )

class GCPLayerNorm(nn.Module):
    """
    Adapted from https://github.com/drorlab/gvp-pytorch
    """

    def __init__(
        self, dims: ScalarVector, eps: float = 1e-8, use_gcp_norm: bool = True
    ):
        super().__init__()
        self.scalar_dims, self.vector_dims = dims
        self.scalar_norm = (
            nn.LayerNorm(self.scalar_dims) if use_gcp_norm else nn.Identity()
        )
        self.use_gcp_norm = use_gcp_norm
        self.eps = eps

    @staticmethod
    @jaxtyped
    @beartype
    def norm_vector(
        v: torch.Tensor, use_gcp_norm: bool = True, eps: float = 1e-8
    ) -> torch.Tensor:
        v_norm = v
        if use_gcp_norm:
            vector_norm = torch.clamp(
                torch.sum(torch.square(v), dim=-1, keepdim=True), min=eps
            )
            vector_norm = torch.sqrt(
                torch.mean(vector_norm, dim=-2, keepdim=True)
            )
            v_norm = v / vector_norm
        return v_norm

    @jaxtyped
    @beartype
    def forward(
        self, x: Union[torch.Tensor, ScalarVector]
    ) -> Union[torch.Tensor, ScalarVector]:
        if isinstance(x, torch.Tensor) and x.shape[0] == 0:
            return x
        elif isinstance(x, ScalarVector) and (
            x.scalar.shape[0] == 0 or x.vector.shape[0] == 0
        ):
            return x
        elif not self.vector_dims:
            return self.scalar_norm(x)
        s, v = x
        return ScalarVector(
            self.scalar_norm(s),
            self.norm_vector(v, use_gcp_norm=self.use_gcp_norm, eps=self.eps),
        )

class GCP(nn.Module):
    def __init__(
        self,
        input_dims: ScalarVector,
        output_dims: ScalarVector,
        nonlinearities: Tuple[Optional[str]] = ("silu", "silu"),
        scalar_out_nonlinearity: Optional[str] = "silu",
        scalar_gate: int = 0,
        vector_gate: bool = True,
        feedforward_out: bool = False,
        bottleneck: int = 1,
        scalarization_vectorization_output_dim: int = 3,
        enable_e3_equivariance: bool = False,
        **kwargs,
    ):
        super().__init__()

        if nonlinearities is None:
            nonlinearities = ("none", "none")

        self.scalar_input_dim, self.vector_input_dim = input_dims
        self.scalar_output_dim, self.vector_output_dim = output_dims
        self.scalar_nonlinearity, self.vector_nonlinearity = (
            get_activations(nonlinearities[0], return_functional=True),
            get_activations(nonlinearities[1], return_functional=True),
        )
        self.scalar_gate, self.vector_gate = scalar_gate, vector_gate
        self.enable_e3_equivariance = enable_e3_equivariance

        if self.scalar_gate > 0:
            self.norm = nn.LayerNorm(self.scalar_output_dim)

        if self.vector_input_dim:
            assert (
                self.vector_input_dim % bottleneck == 0
            ), f"Input channel of vector ({self.vector_input_dim}) must be divisible with bottleneck factor ({bottleneck})"

            self.hidden_dim = (
                self.vector_input_dim // bottleneck
                if bottleneck > 1
                else max(self.vector_input_dim, self.vector_output_dim)
            )

            scalar_vector_frame_dim = (
                scalarization_vectorization_output_dim * 3
            )
            self.vector_down = nn.Linear(
                self.vector_input_dim, self.hidden_dim, bias=False
            )
            self.scalar_out = (
                nn.Sequential(
                    nn.Linear(
                        self.hidden_dim
                        + self.scalar_input_dim
                        + scalar_vector_frame_dim,
                        self.scalar_output_dim,
                    ),
                    get_activations(scalar_out_nonlinearity),
                    nn.Linear(self.scalar_output_dim, self.scalar_output_dim),
                )
                if feedforward_out
                else nn.Linear(
                    self.hidden_dim
                    + self.scalar_input_dim
                    + scalar_vector_frame_dim,
                    self.scalar_output_dim,
                )
            )

            self.vector_down_frames = nn.Linear(
                self.vector_input_dim,
                scalarization_vectorization_output_dim,
                bias=False,
            )

            if self.vector_output_dim:
                self.vector_up = nn.Linear(
                    self.hidden_dim, self.vector_output_dim, bias=False
                )
                if self.vector_gate:
                    self.vector_out_scale = nn.Linear(
                        self.scalar_output_dim, self.vector_output_dim
                    )
        else:
            self.scalar_out = (
                nn.Sequential(
                    nn.Linear(self.scalar_input_dim, self.scalar_output_dim),
                    get_activations(scalar_out_nonlinearity),
                    nn.Linear(self.scalar_output_dim, self.scalar_output_dim),
                )
                if feedforward_out
                else nn.Linear(self.scalar_input_dim, self.scalar_output_dim)
            )

    @jaxtyped
    @beartype
    def create_zero_vector(
        self,
        scalar_rep: Float[
            torch.Tensor, "batch_num_entities merged_scalar_dim"
        ],
    ) -> Float[torch.Tensor, "batch_num_entities o 3"]:
        return torch.zeros(
            scalar_rep.shape[0],
            self.vector_output_dim,
            3,
            device=scalar_rep.device,
        )

    @staticmethod
    @jaxtyped
    @beartype
    def scalarize(
        vector_rep: Float[torch.Tensor, "batch_num_entities 3 3"],
        edge_index: Int64[torch.Tensor, "2 batch_num_edges"],
        frames: Float[torch.Tensor, "batch_num_edges 3 3"],
        node_inputs: bool,
        enable_e3_equivariance: bool,
        dim_size: int,
        node_mask: Optional[Bool[torch.Tensor, "n_nodes"]] = None,
    ) -> Float[torch.Tensor, "effective_batch_num_entities 9"]:
        row, col = edge_index[0], edge_index[1]

        # gather source node features for each `entity` (i.e., node or edge)
        # note: edge inputs are already ordered according to source nodes
        vector_rep_i = vector_rep[row] if node_inputs else vector_rep

        # project equivariant values onto corresponding local frames
        if vector_rep_i.ndim == 2:
            vector_rep_i = vector_rep_i.unsqueeze(-1)
        elif vector_rep_i.ndim == 3:
            vector_rep_i = vector_rep_i.transpose(-1, -2)

        if node_mask is not None:
            edge_mask = node_mask[row] & node_mask[col]
            local_scalar_rep_i = torch.zeros(
                (edge_index.shape[1], 3, 3), device=edge_index.device
            )
            local_scalar_rep_i[edge_mask] = torch.matmul(
                frames[edge_mask], vector_rep_i[edge_mask]
            )
            local_scalar_rep_i = local_scalar_rep_i.transpose(-1, -2)
        else:
            local_scalar_rep_i = torch.matmul(frames, vector_rep_i).transpose(
                -1, -2
            )

        # potentially enable E(3)-equivariance and, thereby, chirality-invariance
        if enable_e3_equivariance:
            # avoid corrupting gradients with an in-place operation
            local_scalar_rep_i_copy = local_scalar_rep_i.clone()
            local_scalar_rep_i_copy[:, :, 1] = torch.abs(
                local_scalar_rep_i[:, :, 1]
            )
            local_scalar_rep_i = local_scalar_rep_i_copy

        # reshape frame-derived geometric scalars
        local_scalar_rep_i = local_scalar_rep_i.reshape(
            vector_rep_i.shape[0], 9
        )

        if node_inputs:
            # for node inputs, summarize all edge-wise geometric scalars using an average
            return torch_scatter.scatter(
                local_scalar_rep_i,
                # summarize according to source node indices due to the directional nature of GCP's equivariant frames
                row,
                dim=0,
                dim_size=dim_size,
                reduce="mean",
            )

        return local_scalar_rep_i

    @jaxtyped
    @beartype
    def vectorize(
        self,
        scalar_rep: Float[
            torch.Tensor, "batch_num_entities merged_scalar_dim"
        ],
        vector_hidden_rep: Float[torch.Tensor, "batch_num_entities 3 n"],
    ) -> Float[torch.Tensor, "batch_num_entities o 3"]:
        vector_rep = self.vector_up(vector_hidden_rep)
        vector_rep = vector_rep.transpose(-1, -2)

        if self.vector_gate:
            gate = self.vector_out_scale(self.vector_nonlinearity(scalar_rep))
            vector_rep = vector_rep * torch.sigmoid(gate).unsqueeze(-1)
        elif not is_identity(self.vector_nonlinearity):
            vector_rep = vector_rep * self.vector_nonlinearity(
                safe_norm(vector_rep, dim=-1, keepdim=True)
            )

        return vector_rep

    @jaxtyped
    @beartype
    def forward(
        self,
        s_maybe_v: Union[
            Tuple[
                Float[torch.Tensor, "batch_num_entities scalar_dim"],
                Float[torch.Tensor, "batch_num_entities m vector_dim"],
            ],
            Float[torch.Tensor, "batch_num_entities merged_scalar_dim"],
        ],
        edge_index: Int64[torch.Tensor, "2 batch_num_edges"],
        frames: Float[torch.Tensor, "batch_num_edges 3 3"],
        node_inputs: bool = False,
        node_mask: Optional[Bool[torch.Tensor, "batch_num_nodes"]] = None,
    ) -> Union[
        Tuple[
            Float[torch.Tensor, "batch_num_entities new_scalar_dim"],
            Float[torch.Tensor, "batch_num_entities n vector_dim"],
        ],
        Float[torch.Tensor, "batch_num_entities new_scalar_dim"],
    ]:
        if self.vector_input_dim:
            scalar_rep, vector_rep = s_maybe_v
            v_pre = vector_rep.transpose(-1, -2)

            vector_hidden_rep = self.vector_down(v_pre)
            vector_norm = safe_norm(vector_hidden_rep, dim=-2)
            merged = torch.cat((scalar_rep, vector_norm), dim=-1)

            # curate direction-robust and (by default) chirality-aware scalar geometric features
            vector_down_frames_hidden_rep = self.vector_down_frames(v_pre)
            scalar_hidden_rep = self.scalarize(
                vector_down_frames_hidden_rep.transpose(-1, -2),
                edge_index,
                frames,
                node_inputs=node_inputs,
                enable_e3_equivariance=self.enable_e3_equivariance,
                dim_size=vector_down_frames_hidden_rep.shape[0],
                node_mask=node_mask,
            )
            merged = torch.cat((merged, scalar_hidden_rep), dim=-1)
        else:
            # bypass updating scalar features using vector information
            merged = s_maybe_v

        scalar_rep = self.scalar_out(merged)

        if not self.vector_output_dim:
            # bypass updating vector features using scalar information
            return self.scalar_nonlinearity(scalar_rep)
        elif self.vector_output_dim and not self.vector_input_dim:
            # instantiate vector features that are learnable in proceeding GCP layers
            vector_rep = self.create_zero_vector(scalar_rep)
        else:
            # update vector features using either row-wise scalar gating with complete local frames or row-wise self-scalar gating
            vector_rep = self.vectorize(scalar_rep, vector_hidden_rep)

        scalar_rep = self.scalar_nonlinearity(scalar_rep)
        return ScalarVector(scalar_rep, vector_rep)

class GCPEmbedding(nn.Module):
    def __init__(
        self,
        edge_input_dims: ScalarVector,
        node_input_dims: ScalarVector,
        edge_hidden_dims: ScalarVector,
        node_hidden_dims: ScalarVector,
        num_atom_types: int = 0,
        nonlinearities: Tuple[Optional[str]] = ("silu", "silu"),
        cfg: DictConfig = None,
        pre_norm: bool = True,
        use_gcp_norm: bool = True,
    ):
        super().__init__()

        if num_atom_types > 0:
            self.atom_embedding = nn.Embedding(num_atom_types, num_atom_types)
        else:
            self.atom_embedding = None

        self.radial_embedding = partial(
            radial.compute_rbf, max_distance=cfg.r_max, num_rbf=cfg.num_rbf
        )

        self.pre_norm = pre_norm
        if pre_norm:
            self.edge_normalization = GCPLayerNorm(
                edge_input_dims, use_gcp_norm=use_gcp_norm
            )
            self.node_normalization = GCPLayerNorm(
                node_input_dims, use_gcp_norm=use_gcp_norm
            )
        else:
            self.edge_normalization = GCPLayerNorm(
                edge_hidden_dims, use_gcp_norm=use_gcp_norm
            )
            self.node_normalization = GCPLayerNorm(
                node_hidden_dims, use_gcp_norm=use_gcp_norm
            )

        self.edge_embedding = GCP(
            edge_input_dims,
            edge_hidden_dims,
            nonlinearities=nonlinearities,
            scalar_gate=cfg.scalar_gate,
            vector_gate=cfg.vector_gate,
            enable_e3_equivariance=cfg.enable_e3_equivariance,
        )

        self.node_embedding = GCP(
            node_input_dims,
            node_hidden_dims,
            nonlinearities=("none", "none"),
            scalar_gate=cfg.scalar_gate,
            vector_gate=cfg.vector_gate,
            enable_e3_equivariance=cfg.enable_e3_equivariance,
        )

    @jaxtyped
    @beartype
    def forward(
        self, batch: Union[Batch, ProteinBatch]
    ) -> Tuple[
        Union[
            Tuple[
                Float[torch.Tensor, "batch_num_nodes h_hidden_dim"],
                Float[torch.Tensor, "batch_num_nodes m chi_hidden_dim"],
            ],
            Float[torch.Tensor, "batch_num_nodes h_hidden_dim"],
        ],
        Union[
            Tuple[
                Float[torch.Tensor, "batch_num_edges e_hidden_dim"],
                Float[torch.Tensor, "batch_num_edges x xi_hidden_dim"],
            ],
            Float[torch.Tensor, "batch_num_edges e_hidden_dim"],
        ],
    ]:
        if self.atom_embedding is not None:
            node_rep = ScalarVector(self.atom_embedding(batch.h), batch.chi)
        else:
            node_rep = ScalarVector(batch.h, batch.chi)

        edge_rep = ScalarVector(batch.e, batch.xi)

        edge_vectors = (
            batch.pos[batch.edge_index[0]] - batch.pos[batch.edge_index[1]]
        )  # [n_edges, 3]
        edge_lengths = torch.linalg.norm(edge_vectors, dim=-1)  # [n_edges, 1]
        edge_rep = ScalarVector(
            torch.cat(
                (edge_rep.scalar, self.radial_embedding(edge_lengths)), dim=-1
            ),
            edge_rep.vector,
        )

        edge_rep = (
            edge_rep.scalar
            if not self.edge_embedding.vector_input_dim
            else edge_rep
        )
        node_rep = (
            node_rep.scalar
            if not self.node_embedding.vector_input_dim
            else node_rep
        )

        if self.pre_norm:
            edge_rep = self.edge_normalization(edge_rep)
            node_rep = self.node_normalization(node_rep)

        edge_rep = self.edge_embedding(
            edge_rep,
            batch.edge_index,
            batch.f_ij,
            node_inputs=False,
            node_mask=getattr(batch, "mask", None),
        )
        node_rep = self.node_embedding(
            node_rep,
            batch.edge_index,
            batch.f_ij,
            node_inputs=True,
            node_mask=getattr(batch, "mask", None),
        )

        if not self.pre_norm:
            edge_rep = self.edge_normalization(edge_rep)
            node_rep = self.node_normalization(node_rep)

        return node_rep, edge_rep

@beartype
def get_GCP_with_custom_cfg(
    input_dims: Any, output_dims: Any, cfg: DictConfig, **kwargs
):
    cfg_dict = copy(OmegaConf.to_container(cfg, throw_on_missing=True))
    cfg_dict["nonlinearities"] = cfg.nonlinearities
    del cfg_dict["scalar_nonlinearity"]
    del cfg_dict["vector_nonlinearity"]

    for key in kwargs:
        cfg_dict[key] = kwargs[key]

    return GCP(input_dims, output_dims, **cfg_dict)

class GCPMessagePassing(nn.Module):
    def __init__(
        self,
        input_dims: ScalarVector,
        output_dims: ScalarVector,
        edge_dims: ScalarVector,
        cfg: DictConfig,
        mp_cfg: DictConfig,
        reduce_function: str = "sum",
        use_scalar_message_attention: bool = True,
    ):
        super().__init__()

        # hyperparameters
        self.scalar_input_dim, self.vector_input_dim = input_dims
        self.scalar_output_dim, self.vector_output_dim = output_dims
        self.edge_scalar_dim, self.edge_vector_dim = edge_dims
        self.conv_cfg = mp_cfg
        self.self_message = self.conv_cfg.self_message
        self.reduce_function = reduce_function
        self.use_scalar_message_attention = use_scalar_message_attention

        scalars_in_dim = 2 * self.scalar_input_dim + self.edge_scalar_dim
        vectors_in_dim = 2 * self.vector_input_dim + self.edge_vector_dim

        # config instantiations
        soft_cfg = copy(cfg)
        soft_cfg.bottleneck = cfg.default_bottleneck

        primary_cfg_GCP = partial(get_GCP_with_custom_cfg, cfg=soft_cfg)
        secondary_cfg_GCP = partial(get_GCP_with_custom_cfg, cfg=cfg)

        # PyTorch modules #
        module_list = [
            primary_cfg_GCP(
                (scalars_in_dim, vectors_in_dim),
                output_dims,
                nonlinearities=cfg.nonlinearities,
                enable_e3_equivariance=cfg.enable_e3_equivariance,
            )
        ]

        for _ in range(self.conv_cfg.num_message_layers - 2):
            module_list.append(
                secondary_cfg_GCP(
                    output_dims,
                    output_dims,
                    enable_e3_equivariance=cfg.enable_e3_equivariance,
                )
            )

        if self.conv_cfg.num_message_layers > 1:
            module_list.append(
                primary_cfg_GCP(
                    output_dims,
                    output_dims,
                    nonlinearities=cfg.nonlinearities,
                    enable_e3_equivariance=cfg.enable_e3_equivariance,
                )
            )

        self.message_fusion = nn.ModuleList(module_list)

        # learnable scalar message gating
        if use_scalar_message_attention:
            self.scalar_message_attention = nn.Sequential(
                nn.Linear(output_dims.scalar, 1), nn.Sigmoid()
            )

    @jaxtyped
    @beartype
    def message(
        self,
        node_rep: ScalarVector,
        edge_rep: ScalarVector,
        edge_index: Int64[torch.Tensor, "2 batch_num_edges"],
        frames: Float[torch.Tensor, "batch_num_edges 3 3"],
        node_mask: Optional[Bool[torch.Tensor, "batch_num_nodes"]] = None,
    ) -> Float[torch.Tensor, "batch_num_edges message_dim"]:
        row, col = edge_index
        vector = node_rep.vector.reshape(
            node_rep.vector.shape[0],
            node_rep.vector.shape[1] * node_rep.vector.shape[2],
        )
        vector_reshaped = ScalarVector(node_rep.scalar, vector)

        s_row, v_row = vector_reshaped.idx(row)
        s_col, v_col = vector_reshaped.idx(col)

        v_row = v_row.reshape(v_row.shape[0], v_row.shape[1] // 3, 3)
        v_col = v_col.reshape(v_col.shape[0], v_col.shape[1] // 3, 3)

        message = ScalarVector(s_row, v_row).concat(
            (edge_rep, ScalarVector(s_col, v_col))
        )

        message_residual = self.message_fusion[0](
            message, edge_index, frames, node_inputs=False, node_mask=node_mask
        )
        for module in self.message_fusion[1:]:
            # exchange geometric messages while maintaining residual connection to original message
            new_message = module(
                message_residual,
                edge_index,
                frames,
                node_inputs=False,
                node_mask=node_mask,
            )
            message_residual = message_residual + new_message

        # learn to gate scalar messages
        if self.use_scalar_message_attention:
            message_residual_attn = self.scalar_message_attention(
                message_residual.scalar
            )
            message_residual = ScalarVector(
                message_residual.scalar * message_residual_attn,
                message_residual.vector,
            )

        return message_residual.flatten()

    @jaxtyped
    @beartype
    def aggregate(
        self,
        message: Float[torch.Tensor, "batch_num_edges message_dim"],
        edge_index: Int64[torch.Tensor, "2 batch_num_edges"],
        dim_size: int,
    ) -> Float[torch.Tensor, "batch_num_nodes aggregate_dim"]:
        row, col = edge_index
        aggregate = torch_scatter.scatter(
            message, row, dim=0, dim_size=dim_size, reduce=self.reduce_function
        )
        return aggregate

    @jaxtyped
    @beartype
    def forward(
        self,
        node_rep: ScalarVector,
        edge_rep: ScalarVector,
        edge_index: Int64[torch.Tensor, "2 batch_num_edges"],
        frames: Float[torch.Tensor, "batch_num_edges 3 3"],
        node_mask: Optional[Bool[torch.Tensor, "batch_num_nodes"]] = None,
    ) -> ScalarVector:
        message = self.message(
            node_rep, edge_rep, edge_index, frames, node_mask=node_mask
        )
        aggregate = self.aggregate(
            message, edge_index, dim_size=node_rep.scalar.shape[0]
        )
        return ScalarVector.recover(aggregate, self.vector_output_dim)

class GCPInteractions(nn.Module):
    def __init__(
        self,
        node_dims: ScalarVector,
        edge_dims: ScalarVector,
        cfg: DictConfig,
        layer_cfg: DictConfig,
        dropout: float = 0.0,
        nonlinearities: Optional[Tuple[Any, Any]] = None,
    ):
        super().__init__()

        # hyperparameters #
        if nonlinearities is None:
            nonlinearities = cfg.nonlinearities
        self.pre_norm = layer_cfg.pre_norm
        self.predict_node_positions = getattr(
            cfg, "predict_node_positions", False
        )
        self.node_positions_weight = getattr(cfg, "node_positions_weight", 1.0)
        self.update_positions_with_vector_sum = getattr(
            cfg, "update_positions_with_vector_sum", False
        )
        reduce_function = "sum"

        # PyTorch modules #

        # geometry-complete message-passing neural network
        message_function = GCPMessagePassing

        self.interaction = message_function(
            node_dims,
            node_dims,
            edge_dims,
            cfg=cfg,
            mp_cfg=layer_cfg.mp_cfg,
            reduce_function=reduce_function,
            use_scalar_message_attention=layer_cfg.use_scalar_message_attention,
        )

        # config instantiations
        ff_cfg = copy(cfg)
        ff_cfg.nonlinearities = nonlinearities
        ff_GCP = partial(get_GCP_with_custom_cfg, cfg=ff_cfg)

        self.gcp_norm = nn.ModuleList(
            [GCPLayerNorm(node_dims, use_gcp_norm=layer_cfg.use_gcp_norm)]
        )
        self.gcp_dropout = nn.ModuleList(
            [GCPDropout(dropout, use_gcp_dropout=layer_cfg.use_gcp_dropout)]
        )

        # build out feedforward (FF) network modules
        hidden_dims = (
            (node_dims.scalar, node_dims.vector)
            if layer_cfg.num_feedforward_layers == 1
            else (4 * node_dims.scalar, 2 * node_dims.vector)
        )
        ff_interaction_layers = [
            ff_GCP(
                (node_dims.scalar * 2, node_dims.vector * 2),
                hidden_dims,
                nonlinearities=("none", "none")
                if layer_cfg.num_feedforward_layers == 1
                else cfg.nonlinearities,
                feedforward_out=layer_cfg.num_feedforward_layers == 1,
                enable_e3_equivariance=cfg.enable_e3_equivariance,
            )
        ]

        interaction_layers = [
            ff_GCP(
                hidden_dims,
                hidden_dims,
                enable_e3_equivariance=cfg.enable_e3_equivariance,
            )
            for _ in range(layer_cfg.num_feedforward_layers - 2)
        ]
        ff_interaction_layers.extend(interaction_layers)

        if layer_cfg.num_feedforward_layers > 1:
            ff_interaction_layers.append(
                ff_GCP(
                    hidden_dims,
                    node_dims,
                    nonlinearities=("none", "none"),
                    feedforward_out=True,
                    enable_e3_equivariance=cfg.enable_e3_equivariance,
                )
            )

        self.feedforward_network = nn.ModuleList(ff_interaction_layers)

        # potentially build out node position update modules
        if self.predict_node_positions:
            # node position update GCP(s)
            position_output_dims = (
                node_dims
                if getattr(cfg, "update_positions_with_vector_sum", False)
                else (node_dims.scalar, 1)
            )
            self.node_position_update_gcp = ff_GCP(
                node_dims,
                position_output_dims,
                nonlinearities=cfg.nonlinearities,
                enable_e3_equivariance=cfg.enable_e3_equivariance,
            )

    @jaxtyped
    @beartype
    def derive_x_update(
        self,
        node_rep: ScalarVector,
        edge_index: Int64[torch.Tensor, "2 batch_num_edges"],
        f_ij: Float[torch.Tensor, "batch_num_edges 3 3"],
        node_mask: Optional[Bool[torch.Tensor, "batch_num_nodes"]] = None,
    ) -> Float[torch.Tensor, "batch_num_nodes 3"]:
        # use vector-valued features to derive node position updates
        node_rep_update = self.node_position_update_gcp(
            node_rep, edge_index, f_ij, node_inputs=True, node_mask=node_mask
        )
        if self.update_positions_with_vector_sum:
            x_vector_update = node_rep_update.vector.sum(1)
        else:
            x_vector_update = node_rep_update.vector.squeeze(1)

        # (up/down)weight position updates
        x_update = x_vector_update * self.node_positions_weight

        return x_update

    @jaxtyped
    @beartype
    def forward(
        self,
        node_rep: Tuple[
            Float[torch.Tensor, "batch_num_nodes node_hidden_dim"],
            Float[torch.Tensor, "batch_num_nodes m 3"],
        ],
        edge_rep: Tuple[
            Float[torch.Tensor, "batch_num_edges edge_hidden_dim"],
            Float[torch.Tensor, "batch_num_edges x 3"],
        ],
        edge_index: Int64[torch.Tensor, "2 batch_num_edges"],
        frames: Float[torch.Tensor, "batch_num_edges 3 3"],
        node_mask: Optional[Bool[torch.Tensor, "batch_num_nodes"]] = None,
        node_pos: Optional[Float[torch.Tensor, "batch_num_nodes 3"]] = None,
    ) -> Tuple[
        Tuple[
            Float[torch.Tensor, "batch_num_nodes hidden_dim"],
            Float[torch.Tensor, "batch_num_nodes n 3"],
        ],
        Optional[Float[torch.Tensor, "batch_num_nodes 3"]],
    ]:
        node_rep = ScalarVector(node_rep[0], node_rep[1])
        edge_rep = ScalarVector(edge_rep[0], edge_rep[1])

        # apply GCP normalization (1)
        if self.pre_norm:
            node_rep = self.gcp_norm[0](node_rep)

        # forward propagate with interaction module
        hidden_residual = self.interaction(
            node_rep, edge_rep, edge_index, frames, node_mask=node_mask
        )

        # aggregate input and hidden features
        hidden_residual = ScalarVector(*hidden_residual.concat((node_rep,)))

        # propagate with feedforward layers
        for module in self.feedforward_network:
            hidden_residual = module(
                hidden_residual,
                edge_index,
                frames,
                node_inputs=True,
                node_mask=node_mask,
            )

        # apply GCP dropout
        node_rep = node_rep + self.gcp_dropout[0](hidden_residual)

        # apply GCP normalization (2)
        if not self.pre_norm:
            node_rep = self.gcp_norm[0](node_rep)

        # update only unmasked node representations and residuals
        if node_mask is not None:
            node_rep = node_rep.mask(node_mask.float())

        # bypass updating node positions
        if not self.predict_node_positions:
            return node_rep, node_pos

        # update node positions
        node_pos = node_pos + self.derive_x_update(
            node_rep, edge_index, frames, node_mask=node_mask
        )

        # update only unmasked node positions
        if node_mask is not None:
            node_pos = node_pos * node_mask.float().unsqueeze(-1)

        return node_rep, node_pos


# -*- coding: utf-8 -*-
"""
Created on Sat Jul 13 20:37:47 2013

@author: shanef
"""

import unittest
import os
import sys
parentdir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, parentdir)
import TVEncoder

class TVEncoderTest(unittest.TestCase):
    def test_processarguments_encodereadonly(self):
        args = []
        args.append(('-e', ''))
        args.append(('-l', ''))
        result = TVEncoder.processarguments(args)

        self.assertTrue(result.doencode)
        self.assertTrue(result.readonly)

    def test_processarguments_encodereadonlyreverse(self):
        args = []
        args.append(('-l', ''))
        args.append(('-e', ''))
        result = TVEncoder.processarguments(args)

        self.assertTrue(result.doencode)
        self.assertTrue(result.readonly)

    def test_processarguments_encode(self):
        args = []
        args.append(('-e', ''))
        result = TVEncoder.processarguments(args)

        self.assertTrue(result.doencode)
        self.assertFalse(result.readonly)

if __name__ == '__main__':
    suite = unittest.TestLoader().loadTestsFromTestCase(TVEncoderTest)
    unittest.TextTestRunner(verbosity=2).run(suite)


import responses
from responses import matchers

@responses.activate
def test_create_token(zpa, session):
    responses.add(
        responses.POST,
        url="https://config.private.zscaler.com/signin",
        json=session,
        status=200,
        match=[
            matchers.urlencoded_params_matcher(
                {
                    "client_id": "1",
                    "client_secret": "yyy",
                }
            ),
            matchers.header_matcher({"Content-Type": "application/x-www-form-urlencoded"}),
        ],
    )

    resp = zpa.session.create_token(client_id="1", client_secret="yyy")

    assert isinstance(resp, str)
    assert resp == "xyz"


#! /env/python
import numpy as np
import matplotlib.pyplot as plt
import sys

from matplotlib.ticker import ScalarFormatter,AutoMinorLocator

plt.rc('font', family='serif', serif='STIXGeneral')

data = {
  'distance_from_centre': {
    'run1': {
      'meas':[],
      'max': [],
      'min': [],
      'color': 'black' # good
    },
    'run2': {
      'meas':[],
      'max': [],
      'min': [],
      'color': 'green'
    },
    'run3': {
      'meas':[],
      'max': [],
      'min': [],
      'color': 'red'
    },
    'run4': {
      'meas':[],
      'max': [],
      'min': [],
      'color': '#FF9955'
    },
    'run5': {
      'meas':[],
      'max': [],
      'min': [],
      'color': 'blue'
    }
  },
  'distance_from_average': {
    'run1': {
      'meas':[],
      'max': [],
      'min': [],
    },
    'run2': {
      'meas':[],
      'max': [],
      'min': [],
    },
    'run3': {
      'meas':[],
      'max': [],
      'min': [],
    },
    'run4': {
      'meas':[],
      'max': [],
      'min': [],
    },
    'run5': {
      'meas':[],
      'max': [],
      'min': [],
    }
  },
}
block = 0
with open('tracking.csv', 'r') as f:
  for row, line in enumerate(f.readlines()):
    val_count = 0
    g_type = 'distance_from_centre'
    for col, val in enumerate(line.split(',')):
      if val != "" and val != '\n':
        val_count += 1
        try:
          num = float(val)
        except ValueError:
          print("got", val)
          num
      else:
        num = np.nan

      run = col
      if col == 5 or col > 10:
        continue
      if col > 4:
        g_type = 'distance_from_average'
        run = col - 6
      # print(block, run, num)
      if block == 0:
        cat = 'meas'
      elif block == 1:
        cat = 'max'
      elif block == 2:
        cat = 'min'

      data[g_type]['run'+str(run+1)][cat].append(num)

    if val_count == 0:
      block += 1

# print(data)

# sys.exit()

dimensions = [390*0.01384*1.0,195*0.01384]
params = {'backend': 'ps',
          'axes.labelsize': 10,
          'text.fontsize': 10,
          'legend.fontsize': 8,
          'axes.linewidth' : 0.5,
          'xtick.labelsize': 10,
          'ytick.labelsize': 10,
          'axes.formatter.limits' : '-3, 3',
          'figure.figsize': dimensions,
          'text.usetex'        : True,
          'text.latex.unicode' : True,
          'lines.markersize' : 4}
plt.rcParams.update(params)

fig, ax = plt.subplots(1)

xs = range(1,len(data['distance_from_average']['run1']['meas'])+1)
ys = data['distance_from_average']['run1']['meas']

errs = []

for item_num in range(40):
  print(item_num)
  errs.append(data['distance_from_centre']['run1']['max'][item_num] - data['distance_from_centre']['run1']['min'][item_num])

print(errs)

for run in range(1,6):
  print(run)
  # ax.plot(xs,data['distance_from_average']['run'+str(run)]['meas'])
  ax.errorbar(xs,list(map(lambda x: x/1000, data['distance_from_centre']['run'+str(run)]['meas'])), xerr=None, yerr=0.546703/2, barsabove=True, fmt="-x", capsize=2, color=data['distance_from_centre']['run'+str(run)]['color'])
# ax.bar(range(1, len(datasource)+1), datasource, color='b', edgecolor="none", align="center", width=0.7)
ax.set_xlim(0, 41)
# ax.set_ylim(-4.370/2,4.370/2)
ax.set_ylim(-1.0, 1.0)

# plt.yticks(range(11))
# plt.xticks(range(len(datasource)+1))
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.yaxis.set_ticks_position('left')
ax.xaxis.set_ticks_position('bottom')
leg = ax.legend(scatterpoints=1, fancybox=True, loc=4)
# leg.get_frame().set_linewidth(0.0)

ax = plt.gca()
plt.grid()
plt.ylabel('Displacement from \npost/trunk mid-point (m)')
plt.xlabel('Post/trunk pair')
plt.tight_layout()
plt.savefig("row_tracking_averages.pdf")

fig.clf()


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import reco as rc

def imm(mod,sens,time):
    for i in range(len(mod)):
        ril[i]=rc.Hit(mod[i],sens[i],time [i])
    return ril

data0=pd.read_csv('hit_times_M0.csv')
data1=pd.read_csv('hit_times_M1.csv')
data2=pd.read_csv('hit_times_M2.csv')
data3=pd.read_csv('hit_times_M3.csv')

print (data0)

'''plt.hist(data0['hit_time'],bins=50)
plt.show()

diff=np.diff(data0['hit_time'].values)
pippo=diff>0
difflog=np.log10(diff[pippo])

plt.hist(difflog,bins=50)
plt.show()'''

#mod_id  det_id   hit_time

ril0=imm(data0['mod_id'].values,data0['det_id'].values,data0['hit_time'].values)
ril1=imm(data1['mod_id'].values,data1['det_id'].values,data1['hit_time'].values)
ril2=imm(data2['mod_id'].values,data2['det_id'].values,data2['hit_time'].values)
ril3=imm(data3['mod_id'].values,data3['det_id'].values,data3['hit_time'].values)

hits=np.append(ril0,ril1)
hits=np.append(hits,ril2)
hits=np.append(hits,ril3)

np.sort(hits)
diffhits=np.diff(hits)
pluto=diffhits>0
loghits=np.log10(diffhits[pluto])

plt.hist(loghits,bins=50)

plt.show()


import math
import numpy as np
import torch
from collections import deque, namedtuple

device = torch.device('cpu')

def rp_set_device(new_device):
    global device
    device = new_device

class ReplayBuffer:
    """Fixed-size buffer to store experience tuples."""

    def __init__(self, buffer_size, batch_size, n_multisteps, gamma, a, separate_experiences):
        """Initialize a ReplayBuffer object.

        Params
        ======
            buffer_size (int): maximum size of buffer
            batch_size (int): size of each training batch
            n_multisteps (int): number of time steps to consider for each experience
            gamma (float): discount factor
            a (float): priority exponent parameter
            separate_experiences (bool): whether to store experiences with no overlap
        """
        self.buffer_size = buffer_size
        self.batch_size = batch_size
        self.n_multisteps = n_multisteps
        self.gamma = gamma
        self.a = a
        self.separate_experiences = bool(separate_experiences)

        self.memory_write_idx = 0
        self._non_leaf_depth = math.ceil(math.log2(buffer_size))
        self._memory_start_idx = 2 ** self._non_leaf_depth
        self._buffer_is_full = False
        self.memory = [None for _ in range(buffer_size)]
        self.priorities_a = np.zeros(buffer_size)
        self.tree = np.zeros(self._memory_start_idx + buffer_size) # starts from index 1, not 0; makes implementation easier and reduces many small computations

        self.multistep_collector = deque(maxlen=n_multisteps)
        self.max_priority_a = 1.
        self.experience = namedtuple("Experience", field_names=["state", "action", "reward", "next_state", "done"])

        self._divisors = np.power(2, np.arange(1, self._non_leaf_depth + 1)).reshape((-1, 1))
        self._discounts = np.power(self.gamma, np.arange(self.n_multisteps + 1))
        self._target_discount = float(self._discounts[-1])

    def add(self, state, action, reward, next_state, done):
        """Add a new experience to memory."""
        e = self.experience(state, action, reward, next_state, done)
        self.multistep_collector.append(e)
        if len(self.multistep_collector) == self.n_multisteps:
            self.memory[self.memory_write_idx] = tuple(self.multistep_collector)

            delta_priority_a = self.max_priority_a - self.priorities_a[self.memory_write_idx]
            tree_idx = self._memory_start_idx + self.memory_write_idx
            self.priorities_a[self.memory_write_idx] = self.max_priority_a
            self.tree[tree_idx] = self.max_priority_a
            # tree_indices = np.floor_divide(tree_idx, self._divisors).reshape((-1,))
            # np.add.at(self.tree, tree_indices, np.tile(delta_priority_a, self._non_leaf_depth))
            for _ in range(self._non_leaf_depth):
                tree_idx = tree_idx // 2
                self.tree[tree_idx] += delta_priority_a

            self.memory_write_idx += 1
            if self.memory_write_idx >= self.buffer_size:
                self._buffer_is_full = True
                self.memory_write_idx = 0

            if self.separate_experiences:
                self.multistep_collector.clear()
        if done:
            self.multistep_collector.clear()

    def sample(self, beta):
        """Randomly sample a batch of experiences from memory.

        Params
        ======
            beta (int or float): parameter used for calculating importance-priority weights

        Returns
        =======
            experiences (Tuple[torch.Tensor]): tuple of (s, a, r, s', done) tuples
            target_discount (float): discount factor for target max-Q value
            is_weights (torch.Tensor): tensor of importance-sampling weights
            indices (np.ndarray): sample indices"""
        indices_not_prepared = True
        sample_value_basis = np.linspace(0, self.tree[1], num=self.batch_size, endpoint=False, dtype=np.float32)
        while indices_not_prepared:
            sample_values = np.add(sample_value_basis, np.multiply(np.random.rand(self.batch_size), sample_value_basis[1]))
            tree_indices = np.ones(self.batch_size, dtype=np.int32)
            try:
                for _ in range(self._non_leaf_depth):
                    left_child_indices = np.multiply(tree_indices, 2)
                    right_child_indices = np.add(left_child_indices, 1)
                    greater_than_left = np.greater(sample_values, self.tree[left_child_indices])
                    sample_values = np.where(greater_than_left, np.subtract(sample_values, self.tree[left_child_indices]), sample_values)
                    tree_indices = np.where(greater_than_left, right_child_indices, left_child_indices)
            except IndexError: # Don't know exactly why it occurs. Suspecting numerical error issues with floating numbers as a probable cause
                continue
            else:
                indices_not_prepared = False
        indices = np.subtract(tree_indices, self._memory_start_idx)

        experiences = tuple(zip(*[self.memory[i] for i in indices if self.memory[i] is not None]))

        first_states = torch.tensor([e[0] for e in experiences[0]], dtype=torch.float, device=device)
        actions      = torch.tensor([e[1] for e in experiences[0]], dtype=torch.long, device=device)
        rewards      = torch.tensor(
                           np.sum(
                               np.multiply(
                                   np.array([[e[2] for e in experiences_step] for experiences_step in experiences]).transpose(), self._discounts[:-1]
                               ), axis=1, keepdims=True
                           ), dtype=torch.float, device=device)
        last_states  = torch.tensor([e[3] for e in experiences[-1]], dtype=torch.float, device=device)
        dones        = torch.tensor([e[4] for e in experiences[-1]], dtype=torch.float, device=device).view(-1, 1)

        is_weights = np.divide(self.priorities_a[indices], self.tree[1])
        is_weights = np.power(np.multiply(is_weights, self.buffer_size if self._buffer_is_full else self.memory_write_idx), -beta)
        is_weights = torch.tensor(np.divide(is_weights, np.max(is_weights)).reshape((-1, 1)), dtype=torch.float, device=device)

        return (first_states, actions, rewards, last_states, dones), self._target_discount, is_weights, indices

    def update_priorities(self, indices, new_priorities):
        """Update the priorities for the experiences of given indices to the given new values.

        Params
        ======
            indices (array_like): indices of experience priorities to update
            new_priorities (array-like): new priority values for given indices"""
        # Remove Duplicate Samples (discard except the first occurence in the array)
        indices, idx_indices = np.unique(indices, return_index=True)
        new_priorities_a = np.power(new_priorities[idx_indices], self.a)

        delta_priority_a = np.subtract(new_priorities_a, self.priorities_a[indices])
        tree_indices = np.add(indices, self._memory_start_idx)
        self.priorities_a[indices] = new_priorities_a
        self.tree[tree_indices] = new_priorities_a
        tree_indices = np.floor_divide(tree_indices, self._divisors).reshape((-1,))
        np.add.at(self.tree, tree_indices, np.tile(delta_priority_a, self._non_leaf_depth))

        self.max_priority_a = np.max(self.priorities_a)

    def reset_multisteps(self):
        self.multistep_collector.clear()

    def __len__(self):
        """Return the current size of internal memory."""
        return self.buffer_size if self._buffer_is_full else self.memory_write_idx


from django.urls import path
from rest_framework import routers
from . import views

# Wire up our API using automatic URL routing.
# Additionally, we include login URLs for the browsable API.
urlpatterns = [
   path('fincas', views.get_fincas),
   path('fincas2', views.get_fincas_2),
   path('finca/<int:pk>', views.finca),

   path('farms_new', views.fincas),
   path('farm_new/create', views.fincasCrear),
   path('farm_new/<int:pk>', views.finca_new)
]
router = routers.DefaultRouter()
#router.register('fincas',views.FincaViewSet,'fincas')
urlpatterns = urlpatterns + router.urls


#!/usr/bin/env python
import bottle
import subprocess
import os

p1 = subprocess.Popen(['ip','addr','show','eth0'],stdout=subprocess.PIPE)
p2 = subprocess.Popen(['sed','-rn',r's/\s*inet\s(([0-9]{1,3}\.){3}[0-9]{1,3}).*/\1/p'],stdin=p1.stdout,stdout=subprocess.PIPE)
p1.stdout.close()
ip_addr = p2.communicate()[0].strip()
p1.wait()

app = bottle.app()

@bottle.route('/')
def root_index():
    return bottle.template('index',ip_addr = ip_addr)

@bottle.route('/json')
def json_reply():
    heads = bottle.request.headers
    bottle.response.content_type = 'application/json'

    response = {'headers':dict(heads),
            'environment':dict(os.environ),
            'response':dict(bottle.response.headers)}
    return response

if __name__=='__main__':
    bottle.debug(True)
    bottle.run(app=app,host='localhost',port=8080)


import torch
from torch.utils.data import Dataset
from .tokenization import *

class NestedNamedEntitiesDataset(Dataset):
    def __init__(self, data, lm_input, word_input=None, char_input=None, padding_length=512,
                 total_layers=16, skip_exceptions=True, max_items=-1, has_outputs=True):
        self.total_layers = total_layers
        self.has_outputs = has_outputs

        self.item_id = []
        self.texts = []
        self.token_offsets = []
        self.masks = []
        self.X_word = []
        self.X_char = []
        self.X_lm_inputs = []
        self.X_lm_attention = []
        self.X_lm_type_ids = []
        self.X_lm_spans = []
        self.Y_entities = [[] for _ in range(total_layers)]
        self.n_skipped = 0

        for i, item in enumerate(data):
            if max_items > 0 and i >= max_items:
                break

            try:
                x_word = word_input.encode(item['tokens'], padding_length) if word_input is not None else []
                x_char = char_input.encode(item['tokens'], padding_length) if char_input is not None else []
                x_lm, x_lm_span = lm_input.encode(item['tokens'], padding_length)

                mask = [1.] * len(item['tokens']) + [0.] * (padding_length - len(item['tokens']))
                self.masks.append(mask)

                token_offsets = get_offsets_from_tokens(item['text'], item['tokens'])
                self.token_offsets.append(token_offsets)

                if 'item_id' in item:
                    self.item_id.append(item['item_id'])
                else:
                    self.item_id.append(i)

                self.texts.append(item['text'])
                self.X_word.append(x_word)
                self.X_char.append(x_char)

                self.X_lm_inputs.append(x_lm[0])
                self.X_lm_attention.append(x_lm[1])
                self.X_lm_type_ids.append(x_lm[2])
                self.X_lm_spans.append(x_lm_span)

                if has_outputs:
                    for i_layer, raw_seq in enumerate(item['layer_outputs']):
                        if i_layer >= total_layers:
                            break
                        padded_seq = raw_seq if len(raw_seq) > 0 else []
                        if len(padded_seq) < padding_length - i_layer:
                            padded_seq = padded_seq + [0] * (padding_length - len(padded_seq) - i_layer)
                        self.Y_entities[i_layer].append(padded_seq)
            except Exception as e:
                # Text-too-long exception
                if skip_exceptions:
                    self.n_skipped += 1
                    continue
                else:
                    raise e

        self._init_tensors(padding_length, total_layers)

    def _init_tensors(self, padding_length, total_layers):
        self.masks = torch.tensor(self.masks, dtype=torch.float)
        self.X_word = torch.tensor(self.X_word, dtype=torch.long)
        self.X_char = torch.tensor(self.X_char, dtype=torch.long)
        self.X_lm_inputs = torch.tensor(self.X_lm_inputs, dtype=torch.long)
        self.X_lm_attention = torch.tensor(self.X_lm_attention, dtype=torch.long)
        self.X_lm_type_ids = torch.tensor(self.X_lm_type_ids, dtype=torch.long)
        self.X_lm_spans = torch.tensor(self.X_lm_spans, dtype=torch.long)
        if self.has_outputs:
            for i, seqs in enumerate(self.Y_entities):
                self.Y_entities[i] = torch.tensor(seqs, dtype=torch.long)

    def __len__(self):
        return len(self.X_word)

    def __getitem__(self, idx):
        data = {
            'masks': self.masks[idx],
            'x_word': self.X_word[idx],
            'x_char': self.X_char[idx],
            'x_lm_input': self.X_lm_inputs[idx],
            'x_lm_attention': self.X_lm_attention[idx],
            'x_lm_type_ids': self.X_lm_type_ids[idx],
            'x_lm_spans': self.X_lm_spans[idx],
        }
        if self.has_outputs:
            for i_layer in range(self.total_layers):
                data['y_target_%d' % i_layer] = self.Y_entities[i_layer][idx]
        return data

    def get_item(self, idx):
        return self.__getitem__(idx)

    def get_item_id(self, idx):
        return self.item_id[idx]

    def get_texts(self, idx):
        return self.texts[idx]

    def get_token_offsets(self, idx):
        return self.token_offsets[idx]

    def get_n_skipped(self):
        return self.n_skipped


def merge_sort(list):
    if len(list) <= 1:
        return list
    # print(list)
    middle = len(list) // 2
    first = merge_sort(list[:middle])
    second = merge_sort(list[middle:])
    return merge_arrays(first, second)

def merge_arrays(leftlist, rightlist):                  # first = leftlist and second = rightlist
    merged_array = []
    l = r = 0                                           # INDEX for leftlist = lf  and   rightlist = rt
    while l < len(leftlist) or r < len(rightlist):

        if l == len(leftlist):                          # IF LEFT IS NOT THERE, (because l = 0 up there) then add r to rigtlist.
            merged_array.append(rightlist[r])           # we add the value
            r += 1                                      # ADDS +1 to index, and moves to next index.
            continue                                    # continues back to while

        if r == len(rightlist):                         # IF RIGHT IS NOT THERE (because r = 0 up there) then add l to leftlist.
            merged_array.append(leftlist[l])
            l += 1
            continue

        if leftlist[l] <= rightlist[r]:                 # Q: By value; if l in leftlist (smaller than / equal to) r in rightlist
            merged_array.append(leftlist[l])            # add leftlist into merged array
            l += 1

        else:
            merged_array.append(rightlist[r])
            r += 1

    return merged_array

# la = [1, 7]
# ra = [2, 5]
# combined = merge_arrays(la, ra)
# print('merge arrays, combined:', combined)

test_list = [1, 4, 3, 5, 2, 6, 7, 8]
result_tl = merge_sort(test_list)
print('merge_sort:',result_tl)


import time
import os
import json
import math
import cv2
import numpy as np

from scipy import ndimage
from glob import glob
from config import inference_img_inp_dir, inference_preprocess_dir,\
    inference_extract_boxes_path, inference_angle_path,\
    inference_rotated_boxes_txt_dir, inference_rotated_img_dir
from craft_text_detector import (
    load_craftnet_model,
    load_refinenet_model,
    get_prediction,
    empty_cuda_cache
)

def extract_boxes(input_imgs_dir, output_json_path):
    img_dir = input_imgs_dir
    output_boxes_json = output_json_path

    img_boxes = {}
    try:
        with open(output_boxes_json, 'r+') as out:
            saved_output = json.load(out)
    except IOError:
        saved_output = {}

    refine_net = load_refinenet_model(cuda=True)
    craft_net = load_craftnet_model(cuda=True)

    start_time_alpha = time.time()
    for img_path in glob(img_dir + '/*'):
        print(img_path)
        img_id = os.path.basename(img_path)
        if img_id in saved_output:
            img_boxes[img_id] = saved_output[img_id]
            continue

        # perform prediction
        start_time = time.time()
        prediction_result = get_prediction(
            image=img_path,
            craft_net=craft_net,
            refine_net=refine_net,
            text_threshold=0.7,
            link_threshold=0.4,
            low_text=0.4,
            cuda=True,
            long_size=1280
        )
        print(img_id, time.time() - start_time)

        boxes = prediction_result['boxes']
        boxes_as_ratios = prediction_result['boxes_as_ratios']
        boxes_dict = {
            'boxes': boxes if type(boxes) is list else boxes.tolist(),
            'boxes_as_ratios': boxes_as_ratios if type(boxes_as_ratios) is list else boxes_as_ratios.tolist(),
        }

        img_boxes[img_id] = boxes_dict
    print('Total time:', time.time() - start_time_alpha)

    with open(output_boxes_json, 'w') as out:
        json.dump(img_boxes, out)

def is_horizontal(boxes: list) -> bool:
    edge1_sum = 0
    edge2_sum = 0
    for box in boxes:
        # box chứa 4 list con (ví dụ [x1, y1]) là 4 đỉnh của box
        edge1_len = math.dist(box[0], box[1])
        edge2_len = math.dist(box[1], box[2])
        edge1_sum += edge1_len
        edge2_sum += edge2_len

    if edge1_sum > edge2_sum:
        return True
    return False

def get_img_angles(input_json_path, output_json_path):
    input_path = input_json_path
    output_path = output_json_path

    with open(input_path) as inp:
        img_boxes = json.load(inp)

    img_angles = {}
    for img_id, img in img_boxes.items():

        # xác định tọa độ các cạnh dài
        if is_horizontal(img['boxes']):
            print(img_id, 'is vertical.')
            long_edges = [(coord[0], coord[1]) for coord in img['boxes']]
        else:
            print(img_id, 'is horizontal.')
            long_edges = [(coord[1], coord[2]) for coord in img['boxes']]

        # tính trung bình góc xoay
        angle_sum = 0
        for (x1, y1), (x2, y2) in long_edges:
            angle = math.degrees(math.atan2(y2 - y1, x2 - x1))
            angle_sum += angle

        rotation_angle = angle_sum / len(long_edges) if len(long_edges) > 0 else 0
        print('Avg angle:', rotation_angle)
        img_angles[img_id] = rotation_angle

    # xuất ra json file
    with open(output_path, 'w') as out:
        json.dump(img_angles, out)

def write_output_craft(boxes, result_file_path):
    result = ''
    for box in boxes:
        line = ','.join((map(str, box)))
        result += line + '\n'  # 100, 200, ..., 123\n
    result = result.rstrip('\n')
    with open(result_file_path, 'w', encoding='utf8') as res:
        res.write(result)

def rotate_img(angle_input_path, img_inupt_dir, craft_boxes_json_path, img_output_dir, craft_txt_out_dir):

    with open(angle_input_path) as inp:
        img_angles_dict = json.load(inp)
    with open(craft_boxes_json_path) as inp:
        craft_boxes_dict = json.load(inp)

    for img_id, angle in img_angles_dict.items():
        try:
            craft_boxes = craft_boxes_dict[img_id]['boxes']
        except KeyError:
            continue

        # đọc ảnh, tính các thông số để xoay ảnh
        img_input_path = img_inupt_dir + '/' + img_id
        img = cv2.imread(img_input_path)
        h, w = img.shape[:2]
        cx, cy = (int(w / 2), int(h / 2))

        # xoay ảnh
        img_rotated = ndimage.rotate(img, angle)
        img_output_path = img_output_dir + '/' + img_id
        cv2.imwrite(img_output_path, img_rotated)

        # xoay craft boxes
        rotated_craft_boxes = []
        for box in craft_boxes:
            rotated_craft_box = []
            for coord in box:
                M = cv2.getRotationMatrix2D((cx, cy), angle, 1.0)
                cos, sin = abs(M[0, 0]), abs(M[0, 1])
                newW = int((h * sin) + (w * cos))
                newH = int((h * cos) + (w * sin))
                M[0, 2] += (newW / 2) - cx
                M[1, 2] += (newH / 2) - cy
                v = [coord[0], coord[1], 1]
                adjusted_coord = np.dot(M, v)
                rotated_craft_box.append((adjusted_coord[0], adjusted_coord[1]))
            rotated_craft_box = [int(x) for t in rotated_craft_box for x in t]
            rotated_craft_boxes.append(rotated_craft_box)

        # xuất txt
        write_output_craft(rotated_craft_boxes, craft_txt_out_dir + '/' + img_id[:-4] + '.txt')
        print(img_id, 'rotated', len(rotated_craft_boxes), 'boxes')

if __name__ == '__main__':
    img_dir = inference_img_inp_dir
    craft_boxes_json = inference_extract_boxes_path
    img_angle_json = inference_angle_path
    rotate_img_dir = inference_rotated_img_dir
    rotated_craft_boxes_dir = inference_rotated_boxes_txt_dir

    extract_boxes(img_dir, craft_boxes_json)
    get_img_angles(craft_boxes_json, img_angle_json)
    rotate_img(img_angle_json, img_dir, craft_boxes_json,
               rotate_img_dir, rotated_craft_boxes_dir)


import matplotlib.pyplot as plt

f = open("dataset/twitter_1000000.txt", "r")

xs = []
ys = []
for l in f:
    data = (l.split("\t")[1]).split(" ")
    xs.append(float(data[0]))
    ys.append(float(data[1]))

plt.scatter(xs,ys,s=1)

plt.title('Dataset preview')
plt.xlabel('x')
plt.ylabel('y')

plt.show()


import sklearn.ensemble

from autosklearn.pipeline.components.regression.adaboost import AdaboostRegressor

from .test_base import BaseRegressionComponentTest

class AdaBoostComponentTest(BaseRegressionComponentTest):

    __test__ = True

    res = dict()
    res["default_boston"] = 0.5951486466070626
    res["default_boston_iterative"] = None
    res["default_boston_sparse"] = 0.18067558132702222
    res["default_boston_iterative_sparse"] = None
    res["default_diabetes"] = 0.250565253614339
    res["default_diabetes_iterative"] = None
    res["default_diabetes_sparse"] = 0.09126705185668416
    res["default_diabetes_iterative_sparse"] = None

    sk_mod = sklearn.ensemble.AdaBoostRegressor

    module = AdaboostRegressor


def reverse_words(s):
    """
    Write a reverseWords function that accepts a string a parameter,
    and reverses each word in the string. Any spaces in the string should be retained.

    Example:

    reverse_words("This is an example!") # returns  "sihT si na !elpmaxe"
    reverse_words("double  spaces") # returns  "elbuod  secaps"
    :param s: a string
    :return: s reversed
    """

    split_s = s.split()
    new_string = ""
    counter = 0
    for i in range(0,len(split_s)):
        while s[counter] == " ":
            new_string += " "
            counter += 1
        current_string = list(split_s[i])
        current_string.reverse()
        for char in current_string:
            new_string += char
        counter += len(current_string)

    return new_string

print(reverse_words("double  spaces"))


"""
在 O(n log n) 时间复杂度和常数级空间复杂度下，对链表进行排序。

示例 1:

输入: 4->2->1->3
输出: 1->2->3->4
示例 2:

输入: -1->5->3->4->0
输出: -1->0->3->4->5

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/sort-list
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。
"""

class ListNode:
    def __init__(self, x):
        self.val = x
        self.next = None

# 一、归并排序（递归）
class Solution:
    def sortList(self, head: ListNode) -> ListNode:
        if not head or not head.next:
            return head  # termination.
        # cut the LinkedList at the mid index.
        slow, fast = head, head.next
        while fast and fast.next:
            fast, slow = fast.next.next, slow.next
        mid, slow.next = slow.next, None  # save and cut.
        # recursive for cutting.
        left, right = self.sortList(head), self.sortList(mid)
        # merge `left` and `right` linked list and return it.
        h = res = ListNode(0)
        while left and right:
            if left.val < right.val:
                h.next, left = left, left.next
            else:
                h.next, right = right, right.next
            h = h.next
        h.next = left if left else right
        return res.next

# 二、归并排序（从底至顶直接合并）
class Solution2:
    def sortList(self, head: ListNode) -> ListNode:
        h, length, step = head, 0, 1
        while h:
            h, length = h.next, length + 1
        res = ListNode(0)
        res.next = head
        # merge the list in different step.
        while step < length:
            pre, h = res, res.next
            while h:
                # get the two merge head `h1`, `h2`
                h1, i = h, step
                while i and h:
                    h, i = h.next, i - 1
                if i:
                    break  # no need to merge because the `h2` is None.
                h2, i = h, step
                while i and h:
                    h, i = h.next, i - 1
                c1, c2 = (
                    step,
                    step - i,
                )  # the `c2`: length of `h2` can be small than the `step`.
                # merge the `h1` and `h2`.
                while c1 and c2:
                    if h1.val < h2.val:
                        pre.next, h1, c1 = h1, h1.next, c1 - 1
                    else:
                        pre.next, h2, c2 = h2, h2.next, c2 - 1
                    pre = pre.next
                pre.next = h1 if c1 else h2
                while c1 > 0 or c2 > 0:
                    pre, c1, c2 = pre.next, c1 - 1, c2 - 1
                pre.next = h
            step *= 2
        return res.next