import time
import re
import logging
import scipy.misc
from scipy.ndimage.filters import gaussian_filter
from fpga import api
import numpy as np
import cv2
import os,sys
from bitstring import BitArray
from PIL import Image
import numpngw

logging.basicConfig(format='%(levelname)-6s[%(filename)s:%(lineno)d] %(message)s'
                    ,level=logging.DEBUG)

class T6(api):

    address = {'mk_routine_rst': 0x00,
               'wire_check_set': 0x01,
                     'param_en': 0x02,
                     'param_in': 0x03,
                   'param_addr': 0x04,
                   'rowmap_set': 0x0A,
                     'chip_rst': 0x10,
                 'wire_check_o': 0x20,
                     'fifo_cnt': 0x21,
                    'param_out': 0x22,
                     'exposure': 0x11,
                      'pat_num': 0x12,
                  'mem_pat_num': 0x14,
                    'sys_start': 0x40,
                     'mem_mode': 0x13,
                      'trigger': 0x53,
                    'fifo_flag': 0x6a,
                       'mem_wr': 0x80,
                      'ldo_spi': 0x81,
                      'cis_spi': 0x82,
                       'mem_rd': 0xA0,
                        'image': 0xB0,
                      'adc2_rd': 0xA3
               }

    param = {      'numPattern' : 0,
                    'Texp_ctrl' : 1,
                    'Tdes2_d' : 2,
                    'Tdes2_w' : 3,
                    'Tmsken_d' : 4,
                    'Tmsken_w' : 5,
                    'Tgsub_w' : 6,
                    'Tgl_Res' : 7,
                    'Tproj_dly' : 8,
                    'Tadd' : 9,
                    'NumRep' : 10,
                    'NumGsub' : 11,
                    'TExpRST' : 12,
                    'Tdrain_w' : 13,
                    'imgCountTrig' : 14,
                    'TdrainR_d' : 15,
                    'TdrainF_d' : 16,
                    'TLedOn' : 17,
                    'Select' : 18,
                 'CAM_TG_pulseWid' : 19,
                 'CAM_TG_HoldTime' : 20,
                    'CAM_TG_delay' : 21,
                       'MASK_SIZE' : 22,
                        'IMG_SIZE' : 23,
                     'T_DEC_SEL_0' : 24,
                     'T_DEC_SEL_1' : 25,
                      'T_DEC_EN_0' : 26,
                      'T_DEC_EN_1' : 27,
                        'T_DONE_1' : 28,
                  'MSTREAM_Select' : 29,
                      'MU_NUM_ROW' : 30,
                    'SUB_IMG_SIZE' : 31,
                    'ADC1_Tcolumn' : 32,
                         'ADC1_T1' : 33,
                       'ADC1_T2_1' : 34,
                       'ADC1_T2_0' : 35,
                         'ADC1_T3' : 36,
                         'ADC1_T4' : 37,
                         'ADC1_T5' : 38,
                         'ADC1_T6' : 39,
                         'ADC1_T7' : 40,
                         'ADC1_T8' : 41,
                         'ADC1_T9' : 42,
                       'ADC1_TADC' : 43,
                    'ADC1_NUM_ROW' : 44,
                       'ADC1_Wait' : 45,
                       'T_MU_wait' : 46,
                    'ADC2_Tcolumn' : 47,
                         'ADC2_T1' : 48,
                       'ADC2_T2_1' : 49,
                       'ADC2_T2_0' : 50,
                         'ADC2_T3' : 51,
                         'ADC2_T4' : 52,
                         'ADC2_T5' : 53,
                         'ADC2_T6' : 54,
                         'ADC2_T7' : 55,
                         'ADC2_T8' : 56,
                         'ADC2_T9' : 57,
                       'ADC2_TADC' : 58,
                    'ADC2_NUM_ROW' : 59,
                       'ADC2_Wait' : 60,
                       'T_MU_wait' : 61,
                          'TrigNo' : 62,
                     'TrigOffTime' : 63,
                      'TrigOnTime' : 64,
                'UNIT_SUB_IMG_SIZE': 65,
                    'N_SUBREADOUTS': 66,
                      'spi_control': 67,
                         'i2c_data': 68,
                      'i2c_control': 69,
                   'ROWADD_INCRMNT': 70,
                     'TrigWaitTime': 71,
                           'TLEDNo': 72,
                          'TLEDOff': 73,
                           'TLEDOn': 74,
                         'TLEDWait': 75
            }

    # voltages for LDO
    voltages = np.array([3.3, 3.3, 3.3, 3.3,  # CS1
                         3.3, 3.3, 3.3, 3.3,  # CS2
                         3.3, 3.3, 3.3, 3.3]) # CS3

    # this is 96-bit on chip. To configure mask upload settings
    cis_spi_mu = [  '0100001000110000', '0000000000000000', # 95:80, 79:64
                    '0000000000000000', '0000000000000000', # 63:48, 47:32
                    '0000000000000000', '1000000000000000', # 31:16, 15:00
                    '0000000000000010', '0000000000000000']

    # this is only 48-bit on chip. To configure all other settings
    cis_spi_ro = [  '0000000000000000', '0000000000000000', # GARBJ, GARBJ
                    '0000000000000000', '0000000000011101', # GARBJ, 47:32
                    '1100000001111100', '1101111101000011', # 31:16, 15:00
                    '0000000000000001', '0000000000000000'] # <slave_select>

    col, row, tab, ch = 108, 320, 2, 3
    #frame_length = 552960
    # frame_length    = col*4*row*tab*2 # t6
    frame_length            = round((480*40*17*2*12*256/255) // 8) #t7
    unit_subframe_length    = round(480*40*17*1*1*256/255/8)

    adc2_container = bytearray(unit_subframe_length)
    adc1_container = bytearray(frame_length)

    BLOCK_SIZE      = 512

    # trying to fix importing t4 from path outside of Cam_API_py and still have it work
    # use module relative paths
    black_img_file  = os.path.join(os.path.dirname(__file__), 't6_black_image_avg.npy')
    bright_img_file = os.path.join(os.path.dirname(__file__), 't6_bright_image_avg.npy')
    gain_mat_file   = os.path.join(os.path.dirname(__file__), 'gain_mat.npy')
    gain_cali_file  = os.path.join(os.path.dirname(__file__), 't6_gain_cali.npy')
    img_map_file    = os.path.join(os.path.dirname(__file__), 't6_img_map.npy')
    img_graymap_file = os.path.join(os.path.dirname(__file__),'t6_img_graymap.npy')
    mask_map_file   = os.path.join(os.path.dirname(__file__), 't6_mask_map.npy')
    adc2_map_file   = os.path.join(os.path.dirname(__file__), 't7_adc2_map_1bit.npy')

    black_flag = False
    black_scale_flag = False
    gain_flag = False
    gain_scale_flag = False
    scale_flag = False
    MaskNumInMem=0
    MaskNumForCam=0
    MaskRepetition=0
    ldo_delay = 0.1

    def __init__(self, bitfile,reConfFPGA=True):

        super(T6, self).__init__(bitfile)
        logging.info("Bitfile: {}".format(bitfile))
        if(reConfFPGA):
            time.sleep(1)
            self.LDOConf(self.voltages)
            time.sleep(1)
            self.CISConf(self.cis_spi_mu)
            time.sleep(1)
            self.CISConf(self.cis_spi_ro)

            time.sleep(5)
            self.wire_in(self.address['mk_routine_rst'], 0x04)
            self.wire_in(self.address['mk_routine_rst'], 0x00)
            #self.wire_in(self.address['mask'], 0xfff)
            self.wire_in(self.address['chip_rst'], 0x01)
            #self.wire_in(self.address['chip_rst'], 0x00)
            self.wire_in(self.address['exposure'], 1)
            self.wire_in(self.address['pat_num'], 1)

            # self.param_set(self.param['imgCountTrig'],3000)
            self.param_set(self.param['TExpRST'], 10000) # 200MHz
            self.param_set(self.param['Tproj_dly'], 100)
            self.param_set(self.param['NumRep'], 1)
            self.param_set(self.param['NumGsub'], 0)
            self.param_set(self.param['Tdrain_w'], 0)

            self.param_set(self.param['Tgsub_w'], 100)
            self.param_set(self.param['TdrainF_d'], 20)
            self.param_set(self.param['TdrainR_d'], 10)

            #DELAY+WIDTH<=15
            self.param_set(self.param['Tmsken_d'], 4)
            self.param_set(self.param['Tmsken_w'], 11)

            self.param_set(self.param['Select'], 0x80000000)  # enable slave mode

        else:
            logging.info("Did not configure FPGA with bitfile")
            self.wire_in(self.address['mk_routine_rst'], 0x04)
            self.wire_in(self.address['mk_routine_rst'], 0x00)
            self.wire_in(self.address['chip_rst'], 0x01)

        self.mapping = self.ImgMapCal()
        # self.mapping = self.ImgGrayMapCal()

        self.mkmapping = self.MaskMapCal()
        if(os.path.exists(self.black_img_file)):
            self.black_img = np.load(self.black_img_file)
        if(os.path.exists(self.gain_cali_file)):
            self.gain_cali = np.load(self.gain_cali_file)
        if(os.path.exists(self.bright_img_file)):
            self.bright_img = np.load(self.bright_img_file)
            #self.coeff_mat = np.load(self.gain_cali_file)
            #self.coeff_mat[324:] = self.coeff_mat[0:324]
        if(os.path.exists(self.adc2_map_file)):
            self.adc2_map_1row = np.load(self.adc2_map_file).flatten()

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

    def UploadMask(self, mask_file, MaskNum):
        """
        mask_file@in: mask file name and path
        MaskNum@in: give mask number which will be set in memory control
        """
        # reset camera
        self.wire_in(self.address['chip_rst'], 0x01)
        logging.info("Camera Stopped.")
        self.fmaskupload(mask_file)
        logging.info("Mask {} uploaded.".format(mask_file))
        self.wire_in(self.address['mem_pat_num'], MaskNum)
        self.MaskNumInMem=MaskNum
        logging.info("Mask #{}".format(MaskNum))

    def UploadMaskDummy(self, mask_file, MaskNum):
        """
    Dummy version of the mask upload which sets all variables except uploading mask to memory
        mask_file@in: mask file name and path
        MaskNum@in: give mask number which will be set in memory control
        """
        # reset camera
        self.wire_in(self.address['chip_rst'], 0x01)
        logging.info("Camera Stopped.")
    # self.fmaskupload(mask_file)
        # memory mask upload mode
        self.wire_in(self.address['mem_mode'], 0x01)
        self.wire_in(self.address['mk_routine_rst'], 0x04)
        self.wire_in(self.address['mk_routine_rst'], 0x01)
        logging.info("Mask {} was ***NOT*** uploaded.".format(mask_file))
        self.wire_in(self.address['mem_pat_num'], MaskNum)
        self.MaskNumInMem=MaskNum
        logging.info("Mask #{}".format(MaskNum))

    def readout_reset(self):
        self.wire_in(self.address['chip_rst'], 0x01)
        self.wire_in(self.address['chip_rst'], 0x00)

    def chip_reset(self):
        self.wire_in(self.address['mk_routine_rst'], 0x04)
        self.wire_in(self.address['mk_routine_rst'], 0x00)
        #self.wire_in(self.address['mask'], 0xfff)
        self.wire_in(self.address['chip_rst'], 0x01)

    def Pause(self):
        self.wire_in(self.address['chip_rst'], 0x01)
        logging.info("Camera Stopped.")

    def Resume(self):
        self.wire_in(self.address['chip_rst'], 0x00)
        logging.info("Camera Stopped.")

    def SetMaskParam(self, MaskNum, MaskRepNum=1):
        """
        MaskNum@in: set mask number applied in camera
        MaskRepNum@in: each mask repetition time in camera
        """
        # reset camera
        assert self.MaskNumInMem>0, "Please upload mask to camera first."
        assert self.MaskNumInMem%MaskNum==0, "CamMask #{} should dividable by MemMask #{}".format(MaskNum, self.MaskNumInMem)
        self.wire_in(self.address['pat_num'], MaskNum)
        logging.info("Camera mask #{}".format(MaskNum))
        self.MaskNumForCam=MaskNum
        self.param_set(self.param['NumRep'], MaskRepNum)
        logging.info("Camera mask repetition #{}".format(MaskRepNum))
        self.MaskRepetition=MaskRepNum

    def SetExposure(self, value):
        """
        value@in: exposure time in us
        """
        assert self.MaskNumForCam != 0, "Please set mask parameter to camera first."
        assert self.MaskRepetition != 0, "Please set mask parameter to camera first."
        logging.info("Setting Camera exposure #{}us ...".format(value))

        maskuploadtime = 26195 # 26195ns for maskupload, 200MHz freq. 5n
        SetAsClk = int((value*1000-maskuploadtime*self.MaskRepetition)/5-2)

        assert SetAsClk > 0, "Exposure time should larger than {}us".format((10+maskuploadtime*self.MaskRepetition)/1000)

        self.wire_in(self.address['exposure'], SetAsClk)
        logging.info("Camera exposure setting finished. #{} clk".format(SetAsClk))

    def getimg(self, mode='raw'):
        """
        mode = 'raw'/'black'
        """
        assert mode == 'raw' or 'black', "No {} mode. 'raw' or 'black'".format(mode)

        if mode == 'raw':
            self.showMode()
        if mode == 'black':
            self.showMode(black_scale=True)

        return self.arrange(self.imread())

    def nothing(self, x):
        pass

    def imshow(self, numbers = None):

        img_name = 'image'
        cv2.namedWindow(img_name)
        cv2.createTrackbar('Zoom', img_name, 10, 30, self.nothing)

        while True:

            factor = cv2.getTrackbarPos('Zoom',img_name)
            if factor == 0: factor = 1

            #img=self.arrange(self.imread())
            img=self.getimg()

            img_scaled = cv2.resize(img,None,fx=factor/10, fy=factor/10,
                                    interpolation = cv2.INTER_LINEAR)

            cv2.imshow(img_name, img_scaled)

            k=cv2.waitKey(1)
            if k==27: break

        cv2.destroyAllWindows()

    def setRowMap(self,rowmapGuide):
        logging.info("rewriting row map")
        set_write   = '0'+'{:015b}'+'{:016b}'
        done_write  = '1'+'{:015b}'+'{:016b}'
        totalRows   = len(rowmapGuide);      i=0;
        for current_map in rowmapGuide:
            from_row  = current_map[0]
            to_row    = current_map[1]
            write_value = int(set_write.format(from_row,to_row),2)
            done_value  = int(done_write.format(from_row,to_row),2)
            self.wire_in(self.address['rowmap_set'],write_value)
            time.sleep(1/1000)
            self.wire_in(self.address['rowmap_set'],done_value)
            time.sleep(1/1000)
            self.wire_in(self.address['rowmap_set'],write_value)
            time.sleep(1/1000)
            print("{}/{} done".format(i,totalRows),end='\r')
        self.wire_in(self.address['rowmap_set'],0)

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

    def showMode(self, black=False, black_scale=False, gain=False, gain_scale=False, scale=False):
        if black or black_scale and os.path.exists(self.black_img_file):
            if black_scale:
                self.black_scale_flag = True
            else:
                self.black_flag = True
            print("Black calliberation On")
        elif gain or gain_scale and os.path.exists(self.gain_cali_file):
            if gain_scale:
                self.gain_scale_flag = True
            else:
                self.gain_flag = True
            print("Black calliberation On")
            print("Gain calliberation On")
        elif scale:
            print("Normal scale On")
            self.scale_flag = True
        else:
            print("Raw image On")

    def LDOConf(self, voltages):
        '''
        Takes in voltage values used to set LDOs on board and configures
        byte stream to send to SPI_master module.
        ----------------------------------------------------------
        Variables

        SlaveSel:  1-hot encoded mask for chip select
        Addr:      8 bit pattern for DAC addr. Ordered as 0,0,0...A2,A1,A0
        Voltage:   8 bit pattern to control LDO voltage. Ordered from Msb -> Lsb
        Mode:      Currently unused in SPI_master. Can be used to change SPI mode later
                   for different CPOL and CPAH configurations.
        ------------------------------------------------------------
        Output to SPI_master

        words are formatted as { SS | Mode | Voltage | Addr } in python
        FPGA flips byte order so words recieved as { Addr | Voltage | Mode | SS } by SPI_master
        '''

        # SPI order as { SS | Mode | Voltage | Addr } per row
        SPI = np.array([
            # CS_1
            '00000001', '00000000', '________', '00000000',
            '00000001', '00000000', '________', '00000001',
            '00000001', '00000000', '________', '00000010',
            '00000001', '00000000', '________', '00000011',
            # CS_2
            '00000010', '00000000', '________', '00000000',
            '00000010', '00000000', '________', '00000001',
            '00000010', '00000000', '________', '00000010',
            '00000010', '00000000', '________', '00000011',
            # CS_3
            '00000100', '00000000', '________', '00000000',
            '00000100', '00000000', '________', '00000001',
            '00000100', '00000000', '________', '00000010',
            '00000100', '00000000', '________', '00000011'
        ])

        # get 0->255 int that gives voltage out
        voltages = np.rint(128/voltages).astype(int)
        # convert int to 8'b binary representation
        for i in range(len(voltages)):
            SPI[4*i+2] = np.binary_repr(voltages[i], width=8)

        # convert to bytes to send out to FPGA
        byte_pattern = bytearray(int(SPI[i],2) for i in range(len(SPI)))
        self.write(self.address['ldo_spi'], byte_pattern)

    def CISConf(self, data):
        self.wire_in(self.address['chip_rst'], 0x01)
        self.wire_in(self.address['chip_rst'], 0x00)
        arr = np.array(list(data[0]), dtype='|S1')
        for i in range(1,len(data),1):
            arr = np.vstack((arr,np.array(list(data[i]), dtype='|S1')))
        arr=arr.reshape(-1,32)
        # correct inverted captured image
        arrange = np.flip(np.arange(32*3),axis=0) # vertical flip
        data = arr[:3,:].flatten()
        data = data[arrange].reshape(-1,32)
        arr = np.vstack((data,arr[3,:]))

        arrange = np.arange(32).reshape(4,-1)[[3,2,1,0],:]
        arr = arr[:,arrange.flatten()]

        p = arr.flatten().tobytes()
        pattern = bytearray(int(p[i:i+8],2) for i in range(0, len(p), 8))
        print(len(pattern))
        self.write(self.address['cis_spi'], pattern)

    def ImgMapCal(self):
        if os.path.exists(self.img_map_file):
            imgmap = np.load(self.img_map_file)
            print("Loading exist image mapping data...")
        else:
            imgmap = np.arange(self.tab*self.row*self.col*self.ch).reshape(-1,3)
            #imgmap = np.arange(self.tab*self.row*107*self.ch).reshape(-1,3)
            imgmap = np.concatenate(np.vsplit(imgmap,self.row*self.tab),axis=1)
            imgmap = imgmap.transpose().reshape(-1,self.col*self.ch*self.tab)
            #imgmap = imgmap.transpose().reshape(-1,107*self.ch*self.tab)
            np.save(self.img_map_file,imgmap)
            print("Calculating image mapping data...")

        # one column shift in bucket 2
        # h,w=imgmap.shape
        # imgmap[:,int(w/2):w-1] = imgmap[:,int(w/2)+1:w]

        return imgmap

    def int2gray(self,n):
        n ^= (n>>1)
        return bin(n)[2:]

    def ImgGrayMapCal(self):
        if os.path.exists(self.img_graymap_file):
            imgmap = np.load(self.img_graymap_file)
            print("Loading exist image mapping data...")
        else:
            imgmap = np.arange(self.tab*self.row*self.col*self.ch).reshape(-1,3)
            #imgmap = np.arange(self.tab*self.row*107*self.ch).reshape(-1,3)
            imgmap = np.concatenate(np.vsplit(imgmap,self.row*self.tab),axis=1)
            imgmap = imgmap.transpose().reshape(-1,self.col*self.ch*self.tab)
            gray = np.concatenate(([int(self.int2gray(i),2) for i in range(256)],[int(self.int2gray(i),2)+256 for i in range(320-256)]))
            imgmap = imgmap[gray,:]

            np.save(self.img_graymap_file,imgmap)
            print("Calculating image mapping data...")

        # one column shift in bucket 2
        h,w=imgmap.shape
        imgmap[:,int(w/2):w-1] = imgmap[:,int(w/2)+1:w]

        return imgmap

    def MaskMapCal(self,width=512, ch_num=16):
        if False:
        # if os.path.exists(self.mask_map_file):
            imgmap = np.load(self.mask_map_file)
            print("Loading exist mask mapping data...")
        else:
            a = np.arange(width)
            # a = np.concatenate((a[64:],a[:64]))
            a = a.reshape(ch_num,-1)
            # ch_map = np.arange(ch_num)
            ch_map = [8,  9, 10, 11, 12, 13, 14, 15,
                      0,  1,  2,  3,  4,  6,  5,  7]#, 16, 17, 18, 19]
            a[:,:] = a[ch_map,:]
            a = a.flatten()
            ch_width = int(width/ch_num)
            imgmap = a[0::ch_width]
            for i in range(1,ch_width):
                imgmap = np.concatenate((imgmap, a[i::ch_width]))
            np.save(self.mask_map_file,imgmap)
            print("Calculating mask mapping data...")
        return imgmap

    def average_cal(self, time=100, save=False, dir='./image/ave/'):
        """
        calculate the average value of image
        """
        if not os.path.exists(dir):
            os.mkdir(dir)

        black_img = np.zeros([self.row,self.col*self.ch*self.tab])

        for i in range(time):
            black_img += self.arrange_raw(self.imread())
        black_img /= time
        print(black_img)
        np.save(self.black_img_file,black_img)

    def black_cal(self, time=100):
        """
        Cover camera lens and calculate the average value of black level
        """
        black_img = np.zeros([self.row,self.col*self.ch*self.tab])
        for i in range(time):
            black_img += self.arrange_raw(self.imread())
        black_img /= time
        print(black_img)
        np.save(self.black_img_file,black_img)

    def gain_cal(self, num_samps=100):
        cal_mat_gain = np.zeros([self.row,self.col*self.ch*self.tab],dtype=np.float32)
        if (os.path.exists(self.black_img_file)==False):
            # print(~os.path.exists("./image/black_level_calib.npy"))
            print("Perform black level calibration first")
        else :
            for i in range(num_samps):
                cal_mat_gain += self.arrange_raw(self.imread())

            cal_mat_gain = cal_mat_gain/num_samps
            cal_mat_gain = self.black_img - cal_mat_gain
            col_gain_mean = np.mean(cal_mat_gain,0)
            med_gain = np.median(col_gain_mean[10:300])
            coeff_mat = med_gain/col_gain_mean
            print(coeff_mat)
            np.save(self.gain_mat_file,coeff_mat)

    def gain_cal2d(self,num_samps=100):
        bright_img = np.zeros([self.row,self.col*self.ch*self.tab],dtype=np.float32)
        if (os.path.exists(self.black_img_file)==False):
            # print(~os.path.exists("./image/black_level_calib.npy"))
            print("Perform black level calibration first")
        else :
            for i in range(num_samps):
                bright_img += self.arrange_raw(self.imread())

            # #Black calibrated averaged output
            # bright_img  = self.black_img - bright_img/num_samps
            # bright_img[bright_img<0] = 0
            # med_gain    = np.median(bright_img[:240,:])
            # gain_arr    = med_gain/bright_img
            # print(med_gain,gain_arr)
            # np.save(self.gain_cali_file,gain_arr)

        np.save(self.bright_img_file, bright_img/num_samps)

    def param_set(self, addr, value):
        self.wire_in(self.address['param_en'], 0)
        self.wire_in(self.address['param_addr'], addr)
        self.wire_in(self.address['param_in'], value)
        self.wire_in(self.address['param_en'], 1)
        self.wire_in(self.address['param_en'], 0)
        if self.wire_out(self.address['param_out']) != value:
            logging.error("Parameter set failed. A.{}/V.{}".format(addr,value))
            exit(1)

    def readout_reset(self):
        self.wire_in(self.address['chip_rst'], 0x01)
        self.wire_in(self.address['chip_rst'], 0x00)

    def wire_in_check(self, addr, value):
        self.wire_in(addr, value)
        self.wire_check(addr, value)

    def wire_check(self, addr, value):
        self.wire_in(self.address['wire_check_set'], addr)
        cnt = 0
        #while self.wire_out(self.address['wire_check_o']) != value:
        #    self.wire_in(addr, value)
        #    cnt += 1
        #    if(cnt > 10):
        #        logging.error("Wire check failed. A.{}/V.{}".format(addr,value))
        #        exit(1)

    def memory_test(self, msk_size=320*64,msk_num=10):
        # set mask number
        self.wire_in(self.address['pat_num'], msk_num)
        self.wire_in(self.address['mem_pat_num'], msk_num)
        # memory readback mode
        self.wire_in(self.address['mem_mode'], 0x00)
        # reset
        self.wire_in(self.address['mk_routine_rst'], 0x04)
        # write mode
        self.wire_in(self.address['mk_routine_rst'], 0x02)

        times = msk_size*msk_num
        target = bytearray(os.urandom(times))
        result = bytearray(times)

        self.block_write(self.address['mem_wr'], self.BLOCK_SIZE, target)

        # read mode
        self.wire_in(self.address['mk_routine_rst'], 0x04)
        self.wire_in(self.address['mk_routine_rst'], 0x01)
        self.block_read(self.address['mem_rd'], self.BLOCK_SIZE, result)

        for i in range(times):
            if result[i] != target[i]:
                logging.error("Memory Test failed. Total:{} No.{}". \
                              format(times,i))
                for j in range(10):
                    logging.error("Memory Test failed. R:{:02x}/T:{:02x}". \
                                  format(result[i+j],target[j]))
                exit(1)
        logging.info("Memory Test passed.")

    def fmaskupload(self, file_name,dummy=False):
        assert isinstance(file_name, str), "Need a file name"
        assert file_name.lower().endswith('.bmp'), "Wrong file type, should be .bmp"

        pattern = self.loadMask(file_name)
        result = bytearray(len(pattern))

        # reset
        self.wire_in(self.address['mk_routine_rst'], 0x04)
        # write mode
        self.wire_in(self.address['mk_routine_rst'], 0x02)

        if(dummy):
            pass
        else:
            self.block_write(self.address['mem_wr'], self.BLOCK_SIZE, pattern)

        # read mode
        self.wire_in(self.address['mk_routine_rst'], 0x01)

        # check ddr3 calibration finished.
        while True:
            self.dev.UpdateTriggerOuts()
            if self.istriggered(self.address['fifo_flag'], 0x08):
                break;
            else:
                logging.info("DDR3 calibrating...")
                time.sleep(1)

        # memory readback mode
        #self.wire_in(self.address['mem_mode'], 0x00)
        #self.wire_in(self.address['mk_routine_rst'], 0x04)
        #self.wire_in(self.address['mk_routine_rst'], 0x01)
        #length = self.block_read(self.address['mem_rd'], self.BLOCK_SIZE, result)
        #logging.info("Read back length: {}".  format(length))

        #for i in range(len(pattern)):
        #    if result[i] != pattern[i]:
        #        logging.error("Memory Test failed. Total:{} No.{}". \
        #                      format(len(pattern),i))
        #        for j in range(10):
        #            logging.error("Memory Test failed. R:{:02x}/T:{:02x}". \
        #                          format(result[i+j],pattern[i+j]))
        #        exit(1)
        #logging.info("Memory Test passed.")

        # memory mask upload mode
        self.wire_in(self.address['mem_mode'], 0x01)
        self.wire_in(self.address['mk_routine_rst'], 0x04)
        self.wire_in(self.address['mk_routine_rst'], 0x01)

    def loadMask(self, filename, row=320, ch_num=16):
        """
        row: 2D image row number
        """

        file = open(filename, 'rb') #open bmp file
        file.seek(10,0)
        offset = int.from_bytes(file.read(4), 'little') #where the first pixel starts
        file.seek(18, 0)
        width = int.from_bytes(file.read(4), 'little')
        file.seek(22,0)
        height = int.from_bytes(file.read(4), 'little')

        print("2D image row size: {}".format(row))
        print("mask data size: {}x{}".format(height, width))
        print("mask numbers: {:f}".format(height/row))

        # img = cv2.imread(filename,CV2.IMREAD_GRAYSCALE).astype(bool)

        print(height%row)
        assert height%row == 0, "Wrong mask data size"
        imageSize = int(width * height / 8) #in bytes. Using 1 bit per pixel (black or white)

        pattern = bytearray(imageSize)
        mask = bytearray(imageSize)

        file.seek(offset, 0)
        mask = file.read(imageSize)
        file.close()

        toBin = BitArray(bytes=mask).bin
        toArray = np.array(list(toBin),dtype='|S1').reshape(height, width)
        p = toArray[:, self.mkmapping]
        p = p.flatten().tobytes()
        pattern = bytearray(int(p[i:i+8],2) for i in range(0, len(p), 8))

        # resize bytearray to fit blockwrite length
        size = (int(imageSize/512)+1)*512 - imageSize
        logging.info("Added Size: {}".format(imageSize/512.))
        #extra = pattern + pattern + pattern+ pattern+ pattern+ pattern+ pattern+ pattern+ bytearray(size)
        extra = pattern+ bytearray(size)
        print("data arragement finished.")

        return pattern
        #return extra

    def arrange1(self, image):
        col, row, tab, ch = 108, 320, 2, 3

        r_margin = np.full((tab, col*ch), 2**16-1)
        c_margin = np.full((row*tab+tab, 2), 2**16-1)
        temp=np.frombuffer(image, dtype=np.uint16)[:tab*row*col*4].reshape(-1,4)
        img=np.concatenate(np.vsplit(temp[:,:ch],row*tab),axis=1).transpose().reshape(-1,col*ch)
        img=np.concatenate((img, r_margin), axis=0)
        img=np.concatenate((c_margin, img), axis=1)

        i = np.arange(row*tab+tab)
        i = np.concatenate((i[0::tab], i[1::tab]))
        #return img[i,:]
        return self.image_scale(img[i,:])

    def arrange(self, image):
        #temp=np.frombuffer(image, dtype=np.uint16)[:tab*row*col*4].reshape(-1,4)
        #img=np.concatenate(np.vsplit(temp[:,:ch],row*tab),axis=1).transpose().reshape(-1,col*ch*tab)
        #img=np.concatenate(np.vsplit(temp[:,:ch],row*tab),axis=1).transpose().reshape(-1,col*ch)
        # get image data
        #temp=np.frombuffer(image, dtype=np.uint16)[:self.tab*self.row*self.col*4]
        temp=np.frombuffer(image, dtype=np.uint16)
        # delete dummy data
        temp = temp.reshape(-1,4)[:,:self.ch].flatten()
        # mapping data
        img = temp[self.mapping]
        #print(" ".join('{:08b}'.format(x) for x in img[1,4:14]))
        #print(img[1][:10])

        logging.debug(np.mean(img[:,6:312]))
        # logging.debug(" ".join('{:014b}'.format(x) for x in img[1,4:14]))

        # correct left right flip in each individual bucket
        img_yflip = np.zeros(img.shape)
        img_yflip[:,:324] = img[:,:324][:,::-1]
        img_yflip[:,324:] = img[:,324:][:,::-1]

        img_yflip = img

        if(self.gain_flag):
            logging.debug("gain")
            return self.image_scale2(img_yflip)
            #return self.image_scale2(img)
        elif(self.gain_scale_flag):
            logging.debug("gain scale")
            return self.image_scale4(img_yflip)
            #return self.image_scale4(img)
        elif(self.black_flag):
            logging.debug("black")
            return self.image_scale1(img_yflip)
            #return self.image_scale1(img)
        elif(self.black_scale_flag):
            logging.debug("black scale")
            return self.image_scale3(img_yflip)
            #return self.image_scale3(img)
        elif(self.scale_flag):
            logging.debug("normal scale")
            return self.image_scale(img_yflip)
            #return self.image_scale(img)
        else:
            logging.debug("raw")
            return img_yflip
            #return img

    def arrange_raw(self, image, rows):
        taps, colsPerADC, ADCsPerCh, ch = 2, 2, 20, 12
        cols = colsPerADC*ADCsPerCh*ch

        img = np.frombuffer(image,dtype=np.uint8)

        img2=img.reshape(-1,8,4)[:,:,::-1]
        img3=np.unpackbits(img2,axis=2).reshape(-1,256) # recover the bits as stored in fifo

        img4 = img3[:,16:] # get rid of padded bits

        noch    = 12 # total number of channels
        nob     = 12 # number of bits per channel
        img5 = img4.reshape(-1,nob,noch,20)         # convert into 12 separate data channels
        # img6 = np.moveaxis(img5,[1,2,3],[3,1,2])[:,:,::-1]
        img6 = np.moveaxis(img5,1,3)[:,:,:,::-1]
        # img6 = img6[:,:,::-1]
        # at this point
        # AXES:
        #    0: Each conversion (#rows x #taps x #cols/ADC)
        #    1: each digital channel
        #    2: columns serialized by each channel
        #    3: the bits for each column

        img7_shape = np.array(img6.shape); img7_shape[-1] = 16
        img7    = np.zeros(img7_shape,dtype=bool)
        img7[:,:,:,-12:] = img6.copy()
        # img8 = np.packbits(img7,axis=-1).view(np.uint16).reshape(img7_shape[:-1])
        img8    = np.packbits(img7,axis=-1).astype(np.uint16)
        img8    = img8[:,:,:,0]*256+img8[:,:,:,1]
        # col_map = np.arange(20)
        # col_map = [19,9,14,4,18,8,13,3,16,6,11,1,17,7,12,2,15,5,10,0]
        col_map = [0, 4, 12, 8, 16, 2, 6, 14, 10, 18, 1, 5, 13, 9, 17, 3, 7, 15, 11, 19]

        ch_map  = [0,1,2,6,5,3,4,7,8,9,10,11]
        img8    = img8[:,ch_map,:]
        img9 = img8[:,:,[col_map]].reshape(rows,taps,colsPerADC,ADCsPerCh*ch)

        img10 = np.zeros((rows,taps,colsPerADC*ADCsPerCh*ch),dtype=np.uint16)
        img10[:,:,0::2] =   img9[:,:,0,:]
        img10[:,:,1::2] =   img9[:,:,1,:]

        img11 = np.concatenate([img10[:,0,:],img10[:,1,:]],axis=1)

        return img11

    def arrange_raw_T7(self,image,rows):
        tics = [['img0',time.time()]]
        ##rows:480, col:680, taps:2
        taps, colsPerADC, ADCsPerCh, ch = 2, 2, 20, 17
        cols = colsPerADC*ADCsPerCh*ch

        img = np.frombuffer(image,dtype=np.uint8)
        tics.append(['img1',time.time()])

        img2=img.reshape(-1,8,4)[:,:,::-1]
        tics.append(['img2',time.time()])

        img3=np.unpackbits(img2,axis=2).reshape(-1,256) # recover the bits as stored in fifo
        tics.append(['img3',time.time()])

        img4 = img3[:,1:] # get rid of padded bits

        #reorder columns from 15 split unevenly to all 20 in a row
        img4_1 = img4.reshape(-1,17,15)
        img4_1 = np.moveaxis(img4_1,0,1)
        img4_1 = img4_1.reshape(-1,(int(np.shape(img4_1)[1]*15/20)),20)
        img4_1 = np.moveaxis(img4_1,0,1)
        img4_1 = img4_1.reshape(-1,340)
        tics.append(['img4',time.time()])

        noch    = 17 # total number of channels
        nob     = 12 # number of bits per channel
        img5 = img4_1.reshape(-1,nob,noch,20)         # convert into 12 separate data channels
        tics.append(['img5',time.time()])

        # img6 = np.moveaxis(img5,[1,2,3],[3,1,2])[:,:,::-1]
        img6 = np.moveaxis(img5,1,3)[:,:,:,::-1]
        tics.append(['img6',time.time()])

        # at this point
        # AXES:
        #    0: Each conversion (#rows x #taps x #cols/ADC)
        #    1: each digital channel
        #    2: columns serialized by each channel
        #    3: the bits for each column

        # img7_shape       = np.array(img6.shape); img7_shape[-1] = 16
        # img7             = np.zeros(img7_shape,dtype=bool)
        # img7[:,:,:,-12:] = img6
        # pad_loc          = np.zeros((4,2),dtype=int);pad_loc[-1,0]=16-nob
        # img7             = np.pad(img6,[[0,0],[0,0],[0,0],[4,0]])
        img7    = img6
        tics.append(['img7',time.time()])
        img8    = np.packbits(img7,axis=-1).astype(np.uint16)
        tics.append(['img8_packbits',time.time()])
        # img8    = img8[:,:,:,0]*256+img8[:,:,:,1]
        img8      = (img8[:,:,:,0]<<4).astype(np.uint16) + (img8[:,:,:,1]>>4).astype(np.uint16)
        # col_map = np.arange(20)
        # col_map = [19,9,14,4,18,8,13,3,16,6,11,1,17,7,12,2,15,5,10,0]
        col_map = [0, 4, 12, 8, 16, 2, 6, 14, 10, 18, 1, 5, 13, 9, 17, 3, 7, 15, 11, 19]

        ch_map  = np.arange(17)
        ch_map  = [0,1,2,3,4,5,6,7,11,10,8,9,12,13,14,15,16]

        img8    = img8[:,ch_map,:]
        img9 = img8[:,:,[col_map]].reshape(rows,taps,colsPerADC,ADCsPerCh*ch)
        tics.append(['img9',time.time()])

        img10 = np.zeros((rows,taps,colsPerADC*ADCsPerCh*ch),dtype=np.uint16)
        img10[:,:,0::2] =   img9[:,:,0,:]
        img10[:,:,1::2] =   img9[:,:,1,:]
        tics.append(['img10',time.time()])

        img11 = np.concatenate([img10[:,0,:],img10[:,1,:]],axis=1)
        tics.append(['img11',time.time()])

        img12 = img11

        # 2x ODDCOL_EN time
        # img12 = np.zeros(img11.shape, dtype=np.uint16)
        # img12[:,1:680:2] = img11[:,0:680:2]
        # img12[:,0:680:2] = img11[:,680::2]
        # img12[:,681::2] = img11[:,1:680:2]
        # img12[:,680::2] = img11[:,681::2]

        # 2x ROWADD cycles
        img12 = np.zeros(img11.shape, dtype=np.uint16)
        img12[0::2,1:680:2] = img11[:240,0:680:2]
        img12[1::2,1:680:2] = img11[:240,680::2]
        img12[0::2,0:680:2] = img11[240:,0:680:2]
        img12[1::2,0:680:2] = img11[240:,680::2]

        img12[0::2,681::2] = img11[:240,1:680:2]
        img12[1::2,681::2] = img11[:240,681::2]
        img12[0::2,680::2] = img11[240:,1:680:2]
        img12[1::2,680::2] = img11[240:,681::2]

        # for i in range(1,len(tics)):
        #     logging.info("arrange_raw_t7 {}:{}".format(tics[i][0],(tics[i][1]-tics[i-1][1])*1000))

        return img12

    def arrange_adc2(self, image, rows=480, ADCsPerCh=20):
        tics = [['img0',time.time()]]
        ##rows:480, col:680, taps:2
        taps, colsPerADC, ch = 1, 2, 17
        cols = colsPerADC*ADCsPerCh*ch
        # rows = 480
        nuImgs = round(len(image)/taps/cols/rows*8*255/256)
        img = np.frombuffer(image,dtype=np.uint8)
        tics.append(['img1',time.time()])

        img2=img.reshape(-1,8,4)[:,:,::-1]
        tics.append(['img2',time.time()])

        img3=np.unpackbits(img2,axis=2).reshape(-1,256) # recover the bits as stored in fifo
        tics.append(['img3',time.time()])

        img4 = img3[:,1:] # get rid of padded bits

        #reorder columns from 15 split unevenly to all 20 in a row
        img4_1 = img4.reshape(-1,17,15)
        img4_1 = np.moveaxis(img4_1,0,1)
        img4_1 = img4_1.reshape(-1,(int(np.shape(img4_1)[1]*15/ADCsPerCh)),ADCsPerCh)
        img4_1 = np.moveaxis(img4_1,0,1)
        img4_1 = img4_1.reshape(-1,17*ADCsPerCh)
        tics.append(['img4',time.time()])

        noch    = 17 # total number of channels
        nob     = 1  # number of bits per channel
        img5 = img4_1.reshape(-1,nob,noch,ADCsPerCh)         # convert into 12 separate data channels
        tics.append(['img5',time.time()])

        # img6 = np.moveaxis(img5,[1,2,3],[3,1,2])[:,:,::-1]
        img6 = np.moveaxis(img5,1,3)[:,:,:,::-1]
        tics.append(['img6',time.time()])

        # img3 = np.unpackbits(img2,axis=2).reshape(-1,256*4).astype(bool)
        # tics.append(['img3',time.time()])
        # img6 = img3[:,self.adc2_map_1row]
        # tics.append(['img6',time.time()])
        # img6=img6.reshape(-1,ch,ADCsPerCh,nob)
        # tics.append(['img6_reshape',time.time()])

        # at this point
        # AXES:
        #    0: Each conversion (#rows x #taps x #cols/ADC)
        #    1: each digital channel
        #    2: columns serialized by each channel
        #    3: the bits for each column

        img7_shape = np.array(img6.shape); img7_shape[-1] = 1; img7_shape[2]=ADCsPerCh
        img7    = np.zeros(img7_shape,dtype=bool)
        img7[:,:,:ADCsPerCh,-nob:] = img6.copy()
        tics.append(['img7',time.time()])

        # img8 = np.packbits(img7,axis=-1).view(np.uint16).reshape(img7_shape[:-1])
        img8    = img7.copy()
        tics.append(['img8',time.time()])

        # col_order = np.arange(ADCsPerCh)
        col_order =  [19, 9, 14, 4, 18, 8,13, 3,  16, 6, 11, 1, 17, 7, 12, 2,15, 5, 10, 0]
        # col_map =  [0, 4, 12, 8, 16, 2, 6, 14, 10, 18, 1, 5, 13, 9, 17, 3, 7, 15, 11, 19]
        col_map = np.argsort(col_order[:ADCsPerCh])[::-1]

        ch_map  = np.arange(ch)
        ch_map  = [0,1,2,3,4,5,6,7,11,10,8,9,12,13,14,15,16]

        img8    = img8[:,:,col_map]
        tics.append(['img8_mapped',time.time()])

        img9 = img8[:,ch_map,:].reshape(nuImgs,rows,taps,colsPerADC,ADCsPerCh*ch)
        tics.append(['img9',time.time()])

        img10 = np.zeros((nuImgs,rows,taps,colsPerADC*ADCsPerCh*ch),dtype=np.uint16)
        img10[:,:,:,0::2] =   img9[:,:,:,0,:]
        img10[:,:,:,1::2] =   img9[:,:,:,1,:]
        tics.append(['img10',time.time()])

        # for i in range(1,len(tics)):
        #     logging.info("arrange_adc2 {}:{}".format(tics[i][0],(tics[i][1]-tics[i-1][1])*1000))

        return img10.reshape(nuImgs,rows,cols)

    def image_scale6(self,image,gain=False,black=True,dynamic=False,max_scale = 3000):

        [row,col,tab] = [480, 680, self.tab]

        x1 = self.black_img.astype(np.float32)
        y1 = np.median(x1)

        y2 = np.zeros((row,col*tab)) # target median

        if(gain):
            # assert black==True, "Black cal must be ON if gain needs to be ON"
            x2 = self.bright_img.astype(np.float32)
            y2_left  = np.median(x2[:,30:col-10].flatten()) # target median left tap
            y2_right = np.median(x2[:,col+30:2*col-10].flatten()) # target median right tap

            y2[:,:col] = y2_left
            y2[:,col:] = y2_right

            # y2 = np.median(x2)
            [num,den] = [(y2-y1),(x2-x1)]
            num[den==0] = 0
            den[den==0] = np.inf
            m = num/den
            c = y1-num/den*x1

            logging.debug(np.mean(m[100:110,100:110]))
            img = m*image + c
            img = y1 - img
            img[img<0] = 0

        elif(black):
            img = self.black_img - image
            img[img<0] = 0
        elif(dynamic):
            percentile = 5
            min_lvl = np.percentile(np.concatenate((image[:,:320],image[:,324:644]),axis=1),5)
            max_lvl = np.percentile(image,95)
            img = 2**16-1 - (np.clip(image, min_lvl, max_lvl)-min_lvl)*(2**16-1)/(max_lvl-min_lvl)
            return img.astype(np.uint16)
        else:
            img = image.astype(np.float32)
        # Clip values between 0 and 3000 and scale to 0 to 2**16-1
        img = (np.clip(img,0,max_scale)*(2**16-1)/max_scale).astype(np.uint16)

        return img

    def imread(self):
        # container = bytearray(self.frame_length)
        cnt = 0
        cnt1=1
        #self.trigger_in(self.address['trigger'], 0)
        while True:

            #time.sleep(0.05)
            self.dev.UpdateTriggerOuts()
            if self.istriggered(self.address['fifo_flag'], 0x02):
                logging.debug("fifo_cnt: "+str(self.wire_out(self.address['fifo_cnt'])))
                self.dev.UpdateTriggerOuts()
                self.read(self.address['image'], self.adc1_container)
                #print("".join('{:02x}'.format(x) for x in container[:8]))
                #print(" ".join('{:08b}'.format(x) for x in container[:8]))
                #toBin = BitArray(bytes=container).bin
                #print(toBin[:16]+" "+toBin[16:32]+" "+toBin[32:48]+" "+toBin[48:64])
                #self.trigger_in(self.address['trigger'], 0)
                return self.adc1_container

            cnt += 1
            #self.state_check();
            if cnt % 1000 == 999:
                bt=self.wire_out(self.address['fifo_cnt'])
                print("fifo_cnt: "+str(bt))
                print(cnt1)
                cnt1 += 1
                print("empty: "+str(self.istriggered(self.address['fifo_flag'], 0x04)))
                print("prg_full: "+str(self.istriggered(self.address['fifo_flag'], 0x02)))
                print("full: "+str(self.istriggered(self.address['fifo_flag'], 0x01)))
                #self.read(self.address['image'], container)

    def adc2_read(self,NSUB):
        cnt = 0
        self.dev.UpdateTriggerOuts()
        while True:
            if self.istriggered(self.address['fifo_flag'], 0x10):
                self.read(self.address['adc2_rd'],self.adc2_container)
                self.dev.UpdateTriggerOuts()

            cnt+=1
            if(cnt%1000==999):
                logging.info('stuck reading adc2 cnt:{:04d}'.format(cnt))
            # t6.read(0xa3,pread)
            return self.adc2_container

    def imsave(self, image, dir, filename,full=False):
        if not os.path.exists(dir):
            os.mkdir(dir)
        #pattern =re.compile(".png")
        #if pattern.match(filename):

        if filename.find(".png")!=-1:
            #print(image.dtype)
            if full:
                numpngw.write_png(dir+filename,image)
            else:
                #scipy.misc.imsave(dir+filename,image)
                cv2.imwrite(dir+filename, image)
        #pattern =re.compile(".npy")
        #if pattern.match(filename):
        if filename.find(".npy")!=-1:
            #print(image.dtype)
            np.save(dir+filename,image)
        if filename.find(".npz")!=-1:
            #print(image.dtype)
            np.savez_compressed(dir+filename,img=image)

if __name__ == "__main__":

    #bitfile = "./bitfile/t4/Reveal_Top_multi3.bit"
    bitfile="bitfile/Reveal_Top_3.bit"
    # bitfile="bitfile/Reveal_Top.bit"

    t6=T6(bitfile)
    # t6.param_set(self.param['Select'], 0x80000000)  # disable slave mode
    # t6.fmaskupload("./maskfile/t4_merge.bmp")
    # t6.set_exposure(10000)
    # t6.set_pattern_num(10)

    input("Press enter after SPI")

    t6.UploadMaskDummy('',10)
    t6.SetMaskParam(5,2)
    t6.SetExposure(2000)
    t6.readout_reset()
    t6.gain_flag = False
    t6.black_flag = False
    t6.black_scale_flag = False
    t6.imshow()


# coding: utf-8

# In[3]:

get_ipython().run_line_magic('config', "InlineBackend.figure_format = 'retina'")
get_ipython().run_line_magic('matplotlib', 'inline')

import numpy as np
import matplotlib.pyplot as plt

def iterate_map(f, x_0, n):
    # extended for multi-dim case
    # trajectory = [x_0] * (n+1)
    trajectory = np.stack([x_0] * (n+1), axis=0)
    trajectory[0] = x_0
    for i in range(1, n + 1):
        # should probably just pass in vector instead of unpacking
        trajectory[i] = f(*trajectory[i-1])

    return trajectory

henon = lambda a, b: lambda x, y: np.asarray((a - (x**2) + b * y, x))
henon_inv = lambda a, b: lambda x, y: np.asarray((y, (x + y**2 - a) / b))

# In[4]:

# def problem_1():
#     trapezoid =
def make_polygon(corners, total_pts):
    pts_per_side = total_pts // len(corners)
    points = np.array([]).reshape(0, 2)

    for i in range(len(corners)):
        (x_start, y_start) = corners[i]
        (x_end, y_end) = corners[(i+1) % len(corners)]

        side_points = np.array([np.linspace(x_start, x_end, pts_per_side),
                                np.linspace(y_start, y_end, pts_per_side)]).T

#         np.linspace(x_start, x_end, pts_per_side)
#         np.interp(,  [x_start, x_end], [y_start, y_end])
        points = np.concatenate((points, side_points), axis=0)

    return points

# In[5]:

def recur(fn, x_in, y_in, remaining):
    if remaining == 0:
        return (x_in, y_in)
    else:
        (x_out, y_out) = fn(x_in, y_in)
        return recur(fn, x_out, y_out, remaining - 1)

def problem_1():
    a = 4
    b = 0.2

    fig1, ax1 = plt.subplots(figsize=(8, 8))
    fig2, ax2 = plt.subplots(figsize=(8, 8))

    ax1.set_aspect('equal')
    ax1.set_ylim(-6, 6)
    ax1.set_xlim(-6, 6)

    ax2.set_aspect('equal')
    ax2.set_ylim(-6, 6)
    ax2.set_xlim(-6, 6)

    trapezoid = make_polygon([[-3, -3], [3, -1.5], [3, 1.5], [-3, 3]], 10**6)
    trapezoid_x = trapezoid[:,0].flatten()
    trapezoid_y = trapezoid[:,1].flatten()
    dist = trapezoid_x + trapezoid_y

    # Plot trapezoid, f1 and f1_inv
    ax1.scatter(trapezoid_x, trapezoid_y, c=dist, s=0.1, cmap=plt.get_cmap('cool'))

    trapezoid_f1 = henon(a, b)(trapezoid_x, trapezoid_y)
    ax1.scatter(trapezoid_f1[0], trapezoid_f1[1], c=dist, s=0.1, cmap=plt.get_cmap('cool'))

    trapezoid_f1_inv = henon_inv(a, b)(trapezoid_x, trapezoid_y)
    ax1.scatter(trapezoid_f1_inv[0], trapezoid_f1_inv[1], c=dist, s=0.1, cmap=plt.get_cmap('cool'))

    # Plot trapezoid, f4 and f4_inv
    ax2.scatter(trapezoid_x, trapezoid_y, s=0.1, color="tomato")

    henon4 = lambda a, b: lambda x, y: recur(henon(a, b), x, y, 4)
    henon4_inv = lambda a, b: lambda x, y: recur(henon_inv(a, b), x, y, 4)

    trapezoid_f4 = henon4(a, b)(trapezoid_x, trapezoid_y)
    trapezoid_f4_inv = henon4_inv(a, b)(trapezoid_x, trapezoid_y)

    ax2.scatter(trapezoid_f4[0], trapezoid_f4[1], s=0.1, color="orange")
    ax2.scatter(trapezoid_f4_inv[0], trapezoid_f4_inv[1], s=1, color="dodgerblue")

# In[6]:

problem_1()

# In[124]:

# .plot is faster than .scatter for large # of data points
def problem_1_fast():
    a = 4
    b = 0.2

    fig1, ax1 = plt.subplots(figsize=(8, 8))
    fig2, ax2 = plt.subplots(figsize=(8, 8))

#     https://matplotlib.org/gallery/text_labels_and_annotations/custom_legends.html#sphx-glr-gallery-text-labels-and-annotations-custom-legends-py
#     color_cycle = plt.get_cmap('Pastel2')(np.linspace(0, 1, 3))
#     color_cycle = ['#E0BBE4', '#957DAD', '#D291BC']
    color_cycle = ['#FFCDCD', '#FFC961', '#FF927A']
    ax1.set_prop_cycle(color=color_cycle)
    ax2.set_prop_cycle(color=color_cycle)

    ax1.set_aspect('equal')
    ax1.set_ylim(-6, 6)
    ax1.set_xlim(-6, 6)

    ax2.set_aspect('equal')
    ax2.set_ylim(-6, 6)
    ax2.set_xlim(-6, 6)

    trapezoid = make_polygon([[-3, -3], [3, -1.5], [3, 1.5], [-3, 3]], 10**6)
    trapezoid_x = trapezoid[:,0].flatten()
    trapezoid_y = trapezoid[:,1].flatten()
    dist = trapezoid_x + trapezoid_y

    # Plot trapezoid, f1 and f1_inv
    ax1.plot(trapezoid_x, trapezoid_y, 'o', markersize=.1)

    trapezoid_f1 = henon(a, b)(trapezoid_x, trapezoid_y)
    ax1.plot(trapezoid_f1[0], trapezoid_f1[1], 'o', markersize=.1)

    trapezoid_f1_inv = henon_inv(a, b)(trapezoid_x, trapezoid_y)
    ax1.plot(trapezoid_f1_inv[0], trapezoid_f1_inv[1], 'o', markersize=.1)

    # Plot trapezoid, f4 and f4_inv
    ax2.plot(trapezoid_x, trapezoid_y, 'o', markersize=.1)

    henon4 = lambda a, b: lambda x, y: recur(henon(a, b), x, y, 4)
    henon4_inv = lambda a, b: lambda x, y: recur(henon_inv(a, b), x, y, 4)

    trapezoid_f4 = henon4(a, b)(trapezoid_x, trapezoid_y)
    trapezoid_f4_inv = henon4_inv(a, b)(trapezoid_x, trapezoid_y)

    ax2.plot(trapezoid_f4[0], trapezoid_f4[1], 'o', markersize=.1)
    ax2.plot(trapezoid_f4_inv[0], trapezoid_f4_inv[1], marker='o', markersize=.1)

# In[125]:

problem_1_fast()

# In[7]:

import functools

def make_polygon_col(corners, total_pts):
    pts_per_side = total_pts // len(corners)
    points = np.array([]).reshape(2, 0)

    for i in range(len(corners)):
        (x_start, y_start) = corners[i]
        (x_end, y_end) = corners[(i+1) % len(corners)]

        side_points = np.array([np.linspace(x_start, x_end, pts_per_side),
                                np.linspace(y_start, y_end, pts_per_side)])

        points = np.concatenate((points, side_points), axis=1)

    return points

def make_ifs(fns, probabilities=[]):
    # create prob thresholds
    if not probabilities:
        probabilities = [1/len(fns)] * len(fns)
    thresholds = functools.reduce(lambda l, el: l + [l[-1] + el], probabilities, [0])[1:]

    def apply_ifs(x): # let x be a vector
        random = np.random.random()
        for i in range(len(thresholds)):
            if random < thresholds[i]:
                return fns[i](x)
        else:
            return fns[len(fns) - 1](x)

    return apply_ifs

def iterate_map_vec(f, x_0, n):
    # extended for multi-dim case
    # trajectory = [x_0] * (n+1)
    trajectory = np.stack([x_0] * (n+1), axis=0)
    trajectory[0] = x_0
    for i in range(1, n + 1):
        # takes in vector instead of params
        trajectory[i] = f(trajectory[i-1])

    return trajectory

def problem_2():

    square_x, square_y = make_polygon_col([[0, 0], [0, 1], [1, 1], [1, 0]], 10**4)
    square = [square_x, square_y]

    f1 = lambda x: np.array((x[0]/2, x[1]/2))
    f2 = lambda x: np.array((x[0]/2 + 1/4, x[1]/2 + 1/2))
    f3 = lambda x: np.array((x[0]/2 + 1/2, x[1]/2))

    ifs = make_ifs([f1, f2, f3])
    trajectory = iterate_map_vec(ifs, [0.5, 0.5], 10**5)

    fig, ax = plt.subplots(figsize=(8, 8))
    ax.set_aspect('equal')

    square_f1 = np.apply_along_axis(f1, 0, square)
    square_f2 = np.apply_along_axis(f2, 0, square)
    square_f3 = np.apply_along_axis(f3, 0, square)

    ax.scatter(square_x, square_y, s=1)
    ax.scatter(square_f1[0], square_f1[1], s=1)
    ax.scatter(square_f2[0], square_f2[1], s=1)
    ax.scatter(square_f3[0], square_f3[1], s=1)
    ax.scatter(trajectory[:,0].flatten(), trajectory[:,1].flatten(), s=0.1, color="mediumpurple")

# In[246]:

problem_2()
# scaling dimension: ln(3) / ln(2)

# In[9]:

# np.apply_along_axis(lambda line: print(line), 0, [[1, 2, 3], [4, 5, 6]])
# np.apply_along_axis(lambda line: print(line), 1, [[1, 2, 3], [4, 5, 6]])

# In[10]:

# deterministic IFS
def make_difs(fns, probabilities=[]):
    # create prob thresholds
    if not probabilities:
        probabilities = [1/len(fns)] * len(fns)
    thresholds = functools.reduce(lambda l, el: l + [l[-1] + el], probabilities, [0])[1:]

    def apply_ifs(x): # let x be a vector of inputs and z at end
        z = x[-1]
        z_shifted = (z*len(fns)) % 1

        for i in range(len(thresholds)):
            if z < thresholds[i]:
                return [*fns[i](x), z_shifted]
        else:
            return [*fns[len(fns) - 1](x), z_shifted]

    return apply_ifs

def problem_3():
    f1 = lambda x: np.array((x[0]/2, x[1]/2))
    f2 = lambda x: np.array((x[0]/2 + 1/4, x[1]/2 + 1/2))
    f3 = lambda x: np.array((x[0]/2 + 1/2, x[1]/2))

    difs = make_difs([f1, f2, f3])
    random_xyz = [np.random.random(), np.random.random(), np.random.random()]

    trajectory = iterate_map_vec(difs, random_xyz, 10**6)

    fig, ax = plt.subplots(figsize=(8, 8))
    ax.set_aspect('equal')

    ax.scatter(trajectory[:,0].flatten(), trajectory[:,1].flatten(), s=0.1, color="mediumpurple")

# In[11]:

problem_3()

# In[12]:

def problem_4():
    f1 = lambda x: np.array((x[0]/3, x[1]/3))
    f2 = lambda x: np.array((x[0]/3 + 2/3, x[1]/3))
    f3 = lambda x: np.array((x[0]/3, x[1]/3 + 2/3))
    f4 = lambda x: np.array((x[0]/3 + 2/3, x[1]/3 + 2/3))

    ifs = make_ifs([f1, f2, f3, f4])

    trajectory = iterate_map_vec(ifs, [0.5, 0.5], 10**5)

    fig, ax = plt.subplots(figsize=(8, 8))
    ax.set_aspect('equal')

    ax.scatter(trajectory[:,0].flatten(), trajectory[:,1].flatten(), s=1)

# In[13]:

problem_4()

# In[14]:

def problem_5e():
    c = (2 * np.sqrt(2))
    f1 = lambda x: np.asarray([x[0]/c - x[1]/c + 1/4, x[0]/c + x[1]/c + 1/4])
    f2 = lambda x: np.asarray([x[0]/2 + 1/2, x[1]/2])

    ifs = make_ifs([f1, f2])

    trajectory = iterate_map_vec(ifs, [0.5, 0.5], 10**5)

    fig, ax = plt.subplots(figsize=(8, 8))
    ax.set_aspect('equal')

    ax.scatter(trajectory[:,0].flatten(), trajectory[:,1].flatten(), s=0.1, alpha=0.8)

# In[15]:

problem_5e()

# In[244]:

def get_boxes(x, n):
#     return np.floor((x) * boxes)
    return (x) // (1 / n) # round towards zero

def problem_6():
    fn = henon(1.4, 0.3)
    trajectory = iterate_map(fn, [0.2, 0.2], 10**6 - 1)
    traj_x = trajectory[:,0].flatten()
    traj_y = trajectory[:,1].flatten()

    fig, ax = plt.subplots(figsize=(8, 8))
    ax.plot(traj_x, traj_y, 'o', markersize=0.1, color='#FF8686')

    n = 10**5
    boxes = n ** 2

    box_coords = np.array([get_boxes(traj_x, n), get_boxes(traj_y, n)]).T
    covering_boxes = len(np.unique(box_coords))

    print(covering_boxes)

    print("Box counting dimension: ")
    print(np.log(covering_boxes) / np.log(n))

# In[245]:

problem_6()


import json
from pathlib import Path
import time
import glob
from argparse import ArgumentParser
import math

import tensorflow as tf
from tensorflow import keras

from airsim_gd.dataset.tf_dataset import TFCVDataset
from airsim_gd.vision.nn.utils import train_utils
from airsim_gd.vision.nn import fast_scnn
from airsim_gd.dataset.utils.data import FEATURE_DESCRIPTION

GEN_MODELS = {'fast_scnn': fast_scnn.generate_model}

class CVTrainer(object):
    # TODO: Load params from config yaml file
    def __init__(self, *, train_dir, val_dir, save_dir,
                 arch, sess_name, input_names, output_names,
                 epochs, seed=None, end_learning_rate, batch_size=16, n_aux=0,
                 max_depth, sim_seg_id_list_path, seg_id_map_path):
        self.train_dir = train_dir
        self.val_dir = val_dir
        self.save_dir = Path(save_dir, sess_name)
        self.save_dir.mkdir(exist_ok=True, parents=True)

        self.arch = arch
        self.input_layer_names = input_names
        self.output_layer_names = output_names

        self.epochs = epochs
        self.seed = seed
        self.end_learning_rate = end_learning_rate
        self.batch_size = batch_size

        with open(Path(sim_seg_id_list_path)) as f:
            self.sim_seg_id_list = json.load(f)

        with open(Path(seg_id_map_path)) as f:
            self.id2rgb_dict = json.load(f)

        self.sess_name = time.ctime(time.time()).replace(':', '.') if sess_name is None else sess_name

        print('\nPrepping training dataset...\n')
        self.train_ds = TFCVDataset(tfr_dir=self.train_dir, seed=self.seed, batch_size=self.batch_size,
                                    max_depth=max_depth, output_names=output_names, seg_id_list=self.sim_seg_id_list).generate_dataset()
        print('\nPrepping validation dataset...\n')
        self.val_ds = TFCVDataset(tfr_dir=self.val_dir, seed=self.seed, batch_size=self.batch_size//2,
                                  max_depth=max_depth, output_names=output_names, seg_id_list=self.sim_seg_id_list).generate_dataset()

        file_list = glob.glob(str(Path(self.train_dir, f'*.tfrecord')))
        tf_example = tf.data.TFRecordDataset(file_list[0])
        tf_example = tf_example.map(lambda x: tf.io.parse_single_example(x, FEATURE_DESCRIPTION))
        tf_example = [data for data in tf_example][0]  # only need one
        self.img_shape = (tf_example['image/height'].numpy(), tf_example['image/width'].numpy(), tf_example['image/channels'].numpy())
        self.n_train_data = len(file_list)

    def get_model(self):
        n_classes = len(self.sim_seg_id_list)

        return GEN_MODELS[self.arch](self.img_shape, n_classes)

    def get_callbacks(self):
        ckpt_path = Path(self.save_dir, 'checkpoints')
        cp_cb = keras.callbacks.ModelCheckpoint(filepath=str(ckpt_path), monitor='val_loss', verbose=1, save_best_only=False)

        es_cb = keras.callbacks.EarlyStopping(monitor='val_loss', mode='min', verbose=1, patience=50,
                                              restore_best_weights=True)
        tb_cb = keras.callbacks.TensorBoard(log_dir=str(self.save_dir), histogram_freq=1, write_graph=True, write_images=False,
                                            update_freq='epoch')
        # img_cb = train_utils.TFImageCallback(model_save_dir=self.save_dir, seg_id_list=self.sim_seg_id_list,
        #                                      input_layer_names=self.input_layer_names, output_layer_names=self.output_layer_names,
        #                                      id22rgb_dict=self.id2rgb_dict)
        return [cp_cb, es_cb, tb_cb]

    def load_model(self):
        return keras.models.load_model(self.save_dir)

    def __call__(self, mode):
        callback_list = self.get_callbacks()
        # steps_per_epoch =

        if mode == 'train':
            print('\nCreating model...\n')
            model, loss_dict, loss_weights = self.get_model()
            decay_steps = self.epochs * self.n_train_data
            learning_rate = keras.optimizers.schedules.PolynomialDecay(0.045, decay_steps, self.end_learning_rate,
                                                                       power=0.5)
            optimizer = keras.optimizers.SGD(momentum=0.9, learning_rate=learning_rate)
            model.compile(loss=loss_dict, loss_weights=loss_weights, optimizer=optimizer, metrics=['accuracy'])

            print('\nModel created. Training now...\n')
            history = model.fit(self.train_ds, epochs=self.epochs, validation_data=self.val_ds, callbacks=callback_list)

        elif mode == 'resume':
            print('\nLoading model...\n')
            model = self.load_model()

            print('\nResuming training...\n')
            history = model.fit(self.train_ds, epochs=self.epochs, validation_data=self.val_ds, callbacks=callback_list)

        else:
            raise ValueError(f"Mode {mode} not recognized.")

        print('\nTraining completed.\n')
        model.save(str(self.save_dir), overwrite=False)
        print('\nModel saved.\n')
        return history

def main(args):
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        try:
            # Currently, memory growth needs to be the same across GPUs
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
            print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
        except RuntimeError as e:
            # Memory growth must be set before GPUs have been initialized
            print(e)

    trainer = CVTrainer(train_dir=args.train_dir, val_dir=args.val_dir, save_dir=args.save_dir,
                        arch=args.arch, sess_name=args.sess_name,
                        epochs=args.epochs, seed=args.seed, end_learning_rate=args.end_learning_rate, batch_size=args.batch_size, n_aux=args.n_aux,
                        max_depth=args.max_depth, sim_seg_id_list_path=args.sim_seg_id_list_path, seg_id_map_path=args.seg_id_map_path,
                        input_names=args.input_names, output_names=args.output_names)

    trainer(args.mode)

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument('--train_dir', type=str, default=Path('../data/nominal/train'))
    parser.add_argument('--val_dir', type=str, default=Path('../data/nominal/val'))
    parser.add_argument('--save_dir', type=str, default=Path.cwd().parent.joinpath("model_training"))
    parser.add_argument('--mode', type=str, choices=['train', 'resume'])
    parser.add_argument('--arch', type=str, default='fast_scnn')
    parser.add_argument('--arch_args', type=str, default=None)
    parser.add_argument('--sess_name', type=str, default=time.ctime(time.time()).replace(':', '.'))
    parser.add_argument('--seed', type=int, default=None)
    parser.add_argument('--sim_seg_id_list_path', type=str, default=Path('sim_seg_id_list.json'))
    parser.add_argument('--seg_id_map_path', type=str, default=Path("airsim_gd/dataset/segmentation_id_maps.json"))
    parser.add_argument('--epochs', type=int, default=1000)
    parser.add_argument('--batch_size', type=int, default=16)
    parser.add_argument('--end_learning_rate', type=float, default=0.00001)
    parser.add_argument('--n_aux', type=int, default=0)
    parser.add_argument('--max_depth', type=float, default=100.0)
    parser.add_argument('--input_names', type=str, nargs='+', default=['input_layer'])
    parser.add_argument('--output_names', type=str, nargs='+', default=None)
    args = parser.parse_args()
    main(args)


#!/usr/bin/env python
import time
import math
import rospy
from std_msgs.msg import Float64
from geometry_msgs.msg import Twist
from robot import RobotCalculator

class DeliveryRobot(object):
    def __init__(self):
        rospy.loginfo("CuriosityRoverAckerMan Initialising...")

        self.calculator = RobotCalculator(
            (19.6, 25.1, 25.1, 19.6), (math.pi / 180, 1.0472))  # (1 degree, 60 degree)
        self.publishers_curiosity_d = {}
        self.controller_ns = "beetlebot"
        self.controller_command = "command"
        self.controllers_list = ["back_wheel_L_joint_velocity_controller",
                                "back_wheel_R_joint_velocity_controller",
                                "front_wheel_L_joint_velocity_controller",
                                "front_wheel_R_joint_velocity_controller",
                                "middle_wheel_L_joint_velocity_controller",
                                "middle_wheel_R_joint_velocity_controller",
                                "sterring_B_L_joint_position_controller",
                                "sterring_B_R_joint_position_controller",
                                "sterring_F_L_joint_position_controller",
                                "sterring_F_R_joint_position_controller"]

        for controller_name in self.controllers_list:
            topic_name = "/"+self.controller_ns+"/" + \
                controller_name+"/"+self.controller_command
            self.publishers_curiosity_d[controller_name] = rospy.Publisher(
                topic_name,
                Float64,
                queue_size=1)

        self.wait_publishers_to_be_ready()
        self.init_publisher_variables()
        self.init_state()

        self.cmd_vel_msg = Twist()
        cmd_vel_topic = "/cmd_vel"
        rospy.Subscriber(cmd_vel_topic, Twist, self.cmd_vel_callback)

        rospy.logwarn("beetlebot_AckerMan...READY")

    def cmd_vel_callback(self, msg):
        self.cmd_vel_msg = msg

    def wait_publishers_to_be_ready(self):

        rate_wait = rospy.Rate(10)
        for controller_name, publisher_obj in self.publishers_curiosity_d.iteritems():
            publisher_ready = False
            while not publisher_ready:
                rospy.logwarn("Checking Publisher for ==>" + str(controller_name))
                pub_num = publisher_obj.get_num_connections()
                publisher_ready = (pub_num > 0)
                rate_wait.sleep()
            rospy.loginfo("Publisher ==>" + str(controller_name) + "...READY")

    def init_publisher_variables(self):
        """
        We create variables for more pythonic access access to publishers
        and not need to access any more
        :return:
        """
        # Get the publishers for wheel speed
        self.back_wheel_L = self.publishers_curiosity_d[self.controllers_list[0]]
        self.back_wheel_R = self.publishers_curiosity_d[self.controllers_list[1]]
        self.front_wheel_L = self.publishers_curiosity_d[self.controllers_list[2]]
        self.front_wheel_R = self.publishers_curiosity_d[self.controllers_list[3]]
        self.middle_wheel_L = self.publishers_curiosity_d[self.controllers_list[4]]
        self.middle_wheel_R = self.publishers_curiosity_d[self.controllers_list[5]]
        # Get the publishers for steering
        self.suspension_steer_B_L = self.publishers_curiosity_d[self.controllers_list[6]]
        self.suspension_steer_B_R = self.publishers_curiosity_d[self.controllers_list[7]]
        self.suspension_steer_F_L = self.publishers_curiosity_d[self.controllers_list[8]]
        self.suspension_steer_F_R = self.publishers_curiosity_d[self.controllers_list[9]]

        # Init Messages
        self.back_wheel_L_velocity_msg = Float64()
        self.back_wheel_R_velocity_msg = Float64()
        self.front_wheel_L_velocity_msg = Float64()
        self.front_wheel_R_velocity_msg = Float64()
        self.middle_wheel_L_velocity_msg = Float64()
        self.middle_wheel_R_velocity_msg = Float64()
        self.suspension_steer_B_L_pos_msg = Float64()
        self.suspension_steer_B_R_pos_msg = Float64()
        self.suspension_steer_F_L_pos_msg = Float64()
        self.suspension_steer_F_R_pos_msg = Float64()

    def init_state(self):
        self.set_turning_radius(None, False)
        self.set_wheels_speed(0.0, None, False)

    def set_turning_radius(self, turn_radius, isRotate):

        r_angle = self.calculator.calculateDegree(turn_radius, isRotate)

        self.suspension_steer_B_L_pos_msg.data = r_angle['3']
        self.suspension_steer_B_R_pos_msg.data = r_angle['6']
        self.suspension_steer_F_L_pos_msg.data = r_angle['1']
        self.suspension_steer_F_R_pos_msg.data = r_angle['4']

        self.suspension_steer_B_L.publish(self.suspension_steer_B_L_pos_msg)
        self.suspension_steer_B_R.publish(self.suspension_steer_B_R_pos_msg)
        self.suspension_steer_F_L.publish(self.suspension_steer_F_L_pos_msg)
        self.suspension_steer_F_R.publish(self.suspension_steer_F_R_pos_msg)

    def set_wheels_speed(self, turning_speed, turning_radius, isRotate):
        velocity = self.calculator.calculateVelocity(
            turning_speed, turning_radius, isRotate)

        self.back_wheel_L_velocity_msg.data = velocity['3']
        self.back_wheel_R_velocity_msg.data = velocity['6']
        self.front_wheel_L_velocity_msg.data = velocity['1']
        self.front_wheel_R_velocity_msg.data = velocity['4']
        self.middle_wheel_L_velocity_msg.data = velocity['2']
        self.middle_wheel_R_velocity_msg.data = velocity['5']

        self.back_wheel_L.publish(self.back_wheel_L_velocity_msg)
        self.back_wheel_R.publish(self.back_wheel_R_velocity_msg)
        self.front_wheel_L.publish(self.front_wheel_L_velocity_msg)
        self.front_wheel_R.publish(self.front_wheel_R_velocity_msg)
        self.middle_wheel_L.publish(self.middle_wheel_L_velocity_msg)
        self.middle_wheel_R.publish(self.middle_wheel_R_velocity_msg)

    def move_with_cmd_vel(self):
        wheel_speed = self.cmd_vel_msg.linear.x
        turning_radius = self.cmd_vel_msg.angular.z
        # Check if delivery-robot want to rotate
        isRotate = True if self.cmd_vel_msg.linear.z != 0 else False

        if turning_radius == 0.0:
            turning_radius = None

        rospy.logdebug("turning_radius="+str(turning_radius) + ",wheel_speed="+str(wheel_speed))
        self.set_turning_radius(turning_radius, isRotate)
        self.set_wheels_speed(wheel_speed, turning_radius, isRotate)

if __name__ == "__main__":
    rospy.init_node("beetlebot_AckerMan_node", log_level=rospy.INFO)
    delivery_robot = DeliveryRobot()
    rate = rospy.Rate(10.0)
    while not rospy.is_shutdown():
        delivery_robot.move_with_cmd_vel()
        rate.sleep()


from pymongo import MongoClient, DESCENDING
import logging

class DataReceiver:
    def __init__(self):
        self.client = MongoClient('mongo')
        self.db = self.client["wazuhl"]
        self.rewards = self.db["rewards"]

    def get_latest_rewards(self, last_n=0):
        query = self.rewards.find(sort=[("step", DESCENDING)],
                                  limit=last_n, projection={"reward": 1,
                                                            "step": 1})
        steps_rewards = ((record["step"], record["reward"]) for record in query)
        if not steps_rewards:
            return [], []
        steps, rewards = zip(*steps_rewards)
        logging.info(steps)
        logging.info(rewards)
        return steps, rewards


import jwt
import datetime

# Clave secreta para firmar y verificar el token (puede ser cualquier valor secreto)
clave_secreta = 'mi_clave_secreta'

# Función para generar un token JWT
def generar_token(datos):
    # Definir fecha de expiración del token (en este caso, 1 hora desde ahora)
    fecha_expiracion = datetime.datetime.utcnow() + datetime.timedelta(hours=1)

    # Crear payload (datos) del token
    payload = {
        'datos': datos,
        'exp': fecha_expiracion
    }

    # Generar el token JWT
    token = jwt.encode(payload, clave_secreta, algorithm='HS256')
    return token

# Función para verificar un token JWT
def verificar_token(token):
    try:
        # Decodificar el token y obtener el payload (datos)
        payload = jwt.decode(token, clave_secreta, algorithms=['HS256'])
        return payload['datos']
    except jwt.ExpiredSignatureError:
        return 'El token ha expirado.'
    except jwt.InvalidTokenError:
        return 'Token inválido.'

# Ejemplo de uso
datos_usuario = {
    'id': 123,
    'nombre': 'Usuario Ejemplo',
    'rol': 'admin'
}

# Generar el token
token_generado = generar_token(datos_usuario)
print('Token generado:', token_generado)

# Verificar el token
datos_verificados = verificar_token(token_generado)
print('Datos verificados:', datos_verificados)


import pytest
import unittest.mock
from podcast_agent.services.supabase import SupabaseService

class TestSupabaseService:
    @pytest.fixture
    def mock_upload_file(self):
        with unittest.mock.patch(
            'podcast_agent.services.supabase.SupabaseService.upload_file',
            new_callable=unittest.mock.MagicMock
        ) as _mock:
            yield _mock

    def test_upload_file(self, mock_upload_file):
        mock_upload_file.return_value = True

        service = SupabaseService()

        result = service.upload_file(
            bucket_name='episode-audio-files',
            file_name='test.mp3'
        )

        mock_upload_file.assert_called_once_with(
            bucket_name='episode-audio-files',
            file_name='test.mp3'
        )
        assert result is not None


from django.db import models
from django.contrib.auth.models import User
from django.db.models.signals import post_save
from django.dispatch import receiver

class UserProfile(models.Model):
    user = models.OneToOneField(User, on_delete=models.CASCADE)
    age = models.IntegerField(default= 20)
    longitude = models.DecimalField(max_digits=20, decimal_places=12,default=0)
    latitude = models.DecimalField(max_digits=20, decimal_places=12, default=0)
    location_text = models.TextField(default="")
    aboutMe = models.TextField(default="")

@receiver(post_save, sender= User)
def create_user_profile(sender,instance,created, **kwargs):
     if created:
         UserProfile.objects.create(user= instance)

@receiver(post_save, sender= User)
def save_user_profile(sender,instance,**kwargs):
    instance.userprofile.save()

class Instruments(models.Model):
    user =models.ForeignKey(User, on_delete=models.CASCADE)
    Instrument = models.CharField(max_length=200)
    experience = models.CharField(max_length=50)

    def __str__(self):
        return "%s %s %s" % (self.user, self.Instrument, self.experience)

class urls(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    url = models.URLField(null=True)


import sys
import json
import pickle
import argparse
from lop.nets.ffnn import FFNN
from lop.nets.linear import MyLinear
from lop.algos.bp import Backprop
from lop.algos.cbp import ContinualBackprop
from lop.utils.miscellaneous import *

def expr(params: {}):
    agent_type = params['agent']
    env_file = params['env_file']
    num_data_points = int(params['num_data_points'])
    to_log = False
    to_log_grad = False
    to_log_activation = False
    beta_1 = 0.9
    beta_2 = 0.999
    weight_decay = 0.0
    accumulate = False
    perturb_scale = 0
    if 'to_log' in params.keys():
        to_log = params['to_log']
    if 'to_log_grad' in params.keys():
        to_log_grad = params['to_log_grad']
    if 'to_log_activation' in params.keys():
        to_log_activation = params['to_log_activation']
    if 'beta_1' in params.keys():
        beta_1 = params['beta_1']
    if 'beta_2' in params.keys():
        beta_2 = params['beta_2']
    if 'weight_decay' in params.keys():
        weight_decay = params['weight_decay']
    if 'accumulate' in params.keys():
        accumulate = params['accumulate']
    if 'perturb_scale' in params.keys():
        perturb_scale = params['perturb_scale']

    num_inputs = params['num_inputs']
    num_features = params['num_features']
    hidden_activation = params['hidden_activation']
    step_size = params['step_size']
    opt = params['opt']
    replacement_rate = params["replacement_rate"]
    decay_rate = params["decay_rate"]
    mt = 10
    util_type='adaptable_contribution'
    init = 'kaiming'
    if "mt" in params.keys():
        mt = params["mt"]
    if "util_type" in params.keys():
        util_type = params["util_type"]
    if "init" in params.keys():
        init = params["init"]

    if agent_type == 'linear':
        net = MyLinear(
            input_size=num_inputs,
        )
    else:
        net = FFNN(
            input_size=num_inputs,
            num_features=num_features,
            hidden_activation=hidden_activation,
        )

    if agent_type == 'bp' or agent_type == 'linear' or agent_type == 'l2':
        learner = Backprop(
            net=net,
            step_size=step_size,
            opt=opt,
            beta_1=beta_1,
            beta_2=beta_2,
            weight_decay=weight_decay,
            to_perturb=(perturb_scale > 0),
            perturb_scale=perturb_scale,
        )
    elif agent_type == 'cbp':
        learner = ContinualBackprop(
            net=net,
            step_size=step_size,
            opt=opt,
            replacement_rate=replacement_rate,
            decay_rate=decay_rate,
            device='cpu',
            maturity_threshold=mt,
            util_type=util_type,
            init=init,
            accumulate=accumulate,
        )

    with open(env_file, 'rb+') as f:
        inputs, outputs, _ = pickle.load(f)

    errs = torch.zeros((num_data_points), dtype=torch.float)
    if to_log: weight_mag = torch.zeros((num_data_points, 2), dtype=torch.float)
    if to_log_grad: grad_mag = torch.zeros((num_data_points, 2), dtype=torch.float)
    if to_log_activation: activation = torch.zeros((num_data_points, ), dtype=torch.float)
    for i in tqdm(range(num_data_points)):
        x, y = inputs[i: i+1], outputs[i: i+1]
        err = learner.learn(x=x, target=y)
        if to_log:
            weight_mag[i][0] = learner.net.layers[0].weight.data.abs().mean()
            weight_mag[i][1] = learner.net.layers[-1].weight.data.abs().mean()
        if to_log_grad:
            grad_mag[i][0] = learner.net.layers[0].weight.grad.data.abs().mean()
            grad_mag[i][1] = learner.net.layers[-1].weight.grad.data.abs().mean()
        if to_log_activation:
            if hidden_activation == 'relu':
                activation[i] = (learner.previous_features[0] == 0).float().mean()
            if hidden_activation == 'tanh':
                activation[i] = (learner.previous_features[0].abs() > 0.9).float().mean()
        errs[i] = err

    data_to_save = {
        'errs': errs.numpy()
    }
    if to_log:
        data_to_save['weight_mag'] = weight_mag.numpy()
    if to_log_grad:
        data_to_save['grad_mag'] = grad_mag.numpy()
    if to_log_activation:
        data_to_save['activation'] = activation.numpy()
    return data_to_save

def main(arguments):
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-c', help="Path to the file containing the parameters for the experiment",
                        type=str, default='temp_cfg/0.json')
    args = parser.parse_args(arguments)
    cfg_file = args.c

    with open(cfg_file, 'r') as f:
        params = json.load(f)

    data = expr(params)

    with open(params['data_file'], 'wb+') as f:
        pickle.dump(data, f)

if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))


