import numpy as np
from astropy.io import fits

try:
  import matplotlib.pyplot as plt
  from matplotlib import colors
  mpl=True
except:
  print("Warning: matplotlib not found, plotting will not work...")
  mpl=False

try:
  from pyds9 import DS9
  pyds9=True
except:
  pyds9=False

fitter=None
fitters=[]
#Had issues getting iminuit to work, curve_fit works fine for now
try:
  import iminuit
  import probfit
  fitter="iminuit"
  fitters.append("iminuit")
  print("Using iminuit as fitter")
except:
  pass
try:
  from scipy.optimize import curve_fit
  if fitter is None:
    print("Iminuit not found, using scipy as fitter...")
    fitter="scipy"
    fitters.append("scipy")
except:
    print("Warning, no fitter found, cannot fit to noise...")

def mad(data):
  return np.median(np.abs(data-np.median(data)))

def plot_2d(data,cmap="rainbow",xlim=None, ylim=None, title=None, units="ADU", vmin=-50, vmax=None):
  if not mpl:
    return None, None
  data=np.array(data)
  fig=plt.figure()
  # mesh=plt.pcolormesh(data-np.min(data),cmap=plt.get_cmap(cmap), norm=colors.LogNorm())
  mesh=plt.pcolormesh(data,cmap=plt.get_cmap(cmap), norm=colors.SymLogNorm(10, vmin=vmin, vmax=vmax))

  ax=mesh.axes
  ax.set_aspect("equal")
  if xlim is not None:
    ax.set_xlim(right=xlim)
  if ylim is not None:
    ax.set_ylim(top=ylim)
  cbar=plt.colorbar()
  cbar.set_label(units)
  plt.xlabel("X"); plt.ylabel("Y");
  plt.title(title)
  plt.show(False)
  return fig, mesh

#For fitting
def gauss(x, norm,mean, sigma):
  return norm*np.exp(-(x-mean)**2/(2.*sigma**2))

#For fitting
def gauss_norm(x, mean, sigma):
  return 1./(sigma*(2*np.pi)**.5)*np.exp(-(x-mean)**2/(2.*sigma**2))

def fit_gauss(data, nbins=None, plot_fit=False, fitter=fitter, fit_range=None, rms_range_factor=8, *args, **kwargs):
  '''
  Wrapper function to perform binned fit of data to func. If given, will fit in the range "fit_range" (should be a tuple/array if [min,max]). Otherwise fits in the range median-rms_range_factor*MAD.
  '''
  if len(fitters) < 1:
    print("Error, no valid fitters loaded, exiting...")
    return None, None
  #Guess for initial parameters and to get a reasonable fitting range
  mean_guess=np.median(data)
  rms_guess=mad(data)
  if fit_range is None:
    fit_range=[mean_guess-rms_guess*rms_range_factor, mean_guess+rms_guess*rms_range_factor]
  fit_data=data[(data > fit_range[0]) & (data < fit_range[1])]

  #Get a reasonable number of bins
  if nbins is None:
    nbins=int((fit_range[1]-fit_range[0])/2)
    print("Bins for fit: " + str(nbins))
    if nbins < 40:
      print("Warning, number of bins is fairly small, fit may not be good...")

  if fitter not in fitters:
    print("Error, selected fitter is not loaded, loading another fitter")
    fitter=fitters[0]

  if fitter=="scipy":
    hist, bin_edges=np.histogram(fit_data, density=True, bins=nbins)
    bin_centres = (bin_edges[:-1] + bin_edges[1:])/2
    coeff, var_matrix = curve_fit(gauss, bin_centres, hist)

    if plot_fit and mpl:
      hist_fit = gauss(bin_centres, *coeff)
      fig=plt.figure()
      plt.plot(bin_centres, hist, label="Data")
      plt.plot(bin_centres, hist_fit, label="Fit")
      plt.show(False)
    #To keep compatibility between scipy fitter (which requires norm parameter A) and iminuit (which does not),
    #don't return the normalization coeff.
    coeff=coeff[1:]
    var_matrix=var_matrix[1:,1:]
    if plot_fit:
      return coeff, var_matrix, fig
    else:
      return coeff, var_matrix

  elif fitter=="iminuit":
    #ext_gauss=probfit.Extended(probfit.gaussian)
    #likelihood=probfit.BinnedLH(ext_gauss, fit_data,bins=nbins, extended=True)#, extended_bound=(-1000,1000))
    likelihood=probfit.BinnedLH(probfit.gaussian, fit_data,bins=nbins)#, extended_bound=(-1000,1000))
    #likelihood=probfit.UnbinnedLH(probfit.gaussian, fit_data, extended=False)
    minuit = iminuit.Minuit(likelihood,mean=mean_guess,sigma=rms_guess, print_level=1)
    minuit.migrad()
    #      minuit.minos()
    if plot_fit:
      fig=plt.figure()
      likelihood.draw(minuit)
    #Probably a better way to handle this
    coeff=[]
    coeff.append(minuit.values[0])
    coeff.append(minuit.values[1])
    var_matrix=[]
    var_matrix.append(minuit.matrix()[0])
    var_matrix.append(minuit.matrix()[1])
#    var_matrix.append(minuit.matrix()[2])
    if plot_fit:
      return coeff, var_matrix, fig
    else:
      return coeff, var_matrix
  else:
    print("Error, fitter not valid, exiting without doing fit...")
    return None, None

class SkipperImage:
  def _reset_derived(self):
    self.image_means=None
    self.image_rmses=None
    self.baseline=None
    self.charge_mask=None

    self.baseline_subtracted=False
    self.noise=None
    self.noise_coeffs=None
    self.noise_errs=None
    self.gen_charge=None
    self.image_gen_charge=None
    return True

  def __init__(self, fname, extension=0):
    self.hdu=fits.open(fname)[extension]
    self.header=self.hdu.header.copy()
    self.ndcms=self.header.get("NDCMS")
    self.nrows=self.header.get("NAXIS2")
    self.ncols=self.header.get("NAXIS1")
    self.ncols_phys=int(self.ncols/self.ndcms)
    self.npix=self.nrows*self.ncols_phys
    self.data=self.hdu.data
    #Total exposure time in seconds
    self.exptot=self.header.get("MEXP")/1000+self.header.get("MREAD")/1000
    self.nskips=self.ndcms
    #Default parameters
    self.pre_skips=0
    self.post_skips=0
    self.invert=False

    #Derived values
    self._reset_derived()

    self.figures=[]
    self.ds9=None

  def set_params(self,nskips=None, pre_skips=None, post_skips=None, invert=None, *args, **kwars):
    '''
    Simple handler function to change parameters for handling the skipper image. Returns number of changed parameters.
    '''
    changed=0
    if nskips is not None and nskips > 0 and nskips != self.nskips:
      self.nskips=nskips
      changed+=1

    if pre_skips is not None and pre_skips >= 0 and pre_skips != self.pre_skips:
      self.pre_skips=pre_skips
      changed+=1

    if post_skips is not None and post_skips >= 0 and post_skips != self.post_skips:
      self.post_skips=post_skips
      changed+=1

    if invert is not None and type(invert) is bool and invert != self.invert:
      self.invert=invert
      changed+=1

    total_skips=self.pre_skips+self.post_skips
    if total_skips >= self.ndcms:
      print("Warning, trying to skip more skips than exist, reducing skipped skips...")
      #Reduced both proportional to the original setting
      self.pre_skips=int((self.ndcms-1)*self.pre_skips/(total_skips))
      self.post_skips=int((self.ndcms-1)*self.post_skips/(total_skips))
      print("New pre_skips is: " + str(self.pre_skips) + ". New post_skips is: " +str(self.post_skips))

    if changed>0:
      #All our old computed values become useless, so clear them out
      self._reset_derived()

    return changed

  def combine_skips(self,recompute=False,*args, **kwargs):
    '''
    Creates an image from the raw skipper data, skipping the first "pre_skips" and last "post_skip" skips.
    '''
    new_params=self.set_params(*args, **kwargs)

    if new_params==0 and self.image_means is not None and not recompute:
      return False

    means=np.zeros((self.nrows,self.ncols_phys))
    rmses=np.zeros((self.nrows,self.ncols_phys))

    #Actually make the image
    for i in range(self.ncols_phys):
      xslice=np.s_[i*self.ndcms+self.pre_skips:(i+1)*self.ndcms-self.post_skips]
      means[:,i]=np.mean(self.data[:,xslice],axis=1)
      rmses[:,i]=np.std(self.data[:,xslice],axis=1)

    if self.invert:
      means=-1*means
    #Store the result
    self.image_means=means
    self.image_rmses=rmses
    self.baseline_subtracted=False

    return True

  def write_combined_fits(self,fname, *args, **kwargs):
    '''
    Write the combined skips to a .fits file.
    '''

    header=self.header.copy()
    header["NAXIS1"]=self.ncols/self.ndcms

    self.combine_skips(*args,**kwargs)

    mean_hdu=fits.PrimaryHDU(self.image_means)
    mean_hdu.header=header
    rms_hdu=fits.ImageHDU(self.image_rmses, header, "RMSES")

    hdul=fits.HDUList([mean_hdu, rms_hdu])
    hdul.writeto(fname,clobber=False)

  def compute_baseline(self, *args, **kwargs):
    '''
    Computes the baseline for the CCD image, as the median value along each row and column. Also forms the baseline per pixel as the sum of the baselines in x and y.
    '''
    new_image=self.combine_skips(*args, **kwargs)

    if self.baseline is not None and not new_image:
      return False
    self.baseline=np.zeros(self.image_means.shape)

    self.baseline_x=np.median(self.image_means-self.baseline, axis=0)
    self.baseline[:,:]+=self.baseline_x[np.newaxis,:]
    self.baseline_y=np.median(self.image_means-self.baseline, axis=1)
    self.baseline[:,:]+=self.baseline_y[:,np.newaxis]
    return True

  def subtract_baseline(self, *args, **kwargs):
    '''
    Subtracts baseline from the image. Computes the baseline and skipper means, if not already done.
    '''
    new_image=self.compute_baseline(*args, **kwargs)
    if self.baseline_subtracted and not new_image:
      return False

    self.image_means-=self.baseline
    self.baseline_subtracted=True
    return True

  def draw_image(self, cmap="spectral", *args,**kwargs):
    '''
    Function to draw the image after combining skips.
    '''

    if not mpl:
      print("Error, matplotlib not found, cannot draw image...")
      return False

    self.combine_skips(*args, **kwargs)
    if pyds9:
      if self.ds9 is None:
        self.ds9=DS9()
      print("Drawing to DS9")
      self.ds9.set_np2arr(self.image_means)
      self.ds9.set("scale histequ")
      return True
    else:
      print("Drawing using matplotlib")
      fig, mesh=plot_2d(self.image_means)
      self.figures.append(fig)
      return True

  def compute_charge_mask(self, *args, **kwargs):
    '''
    Not implemented
    '''
    self.compute_baseline(*args, **kwargs)

    if self.baseline_subtracted:
      pass

  def compute_noise(self,nbins=None, plot_fit=False,*args, **kwargs):
    '''
    Performs a binned fit to the data to estimate the pixel noise of the CCD image
    '''

    #self.combine_skips(*args, **kwargs)
    #Need to subtract the baseline for the fit to work properly
    #TODO: fix that so this doesn't need to be done (maybe split left/right sides of image? Estimate mu?)
    self.subtract_baseline(*args, **kwargs)
    plot_fit=plot_fit and mpl

    fit_results=fit_gauss(self.image_means.flatten(), nbins=nbins, plot_fit=plot_fit, *args, **kwargs)
    coeff=fit_results[0]
    var_matrix=fit_results[1]
    if plot_fit:
      self.figures.append(fit_results[2])
    self.noise=coeff[-1]
    self.noise_coeffs=coeff
    self.noise_errs=var_matrix

    return coeff, var_matrix

  def compute_gen_charge(self, *args, **kwargs):
    '''
    Estimates the amount of spurious charge generated in each pixel from the skipping procedure
    '''
    self.set_params(*args, **kwargs)
    self.image_gen_charge=np.zeros((self.nrows, self.ncols_phys))
    #TODO: reimplement this using slices to eliminate the loop
    for i in range(self.ncols_phys):
      x1=i*self.ndcms+self.pre_skips
      x2=(i+1)*self.ndcms-self.post_skips-1
      self.image_gen_charge[:,i]=self.data[:,x2]-self.data[:,x1]
    self.gen_charge=np.median(self.image_gen_charge)
    return self.gen_charge

  def compute_charge_loss(*args, **kwargs):
    self.set_params(*args, **kwargs)

  def compute_statistics(self, *args, **kwargs):
    self.combine_skips(*args, **kwargs)
    self.compute_gen_charge(*args, **kwargs)
    self.compute_charge_loss(*args, **kwargs)


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

