import os
import sys
import io
import threading
import queue

import connection
from type import Uint8, Uint16, OpaqueUint16, OpaqueLength
from disp import hexdump

from protocol_tlscontext import TLSContext
from protocol_types import ContentType, HandshakeType
from protocol_recordlayer import TLSPlaintext, TLSCiphertext, TLSInnerPlaintext
from protocol_handshake import Handshake
from protocol_hello import ClientHello
from protocol_ciphersuite import CipherSuites, CipherSuite
from protocol_extensions import Extensions, Extension, ExtensionType
from protocol_ext_version import SupportedVersions, \
    ProtocolVersions, ProtocolVersion
from protocol_ext_supportedgroups import NamedGroupList, NamedGroups, NamedGroup
from protocol_ext_signature import SignatureSchemeList, \
    SignatureSchemes, SignatureScheme
from protocol_ext_keyshare import KeyShareHello, KeyShareEntrys, KeyShareEntry
from protocol_authentication import Finished, Hash, OpaqueHash
from protocol_alert import Alert, AlertLevel, AlertDescription

from crypto_x25519 import x25519
import crypto_hkdf as hkdf

ctx = TLSContext('client')

# === Key Exchange Parameters ===

# ECDH with X25519, with Forward Secrecy available
dhkex_class = x25519
# client_private_key
secret_key = os.urandom(32)
# client_public_key
public_key = dhkex_class(secret_key)

# === Client Hello ====

client_hello = Handshake(
    msg_type=HandshakeType.client_hello,
    msg=ClientHello(
        # Cipher Suites
        cipher_suites=CipherSuites([
            CipherSuite.TLS_CHACHA20_POLY1305_SHA256,
            CipherSuite.TLS_EMPTY_RENEGOTIATION_INFO_SCSV
        ]),
        # Extensions
        extensions=Extensions([
            # Supported-Versions ext
            # Here, we only support TLS 1.3
            Extension(
                extension_type=ExtensionType.supported_versions,
                extension_data=SupportedVersions(
                    versions=ProtocolVersions([
                        ProtocolVersion.TLS13
                    ])
                )
            ),
            # Supported-Groups ext
            # Here, we only support X25519 as key exchange algorithm
            Extension(
                extension_type=ExtensionType.supported_groups,
                extension_data=NamedGroupList(
                    named_group_list=NamedGroups([
                        NamedGroup.x25519
                    ])
                )
            ),
            # Signature-Algorithms ext
            # Here,
            Extension(
                extension_type=ExtensionType.signature_algorithms,
                extension_data=SignatureSchemeList(
                    supported_signature_algorithms=SignatureSchemes([
                        SignatureScheme.rsa_pss_rsae_sha256,
                        SignatureScheme.rsa_pss_rsae_sha384,
                        SignatureScheme.rsa_pss_rsae_sha512,
                    ])
                )
            ),
            # Key-Share ext
            # Here, we only support X25519 as key exchange algorithm
            Extension(
                extension_type=ExtensionType.key_share,
                extension_data=KeyShareHello(
                    shares=KeyShareEntrys([
                        KeyShareEntry(
                            group=NamedGroup.x25519,
                            key_exchange=OpaqueUint16(public_key)
                        )
                    ])
                )
            )
        ])
    )
)
ctx.append_msg(client_hello)

# Create TLSPlaintext to include Handshake message
tlsplaintext = TLSPlaintext.create(ContentType.handshake, client_hello)
print(tlsplaintext)
print('[>>>] Send:')
print(hexdump(bytes(tlsplaintext)))

# Connect to server & send TLSPlaintext
client_conn = connection.ClientConnection('localhost', 50007)
client_conn.send_msg(bytes(tlsplaintext))

is_recv_serverhello = False
is_recv_finished = False

print("=== Handshake ===")

# Wait for server response
while True:
    # Blocking receive
    buf = None
    while not buf:
        buf = client_conn.recv_msg(setblocking=True)

    # Upon receiving a message, print it, and turn it into a stream
    print('[<<<] Recv:')
    print(hexdump(buf))
    stream = io.BytesIO(buf)

    # Parse TLSPlaintext
    while True:
        firstbyte = stream.read(1)
        if firstbyte == b'':
            break
        stream.seek(-1, io.SEEK_CUR)
        # Get ContentType
        content_type = ContentType(
            Uint8(int.from_bytes(firstbyte, byteorder='big')))

        # Handle Alert
        if content_type == ContentType.alert:
            tlsplaintext = TLSPlaintext.from_fs(stream)
            for alert in tlsplaintext.get_messages():
                print('[-] Recv Alert!')
                print(alert)
            sys.exit(1)

        # Handle 1st ServerHello
        elif not is_recv_serverhello:
            # ServerHello
            tlsplaintext = TLSPlaintext.from_fs(stream)
            for msg in tlsplaintext.get_messages():
                print('[*] ServerHello!')
                print(msg)
                ctx.append_msg(msg)

            # Produce shared key. Using key-exchange alg, and client secret key
            # Similar to pre-master secret in TLS 1.2
            ctx.set_key_exchange(dhkex_class, secret_key)
            Hash.length = ctx.hash_size
            print('[+] shared key:', ctx.shared_key.hex())

            # Key Schedule, produce early secret & handshake secret
            ctx.key_schedule_in_handshake()

            is_recv_serverhello = True

        # Ignore ChangeCipherSpec message
        elif content_type == ContentType.change_cipher_spec:
            # ChangeCipherSpec
            change_cipher_spec = TLSPlaintext.from_fs(stream)
            print(change_cipher_spec)

        # Handle EncryptedExtensions, Certificate, CertificateVerify, Finished
        elif content_type == ContentType.application_data:
            # EncryptedExtensions, Certificate, CertificateVerify, Finished
            print("Got!")

            tlsplaintext = TLSCiphertext.from_fs(stream) \
                                        .decrypt(ctx.server_traffic_crypto)
            # print(tlsplaintext)
            for msg in tlsplaintext.get_messages():
                ctx.append_msg(msg)
                print(msg)

                if msg.msg_type == HandshakeType.finished:
                    print('[*] Received Finished!')
                    is_recv_finished = True
                    break

    if is_recv_finished:
        break

# Verify Server-Finished
msgs_byte = b''.join(ctx.tls_messages_bytes[:-1])   # exclude Finished
finished_key = hkdf.HKDF_expand_label(
    ctx.server_hs_traffic_secret, b'finished', b'', ctx.hash_size, ctx.hash_name)
expected_verify_data = hkdf.secure_HMAC(
    finished_key, hkdf.transcript_hash(msgs_byte, ctx.hash_name), ctx.hash_name)
actual_verify_data = \
    ctx.tls_messages.get(
        HandshakeType.finished).msg.verify_data.get_raw_bytes()

if actual_verify_data != expected_verify_data:
    print('decrypt_error!')
    # TODO: create and send Alert msg
    sys.exit(0)

# Send Client-Finished
msgs_byte = ctx.get_messages_byte()
finished_key = hkdf.HKDF_expand_label(
    ctx.client_hs_traffic_secret, b'finished', b'', ctx.hash_size, ctx.hash_name)
verify_data = hkdf.secure_HMAC(
    finished_key, hkdf.transcript_hash(msgs_byte, ctx.hash_name), ctx.hash_name)

finished = Handshake(
    msg_type=HandshakeType.finished,
    msg=Finished(
        verify_data=OpaqueHash(bytes(verify_data))
    )
)
print(finished)

tlsciphertext = TLSPlaintext.create(ContentType.handshake, finished) \
    .encrypt(ctx.client_traffic_crypto)
print(tlsciphertext)
print(hexdump(bytes(tlsciphertext)))
client_conn.send_msg(bytes(tlsciphertext))

# Key Schedule
ctx.key_schedule_in_app_data()

# Handsake is done.
# Now, we can send application data

loop_keyboard_input = True

def read_keyboard_input(inputQueue):
    print('Ready for keyboard input:')
    while loop_keyboard_input:
        input_str = input()
        inputQueue.put(input_str + "\n")

inputQueue = queue.Queue()
inputThread = threading.Thread(target=read_keyboard_input,
                               args=(inputQueue,), daemon=True)
inputThread.start()

print("=== Application Data ===")

try:
    while True:

        buf = None
        while not buf:
            buf = client_conn.recv_msg(setblocking=False)

            # Send keyboard input
            if inputQueue.qsize() > 0:
                input_byte = inputQueue.get().encode()
                tlsciphertext = \
                    TLSPlaintext.create(ContentType.application_data, input_byte) \
                    .encrypt(ctx.client_app_data_crypto)
                print(tlsciphertext)
                print('[>>>] Send:')
                print(hexdump(bytes(tlsciphertext)))

                client_conn.send_msg(bytes(tlsciphertext))

        print('[<<<] Recv:')
        print(hexdump(buf))

        stream = io.BytesIO(buf)

        while True:
            firstbyte = stream.read(1)
            if firstbyte == b'':
                break
            stream.seek(-1, io.SEEK_CUR)

            content_type = \
                ContentType(Uint8(int.from_bytes(firstbyte, byteorder='big')))

            # Handle Alert
            if content_type == ContentType.alert:
                tlsplaintext = TLSPlaintext.from_fs(stream)
                for alert in tlsplaintext.get_messages():
                    print('[-] Recv Alert!')
                    print(alert)
                sys.exit(1)

            # Handle encrypted application data
            elif content_type == ContentType.application_data:
                obj = TLSCiphertext.from_fs(stream) \
                    .decrypt(ctx.server_app_data_crypto)
                print(obj)

                if isinstance(obj.fragment, Handshake):
                    # New Session Ticket
                    print('[+] New Session Ticket arrived!')
                    ctx.append_msg(obj)

                else:
                    print(bytes(obj.fragment))

except KeyboardInterrupt:
    print('\nBye!')

# Closure Alert
closure_alert = Alert(
    level=AlertLevel.fatal,
    description=AlertDescription.close_notify
)

tlsciphertext = TLSPlaintext.create(ContentType.alert, closure_alert) \
    .encrypt(ctx.client_app_data_crypto)
print(tlsciphertext)
print(hexdump(bytes(tlsciphertext)))
client_conn.send_msg(bytes(tlsciphertext))

loop_keyboard_input = False

client_conn.close()


# -*- coding: utf-8 -*-
"""
Created on Sun May 31 16:13:35 2020

@author: kangsun
"""
import numpy as np
import datetime as dt
import struct
import os
import logging
from scipy.io import loadmat
from netCDF4 import Dataset
from scipy.interpolate import interp1d