import numpy as np
from pathlib import Path
import shutil
import subprocess
import os
import h5py

from openquake.hazardlib.source.point import PointSource
from openquake.hazardlib.geo.point import Point
from openquake.hazardlib.geo.nodalplane import NodalPlane
from openquake.hazardlib.pmf import PMF
from openquake.hazardlib.scalerel.point import PointMSR
from openquake.hazardlib.mfd.evenly_discretized import EvenlyDiscretizedMFD
from openquake.hazardlib import sourcewriter
from openquake.hazardlib.tom import PoissonTOM, NegativeBinomialTOM

def run_model(folder_name):
    """
    Runs a folder where the openquake files are found: Job.ini, source model, source logic tree, grid and gmpe logic tree
    :param folder_name: Location of Openquake source files
    :return:
    """
    job_path = os.path.join(folder_name, 'job.ini')
    print('Running Openquake')

    oq_proc = subprocess.call(['oq', 'engine', '--run', job_path], stdout=subprocess.PIPE)
    if oq_proc == 0:
        # If success, gets calculation id
        process_2 = subprocess.run(['oq', 'engine', '--lhc'], stdout=subprocess.PIPE)
        out = str(process_2).split(r'\n')[-2]
        line = out.split()
        print(line)
        calc_id = int(line[0])
        src = os.path.join(Path.home(), 'oqdata', f'calc_{calc_id}.hdf5')
        # Copy calculation hdf5 database into model folder
        shutil.copy(src, folder_name)
        print(f'Done. hdf5 database copied to {folder_name}')
    else:
        #if oq throws error
        raise Exception('Run failed')

