''' ############################ NIRC2 Image Processing ##############################
                             written by Logan Pearce, 2019
    ##################################################################################
    Prepare raw NIRC2 images for science by linearizing, dark subtracting, flat dividing,
    and interpolating over bad pixels.  Python 2.7 based. This script must be located in
    the same directory as the files nirc2.512.512.badpix and nirc2.1024.1024.badpix

# Requires:
#   python packages astropy
#   files nirc2.512.512.badpix and nirc2.1024.1024.badpix
#
# Input:
#   Directory containing NIRC2 science images
#   Directory containing associated NIRC2 dark and flat calibration images
#
# Output:
#      science_image_original_filename_LDFC: linearized, dark subtracted, flat corrected .fits file
#      science_image_original_filename_LDFBC: same as above with bad pixels interpolated over .fits file
#
# Useage:
# image_process.py [-h] [path_to_science_image_directory] [-d ANG_DIAM] [-c CONDITIONS]
#                      [-t TABLE]
#                      filename
#
# positional arguments:
# images                the path to the directory containing science images,
#                       including final '/'
# cals                  the path to directory containing calibration images,
                        including final '/'

# optional arguments:
# -h, --help            show this help message and exit
# -f MASTER_FLAT, --master_flat MASTER_FLAT
#                       input a previously made master flat image.
# -d MASTER_DARK, --master_dark MASTER_DARK
#                       input a previously made master dark image.
#
# examples:
#    python image_process.py /DSTuc/science/2017/ /DSTuc/cal/2017   <- Calibrate all science images of DS Tuc in the 2017 observation
#                            by looking through all of the calibration files from that observation and finding the right ones to apply
#    python image_process.py /DSTuc/science/2017/ /DSTuc/cal/2017 -f flat.20170702.4000.3.5.2.fits   <- Calibrate the
#                            science images by finding the appropriate dark frames, but using this previously made master flat frame.

'''

import numpy as np
from astropy import units as u
from astropy.io import fits
from astropy.io.fits import getheader
import os
import time
import argparse
import warnings
warnings.filterwarnings("ignore")

def mode(a):
    '''Counts the number of unique occurances of values and returns which value occurs the most.
    Made this because the scipy mode function was so very slow on these files.'''
    nums,counts = np.unique(a, return_counts=True)
    return nums[np.argmax(counts)]

def make_master_dark(list):
    ''' Makes a median-combined and linearized master dark file from a stack of calibration images
        Inputs:
            list (array): array of calibration filenames to put into stack
        Returns:
            masterdark (2d array): 2d image array of master dark
    '''
    # Linearizing coefficients:
    coeff = [1.001,-6.9e-6,-0.70e-10]
    # 10 darks is more than enough to make a master:
    if len(list) > 10:
        length = 10
    else:
        length = len(list)
    # Initialize master dark image:
    darkstack = np.zeros((length,imsize1,imsize2))
    # For each dark image;
    for line,i in zip(list,range(length)):
        # Open the header and pull out coadds:
        drkhdr = getheader(line, ignore_missing_end=True)
        coadds = float(drkhdr['coadds'])
        # Open the data file:
        dark1i = fits.open(line, ignore_missing_end=True)
        dark1s = dark1i[0].data
        # Linearize the dark frame:
        norm1 = coeff[0]+(coeff[1]*(dark1s/coadds))+(coeff[2]*((dark1s/coadds)**2))
        dark1 = dark1s/norm1
        # Place into the stack:
        darkstack[i]= dark1

    # Median combine the corrected dark frames into master dark frame:
    masterdark = np.median(darkstack,axis=0)
    fits.writeto('masterdark.fits',masterdark,drkhdr,overwrite=True)
    return masterdark

def make_master_flat(list,masterdark):
    ''' Makes a median-combined, linearized, and dark subtracted master flat from a stack of calibration images
        Inputs:
            list (array): array of calibration filenames to put into stack
        Returns:
            masterflat (2d array): 2d image array of master flat
    '''
    # Linearizing coefficients:
    coeff = [1.001,-6.9e-6,-0.70e-10]
    # 10 flats is more than enough to make a master:
    if len(list) > 10:
        length = 10
    else:
        length = len(list)
    # Initialize flat stack:
    flatstack = np.zeros((length,imsize1,imsize2))
    # For each flat frame:
    for line,i in zip(list,range(length)):
        # Open the header and get the coadds:
        flthdr = getheader(line, ignore_missing_end=True)
        coadds = float(flthdr['coadds'])
        # Open the image:
        flat1i = fits.open(line, ignore_missing_end=True)
        flat1s = flat1i[0].data
        # Linearize:
        norm1 = coeff[0]+(coeff[1]*(flat1s/coadds))+(coeff[2]*((flat1s/coadds)**2))
        flat1l = flat1s/norm1
        # Dark subtract the flat frame:
        flat1=flat1l-masterdark
        # Normalize the flat:
        flat1 = flat1/mode(flat1)
        # Add it to the stack:
        flatstack[i]= flat1
    # Median combine the stack:
    masterflat = np.median(flatstack,axis=0)
    # Normalize the master:
    masterflat = masterflat / np.mean(masterflat)
    # Write out calibrated science frame:
    fits.writeto('masterflat.fits',masterflat,flthdr,overwrite=True)
    return masterflat

def calibrate_science_image(image,master_dark,master_flat,newfilename):
    ''' Makes linearizes, dark-subtracts, and flat-divides science frames
        Inputs:
            image (2d array): image to be calibrated
            masterdark,masterflat (2d array): master dark and flat arrays
            newfilename (str): name for calibrated image file
        Returns:
            flat_div (2d array): calibrated image
            imhdr (fits header): annotated header
    '''
    # Linearizing coefficients:
    coeff = [1.001,-6.9e-6,-0.70e-10]
    # Open science image:
    im1 = fits.open(image, ignore_missing_end=True)
    im = im1[0].data
    # Get coadds:
    imhdr = getheader(image, ignore_missing_end=True)
    coadds = float(imhdr['coadds'])
    # Linearize science frame:
    norm1 = coeff[0]+(coeff[1]*(im/coadds))+(coeff[2]*((im/coadds)**2))
    im = im/norm1
    # Dark subtract and flat divide image:
    flat_div = (im - master_dark) / master_flat
    # Add a comment to the header:
    imhdr['COMMENT'] = '         Linearized, dark subtracted and flat corrected on '+time.strftime("%m/%d/%Y")+ ' By Logan A. Pearce'
    # Write out calibrated science frame:
    fits.writeto(newfilename,flat_div,imhdr,overwrite=True)
    return flat_div, imhdr

def bad_pixel_fix(image,hdr,sizex,sizey,newfilename,s=3):
    ''' Fixes bad pixels by interpolating from surrounding pixels.
        Copied from:
        http://mtham.ucolick.org/egates/2016GradWorkshop/PDFs/DataReduction/DataReductionProcedures-Python-2016.pdf
        Inputs:
            image (2d array): image to be fixed
            hdr (fits header): header of image for getting
            sizex, sizey (float): dimensions of image to be fixed
            newfilename (str): new file name
            s (int): radius of pixels to use in interpolation (default = 3)
        Returns:
            flat_div (2d array): pixel corrected image
    '''

    # Depending on the size of the image, load the correct bad pixel list:
    #     (Python indicies are (row,column) so the image indicies are (y,x) in python; additionally python begins with
    #      initial index of 0, while image data begins at 1)
    if sizex == 1024 and sizey == 1024:
        badpix = np.loadtxt(open("nirc2.1024.1024.badpix","rb"))
        badpix = badpix.astype(int)
    elif sizex == 512 and sizey == 512:
        badpix = np.loadtxt(open("nirc2.512.512.badpix","rb"))
        badpix = badpix.astype(int)
    else:
        print "I don't have a bad pixel list to match this image size.  I only have 1024x1024 or 512x512.  \
        Check your image sizes please."
        quit()

    # Make a boolean mask the same size as the image and set all initial values to False:
    mask = np.ma.make_mask(image,copy=True,shrink=True, dtype=np.bool)
    mask[:,:] = False
    # For each bad pixel in the list, set the value of the mask to true:
    for i in badpix:
        mask[i[1]-1][i[0]-1] = True
    # Set the value of all bad pixels in the image to nan:
    mdata = np.ma.masked_array(image,mask=mask,fill_value=np.nan)
    # Make a new array as a copy of original image:
    badpixelfixed = image.copy()
    # For each x and y value, loop through image and replace the value of all "nan"
    # pixels with the mean of the four pixels on either side of the bad one:
    for i in range(0,mdata.shape[0]):
        for j in range(0,mdata.shape[1]):
            if np.math.isnan(mdata[i,j]):
                x1 = i-s
                x2 = i+s+1
                y1 = j-s
                y2 = j+s+1
                if x1<0:
                    x1 = 0
                if x2>mdata.shape[0]:
                    x2=mdata.shape[0]
                if y1<0:
                    y1 = 0
                if y2>mdata.shape[1]:
                    y2 = mdata.shape[1]
                badpixelfixed[i,j] = np.mean(mdata[x1:x2,y1:y2])
    # Add a comment to the header:
    hdr['COMMENT'] = '         Bad pixels fixed on '+time.strftime("%m/%d/%Y")+ ' By Logan A. Pearce'
    # Write out the bad pixel corrected image to a new fits file:
    fits.writeto(newfilename,badpixelfixed,hdr,overwrite=True)
    return badpixelfixed