def F_isrf_convolve_fft(w1,s1,w2,isrf_w,isrf_dw0,isrf_lut0):
    """
    astropy.covolute.convolve_fft-based convolution using wavelength-dependent isrf
    w1:
        high-resolution wavelength
    s1:
        high-resolution spectrum
    w2:
        low-resolution wavelength
    isrf_w:
        center wavelength grid of isrf lut
    isrf_dw:
        wavelength grid on which isrfs are defined
    isrf_lut:
        instrument spectral response function look up table
    """
    from astropy.convolution import convolve_fft
    from scipy.interpolate import RegularGridInterpolator
    from math import isclose
    if isrf_lut0.shape != (len(isrf_w),len(isrf_dw0)):
        raise ValueError('isrf_lut dimension incompatible!')
        return np.full(w2.shape,np.nan)
    # make sure w1 and isrf_dw have the same resolution
    w1_step = np.median(np.diff(w1))
    isrf_dw_step = np.median(np.diff(isrf_dw0))
    if isclose(w1_step,isrf_dw_step):
        isrf_dw = isrf_dw0
        isrf_lut = isrf_lut0
    else:# if not, interpolate isrf to a compatible grid with w1
        isrf_dw_min = np.min(isrf_dw0)
        isrf_dw_max = -isrf_dw_min
        isrf_dw = np.linspace(isrf_dw_min,isrf_dw_max,int((isrf_dw_max-isrf_dw_min)/w1_step)+1)
        isrf_lut = np.zeros((len(isrf_w),len(isrf_dw)))
        for (iw,w) in enumerate(isrf_w):
            interp_func = interp1d(isrf_dw0,isrf_lut0[iw,:])
            isrf_lut[iw,:] = interp_func(isrf_dw)
    # note that the isrf is flipped: convolution is the mirror-image of kernel averaging
    s2_fft_lut = np.array([convolve_fft(s1,isrf_lut[iw,::-1]) for (iw,w) in enumerate(isrf_w)])
    inter_func = RegularGridInterpolator((isrf_w,w1),s2_fft_lut,bounds_error=False)
    return inter_func((w2,w2))

# below ReadSeq routines are from Jonathan Franklin
class Sequence():
    pass

class FrameMeta():
    pass

def ParseFrame(framenum,height,width,bindata):
#    print(framenum)
    Meta = FrameMeta()

    # Calc size of actual image minus meta data
    numpixels = (height-1)*width
    fmt = '<{}h'.format(numpixels)
    fmtsize = struct.calcsize(fmt)
    dataraw = struct.unpack_from(fmt,bindata[:fmtsize])
    data = np.reshape(dataraw,(height-1,width))

    # Grab time stamp
    temp = struct.unpack_from('<lh',bindata[-6:])
    Meta.timestamp = dt.datetime(1970,1,1) + \
              dt.timedelta(seconds=temp[-2],microseconds=temp[-1]*1000)
    # Grab Meta Data -- Not all of this seems to have real data in it.
    metaraw = bindata[fmtsize:-6]
    temp = struct.unpack_from('<{}h'.format(width),metaraw)
    metaraw = struct.pack('>{}h'.format(width),*temp)

    Meta.partNum = metaraw[2:34].decode('ascii').rstrip('\x00')
    Meta.serNum = metaraw[34:48].decode('ascii').rstrip('\x00')
    Meta.fpaType = metaraw[48:64].decode('ascii').rstrip('\x00')
#    print(Meta.partNum)
#    print(Meta.serNum)
#    print(Meta.fpaType)
    Meta.crc = struct.unpack_from('>I',metaraw[64:68])[0]
    Meta.frameCounter = struct.unpack_from('>i',metaraw[68:72])[0]
    Meta.frameTime = struct.unpack_from('>f',metaraw[72:76])[0]
    Meta.intTime = struct.unpack_from('>f',metaraw[76:80])[0]
    Meta.freq = struct.unpack_from('>f',metaraw[80:84])[0]
    Meta.boardTemp = struct.unpack_from('>f',metaraw[120:124])[0]
    Meta.rawNUC = struct.unpack_from('>H',metaraw[124:126])[0]
    Meta.colOff = struct.unpack_from('>h',metaraw[130:132])[0]
    Meta.numCols = struct.unpack_from('>h',metaraw[132:134])[0] + 1
    Meta.rowOff = struct.unpack_from('>h',metaraw[136:138])[0]
    Meta.numRows = struct.unpack_from('>h',metaraw[138:140])[0] + 1
    timelist = struct.unpack_from('>7h',metaraw[192:206])
    Meta.yr = timelist[0]
    Meta.dy = timelist[1]
    Meta.hr = timelist[2]
    Meta.mn = timelist[3]
    Meta.sc = timelist[4]
    Meta.ms = timelist[5]
    Meta.microsec = timelist[6]
    Meta.fpaTemp = struct.unpack_from('>f',metaraw[476:480])[0]
    Meta.intTimeTicks = struct.unpack_from('>I',metaraw[142:146])[0]

    return [data,Meta]

def ReadSeq(seqfile):
    ## Pull camera (ch4 v o2) and sequence timestamp from filename.
    temp = seqfile.split('/')[-1]
    temp = temp.split('_camera_')
    Seq = Sequence()
    Seq.Camera = temp[0]
    Seq.SeqTime = dt.datetime.strptime(temp[1].strip('.seq'),'%Y_%m_%d_%H_%M_%S')

    ## Open file for binary read
    fin = open(seqfile,'rb')
    binhead = fin.read(8192)

    ## Grab meta data for sequence file.
    temp = struct.unpack('<9I',binhead[548:548+36])
    Seq.ImageWidth = temp[0]
    Seq.ImageHeight = temp[1]
    Seq.ImageBitDepth = temp[2]
    Seq.ImageBitDepthTrue = temp[3]
    Seq.ImageSizeBytes = temp[4]
    Seq.NumFrames = temp[6]
    Seq.TrueImageSize = temp[8]
    Seq.NumPixels = Seq.ImageWidth*Seq.ImageHeight

    ## Read raw frames
    rawframes = [fin.read(Seq.TrueImageSize) for i in range(Seq.NumFrames)]
    fin.close()

    ## Process each frame -- dropping filler bytes before passing raw
    # print('Reading {} frames'.format(Seq.NumFrames))
    frames = [ParseFrame(i,Seq.ImageHeight,Seq.ImageWidth,
                         rawframes[i][:Seq.ImageSizeBytes+6]) \
                    for i in range(Seq.NumFrames)]

    Data = np.array([dd[0] for dd in frames])
    Meta = [dd[1] for dd in frames]

    return(Data,Meta,Seq)
# above ReadSeq routines are from Jonathan Franklin

class Dark():
    """
    place holder for an object for dark frame
    """
    pass

class Granule():
    """
    place holder for an object for calibrated level 1b granule
    """
    pass

class MethaneAIR_L1(object):
    """
    level 0 to level 1b processor for MethaneAIR
    """
    def __init__(self,whichBand,l0DataDir,l1DataDir,
                 badPixelMapPath,radCalPath,wavCalPath,windowTransmissionPath):
        """
        whichBand:
            CH4 or O2
        l0DataDir:
            directory of level 0 data
        l1DataDir:
            directory of level 1 data
        badPixelMapPath:
            path to the bad pixel map
        radCalPath:
            path to the radiometric calibration coefficients
        wavCalPath:
            path to wavelength calibration coefficients
        windowTransmissionPath:
            path to the window transmission
        """
        self.whichBand = whichBand
        self.logger = logging.getLogger(__name__)
        self.logger.info('creating an instance of level0->level1 for MethaneAIR '
                         +whichBand+ ' band')
        self.l0DataDir = l0DataDir
        if not os.path.isdir(l1DataDir):
            self.logger.warning(l1DataDir,' does not exist, creating one')
            os.mkdir(l1DataDir)
        self.l1DataDir = l1DataDir

        self.ncol = 1024
        self.nrow = 1280

        # bad pixel map
        self.logger.info('loading bad pixel map')
        self.badPixelMap = np.genfromtxt(badPixelMapPath,delimiter=',',dtype=np.int)
        self.logger.info('loading radiometric calibration coefficients')
        d = loadmat(radCalPath)['coef']

        # pad zero intercept and flip poly coeff order to be compatible with polyval
        radCalCoef = np.concatenate((np.zeros((*d.shape[0:2], 1)),d),axis=2)
        self.radCalCoef = np.flip(radCalCoef,axis=2)

        # wavelength calibration
        self.logger.info('loading wavelength calibration coefficients')
        spec_cal_fid = Dataset(wavCalPath)
        wavCalCoef = spec_cal_fid['pix2nm_polynomial_zero2four'][:]
        if wavCalCoef.shape[0] < wavCalCoef.shape[1]:
            self.logger.warning('you appear to be using an old version wavcal! tranposing...')
            self.wavCalCoef = wavCalCoef.T

        # wavelength calibration is independent of granules for now
        wavelength = np.array([np.polyval(self.wavCalCoef[i,::-1],np.arange(1,self.ncol+1))\
                           for i in range(self.nrow)])
        d = loadmat(windowTransmissionPath)
        # interpolate window transmission to detector wavelength
        f = interp1d(d['wavelength_nm'].squeeze(),d['transmission'].squeeze(),
                     fill_value='extrapolate')
        windowTransmission = f(wavelength)
        wavelength = np.ma.masked_array(wavelength,np.isnan(wavelength))
        windowTransmission = np.ma.masked_array(windowTransmission,
                                                np.isnan(windowTransmission))
        self.wavelength = wavelength
        self.windowTransmission = windowTransmission

    def F_stray_light_input(self,strayLightKernelPath,rowExtent,colExtent,
                            rowCenterMask=5,colCenterMask=7,nDeconvIter=1):
        """
        options for stray light correction
        strayLightKernelPath:
            path to the stray light deconvolution kernel
        rowExtent:
            max row difference away from center (0)
        colExtent:
            max column difference away from center (0)
        rowCenterMask:
            +/-rowCenterMask rows of the kernel will be masked to be 0
        colCenterMask:
            +/-colCenterMask columns of the kernel will be masked to be 0
        nDeconvIter:
            number of iterations in the Van Crittert deconvolution
        """
        d = loadmat(strayLightKernelPath)
        rowFilter = (d['rowAllGrid'].squeeze() >= -np.abs(rowExtent)) & \
        (d['rowAllGrid'].squeeze() <= np.abs(rowExtent))
        colFilter = (d['colAllGrid'].squeeze() >= -np.abs(colExtent)) & \
        (d['colAllGrid'].squeeze() <= np.abs(colExtent))
        reducedRow = d['rowAllGrid'].squeeze()[rowFilter]
        reducedCol = d['colAllGrid'].squeeze()[colFilter]
        strayLightKernel = d['medianStrayLight'][np.ix_(rowFilter,colFilter)]
        strayLightKernel[strayLightKernel<0] = 0.
        strayLightKernel = strayLightKernel/np.nansum(strayLightKernel)
        centerRowFilter = (reducedRow >= -np.abs(rowCenterMask)) & \
        (reducedRow <= np.abs(rowCenterMask))
        centerColFilter = (reducedCol >= -np.abs(colCenterMask)) & \
        (reducedCol <= np.abs(colCenterMask))
        strayLightKernel[np.ix_(centerRowFilter,centerColFilter)] = 0
        strayLightFraction = np.nansum(strayLightKernel)

        self.strayLightKernel = strayLightKernel
        self.strayLightFraction = strayLightFraction
        self.nDeconvIter = nDeconvIter

    def F_grab_dark(self,darkFramePath):
        """
        function to load dark measurement level 0 file
        darkFramePath:
            path to the file
        output is a Dark object
        """
        self.logger.info('loading dark seq file '+darkFramePath)
        (Data,Meta,Seq) = ReadSeq(darkFramePath)
        Data = np.transpose(Data,(2,1,0))