import pretend

import pytest

from http11.parser import Callback, Error, HTTPParser, MessageType

def test_basic():
    parser = HTTPParser()
    assert parser.parse(b"GET / HTTP/1.1\r\n\r\n") == 18
    assert parser.finished
    assert not parser.errored

@pytest.mark.parametrize(
    ("callback", "expected"),
    [
        (Callback.request_method, b"GET"),
        (Callback.request_uri, b"/foobar/"),
        (Callback.http_version, b"HTTP/1.1"),
    ],
)
def test_element_callback_request(callback, expected):
    func = pretend.call_recorder(lambda data: None)

    parser = HTTPParser()
    parser.callback(callback, func)
    parser.parse(b"GET /foobar/ HTTP/1.1\r\n\r\n")

    assert parser.finished
    assert not parser.errored
    assert func.calls == [pretend.call(expected)]

@pytest.mark.parametrize(
    ("callback", "expected"),
    [
        (Callback.reason_phrase, b"OK"),
        (Callback.http_version, b"HTTP/1.1"),
    ],
)
def test_element_callback_response(callback, expected):
    func = pretend.call_recorder(lambda data: None)

    parser = HTTPParser()
    parser.callback(callback, func)
    parser.parse(b"HTTP/1.1 200 OK\r\n\r\n")

    assert parser.finished
    assert not parser.errored
    assert func.calls == [pretend.call(expected)]

def test_status_code_callback():
    func = pretend.call_recorder(lambda status_code: None)

    parser = HTTPParser()
    parser.callback(Callback.status_code, func)
    parser.parse(b"HTTP/1.1 200 OK\r\n\r\n")

    assert parser.finished
    assert not parser.errored
    assert func.calls == [pretend.call(200)]

def test_header_callback():
    func = pretend.call_recorder(lambda name, value: None)

    parser = HTTPParser()
    parser.callback(Callback.http_header, func)
    parser.parse(b"HTTP/1.1 200 OK\r\nFoo: Bar\r\n\r\n")

    assert parser.finished
    assert not parser.errored
    assert func.calls == [pretend.call(b"Foo", b"Bar")]

@pytest.mark.parametrize(
    ("callback", "data"),
    [
        (Callback.request_method, b"GET / HTTP/1.1\r\n\r\n"),
        (Callback.request_uri, b"GET / HTTP/1.1\r\n\r\n"),
        (Callback.http_version, b"GET / HTTP/1.1\r\n\r\n"),
        (Callback.reason_phrase, b"HTTP/1.1 200 OK\r\n\r\n"),
        (Callback.status_code, b"HTTP/1.1 200 OK\r\n\r\n"),
        (Callback.http_header, b"HTTP/1.1 200 OK\r\nFoo: Bar\r\n\r\n"),
    ],
)
def test_callback_error(callback, data):
    @pretend.call_recorder
    def raiser(*args, **kwargs):
        raise Exception

    parser = HTTPParser()
    parser.callback(callback, raiser)
    parser.parse(data)

    assert parser.finished
    assert parser.errored
    assert parser.error == Error.General

def test_error_is_none():
    parser = HTTPParser()

    assert not parser.finished
    assert not parser.errored
    assert parser.error is None

    parser.parse(b"HTTP/2.0 200 OK\r\n\r\n")

    assert parser.finished
    assert parser.errored
    assert parser.error == Error.BadVersion

def test_unknown_callback():
    parser = HTTPParser()

    with pytest.raises(ValueError):
        parser.callback("wat", lambda: None)

def test_callback_decorator():
    parser = HTTPParser()

    @parser.callback(Callback.http_version)
    @pretend.call_recorder
    def noop(*args, **kwargs):
        pass

    parser.parse(b"HTTP/1.1 200 OK\r\n\r\n")

    assert parser.finished
    assert not parser.errored

@pytest.mark.parametrize(
    ("message", "mtype"),
    [
        (b"GET / HTTP/1.1\r\n\r\n", MessageType.Request),
        (b"HTTP/1.1 200 OK\r\n\r\n", MessageType.Response),
    ],
)
def test_callback_message_type(message, mtype):
    parser = HTTPParser()
    assert parser.type is None
    parser.parse(message)
    assert parser.type == mtype

def test_message_reset():
    parser = HTTPParser()

    assert not parser.finished
    assert not parser.errored
    assert parser.error is None
    assert parser.type is None

    parser.parse(b"HTTP/1.1 200 OK\r\n\r\n")

    assert parser.finished
    assert not parser.errored
    assert parser.error is None
    assert parser.type is MessageType.Response

    parser.reset()

    assert not parser.finished
    assert not parser.errored
    assert parser.error is None
    assert parser.type is None

    parser.parse(b"GET / HTTP/1.1\r\n\r\n")

    assert parser.finished
    assert not parser.errored
    assert parser.error is None
    assert parser.type is MessageType.Request


from django import template
register = template.Library()

@register.filter(name='categorySelected')
def categorySelected(category,cat_name):
	if category == cat_name:
		return True
	else:
		return False

@register.filter(name='countLC')
def countLC(dict,id):
	newList = []
	for i in dict:
		if i.post.id == id:
			newList.append(i.post.id)
	return len(newList)

@register.filter(name='hasLike')
def hasLike(allLike,loggedInUserId):
	for like in allLike:
		if like.user.id == loggedInUserId:
			return True
	return False

@register.filter(name='hasFriend')
def hasFriend(friendList,loggedInUserId):
	for friend in friendList:
		if friend.receiver.id == loggedInUserId or friend.sender.id == loggedInUserId:
			return True
	return False

@register.filter(name='loggedInUserReqFriend')
def loggedInUserReqFriend(friendList,profile_id):
	for friend in friendList:
		if friend.sender.id == profile_id:
			return True
	return False

@register.filter(name='messageFilter')
def messageFilter(msgObj):
	mList = []


#python黑箱：输入与输出

#input() 函数暂停程序运行，同时等待键盘输入；直到回车被按下，函数的参数即为提示语，输入的类型永远是字符串型（str）。
#print() 函数则接受字符串、数字、字典、列表甚至一些自定义类的输出。

#生产环境中，强制类型转换加上try，except

#文件输入输出
#      NLP（自然语言处理）

#读取文件；
# 去除所有标点符号和换行符，并把所有大写变成小写；
# 合并相同的词，统计每个词出现的频率，并按照词频从大到小排序；
# 将结果按行输出到文件 out.txt。

import re

def parse(text):
    # 使用正则表达式去除标点符号和换行符
    text = re.sub(r'[^\w ]', ' ', text)

    # 转为小写
    text = text.lower()

    # 生成所有单词的列表
    word_list = text.split(' ')

    # 去除空白单词
    word_list = filter(None, word_list)

    # 生成单词和词频的字典
    word_cnt = {}
    for word in word_list:
        if word not in word_cnt:
            word_cnt[word] = 0
        word_cnt[word] += 1

    # 按照词频排序
    sorted_word_cnt = sorted(word_cnt.items(), key=lambda kv: kv[1], reverse=True)

    return sorted_word_cnt

with open('in.txt', 'r') as fin:
    text = fin.read()

word_and_freq = parse(text)

with open('out.txt', 'w') as fout:
    for word, freq in word_and_freq:
        fout.write('{} {}\n'.format(word, freq))

#JSON序列化

import json

params = {
    'symbol': '123456',
    'type': 'limit',
    'price': 123.4,
    'amount': 23
}

params_str = json.dumps(params)

print('after json serialization')
print('type of params_str = {}, params_str = {}'.format(type(params_str), params))

original_params = json.loads(params_str)

print('after json deserialization')
print('type of original_params = {}, original_params = {}'.format(type(original_params), original_params))
print("-----------------------------------------------------------------------------")

#文件JSON

params = {
    'symbol': '123456',
    'type': 'limit',
    'price': 123.4,
    'amount': 23
}

with open('params.json', 'w') as fout:
    params_str = json.dump(params, fout)

with open('params.json', 'r') as fin:
    original_params = json.load(fin)

print('after json deserialization')
print('type of original_params = {}, original_params = {}'.format(type(original_params), original_params))

#剩下两个问题，谷歌工具Protocol Buffer和  如何读取大信息量的文件

#读取大文件

large_size =100

with open("in.txt","r")as fin:
    text =""
    while True:
        str =fin.read(large_size)
        if not str:
            break
        text +=str
    # print(text)


import logging
import sys
from constants import *
from common import c_word_w, c_nof_complex, ceil_pow2, ceil_log2
from base_object import BaseObject
from numpy import mod

logger = logging.getLogger('main.periph.field')