# Pull out arguments:
parser = argparse.ArgumentParser()
# Required positional arguments:
parser.add_argument("images", help="the path to the directory containing science images, including final '/' ", type=str)
parser.add_argument("cals", help="the path to directory containing calibration images, including final '/' ", type=str)
# Optional positional arguments"
parser.add_argument("-f","--master_flat", help="input a previously made master flat image.",type=str)
parser.add_argument("-d","--master_dark", help="input a previously made master dark image.",type=str)

args = parser.parse_args()
# Point at the directory containing the images to processed:
images_directory = args.images
# Point at the directory containing all calibration files:
cals_directory = args.cals

os.system('ls '+images_directory+'*.fits > images_list')
os.system('ls '+cals_directory+'*.fits > cals_list')

with open('images_list') as f:
    images = f.read().splitlines()
with open('cals_list') as f:
    cals = f.read().splitlines()

# For each science image:
for image in images:
    print 'Calibrating ',image
    # Open the header:
    imhdr = getheader(image, ignore_missing_end=True)
    # Pull out the relevant header parameters:
    mm,ss,ii,cc= imhdr['MULTISAM'],imhdr['SAMPMODE'],imhdr['ITIME'],imhdr['COADDS']
    imfilt=imhdr['FILTER']
    #print imfilt
    imsize1,imsize2 = imhdr['NAXIS1'],imhdr['NAXIS2']
    #print imsize1,imsize2
    try:
        imkoaid = imhdr['KOAID']
        obsdate = imkoaid.split('.')[1]
    except:
        imkoaid = imhdr['OBJECT'].replace(" ","_")
        obsdate = imhdr['DATE-OBS'].replace("-", "")
    print imkoaid,obsdate

    if args.master_dark:
        pass
    else:
        # Go through each calibration file and find the matching cal files:
        darklist = []
        flatlist = []

        for cal in cals:
            hdr = getheader(cal, ignore_missing_end=True)
            m,s,itime,c = hdr['MULTISAM'],hdr['SAMPMODE'],hdr['ITIME'],hdr['COADDS']
            filt = hdr['FILTER']
            print filt
            try:
                caltype = hdr['KOAIMTYP']
                koaid = hdr['KOAID']
                caldate = koaid.split('.')[1]
            except:
                caltype = hdr['OBJECT']
                caldate = hdr['DATE-OBS'].replace("-", "")
            size1,size2 = hdr['NAXIS1'], hdr['NAXIS2']
            #print size1,size2
            if 'dark' in caltype:
                # If dark, match multisam, sampmode, itime, coadds, and the date
                if m==mm and s==ss and itime==ii and c==cc and size1==imsize1 and size2==imsize2 and caldate==obsdate:
                    darklist.append(cal)
            if 'flat' in caltype:
                # If flat, match filter and date:
                if filt==imfilt and size1==imsize1 and size2==imsize2: #and caldate==obsdate:
                    flatlist.append(cal)
    print darklist
    print flatlist
    # Make master dark:
    if args.master_dark:
        masterdark = fits.open(args.master_dark)[0].data
    else:
        masterdark = make_master_dark(darklist)

    # Make master flat:
    if args.master_flat:
        masterflat = fits.open(args.master_flat)[0].data
    else:
        masterflat = make_master_flat(flatlist,masterdark)

    # Calibrate science frame:
    newfilename = '..'+image.split('.')[2]+'.'+image.split('.')[3]+'.'+image.split('.')[4]+'.LDFC.'+image.split('.')[-1]
    image_ldfc, image_ldfc_hdr = calibrate_science_image(image,masterdark,masterflat,newfilename)

    # Fix bad pixels:
    newfilename = '..'+image.split('.')[2]+'.'+image.split('.')[3]+'.'+image.split('.')[4]+'.LDFBC.'+image.split('.')[-1]
    image_ldfbc = bad_pixel_fix(image_ldfc,image_ldfc_hdr,size1,size2,newfilename,s=3)
    print 'done... next'

os.system('rm images_list')
if args.master_dark:
    pass
else:
    os.system('rm cals_list')


import numpy as np
from grasp.models.robots.robot import Robot
from grasp.utils.mjcf_utils import xml_path_completion

class Sawyer(Robot):

    def __init__(self):
        super().__init__(xml_path_completion("Baxter/baxter/master.xml"))

    def set_base_xpos(self, pos):
        pass

    @property
    def dof(self):
        # 14 + 4 from grippers
        return 18

    @property
    def joints(self):
        out = []
        for s in ['right_', 'left_']:
            out.extend((s + a for a in ["s0", "s1", "e0", "e1", "w0", "w1", "w2"]))

        return out

    @property
    def gripper_joints(self):
        out = []
        for s in ['r_gripper_', 'l_gripper_']:
            out.extend((s + a for a in ['l_finger_joint', 'r_finger_joint']))
        return out

    @property
    def init_qpos(self):
        # joint positions
        # normal_joint = np.array([
        #     0.535, -0.093, 0.038, 0.166, 0.643, 1.960, -1.297,
        #     -0.518, -0.026, -0.076, 0.175, -0.748, 1.641, -0.158])

        joints = np.array([
            0., 0., 0., 0., 0., 0., 0.,
            -0.98473254,  0.05363101,  0.97193983, -0.57413665, -1.02427219, 1.83454327,  0.88820527,
            0.,0.,0.,0.
        ])
        return joints

    @property
    def gripper_init_qpos(self):
        return np.array([0, 0, 0, 0])


from __future__ import division

import redis
import msgpack
import numpy as np

import math
import torch
from torch.autograd import Variable

import sys
from onmt.Dataset import Batch
from onmt.Constants import *