#        badPixelMap3D = np.repeat(self.badPixelMap[...,np.newaxis],Seq.NumFrames,axis=2)
#        darkData = np.nanmean(np.ma.masked_where(badPixelMap3D!=0,Data),axis=2)
#        darkStd = np.nanstd(np.ma.masked_where(badPixelMap3D!=0,Data),axis=2)
        darkData = np.nanmean(Data,axis=2)
        darkStd = np.nanstd(Data,axis=2)
#        darkData[self.badPixelMap!=0] = np.nan
#        darkStd[self.badPixelMap!=0] = np.nan
        dark = Dark()
        dark.data = darkData
        dark.noise = darkStd
        dark.nFrame = Seq.NumFrames
        dark.frameTime = np.array([Meta[i].frameTime for i in range(Seq.NumFrames)])
        dark.seqDateTime = Seq.SeqTime
        dark.frameDateTime = np.array([Meta[i].timestamp for i in range(Seq.NumFrames)])
        return dark

#    def F_error_input(self,darkOffsetDN=1500,ePerDN=4.6,readOutError=40):
#        """
#        try realistic error estimation parameters
#        removed given Jenna Samra's email on 2020/6/9 10:47 pm
#        """
#        self.darkOffsetDN = darkOffsetDN
#        self.ePerDN = ePerDN
#        self.readOutError = readOutError

    def F_granule_processor(self,granulePath,dark,ePerDN=4.6,timeOnly=False):
        """
        working horse
        granulePath:
            path to a granule of level 0 data
        dark:
            a Dark object. Most useful part is dark.data, a nrow by ncol matrix
            of dark frame for subtraction
        ePerDN:
            electrons per DN number
        timeOnly:
            if only save time stamps, no worries about other data
        output is a Granule object
        """
        if timeOnly:
            self.logger.info('loading seq file '+granulePath)
            (Data,Meta,Seq) = ReadSeq(granulePath)
            granuleFrameTime = np.array([Meta[i].frameTime for i in range(Seq.NumFrames)])
            granule = Granule()
            granule.nFrame = Seq.NumFrames
            granule.frameTime = granuleFrameTime
            granule.seqDateTime = Seq.SeqTime
            granule.frameDateTime = np.array([Meta[i].timestamp for i in range(Seq.NumFrames)])
            return granule
        darkData = dark.data
        darkStd = dark.noise
#        darkOffsetDN = self.darkOffsetDN
#        ePerDN = self.ePerDN
#        readOutError = self.readOutError
        self.logger.info('loading seq file '+granulePath)
        (Data,Meta,Seq) = ReadSeq(granulePath)
        Data = np.float32(np.transpose(Data,(2,1,0)))
        # estimate noise
        self.logger.info('estimate noise')
        Noise = np.sqrt((Data-darkData[...,np.newaxis])*ePerDN+(darkStd[...,np.newaxis]*ePerDN)**2)/ePerDN
#        Noise = np.sqrt((Data-darkOffsetDN)*ePerDN+readOutError**2)/ePerDN
        # remove dark
        Data = Data-darkData[...,np.newaxis]
        # mask bad pixels
        self.logger.info('mask bad pixels')
        badPixelMap3D = np.repeat(self.badPixelMap[...,np.newaxis],Seq.NumFrames,axis=2)
        Data[badPixelMap3D!=0] = np.nan
        Noise[badPixelMap3D!=0] = np.nan
#        Data = np.ma.masked_array(Data,
#                                  ((badPixelMap3D!=0) | (np.isnan(Data))) )
#        Noise = np.ma.masked_array(Noise,
#                                  ((badPixelMap3D!=0) | (np.isnan(Data))) )
        # normalize frame time, DN/s
        granuleFrameTime = np.array([Meta[i].frameTime for i in range(Seq.NumFrames)])
        Data = Data/granuleFrameTime[np.newaxis,np.newaxis,:]
        Noise = Noise/granuleFrameTime[np.newaxis,np.newaxis,:]
        # rad cal
        self.logger.info('radiometric calibration')
        for i in range(self.nrow):
            for j in range(self.ncol):
                if np.isnan(self.radCalCoef[i,j,0]):
                    Data[i,j,:] = np.nan
                    Noise[i,j,:] = np.nan
                    continue
                Data[i,j,:] = np.polyval(self.radCalCoef[i,j,:],Data[i,j,:])
                Noise[i,j,:] = np.polyval(self.radCalCoef[i,j,:],Noise[i,j,:])