def parse_db(folder):
    """
    Reads a hdf5 database and parses the mean curves.  Intensity measure levels (imtl) are fixed.
    :param folder:
    :return:
    """
    files = list(os.walk(folder))[0][2]
    calcs = [i for i in files if 'calc' in i]
    ids = [int(i.split('_')[1].split('.')[0]) for i in calcs]
    calc_path = os.path.join(folder, calcs[int(np.argmax(ids))])

    db = h5py.File(os.path.join(folder, calc_path), 'r')

    grid = np.stack((db['sitecol']['lon'], db['sitecol']['lat'])).T

    hcurves = db['hcurves-stats'][:, 0, 0, :]

    imtl = np.logspace(-2, 0.2, 30)   # Same as all models' job.ini

    data = {'grid': grid,
            'hcurves': hcurves,
            'imtl': imtl}
    return data

class hazard_model:

    def __init__(self, name=None):
        self.name = name
        self.cells = np.array(0, np.dtype(float))
        self.rates = np.array(0, np.dtype(float))
        self.dispersion = np.array(0, np.dtype(float))
        self.rates_mbin = np.array(0, np.dtype(float))
        self.sources = []

        self.mmin = 5.
        self.mmax = 8.
        self.mbin = 0.1
        self.magnitudes = np.arange( self.mmin,  self.mmax + self.mbin,  self.mbin)

    @property
    def nsources(self):
        return self.cells.shape[0]

    def read_forecast(self, filename, eepas=False):
        """
        Parses forecast files
        :param forecast_name: path to forecast
        :param eepas: Flag to indicate if the forecast is just a spatial PDF (False) or, magnitude bins spatial PDFs (True)

        """
        if eepas is False:
            data = np.genfromtxt(filename, delimiter=',', skip_header=1)
            self.cells = np.array([[i[0] + i[1], i[2] + i[3]] for i in data[:, :4]])/2.
            self.rates = data[:, 8]
            self.dispersion = data[:, 9]

        else:
            data = np.genfromtxt(filename, delimiter=',', skip_header=1)
            self.cells = np.array([[i[0] + i[1], i[2] + i[3]] for i in data[:, :4]])/2.
            self.rates_mbin = data[:, 6:-1]
            self.rates = np.sum(self.rates_mbin, axis=1)

            self.dispersion = data[:, -1]

    def scale(self, scale):
        """
        Scales rates and magnitude rates
        :param scale:
        :return:
        """
        self.rates *= scale
        self.rates_mbin *= scale

    def set_trunc_gr(self, bval, eepas=False):
        """
        Projects spatial rates using a truncated Gutenberg-Richter magnitude-frequency-distribution into magnitude bins.
        If the model contains EEPAS, scales the magnitude bins pdf scaling each magnitude PDF by the weights of a truncated GR distribution
        :param bval: b-value of the distribution
        :return:
        """

        mags = self.magnitudes
        mbin = self.mbin
        mmin = self.mmin
        mmax = self.mmax
        rates = self.rates

        weights = (10 ** (- bval * (mags - mbin / 2.)) - 10 ** (- bval * (mags + mbin / 2.))) / (
                   10 ** (- bval * (mmin - mbin / 2.)) - 10 ** (- bval * (mmax + mbin / 2.)))

        if eepas is False:
            rates_mbin = np.zeros((rates.shape[0], mags.shape[0]))
            for i, mu in enumerate(rates):
                rates_mbin[i, :] = mu * weights
            self.rates_mbin = rates_mbin

        else:
            rates_mbin = self.rates_mbin

            for i, mu_m in enumerate(rates_mbin):
                self.rates_mbin[i, :] = weights * mu_m
                self.rates[i] = np.sum(self.rates_mbin[i, :])

    def get_point_srcs(self,
                       time_span=1,
                       tect_region='Active Shallow Crust',
                       rupt_mesh_spacing=5.0,
                       magnitude_scale_rel=PointMSR(),
                       rupt_aspect_ratio=1.0,
                       upper_seis_depth=0,
                       lower_seis_depth=30,
                       nodal_dist=[(1, NodalPlane(0, 90, 0.))],
                       hypo_dist=((0.5, 10.0), (0.5, 30.0))):
        """
        Creates the Openquake objects necessary to write a source nrml file. Assign a Negative Binomial temporal model
        if the dispersion is greater than 0, and a Poisson temporal model if 0.
        """
        npd = PMF(nodal_dist)
        hd = PMF(hypo_dist)

        for n in range(self.nsources):
            mfd = EvenlyDiscretizedMFD(self.mmin, self.mbin, self.rates_mbin[n, :])
            if self.dispersion[n] != 0:
                mu_tot = self.rates[n]
                alpha = self.dispersion[n]
                tom = NegativeBinomialTOM(time_span, mu_tot, alpha)
            else:
                tom = PoissonTOM(1)

            source = PointSource(source_id='%05i' % n,
                                 name='point%05i' % n,
                                 tectonic_region_type=tect_region,
                                 mfd=mfd,
                                 rupture_mesh_spacing=rupt_mesh_spacing,
                                 magnitude_scaling_relationship=magnitude_scale_rel,
                                 rupture_aspect_ratio=rupt_aspect_ratio,
                                 temporal_occurrence_model=tom,
                                 upper_seismogenic_depth=upper_seis_depth,
                                 lower_seismogenic_depth=lower_seis_depth,
                                 location=Point(self.cells[n, 0],
                                                self.cells[n, 1]),
                                 nodal_plane_distribution=npd,
                                 hypocenter_distribution=hd)
            self.sources.append(source)

    def write_source(self, source_path):

        sourcewriter.write_source_model(source_path,
                                        self.sources,
                                        name=self.name,
                                        investigation_time=1)