class RedisDataset(object):
    """
    """
    def __init__(self, name, batchSize, cuda, host="localhost", port=6379, db=0, reverse=False, volatile=False, r2l=False):
        """
        """
        self.cuda = cuda
        self.name = name
        self.batchSize = batchSize
        self.volatile = volatile
        self.reverse = reverse ### reverse translation direction
        #self.redis = LoadRedis(name, host, port)
        #self.redis.load(dataPath, batchSize)
        self.r = redis.Redis(host=host, port=port, db=db)
        self.r2l = r2l

        self.numBucket = int(self.r.get(self.name+"_"+"size"))
        assert self.numBucket is not None and self.numBucket > 0
        self.bucket_size = int(self.r.get("bucket_size"))
        assert self.bucket_size > 0 and self.batchSize % self.bucket_size == 0
        self.batchBucket = int(self.batchSize / self.bucket_size)
        self.numBatches = math.ceil(float(self.numBucket) / self.batchBucket)
        indices = np.random.permutation(self.numBucket).tolist()
        ##  for sanity check
        #indices = list(np.arange(self.numBucket))

        self.indices = [ name + "_" + str(indice) for indice in indices]

    def _batchify(self, data, align_right=False,
                  include_lengths=False, dtype="text", features=None):
        if dtype == "text":
            lengths = [len(x) for x in data]
            max_length = max(lengths)
            if features:
                num_features = len(features)
                out = torch.LongTensor(len(data), max_length, num_features + 1) \
                             .fill_(PAD)
                assert (len(data) == len(features[0])), \
                    ("%s %s" % (len(data[0]), len(features[0])))
            else:
                out = torch.LongTensor(len(data), max_length) \
                             .fill_(PAD)
            for i in range(len(data)):
                data_length = len(data[i])
                offset = max_length - data_length if align_right else 0
                if features:
                    out[i].narrow(0, offset, data_length)[:, 0].copy_(torch.LongTensor(data[i]))
                    for j in range(num_features):
                        out[i].narrow(0, offset, data_length)[:, j+1] \
                             .copy_(torch.LongTensor(features[j][i]))
                    #print('num features ', num_features)
                else:
                    out[i].narrow(0, offset, data_length).copy_(torch.LongTensor(data[i]))

            if include_lengths:
                return out, lengths
            else:
                return out
        elif dtype == "img":
            heights = [x.size(1) for x in data]
            max_height = max(heights)
            widths = [x.size(2) for x in data]
            max_width = max(widths)

            out = data[0].new(len(data), 3, max_height, max_width).fill_(0)
            for i in range(len(data)):
                data_height = data[i].size(1)
                data_width = data[i].size(2)
                height_offset = max_height - data_height if align_right else 0
                width_offset = max_width - data_width if align_right else 0
                out[i].narrow(1, height_offset, data_height) \
                      .narrow(2, width_offset, data_width).copy_(data[i])
            return out, widths

    #@profile
    def __getitem__(self, index):
        assert index < self.numBatches, "%d > %d" %(index, self.numBatches)
        src = []
        tgt = []
        align = []
        feature = []
        #print 'batchBucket ', self.batchBucket
        s = self.batchBucket * index
        e = self.batchBucket * (index + 1)
        #print "slice"
        #print type(s), " : ", s
        #print type(e), " : ", e
        #print "---slice"
        redis_indices = self.indices[s:e]
        datas = self.r.mget(redis_indices)
        for d in datas:
            data = msgpack.unpackb(d)
            src += data[0]
            if len(data) >= 2:
                tgt_datas = data[1]
                if self.r2l:
                    reverse_tgt_datas = []
                    for tgt_data in tgt_datas:
                        main_tgt_data = tgt_data[1:-1]
                        main_tgt_data.reverse()
                        reverse_tgt_datas += [[tgt_data[0]] + main_tgt_data + [tgt_data[-1]]]
                    tgt_datas = reverse_tgt_datas
                tgt += tgt_datas
                if len(data) == 3:
                    align += data[2]
                    if len(data) == 4:
                        feature += data[3]
        batch_size = len(src)
        if len(tgt) >= 1:
            assert len(src) == len(tgt), "size of data from redis does not matches"

        features = None
        if len(feature) > 0:
            features = [feature]
        if self.reverse:
            #print "training in reverse direction"
            src, tgt = tgt, src
            align = [np.array(alignment).transpose().tolist() for alignment in align]

        srcBatch, lengths = self._batchify(src, align_right=False, include_lengths=True, dtype="text", features=features)
        if srcBatch.dim() == 2:
            srcBatch = srcBatch.unsqueeze(2)
        if tgt is not None:
            tgtBatch = self._batchify(tgt, dtype="text")
        else:
            tgtBatch = None
        alignment = None
        if len(align) > 0:
            src_len = srcBatch.size(1)
            tgt_len = tgtBatch.size(1)
            #assert src_len == len(align)
            #assert tgt_len == len(align[0])
            alignment = torch.ByteTensor(tgt_len, batch_size, src_len).fill_(0)
        #    #alignment = torch.ByteTensor(tgt_len, batch_size, src_len).fill_(0).pin_memory()
            for i in range(len(align)):
                region = torch.ByteTensor(align[i])
                alignment[1:region.size(1)+1, i, :region.size(0)] = region.t()
            alignment = alignment.float()
            if self.cuda:
                alignment = alignment.cuda()
        # tgt_len x batch x src_len
        lengths = torch.LongTensor(lengths)
        #lengths = torch.LongTensor(lengths).pin_memory()
        indices = range(len(srcBatch))
        # within batch sorting by decreasing length for variable length rnns
        lengths, perm = torch.sort(lengths, 0, descending=True)
        indices = [indices[p] for p in perm]
        srcBatch = [srcBatch[p] for p in perm]
        if tgtBatch is not None:
            tgtBatch = [tgtBatch[p] for p in perm]
        #if alignment is not None:
        #    alignment = alignment.transpose(0, 1)[
        #        perm.type_as(alignment).long()]
        #    alignment = alignment.transpose(0, 1).contiguous()

        def wrap(b, dtype="text"):
            if b is None:
                return b
            b = torch.stack(b, 0)
            if dtype == "text":
                b = b.transpose(0, 1).contiguous()
            if self.cuda:
                b = b.cuda()
            b = Variable(b, volatile=self.volatile)
            return b
        # Wrap lengths in a Variable to properly split it in DataParallel
        lengths = lengths.view(1, -1)
        lengths = Variable(lengths, volatile=self.volatile)

        return Batch(wrap(srcBatch, "text"),
                     wrap(tgtBatch, "text"),
                     lengths,
                     indices,
                     batch_size,
                     alignment=alignment)
    def __len__(self):
        return self.numBatches
if __name__ == "__main__":
    torch.cuda.set_device(2)
    test = RedisDataset("train", 128, True, volatile=False, db=1)
    import time
    BATCH = len(test)
    start = time.time()
    for idx in range(50):
        batch = test[idx]
    end = time.time() - start
    #print 'elapsed time : ',  end


# -*- coding: utf-8 -*-
# @Time    : 2021/9/30 14:32
# @Author  : He Ruizhi
# @File    : go_engine.py
# @Software: PyCharm

from GymGo.gym_go import govars, gogame
from typing import Union, List, Tuple
import numpy as np
from scipy import ndimage

surround_struct = np.array([[0, 1, 0],
                            [1, 0, 1],
                            [0, 1, 0]])

eye_struct = np.array([[1, 1, 1],
                       [1, 0, 1],
                       [1, 1, 1]])

corner_struct = np.array([[1, 0, 1],
                          [0, 0, 0],
                          [1, 0, 1]])
BLACK = govars.BLACK
WHITE = govars.WHITE