#        newData = np.zeros(Data.shape)
#        newNoise = np.zeros(Data.shape)
#        for ipoly in range(self.radCalCoef.shape[-1]):
#            newData = newData+self.radCalCoef[...,ipoly][...,np.newaxis]*\
#            np.power(Data,ipoly+1)
#            newNoise = newNoise+self.radCalCoef[...,ipoly][...,np.newaxis]*\
#            np.power(Noise,ipoly+1)
#        Data = newData
#        Noise = newNoise

        if hasattr(self,'strayLightKernel'):
            self.logger.info('proceed with stray light correction')
            nDeconvIter = self.nDeconvIter
            strayLightKernel = self.strayLightKernel
            strayLightFraction = self.strayLightFraction
            from astropy.convolution import convolve_fft
            for iframe in range(Seq.NumFrames):
                tmpData = Data[:,:,iframe]
                for iDeconvIter in range(nDeconvIter):
                    tmpData = (Data[:,:,iframe]-convolve_fft(tmpData,strayLightKernel,normalize_kernel=False))\
                    /(1-strayLightFraction)
                Data[:,:,iframe] = tmpData
        else:
            self.logger.info('no data for stray light correction')
        # flip column order
        Data = Data[:,::-1,:]
        Noise = Noise[:,::-1,:]
        # window transmission correction
        Data = Data/self.windowTransmission[...,np.newaxis]
        Noise = Noise/self.windowTransmission[...,np.newaxis]
        granule = Granule()
        granule.data = Data
        granule.noise = Noise
        granule.nFrame = Seq.NumFrames
        granule.frameTime = granuleFrameTime
        granule.seqDateTime = Seq.SeqTime
        granule.frameDateTime = np.array([Meta[i].timestamp for i in range(Seq.NumFrames)])
        return granule

    def F_block_reduce_granule(self,granule,acrossTrackAggregation=6,
                               alongTrackAggregation=10,ifKeepTail=True):
        """
        block-reduce gradule by aggregating in across/along track or row/frame dimensions
        graule:
            outputs from F_granule_processor, a Granule object
        across/alongTrackAggregation:
            as indicated by name
        ifKeepTail:
            whether to aggregate the leftover frames after block reduce
        """
        if acrossTrackAggregation == 1 and alongTrackAggregation == 1:
            self.logger.warning('No aggregation needs to be done. Why are you calling this function?')
            return granule
        from astropy.nddata.utils import block_reduce
        newGranule = Granule()
        newGranule.seqDateTime = granule.seqDateTime
        nFootprint = np.floor(granule.data.shape[0]/acrossTrackAggregation).astype(np.int)
        nTailRow = (granule.data.shape[0]-nFootprint*acrossTrackAggregation).astype(np.int)
        self.logger.info('%d'%granule.data.shape[0]+' rows will be reduced to %d'%nFootprint+' footprints')
        self.logger.info('The last %d'%nTailRow+' rows will be thrown away')
        nFrameAggregated = np.floor(granule.data.shape[2]/alongTrackAggregation).astype(np.int)
        nTailFrame = (granule.data.shape[2]-nFrameAggregated*alongTrackAggregation).astype(np.int)
        self.logger.info('%d'%granule.data.shape[2]+' frames in the granule will be reduced to %d'%nFrameAggregated+' aggregated frames')

        if not ifKeepTail or nTailFrame == 0:
            self.logger.info('The last %d'%nTailFrame+' frames will be thrown away')
            newGranule.data = block_reduce(granule.data,
                                           (acrossTrackAggregation,1,alongTrackAggregation),
                                           func=np.nanmean)
            sumSquare = block_reduce(np.power(granule.noise,2),
                                    (acrossTrackAggregation,1,alongTrackAggregation),
                                    func=np.nansum)
            newGranule.noise = np.sqrt(sumSquare/acrossTrackAggregation/alongTrackAggregation)
            newGranule.frameDateTime = block_reduce(granule.frameDateTime,
                                                    alongTrackAggregation,
                                                    func=np.nanmax)
            newGranule.frameTime = block_reduce(granule.frameTime,
                                                alongTrackAggregation,
                                                func=np.nansum)
            newGranule.nFrame = nFrameAggregated
        else:
            self.logger.info('The last %d'%nTailFrame+' frame will be aggregated')
            data1 = block_reduce(granule.data,
                                (acrossTrackAggregation,1,alongTrackAggregation),
                                func=np.nanmean)
            data2 = block_reduce(granule.data[...,-nTailFrame:],
                                 (acrossTrackAggregation,1,nTailFrame),
                                 func=np.nanmean)
            self.logger.debug('data1 shape is %d'%data1.shape[0]+', %d'%data1.shape[1]+', %d'%data1.shape[2])
            self.logger.debug('data2 shape is %d'%data2.shape[0]+', %d'%data2.shape[1]+', %d'%data2.shape[2])
            newGranule.data = np.concatenate((data1,data2),axis=2)
            sumSquare = block_reduce(np.power(granule.noise,2),
                                     (acrossTrackAggregation,1,alongTrackAggregation),
                                     func=np.nansum)
            noise1 = np.sqrt(sumSquare/acrossTrackAggregation/alongTrackAggregation)
            sumSquare = block_reduce(np.power(granule.noise[...,-nTailFrame:],2),
                                     (acrossTrackAggregation,1,nTailFrame),
                                     func=np.nansum)
            noise2 = np.sqrt(sumSquare/acrossTrackAggregation/nTailFrame)
            newGranule.noise = np.concatenate((noise1,noise2),axis=2)
            frameDateTime1 = block_reduce(granule.frameDateTime,
                                          alongTrackAggregation,
                                          func=np.nanmax)
            frameDateTime2 = block_reduce(granule.frameDateTime[-nTailFrame:],
                                          nTailFrame,
                                          func=np.nanmax)
            newGranule.frameDateTime = np.concatenate((frameDateTime1,frameDateTime2))
            frameTime1 = block_reduce(granule.frameTime,
                                      alongTrackAggregation,
                                      func=np.nansum)
            frameTime2 = block_reduce(granule.frameTime[-nTailFrame:],
                                      nTailFrame,
                                      func=np.nansum)
            newGranule.frameTime = np.concatenate((frameTime1,frameTime2))
            newGranule.nFrame = nFrameAggregated+1
            if len(newGranule.frameTime) != newGranule.nFrame:
                self.logger.error('this should not happen!')
        return newGranule

    def F_cut_granule(self,granule,granuleSeconds=10):
        """
        cut a granule into a list of granules with shorter, regular-time intervals
        graule:
            outputs from F_granule_processor, a Granule object
        granuleSeconds:
            length of cut granule in s
        """
        if hasattr(granule,'data'):
            data = granule.data
            noise = granule.noise
        #nFrame = granule.nFrame
        frameTime = granule.frameTime
        seqDateTime = granule.seqDateTime
        frameDateTime = granule.frameDateTime
        minDateTime = np.min(frameDateTime)
        maxDateTime = np.max(frameDateTime)
        minSecond = (minDateTime-minDateTime.replace(hour=0,minute=0,second=0,microsecond=0)).total_seconds()
        startSecond = np.floor(minSecond/granuleSeconds)*granuleSeconds
        startDateTime = minDateTime.replace(hour=0,minute=0,second=0,microsecond=0)+dt.timedelta(seconds=startSecond)
        nGranule = np.floor((maxDateTime-startDateTime).total_seconds()/granuleSeconds)+1
        nGranule = np.int16(nGranule)
        secondList = np.arange(nGranule+1)*np.float(granuleSeconds)
        granuleEdgeDateTimeList = np.array([startDateTime+dt.timedelta(seconds=secondList[i])
        for i in range(nGranule+1)])

        self.logger.info('cutting granule into %d'%nGranule+' shorter granules with length %d'%granuleSeconds +' seconds')
        granuleList = np.ndarray(shape=(nGranule),dtype=np.object_)
        for i in range(nGranule):
            g0 = Granule()
            g0.seqDateTime = seqDateTime
            f = (frameDateTime >= granuleEdgeDateTimeList[i]) &\
            (frameDateTime < granuleEdgeDateTimeList[i+1])
            g0.nFrame = np.int16(np.sum(f))
            g0.frameDateTime = frameDateTime[f]
            g0.frameTime = frameTime[f]
            if hasattr(granule,'data'):
                g0.data = data[...,f]
                g0.noise = noise[...,f]
            granuleList[i] = g0
        return granuleList

    def F_save_L1B_time_only(self,granule,headerStr='MethaneAIR_L1B_CH4_timeonly_'):
        """
        save only the time stamps for calibrated data
        granule:
            a Granule object generated by F_granule_processor or F_cut_granule
        headerStr:
            string that is different from the actual l1b files with real data
        """
        from scipy.io import savemat
        l1FilePath = os.path.join(self.l1DataDir,headerStr
                                  +np.min(granule.frameDateTime).strftime('%Y%m%dT%H%M%S')+'_'
                                  +np.max(granule.frameDateTime).strftime('%Y%m%dT%H%M%S')+'_'
                                  +dt.datetime.now().strftime('%Y%m%dT%H%M%S')
                                  +'.mat')
        GEOS_5_tau = np.array([(granule.frameDateTime[i]-dt.datetime(1985,1,1,0,0,0)).total_seconds()/3600.
                               for i in range(granule.nFrame)])
        granuleYear = np.array([granule.frameDateTime[i].year for i in range(granule.nFrame)])
        granuleMonth = np.array([granule.frameDateTime[i].month for i in range(granule.nFrame)])
        granuleDay = np.array([granule.frameDateTime[i].day for i in range(granule.nFrame)])
        granuleHour = np.array([granule.frameDateTime[i].hour for i in range(granule.nFrame)])
        granuleMinute = np.array([granule.frameDateTime[i].minute for i in range(granule.nFrame)])
        granuleSecond = np.array([granule.frameDateTime[i].second for i in range(granule.nFrame)])
        granuleMicrosecond = np.array([granule.frameDateTime[i].microsecond for i in range(granule.nFrame)])
        self.logger.info('saving time only .mat L1B file '+l1FilePath)
        savemat(l1FilePath,{'GEOS_5_tau':GEOS_5_tau,
                            'granuleYear':granuleYear,
                            'granuleMonth':granuleMonth,
                            'granuleDay':granuleDay,
                            'granuleHour':granuleHour,
                            'granuleMinute':granuleMinute,
                            'granuleSecond':granuleSecond,
                            'granuleMicrosecond':granuleMicrosecond})

    def F_save_L1B_mat(self,granule,headerStr='MethaneAIR_L1B_CH4_',radianceOnly=True):
        """
        save calibrated data to level 1b file in .mat format for quick view
        granule:
            a Granule object generated by F_granule_processor or F_cut_granule
        headerStr:
            'MethaneAIR_L1B_CH4_' or 'MethaneAIR_L1B_O2_' or 'MethaneAIR_L1B_'
        """
        from scipy.io import savemat
        l1FilePath = os.path.join(self.l1DataDir,headerStr
                                  +np.min(granule.frameDateTime).strftime('%Y%m%dT%H%M%S')+'_'
                                  +np.max(granule.frameDateTime).strftime('%Y%m%dT%H%M%S')+'_'
                                  +dt.datetime.now().strftime('%Y%m%dT%H%M%S')
                                  +'.mat')
        GEOS_5_tau = np.array([(granule.frameDateTime[i]-dt.datetime(1985,1,1,0,0,0)).total_seconds()/3600.
                               for i in range(granule.nFrame)])
        granuleYear = np.array([granule.frameDateTime[i].year for i in range(granule.nFrame)])
        granuleMonth = np.array([granule.frameDateTime[i].month for i in range(granule.nFrame)])
        granuleDay = np.array([granule.frameDateTime[i].day for i in range(granule.nFrame)])
        granuleHour = np.array([granule.frameDateTime[i].hour for i in range(granule.nFrame)])
        granuleMinute = np.array([granule.frameDateTime[i].minute for i in range(granule.nFrame)])
        granuleSecond = np.array([granule.frameDateTime[i].second for i in range(granule.nFrame)])
        granuleMicrosecond = np.array([granule.frameDateTime[i].microsecond for i in range(granule.nFrame)])
        wavelength = np.tile(self.wavelength[...,np.newaxis],granule.nFrame).transpose([0,2,1])
        data = granule.data.transpose([0,2,1])
        noise = granule.noise.transpose([0,2,1])
        # flip row order if O2 camera
        if self.whichBand == 'O2':
            wavelength = wavelength[::-1,...]
            data = data[::-1,...]
            noise = noise[::-1,...]
        self.logger.info('saving .mat L1B file '+l1FilePath)
        if radianceOnly:
            savemat(l1FilePath,{'radiance':np.asfortranarray(data).astype(np.float32),
                            'GEOS_5_tau':GEOS_5_tau,
                            'granuleYear':granuleYear,
                            'granuleMonth':granuleMonth,
                            'granuleDay':granuleDay,
                            'granuleHour':granuleHour,
                            'granuleMinute':granuleMinute,
                            'granuleSecond':granuleSecond,
                            'granuleMicrosecond':granuleMicrosecond})
            return
        savemat(l1FilePath,{'wavelength':np.asfortranarray(wavelength),
                            'radiance':np.asfortranarray(data),
                            'radiance_error':np.asfortranarray(noise),
                            'GEOS_5_tau':GEOS_5_tau,
                            'granuleYear':granuleYear,
                            'granuleMonth':granuleMonth,
                            'granuleDay':granuleDay,
                            'granuleHour':granuleHour,
                            'granuleMinute':granuleMinute,
                            'granuleSecond':granuleSecond,
                            'granuleMicrosecond':granuleMicrosecond})

    def F_save_L1B(self,granule,headerStr='MethaneAIR_L1B_CH4_'):
        """
        save calibrated data to level 1b file
        granule:
            a Granule object generated by F_granule_processor or F_cut_granule
        headerStr:
            'MethaneAIR_L1B_CH4_' or 'MethaneAIR_L1B_O2_' or 'MethaneAIR_L1B_'
        """
        from pysplat import level1
        l1FilePath = os.path.join(self.l1DataDir,headerStr
                                  +np.min(granule.frameDateTime).strftime('%Y%m%dT%H%M%S')+'_'
                                  +np.max(granule.frameDateTime).strftime('%Y%m%dT%H%M%S')+'_'
                                  +dt.datetime.now().strftime('%Y%m%dT%H%M%S')
                                  +'.nc')
        GEOS_5_tau = np.array([(granule.frameDateTime[i]-dt.datetime(1985,1,1,0,0,0)).total_seconds()/3600.
                               for i in range(granule.nFrame)])
        wavelength = np.tile(self.wavelength[...,np.newaxis],granule.nFrame).transpose([0,2,1])
        data = granule.data.transpose([0,2,1])
        noise = granule.noise.transpose([0,2,1])
        # flip row order if O2 camera
        if self.whichBand == 'O2':
            wavelength = wavelength[::-1,...]
            data = data[::-1,...]
            noise = noise[::-1,...]
        l1 = level1(l1FilePath,
                    lon=np.zeros((self.nrow,granule.nFrame)),
                    lat=np.zeros((self.nrow,granule.nFrame)),
                    obsalt=np.zeros(granule.nFrame),
                    time=GEOS_5_tau)
        l1.add_radiance_band(wvl=wavelength,
                             rad=data,
                             rad_err=noise)
        l1.close()


import numpy as np

#####################################################################
# All plots, models, config file will be stored here
OutputDirName = "./Results/Output_Merged2GsfID_EB_2017_dask"
xgbmodeloutname = "XGBoostModel_Merged2GsfID_EB.root"
Pickle_signal = "./data/Merged-ID/Dataframe_MergedID_HDalitz_eeg_allMass_EB_2017.pkl"
Pickle_bkg = "./data/Merged-ID/Dataframe_MergedID_DYJets_GJets_QCD_EB_2017.pkl"
Clfname = "MergedID"
#####################################################################
# prepare sample
testsize = 0.2  #(0.2 means 20%)
CommonCut = "(elePresel == True) and (nVtx > 1) and (nGsfMatchToReco >= 2) and (Class != 'Merged-1Gsf')"
#####################################################################
Classes = ["Merged-2Gsf", "DYJets", "GJets", "QCD"]
ClassColors = ["#377eb8", "#3A9679", "#E16262", "#787A91"]
MVA = "XGB"

featureplotparam_json = "FeaturePlotParam_MergedID_EB.json"
features = [
    "rho",
    "eleSCEta",
    "eleSCRawEn",

    "eleD0",
    "eleDz",
    "elePtError",
    "eleHoverE",
    "eleEoverP",
    "eleEoverPInv",
    "eledPhiAtVtx",
    "eleSigmaIEtaIEtaFull5x5",
    "elePFChIso",
    "elePFPhoIso",
    "eleR9Full5x5",
    "eleTrkdxy",
    "eleMissHits",
    "eleConvVeto",
    "gsfPtRatio",
    "gsfDeltaR",
    "gsfRelPtRatio"
]

RandomState = 42
param = {
    "objective": "multi:softprob",
    "num_class": len(Classes),
    "eval_metric": "mlogloss",
    "tree_method": "gpu_hist",
    "random_state": RandomState,
    "eta": 0.2,
    "max_delta_step": 5,
    "max_depth": 7,
    "min_child_weight": 40,
    "gamma": 1,
    "subsample": 0.5
}

# number of trees
num_boost_round = 500
early_stopping_rounds = 20