class Field(BaseObject):
    """ A field defines data at certain address or an array of addresses
    """
    def __init__(self, name, settings=None):
        super().__init__()

        self.name(name)

        # self._valid_keys = ['width', 'bit_offset', 'access_mode', 'side_effect', 'address_offset', 'number_of_fields',
                            # 'reset_value', 'software_value', 'radix', 'field_description', 'field_name', 'default', 'user_width']
        self._valid_dict = {'width':{'max':32, 'min':1}, 'bit_offset':{'max':UNASSIGNED_BIT,'min':0}, 'access_mode':{}, 'interface':{}, 'side_effect':{},'address_offset': {'max':16384,'min':0, 'word_aligned':True},
                            'number_of_fields':{'max':262144,'min':'1'},'reset_value':{'max':131071},'software_value':{},'radix':{},'field_description':{},'field_name':{},'user_width':{'max':2048,'min':32}}
        self._args.update({'width'            : DEFAULT_WIDTH,
                           'bit_offset'       : DEFAULT_BIT_OFFSET,
                           'access_mode'      : DEFAULT_ACCESS_MODE,
                           'side_effect'      : DEFAULT_SIDE_EFFECT,
                           'address_offset'   : DEFAULT_ADDRESS_OFFSET,
                           'number_of_fields' : DEFAULT_NUMBER_OF_FIELDS,
                           'reset_value'      : RESET_VALUE,
                           'software_value'   : DEFAULT_SOFTWARE_VALUE,
                           'radix'            : DEFAULT_RADIX,
                           'field_description': DEFAULT_DESCRIPTION,
                           'interface'        : DEFAULT_INTERFACE,
                           'group_name'       : None})

        if settings is not None:
            # print(name)
            # print(settings)
            for key, val in settings.items():
                if key in self._valid_dict.keys(): #self._valid_keys:

                    if key == 'access_mode' and val.upper() not in VALID_ACCESS_MODES:
                        logger.error("Field.__init__(), Not a valid acces_mode '%s'", val)
                        self.success = False
                    elif key == 'side_effect' and val.upper() not in VALID_SIDE_EFFECTS:
                        logger.error("Field.__init__(), Not a valid side_effect '%s'", val)
                        self.success = False
                    else:
                        self.set_kv(key, val)
                else:
                    logger.error("Field.__init__(), Not a valid key '%s'", key)
                    self.success = False

    def __lt__(self, other):
        return self.address_offset() < other.address_offset()

    def number_of_fields(self, val=None):
        """ set/get number of fields """
        if val is not None:
            self.set_kv('number_of_fields', val)
            return
        return self._as_int('number_of_fields')

    def group_name(self, val=None):
        """ set/get group name """
        if val is not None:
            self.set_kv('group_name', val)
            return
        return self._as_str('group_name')

    def width(self, val=None):
        """ set/get width of field
        val: if not None set width of field
        return: actual width of field """
        if val is not None:
            return self.set_kv('width', val)
        _val = self._as_int('width')

        if not self.is_valid('width'):
            return None

        return self._as_int('width')

    def bit_offset(self, val=None):
        """ set/get bit_offset of field
        val: if not None set bit_offset of field
        return: actual bit_offset of field """
        if val is not None:
            return self.set_kv('bit_offset', val)
        if not self.is_valid('bit_offset'):
            return None
        return self._as_int('bit_offset')

    def access_mode(self, val=None):
        """ set/get access_mode of field
        val: if not None and a valid mode set access_mode of field
        return: actual access_mode of field """
        if val is not None:
            if val.upper() in VALID_ACCESS_MODES:
                return self.set_kv('access_mode', val.upper())
            else:
                logger.error("unknown access_mode '%s'", val)
                self.success = False
                return False
        return self._as_str('access_mode').upper()

    def side_effect(self, val=None):
        """ set/get side_effect of field
        val: if not None and a valid side_effect set side_effect of field
        return: actual side_effect of field """
        if val is not None:
            if val.upper() != 'NONE':
                vals = val.split(',')
                for _val in vals:
                    _val = _val.strip()
                    if _val.upper() not in VALID_SIDE_EFFECTS:
                        logger.error("unknown side_effect '%s'", _val)
                        self.success = False
                        return False
                return self.set_kv('side_effect', val.upper())
        return self._as_str('side_effect').upper()

    def address_offset(self, val=None):
        """ set/get address offset of field
        val: if not None set address offset of field
        return: actual address offset of field """
        if val is not None:
            return self.set_kv('address_offset', val)
        if not self.is_valid('address_offset'):
            return None
        return self._as_int('address_offset')

    def reset_value(self, val=None):
        """ set/get default hardware reset value of field
        val: if not None set default value of field
        return: active hardware reset value of field """
        if val is not None:
            return self.set_kv('reset_value', val)
        if not self.is_valid('reset_value'):
            return None
        return self._as_int('reset_value')

    def software_value(self, val=None):
        """ set/get software reset value of field
        val: if not None set default value of field
        return: active software reset value of field """
        if val is not None:
            return self.set_kv('software_value', val)
        if not self.is_valid('software_value'):
            return None
        return self._as_int('software_value')

    def radix(self, val=None):
        """ set/get radix value of field
        val: if not None set default value of field
        return: active radix value of field """
        if val is not None:
            if val.upper() in VALID_RADIXS:
                return self.set_kv('radix', val.upper())
            else:
                logger.error("unknown radix '%s'", val)
                self.success = False
                return False
        return self._as_int('radix')

    def field_description(self, val=None):
        """ set/get description of field
        val: if not None set description of field
        return: description of field """
        if val is not None:
            return self.set_kv('field_description', val)
        return self._as_str('field_description')

    def user_width(self, val=None):
        """ set/get width of RAM MM bus side interface  """
        if val is not None:
            self.set_kv('user_width', val)
            return
        return self._as_int('user_width', default=32)

    def interface(self, val=None):
        """ interface is used for rams to select between simple addr/data and full AXI4 """
        if val is not None:
            self.set_kv('interface', val)
            return
        return self._as_str('interface')

    def base_address(self, val = None):
        if val is not None:
            if mod(val, WIDTH_IN_BYTES):  # don't need check here if tool calcs are correct
                logger.error("Base address for field {} is not word aligned".format(field.name()))
            self.set_kv('base_addr', val)
            return
        return self._as_int('base_addr', default=0)

    def is_valid(self, yaml_key):
        """ Check value falls within min max range for keys that are ints """
        min = self._valid_dict[yaml_key].get('min', None)
        max = self._valid_dict[yaml_key].get('max', None)
        check_alignment = self._valid_dict[yaml_key].get('word_aligned', False)

        if yaml_key in ['software_value','reset_value']:
            min =  0
            max = int(1<<self._args.get('width', 32))-1

        _val = self._as_int(yaml_key)
        if not isinstance(_val, int):
            return True # may not have been evaluated from parameter yet, skip

        if check_alignment:
            if mod(_val, WIDTH_IN_BYTES):
                logger.error("Address offset for field {} is not word aligned".format(self.name()))
                sys.exit()

        if min is not None:
            if _val < min:
                logger.error("Field '{}': Value {} for '{}' key is outside of supported range, min is {}. Value reset to None".format(self.name(), _val, yaml_key, min))
                return False
        if max is not None:
            if _val > max:
                logger.error("Field '{}': Value {} for '{}' key is outside of supported range, max is {}. Value reset to None".format(self.name(), _val, yaml_key, max))
                return False
        return True

    # TODO: calc size in bytes
    # def get_byte_size(self):

    #def address_length(self, val=None):
    #    """ set/get address length of field
    #    val: if not None set address length of field
    #    return: active address_length of field """
    #    if val is not None:
    #        return self.set_kv('address_length', val)
    #    return self._as_int('address_length')


"""
    EXTERNAL LIBRARIES
"""

import sys
import numpy as np
import random
import matplotlib.pyplot as plt
import math

"""
    INTERNAL UTIL LIBRARIES
"""

# Add parent directory to the sys path
sys.path.append('../')
# Import any python modules from parent directory
from utils.SingleLayerNetwork import SingleLayerNetwork
from utils.auxfunctions import plotModel

"""
    The perceptron learning algorithm is suitable for learning linearly
    separable binary functions of binary inputs and is guaranteed to
    find a solution if the classes are linearly separable.

    Unfortunately, if the classes are not linearly separable, the
    algorithm may not even converge and is unlikely to produce the best
    solution when it does.
"""

class PerceptronLearningNetwork(SingleLayerNetwork):

    def __init__(self, numInputs, numOutputs, randomize=False, bias=False):
        """
            Construct a single-layer neural network given the number of input
            nodes and the number of output nodes. This network will train
            using the perceptron learning algorithm. The activation function
            used in this network will be a linear threshold unit where
            y = -1 if u <= 0
            y = +1 if u > 0     where u is the output of the weighted sum of inputs

            Input, output, and target values are assumed to be +- 1
        """
        # Define activation function
        def ltu(x):
            if x <= 0:
                return -1
            return 1
        # Parameters
        self.bias = bias
        self.activationFunc = ltu
        # Set weights according to randomize parameter
        if bias:
            self.weights = np.zeros((numOutputs, numInputs+1))
        else:
            self.weights = np.zeros((numOutputs, numInputs))
        if randomize:
            for i in range(self.weights.shape[0]):
                for j in range(self.weights.shape[1]):
                    self.weights[i, j] = random.random() * 2.0 - 1.0

    def train(self, x, y, epochs=1, lr=0.1, verbose=True):
        """
            Train the network given a batch of inputs, x, and their corresponding
            target outputs, y. Run the perceptron learning algorithm on the training
            set for the number of times specified by the epochs parameter.
        """
        numTrainingPoints = len(x)
        for e in range(epochs):
            # Set accuracy at beginning of epoch to 0s
            accuracy = 0
            # Compute the output for all training points
            allOutputs = self.evaluate(x)
            for i in range(numTrainingPoints):
                # Grab the inputs for the specific training point
                trainingPointInputs = x[i]
                # Grab the output for the specific training point
                trainingPointOutput = allOutputs[i]
                # Get the target outputs for the specific training point
                targets = y[i]
                # Compare each output 1 by 1
                for outputIndex in range(len(trainingPointOutput)):
                    # Grab specific output and corresponding target value
                    targetVal = targets[outputIndex]
                    outputVal = trainingPointOutput[outputIndex]
                    # If the outputs match, increment accuracy
                    if targetVal == outputVal:
                        accuracy += 1
                        continue
                    # Else, update the weights
                    else:
                        # For each input weight, compute its delta change, and then apply the change
                        for inputWeightIndex in range(len(self.weights[outputIndex])):
                            # If the inputWeightIndex is in the range of values for inputs, use the input at that index
                            if inputWeightIndex < len(trainingPointInputs):
                                trainingPointInput = trainingPointInputs[inputWeightIndex]
                            # Else, that value is the bias, and the input should be constant 1.0
                            else:
                                trainingPointInput = 1.0
                            # Compute delta w and apply the change
                            deltaW = 2.0 * lr * targetVal * trainingPointInput
                            self.weights[outputIndex, inputWeightIndex] += deltaW
            # Compute accuracy
            accuracy /= numTrainingPoints
            # If verbose == True, print accuracy for each training epoch
            if verbose:
                print('Epoch ' + str(e+1) + ' / ' + str(epochs) + ' Accuracy: ' + str(accuracy))
            # Return final accuracy
        return accuracy