class GoEngine:
    def __init__(self, board_size: int = 9,
                 komi=7.5,
                 record_step: int = 4,
                 state_format: str = "separated",
                 record_last: bool = True):
        """
        围棋引擎初始化

        :param board_size: 棋盘大小，默认为9
        :param komi: 黑棋贴目数，默认黑贴7.5目（3又3/4子）
        :param record_step: 记录棋盘历史状态步数，默认为4
        :param state_format: 记录棋盘历史状态格式
                            【separated：黑白棋子分别记录在不同的矩阵中，[黑棋，白棋，下一步落子方，上一步落子位置(可选)]】
                            【merged：黑白棋子记录在同一个矩阵中，[棋盘棋子分布(黑1白-1)，下一步落子方，上一步落子位置(可选)]】
        :param record_last: 是否记录上一步落子位置
        """
        assert state_format in ["separated", "merged"],\
            "state_format can only be 'separated' or 'merged', but received: {}".format(state_format)

        self.board_size = board_size
        self.komi = komi
        self.record_step = record_step
        self.state_format = state_format
        self.record_last = record_last
        self.current_state = gogame.init_state(board_size)
        # 保存棋盘状态，用于悔棋
        self.board_state_history = []
        # 保存历史动作，用于悔棋
        self.action_history = []

        if state_format == "separated":
            record_step *= 2
        self.state_channels = record_step + 2 if record_last else record_step + 1
        self.board_state = np.zeros((self.state_channels, board_size, board_size))
        self.done = False

    def reset(self) -> np.ndarray:
        """重置current_state, board_state, board_state_history, action_history"""
        self.current_state = gogame.init_state(self.board_size)
        self.board_state = np.zeros((self.state_channels, self.board_size, self.board_size))
        self.board_state_history = []
        self.action_history = []
        self.done = False
        return np.copy(self.current_state)

    def step(self, action: Union[List[int], Tuple[int], int, None]) -> np.ndarray:
        """
        围棋落子

        :param action: 下一步落子位置
        :return:
        """
        assert not self.done
        if isinstance(action, tuple) or isinstance(action, list) or isinstance(action, np.ndarray):
            assert 0 <= action[0] < self.board_size
            assert 0 <= action[1] < self.board_size
            action = self.board_size * action[0] + action[1]
        elif isinstance(action, int):
            assert 0 <= action <= self.board_size ** 2
        elif action is None:
            action = self.board_size ** 2

        self.current_state = gogame.next_state(self.current_state, action, canonical=False)
        # 更新self.board_state
        self.board_state = self._update_state_step(action)
        # 存储历史状态
        self.board_state_history.append(np.copy(self.current_state))
        # 存储历史动作
        self.action_history.append(action)
        self.done = gogame.game_ended(self.current_state)
        return np.copy(self.current_state)

    def _update_state_step(self, action: int) -> np.ndarray:
        """
        更新self.board_state，须在更新完self.current_state之后更新self.board_state

        :param action: 下一步落子位置，1d-action
        :return:
        """
        if self.state_format == "separated":
            # 根据上一步落子方更新self.board_state(因为self.current_state已经更新完毕)
            if self.turn() == govars.WHITE:
                # 根据更新过后的self.current_state，下一步落子方为白方，则上一步落子方为黑方
                self.board_state[:self.record_step - 1] = np.copy(self.board_state[1:self.record_step])
                self.board_state[self.record_step - 1] = np.copy(self.current_state[govars.BLACK])
            else:
                # 根据更新过后的self.current_state，下一步落子方为黑方，则上一步落子方为白方
                self.board_state[self.record_step: self.record_step * 2 - 1] = \
                    np.copy(self.board_state[self.record_step + 1: self.record_step * 2])
                self.board_state[self.record_step * 2 - 1] = np.copy(self.current_state[govars.WHITE])
        elif self.state_format == "merged":
            self.board_state[:self.record_step - 1] = np.copy(self.board_state[1:self.record_step])
            current_state = self.current_state[[govars.BLACK, govars.WHITE]]
            current_state[govars.WHITE] *= -1
            self.board_state[self.record_step - 1] = np.sum(current_state, axis=0)

        if self.record_last:
            # 更新下一步落子方
            self.board_state[-2] = np.copy(self.current_state[govars.TURN_CHNL])
            # 更新上一步落子位置
            self.board_state[-1] = np.zeros((self.board_size, self.board_size))
            # 上一步不为pass
            if action != self.board_size ** 2:
                # 将action转换成position
                position = action // self.board_size, action % self.board_size
                self.board_state[-1, position[0], position[1]] = 1
        else:
            # 更新下一步落子方
            self.board_state[-1] = np.copy(self.current_state[govars.TURN_CHNL])
        return np.copy(self.board_state)

    def get_board_state(self) -> np.ndarray:
        """用于训练神经网络的棋盘状态矩阵"""
        return np.copy(self.board_state)

    def game_ended(self) -> bool:
        """游戏是否结束"""
        return self.done

    def winner(self) -> int:
        """获胜方，游戏未结束返回-1"""
        if not self.done:
            return -1
        else:
            winner = self.winning()
            winner = govars.BLACK if winner == 1 else govars.WHITE
            return winner

    def action_valid(self, action) -> bool:
        """判断action是否合法"""
        return self.valid_moves()[action]

    def valid_move_idcs(self) -> np.ndarray:
        """下一步落子有效位置的id"""
        valid_moves = self.valid_moves()
        return np.argwhere(valid_moves).flatten()

    def advanced_valid_move_idcs(self) -> np.ndarray:
        """下一步落子的非真眼有效位置的id"""
        advanced_valid_moves = self.advanced_valid_moves()
        return np.argwhere(advanced_valid_moves).flatten()

    def uniform_random_action(self) -> np.ndarray:
        """随机选择落子位置"""
        valid_move_idcs = self.valid_move_idcs()
        return np.random.choice(valid_move_idcs)

    def advanced_uniform_random_action(self) -> np.ndarray:
        """不填真眼的随机位置"""
        advanced_valid_move_idcs = self.advanced_valid_move_idcs()
        return np.random.choice(advanced_valid_move_idcs)

    def turn(self) -> int:
        """下一步落子方"""
        return gogame.turn(self.current_state)

    def valid_moves(self) -> np.ndarray:
        """下一步落子的有效位置"""
        return gogame.valid_moves(self.current_state)

    def advanced_valid_moves(self):
        """下一步落子的非真眼有效位置"""
        valid_moves = 1 - self.current_state[govars.INVD_CHNL]
        eyes_mask = 1 - self.eyes()
        return np.append((valid_moves * eyes_mask).flatten(), 1)

    def winning(self):
        """
        当游戏结束之后，从黑方角度看待，上一步落子后，哪一方胜利
        黑胜：1 白胜：-1
        """
        return gogame.winning(self.current_state, self.komi)

    def areas(self):
        """black_area, white_area"""
        return gogame.areas(self.current_state)

    def eyes(self):
        """
        下一步落子方的真眼位置
        1.如果在角上或者边上，则需要对应8个最近位置均有下一步落子方的棋子；
        2.如果不在边上和角上，则需要对应4个最近边全有下一步落子方的棋子，且至少有三个角有下一步落子方的棋子；
        3.所判断的位置没有棋子
        """
        board_shape = self.current_state.shape[1:]

        side_mask = np.zeros(board_shape)
        side_mask[[0, -1], :] = 1
        side_mask[:, [0, -1]] = 1
        nonside_mask = 1 - side_mask

        # 下一步落子方
        next_player = self.turn()
        # next_player的棋子分布矩阵
        next_player_pieces = self.current_state[next_player]
        # 棋盘所有有棋子的分布矩阵，有棋子则相应位置为1
        all_pieces = np.sum(self.current_state[[govars.BLACK, govars.WHITE]], axis=0)
        # 棋盘上所有空交叉点的分布矩阵，空交叉点位置为1
        empties = 1 - all_pieces

        # 对于边角位置
        side_matrix = ndimage.convolve(next_player_pieces, eye_struct, mode='constant', cval=1) == 8
        side_matrix = side_matrix * side_mask
        # 对于非边角位置
        nonside_matrix = ndimage.convolve(next_player_pieces, surround_struct, mode='constant', cval=1) == 4
        nonside_matrix *= ndimage.convolve(next_player_pieces, corner_struct, mode='constant', cval=1) > 2
        nonside_matrix = nonside_matrix * nonside_mask

        return empties * (side_matrix + nonside_matrix)

    def all_symmetries(self) -> List[np.ndarray]:
        """board_state的8种等价表示"""
        return gogame.all_symmetries(np.copy(self.board_state))

    @staticmethod
    def array_symmetries(array: np.ndarray) -> List[np.ndarray]:
        """
        指定array的8种旋转表示

        :param array: A (C, BOARD_SIZE, BOARD_SIZE) numpy array, where C is any number
        :return:
        """
        return gogame.all_symmetries(array)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function
import cv2
import numpy as np
import mxnet as mx
import os
os.environ["MXNET_CUDNN_AUTOTUNE_DEFAULT"] = '0'
from sklearn.preprocessing import normalize
from .iface_config import IFaceConfig

class IFaceEmbedding(object):
    def __init__(self, image_size=(112,112)):
        self.config = IFaceConfig()
        self.image_size = image_size
        self.init_model(self.config.MODEL_R100)

    def init_model(self, model_str, layer="fc1"):
        if self.config.USE_CUDA:
            ctx = mx.gpu(0)
        else:
            ctx = mx.cpu(0)

        _vec = model_str.split(',')
        assert len(_vec) == 2
        prefix = _vec[0]
        epoch = int(_vec[1])
        sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
        all_layers = sym.get_internals()
        sym = all_layers[layer+'_output']
        model = mx.mod.Module(symbol=sym, context=ctx, label_names=None)
        model.bind(data_shapes=[('data', (1, 3, self.image_size[0],
                                          self.image_size[1]))])
        model.set_params(arg_params, aux_params)
        self.model = model

    def get_feature(self, face_region):
        face_region = cv2.resize(face_region, self.image_size)
        face_region = cv2.cvtColor(face_region, cv2.COLOR_BGR2RGB)
        face_region = np.transpose(face_region, (2,0,1))
        input_blob = np.expand_dims(face_region, axis=0)
        data = mx.nd.array(input_blob)
        db = mx.io.DataBatch(data=(data,))
        self.model.forward(db, is_train=False)
        embedding = self.model.get_outputs()[0].asnumpy()
        embedding = normalize(embedding).flatten()
        return embedding

    def test_image(self):
        imgs = [cv2.imread(_img_path) for _img_path in self.config.TEST_IMAGES]
        embeddings = [self.get_feature(img) for img in imgs]
        e1, e2 = embeddings
        dist = np.sum(np.square(e1-e2))
        return e1, e2, dist

def main():
    iface = IFaceEmbedding()
    imgs = [cv2.imread(_img_path) for _img_path in iface.config.TEST_IMAGES]
    embeddings = [iface.get_feature(img) for img in imgs]
    e1, e2 = embeddings
    dist = np.sum(np.square(e1-e2))
    print(dist)

if __name__ == '__main__':
    main()


import numpy as np
import torch
import numbers
import math