#####################################################################
# Reweighting scheme
# ! Notice that "intwei" should be included in the input pkl files
# * instwei: (relative xs) ggF as 1, for it being the production with the largest xs and VBF as xs_VBF/xs_ggF
Reweighing = "Nothing"
ptbins = [0, 10, 20, 30, 40, 50, 60, 100, 150]
etabins = np.linspace(-1.5, 1.5, num = 11)
ptwtvar = "eleCalibPt"
etawtvar = "eleEta"


import json
from twilio.rest import Client
from twilio.twiml.messaging_response import MessagingResponse
from twilio.twiml.voice_response import VoiceResponse
from flask import Flask, request, abort
from functools import wraps
import NumberList

app = Flask(__name__)

with open("API_Keys.json") as json_file:
    file = json.load(json_file)
    _account_sid = file["SID"]
    _auth_token = file["Token"]
    client = Client(_account_sid, _auth_token)

roster = "2020-2021 Roster.csv"
remind_message = "Hey {}, this is a reminder that your {} class starts in 10 mins at {} PM (PST). "\
                 "Please be sure to come on time. If you would like to opt out of these reminder " \
                 "messages, reply with 'STOP'."
oh_message = "Hey {}, this is a reminder that {} has office hours in 10 mins from {} (PST). Feel " \
             "free to drop by if you need any additional help. If you would like to opt out of " \
             "these reminder messages, reply with 'STOP'."

def validate_cron_request(f):
    """
    Validates that incoming requests genuinely originated from GAE Cron Scheduler
    """
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if request.headers.get('X-AppEngine-Cron') is None:
            app.logger.info("Error — Not sent from Cron Job")
            return abort(403)
        else:
            return f(*args, **kwargs)

    return decorated_function

@app.route("/text-reply", methods=["GET", "POST"])
def text_reply():
    response = MessagingResponse()
    form = "https://website.com/attendance"
    response.message(f"Sorry, staff does not receive replies to this "
                     f"number. If you are going to be absent, please "
                     f"fill out this form: {form}")
    return str(response)

@app.route("/voice-reply", methods=["GET", "POST"])
def phone_reply():
    response = VoiceResponse()
    response.say("Sorry, staff does not actively monitor this number for "
                 "calls. If you have any questions please email us. Thank you!",
                 voice='alice')
    return str(response)

@app.route("/send-test")
@validate_cron_request
def send_test():
    number_dict = {"Firstname Lastname": "+1234567890"}
    for person in number_dict:
        msg = client.messages.create(
            body=remind_message.format(person, "English", "3:00 AM"),
            from_="+1234567890",
            to=number_dict[person]
        )
        return msg.sid

@app.route("/send/language/<path:class_name>")
@validate_cron_request
def send_class_text(class_name):
    app.logger.info(class_name)
    number_dict = NumberList.get_lang_dict(csv=roster, class_name=class_name)
    for person in number_dict:
        class_title = person["Language Class"]
        language = class_title[0: class_title.find(":")]

        class_time = class_title.split()[-1]
        time = class_time[0: class_time.find("-")]
        msg = client.messages.create(
            body=remind_message.format(person["First Name"], language, time),
            from_="+1234567890",
            to=person["Phone Number"]
        )
        # print(remind_message.format(person["First Name"], language, time))
    return class_name

@app.route("/send/leadership/<path:cohort>")
@validate_cron_request
def send_leadership_text(cohort):
    time = "2:45" if cohort == "ARABIC" else "3:45"
    number_dict = NumberList.get_leadership_dict(csv=roster, leadership_cohort=cohort)
    for person in number_dict:
        msg = client.messages.create(
            body=remind_message.format(person["First Name"], "Leadership", time),
            from_="+1234567890",
            to=person["Phone Number"]
        )
    return cohort

@app.route("/send/office_hours/<path:teacher>")
@validate_cron_request
def send_office_hours_text(teacher):
    number_dict = NumberList.get_office_hours_dict(csv=roster, teacher=teacher)
    hours = {"Teacher1": "2:45 PM - 3:45 PM", "Teacher2": "11:00 AM - 12:00 PM",
             "Teacher3": "3:00 PM - 4:00 PM", "Teacher4": "12:00 PM - 1:00 PM ",
             "Teacher5": "5:20 PM - 5:50 PM", "Teacher6": "2:45 PM - 3:45 PM"}
    for person in number_dict:
        msg = client.messages.create(
            body=oh_message.format(person["First Name"], teacher, hours[teacher]),
            from_="+1234567890",
            to=person["Phone Number"]
        )
    return teacher

def main():
    print("Hello World")

if __name__ == '__main__':
    app.run(port=3000, debug=True)
    main()


#! /usr/bin/env python3
# -*- coding: utf-8 -*-

# Time client program

from socket import *
s = socket(AF_INET, SOCK_STREAM)  # Create TCP socket
s.connect(("garfield.ltas.ulg.ac.be", 8888))  # Connect to server
tm = s.recv(1024)  # Receive up to 1024 bytes
s.close()  # Close connection
print("The time is", tm)

input()


import logging.config
import random
import string
import time
from datetime import datetime

import ccxt
import requests

from gbot import config

# 日志配置
logger = logging.getLogger("bot")
DEFAULT_LOGGING = {
    "version": 1,
    "disable_existing_loggers": False,
    "formatters": {
        "simple": {
            "format": "%(asctime)s %(levelname)s %(message)s",
            "datefmt": "%Y-%m-%d %H:%M:%S",
        },
    },
    "handlers": {
        "console": {
            "level": config.LOG_LEVEL,
            "class": "logging.StreamHandler",
            "formatter": "simple",
        }
    },
    "loggers": {"bot": {"handlers": ["console"], "level": config.LOG_LEVEL}},
}
logging.config.dictConfig(DEFAULT_LOGGING)

# API
GRID_PARAMETERS_REQ_PATH = "/grids/parameters/"
GRIDS_REQ_PATH = "/grids/"
ROBOT_LOG_REQ_PATH = "/robots/logs/"
ROBOT_CONFIG_REQ_PATH = "/robots/{robot_id}/config/".format(robot_id=config.ROBOT_ID)
ROBOT_REQ_PATH = "/robots/{robot_id}/".format(robot_id=config.ROBOT_ID)

class HttpClient(object):
    headers = {"Authorization": f"Token {config.TOKEN}"}

    def __init__(self):
        self.base_url: str = config.SERVER_BASE_URL.rstrip("/")

    def fetch_grid_parameters(self):
        parameters = self._request("GET", GRID_PARAMETERS_REQ_PATH)
        return parameters

    def fetch_grids(self):
        grids = self._request("GET", GRIDS_REQ_PATH)
        return grids

    def fetch_robot_config(self):
        cfg = self._request("GET", ROBOT_CONFIG_REQ_PATH)
        return cfg

    def update_grids(self, grids):
        self._request("PATCH", GRIDS_REQ_PATH, json=grids)

    def update_robot(self, data):
        self._request("PATCH", ROBOT_REQ_PATH, json=data)

    def create_log(self, msg):
        data = {"msg": msg, "robot_id": config.ROBOT_ID}
        self._request("POST", ROBOT_LOG_REQ_PATH, json=data)

    def _request(
        self, method, req_path, headers=None, params=None, data=None, json=None
    ):
        url = self.base_url + req_path
        response = requests.request(
            method,
            url,
            headers=headers or self.headers,
            params=params,
            data=data,
            json=json,
        )
        response.raise_for_status()
        return response.json()

class GridBot(object):
    pass