# Test the network's learning capabilities
if __name__ == '__main__':
    """"
        Test the network on the AND and XOR Boolean functions. Show training
        progress in graphs.
    """
    # Params
    EPOCHS = 10
    NUM_GRAPHS = 10

    NUM_EPOCHS_PER_GRAPH = math.ceil(EPOCHS / NUM_GRAPHS)
    # Create AND training data
    and_train_x = np.array([
        [-1, -1],
        [-1, 1],
        [1, -1],
        [1, 1]
    ])
    and_train_y = np.array([
        [-1],
        [-1],
        [-1],
        [1]
    ])
    # Create XOR training data
    xor_train_x = np.array([
        [-1, -1],
        [-1, 1],
        [1, -1],
        [1, 1]
    ])
    xor_train_y = np.array([
        [-1],
        [1],
        [1],
        [-1]
    ])
    # Create networks
    and_pln = PerceptronLearningNetwork(2, 1, randomize=True, bias=True)
    xor_pln = PerceptronLearningNetwork(2, 1, randomize=True, bias=True)

    # Create the suplots
    fig, axes = plt.subplots(2, NUM_GRAPHS)
    fig.suptitle('Perceptron Learning Algorithm Training on AND and XOR Boolean Functions',
                 fontsize=18)

    # Label the subplots
    for x in range(2):
        for y in range(1, NUM_GRAPHS+1):
            if x == 0:
                axes[x][y-1].set_title('XOR Model at\n' + str(y * int(NUM_EPOCHS_PER_GRAPH)) + ' Epochs')
            else:
                axes[x][y-1].set_title('AND Model at\n' + str(y * int(NUM_EPOCHS_PER_GRAPH)) + ' Epochs')

    # Train each model on their respective datasets and show progress in the form
    # of decision boundaries for every N epochs
    for i in range(NUM_GRAPHS):
        # Train each model for N more epochs
        xor_pln.train(xor_train_x, xor_train_y, epochs=NUM_EPOCHS_PER_GRAPH, lr=0.05, verbose=False)
        and_pln.train(and_train_x, and_train_y, epochs=NUM_EPOCHS_PER_GRAPH, lr=0.05, verbose=False)

        # Plot the model
        plotModel(xor_pln, axes[0, i], [-1.2, 1.2, -1.2, 1.2], [-1.0, 1.0])
        axes[0, i].set_xlim(-1.2, 1.2)
        axes[0, i].set_ylim(-1.2, 1.2)
        axes[0, i].set_aspect('equal', adjustable='box')

        plotModel(and_pln, axes[1, i], [-1.2, 1.2, -1.2, 1.2], [-1.0, 1.0])
        axes[1, i].set_xlim(-1.2, 1.2)
        axes[1, i].set_ylim(-1.2, 1.2)
        axes[1, i].set_aspect('equal', adjustable='box')

        # Plot the actual data for each model
        axes[0, i].scatter([-1, 1], [1, -1], c='blue', s=150)
        axes[0, i].scatter([-1, 1], [-1, 1], c='red', s=150)

        axes[1, i].scatter([1], [1], c='blue', s=150)
        axes[1, i].scatter([-1, -1, 1], [-1, 1, -1], c='red', s=150)

    # Plot the graph
    figManager = plt.get_current_fig_manager()
    figManager.window.showMaximized()
    plt.show()

    """
        14 of the 16 Boolean functions are linearly separable.
        Let's show all of them and how Single-Layer Networks perform when
        using the Perceptron Learning Algorithm.
    """
    # Params
    EPOCHS_PER_MODEL = 20

    # Define class to hold info about each Boolean function
    class BooleanFunction:
        def __init__(self, true_x, false_x, name):
            self.true_x = true_x
            self.false_x = false_x
            self.name = name

        def getTrainingData(self):
            training_x = []
            training_y = []
            for pair in self.true_x:
                training_x.append(pair)
                training_y.append([1])
            for pair in self.false_x:
                training_x.append(pair)
                training_y.append([-1])
            return np.array(training_x), np.array(training_y)

    # Define each of the 16 boolean functions
    functions = []
    functions.append(
        BooleanFunction(
            [],
            [[-1, -1], [-1, 1], [1, -1], [1,1]],
            'FALSE'
        )
    )
    functions.append(
        BooleanFunction(
            [[1, 1]],
            [[-1, -1], [-1, 1], [1, -1]],
            'AND'
        )
    )
    functions.append(
        BooleanFunction(
            [[1, -1]],
            [[-1, -1], [-1, 1], [1, 1]],
            'A AND NOT B'
        )
    )
    functions.append(
        BooleanFunction(
            [[1, -1], [1, 1]],
            [[-1, -1], [-1, 1]],
            'A'
        )
    )
    functions.append(
        BooleanFunction(
            [[-1, 1]],
            [[-1, -1], [1, -1], [1, 1]],
            'NOT A AND B'
        )
    )
    functions.append(
        BooleanFunction(
            [[-1, 1], [1, 1]],
            [[-1, -1], [1, -1]],
            'B'
        )
    )
    functions.append(
        BooleanFunction(
            [[-1, 1], [1, -1]],
            [[-1, -1], [1, 1]],
            'XOR'
        )
    )
    functions.append(
        BooleanFunction(
            [[-1, 1], [1, -1], [1, 1]],
            [[-1, -1]],
            'OR'
        )
    )
    functions.append(
        BooleanFunction(
            [[-1, -1]],
            [[-1, 1], [1, -1], [1, 1]],
            'NOR'
        )
    )
    functions.append(
        BooleanFunction(
            [[-1, -1], [1, 1]],
            [[-1, 1], [1, -1]],
            'XNOR'
        )
    )
    functions.append(
        BooleanFunction(
            [[-1, -1], [1, -1]],
            [[-1, 1], [1, 1]],
            'NOT B'
        )
    )
    functions.append(
        BooleanFunction(
            [[-1, -1], [1, -1], [1, 1]],
            [[-1, 1]],
            'A OR NOT B'
        )
    )
    functions.append(
        BooleanFunction(
            [[-1, -1], [-1, 1]],
            [[1, -1], [1, 1]],
            'NOT A'
        )
    )
    functions.append(
        BooleanFunction(
            [[-1, -1], [-1, 1], [1, 1]],
            [[1, -1]],
            'NOT A OR B'
        )
    )
    functions.append(
        BooleanFunction(
            [[-1, -1], [-1, 1], [1, -1]],
            [[1, 1]],
            'NAND'
        )
    )
    functions.append(
        BooleanFunction(
            [[-1, -1], [-1, 1], [1, -1], [1, 1]],
            [],
            'TRUE'
        )
    )

    # Create the pyplot
    fig, axes = plt.subplots(2, 8)
    fig.suptitle('14 of 16 Boolean functions are linearly separable.\n' +
                'These graphs display decision boundaries formed by training\n' +
                'Single-Layer Networks using Perceptron Learning Algorithm.',
                fontsize=14)

    # Print out a warning to user that this may take a while, but progress will be output below
    print('\n\nWARNING: This may take a while. Progress will be output below:\n\n')

    # Create models for each of the functions, and plot the trained models
    f, x, y = 0, 0, 0
    for function in functions:
        # Create Single Layer Network
        model = PerceptronLearningNetwork(2, 1, randomize=True, bias=True)

        # Train the model on its dataset
        training_x, training_y = function.getTrainingData()
        model.train(training_x, training_y, epochs=EPOCHS_PER_MODEL, lr=0.5, verbose=False)

        # Print progress
        print('PROGRESS: ' + str(f+1) + ' / 16')

        # Plot the trained model in the appropriate subplot
        plotModel(model, axes[x, y], [-1.2, 1.2, -1.2, 1.2], valRange=[-1.0, 1.0])
        axes[x, y].set_xlim(-1.2, 1.2)
        axes[x, y].set_ylim(-1.2, 1.2)
        axes[x, y].set_aspect('equal', adjustable='box')
        axes[x, y].set_title(function.name)

        # Plot the points in the dataset
        for pair in function.true_x:
            axes[x, y].scatter(pair[0], pair[1], color='blue', s=150)
        for pair in function.false_x:
            axes[x, y].scatter(pair[0], pair[1], color='red', s=150)

        # Increment the trackers
        f += 1
        x += 1
        if x >= 2:
            x = 0
            y += 1

    # Plot the graph
    figManager = plt.get_current_fig_manager()
    figManager.window.showMaximized()
    plt.show()

3. Single-Layer Networks/3.4 Learning Rules for Single-Layer Networks/PerceptronLearningAlgorithm.py


from enum import IntEnum
import functools
from itertools import zip_longest
from queue import deque

DEBUG = False

class AccesMode(IntEnum):
    READ = 0
    WRITE = 1

class AddressingMode(IntEnum):
    POSITION = 0
    IMMEDIATE = 1

def debug_fun(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kargs):
        #if DEBUG:
        #    print(f'DEBUG fn:{fn.__name__}, args:{args}, kargs={kargs}')
        #    return fn(*args, **kargs)
        return fn(*args, **kargs)

    return wrapper

class IntCode:

    def __init__(self, prg, cpu_id=-1):
        self.pc = 0
        self.input = deque()
        self.output = deque()
        self.last_output = -1
        self.is_running = True
        self.is_waiting = False
        self.is_halted = False
        self.cpu_id = cpu_id
        self.prg = prg.copy()
        self.instructions = {
            1:(self.op_add, AccesMode.READ, AccesMode.READ, AccesMode.WRITE),
            2:(self.op_mul, AccesMode.READ, AccesMode.READ, AccesMode.WRITE),
            3:(self.op_input, AccesMode.WRITE, ),
            4:(self.op_output, AccesMode.READ,),
            5:(self.op_jmp_not_zero, AccesMode.READ, AccesMode.READ),
            6:(self.op_jmp_zero, AccesMode.READ, AccesMode.READ),
            7:(self.op_less_than, AccesMode.READ, AccesMode.READ, AccesMode.WRITE),
            8:(self.op_equals, AccesMode.READ, AccesMode.READ, AccesMode.WRITE),
            99:(self.op_halt,),
        }
        self.value_fetchers = {
            AddressingMode.POSITION:self.position_fetcher,
            AddressingMode.IMMEDIATE:self.immediate_fetcher,
        }
        self.flags = {
            'jmp':False
        }

    def __repr__(self):
        return f'<{self.__class__.__name__} #{self.cpu_id} with input {list(self.input)}'\
            f' and output {list(self.output)}' \
            f' is {"not " if not self.is_running else ""}running' \
            f', is {"not " if not self.is_halted else ""}halted'\
            f', is {"not " if not self.is_waiting else ""}waiting' \
            f', last output was {self.last_output}' \
            f', pc is {self.pc}'\
            f', prg is {self.prg}>'

    @staticmethod
    @debug_fun
    def decode_opcode(opcode):
        if opcode <= 99:
            return ([opcode, ])
        _s=str(opcode)
        # _s[-2:] gets 2 rightmost digits
        # _s[:-2][::-1] gets leftmost digits and reverse them
        return ([int(_s[-2:])] + list(map(lambda i:
                                          AddressingMode(int(i)),
                                          list(_s[:-2][::-1]))))

    @debug_fun
    def position_fetcher(self, addr):
        return self.prg[addr]

    @debug_fun
    def immediate_fetcher(self, val):
        return val

    @debug_fun
    def fn_fetch_value(self, mode):
        return self.value_fetchers[mode[0] + mode[1]]

    @debug_fun
    def fetch_values(self, modes):
        res = []
        for mode in modes:
            res.append(self.fn_fetch_value(mode)(self.prg[self.pc + 1]))
            self.inc()
        return res

    @debug_fun
    def inc(self, i=1):
        self.pc += i

    @debug_fun
    def jmp(self, addr):
        self.pc = addr

    @debug_fun
    def op_add(self, operands):
        v1, v2, addr = operands
        self.prg[addr] = v1 + v2
        self.inc()

    @debug_fun
    def op_mul(self, operands):
        v1, v2, addr = operands
        self.prg[addr] = v1 * v2
        self.inc()

    @debug_fun
    def op_input(self, operands):
        addr = operands.pop()
        if len(self.input) == 0:
            self.is_waiting = True
            self.inc(-1)
            return
        self.prg[addr] = self.input.popleft()
        self.inc()

    @debug_fun
    def op_output(self, operands):
        v = operands.pop()
        self.output.append(v)
        self.last_output = self.output[-1]
        self.inc()
        self.is_waiting = True

    @debug_fun
    def op_jmp_not_zero(self, operands):
        v1, addr = operands
        if v1 != 0:
            self.jmp(addr)
        else:
            self.inc()

    @debug_fun
    def op_jmp_zero(self, operands):
        v1, addr = operands
        if v1 == 0:
            self.jmp(addr)
        else:
            self.inc()

    @debug_fun
    def op_less_than(self, operands):
        v1, v2, addr = operands
        self.prg[addr] = 1 if v1 < v2 else 0
        self.inc()

    @debug_fun
    def op_equals(self, operands):
        v1, v2, addr = operands
        self.prg[addr] = 1 if v1 == v2 else 0
        self.inc()

    @debug_fun
    def op_halt(self, _):
        self.is_running = False
        self.is_halted = True

    def run(self, input):
        self.input = deque(input)
        if self.is_halted:
            print("IntCode computer is halted.")
            return
        self.is_waiting = False
        if DEBUG: print(f'Running {self}')
        while self.is_running and not self.is_waiting:
            self.flags['jmp'] = False
            # Get OpCode and addressing modes
            opcode, *addressing_modes=IntCode.decode_opcode(self.prg[self.pc])
            # Get operation and its operands access modes
            op, *opmodes = self.instructions[opcode]
            if DEBUG: print(80*'*')
            if DEBUG: print(f'DEBUG CPU #{self.cpu_id} pc:{self.pc}, prg:{self.prg}')
            if DEBUG: print(f'DEBUG CPU #{self.cpu_id} opcode: {opcode}, op:{op.__name__}')
            if DEBUG: print(f'DEBUG CPU #{self.cpu_id} opcode:{opmodes}, addressing_modes:{addressing_modes}')
            # Associate each operand acces mode with its addressing mode
            modes = list(zip_longest(opmodes, addressing_modes,
                                     fillvalue=AddressingMode.POSITION))
            if DEBUG: print(f'DEBUG CPU #{self.cpu_id} mode:{modes}')
            op(self.fetch_values(modes))
            if DEBUG: print(f'State {self}')
        return self.last_output