class IncConfusionMeter:
    """Maintains a confusion matrix for a given calssification problem.
    The ConfusionMeter constructs a confusion matrix for a multi-class
    classification problems. It does not support multi-label, multi-class problems:
    for such problems, please use MultiLabelConfusionMeter.
    Args:
        k (int): number of classes in the classification problem
        normalized (boolean): Determines whether or not the confusion matrix
            is normalized or not
    """
    def __init__(self, k, increments, normalized=False):
        self.conf = np.ndarray((k, k), dtype=np.int32)
        self.normalized = normalized
        self.increments = increments
        self.cum_increments = [0] + [sum(increments[:i + 1]) for i in range(len(increments))]
        self.k = k
        self.reset()

    def reset(self):
        self.conf.fill(0)

    def add(self, predicted, target):
        """Computes the confusion matrix of K x K size where K is no of classes
        Args:
            predicted (tensor): Can be an N x K tensor of predicted scores obtained from
                the model for N examples and K classes or an N-tensor of
                integer values between 0 and K-1.
            target (tensor): Can be a N-tensor of integer values assumed to be integer
                values between 0 and K-1 or N x K tensor, where targets are
                assumed to be provided as one-hot vectors
        """
        if isinstance(predicted, torch.Tensor):
            predicted = predicted.cpu().numpy()
        if isinstance(target, torch.Tensor):
            target = target.cpu().numpy()

        assert predicted.shape[0] == target.shape[0], \
            'number of targets and predicted outputs do not match'

        if np.ndim(predicted) != 1:
            assert predicted.shape[1] == self.k, \
                'number of predictions does not match size of confusion matrix'
            predicted = np.argmax(predicted, 1)
        else:
            assert (predicted.max() < self.k) and (predicted.min() >= 0), \
                'predicted values are not between 1 and k'

        onehot_target = np.ndim(target) != 1
        if onehot_target:
            assert target.shape[1] == self.k, \
                'Onehot target does not match size of confusion matrix'
            assert (target >= 0).all() and (target <= 1).all(), \
                'in one-hot encoding, target values should be 0 or 1'
            assert (target.sum(1) == 1).all(), \
                'multi-label setting is not supported'
            target = np.argmax(target, 1)
        else:
            assert (predicted.max() < self.k) and (predicted.min() >= 0), \
                'predicted values are not between 0 and k-1'

        # hack for bincounting 2 arrays together
        x = predicted + self.k * target
        bincount_2d = np.bincount(x.astype(np.int32), minlength=self.k**2)
        assert bincount_2d.size == self.k**2
        conf = bincount_2d.reshape((self.k, self.k))

        self.conf += conf

    def value(self):
        """
        Returns:
            Confustion matrix of K rows and K columns, where rows corresponds
            to ground-truth targets and columns corresponds to predicted
            targets.
        """
        conf = self.conf.astype(np.float32)
        new_conf = np.zeros([len(self.increments), len(self.increments) + 2])
        for i in range(len(self.increments)):
            idxs = range(self.cum_increments[i], self.cum_increments[i + 1])
            new_conf[i, 0] = conf[idxs, idxs].sum()
            new_conf[i, 1] = conf[self.cum_increments[i]:self.cum_increments[i + 1],
                                  self.cum_increments[i]:self.cum_increments[i + 1]].sum() - new_conf[i, 0]
            for j in range(len(self.increments)):
                new_conf[i, j + 2] = conf[self.cum_increments[i]:self.cum_increments[i + 1],
                                          self.cum_increments[j]:self.cum_increments[j + 1]].sum()
        conf = new_conf
        if self.normalized:
            return conf / conf[:, 2:].sum(1).clip(min=1e-12)[:, None]
        else:
            return conf

class ClassErrorMeter:
    def __init__(self, topk=[1], accuracy=False):
        super(ClassErrorMeter, self).__init__()
        self.topk = np.sort(topk)
        self.accuracy = accuracy
        self.reset()

    def reset(self):
        self.sum = {v: 0 for v in self.topk}
        self.n = 0

    def add(self, output, target):
        if isinstance(output, np.ndarray):
            output = torch.Tensor(output)
        if isinstance(target, np.ndarray):
            target = torch.Tensor(target)
        # if torch.is_tensor(output):
        #     output = output.cpu().squeeze().numpy()
        # if torch.is_tensor(target):
        #     target = target.cpu().squeeze().numpy()
        # elif isinstance(target, numbers.Number):
        #     target = np.asarray([target])
        # if np.ndim(output) == 1:
        #     output = output[np.newaxis]
        # else:
        #     assert np.ndim(output) == 2, \
        #         'wrong output size (1D or 2D expected)'
        #     assert np.ndim(target) == 1, \
        #         'target and output do not match'
        # assert target.shape[0] == output.shape[0], \
        #     'target and output do not match'
        topk = self.topk
        maxk = int(topk[-1])  # seems like Python3 wants int and not np.int64
        no = output.shape[0]

        pred = output.topk(maxk, 1, True, True)[1]
        correct = pred == target.unsqueeze(1).repeat(1, pred.shape[1])
        # pred = torch.from_numpy(output).topk(maxk, 1, True, True)[1].numpy()
        # correct = pred == target[:, np.newaxis].repeat(pred.shape[1], 1)

        for k in topk:
            self.sum[k] += no - correct[:, 0:k].sum()
        self.n += no

    def value(self, k=-1):
        if k != -1:
            assert k in self.sum.keys(), \
                'invalid k (this k was not provided at construction time)'
            if self.n == 0:
                return float('nan')
            if self.accuracy:
                return (1. - float(self.sum[k]) / self.n) * 100.0
            else:
                return float(self.sum[k]) / self.n * 100.0
        else:
            return [self.value(k_) for k_ in self.topk]

class AverageValueMeter:
    def __init__(self):
        super(AverageValueMeter, self).__init__()
        self.reset()
        self.val = 0

    def add(self, value, n=1):
        self.val = value
        self.sum += value
        self.var += value * value
        self.n += n

        if self.n == 0:
            self.mean, self.std = np.nan, np.nan
        elif self.n == 1:
            self.mean, self.std = self.sum, np.inf
            self.mean_old = self.mean
            self.m_s = 0.0
        else:
            self.mean = self.mean_old + (value - n * self.mean_old) / float(self.n)
            self.m_s += (value - self.mean_old) * (value - self.mean)
            self.mean_old = self.mean
            self.std = math.sqrt(self.m_s / (self.n - 1.0))

    def value(self):
        return self.mean, self.std

    def reset(self):
        self.n = 0
        self.sum = 0.0
        self.var = 0.0
        self.val = 0.0
        self.mean = np.nan
        self.mean_old = 0.0
        self.m_s = 0.0
        self.std = np.nan


import datetime
import numpy as np
import os
import os.path as osp
import glob
import cv2
import insightface

assert insightface.__version__>='0.4'

def detect_person(img, detector):
    bboxes, kpss = detector.detect(img)
    bboxes = np.round(bboxes[:,:4]).astype(np.int)
    kpss = np.round(kpss).astype(np.int)
    kpss[:,:,0] = np.clip(kpss[:,:,0], 0, img.shape[1])
    kpss[:,:,1] = np.clip(kpss[:,:,1], 0, img.shape[0])
    vbboxes = bboxes.copy()
    vbboxes[:,0] = kpss[:, 0, 0]
    vbboxes[:,1] = kpss[:, 0, 1]
    vbboxes[:,2] = kpss[:, 4, 0]
    vbboxes[:,3] = kpss[:, 4, 1]
    return bboxes, vbboxes

if __name__ == '__main__':
    import glob
    detector = insightface.model_zoo.get_model('scrfd_person_2.5g.onnx', download=True)
    detector.prepare(0, nms_thresh=0.5, input_size=(640, 640))
    img_paths = glob.glob('data/images/*.jpg')
    for img_path in img_paths:
        img = cv2.imread(img_path)
        bboxes, vbboxes = detect_person(img, detector)
        for i in range(bboxes.shape[0]):
            bbox = bboxes[i]
            vbbox = vbboxes[i]
            x1,y1,x2,y2 = bbox
            vx1,vy1,vx2,vy2 = vbbox
            cv2.rectangle(img, (x1,y1)  , (x2,y2) , (0,255,0) , 1)
            alpha = 0.8
            color = (255, 0, 0)
            for c in range(3):
                img[vy1:vy2,vx1:vx2,c] = img[vy1:vy2, vx1:vx2, c]*alpha + color[c]*(1.0-alpha)
            cv2.circle(img, (vx1,vy1) , 1, color , 2)
            cv2.circle(img, (vx1,vy2) , 1, color , 2)
            cv2.circle(img, (vx2,vy1) , 1, color , 2)
            cv2.circle(img, (vx2,vy2) , 1, color , 2)
        filename = img_path.split('/')[-1]
        cv2.imwrite('./outputs/%s'%filename, img)