class BybitGridBot(GridBot):
    id = "bybit"
    name = "Bybit"
    ccxt_symbols = {
        "BTCUSD": "BTC/USD",
        "ETHUSD": "ETH/USD",
        "EOSUSD": "EOS/USD",
    }
    ccxt_exchange_class = ccxt.bybit

    def __init__(self, pair):
        self.pair = pair
        self.ccxt_symbol: str = self.ccxt_symbols[self.pair]
        self.ccxt_exchange = self.ccxt_exchange_class(
            {
                "enableRateLimit": True,
                "options": {"adjustForTimeDifference": True, "recvWindow": 10 * 1000},
            }
        )
        self.client = HttpClient()
        self.grids = []
        self.position = {
            "dir": 0,
            "qty": 0.0,
            "avg_price": 0.0,
            "liq_price": 0.0,
            "unrealised_pnl": 0.0,
            "leverage": 0.0,
        }
        self.balance = 0.0
        self._api_key = ""
        self._secret = ""
        self._synced_order_set = set()
        self._grid_id_index_map = {}
        now = self.timestamp_millisecond()
        self._timestamp_markers = {
            "earliest_open_order": now,
            "reset_open_order": now,
        }
        self._order_sync_window = 1 * 60 * 60 * 1000
        self._open_orders_reset_interval = 30 * 60 * 1000

    def sync_config(self):
        cfg = self.client.fetch_robot_config()
        self._api_key = cfg["credential_keys"]["api_key"]
        self._secret = cfg["credential_keys"]["secret"]

    def load_grids(self):
        grids = self.client.fetch_grids()
        self.grids = grids
        for i, grid in enumerate(grids):
            self._grid_id_index_map[grid["id"]] = i

    def place_order(self, *, order_dir, qty, price, cid=""):
        dir_mapping = {
            -1: "sell",
            1: "buy",
        }
        self.ccxt_exchange.create_order(
            symbol=self.ccxt_symbol,
            type="limit",
            side=dir_mapping[order_dir],
            amount=qty,
            price=price,
            params={"time_in_force": "PostOnly", "order_link_id": cid},
        )

    def enable_test_net(self):
        self.ccxt_exchange.set_sandbox_mode(enabled=True)

    def place_orders_batch(self, orders):
        for order_args in orders:
            self.place_order(**order_args)

    def _find_grid_order(self, orders, grid_id):
        for order in orders:
            cid: str = order["info"]["order_link_id"]
            if "-" in cid and grid_id == int(cid.split("-", 1)[0]):
                return order

    def ensure_order(self):
        now = self.timestamp_millisecond()
        if (
            now
            > self._timestamp_markers["reset_open_order"]
            + self._open_orders_reset_interval
        ):
            self.ccxt_exchange.cancel_all_orders(symbol=self.ccxt_symbol)
            self._timestamp_markers["reset_open_order"] = now

        last_price = self.get_last_price()
        bid, ask = self.get_best_price()
        open_orders = self.ccxt_exchange.fetch_open_orders(
            symbol=self.ccxt_symbol, limit=50
        )
        order_args_list = []
        for grid in self.grids:
            grid_id = grid["id"]
            if self._find_grid_order(open_orders, grid_id):
                continue

            if grid["holding"]:
                order_args = {
                    "order_dir": 1,
                    "qty": grid["entry_qty"],
                    "price": min(
                        grid["exit_price"], ask - 0.5, self.position["avg_price"]
                    ),
                }
            else:
                if grid["entry_price"] <= last_price:
                    continue

                order_args = {
                    "order_dir": -1,
                    "qty": grid["entry_qty"],
                    "price": grid["entry_price"],
                }
            suffix = self._get_random_string(k=10)
            order_args["cid"] = f"{grid_id}-{suffix}"
            order_args_list.append(order_args)

        self.place_orders_batch(order_args_list)
        if len(order_args_list) > 0:
            open_orders = self.ccxt_exchange.fetch_open_orders(
                symbol=self.ccxt_symbol, limit=50
            )
            self._timestamp_markers["earliest_open_order"] = min(
                open_orders, key=lambda x: x["timestamp"]
            )["timestamp"]

    def sync_balance(self):
        result = self.ccxt_exchange.fetch_balance()
        self.balance = result["BTC"]["total"]

    def sync_position(self):
        res = self.ccxt_exchange.private_get_position_list({"symbol": self.pair})
        pos = {
            "dir": 0,
            "qty": 0.0,
            "avg_price": 0.0,
            "liq_price": 0.0,
            "unrealised_pnl": 0.0,
            "leverage": 0.0,
        }
        if res["result"]["side"] == "None":
            self.position.update(pos)
            return

        qty = res["result"]["size"]
        dir_mapping = {
            "buy": 1,
            "sell": -1,
        }
        pos_dir = dir_mapping[res["result"]["side"].lower()]
        avg_price = round(float(res["result"]["entry_price"]), 2)  # $
        liq_price = round(float(res["result"]["liq_price"]), 2)
        unrealised_pnl = res["result"]["unrealised_pnl"]
        leverage = res["result"]["effective_leverage"]
        self.position.update(
            {
                "dir": pos_dir,
                "qty": qty,
                "avg_price": avg_price,
                "liq_price": liq_price,
                "unrealised_pnl": unrealised_pnl,
                "leverage": leverage,
            }
        )

    def sync_grids(self):
        since = self._timestamp_markers["earliest_open_order"] - self._order_sync_window
        orders = self.ccxt_exchange.fetch_closed_orders(
            symbol=self.ccxt_symbol,
            since=since,
            limit=50,
            params={"order_status": "Filled"},
        )
        orders = sorted(orders, key=lambda x: x["timestamp"])
        for order in orders:
            if order["id"] in self._synced_order_set:
                continue

            cid: str = order["info"]["order_link_id"]
            if "-" not in cid:
                continue

            grid_id = int(cid.split("-", 1)[0])
            grid_index = self._grid_id_index_map.get(grid_id, -1)
            if grid_index < 0:
                continue

            grid = self.grids[grid_index]
            if order["side"] == "sell":
                grid["filled_qty"] = order["cost"]
                grid["holding"] = True

            if order["side"] == "buy":
                grid["filled_qty"] = 0
                grid["holding"] = False

            logger.info(f"已同步第 {grid_index} 层网格，网格状态={grid['holding']}")
            self._synced_order_set.add(order["id"])

        self.client.update_grids(self.grids)

    def get_last_price(self):
        ticker = self.ccxt_exchange.fetch_ticker(symbol=self.ccxt_symbol)
        last_price = ticker["last"]
        return last_price

    def get_best_price(self):
        order_book = self.ccxt_exchange.fetch_l2_order_book(
            symbol=self.ccxt_symbol, limit=5
        )
        best_bid = order_book["bids"][0][0]
        best_ask = order_book["asks"][0][0]
        return best_ask, best_bid

    def pre_trade(self):
        if config.TEST_NET:
            self.enable_test_net()
        self.sync_config()
        self.auth()
        self.ccxt_exchange.load_markets()
        self.ccxt_exchange.cancel_all_orders(symbol=self.ccxt_symbol)
        self.sync_position()
        self.sync_balance()
        self.load_grids()

    def trade(self):
        msg = "正在设置交易环境"
        logger.info(msg)
        self.client.create_log(msg)
        self.pre_trade()
        cnt = 1
        while True:
            try:
                msg = f"开始第 {cnt} 轮交易"
                logger.info(msg)
                self.client.create_log(msg)
                self.sync_position()
                self.sync_balance()
                self.sync_grids()
                self.ensure_order()
                time.sleep(config.TRADE_INTERVAL)
                cnt += 1
            except Exception as e:
                logger.exception(e, exc_info=True)

    def auth(self):
        self.ccxt_exchange.apiKey = self._api_key
        self.ccxt_exchange.secret = self._secret

    @staticmethod
    def _get_random_string(k=10):
        letters = string.ascii_letters
        return "".join(random.choices(letters, k=k))

    @staticmethod
    def timestamp_millisecond(dt=None, delta=None):
        if not dt:
            dt = datetime.now()

        if delta is not None:
            dt = dt + delta

        return int(dt.timestamp() * 1000)

if __name__ == "__main__":
    bot = BybitGridBot(pair="BTCUSD")
    bot.trade()


from typing import Dict, List

from datawagon.objects.source_file_metadata import SourceFileMetadata

class FileUtils(object):
    def group_by_table_name(
        self,
        file_info_list: List[SourceFileMetadata],
    ) -> Dict[str, List[SourceFileMetadata]]:
        grouped_files = {}
        for file_info in file_info_list:
            if file_info.table_name not in grouped_files:
                grouped_files[file_info.table_name] = [file_info]
            else:
                grouped_files[file_info.table_name].append(file_info)
        return grouped_files

    def check_for_duplicate_files(
        self, file_info_list: List[SourceFileMetadata]
    ) -> List[SourceFileMetadata]:
        file_names = [
            file_info.file_name_without_extension for file_info in file_info_list
        ]

        duplicate_file_names = set(
            [file_name for file_name in file_names if file_names.count(file_name) > 1]
        )

        duplicate_files = [
            file_info
            for file_info in file_info_list
            if file_info.file_name_without_extension in duplicate_file_names
        ]

        return duplicate_files

    def check_for_different_file_versions(
        self, file_info_list: List[SourceFileMetadata]
    ) -> List[List[SourceFileMetadata]]:
        grouped_files = self.group_by_table_name(file_info_list)
        different_file_versions = []

        for table_name, file_infos in grouped_files.items():
            file_versions = {file_info.file_version for file_info in file_infos}
            if len(file_versions) > 1:
                different_file_versions.append(file_infos)

        return different_file_versions


from tinycss2 import parse_component_value_list

__all__ = ['parse']

SUPPORTED_PSEUDO_ELEMENTS = {
    # As per CSS Pseudo-Elements Module Level 4
    'first-line', 'first-letter', 'prefix', 'postfix', 'selection',
    'target-text', 'spelling-error', 'grammar-error', 'before', 'after',
    'marker', 'placeholder', 'file-selector-button',
    # As per CSS Generated Content for Paged Media Module
    'footnote-call', 'footnote-marker',
    # As per CSS Scoping Module Level 1
    'content', 'shadow',
}

def parse(input, namespaces=None, forgiving=False, relative=False):
    """Yield tinycss2 selectors found in given ``input``.

    :param input:
        A string, or an iterable of tinycss2 component values.

    """
    if isinstance(input, str):
        input = parse_component_value_list(input)
    tokens = TokenStream(input)
    namespaces = namespaces or {}
    try:
        yield parse_selector(tokens, namespaces, relative)
    except SelectorError as exception:
        if forgiving:
            return
        raise exception
    while 1:
        next = tokens.next()
        if next is None:
            return
        elif next == ',':
            try:
                yield parse_selector(tokens, namespaces, relative)
            except SelectorError as exception:
                if not forgiving:
                    raise exception
        else:
            if not forgiving:
                raise SelectorError(next, f'unexpected {next.type} token.')

def parse_selector(tokens, namespaces, relative=False):
    tokens.skip_whitespace_and_comment()
    if relative:
        peek = tokens.peek()
        if peek in ('>', '+', '~'):
            initial_combinator = peek.value
            tokens.next()
        else:
            initial_combinator = ' '
        tokens.skip_whitespace_and_comment()
    result, pseudo_element = parse_compound_selector(tokens, namespaces)
    while 1:
        has_whitespace = tokens.skip_whitespace()
        while tokens.skip_comment():
            has_whitespace = tokens.skip_whitespace() or has_whitespace
        selector = Selector(result, pseudo_element)
        if relative:
            selector = RelativeSelector(initial_combinator, selector)
        if pseudo_element is not None:
            return selector
        peek = tokens.peek()
        if peek is None or peek == ',':
            return selector
        elif peek in ('>', '+', '~'):
            combinator = peek.value
            tokens.next()
        elif has_whitespace:
            combinator = ' '
        else:
            return selector
        compound, pseudo_element = parse_compound_selector(tokens, namespaces)
        result = CombinedSelector(result, combinator, compound)

def parse_compound_selector(tokens, namespaces):
    type_selectors = parse_type_selector(tokens, namespaces)
    simple_selectors = type_selectors if type_selectors is not None else []
    while 1:
        simple_selector, pseudo_element = parse_simple_selector(
            tokens, namespaces)
        if pseudo_element is not None or simple_selector is None:
            break
        simple_selectors.append(simple_selector)

    if simple_selectors or (type_selectors, pseudo_element) != (None, None):
        return CompoundSelector(simple_selectors), pseudo_element

    peek = tokens.peek()
    peek_type = peek.type if peek else 'EOF'
    raise SelectorError(peek, f'expected a compound selector, got {peek_type}')

def parse_type_selector(tokens, namespaces):
    tokens.skip_whitespace()
    qualified_name = parse_qualified_name(tokens, namespaces)
    if qualified_name is None:
        return None

    simple_selectors = []
    namespace, local_name = qualified_name
    if local_name is not None:
        simple_selectors.append(LocalNameSelector(local_name))
    if namespace is not None:
        simple_selectors.append(NamespaceSelector(namespace))
    return simple_selectors

def parse_simple_selector(tokens, namespaces):
    peek = tokens.peek()
    if peek is None:
        return None, None
    if peek.type == 'hash' and peek.is_identifier:
        tokens.next()
        return IDSelector(peek.value), None
    elif peek == '.':
        tokens.next()
        next = tokens.next()
        if next is None or next.type != 'ident':
            raise SelectorError(next, f'Expected a class name, got {next}')
        return ClassSelector(next.value), None
    elif peek.type == '[] block':
        tokens.next()
        attr = parse_attribute_selector(TokenStream(peek.content), namespaces)
        return attr, None
    elif peek == ':':
        tokens.next()
        next = tokens.next()
        if next == ':':
            next = tokens.next()
            if next is None or next.type != 'ident':
                raise SelectorError(
                    next, f'Expected a pseudo-element name, got {next}')
            value = next.lower_value
            if value not in SUPPORTED_PSEUDO_ELEMENTS:
                raise SelectorError(
                    next, f'Expected a supported pseudo-element, got {value}')
            return None, value
        elif next is not None and next.type == 'ident':
            name = next.lower_value
            if name in ('before', 'after', 'first-line', 'first-letter'):
                return None, name
            else:
                return PseudoClassSelector(name), None
        elif next is not None and next.type == 'function':
            name = next.lower_name
            if name in ('is', 'where', 'not', 'has'):
                return parse_logical_combination(next, namespaces, name), None
            else:
                return (
                    FunctionalPseudoClassSelector(name, next.arguments), None)
        else:
            raise SelectorError(next, f'unexpected {next} token.')
    else:
        return None, None