if __name__ == '__main__':
    # Day 2
    test='Day 2 - Test 1'
    IntCode([1,9,10,3,2,3,11,0,99,30,40,50]).run([])
    print(f'{test} OK')
    test='Day 2 - Test 2'
    IntCode([1,0,0,0,99]).run([])
    print(f'{test} OK')
    test='Day 2 - Test 3'
    IntCode([2,3,0,3,99]).run([])
    print(f'{test} OK')
    test='Day 2 - Test 4'
    IntCode([2,4,4,5,99,0]).run([])
    print(f'{test} OK')
    test='Day 2 - Test 5'
    IntCode([1,1,1,4,99,5,6,0,99]).run([])
    print(f'{test} OK')

    # Day 5 tests
    # Equals 8 ?
    test='Day 5 - Test 1.1 - Input equals 8? (Position mode)'
    assert 1 == IntCode([3,9,8,9,10,9,4,9,99,-1,8]).run([8]), f'{test} KO'
    print(f'{test} OK')
    test='Day 5 - Test 1.2 - Input equals 8? (Position mode)'
    assert 0 == IntCode([3,9,8,9,10,9,4,9,99,-1,8]).run([5]), f'{test} KO'
    print(f'{test} OK')

    # Less than 8 ?
    test='Day 5 - Test 2.1 - Input less than 8? (Position mode)'
    assert 0 == IntCode([3,9,7,9,10,9,4,9,99,-1,8]).run([8]), f'{test} KO'
    print(f'{test} OK')
    test='Day 5 - Test 2.2 - Input less than 8? (Position mode)'
    assert 1 == IntCode([3,9,7,9,10,9,4,9,99,-1,8]).run([5]), f'{test} KO'
    print(f'{test} OK')

    # Equals 8 ?
    test='Day 5 - Test 3.1 - Input equals 8? (Immediate mode)'
    assert 1 == IntCode([3,3,1108,-1,8,3,4,3,99]).run([8]), f'{test} KO'
    print(f'{test} OK')
    test='Day 5 - Test 3.2 - Input equals 8? (Immediate mode)'
    assert 0 == IntCode([3,3,1108,-1,8,3,4,3,99]).run([5]), f'{test} KO'
    print(f'{test} OK')

    # Less than 8 ?
    test='Day 5 - Test 4.1 - Input less than 8? (Immediate mode)'
    assert 0 == IntCode([3,3,1107,-1,8,3,4,3,99]).run([8]), f'{test} KO'
    print(f'{test} OK')
    test='Day 5 - Test 4.2 - Input less than 8? (Immediate mode)'
    assert 1 == IntCode([3,3,1107,-1,8,3,4,3,99]).run([5]), f'{test} KO'
    print(f'{test} OK')

    # Zero ?
    test='Day 5 - Test 5.1 - Input is 0? (Position mode)'
    assert 0 == IntCode([3,12,6,12,15,1,13,14,13,4,13,99,-1,0,1,9]).run([0]), f'{test} KO'
    print(f'{test} OK')
    test='Day 5 - Test 5.2 - Input is 0? (Position mode)'
    assert 1 == IntCode([3,12,6,12,15,1,13,14,13,4,13,99,-1,0,1,9]).run([1]), f'{test} KO'
    print(f'{test} OK')

    # Zero ?
    test='Day 5 - Test 6.1 - Input is 0? (Immediate mode)'
    assert 0 == IntCode([3,3,1105,-1,9,1101,0,0,12,4,12,99,1]).run([0]), f'{test} KO'
    print(f'{test} OK')
    test='Day 5 - Test 6.2 - Input is 0? (Immediate mode)'
    assert 1 == IntCode([3,3,1105,-1,9,1101,0,0,12,4,12,99,1]).run([1]), f'{test} KO'
    print(f'{test} OK')

    # Compare to 8
    test='Day 5 - Test 7.1 - Compare to 8'
    assert 999 == IntCode([3,21,1008,21,8,20,1005,20,22,107,8,21,20,1006,20,31,
             1106,0,36,98,0,0,1002,21,125,20,4,20,1105,1,46,104,
             999,1105,1,46,1101,1000,1,20,4,20,1105,1,46,98,99]).run([7]), f'{test} KO'
    print(f'{test} OK')
    test='Day 5 - Test 7.2 - Compare to 8'
    assert 1000 == IntCode([3,21,1008,21,8,20,1005,20,22,107,8,21,20,1006,20,31,
             1106,0,36,98,0,0,1002,21,125,20,4,20,1105,1,46,104,
             999,1105,1,46,1101,1000,1,20,4,20,1105,1,46,98,99]).run([8]), f'{test} KO'
    print(f'{test} OK')
    test='Day 5 - Test 7.3 - Compare to 8'
    assert 1001 == IntCode([3,21,1008,21,8,20,1005,20,22,107,8,21,20,1006,20,31,
             1106,0,36,98,0,0,1002,21,125,20,4,20,1105,1,46,104,
             999,1105,1,46,1101,1000,1,20,4,20,1105,1,46,98,99]).run([9]), f'{test} KO'
    print(f'{test} OK')

    # Amplifier
    test='Day 7 - Test 1.1 - Amplifier Feedback'
    assert 5 == IntCode([3,26,1001,26,-4,26,3,27,1002,27,2,27,1,27,26,
                         27,4,27,1001,28,-1,28,1005,28,6,99,0,0,5]).run([9, 0]), f'{test} KO'
    assert 7 == IntCode([3,26,1001,26,-4,26,3,27,1002,27,2,27,1,27,26,
                         27,4,27,1001,28,-1,28,1005,28,6,99,0,0,5]).run([9, 1]), f'{test} KO'


r"""Pipeline for loading VCF files to BigQuery.

Run locally:
python -m gcp_variant_transforms.vcf_to_bq \
  --input_pattern <path to VCF file(s)> \
  --output_table projectname:bigquerydataset.tablename

Run on Dataflow:
python -m gcp_variant_transforms.vcf_to_bq \
  --input_pattern gs://bucket/vcfs/vcffile.vcf \
  --output_table projectname:bigquerydataset.tablename \
  --project projectname \
  --staging_location gs://bucket/staging \
  --temp_location gs://bucket/temp \
  --job_name vcf-to-bq \
  --setup_file ./setup.py \
  --runner DataflowRunner
"""

from __future__ import absolute_import

import argparse  # pylint: disable=unused-import
import logging
import sys
import tempfile
from typing import List, Optional  # pylint: disable=unused-import

import apache_beam as beam
from apache_beam import pvalue  # pylint: disable=unused-import
from apache_beam.io import filesystems
from apache_beam.options import pipeline_options

from gcp_variant_transforms import pipeline_common
from gcp_variant_transforms.beam_io import vcfio
from gcp_variant_transforms.libs import metrics_util
from gcp_variant_transforms.libs import processed_variant
from gcp_variant_transforms.libs import vcf_header_parser
from gcp_variant_transforms.libs import variant_partition
from gcp_variant_transforms.libs.annotation.vep import vep_runner_util
from gcp_variant_transforms.libs.variant_merge import merge_with_non_variants_strategy
from gcp_variant_transforms.libs.variant_merge import move_to_calls_strategy
from gcp_variant_transforms.libs.variant_merge import variant_merge_strategy  # pylint: disable=unused-import
from gcp_variant_transforms.options import variant_transform_options
from gcp_variant_transforms.transforms import annotate_files
from gcp_variant_transforms.transforms import sample_info_to_bigquery
from gcp_variant_transforms.transforms import combine_call_names
from gcp_variant_transforms.transforms import densify_variants
from gcp_variant_transforms.transforms import extract_input_size
from gcp_variant_transforms.transforms import filter_variants
from gcp_variant_transforms.transforms import infer_headers
from gcp_variant_transforms.transforms import merge_headers
from gcp_variant_transforms.transforms import merge_variants
from gcp_variant_transforms.transforms import partition_variants
from gcp_variant_transforms.transforms import variant_to_avro
from gcp_variant_transforms.transforms import variant_to_bigquery
from gcp_variant_transforms.transforms import write_variants_to_shards

_COMMAND_LINE_OPTIONS = [
    variant_transform_options.VcfReadOptions,
    variant_transform_options.AvroWriteOptions,
    variant_transform_options.BigQueryWriteOptions,
    variant_transform_options.AnnotationOptions,
    variant_transform_options.FilterOptions,
    variant_transform_options.MergeOptions,
    variant_transform_options.PartitionOptions,
    variant_transform_options.ExperimentalOptions,
]

_ESTIMATE_SIZES_FILE_NAME = 'estimate-sizes'
_ESTIMATE_SIZES_JOB_NAME = 'estimate-input-size'
_MERGE_HEADERS_FILE_NAME = 'merged_headers.vcf'
_MERGE_HEADERS_JOB_NAME = 'merge-vcf-headers'
_ANNOTATE_FILES_JOB_NAME = 'annotate-files'
_SHARD_VCF_FILES_JOB_NAME = 'shard-files'
_SHARDS_FOLDER = 'shards'
_GCS_RECURSIVE_WILDCARD = '**'

def _read_variants(all_patterns,  # type: List[str]
                   pipeline,  # type: beam.Pipeline
                   known_args,  # type: argparse.Namespace
                   pipeline_mode  # type: int
                  ):
  # type: (...) -> pvalue.PCollection
  """Helper method for returning a PCollection of Variants from VCFs."""
  representative_header_lines = None
  if known_args.representative_header_file:
    representative_header_lines = vcf_header_parser.get_metadata_header_lines(
        known_args.representative_header_file)
  return pipeline_common.read_variants(
      pipeline,
      all_patterns,
      pipeline_mode,
      known_args.allow_malformed_records,
      representative_header_lines,
      vcfio.VcfParserType[known_args.vcf_parser])

def _get_variant_merge_strategy(known_args  # type: argparse.Namespace
                               ):
  # type: (...) -> Optional[variant_merge_strategy.VariantMergeStrategy]
  merge_options = variant_transform_options.MergeOptions
  if (not known_args.variant_merge_strategy or
      known_args.variant_merge_strategy == merge_options.NONE):
    return None
  elif known_args.variant_merge_strategy == merge_options.MOVE_TO_CALLS:
    return move_to_calls_strategy.MoveToCallsStrategy(
        known_args.info_keys_to_move_to_calls_regex,
        known_args.copy_quality_to_calls,
        known_args.copy_filter_to_calls)
  elif (known_args.variant_merge_strategy ==
        merge_options.MERGE_WITH_NON_VARIANTS):
    return merge_with_non_variants_strategy.MergeWithNonVariantsStrategy(
        known_args.info_keys_to_move_to_calls_regex,
        known_args.copy_quality_to_calls,
        known_args.copy_filter_to_calls)
  else:
    raise ValueError('Merge strategy is not supported.')

def _add_inferred_headers(all_patterns,  # type: List[str]
                          pipeline,  # type: beam.Pipeline
                          known_args,  # type: argparse.Namespace
                          merged_header,  # type: pvalue.PCollection
                          pipeline_mode  # type: int
                         ):
  # type: (...) -> pvalue.PCollection
  annotation_fields_to_infer = (known_args.annotation_fields if
                                known_args.infer_annotation_types else [])
  inferred_headers = (
      _read_variants(all_patterns,
                     pipeline,
                     known_args,
                     pipeline_mode)
      | 'FilterVariants' >> filter_variants.FilterVariants(
          reference_names=known_args.reference_names)
      | 'InferHeaderFields' >> infer_headers.InferHeaderFields(
          pvalue.AsSingleton(merged_header),
          known_args.allow_incompatible_records,
          known_args.infer_headers,
          annotation_fields_to_infer))
  merged_header = (
      (inferred_headers, merged_header)
      | 'FlattenHeaders' >> beam.Flatten()
      | 'MergeHeadersFromVcfAndVariants' >> merge_headers.MergeHeaders(
          known_args.split_alternate_allele_info_fields,
          known_args.allow_incompatible_records))
  return merged_header

def _shard_variants(known_args, pipeline_args, pipeline_mode):
  # type: (argparse.Namespace, List[str], int) -> List[str]
  """Reads the variants and writes them to VCF shards.

  Returns:
   The VCF shards directory.
  """
  options = pipeline_options.PipelineOptions(pipeline_args)
  google_cloud_options = options.view_as(pipeline_options.GoogleCloudOptions)
  shard_files_job_name = pipeline_common.generate_unique_name(
      _SHARD_VCF_FILES_JOB_NAME)
  _update_google_cloud_job_name(google_cloud_options, shard_files_job_name)
  vcf_shards_output_dir = filesystems.FileSystems.join(
      known_args.annotation_output_dir, _SHARDS_FOLDER)
  with beam.Pipeline(options=options) as p:
    variants = _read_variants(
        known_args.all_patterns, p, known_args, pipeline_mode)
    call_names = (variants
                  | 'CombineCallNames' >>
                  combine_call_names.CallNamesCombiner())
    _ = (variants
         | 'DensifyVariants' >> densify_variants.DensifyVariants(
             beam.pvalue.AsSingleton(call_names))
         | 'WriteToShards' >> write_variants_to_shards.WriteToShards(
             vcf_shards_output_dir,
             beam.pvalue.AsSingleton(call_names),
             known_args.number_of_variants_per_shard))

  return [vep_runner_util.format_dir_path(vcf_shards_output_dir) +
          _GCS_RECURSIVE_WILDCARD]