"""Crawl game URLs in matchday data."""

import datetime
import urllib3
import certifi
import numpy as np
import pandas as pd
import multiprocessing as mp

from bs4 import BeautifulSoup
from bld.project_paths import project_paths_join as ppj

def get_soup_obj(url):
    """
    Returns the soup object from a given url.
    """

    http = urllib3.PoolManager(
        cert_reqs='CERT_REQUIRED', ca_certs=certifi.where())
    page_request = http.request("GET", url)
    soup = BeautifulSoup(page_request.data, 'lxml')

    return soup

def get_mtchdy_data(league_url):
    """
    Returns all game URLs for a given matchday URL. Moreover, it is
    indicated whether the respective game already took place and if not, when
    they will be played.
    """

    # Get games soup.
    matchday_soup = get_soup_obj(league_url)
    games = matchday_soup.findAll(
        "td", {"class": "liga_spielplan_container"})

    # Get game information.
    matchday_df = pd.DataFrame()
    for j, game in enumerate(games):
        game_dict = {}
        game_dict["game_url"] = game.a["href"]
        if game.find("div", {"liga_spieltag_vorschau_datum_content_ergebnis"}) != None:
            game_dict["doable"] = 1
            game_dict["reason"] = 0
        elif game.find("div", {"class": "liga_spieltag_vorschau_datum_content"}) != None:
            game_dict["doable"] = 0
            game_dict["reason"] = game.find(
                "div", {"class": "liga_spieltag_vorschau_datum_content"}).text
        else:
            game_dict["doable"] = 0
            game_dict["reason"] = np.nan

        matchday_df = matchday_df.append(game_dict, ignore_index=True)

    if len(matchday_df) != 0:
        matchday_df["mtchdy_url"] = league_url
        return matchday_df
    else:
        pass

def main():
    # Load league data.
    league_df = pd.read_csv(
        ppj("OUT_DATA_FOOTBALL", "matchday_data.csv"), encoding='cp1252')

    # Scrape matchday data via multiprocessing.
    df_list = []
    with mp.Pool() as pool:
        out = pool.map(get_mtchdy_data, league_df.mtchdy_url.values)
        df_list.extend(out)

    # Create df from df list.
    game_df = pd.concat(df_list)

    # Merge to league data and save as CSV.
    game_df = game_df.merge(league_df, on="mtchdy_url")
    game_df.to_csv(ppj("OUT_DATA_FOOTBALL", "game_urls.csv"), index=False)

if __name__ == '__main__':
    main()