def parse_logical_combination(matches_any_token, namespaces, name):
    forgiving = True
    relative = False
    if name == 'is':
        selector_class = MatchesAnySelector
    elif name == 'where':
        selector_class = SpecificityAdjustmentSelector
    elif name == 'not':
        forgiving = False
        selector_class = NegationSelector
    elif name == 'has':
        relative = True
        selector_class = RelationalSelector

    selectors = [
        selector for selector in
        parse(matches_any_token.arguments, namespaces, forgiving, relative)
        if selector.pseudo_element is None]
    return selector_class(selectors)

def parse_attribute_selector(tokens, namespaces):
    tokens.skip_whitespace()
    qualified_name = parse_qualified_name(
        tokens, namespaces, is_attribute=True)
    if qualified_name is None:
        next = tokens.next()
        raise SelectorError(next, f'expected attribute name, got {next}')
    namespace, local_name = qualified_name

    tokens.skip_whitespace()
    peek = tokens.peek()
    if peek is None:
        operator = None
        value = None
    elif peek in ('=', '~=', '|=', '^=', '$=', '*='):
        operator = peek.value
        tokens.next()
        tokens.skip_whitespace()
        next = tokens.next()
        if next is None or next.type not in ('ident', 'string'):
            next_type = 'None' if next is None else next.type
            raise SelectorError(
                next, f'expected attribute value, got {next_type}')
        value = next.value
    else:
        raise SelectorError(
            peek, f'expected attribute selector operator, got {peek}')

    tokens.skip_whitespace()
    next = tokens.next()
    case_sensitive = None
    if next is not None:
        if next.type == 'ident' and next.value.lower() == 'i':
            case_sensitive = False
        elif next.type == 'ident' and next.value.lower() == 's':
            case_sensitive = True
        else:
            raise SelectorError(next, f'expected ], got {next.type}')
    return AttributeSelector(
        namespace, local_name, operator, value, case_sensitive)

def parse_qualified_name(tokens, namespaces, is_attribute=False):
    """Return ``(namespace, local)`` for given tokens.

    Can also return ``None`` for a wildcard.

    The empty string for ``namespace`` means "no namespace".

    """
    peek = tokens.peek()
    if peek is None:
        return None
    if peek.type == 'ident':
        first_ident = tokens.next()
        peek = tokens.peek()
        if peek != '|':
            namespace = '' if is_attribute else namespaces.get(None, None)
            return namespace, (first_ident.value, first_ident.lower_value)
        tokens.next()
        namespace = namespaces.get(first_ident.value)
        if namespace is None:
            raise SelectorError(
                first_ident,
                f'undefined namespace prefix: {first_ident.value}')
    elif peek == '*':
        next = tokens.next()
        peek = tokens.peek()
        if peek != '|':
            if is_attribute:
                raise SelectorError(
                    next, f'expected local name, got {next.type}')
            return namespaces.get(None, None), None
        tokens.next()
        namespace = None
    elif peek == '|':
        tokens.next()
        namespace = ''
    else:
        return None

    # If we get here, we just consumed '|' and set ``namespace``
    next = tokens.next()
    if next.type == 'ident':
        return namespace, (next.value, next.lower_value)
    elif next == '*' and not is_attribute:
        return namespace, None
    else:
        raise SelectorError(next, f'expected local name, got {next.type}')

class SelectorError(ValueError):
    """A specialized ``ValueError`` for invalid selectors."""

class TokenStream:
    def __init__(self, tokens):
        self.tokens = iter(tokens)
        self.peeked = []  # In reversed order

    def next(self):
        if self.peeked:
            return self.peeked.pop()
        else:
            return next(self.tokens, None)

    def peek(self):
        if not self.peeked:
            self.peeked.append(next(self.tokens, None))
        return self.peeked[-1]

    def skip(self, skip_types):
        found = False
        while 1:
            peek = self.peek()
            if peek is None or peek.type not in skip_types:
                break
            self.next()
            found = True
        return found

    def skip_whitespace(self):
        return self.skip(['whitespace'])

    def skip_comment(self):
        return self.skip(['comment'])

    def skip_whitespace_and_comment(self):
        return self.skip(['comment', 'whitespace'])

class Selector:
    def __init__(self, tree, pseudo_element=None):
        self.parsed_tree = tree
        self.pseudo_element = pseudo_element
        if pseudo_element is None:
            #: Tuple of 3 integers: http://www.w3.org/TR/selectors/#specificity
            self.specificity = tree.specificity
        else:
            a, b, c = tree.specificity
            self.specificity = a, b, c + 1

    def __repr__(self):
        pseudo = f'::{self.pseudo_element}' if self.pseudo_element else ''
        return f'{self.parsed_tree!r}{pseudo}'

class RelativeSelector:
    def __init__(self, combinator, selector):
        self.combinator = combinator
        self.selector = selector

    @property
    def specificity(self):
        return self.selector.specificity

    @property
    def pseudo_element(self):
        return self.selector.pseudo_element

    def __repr__(self):
        return (
            f'{self.selector!r}' if self.combinator == ' '
            else f'{self.combinator} {self.selector!r}')

class CombinedSelector:
    def __init__(self, left, combinator, right):
        #: Combined or compound selector
        self.left = left
        # One of `` `` (a single space), ``>``, ``+`` or ``~``.
        self.combinator = combinator
        #: compound selector
        self.right = right

    @property
    def specificity(self):
        a1, b1, c1 = self.left.specificity
        a2, b2, c2 = self.right.specificity
        return a1 + a2, b1 + b2, c1 + c2

    def __repr__(self):
        return f'{self.left!r}{self.combinator}{self.right!r}'

class CompoundSelector:
    def __init__(self, simple_selectors):
        self.simple_selectors = simple_selectors

    @property
    def specificity(self):
        if self.simple_selectors:
            # zip(*foo) turns [(a1, b1, c1), (a2, b2, c2), ...]
            # into [(a1, a2, ...), (b1, b2, ...), (c1, c2, ...)]
            return tuple(map(sum, zip(
                *(sel.specificity for sel in self.simple_selectors))))
        else:
            return 0, 0, 0

    def __repr__(self):
        return ''.join(map(repr, self.simple_selectors))

class LocalNameSelector:
    specificity = 0, 0, 1

    def __init__(self, local_name):
        self.local_name, self.lower_local_name = local_name

    def __repr__(self):
        return self.local_name

class NamespaceSelector:
    specificity = 0, 0, 0

    def __init__(self, namespace):
        #: The namespace URL as a string,
        #: or the empty string for elements not in any namespace.
        self.namespace = namespace

    def __repr__(self):
        if self.namespace == '':
            return '|'
        else:
            return f'{{{self.namespace}}}|'

class IDSelector:
    specificity = 1, 0, 0

    def __init__(self, ident):
        self.ident = ident

    def __repr__(self):
        return f'#{self.ident}'

class ClassSelector:
    specificity = 0, 1, 0

    def __init__(self, class_name):
        self.class_name = class_name

    def __repr__(self):
        return f'.{self.class_name}'

class AttributeSelector:
    specificity = 0, 1, 0

    def __init__(self, namespace, name, operator, value, case_sensitive):
        self.namespace = namespace
        self.name, self.lower_name = name
        #: A string like ``=`` or ``~=``, or None for ``[attr]`` selectors
        self.operator = operator
        #: A string, or None for ``[attr]`` selectors
        self.value = value
        #: ``True`` if case-sensitive, ``False`` if case-insensitive, ``None``
        #: if depends on the document language
        self.case_sensitive = case_sensitive

    def __repr__(self):
        namespace = '*|' if self.namespace is None else f'{{{self.namespace}}}'
        case_sensitive = (
            '' if self.case_sensitive is None else
            f' {"s" if self.case_sensitive else "i"}')
        return (
            f'[{namespace}{self.name}{self.operator}{self.value!r}'
            f'{case_sensitive}]')

class PseudoClassSelector:
    specificity = 0, 1, 0

    def __init__(self, name):
        self.name = name

    def __repr__(self):
        return ':' + self.name

class FunctionalPseudoClassSelector:
    specificity = 0, 1, 0

    def __init__(self, name, arguments):
        self.name = name
        self.arguments = arguments

    def __repr__(self):
        return f':{self.name}{tuple(self.arguments)!r}'

class NegationSelector:
    def __init__(self, selector_list):
        self.selector_list = selector_list

    @property
    def specificity(self):
        if self.selector_list:
            return max(selector.specificity for selector in self.selector_list)
        else:
            return (0, 0, 0)

    def __repr__(self):
        return f':not({", ".join(repr(sel) for sel in self.selector_list)})'

class RelationalSelector:
    def __init__(self, selector_list):
        self.selector_list = selector_list

    @property
    def specificity(self):
        if self.selector_list:
            return max(selector.specificity for selector in self.selector_list)
        else:
            return (0, 0, 0)

    def __repr__(self):
        return f':has({", ".join(repr(sel) for sel in self.selector_list)})'

class MatchesAnySelector:
    def __init__(self, selector_list):
        self.selector_list = selector_list

    @property
    def specificity(self):
        if self.selector_list:
            return max(selector.specificity for selector in self.selector_list)
        else:
            return (0, 0, 0)

    def __repr__(self):
        return f':is({", ".join(repr(sel) for sel in self.selector_list)})'

class SpecificityAdjustmentSelector:
    def __init__(self, selector_list):
        self.selector_list = selector_list

    @property
    def specificity(self):
        return (0, 0, 0)

    def __repr__(self):
        return f':where({", ".join(repr(sel) for sel in self.selector_list)})'


import numpy as np
from pytorch_transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch
import torch.nn.functional as F

from general_beamsearch import LanguageModel, LMScore # for writing a general LM