def _get_input_dimensions(known_args, pipeline_args):
  pipeline_mode = pipeline_common.get_pipeline_mode(
      known_args.all_patterns,
      known_args.optimize_for_large_inputs)
  beam_pipeline_options = pipeline_options.PipelineOptions(pipeline_args)
  google_cloud_options = beam_pipeline_options.view_as(
      pipeline_options.GoogleCloudOptions)

  estimate_sizes_job_name = pipeline_common.generate_unique_name(
      _ESTIMATE_SIZES_JOB_NAME)
  if google_cloud_options.job_name:
    google_cloud_options.job_name += '-' + estimate_sizes_job_name
  else:
    google_cloud_options.job_name = estimate_sizes_job_name
  temp_directory = google_cloud_options.temp_location or tempfile.mkdtemp()
  temp_estimated_input_size_file_name = '-'.join(
      [google_cloud_options.job_name,
       _ESTIMATE_SIZES_FILE_NAME])
  temp_estimated_input_size_file_path = filesystems.FileSystems.join(
      temp_directory, temp_estimated_input_size_file_name)
  with beam.Pipeline(options=beam_pipeline_options) as p:
    estimates = pipeline_common.get_estimates(
        p, pipeline_mode, known_args.all_patterns)

    files_size = (estimates
                  | 'GetFilesSize' >> extract_input_size.GetFilesSize())
    file_count = (estimates
                  | 'CountAllFiles' >> beam.combiners.Count.Globally())
    sample_map = (estimates
                  | 'ExtractSampleMap' >> extract_input_size.GetSampleMap())
    estimated_value_count = (sample_map
                             | extract_input_size.GetEstimatedValueCount())
    estimated_sample_count = (sample_map
                              | extract_input_size.GetEstimatedSampleCount())
    estimated_variant_count = (estimates
                               | 'GetEstimatedVariantCount'
                               >> extract_input_size.GetEstimatedVariantCount())
    _ = (estimated_variant_count
         | beam.ParDo(extract_input_size.print_estimates_to_file,
                      beam.pvalue.AsSingleton(estimated_sample_count),
                      beam.pvalue.AsSingleton(estimated_value_count),
                      beam.pvalue.AsSingleton(files_size),
                      beam.pvalue.AsSingleton(file_count),
                      temp_estimated_input_size_file_path))

  with filesystems.FileSystems.open(temp_estimated_input_size_file_path) as f:
    estimates = f.readlines()
  if len(estimates) != 5:
    raise ValueError('Exactly 5 estimates were expected in {}.'.format(
        temp_estimated_input_size_file_path))

  known_args.estimated_variant_count = int(estimates[0].strip())
  known_args.estimated_sample_count = int(estimates[1].strip())
  known_args.estimated_value_count = int(estimates[2].strip())
  known_args.files_size = int(estimates[3].strip())
  known_args.file_count = int(estimates[4].strip())

def _annotate_vcf_files(all_patterns, known_args, pipeline_args):
  # type: (List[str], argparse.Namespace, List[str]) -> str
  """Annotates the VCF files using VEP.

  Returns:
    The annotated VCF files directory.
  """
  options = pipeline_options.PipelineOptions(pipeline_args)
  google_cloud_options = options.view_as(pipeline_options.GoogleCloudOptions)
  annotate_files_job_name = pipeline_common.generate_unique_name(
      _ANNOTATE_FILES_JOB_NAME)
  _update_google_cloud_job_name(google_cloud_options, annotate_files_job_name)

  with beam.Pipeline(options=options) as p:
    _ = (p
         | beam.Create(all_patterns)
         | 'AnnotateShards' >> beam.ParDo(
             annotate_files.AnnotateFile(known_args, pipeline_args)))
  if known_args.annotation_fields:
    known_args.annotation_fields.append(known_args.vep_info_field)
  else:
    known_args.annotation_fields = [known_args.vep_info_field]
  # TODO(bashir2): The VEP runner by default runs VEP with --allele_number hence
  # we turn on this feature here. However, this might be inconsistent with other
  # annotation fields that are originally present in input files, if they do not
  # have ALLELE_NUM annotation. The fix is to make annotation ALT matching
  # smarter to fall back on other matching methods if ALLELE_NUM is not present.
  # When this is implemented, we may even consider removing use_allele_num flag
  # and always start by checking if ALLELE_NUM is present.
  known_args.use_allele_num = True
  return vep_runner_util.get_output_pattern(known_args.annotation_output_dir)

def _update_google_cloud_job_name(google_cloud_options, job_name):
  if google_cloud_options.job_name:
    google_cloud_options.job_name += '-' + job_name
  else:
    google_cloud_options.job_name = job_name

def _merge_headers(known_args, pipeline_args,
                   pipeline_mode, annotated_vcf_pattern=None):
  # type: (str, argparse.Namespace, List[str], int, str) -> None
  """Merges VCF headers using beam based on pipeline_mode."""
  if known_args.representative_header_file:
    return
  options = pipeline_options.PipelineOptions(pipeline_args)

  # Always run pipeline locally if data is small.
  if (pipeline_mode == pipeline_common.PipelineModes.SMALL and
      not known_args.infer_headers and not known_args.infer_annotation_types):
    options.view_as(pipeline_options.StandardOptions).runner = 'DirectRunner'

  google_cloud_options = options.view_as(pipeline_options.GoogleCloudOptions)
  merge_headers_job_name = pipeline_common.generate_unique_name(
      _MERGE_HEADERS_JOB_NAME)
  if google_cloud_options.job_name:
    google_cloud_options.job_name += '-' + merge_headers_job_name
  else:
    google_cloud_options.job_name = merge_headers_job_name

  temp_directory = google_cloud_options.temp_location or tempfile.mkdtemp()
  temp_merged_headers_file_name = '-'.join([google_cloud_options.job_name,
                                            _MERGE_HEADERS_FILE_NAME])
  temp_merged_headers_file_path = filesystems.FileSystems.join(
      temp_directory, temp_merged_headers_file_name)

  with beam.Pipeline(options=options) as p:
    headers = pipeline_common.read_headers(
        p, pipeline_mode, known_args.all_patterns)
    merged_header = pipeline_common.get_merged_headers(
        headers,
        known_args.split_alternate_allele_info_fields,
        known_args.allow_incompatible_records)
    if annotated_vcf_pattern:
      merged_header = pipeline_common.add_annotation_headers(
          p, known_args, pipeline_mode, merged_header,
          annotated_vcf_pattern)
    if known_args.infer_headers or known_args.infer_annotation_types:
      infer_headers_input_pattern = (
          [annotated_vcf_pattern] if
          annotated_vcf_pattern else known_args.all_patterns)
      merged_header = _add_inferred_headers(infer_headers_input_pattern, p,
                                            known_args, merged_header,
                                            pipeline_mode)

    pipeline_common.write_headers(merged_header, temp_merged_headers_file_path)
    known_args.representative_header_file = temp_merged_headers_file_path

def _validate_annotation_pipeline_args(known_args, pipeline_args):
  match_results = filesystems.FileSystems.match(['{}*'.format(
      vep_runner_util.format_dir_path(known_args.annotation_output_dir))])
  if match_results and match_results[0].metadata_list:
    raise ValueError('Output directory {} already exists.'.format(
        known_args.annotation_output_dir))

  flags_dict = pipeline_options.PipelineOptions(pipeline_args).get_all_options()
  expected_flags = ['max_num_workers', 'num_workers']
  for flag in expected_flags:
    if flag in flags_dict and flags_dict[flag] > 0:
      return
  raise ValueError('Could not find any of {} with a valid value among pipeline '
                   'flags {}'.format(expected_flags, flags_dict))

def _run_annotation_pipeline(known_args, pipeline_args):
  # type: (argparse.Namespace, List[str]) -> str
  annotated_vcf_pattern = None
  if known_args.run_annotation_pipeline:
    _validate_annotation_pipeline_args(known_args, pipeline_args)
    known_args.omit_empty_sample_calls = True

    files_to_be_annotated = known_args.all_patterns
    if known_args.shard_variants:
      pipeline_mode = pipeline_common.get_pipeline_mode(
          files_to_be_annotated,
          known_args.optimize_for_large_inputs)
      files_to_be_annotated = _shard_variants(known_args,
                                              pipeline_args,
                                              pipeline_mode)
    annotated_vcf_pattern = _annotate_vcf_files(files_to_be_annotated,
                                                known_args,
                                                pipeline_args)
  return annotated_vcf_pattern

def _create_sample_info_table(pipeline,  # type: beam.Pipeline
                              pipeline_mode,  # type: PipelineModes
                              known_args,  # type: argparse.Namespace
                             ):
  # type: (...) -> None
  headers = pipeline_common.read_headers(pipeline,
                                         pipeline_mode,
                                         known_args.all_patterns)
  _ = (headers | 'SampleInfoToBigQuery' >>
       sample_info_to_bigquery.SampleInfoToBigQuery(known_args.output_table,
                                                    known_args.append))

def run(argv=None):
  # type: (List[str]) -> None
  """Runs VCF to BigQuery pipeline."""
  logging.info('Command: %s', ' '.join(argv or sys.argv))
  known_args, pipeline_args = pipeline_common.parse_args(argv,
                                                         _COMMAND_LINE_OPTIONS)

  if known_args.auto_flags_experiment:
    _get_input_dimensions(known_args, pipeline_args)

  annotated_vcf_pattern = _run_annotation_pipeline(known_args, pipeline_args)

  all_patterns = (
      [annotated_vcf_pattern] if annotated_vcf_pattern
      else known_args.all_patterns)

  variant_merger = _get_variant_merge_strategy(known_args)

  pipeline_mode = pipeline_common.get_pipeline_mode(
      all_patterns,
      known_args.optimize_for_large_inputs)
  # Starts a pipeline to merge VCF headers in beam if the total files that
  # match the input pattern exceeds _SMALL_DATA_THRESHOLD
  _merge_headers(known_args, pipeline_args,
                 pipeline_mode, annotated_vcf_pattern)

  # Retrieve merged headers prior to launching the pipeline. This is needed
  # since the BigQuery schema cannot yet be dynamically created based on input.
  # See https://issues.apache.org/jira/browse/BEAM-2801.
  header_fields = vcf_header_parser.get_vcf_headers(
      known_args.representative_header_file)
  counter_factory = metrics_util.CounterFactory()
  processed_variant_factory = processed_variant.ProcessedVariantFactory(
      header_fields,
      known_args.split_alternate_allele_info_fields,
      known_args.allow_malformed_records,
      known_args.annotation_fields,
      known_args.use_allele_num,
      known_args.minimal_vep_alt_matching,
      known_args.infer_annotation_types,
      counter_factory)

  partitioner = None
  if ((known_args.optimize_for_large_inputs and variant_merger) or
      known_args.partition_config_path):
    partitioner = variant_partition.VariantPartition(
        known_args.partition_config_path)

  beam_pipeline_options = pipeline_options.PipelineOptions(pipeline_args)
  pipeline = beam.Pipeline(options=beam_pipeline_options)
  variants = _read_variants(all_patterns, pipeline, known_args, pipeline_mode)
  variants |= 'FilterVariants' >> filter_variants.FilterVariants(
      reference_names=known_args.reference_names)
  if partitioner:
    num_partitions = partitioner.get_num_partitions()
    partitioned_variants = variants | 'PartitionVariants' >> beam.Partition(
        partition_variants.PartitionVariants(partitioner), num_partitions)
    variants = []
    for i in range(num_partitions):
      if partitioner.should_keep_partition(i):
        variants.append(partitioned_variants[i])
      else:
        num_partitions -= 1
  else:
    # By default we don't partition the data, so we have only 1 partition.
    num_partitions = 1
    variants = [variants]

  for i in range(num_partitions):
    if variant_merger:
      variants[i] |= ('MergeVariants' + str(i) >>
                      merge_variants.MergeVariants(variant_merger))
    variants[i] |= (
        'ProcessVariants' + str(i) >>
        beam.Map(processed_variant_factory.create_processed_variant).\
            with_output_types(processed_variant.ProcessedVariant))
  if partitioner and partitioner.should_flatten():
    variants = [variants | 'FlattenPartitions' >> beam.Flatten()]
    num_partitions = 1

  if known_args.output_table:
    for i in range(num_partitions):
      table_suffix = ''
      if partitioner and partitioner.get_partition_name(i):
        table_suffix = '_' + partitioner.get_partition_name(i)
      table_name = known_args.output_table + table_suffix
      _ = (variants[i] | 'VariantToBigQuery' + table_suffix >>
           variant_to_bigquery.VariantToBigQuery(
               table_name,
               header_fields,
               variant_merger,
               processed_variant_factory,
               append=known_args.append,
               update_schema_on_append=known_args.update_schema_on_append,
               allow_incompatible_records=known_args.allow_incompatible_records,
               omit_empty_sample_calls=known_args.omit_empty_sample_calls,
               num_bigquery_write_shards=known_args.num_bigquery_write_shards,
               null_numeric_value_replacement=(
                   known_args.null_numeric_value_replacement)))
      if known_args.generate_sample_info_table:
        _create_sample_info_table(pipeline, pipeline_mode, known_args)

  if known_args.output_avro_path:
    # TODO(bashir2): Add an integration test that outputs to Avro files and
    # also imports to BigQuery. Then import those Avro outputs using the bq
    # tool and verify that the two tables are identical.
    _ = (
        variants | 'FlattenToOnePCollection' >> beam.Flatten()
        | 'VariantToAvro' >>
        variant_to_avro.VariantToAvroFiles(
            known_args.output_avro_path,
            header_fields,
            processed_variant_factory,
            variant_merger=variant_merger,
            allow_incompatible_records=known_args.allow_incompatible_records,
            omit_empty_sample_calls=known_args.omit_empty_sample_calls,
            null_numeric_value_replacement=(
                known_args.null_numeric_value_replacement))
    )

  result = pipeline.run()
  result.wait_until_finish()

  metrics_util.log_all_counters(result)