import numpy

from .blend import *
from .blend_detail import *

blend = {1: 0,
         # 2 : 0,
         6: 0,
         7: 0,
         8: 0,
         9: 0,
         22: 0,
         #23 : 0,
         25: 0,
         30: 0,
         38: 0,
         37: 0,
         38: 0,
         39: 0,
         40: 0,
         41: 0,
         43: 0,
         44: 0,
         45: 0,
         46: 0,
         47: 0,

         53: 0,
         56: 0,
         58: 57,
         #59 : 123, no idea what mud blends into
         #60 : 59,
         63: 0,
         64: 0,
         65: 0,
         66: 0,
         67: 0,
         68: 0,
         #70 : 59,
         #75:59,
         76: 57,
         107: 0,
         108: 0,
         111: 0,
         112: 0,
         116: 0,
         117: 0,
         118: 0,
         119: 0,
         120: 0,
         121: 0,
         122: 0,
         123: 0,
         130: 0,
         131: 0,
         140: 0,

         }
blendto = {}
for x in range(65537):
    blendto[x] = [x]
for key in blend:
    blendto[blend[key]].append(key)
grass = {2: 0,
         23: 0,
         60: 59,
         70: 59,
         109: 0,

         }
grassto = {}
for x in range(255):
    grassto[x] = [x]
for key in grass:
    grassto[grass[key]].append(key)
    blendto[grass[key]].append(key)
# blendto[0].append(1)#dirt blends to stone
#fail = []
def calc(header, tiles, rmap):
    blendmap = numpy.zeros((header["width"], header["height"], 2), dtype=numpy.uint16)
    wblendmap = numpy.zeros((header["width"], header["height"], 2), dtype=numpy.uint16)
    for x in range(1, header["width"] - 1):
        for y in range(1, header["height"] - 1):
            wid = tiles[x][y][1]
            wblendmap[x][y] = wmix[(wid == tiles[x - 1][y][1], wid == tiles[x + 1][y][1],
                                    wid == tiles[x][y - 1][1], wid == tiles[x][y + 1][1])][rmap[x, y]]

            tid = tiles[x][y][0]
            blendmap[x][y] = mix[(tid == tiles[x - 1][y][0], tid == tiles[x + 1][y][0],
                                  tid == tiles[x][y - 1][0], tid == tiles[x][y + 1][0])][rmap[x, y]]
    return blendmap, wblendmap

def get(tiles, blendmap, wblendmap, rmap, x, y):
    wid = tiles[x][y][1]
    wblendmap[x][y] = wmix[(wid == tiles[x - 1][y][1], wid == tiles[x + 1][y][1],
                            wid == tiles[x][y - 1][1], wid == tiles[x][y + 1][1])][rmap[x, y]]

    tid = tiles[x][y][0]
    blendmap[x][y] = mix[(tid == tiles[x - 1][y][0], tid == tiles[x + 1][y][0],
                          tid == tiles[x][y - 1][0], tid == tiles[x][y + 1][0])][rmap[x, y]]

def get_tile(tiles, blendmap, rmap, x, y):
    tid = tiles[x][y][0]
    blendmap[x][y] = mix[(tid == tiles[x - 1][y][0], tid == tiles[x + 1][y][0],
                          tid == tiles[x][y - 1][0], tid == tiles[x][y + 1][0])][rmap[x, y]]

def get_wall(tiles, wblendmap, rmap, x, y):
    wid = tiles[x][y][1]
    wblendmap[x][y] = wmix[(wid == tiles[x - 1][y][1], wid == tiles[x + 1][y][1],
                            wid == tiles[x][y - 1][1], wid == tiles[x][y + 1][1])][rmap[x, y]]

def get_tile_detail(tiles, blendmap, rmap, x, y):
    tid = tiles[x][y][0]
    if tid in grass:
        bid = grass[tid]
        if tid == tiles[x - 1][y][0]:
            o = 1
        elif bid == tiles[x - 1][y][0]:
            o = 2
        else:
            o = 0
        if tid == tiles[x + 1][y][0]:
            p = 1
        elif bid == tiles[x + 1][y][0]:
            p = 2
        else:
            p = 0
        if tid == tiles[x][y - 1][0]:
            z = 1
        elif bid == tiles[x][y - 1][0]:
            z = 2
        else:
            z = 0
        if tid == tiles[x][y + 1][0]:
            w = 1
        elif bid == tiles[x][y + 1][0]:
            w = 2
        else:
            w = 0
        blendmap[x][y] = grass_detail[(o, p, z, w)][rmap[x, y]]
    elif tid in blend:
        bid = blend[tid]

        if tid == tiles[x - 1][y][0]:
            o = 1
        elif bid == tiles[x - 1][y][0]:
            o = 2
        else:
            o = 0
        if tid == tiles[x + 1][y][0]:
            p = 1
        elif bid == tiles[x + 1][y][0]:
            p = 2
        else:
            p = 0
        if tid == tiles[x][y - 1][0]:
            z = 1
        elif bid == tiles[x][y - 1][0]:
            z = 2
        else:
            z = 0
        if tid == tiles[x][y + 1][0]:
            w = 1
        elif bid == tiles[x][y + 1][0]:
            w = 2
        else:
            w = 0
        blendmap[x][y] = mix_detail[(o, p, z, w)][rmap[x, y]]

    else:
        accept = blendto[tid]
        blendmap[x][y] = mix[(tiles[x - 1][y][0] in accept, tiles[x + 1][y][0] in accept,
                              tiles[x][y - 1][0] in accept, tiles[x][y + 1][0] in accept)][rmap[x, y]]


import sys
import binascii
import zlib
import cv2
import numpy as np

ENTER_KEY_WIN = 13
ENTER_KEY_LINUX = 10
ESC_KEY = 27

def main():
    global ENTER_KEY_WIN, ENTER_KEY_LINUX, ESC_KEY
    if len(sys.argv) <= 1:
        print("usage: pic2bin.py <imagefile>", file=sys.stderr)
        return 1

    windowname = sys.argv[1] + "    Enter : save    Esc : quit"

    img = cv2.imread(sys.argv[1], cv2.IMREAD_COLOR)
    if img is None:
        print("Image open error", file=sys.stderr)
        return 1
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    height, width = gray.shape
    if height != 64:
        print("Image height is invalid", file=sys.stderr)
        return 1

    cv2.namedWindow(windowname, cv2.WINDOW_AUTOSIZE | cv2.WINDOW_GUI_NORMAL)

    lut = np.arange(256, dtype="float32")
    lut = pow((lut / 255), 2.2)

    binalized = np.zeros((height, width, 1), np.uint8)

    cumulative_err = 0
    calc_done = False
    while True:
        if not calc_done:
            for y in range(height):
                for x in range(width):
                    linear_val = lut[gray[y, x]]
                    if linear_val + cumulative_err < 0.5:
                        cumulative_err += linear_val
                        binalized[y, x] = 0
                    else:
                        cumulative_err -= 1.0 - linear_val
                        binalized[y, x] = 255
            calc_done = True

        frame = cv2.cvtColor(binalized, cv2.COLOR_GRAY2BGR)
        frame = cv2.vconcat([img, frame])
        cv2.imshow(windowname, frame)

        key = cv2.waitKey(100)
        if (key == ENTER_KEY_WIN) | (key == ENTER_KEY_LINUX):
            cv2.destroyAllWindows()
            break
        elif key == ESC_KEY:
            return 0

    img_out = cv2.rotate(binalized, cv2.ROTATE_90_CLOCKWISE)
    height, width = img_out.shape

    encoded = b""
    for y in range(height):
        aggregated = 0
        for shift, x in enumerate(range(width - 1, -1, -1)):
            if img_out[y, x] <= 127:
                aggregated += 1 << shift
        line = aggregated.to_bytes(8, "big")
        encoded += line
        print(binascii.hexlify(line))

    with open("hello.bin", "wb") as f:
        f.write(zlib.compress(encoded))
    return 0

sys.exit(main())


#! /usr/env/bin python