class GPTLanguageModel(LanguageModel):
    def __init__(self, model_path='gpt2', tokenizer_path='gpt2'):
        self.device = torch.device("cuda" if torch.cuda.is_available()  else "cpu")
        self.model = GPT2LMHeadModel.from_pretrained(model_path) # default loads from huggingface AWS servers. can cache if behind a firewall
        self.model.eval() # set model to eval mode because we're not fine tuning
        self.model.to(self.device)
        print('loaded model')

        self.tokenizer = GPT2Tokenizer.from_pretrained(tokenizer_path)
        print('loaded tokenizer')

    def encode(self, word):
        return self.tokenizer.encode(word)

    def decode(self, tokens):
        return self.tokenizer.decode(tokens)

    def score(self, context, words, past=None):
        '''
        Return log probabilities for extending context

        context - a list of GPT tokens that comprise the prior history
        words - a list of possible words.
        past - previous logits of GPT for that beam for efficiency. Currently not used, could be an efficiency gain

        return a list of LMScores in the same order as words
        '''

        if len(words) == 0:
            raise ValueError('No extensions proposed')
        # elif len(words) == 1:
        #   in theory, if there's only one possibility, we could skip the LM since the ranking won't change,
        #   but with the current implementation, we still need to run it because the score won't get accumulated otherwise,
        #   which could make a difference for comparing e.g. "saw" and "stared" in the sentence "I ___ it".
        #   This implementation could be improved by queuing up words with only one possibility and passing them to `words` at once

        # Can't run GPT on sentences with length 1. Prepend EOS
        if not context:
            context = [self.tokenizer.vocab_size - 1]
        # Limit the size of context because GPT can only handle 512 wordpieces
        context = context[-64:]

        tokenized_input = [self.encode(word) for word in words]

        # Pad tokenized_input with zeros to make it rectangular
        word_lengths = np.array([len(word) for word in tokenized_input]) # number of wordpieces in each candidate word
        np_input = np.zeros((len(word_lengths), max(word_lengths)), dtype=np.int32)
        for i, sent in enumerate(tokenized_input):
            np_input[i,:len(tokenized_input[i])] = tokenized_input[i]

        batch_size = 128 # batch up compute so we don't run out of GPU memory
        log_probs = np.zeros(np_input.shape[0]) # accumulate results in this

        with torch.no_grad():
            # Run context to get and cache past. This prevents runtime from growing as the context gets longer.
            # This could be saved from previous rounds, but that's probably unnecessary
            past_input = torch.tensor([context], device=self.device)
            past_logits, past = self.model(past_input)

            # Sort data by length
            order = np.argsort(-word_lengths)

            for batch_start in range(0, np_input.shape[0], batch_size):
                # print('Allocated memory inner loop', torch.cuda.memory_allocated(self.device), 'max allocated', torch.cuda.max_memory_allocated(self.device))
                batch_indices = order[batch_start:batch_start+batch_size]
                batch_lengths = word_lengths[batch_indices]
                batch_inputs = np_input[batch_indices,:max(batch_lengths)]
                torch_batch_lengths = torch.tensor(batch_lengths, dtype=torch.int32, device=self.device)

                true_batch_size = len(batch_lengths) # the actual size of the batch, accounting for the final batch being smaller

                # Evaluate on present state
                present_input = torch.tensor(batch_inputs, dtype=torch.long, device=self.device)

                # actually do GPT compute only if necessary. otherwise can just take it from final logits of previous step
                if max(batch_lengths) > 1:
                    # Run the model with each new possibility (batched). We do a lot of unnecessary computation here
                    # Every possibility is batched, even single word ones for which we know the answer
                    past_ = torch.stack(past)
                    past_ = past_.expand(-1,-1,true_batch_size,-1,-1,-1)
                    present_logits, _present = self.model(present_input, past=past_)

                    # Calculate loss for each possible completion.
                    logits = torch.cat((past_logits[:, -1, :].expand(true_batch_size, 1, -1), present_logits[:, :-1]), dim=1).transpose(1,2)
                    targets = present_input

                    loss_fcn = torch.nn.CrossEntropyLoss(reduction='none')
                    losses = loss_fcn(logits, targets)
                    for b in range(true_batch_size):
                        losses[b,batch_lengths[b]:] = 0
                    log_probs[batch_indices] = losses.sum(axis=1).detach().cpu().numpy()

                # skip extra GPT call for batches of words of length 1
                else:
                    logits = past_logits[:, -1, :].expand(true_batch_size, -1)
                    targets = present_input[:,0]

                    loss_fcn = torch.nn.CrossEntropyLoss(reduction='none')
                    losses = loss_fcn(logits, targets)
                    log_probs[batch_indices] = losses.detach().cpu().numpy()

        result = [LMScore(tokens=tokens, score=lm_prob) for tokens, lm_prob in zip(tokenized_input, -log_probs)]
        return result

def test():
    print('====== testing GPT lm ========')
    from gpt_lm import GPTLanguageModel
    from general_beamsearch import Lattice, Beam, beam_search

    l = Lattice()
    l.from_carmel_lattice('lattices/test2.lattice')
    gpt = GPTLanguageModel()
    Beam.tokenizer = gpt.tokenizer # for debugging
    beams = beam_search(gpt, l, beam_width=16)
    for b in beams:
        print(str(b))

if __name__ == '__main__':
    test()


import pandas as pd
import matplotlib
import seaborn as sns
import matplotlib.pyplot as plt
temp = pd.read_csv(r'tempYearly.csv')
rain = pd.read_csv(r'rainYearly.csv')
print(temp)
print(rain)

temp_f = temp.query('Temperature < 40 & Temperature > 0')
print(temp_f)
rain_f = rain.query('Rainfall < 6 & Rainfall > 0')
#temp_f.plot.scatter(x='Year',y='Temperature',label='Temperature and Year')
#plt.show()
#rain_f.plot.scttter(x='Year',y='Rainfall',label='Rainfall and Year')
#plt.show()
df_merge = pd.merge(temp_f,rain_f,on='Year',how = 'inner')
print(df_merge.sort_values(by = 'Rainfall',ascending = False))
print(df_merge.sort_values(by = 'Temperature',ascending = False))

sns.set(rc={'figure.figsize':(12,6)})
sns.jointplot('Rainfall', 'Temperature', data=df_merge,kind='reg')
plt.show()


# coding: utf-8
# Author: gjwei
class Solution(object):
    def maxProfit(self, k, prices):
        """
        :type k: int
        :type prices: List[int]
        :rtype: int
        """
        if len(prices) < 2:
            return 0

        sells = [0 for _ in range(k)]
        buys = [-(1 << 31) for _ in range(k)]

        for i in range(len(prices)):
            for j in range(k):
                if j == 0:
                    buys[j] = max(buys[j],  - prices[i])
                    sells[j] = max(sells[j], buys[j] + prices[i])
                else:
                    buys[j] = max(buys[j], sells[j - 1] - prices[i])
                    sells[j] = max(sells[j], buys[j] + prices[i])

        return sells[-1]

leetcode/188. Best Time to Buy and Sell Stock IV .py


import requests
from bs4 import BeautifulSoup
import sys

PATH = '/mnt/c/Users/Ryusuke/wsl_home/pro/html/shiryo/'

def main():
    # コマンドライン引数からURL取得
    url = sys.argv[1]

    # コマンドライン引数からファイル名取得
    filename = PATH + sys.argv[2]

    # HTML取得
    r = requests.get(url)

    # HTMLの書き込み
    with open(filename, mode='w') as f:
        f.write(r.text)

if __name__ == '__main__':
    main()


"""
To create mock data for this, do

python src/categorical_from_binary/experiments/sims/run_sims.py  --seed 0 --test_run true  --list_of_scales_for_predictive_categories  0.1 2.0
python src/categorical_from_binary/experiments/sims/run_sims.py  --seed 1 --test_run true  --list_of_scales_for_predictive_categories  0.1 2.0

This module should be run locally after scp'ing the data from `run_sims` to a local results_dir
"""

import numpy as np

np.set_printoptions(precision=3, suppress=True)

import seaborn as sns

sns.set(style="darkgrid")
from matplotlib import pyplot as plt

from categorical_from_binary.experiments.compare_to_mle.results_collector import (
    make_df_from_directory_taking_one_file_for_each_seed,
)

# results_dir="./tmp"
# need to scp -r mwojno01@login.cluster.tufts.edu:/cluster/tufts/hugheslab/mwojno01/data/results/sims/
# from results_dir_cluster to results_dir_local
results_dir_cluster = "/cluster/tufts/hugheslab/mwojno01/data/results/sims/"
results_dir_local = "data/results/sims_results_in_icml_draft/ib_probit/"
df = make_df_from_directory_taking_one_file_for_each_seed(
    results_dir_local, prefix_for_filenames_of_interest="sim_results_CAVI_vs_MLE_"
)
# TODO : time for IB CAVI increases with dataset sizes (up to 30 min!), whereas MLE does not do this...(never exeeds 1 min).
# Should we be concerned? Guess its not a big deal since MLE's problem is with smaller datasets, not bigger ones.  Still not ideal, though.

###
# now try to add in plotting stuff from the old module
###
x_label = "Mean covariate-conditional entropy"  # "Mean covariate-conditional category entropy"
y_label = "Mean holdout log likelihood"
outcome_of_interest_postfix = "log_like"
dir_for_saving_figures = "tmp/"
new_legend_title = "Prediction Method"
# Warning! The new labels need to be in the same order as the plot's default legend.  TODO: Fix this.

fig, axes = plt.subplots(
    1,
    2,
    figsize=(5, 3),
    sharey=True,
    sharex=True,
    squeeze=True,
    constrained_layout=True,
)
Ks = list(set(df["K"]))
Ns = list(set(df["N"]))
Ms = list(set(df["M"]))
K = Ks[0]
M = Ms[0]

legend = True
for (n, N) in enumerate(reversed(Ns)):
    ## select context of interest for this subplot
    dff = df.query(f"K == {K} & N == {N} & M == {M}")

    ## compute the mean covariate conditional entropy over seeds
    dff["mean_covariate_conditional_entropy_over_seeds"] = ""
    for scale in set(dff.scale_for_predictive_categories):
        mean_covariate_conditional_entropy_over_seeds = dff[
            dff.scale_for_predictive_categories == scale
        ].mean_covariate_conditional_entropy_of_true_category_probabilities.mean()
        dff.loc[
            dff.scale_for_predictive_categories == scale,
            "mean_covariate_conditional_entropy_over_seeds",
        ] = mean_covariate_conditional_entropy_over_seeds

    # mostly doing this so x axis isn't overly precise in the plot.
    dff = dff.round(decimals=3)

    dff = dff.loc[
        :,
        dff.columns.str.endswith(outcome_of_interest_postfix)
        | dff.columns.str.startswith("mean_covariate_conditional_entropy_over_seeds"),
    ]

    # 'melting' here basically takes a bunch of different columns and makes them different factors
    # for a single column.  So the dataframe gets longer and thinner.
    df_melted = dff.melt(
        "mean_covariate_conditional_entropy_over_seeds",
        var_name="method",
        value_name=outcome_of_interest_postfix,
    )

    g = sns.lineplot(
        x="mean_covariate_conditional_entropy_over_seeds",
        y=outcome_of_interest_postfix,
        hue="method",
        data=df_melted,
        legend=True,
        estimator="mean",
        ax=axes[n],
    )
    handles, labels = axes[n].get_legend_handles_labels()
    axes[n].legend().set_visible(False)
    axes[n].set(xlabel=None)
    axes[n].set(ylabel=None)

fig.supylabel(y_label)
fig.supxlabel(x_label)

# TODO: These should NOT be hardcoded!
cols = ["N=1P", "N=100P"]

pad = 6  # in points

for ax, col in zip(axes, cols):
    ax.annotate(
        col,
        xy=(0.5, 1),
        xytext=(0, pad),
        xycoords="axes fraction",
        textcoords="offset points",
        ha="center",
        va="baseline",
        fontsize=14,
    )

# ### Get labels from last axes and make legend
# handles, labels = axes[0].get_legend_handles_labels()
# h,l = ax.get_legend_handles_labels()

new_legend_labels = [
    "Generating Process",
    "Baserate Frequency",
    "Softmax (via MLE)",
    "CB-Probit (via IB-CAVI)",
]
new_labels = new_legend_labels

fig.legend(
    [handles[0], handles[3], handles[2], handles[1]],
    [new_labels[0], new_labels[3], new_labels[2], new_labels[1]],
    loc="upper left",
    bbox_to_anchor=(0.57, 0.405, 0.5, 0.1),
    fontsize="small",
)
plt.show()