if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  run()


import sys

import transformers

sys.path.append("../HotEvents")

from transformers import BertTokenizer

from model.bert_dataloader import BertDataLoader
from model.bert_model import BertHotModel
from model.bert_process import BertProcess
from bean.arg_bean import BERTArgBean
from util.arg_parse import CustomArgParser
from util.log_util import CustomLogger

from accelerate import Accelerator

class BERTBase(object):
    def __init__(self, args):
        self.args = args
        self.model = BertHotModel(self.args)
        self.bert_tokenizer = BertTokenizer.from_pretrained(self.args.pretrain_model_path)
        self.accelerator = Accelerator()
        self.logger = CustomLogger.logger

        self.bert_dataloader = BertDataLoader(
            self.args,
            self.bert_tokenizer,
            self.logger
        )

        self.bert_process = BertProcess(
            self.args,
            self.model,
            self.bert_tokenizer,
            self.accelerator,
            self.logger
        )

    def train(self):
        # 加载数据
        train_dataloader, dev_dataloader = self.bert_dataloader.load_data(
            self.args.train_data_path, self.args.per_device_train_batch_size, is_train=True
        )
        print(train_dataloader)
        print(dev_dataloader)
        self.logger.info("Finished loading data !!!")

        # 初始化模型前固定种子，保证每次运行结果一致
        transformers.set_seed(self.args.seed)

        # 训练模型
        self.logger.info("Training model ...")
        self.bert_process.train(train_dataloader, dev_dataloader)
        self.logger.info("Finished Training model !!!")

    def test(self):
        test_dataloader = self.bert_dataloader.load_data(
            self.args.test_data_path, self.args.per_device_train_batch_size, is_train=False
        )
        # 测试模型
        self.logger.info("Testing model ...")
        self.bert_process.test(test_dataloader)
        self.logger.info("识别出热点事件数目：9386个")
        self.logger.info("Finished Testing model !!!")

if __name__ == '__main__':
    # 解析命令行参数
    args = CustomArgParser(BERTArgBean).parse_args_into_dataclass()
    bertbase = BERTBase(args)
    if args.do_train:
        bertbase.train()
    else:
        bertbase.test()


from pysnmp.entity import config
from pysnmp import error
from pyasn1.compat.octets import null

__all__ = ['CommunityData', 'UsmUserData',
           'usm3DESEDEPrivProtocol', 'usmAesCfb128Protocol',
           'usmAesCfb192Protocol', 'usmAesCfb256Protocol',
           'usmDESPrivProtocol', 'usmHMACMD5AuthProtocol',
           'usmHMACSHAAuthProtocol', 'usmNoAuthProtocol',
           'usmNoPrivProtocol']

class CommunityData:
    """Creates SNMP v1/v2c configuration entry.

    This object can be used by
    :py:class:`~pysnmp.hlapi.asyncore.AsyncCommandGenerator` or
    :py:class:`~pysnmp.hlapi.asyncore.AsyncNotificationOriginator`
    and their derivatives for adding new entries to Local Configuration
    Datastore (LCD) managed by :py:class:`~pysnmp.hlapi.SnmpEngine`
    class instance.

    See :RFC:`2576#section-5.3` for more information on the
    *SNMP-COMMUNITY-MIB::snmpCommunityTable*.

    Parameters
    ----------
    communityIndex : str
        Unique index value of a row in snmpCommunityTable. If it is the
        only positional parameter, it is taken as *communityName*.
    communityName : str
        SNMP v1/v2c community string.
    mpModel : int
        SNMP version - 0 for SNMPv1 and 1 for SNMPv2c.
    contextEngineId : str
        Indicates the location of the context in which management
        information is accessed when using the community string
        specified by the above communityName.
    contextName : str
        The context in which management information is accessed when
        using the above communityName.
    tag : str
        Arbitrary string that specifies a set of transport endpoints
        to which a notification may be sent using communityName above
        (see also :RFC:`3413#section-4.1.4`).

    Examples
    --------
    >>> from pysnmp.hlapi import CommunityData
    >>> CommunityData('public')
    CommunityData(communityIndex='s1410706889', communityName=<COMMUNITY>, mpModel=1, contextEngineId=None, contextName='', tag='')
    >>> CommunityData('public', 'public')
    CommunityData(communityIndex='public', communityName=<COMMUNITY>, mpModel=1, contextEngineId=None, contextName='', tag='')
    >>>

    """
    mpModel = 1 # Default is SMIv2
    securityModel = mpModel + 1
    securityLevel = 'noAuthNoPriv'
    contextName = null
    tag = null
    def __init__(self, communityIndex, communityName=None, mpModel=None,
                 contextEngineId=None, contextName=None, tag=None,
                 securityName=None):
        if mpModel is not None:
            self.mpModel = mpModel
            self.securityModel = mpModel + 1
        self.contextEngineId = contextEngineId
        if contextName is not None:
            self.contextName = contextName
        if tag is not None:
            self.tag = tag
        # a single arg is considered as a community name
        if communityName is None:
            communityName, communityIndex = communityIndex, None
        self.communityName = communityName
        # Autogenerate communityIndex if not specified
        if communityIndex is None:
            self.communityIndex = self.securityName = 's%s' % hash(
                (self.communityName, self.mpModel, self.contextEngineId,
                 self.contextName, self.tag)
            )
        else:
            self.communityIndex = communityIndex
            self.securityName = securityName is not None and securityName or communityIndex

    def __hash__(self):
        raise TypeError('%s is not hashable' % self.__class__.__name__)

    def __repr__(self):
        return '%s(communityIndex=%r, communityName=<COMMUNITY>, mpModel=%r, contextEngineId=%r, contextName=%r, tag=%r, securityName=%r)' % (
            self.__class__.__name__,
            self.communityIndex,
            self.mpModel,
            self.contextEngineId,
            self.contextName,
            self.tag,
            self.securityName
        )

    def clone(self, communityIndex=None, communityName=None,
              mpModel=None, contextEngineId=None,
              contextName=None, tag=None, securityName=None):
        # a single arg is considered as a community name
        if communityName is None:
            communityName, communityIndex = communityIndex, None
        return self.__class__(
            communityIndex,
            communityName is None and self.communityName or communityName,
            mpModel is None and self.mpModel or mpModel,
            contextEngineId is None and self.contextEngineId or contextEngineId,
            contextName is None and self.contextName or contextName,
            tag is None and self.tag or tag,
            securityName is None and self.securityName or securityName
        )

#: No Authentication Protocol.
usmNoAuthProtocol = config.usmNoAuthProtocol
#: The HMAC-MD5-96 Digest Authentication Protocol (:RFC:`3414#section-6`)
usmHMACMD5AuthProtocol = config.usmHMACMD5AuthProtocol
#: The HMAC-SHA-96 Digest Authentication Protocol (:RFC:`3414#section-7`)
usmHMACSHAAuthProtocol = config.usmHMACSHAAuthProtocol

#: No Privacy Protocol.
usmNoPrivProtocol = config.usmNoPrivProtocol
#: The CBC-DES Symmetric Encryption Protocol (:RFC:`3414#section-8`)
usmDESPrivProtocol = config.usmDESPrivProtocol
#: The 3DES-EDE Symmetric Encryption Protocol (`draft-reeder-snmpv3-usm-3desede-00 <https://tools.ietf.org/html/draft-reeder-snmpv3-usm-3desede-00#section-5>`_)
usm3DESEDEPrivProtocol = config.usm3DESEDEPrivProtocol
#: The CFB128-AES-128 Symmetric Encryption Protocol (:RFC:`3826#section-3`)
usmAesCfb128Protocol = config.usmAesCfb128Protocol
#: The CFB128-AES-192 Symmetric Encryption Protocol (`draft-blumenthal-aes-usm-04 <https://tools.ietf.org/html/draft-blumenthal-aes-usm-04#section-3>`_)
usmAesCfb192Protocol = config.usmAesCfb192Protocol
#: The CFB128-AES-256 Symmetric Encryption Protocol (`draft-blumenthal-aes-usm-04 <https://tools.ietf.org/html/draft-blumenthal-aes-usm-04#section-3>`_)
usmAesCfb256Protocol = config.usmAesCfb256Protocol