"""
This module contains code for creating a dark current reference file.

Input:
    List of exposures at a calibration state from which you wish to create
    dark current reference files. For NIR detectors, this should be after:

    0) group_scale
    1) dq_init
    2) saturation
    3) ipc* - not currently run
    4) superbias
    5) refpix
    6) linearity
    7) persistence

    For MIRI, this should be after:

    0) group_scale
    1) dq_init
    2) saturation
    3) linearity
    4) rscd

    Also, versions of these files that have been run up through the JUMP
    step must also be present. This is because we need to know where the
    CR hits are in the data.

    ********************POSSIBLE INPUT CHANGE****************************
    (An alternate plan: Run the pipeline once on the input files. Run up
    through the JUMP step. For NIR data, skip the dark subtraction step.
    For MIRI, skip refpix and dark subtraction steps. Then the CR flags
    will be present, and the data will be in the proper state to create
    dark current reference files.)
    ********************POSSIBLE INPUT CHANGE****************************

    All files should be from the same instrument/detector/aperture/readpattern

Ouputs:
    A dark current reference file in the proper format
    A series of images showing for each pixel in each group, how many
        measurements of the signal were used to calculate the mean
        dark value.

Function:
Get file headers from each file. Check for consistency.
Extract CR flags from the JUMP version of the files.
    From these, create arrays giving the group number of the first cosmic
    ray hit in each pixel for each integration
Create images of how many dark current measurements are considered good
    for each pixel and each group. All groups after the first CR hit in a
    pixel are considered bad and not used.
Read in the data. To avoid memory problems only read in N groups at a time,
    where N is determined from user inputs.
Flag all groups that occur between the first CR hit and the end of the integration
    as bad.
From the remaining good data, calculate the sigma-clipped mean dark signal
    in each pixel for each group (i.e. take the mean across integrations).
    Also calculate the sigma-clipped stdard deviation.
Save the sigma-clipped mean as the dark data, and the sigma-clipped stdev
    as the uncertainty in a DarkModel or MIRIDarkModel instance. The
    DQ array is currently all zeros, anticipating bad pixels being placed
    in the bad pixel reference file.
"""

from astropy.io import fits
from astropy.stats import sigma_clip
import copy
from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.cm as cm
import matplotlib.pyplot as plt
import numpy as np
import os

from jwst.datamodels import DarkModel, DarkMIRIModel

from jwst_reffiles.utils import dq_flags
from jwst_reffiles.utils import crds_checks
from jwst_reffiles.utils.file_utils import read_ramp, get_file_header

inst_abbreviations = {'nircam': 'NRC', 'niriss': 'NIS', 'fgs': 'FGS', 'miri': 'MIR', 'nirspec': 'NRS'}