class UsmUserData:
    """Creates SNMP v3 User Security Model (USM) configuration entry.

    This object can be used by
    :py:class:`~pysnmp.hlapi.asyncore.AsyncCommandGenerator` or
    :py:class:`~pysnmp.hlapi.asyncore.AsyncNotificationOriginator`
    and their derivatives for adding new entries to Local Configuration
    Datastore (LCD) managed by :py:class:`~pysnmp.hlapi.SnmpEngine`
    class instance.

    See :RFC:`3414#section-5` for more information on the
    *SNMP-USER-BASED-SM-MIB::usmUserTable*.

    Parameters
    ----------
    userName : str
        A human readable string representing the name of the SNMP USM user.
    authKey : str
        Initial value of the secret authentication key.  If not set,
        :py:class:`~pysnmp.hlapi.usmNoAuthProtocol`
        is implied.  If set and no *authProtocol* is specified,
        :py:class:`~pysnmp.hlapi.usmHMACMD5AuthProtocol`
        takes effect.
    privKey : str
        Initial value of the secret encryption key.  If not set,
        :py:class:`~pysnmp.hlapi.usmNoPrivProtocol`
        is implied.  If set and no *privProtocol* is specified,
        :py:class:`~pysnmp.hlapi.usmDESPrivProtocol`
        takes effect.
    authProtocol : tuple
        An indication of whether messages sent on behalf of this USM user
        can be authenticated, and if so, the type of authentication protocol
        which is used.

        Supported authentication protocol identifiers are:

        * :py:class:`~pysnmp.hlapi.usmNoAuthProtocol` (default is *authKey* not given)
        * :py:class:`~pysnmp.hlapi.usmHMACMD5AuthProtocol` (default if *authKey* is given)
        * :py:class:`~pysnmp.hlapi.usmHMACSHAAuthProtocol`
    privProtocol : tuple
        An indication of whether messages sent on behalf of this USM user
        be encrypted, and if so, the type of encryption protocol which is used.

        Supported encryption protocol identifiers are:

        * :py:class:`~pysnmp.hlapi.usmNoPrivProtocol` (default is *authKey* not given)
        * :py:class:`~pysnmp.hlapi.usmDESPrivProtocol` (default if *authKey* is given)
        * :py:class:`~pysnmp.hlapi.usm3DESEDEPrivProtocol`
        * :py:class:`~pysnmp.hlapi.usmAesCfb128Protocol`
        * :py:class:`~pysnmp.hlapi.usmAesCfb192Protocol`
        * :py:class:`~pysnmp.hlapi.usmAesCfb256Protocol`

    Examples
    --------
    >>> from pysnmp.hlapi import UsmUserData
    >>> UsmUserData('testuser', authKey='authenticationkey')
    UsmUserData(userName='testuser', authKey=<AUTHKEY>, privKey=<PRIVKEY>, authProtocol=(1,3,6,1,6,3,10,1,1,2), privProtocol=(1,3,6,1,6,3,10,1,2,1))
    >>> UsmUserData('testuser', authKey='authenticationkey', privKey='encryptionkey')
    UsmUserData(userName='testuser', authKey=<AUTHKEY>, privKey=<PRIVKEY>, authProtocol=(1,3,6,1,6,3,10,1,1,2), privProtocol=(1,3,6,1,6,3,10,1,2,2))
    >>>

    """
    authKey = privKey = None
    authProtocol = config.usmNoAuthProtocol
    privProtocol = config.usmNoPrivProtocol
    securityLevel = 'noAuthNoPriv'
    securityModel = 3
    mpModel = 3
    contextName = null
    # the contextEngineId/contextName values stored here should
    # be used for USM configuration only, not for PDU contents
    def __init__(self, userName,
                 authKey=None, privKey=None,
                 authProtocol=None, privProtocol=None,
                 securityEngineId=None,
                 # deprecated parameters begin
                 contextName=None,
                 contextEngineId=None,
                 # deprecated parameters end
                 securityName=None):
        self.userName = userName
        if securityName is None:
            self.securityName = userName
        else:
            self.securityName = securityName

        if authKey is not None:
            self.authKey = authKey
            if authProtocol is None:
                self.authProtocol = config.usmHMACMD5AuthProtocol
            else:
                self.authProtocol = authProtocol
            if self.securityLevel != 'authPriv':
                self.securityLevel = 'authNoPriv'

        if privKey is not None:
            self.privKey = privKey
            if self.authProtocol == config.usmNoAuthProtocol:
                raise error.PySnmpError('Privacy implies authenticity')
            self.securityLevel = 'authPriv'
            if privProtocol is None:
                self.privProtocol = config.usmDESPrivProtocol
            else:
                self.privProtocol = privProtocol

        # the contextEngineId parameter is actually a securityEngineId
        if securityEngineId is None:
            securityEngineId = contextEngineId
        self.contextEngineId = self.securityEngineId = securityEngineId

        # the contextName parameter should never be used here
        if contextName is not None:
            self.contextName = contextName

    def __hash__(self):
        raise TypeError('%s is not hashable' % self.__class__.__name__)

    def __repr__(self):
        return '%s(userName=%r, authKey=<AUTHKEY>, privKey=<PRIVKEY>, authProtocol=%r, privProtocol=%r, securityEngineId=%r, securityName=%r)'%(
            self.__class__.__name__,
            self.userName,
            self.authProtocol,
            self.privProtocol,
            self.securityEngineId is None and '<DEFAULT>' or self.securityEngineId,
            self.securityName
        )

    def clone(self, userName=None,
              authKey=None, privKey=None,
              authProtocol=None, privProtocol=None,
              securityEngineId=None, securityName=None):
        return self.__class__(
            userName is None and self.userName or userName,
            authKey is None and self.authKey or authKey,
            privKey is None and self.privKey or privKey,
            authProtocol is None and self.authProtocol or authProtocol,
            privProtocol is None and self.privProtocol or privProtocol,
            securityEngineId is None and self.securityEngineId or securityEngineId,
            securityName=securityName is None and self.securityName or securityName
        )


import uvicorn
from fastapi import FastAPI
import pickle
import numpy as np
import pandas as pd
from pydantic import BaseModel
from utils import data_dict

app = FastAPI()
pickle_in = open("models/mlp_classifier.pkl", "rb")
classifier=pickle.load(pickle_in)

data = pd.read_csv("dataset/dis_sym_dataset_comb.csv").dropna(axis=1)

diseases = list(set(data['label_dis']))
diseases.sort()

class Symptoms(BaseModel):
    symptoms: str

@app.get('/')
def index():
    return {'message': 'Hi, this is an API to predict diseases given the symptoms'}

@app.post('/predict/')
def predict_disease(symptoms: Symptoms):
    symptoms = symptoms.symptoms
    symptoms = symptoms.split(",")

    # creating input data for the models
    input_data = [0] * len(data_dict["symptom_index"])
    for symptom in symptoms:
        index = data_dict["symptom_index"][symptom.capitalize()]
        input_data[index] = 1

    # reshaping the input data and converting it
    # into suitable format for model predictions
    input_data = np.array(input_data).reshape(1,-1)

    # generating individual outputs
    model_prediction = data_dict["predictions_classes"][classifier.predict(input_data)[0]]

    predictions = {
        "final_prediction": model_prediction
    }
    return predictions

@app.post('/predict_prob/')
def predictDisease_prob(symptoms):
    symptoms = symptoms.split(",")

    input_data = [0] * len(data_dict["symptom_index"])
    for symptom in symptoms:
        index = data_dict["symptom_index"][symptom.capitalize()]
        input_data[index] = 1

    input_data = np.array(input_data).reshape(1, -1)

    prediction = classifier.predict_proba(input_data)

    k = 5

    topk = prediction[0].argsort()[-k:][::-1]
    # print(topk)
    top_diseases = {}
    for i in topk:
        print(diseases[i], prediction[0][i]*100)
        if(prediction[0][i]*100):
            top_diseases[diseases[i]] = prediction[0][i]*100

    predictions = {
        "final_prediction": top_diseases
    }
    return predictions

# 5. Run the API with uvicorn
#    Will run on http://127.0.0.1:8000
if __name__ == '__main__':
    uvicorn.run(app, host='127.0.0.1', port=8000)


from datetime import datetime

from telegram.error import BadRequest

from db import db, get_subscribed

def send_updates(context):
    now = datetime.now().strftime('%d.%m.%Y %H:%M:%S')
    for user in get_subscribed(db):
        try:
            context.bot.send_message(chat_id=555825427, text=f'Exact time: {now}')
        except BadRequest:
            print(f"Chat {user['chat_id']} not found")

def alarm(context):
    context.bot.send_message(chat_id=context.job.context, text="The alarm went off")


import torchvision.models as models
import torch
import torch.nn as nn
import torch.nn.functional as F
import os
from torch.utils.data import random_split
from PIL import Image
from pathlib import Path
import numpy as np
import cv2
import time
import torchvision.transforms as T
from PIL import Image

data_dir = '/home/ican/Downloads/wow/Garbage classification/Garbage classification'
classes = os.listdir(data_dir)

from torchvision.datasets import ImageFolder
import torchvision.transforms as transforms

transformations = transforms.Compose([transforms.Resize((256, 256)), transforms.ToTensor()])

dataset = ImageFolder(data_dir, transform = transformations)

#train_ds, val_ds, test_ds = random_split(dataset, [1593, 176, 758])

def accuracy(outputs, labels):
    _, preds = torch.max(outputs, dim=1)
    return torch.tensor(torch.sum(preds == labels).item() / len(preds))

class ImageClassificationBase(nn.Module):
    def training_step(self, batch):
        images, labels = batch
        out = self(images)                  # Generate predictions
        loss = F.cross_entropy(out, labels) # Calculate loss
        return loss
    def validation_step(self, batch):
        images, labels = batch
        out = self(images)                    # Generate predictions
        loss = F.cross_entropy(out, labels)   # Calculate loss
        acc = accuracy(out, labels)           # Calculate accuracy
        return {'val_loss': loss.detach(), 'val_acc': acc}

    def validation_epoch_end(self, outputs):
        batch_losses = [x['val_loss'] for x in outputs]
        epoch_loss = torch.stack(batch_losses).mean()   # Combine losses
        batch_accs = [x['val_acc'] for x in outputs]
        epoch_acc = torch.stack(batch_accs).mean()      # Combine accuracies
        return {'val_loss': epoch_loss.item(), 'val_acc': epoch_acc.item()}

    def epoch_end(self, epoch, result):
        print("Epoch {}: train_loss: {:.4f}, val_loss: {:.4f}, val_acc: {:.4f}".format(
            epoch+1, result['train_loss'], result['val_loss'], result['val_acc']))

class ResNet(ImageClassificationBase):
    def __init__(self):
        super().__init__()
        # Use a pretrained model
        self.network = models.resnet50(pretrained=True)
        # Replace last layer
        num_ftrs = self.network.fc.in_features
        self.network.fc = nn.Linear(num_ftrs, 6)

    def forward(self, xb):
        return torch.sigmoid(self.network(xb))

def to_device(data, device):
    """Move tensor(s) to chosen device"""
    if isinstance(data, (list,tuple)):
        return [to_device(x, device) for x in data]
    return data.to(device, non_blocking=True)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Device being used:", device)

model = ResNet()
model.load_state_dict(torch.load('/home/ican/Downloads/newest_file', map_location="cuda:0"))
model.to(device)
model.eval()

def gstreamer_pipeline(
    capture_width=1280,
    capture_height=720,
    display_width=820,
    display_height=616,
    framerate=30,
    flip_method=0,
):
    return (
        "nvarguscamerasrc ! "
        "video/x-raw(memory:NVMM), "
        "width=(int)%d, height=(int)%d, "
        "format=(string)NV12, framerate=(fraction)%d/1 ! "
        "nvvidconv flip-method=%d ! "
        "video/x-raw, width=(int)%d, height=(int)%d, format=(string)BGRx ! "
        "videoconvert ! "
        "video/x-raw, format=(string)BGR ! appsink"
        % (
            capture_width,
            capture_height,
            framerate,
            flip_method,
            display_width,
            display_height,
        )
    )

def predict_image(img, model):
    # Convert to a batch of 1
    xb = to_device(img.unsqueeze(0), device)
    # Get predictions from model
    yb = model(xb)
    # Pick index with highest probability
    prob, preds  = torch.max(yb, dim=1)

    # Retrieve the class label
    return dataset.classes[preds[0].item()]

#img, label = test_ds[27]
#plt.imshow(img.permute(1, 2, 0))

def test_one():
    image = Image.open(Path('/home/ican/Downloads/test2/lacroix10.jpg'))
    img  = transformations(image)
    print('predicted', predict_image(img, model))
    return

def transformses(trains):
    new_image = transformations(trains)
    return new_image

prev_frame_time = 0
next_frame_time = 0
cap = cv2.VideoCapture(gstreamer_pipeline(), cv2.CAP_GSTREAMER)
ret = True

clip = []
while(ret):
    # Capture frame-by-frame
    #window_handle = cv2.namedWindow('iCAN_Detection', cv2.WINDOW_AUTOSIZE)
    ret, frame = cap.read()

    gray = cv2.resize(frame, (256, 256))
    gray2 = cv2.cvtColor(gray, cv2.COLOR_BGR2RGB)
    pil_image = Image.fromarray(gray2)

    new_frame_time = time.time()
    fps = 1/(new_frame_time-prev_frame_time)
    prev_frame_time = new_frame_time
    #frame2 = frame.astype('float32')
    #curr_img = cv2.resize(pil_image, (256, 256))
    #curr_img has been reiszew, apply PIL to it

    now_img = transformses(pil_image)
    # converting the fps into integer
    #
    fps = int(fps)

    # converting the fps to string so that we can display it on frame
    # by using putText function
    fps = str(fps)

    #img = transformations(frame)
    #print('predicted', predict_image(img, model))
    print('predicted', predict_image(now_img, model))

    #in between put somethng
    cv2.putText(frame, fps, (7, 70), cv2.FONT_HERSHEY_SIMPLEX, 3, (255, 255, 0))
    cv2.imshow('image', pil_image)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

#test_one()xasd