class Dark():
    def __init__(self, file_list=[], max_equivalent_groups=60, sigma_threshold=3, pedigree=None,
                 descrip=None, author=None, use_after=None, history=None, output_file=None,
                 contribution_file='good_signals_per_group.pdf'):
        self.file_list = file_list
        self.sigma_threshold = sigma_threshold
        self.pedigree = pedigree
        self.descrip = descrip
        self.author = author
        self.use_after = use_after
        self.history = history
        self.output_file = output_file
        self.contribution_file = contribution_file

        # If an output directory is specified, save that here, so that
        # ancillary outputs can be saved to the same location
        if self.output_file is not None:
            self.output_dir = os.path.dirname(self.output_file)
        else:
            self.output_dir = './'

        # If no output directory is given for the contribution file, then
        # save it to the same location as the reference file
        if os.path.dirname(self.contribution_file) == '':
            self.contribution_file = os.path.join(self.output_dir, self.contribution_file)

        # Check that all necessary entries are in the dictionary to be used
        # to populate the header of the final reference file
        #self.check_reffile_info()

        # Get metadata from first input dark.
        self.metadata = get_file_header(self.file_list[0], extension='PRIMARY')

        # Use array size and number of
        # groups as a guide to determine if and how to split up the data
        # into chunks, such that one chunk is read in from all input files
        # and calculations performed without using too much memory
        if self.metadata['INSTRUME'] == 'MIRI':
            x_fullframe_dim = 1024
            y_fullframe_dim = 1024
        else:
            x_fullframe_dim = 2048
            y_fullframe_dim = 2048
        pix_per_fullframe_group = x_fullframe_dim * y_fullframe_dim

        pix_per_group = self.metadata['SUBSIZE1'] * self.metadata['SUBSIZE2']
        pix_ratio = np.int(pix_per_fullframe_group / pix_per_group)

        # The number of groups that can be read in from each file is the
        # user-entered max number of groups, divided by the ratio of the
        # number of pixels in a full frame exposure to the number of pixels
        # in the input aperture, divided by the number of files, and
        # finally divided by 2, since we will also need to read in the
        # cosmic ray flags associated with each file
        delta_groups = np.int(max_equivalent_groups / pix_ratio / len(self.file_list)) # / 2.)
        if delta_groups == 0:
            raise ValueError(('ERROR: max_equivalent_groups incompatible with the number/size '
                              'of input files. Even reading just a single group at a time from '
                              'all input files would result in more than max_equivalent_groups '
                              'groups in memory at one time. Increase max_equivalent_groups or '
                              'decrease the number of input exposures.'))
        self.group_index = np.arange(0, self.metadata['NGROUPS'], delta_groups)
        self.group_index = np.append(self.group_index, self.metadata['NGROUPS'])

        # Create the reference file
        self.create_reffile()

    def consistency_check(self):
        """Make sure that all input files have the same dimensions,
        number of groups, etc
        """
        keywords_to_check = ['INSTRUME', 'DETECTOR', 'SUBARRAY', 'READPATT', 'SUBSIZE1',
                             'SUBSIZE2', 'NGROUPS']
        for filename in self.file_list[1:]:
            metadata = get_file_header(filename, extension='PRIMARY')
            for kw in keywords_to_check:
                try:
                    if metadata[kw] != self.metadata[kw]:
                        raise ValueError(("ERROR: inconsistent input files. The first input "
                                          "file contains: \n{}: {}, while {} contains: \n{}: {}"
                                          .format(kw, self.metadata[kw], filename, kw, metadata[kw])))
                except KeyError:
                    raise KeyError("Header keyword {} not found in file {}".format(kw, filename))

    def create_num_contributions_plots(self, clean_signals, bad_signals):
        """Assume PDF
        """
        # Make sure the output file for the plots is a PDF
        if self.contribution_file[-4:].lower() != '.pdf':
            self.contribution_file = '{}.pdf'.format(self.contribution_file)

        ngroups, ysize, xsize = clean_signals.shape
        image_max = np.max(clean_signals)

        # Save all plots to a single PDF
        pdf = PdfPages(self.contribution_file)
        for group in range(ngroups):
            fig = plt.figure(figsize=(9, 9))
            ax1 = fig.add_subplot(1, 1, 1)
            im = ax1.imshow(clean_signals[group, :, :], extent=[0, xsize, 0, ysize], interpolation='None',
                            cmap=cm.RdYlGn, origin='lower', vmin=0, vmax=image_max)
            plt.colorbar(im)

            ax1.set_title('Signal Values Used to Create Mean Dark, Group {}'.format(group))
            fig.tight_layout()
            pdf.savefig(fig)
            plt.close(fig)
        pdf.close()
        print(('Images of the number of signals contributing to the mean dark measurement saved to {}'
               .format(self.contribution_file)))

    def create_reffile(self):
        """MAIN FUNCTION
        Creates a dark current reference file from a list of input dark
        current exposures
        """
        # Make sure reference file metadata have been provided
        if self.pedigree is None:
            raise ValueError("ERROR: 'pedigree' must be provided in order to save a valid reference file.")
        if self.descrip is None:
            raise ValueError("ERROR: 'descrip' must be provided in order to save a valid reference file.")
        if self.author is None:
            raise ValueError("ERROR: 'author' must be provided in order to save a valid reference file.")
        if self.use_after is None:
            raise ValueError("ERROR: 'use_after' must be provided in order to save a valid reference file.")
        if self.history is None:
            raise ValueError(("WARNING: No 'history' provided for the reference file. The list of files "
                              "used to create the reference file will be added as a HISTORY entry, but "
                              "nothing else."))

        # Very basic checks on reference file metadata
        crds_checks.validate_pedigree(self.pedigree)
        crds_checks.validate_useafter(self.use_after)

        # Data consistency check
        self.consistency_check()

        # Read in the DQ arrays from the Jump files and create a map of
        # group numbers associated with the first CR hit in the integration
        cr_indexes = self.prepare_cr_maps()

        # Collect info on the number of remaining good pixels
        for filenum, filename in enumerate(self.file_list):
            # Number of good groups per integration
            num_good = copy.deepcopy(cr_indexes[filename])

            # Set pixels with no hits as good in all groups
            num_good[np.isnan(num_good)] = self.metadata['NGROUPS']

            # Number of bad groups per integration
            num_bad = self.metadata['NGROUPS'] - num_good

            if filenum == 0:
                all_good = copy.deepcopy(num_good)
                all_bad = copy.deepcopy(num_bad)
            else:
                all_good = np.vstack([all_good, num_good])
                all_bad = np.vstack([all_bad, num_bad])

        # Translate into the number of good signal values
        # for each pixel in each group (i.e. look across
        # integrations)
        good_counter, bad_counter = dq_flags.signals_per_group(all_good, self.metadata['NGROUPS'])

        # Save num_good and num_bad as images
        counter_file = '{}fits'.format(self.contribution_file[0: -3])
        print(('Save the arrays listing the number of good signals for each pixel/group '
               'into {}'.format(counter_file)))
        hh0 = fits.PrimaryHDU(good_counter)
        #hh1 = fits.ImageHDU(bad_counter)
        hlist = fits.HDUList([hh0])  #, hh1])
        hlist.writeto(counter_file, overwrite=True)

        # Save some memory
        del num_good
        del num_bad
        del all_good
        del all_bad

        # What we really want is an image of the number of good signal
        # values for every pixel in every group
        self.create_num_contributions_plots(good_counter, bad_counter)

        # Get a map of reference pixels
        #refpix_map = dq_flags.create_refpix_map(self.file_list[0])

        # Create arrays to hold the final dark current data and uncertainty
        final_dark = np.zeros((self.metadata['NGROUPS'], self.metadata['SUBSIZE2'], self.metadata['SUBSIZE1']))
        final_dark_dev = np.zeros((self.metadata['NGROUPS'], self.metadata['SUBSIZE2'], self.metadata['SUBSIZE1']))
        # Read in the data
        for i in range(len(self.group_index) - 1):
            groups = np.arange(self.group_index[i], self.group_index[i + 1])
            for file_index, filename in enumerate(self.file_list):

                # Get metadata so we know how many integrations are in
                # the file, so that we can read them all in.
                metadata = get_file_header(self.file_list[0], extension='PRIMARY')
                for integration in range(metadata['NINTS']):

                    print("Reading in integration: {}, groups: {} - {} from file: {}".format(integration,
                                                                                self.group_index[i],
                                                                                self.group_index[i+1],
                                                                                filename))

                    # Read in the dark data
                    file_data = read_ramp(filename, integration, self.group_index[i],
                                          self.group_index[i + 1], extension='SCI')

                    jump_map = cr_indexes[filename][integration, :, :]

                    # Pixels/groups that come after a CR hit are set to NaN
                    group_arr = np.zeros_like(file_data)
                    for idx, g in enumerate(groups):
                        group_arr[idx, :, :] = g
                    diff = group_arr - jump_map
                    hits = np.where(diff >= 0)
                    file_data[hits[0], hits[1], hits[2]] = np.nan

                    # Place data in full container
                    if integration == 0:
                        data = np.expand_dims(file_data, axis=0)
                    else:
                        data = np.vstack([data, np.expand_dims(file_data, axis=0)])

                if file_index == 0:
                    all_data = copy.deepcopy(data)
                else:
                    all_data = np.vstack([all_data, data])

            # Free up some memory
            del data

            # Now calculate the sigma_clipped mean for each pixel through
            # all the files
            clipped_data = sigma_clip(all_data, sigma=self.sigma_threshold, axis=0, masked=False)
            mean_dark = np.nanmean(clipped_data, axis=0)
            dev_dark = np.nanstd(clipped_data, axis=0)

            # Pixels with no good signals will have values of NaN in mean_dark
            # Set these to zero.
            zero_values = np.isnan(mean_dark)
            mean_dark[zero_values] = 0.
            dev_dark[zero_values] = 0.

            # Place the mean and stdev of the dark signals into the correct
            # location within the final dark arrays that contain all groups
            final_dark[self.group_index[i]: self.group_index[i + 1], :, :] = mean_dark
            final_dark_dev[self.group_index[i]: self.group_index[i + 1], :, :] = dev_dark

        # DQ array associated with dark data
        # All bad pix are put in bad pixel mask, so dark reffile DQ array
        # is empty?
        yd, xd = final_dark.shape[-2:]
        dq_dark = np.zeros((yd, xd)).astype(np.int8)

        # Save reference file
        self.save_reffile(final_dark, final_dark_dev, dq_dark)

    def prepare_cr_maps(self):
        """Read in groupdq extension from jump files, extract only the
        JUMP_DET flags, and collapse into an array that lists the group
        number of the first CR hit for each pixel in each integration

        Returns
        -------
        hit_indexes : dict
            key values are the filenames associated with the data. Each
            value is a 3D array (integ, y, x) of group numbers of CR hits
        """
        hit_indexes = {}
        for filename in self.file_list:
            if not os.path.isfile(filename):
                raise FileNotFoundError("ERROR: Jump file {} not found.".format(filename))

            print('Getting CR map from file {}'.format(filename))

            # Read in the entire exposure from the GROUPDQ extension
            groupdq = read_ramp(filename, extension='GROUPDQ')

            # Keep only the JUMP_DET flags
            file_cr_map = dq_flags.flag_map(groupdq, 'JUMP_DET')

            # Translate CR maps per group into one array per integration.
            # The pixel value is the group number of the first CR hit
            first_hits = dq_flags.collapse_cr_map(file_cr_map)

            print('Temporarily save CR maps for development, into CR_maps.fits')
            hh0 = fits.PrimaryHDU(file_cr_map.astype(np.int))
            hh1 = fits.ImageHDU(first_hits)
            hlist = fits.HDUList([hh0, hh1])
            hlist.writeto('CR_maps.fits', overwrite=True)

            # Update the CR index map for this file
            hit_indexes[filename] = first_hits

        return hit_indexes

    def save_reffile(self, data, uncertainty, dq):
        """Save the given data in the official reference file format

        Parameters
        ----------
        data : numpy.ndarray
            3D array containing the mean dark current signals

        uncertainty : numpy.ndarray
            3D array containing associated uncertainties

        dq : numpy.ndarray
            2D array containing associated data quality flags
        """
        if self.metadata['INSTRUME'].upper() == 'MIRI':
            model = DarkMIRIModel()
        else:
            model = DarkModel()

        # Insert data
        model.data = data
        model.err = uncertainty
        model.dq = dq

        # Metadata copied from input files
        model.meta.reftype = 'DARK'
        model.meta.instrument.name = self.metadata['INSTRUME'].upper()
        model.meta.instrument.detector = self.metadata['DETECTOR'].upper()
        #try:
        #    model.meta.exposure.type = self.metadata['EXP_TYPE'].upper()
        #except KeyError:
        model.meta.exposure.type = '{}_DARK'.format(inst_abbreviations[self.metadata['INSTRUME'].lower()])
        model.meta.exposure.readpatt = self.metadata['READPATT'].upper()
        model.meta.exposure.nframes = self.metadata['NFRAMES']
        model.meta.exposure.ngroups = self.metadata['NGROUPS']
        model.meta.exposure.groupgap = self.metadata['GROUPGAP']
        model.meta.subarray.name = self.metadata['SUBARRAY']
        model.meta.subarray.xstart = self.metadata['SUBSTRT1']
        model.meta.subarray.ystart = self.metadata['SUBSTRT2']
        model.meta.subarray.xsize = self.metadata['SUBSIZE1']
        model.meta.subarray.ysize = self.metadata['SUBSIZE2']
        model.meta.subarray.fastaxis = self.metadata['FASTAXIS']
        model.meta.subarray.slowaxis = self.metadata['SLOWAXIS']

        # Instrument-specific metadata
        if self.metadata['INSTRUME'].upper() != 'MIRI':
            try:
                model.meta.exposure.gain_factor = self.metadata['GAINFACT']
            except KeyError:
                model.meta.exposure.gain_factor = 1.0
        else:
            model.meta.exposure.preadpatt = self.metadata['P_READPA']

        # User-provided metadata
        model.meta.pedigree = self.pedigree
        model.meta.descrip = self.descrip
        model.meta.author = self.author
        model.meta.useafter = self.use_after

        # Add HISTORY entries
        if self.history is not None:
            model.history.append(self.history)
        model.history.append(("{} {} {} dark current file, created from data in files: "
                              .format(self.metadata['INSTRUME'].upper(), self.metadata['DETECTOR'].upper(),
                                      self.metadata['SUBARRAY'])))
        for filename in self.file_list:
            model.history.append(filename)

        # Create output name if necessary
        if self.output_file is None:
            self.output_file = '{}_{}_{}_dark_ref_file.fits'.format(self.metadata['INSTRUME'],
                                                                    self.metadata['DETECTOR'],
                                                                    self.metadata['SUBARRAY'])
        model.save(self.output_file)
        print('Dark current reference file saved to {}'.format(self.output_file))