import torch
import torch.nn as nn

eps=1e-8

def sinkhorn(M,r,c,iteration):
    p = torch.softmax(M, dim=-1)
    u = torch.ones_like(r)
    v = torch.ones_like(c)
    for _ in range(iteration):
        u = r / ((p * v.unsqueeze(-2)).sum(-1) + eps)
        v = c / ((p * u.unsqueeze(-1)).sum(-2) + eps)
    p = p * u.unsqueeze(-1) * v.unsqueeze(-2)
    return p

def sink_algorithm(M,dustbin,iteration):
    M = torch.cat([M, dustbin.expand([M.shape[0], M.shape[1], 1])], dim=-1)
    M = torch.cat([M, dustbin.expand([M.shape[0], 1, M.shape[2]])], dim=-2)
    r = torch.ones([M.shape[0], M.shape[1] - 1],device='cuda')
    r = torch.cat([r, torch.ones([M.shape[0], 1],device='cuda') * M.shape[1]], dim=-1)
    c = torch.ones([M.shape[0], M.shape[2] - 1],device='cuda')
    c = torch.cat([c, torch.ones([M.shape[0], 1],device='cuda') * M.shape[2]], dim=-1)
    p=sinkhorn(M,r,c,iteration)
    return p

def seeding(nn_index1,nn_index2,x1,x2,topk,match_score,confbar,nms_radius,use_mc=True,test=False):

    #apply mutual check before nms
    if use_mc:
        mask_not_mutual=nn_index2.gather(dim=-1,index=nn_index1)!=torch.arange(nn_index1.shape[1],device='cuda')
        match_score[mask_not_mutual]=-1
    #NMS
    pos_dismat1=((x1.norm(p=2,dim=-1)**2).unsqueeze_(-1)+(x1.norm(p=2,dim=-1)**2).unsqueeze_(-2)-2*(x1@x1.transpose(1,2))).abs_().sqrt_()
    x2=x2.gather(index=nn_index1.unsqueeze(-1).expand(-1,-1,2),dim=1)
    pos_dismat2=((x2.norm(p=2,dim=-1)**2).unsqueeze_(-1)+(x2.norm(p=2,dim=-1)**2).unsqueeze_(-2)-2*(x2@x2.transpose(1,2))).abs_().sqrt_()
    radius1, radius2 = nms_radius * pos_dismat1.mean(dim=(1,2),keepdim=True), nms_radius * pos_dismat2.mean(dim=(1,2),keepdim=True)
    nms_mask = (pos_dismat1 >= radius1) & (pos_dismat2 >= radius2)
    mask_not_local_max=(match_score.unsqueeze(-1)>=match_score.unsqueeze(-2))|nms_mask
    mask_not_local_max=~(mask_not_local_max.min(dim=-1).values)
    match_score[mask_not_local_max] = -1

    #confidence bar
    match_score[match_score<confbar]=-1
    mask_survive=match_score>0
    if test:
        topk=min(mask_survive.sum(dim=1)[0]+2,topk)
    _,topindex = torch.topk(match_score,topk,dim=-1)#b*k
    seed_index1,seed_index2=topindex,nn_index1.gather(index=topindex,dim=-1)
    return seed_index1,seed_index2

class PointCN(nn.Module):
    def __init__(self, channels,out_channels):
        nn.Module.__init__(self)
        self.shot_cut = nn.Conv1d(channels, out_channels, kernel_size=1)
        self.conv = nn.Sequential(
            nn.InstanceNorm1d(channels, eps=1e-3),
            nn.SyncBatchNorm(channels),
            nn.ReLU(),
            nn.Conv1d(channels, channels, kernel_size=1),
            nn.InstanceNorm1d(channels, eps=1e-3),
            nn.SyncBatchNorm(channels),
            nn.ReLU(),
            nn.Conv1d(channels, out_channels, kernel_size=1)
        )

    def forward(self, x):
        return self.conv(x) + self.shot_cut(x)

class attention_propagantion(nn.Module):

    def __init__(self,channel,head):
        nn.Module.__init__(self)
        self.head=head
        self.head_dim=channel//head
        self.query_filter,self.key_filter,self.value_filter=nn.Conv1d(channel,channel,kernel_size=1),nn.Conv1d(channel,channel,kernel_size=1),\
                                                            nn.Conv1d(channel,channel,kernel_size=1)
        self.mh_filter=nn.Conv1d(channel,channel,kernel_size=1)
        self.cat_filter=nn.Sequential(nn.Conv1d(2*channel,2*channel, kernel_size=1), nn.SyncBatchNorm(2*channel), nn.ReLU(),
                                      nn.Conv1d(2*channel, channel, kernel_size=1))

    def forward(self,desc1,desc2,weight_v=None):
        #desc1(q) attend to desc2(k,v)
        batch_size=desc1.shape[0]
        query,key,value=self.query_filter(desc1).view(batch_size,self.head,self.head_dim,-1),self.key_filter(desc2).view(batch_size,self.head,self.head_dim,-1),\
                        self.value_filter(desc2).view(batch_size,self.head,self.head_dim,-1)
        if weight_v is not None:
            value=value*weight_v.view(batch_size,1,1,-1)
        score=torch.softmax(torch.einsum('bhdn,bhdm->bhnm',query,key)/ self.head_dim ** 0.5,dim=-1)
        add_value=torch.einsum('bhnm,bhdm->bhdn',score,value).reshape(batch_size,self.head_dim*self.head,-1)
        add_value=self.mh_filter(add_value)
        desc1_new=desc1+self.cat_filter(torch.cat([desc1,add_value],dim=1))
        return desc1_new

class hybrid_block(nn.Module):
    def __init__(self,channel,head):
        nn.Module.__init__(self)
        self.head=head
        self.channel=channel
        self.attention_block_down = attention_propagantion(channel, head)
        self.cluster_filter=nn.Sequential(nn.Conv1d(2*channel,2*channel, kernel_size=1), nn.SyncBatchNorm(2*channel), nn.ReLU(),
                                         nn.Conv1d(2*channel, 2*channel, kernel_size=1))
        self.cross_filter=attention_propagantion(channel,head)
        self.confidence_filter=PointCN(2*channel,1)
        self.attention_block_self=attention_propagantion(channel,head)
        self.attention_block_up=attention_propagantion(channel,head)

    def forward(self,desc1,desc2,seed_index1,seed_index2):
        cluster1, cluster2 = desc1.gather(dim=-1, index=seed_index1.unsqueeze(1).expand(-1, self.channel, -1)), \
                             desc2.gather(dim=-1, index=seed_index2.unsqueeze(1).expand(-1, self.channel, -1))

        #pooling
        cluster1, cluster2 = self.attention_block_down(cluster1, desc1), self.attention_block_down(cluster2, desc2)
        concate_cluster=self.cluster_filter(torch.cat([cluster1,cluster2],dim=1))
        #filtering
        cluster1,cluster2=self.cross_filter(concate_cluster[:,:self.channel],concate_cluster[:,self.channel:]),\
                        self.cross_filter(concate_cluster[:,self.channel:],concate_cluster[:,:self.channel])
        cluster1,cluster2=self.attention_block_self(cluster1,cluster1),self.attention_block_self(cluster2,cluster2)
        #unpooling
        seed_weight=self.confidence_filter(torch.cat([cluster1,cluster2],dim=1))
        seed_weight=torch.sigmoid(seed_weight).squeeze(1)
        desc1_new,desc2_new=self.attention_block_up(desc1,cluster1,seed_weight),self.attention_block_up(desc2,cluster2,seed_weight)
        return desc1_new,desc2_new,seed_weight

class matcher(nn.Module):
    def __init__(self,config):
        nn.Module.__init__(self)
        self.seed_top_k=config.seed_top_k
        self.conf_bar=config.conf_bar
        self.seed_radius_coe=config.seed_radius_coe
        self.use_score_encoding=config.use_score_encoding
        self.detach_iter=config.detach_iter
        self.seedlayer=config.seedlayer
        self.layer_num=config.layer_num
        self.sink_iter=config.sink_iter

        self.position_encoder = nn.Sequential(nn.Conv1d(3, 32, kernel_size=1) if config.use_score_encoding else nn.Conv1d(2, 32, kernel_size=1),
                                            nn.SyncBatchNorm(32),nn.ReLU(),
                                            nn.Conv1d(32, 64, kernel_size=1), nn.SyncBatchNorm(64),nn.ReLU(),
                                            nn.Conv1d(64, 128, kernel_size=1), nn.SyncBatchNorm(128),nn.ReLU(),
                                            nn.Conv1d(128, 256, kernel_size=1), nn.SyncBatchNorm(256),nn.ReLU(),
                                            nn.Conv1d(256, config.net_channels, kernel_size=1))

        self.hybrid_block=nn.Sequential(*[hybrid_block(config.net_channels, config.head) for _ in range(config.layer_num)])
        self.final_project = nn.Conv1d(config.net_channels, config.net_channels, kernel_size=1)
        self.dustbin=nn.Parameter(torch.tensor(1.5,dtype=torch.float32))

        #if reseeding
        if len(config.seedlayer)!=1:
            self.mid_dustbin=nn.ParameterDict({str(i):nn.Parameter(torch.tensor(2,dtype=torch.float32)) for i in config.seedlayer[1:]})
            self.mid_final_project = nn.Conv1d(config.net_channels, config.net_channels, kernel_size=1)

    def forward(self,data,test_mode=True):
        x1, x2, desc1, desc2 = data['x1'][:,:,:2], data['x2'][:,:,:2], data['desc1'], data['desc2']
        desc1, desc2 = torch.nn.functional.normalize(desc1,dim=-1), torch.nn.functional.normalize(desc2,dim=-1)
        if test_mode:
            encode_x1,encode_x2=data['x1'],data['x2']
        else:
            encode_x1,encode_x2=data['aug_x1'], data['aug_x2']

        #preparation
        desc_dismat=(2-2*torch.matmul(desc1,desc2.transpose(1,2))).sqrt_()
        values,nn_index=torch.topk(desc_dismat,k=2,largest=False,dim=-1,sorted=True)
        nn_index2=torch.min(desc_dismat,dim=1).indices.squeeze(1)
        inverse_ratio_score,nn_index1=values[:,:,1]/values[:,:,0],nn_index[:,:,0]#get inverse score

        #initial seeding
        seed_index1,seed_index2=seeding(nn_index1,nn_index2,x1,x2,self.seed_top_k[0],inverse_ratio_score,self.conf_bar[0],\
                                self.seed_radius_coe,test=test_mode)

        #position encoding
        desc1,desc2=desc1.transpose(1,2),desc2.transpose(1,2)
        if not self.use_score_encoding:
            encode_x1,encode_x2=encode_x1[:,:,:2],encode_x2[:,:,:2]
        encode_x1,encode_x2=encode_x1.transpose(1,2),encode_x2.transpose(1,2)
        x1_pos_embedding, x2_pos_embedding = self.position_encoder(encode_x1), self.position_encoder(encode_x2)
        aug_desc1, aug_desc2 = x1_pos_embedding + desc1, x2_pos_embedding + desc2

        seed_weight_tower,mid_p_tower,seed_index_tower,nn_index_tower=[],[],[],[]
        seed_index_tower.append(torch.stack([seed_index1, seed_index2],dim=-1))
        nn_index_tower.append(nn_index1)

        seed_para_index=0
        for i in range(self.layer_num):
            #mid seeding
            if i in self.seedlayer and i!= 0:
                seed_para_index+=1
                aug_desc1,aug_desc2=self.mid_final_project(aug_desc1),self.mid_final_project(aug_desc2)
                M=torch.matmul(aug_desc1.transpose(1,2),aug_desc2)
                p=sink_algorithm(M,self.mid_dustbin[str(i)],self.sink_iter[seed_para_index-1])
                mid_p_tower.append(p)
                #rematching with p
                values,nn_index=torch.topk(p[:,:-1,:-1],k=1,dim=-1)
                nn_index2=torch.max(p[:,:-1,:-1],dim=1).indices.squeeze(1)
                p_match_score,nn_index1=values[:,:,0],nn_index[:,:,0]
                #reseeding
                seed_index1, seed_index2 = seeding(nn_index1,nn_index2,x1,x2,self.seed_top_k[seed_para_index],p_match_score,\
                                                    self.conf_bar[seed_para_index],self.seed_radius_coe,test=test_mode)
                seed_index_tower.append(torch.stack([seed_index1, seed_index2],dim=-1)), nn_index_tower.append(nn_index1)
                if not test_mode and data['step']<self.detach_iter:
                    aug_desc1,aug_desc2=aug_desc1.detach(),aug_desc2.detach()

            aug_desc1, aug_desc2,seed_weight=self.hybrid_block[i](aug_desc1, aug_desc2,seed_index1,seed_index2)
            seed_weight_tower.append(seed_weight)

        aug_desc1,aug_desc2 = self.final_project(aug_desc1), self.final_project(aug_desc2)
        cmat = torch.matmul(aug_desc1.transpose(1, 2), aug_desc2)
        p = sink_algorithm(cmat, self.dustbin,self.sink_iter[-1])
        #seed_weight_tower: l*b*k
        #seed_index_tower: l*b*k*2
        #nn_index_tower: seed_l*b
        return {'p':p,'seed_conf':seed_weight_tower,'seed_index':seed_index_tower,'mid_p':mid_p_tower,'nn_index':nn_index_tower}


from typing import TYPE_CHECKING
from magicgui import magic_factory
if TYPE_CHECKING:
    import napari

def connect_events(widget):
    """
    Connect widget events to make some visible/invisible depending on others
    """

    def toggle_median_n_widget(event):
        widget.median_n.visible = event
    # Connect events
    widget.apply_median.changed.connect(toggle_median_n_widget)
    # Intial visibility states
    widget.median_n.visible = False

    def toggle_tile_widget(event):
        widget.tile.visible = event
    # Connect events
    widget.tile.changed.connect(toggle_tile_widget)
    # Intial visibility states
    widget.tile.visible = False

@magic_factory(widget_init=connect_events)
def phasor_plot(image_layer: "napari.layers.Image",
                harmonic: int = 1,
                threshold: int = 0,
                apply_median: bool = False,
                median_n: int = 1,
                tile: bool = False,
                vertical_dim: int = 1024,
                horizontal_dim: int = 1024,
                vertical_im: int = 1,
                horizontal_im: int = 1,
                v_overlapping_per: float = 0.05,
                h_overlapping_per: float = 0.05,
                store_dir: bool = False, # set true for bidirectional storing
                napari_viewer: "napari.Viewer" = None) -> None:
    """Calculate phasor components from HSI image and plot them.
    Parameters
    ----------
    image_layer : napari.layers.Image
        napari image layer with HSI data.
    harmonic : int, optional
        the harmonic to display in the phasor plot, by default 1
    threshold : int, optional
        pixels with summed intensity below this threshold will be discarded, by default 0
    apply_median : bool, optional
        apply median filter to image before phasor calculation, by default False (median_n is ignored)
    median_n : int, optional
        number of iterations of median filter, by default 1
    napari_viewer : napari.Viewer, optional
        napari viewer instance, by default None
    """

    import numpy as np
    import pandas as pd
    from napari.layers import Labels
    import dask.array as da

    from napari_hsi_phasor.hsitools import phasor, median_filter, tilephasor, stitching
    from napari_hsi_phasor._plotter import PhasorPlotterWidget
    from skimage.segmentation import relabel_sequential

    image = image_layer.data

    if tile:
        dc, g, s = tilephasor(image, vertical_dim, horizontal_dim, harmonic=harmonic)
        dc = stitching(dc, vertical_im, horizontal_im, h_overlapping_per, v_overlapping_per, store_dir)
        g = stitching(g, vertical_im, horizontal_im, h_overlapping_per, v_overlapping_per, store_dir)
        s = stitching(s, vertical_im, horizontal_im, h_overlapping_per, v_overlapping_per, store_dir)

    else:
        dc, g, s = phasor(image, harmonic=harmonic)

    if apply_median:
        g = median_filter(g, median_n)
        s = median_filter(s, median_n)

    dc = dc[np.newaxis, np.newaxis, :, :]
    g = g[np.newaxis, np.newaxis, :, :]
    s = s[np.newaxis, np.newaxis, :, :]

    space_mask = dc > threshold
    label_image = np.arange(np.prod(dc.shape)).reshape(dc.shape) + 1
    label_image[~space_mask] = 0
    label_image = relabel_sequential(label_image)[0]
    label_column = np.ravel(label_image[space_mask])

    g_flat_masked = np.ravel(g[space_mask])
    s_flat_masked = np.ravel(s[space_mask])
    if isinstance(g, da.Array):
        g_flat_masked.compute_chunk_sizes()
        s_flat_masked.compute_chunk_sizes()

    phasor_components = pd.DataFrame({
        'label': np.ravel(label_image[space_mask]),
        'G': g_flat_masked,
        'S': s_flat_masked})
    table = phasor_components

    frame = np.arange(dc.shape[0])
    frame = np.repeat(frame, np.prod(dc.shape[1:]))
    table['frame'] = frame[space_mask.ravel()]

    # The layer has to be created here so the plotter can be filled properly
    # below. Overwrite layer if it already exists.
    for layer in napari_viewer.layers:
        if (isinstance(layer, Labels)) & (layer.name == 'Labelled_pixels_from_' + image_layer.name):
            labels_layer = layer
            labels_layer.data = label_image
            labels_layer.features = table
            break
    else:
        labels_layer = napari_viewer.add_labels(label_image,
                                                name='Labelled_pixels_from_' + image_layer.name,
                                                features=table,
                                                scale=image_layer.scale[1:],
                                                visible=False)

    # Check if plotter was alrerady added to dock_widgets
    # TO DO: avoid using private method access to napari_viewer.window._dock_widgets (will be deprecated)
    dock_widgets_names = [key for key,
                                  value in napari_viewer.window._dock_widgets.items()]
    if 'Phasor Plotter Widget (napari-hsi-phasor)' not in dock_widgets_names:
        plotter_widget = PhasorPlotterWidget(napari_viewer)
        napari_viewer.window.add_dock_widget(
            plotter_widget, name='Phasor Plotter Widget (napari-hsi-phasor)')
    else:
        widgets = napari_viewer.window._dock_widgets['Phasor Plotter Widget (napari-hsi-phasor)']
        plotter_widget = widgets.findChild(PhasorPlotterWidget)

    # Get labels layer with labelled pixels (labels)
    plotter_widget.labels_select.value = [
        choice for choice in plotter_widget.labels_select.choices if choice.name.startswith("Labelled_pixels")][0]
    # Set G and S as features to plot (update_axes_list method clears Comboboxes)
    plotter_widget.plot_x_axis.setCurrentIndex(1)
    plotter_widget.plot_y_axis.setCurrentIndex(2)
    plotter_widget.plotting_type.setCurrentIndex(1)

    # Show parent (PlotterWidget) so that run function can run properly
    plotter_widget.parent().show()
    # Disconnect selector to reset collection of points in plotter
    # (it gets reconnected when 'run' method is run)
    plotter_widget.graphics_widget.selector.disconnect()
    plotter_widget.run(labels_layer.features,
                       plotter_widget.plot_x_axis.currentText(),
                       plotter_widget.plot_y_axis.currentText())
    plotter_widget.axes_limits()
    return


import csv
import datetime
import re
from django.core.management.base import BaseCommand
from phones.models import Phone

def str_to_date(str_date):
    pattern = re.compile(r'\d+')
    result = pattern.findall(str_date)
    correct_date = datetime.date(int(result[0]), int(result[1]), int(result[2]))
    return correct_date

class Command(BaseCommand):
    def add_arguments(self, parser):
        pass

    def handle(self, *args, **options):
        with open('phones.csv') as f:
            phones = csv.DictReader(f, delimiter=';')
            for phone in phones:
                phone.pop(None)
                phone['release_date'] = str_to_date(phone['release_date'])
                slugged_name = phone['name'].replace(' ', '-')
                p = Phone(slug=slugged_name, **phone)
                p.save()


import time

from BTrees.OOBTree import OOBTree
from clouddrive import db
from clouddrive import configurator
import requests
import transaction
from urllib.parse import urlencode

LOGIN_URL = 'https://www.amazon.com/ap/oa'
AUTH_URL = 'https://api.amazon.com/auth/o2/token'

SCOPES = [
    'clouddrive:read_all',
    'clouddrive:write'
]
DRIVE_ENDPOINT = 'https://drive.amazonaws.com'

def get_login_url(redirect_uri):
    params = {
        'client_id': configurator.CLIENT_ID,
        'scope': ' '.join(SCOPES),
        'redirect_uri': redirect_uri,
        'response_type': 'code'
    }
    return '%s?%s' % (
        LOGIN_URL,
        urlencode(params))

def authorize(code, redirect_uri):
    params = {
        'grant_type': 'authorization_code',
        'code': code,
        'client_id': configurator.CLIENT_ID,
        'client_secret': configurator.CLIENT_SECRET,
        'redirect_uri': redirect_uri
    }
    resp = requests.post(AUTH_URL, data=params)
    store_credentials(resp.json())
    return resp.json()

def refresh():
    tokens = get_credentials()
    params = {
        'grant_type': 'refresh_token',
        'refresh_token': tokens['refresh_token'],
        'client_id': configurator.CLIENT_ID,
        'client_secret': configurator.CLIENT_SECRET
    }
    resp = requests.post(AUTH_URL, data=params)
    store_credentials(resp.json())

def store_credentials(data):
    root = db.get()
    root['credentials'] = data
    transaction.commit()

def get_credentials():
    root = db.get()
    if 'credentials' in root:
        return root['credentials']

def store_endpoint():
    root = db.get()
    uri = '%s/drive/v1/account/endpoint' % DRIVE_ENDPOINT
    creds = get_credentials()
    resp = requests.get(uri, headers={
        'Authorization': 'Bearer ' + creds['access_token']
    })
    if 'metadata' in root:
        metadata = root['metadata']
    else:
        metadata = root['metadata'] = OOBTree()
    data = resp.json()
    if data.get('message') == 'Token has expired':
        refresh()
        return store_endpoint()
    metadata['endpoint'] = data
    metadata['endpoint_last_retrieved'] = time.time()
    transaction.commit()

def call(path, endpoint_type='content', method='GET', body=None,
         body_type='json', args=None):
    root = db.get()

    metadata = root.get('metadata', {})
    endpoint = metadata.get('endpoint', {})

    if endpoint.get('message') == 'Token has expired' or metadata == {}:
        refresh()
        store_endpoint()
        root._p_jar.sync()
        metadata = root.get('metadata', {})
        endpoint = metadata.get('endpoint', {})

    if endpoint_type == 'content':
        endpoint = endpoint.get('contentUrl')
    else:
        endpoint = endpoint.get('metadataUrl')
    if not endpoint:
        return
    uri = '%s/%s' % (endpoint.rstrip('/'), path.lstrip('/'))
    creds = get_credentials()
    meth = requests.get
    if method == 'POST':
        meth = requests.post
    elif method == 'PUT':
        meth = requests.put
    if args is None:
        args = {}
    if body:
        args[body_type] = body
    resp = meth(uri, headers={
        'Authorization': 'Bearer ' + creds['access_token']
    }, **args)
    if resp.status_code == 401:
        refresh()
        return call(path, endpoint_type, method, body, body_type, args)
    return resp

def list_files():
    return call('nodes?filters=kind:FOLDER', 'metadata').json()

def get_root_folder():
    return call('nodes?filters=kind:FOLDER AND isRoot:true', 'metadata').json()


# %%
from statistics import mean
input = open('input.txt')
lines = input.read().splitlines()
input.close()
# %%
binaries = [list(map(int, list(line))) for line in lines]
transposed = list(map(list, zip(*binaries)))
means = [mean(line) for line in transposed]
gamma_binary = [round(x) for x in means]
epsilon_binary = [1-round(x) for x in means]
# %%
binary = [2**i for i in range(len(gamma_binary))]
binary.reverse()
gamma = 0
epsilon = 0
for index, decimal in enumerate(binary):
    gamma += decimal * gamma_binary[index]
    epsilon += decimal * epsilon_binary[index]
print(gamma * epsilon)
# %%


import math
import torch
import torch.nn as nn

from fairseq.data.data_utils import collate_tokens
from transformers import BartTokenizer

class ConditionalSequenceGenerator:
    """Conditional sequence generator for calculating prior and posterior probability."""
    def __init__(self, bart):
        self.bart = bart
        self.tokenizer = BartTokenizer.from_pretrained("facebook/bart-large", local_files_only=False)

        self.encode_func = bart.encode
        self.decode_func = bart.decode
        self.max_positions = bart.max_positions
        if type(self.max_positions) == int:
            self.max_positions = [self.max_positions]
        self.encode_line = bart.task.source_dictionary.encode_line

        self._initialize()

    def _initialize(self):
        """Set BART model to evaluation mode."""
        self.bart.cuda()
        self.bart.eval()
        self.bart.half()

    def tokenize_target(self, input_str, left_pad=False, append_eos=False):
        """BPE-encode a sentence (or multiple sentences).

        Args:
            input_str (str or List[str]): input sentence to be tokenized.
            left_pad (bool): self-explained.

        Return:
            prev_output_tokens (torch.Tensor): [batch_size, length]
            target (torch.Tensor): [batch_size, length]
            tgt_lengths (torch.Tensor): [batch_size]

        """
        if type(input_str) == type(''):
            input_str = [input_str]

        prev_ids, tgt_ids = [], []

        for ins in input_str:
            tokens = self.bart.bpe.encode(ins)  # <mask>: 1279 27932 29
            calibration = 1
            if len(tokens.split(" ")) > min(self.max_positions) - calibration:
                tokens = " ".join(tokens.split(" ")[: min(self.max_positions) - calibration])

            if append_eos:
                tokens = "<s> " + tokens
            prev_tokens = "</s> " + tokens
            tgt_tokens = tokens + " </s>"

            tgt_ids.append(self.encode_line(tgt_tokens, append_eos=False).long())
            prev_ids.append(self.encode_line(prev_tokens, append_eos=False).long())

        prev_output_tokens = collate_tokens(prev_ids, pad_idx=1, left_pad=left_pad).cuda()
        target = collate_tokens(tgt_ids, pad_idx=1, left_pad=left_pad).cuda()
        tgt_lengths = torch.sum(target != 1, dim=1).cuda()

        return prev_output_tokens, target, tgt_lengths

    def tokenize(self, input_str, append_bos=False, append_eos=True, left_pad=True):
        """BPE-encode a sentence (or multiple sentences).

        Args:
            input_str (str or List[str]): input sentence to be tokenized.
            append_bos (bool): self-explained.
            append_eos (bool): self-explained.

        Return:
            input_ids (torch.Tensor): [batch_size, length]
            src_lengths (torch.Tensor): [batch_size]
        """
        if type(input_str) == type(''):
            input_str = [input_str]

        input_ids = []
        for ins in input_str:
            tokens = self.bart.bpe.encode(ins)  # <mask>: 1279 27932 29
            calibration = sum([append_bos, append_eos])
            if len(tokens.split(" ")) > min(self.max_positions) - calibration:
                tokens = " ".join(tokens.split(" ")[: min(self.max_positions) - calibration])

            tokens = "<s> " + tokens if append_bos else tokens
            tokens = tokens + " </s>" if append_eos else tokens
            ids = self.encode_line(tokens, append_eos=False).long()
            input_ids.append(ids)

        input_ids = collate_tokens(input_ids, pad_idx=1, left_pad=left_pad).cuda()
        input_lengths = torch.sum(input_ids != 1, dim=1).cuda()

        return input_ids, input_lengths

    def tokenize_with_mask(self, input_str):
        """Tokenize sentence with a special <mask> token in it.

        Args:
            input_str (str or List[str]): input sentence to be tokenized.

        Return:
            input_ids (torch.Tensor): [batch_size, length]
            src_lengths (torch.Tensor): [batch_size]
        """
        input_ids = self.tokenizer(input_str, return_tensors='pt', padding=True)['input_ids'].cuda()
        input_lengths = torch.sum(input_ids != 1, dim=1).cuda()
        return input_ids, input_lengths

    def encode_decode(self, src_input, tgt_input, mask_filling=False):
        """
        Args:
            src_input: (List[str])
            tgt_input: (List[str])

        """
        if mask_filling:
            src_tokens, src_lengths = self.tokenize_with_mask(src_input)
            prev_output_tokens, target, tgt_lengths = self.tokenize_target(tgt_input, left_pad=False, append_eos=True)
        else:
            src_tokens, src_lengths = self.tokenize(src_input, append_bos=False)
            prev_output_tokens, target, tgt_lengths = self.tokenize_target(tgt_input, left_pad=False)

        with torch.no_grad():
            encoder_out = self.bart.model.encoder(src_tokens, src_lengths=src_lengths)
            decoder_out = self.bart.model.decoder(prev_output_tokens, encoder_out=encoder_out, features_only=False)

            probs = nn.functional.softmax(decoder_out[0], dim=-1)
            tgt_token_probs = torch.gather(probs, 2, target.unsqueeze(-1)).squeeze(2)

            # mask <pad> with probability 1.0
            max_tgt_length = tgt_lengths.max().item()
            tgt_lengths = tgt_lengths - 1
            tgt_mask = torch.arange(max_tgt_length)[None, :].cuda() < tgt_lengths[:, None]
            tgt_token_probs.masked_fill_(tgt_mask == False, 1.0)

        return tgt_token_probs, target

    def generate(self, src_input, tgt_input=None):
        """Conditional generation.

        Args:
            src_input (str or List[str]): input source sentence to be tokenized.
            tgt_input (str or List[str]): input target sentence to be tokenized.
        """
        input_ids, lengths = self.tokenize(src_input, append_bos=False)

        target_ids = None
        if tgt_input is not None:
            assert len(src_input) == len(tgt_input), "source & target length should match."
            target_ids, _ = self.tokenize(tgt_input, append_bos=False, left_pad=False)

        with torch.no_grad():
            encoder_output = self.encode_sequence(input_ids, lengths)
            decoder_output = self.decode_sequence(encoder_output,
                                                  target_ids=target_ids,
                                                  prefix_tokens=[2])
        return decoder_output

    def mask_filling(self, src_input, tgt_input=None):
        """
        Filling the mask in sentence(s).
        """
        input_ids, lengths = self.tokenize_with_mask(src_input)

        target_ids = None
        if tgt_input is not None:
            assert len(src_input) == len(tgt_input), "source & target length should match."
            target_ids, _ = self.tokenize(tgt_input, left_pad=False)

        with torch.no_grad():
            encoder_output = self.encode_sequence(input_ids, lengths)
            decoder_output = self.decode_sequence(encoder_output,
                                                  target_ids=target_ids,
                                                  prefix_tokens=[2, 0])
        return decoder_output

    def encode_sequence(self, input_ids, lengths):
        return self.bart.model.encoder(input_ids, src_lengths=lengths)

    def decode_sequence(
        self,
        encoder_out,
        target_ids=None,
        min_decode_step=3,
        max_decode_step=100,
        pad_id=1,
        eos_id=2,
        prefix_tokens=[2, 0],
    ):
        batch_size = encoder_out['encoder_padding_mask'][0].shape[0]
        init_input = torch.tensor([prefix_tokens] * batch_size, dtype=torch.long).cuda()
        token_probs, tokens = None, [[] for i in range(batch_size)]
        end_mask = torch.tensor([False] * batch_size).cuda()

        softmax = nn.Softmax(dim=1)
        for step in range(max_decode_step):
            decoder_outputs = self.bart.model.decoder(init_input, encoder_out, features_only=False)
            logits = decoder_outputs[0][:, -1, :]  # logits: [batch_size, vocab]
            attn = decoder_outputs[1]['attn'][0]  # [batch_size, prev_token_len, src_token_len]

            if step + 1 < min_decode_step:
                logits[:, eos_id] = -math.inf  # mask <EOS> token when within minimal step
            logits[:, pad_id], logits[:, 0] = -math.inf, -math.inf  # never select <PAD> & <BOS> token
            probs = softmax(logits)  # probs: [batch_size, vocab]

            # select tokens
            if target_ids is not None:
                selected_token = target_ids[:, step]
            else:
                value, indices = torch.topk(probs, 5, dim=1)
                selected_token = indices[:, 0]

            selected_token = selected_token.masked_fill(end_mask, pad_id)
            init_input = torch.cat([init_input, selected_token.unsqueeze(1)], dim=-1)

            probs = torch.gather(probs, 1, selected_token.unsqueeze(1)).detach()
            probs = probs.masked_fill(end_mask.unsqueeze(1), 1.0)

            # str & probability
            token_probs = probs if token_probs is None else torch.cat([token_probs, probs], dim=-1)
            for t, s in zip(tokens, selected_token):
                t.append(self.decode_func(s.unsqueeze(0)) if s.item() != pad_id else '<pad>')

            # stop generation when all finished
            end_mask = torch.logical_or(end_mask, selected_token == eos_id)
            if end_mask.sum().item() == batch_size:
                break

        return init_input, tokens, token_probs


__author__ = " Parya Derakhshan "

import numpy as np
from collections import OrderedDict
from inout import Error,Log

EPSILON = '~'
START,END='START','EOF'
ACCEPT,ERROR,POP='ACCEPT','ERROR','POP'
errors = {
    'll1':'GrammarError',
    'prsr':'ParserError [LR1]'
}

class Item(): # variable --> [completed] . [expected] , [lookahead]
    def __init__(self, variable, completed, expected, lookahead):

        '''handle epsilon line an empty symbol'''
        if completed == EPSILON: completed=[]
        if expected == EPSILON: expected=[]
        if lookahead == EPSILON: lookahead=[]
        if EPSILON in completed: completed.remove(EPSILON)
        if EPSILON in expected: expected.remove(EPSILON)

        self.lookahead=[]
        for l in lookahead:
            if l == EPSILON or l==[]: continue
            '''delete repeates from lookahead'''
            if not l in self.lookahead: self.lookahead.append(l)

        self.variable,self.completed,self.expected=variable, completed, expected

    def same_rule(self, item):
        '''check if variable --> completed . expected are the same'''
        return self.variable==item.variable and self.completed==item.completed and self.expected==item.expected

    def equal(self,item):
        '''check if items are the same'''
        return self.same_rule(item) and self.lookahead==item.lookahead

    def print_item(self):
        '''return( variable --> completed . expected , lookahead)'''
        return '{}-->{}.{},{}'.format(self.variable, self.completed, self.expected, self.lookahead)

class State(): # header:item, closure:[items], edge:[(label,state)], id: int
    count=0 #state id
    def __init__(self, item):
        self.header = item
        self.closure, self.edge = [], []
        self.id = State.count
        State.count+=1

    def add_node(self,symbol,state):
        '''add (label,neighbor) to edge'''
        self.edge.append((symbol,state))

    def add_closure(self, item):
        '''add new closure item if it doesn't already exist'''
        for c in self.closure:
            if item.same_rule(c):
                '''if there's an item in the closure with different lookaheads just update lookaheads'''
                tmp = c.lookahead+item.lookahead
                c.lookahead = list(set(tmp))
                return True
        self.closure.append(item)

    def get_edges(self):
        '''get neighbor states based on the first symbol in [expected]'''
        header,closures=self.header, self.closure
        if not header.expected: return[] #leaf
        edges, items = [], ([header] + closures)
        for item in items:
            if len(item.expected)>0:
                tmp = item.completed + [item.expected[0]]
                item_new = Item(item.variable,tmp,item.expected[1:],item.lookahead)
                edges.append((item.expected[0],item_new))
        return edges

    def print_state(self):
        '''return id, header, number of edges'''
        r =('{}) \n\t{}\t:{} edges\n{}'.format(self.id,self.header.print_item(),len(self.edge),('_'*20)))
        '''return closure items'''
        for i in self.closure:
            r+=('{}\n\n'.format(i.print_item()))
        r+=('\n'*2)
        return r

class LR1():
    def __init__(self,grammar):
        self.productions, self.variables, self.terminals = self.read_grammar(grammar)
        if EPSILON in self.terminals: self.terminals.remove(EPSILON)
        self.symbols = self.variables+self.terminals
        self.epsilons = self.mark_epsilon()
        self.firsts = self.first()
        self.automata = self.graph()

    def components(self):
        return self.symbols,self.variables, self.automata, self.rules

    def read_grammar(self,lines): #read grammar.txt and sets terminals,variables and rules
        if len(lines)==0: Error().parser(errors['ll1'],'Grammar Not Found')
        productions,productions_all=OrderedDict(()),OrderedDict(())
        variables,terminals,words = [],[],[]
        for i in range(len(lines)):
            line = lines[i]
            line.strip()
            left,right = line.split('->')
            left = left.replace(' ','')
            if not left in variables: variables.append(left)
            key = str(i)+'.'+left
            productions_all[key]=[]
            for r in right.split(' '):
                r= r.strip()
                if r.isspace() or not r: continue
                if r ==EPSILON : productions_all[key]=r[0]
                else: productions_all[key].append(r)
                words.append(r)
        for w in words:
            if w not in terminals and w not in variables: terminals.append(w)
        for v in variables:
            productions[v]=[]
            for key,value in productions_all.items():
                if value[0]==EPSILON: value=value[0]
                if key.split('.')[1]==v: productions[v].append(value)
        self.rules = productions_all
        return productions,variables,terminals

    def mark_epsilon(self): #marking algorithm for variables that can drive epsilon
        marked = []
        variables = self.variables
        for v in variables:
            derivations = self.productions[v]
            if EPSILON in derivations: marked.append(v)
        for i in range(len(variables)):
            v = variables[i]
            if v in marked: continue
            derivations = self.productions[v]
            for der in derivations:
                all_marked =True
                for s in der:
                    if (s in self.terminals) or (s not in marked):
                        all_marked = False
                        break
                if all_marked:
                    i = 0
                    marked.append(v)
                    break
        return marked

    def relation_F(self): # relation F for firsts (matrix output)
        n = len(self.symbols)
        f1 = np.full((n, n), 0)
        for v in self.variables:
            i = self.symbols.index(v)
            derivations = self.productions[v]
            for der in derivations:
                if der[0] == EPSILON: continue
                j = self.symbols.index(der[0])
                if f1[i,j]==1: return Error().parser(errors['ll1'],'Condition I')
                f1[i,j]=1
        F, F_sum=[f1], f1
        while(F[-1].any()):
            f_p = np.linalg.matrix_power(f1,(len(F)+1))
            F.append(f_p)
            F_sum = np.add(F_sum,f_p)
            if F_sum.max()>1: return Error().parser(errors['ll1'],'Condition I')
        self.F = F_sum

    def first(self): # firsts for variables (dictionary output)
        self.relation_F()
        first = {}
        n_v = len(self.variables)
        for i in range(n_v):
            v = self.variables[i]
            first[v]=[]
            row =self.F[i,n_v:]
            for j in range(len(row)):
                if row[j]: first[v].append(self.terminals[j])
            if v in self.epsilons:
                first[v].append(EPSILON)
        return first

    def rule_index(self,var,der): # index)variable->derivation
        '''get rule index based on the variable and derivations in it'''
        for key,derivation in self.rules.items():
            index,variable = key.split('.')
            if variable == var and derivation==der: return index
        return ERROR

    def rule_derivation(self,i): # index)variable->derivation
        '''get rule's variable and derivations based on its index'''
        for key,derivation in self.rules.items():
            index,variable = key.split('.')
            if index == i: return variable, derivation
        return ERROR

    def expected_lookahead(self, expected, lookahead):
        '''for item with .[expected],[lookahead] returns closures:[(var,lookaheads)]'''

        '''leaf with no closure'''
        if not expected or expected[0] in self.terminals:  return 0

        '''no additional lookahead and just one variable'''
        if len(expected)==1: return[(expected[0],lookahead)]

        '''add the first expected variable and all possible lookaheads'''
        firsts=[(expected[0],self.lookahead(expected[1:],lookahead))]

        '''check for epsilon driven variables after the first expected symbol'''
        if expected[0] in self.epsilons:
            follow = self.expected_lookahead(expected[1:],lookahead)
            if follow: firsts += follow

        return firsts

    def lookahead(self,expected,lookahead):
        '''for item with .[expected[1:]],[lookahead] returns [lookaheads]'''

        '''leaf with no closures'''
        if not expected: return []

        '''the first symbol is a terminal'''
        if expected[0] in self.terminals: return [expected[0]]

        '''add the first variables firsts'''
        firsts, reach_end = self.firsts[expected[0]], True

        '''check for epsilon driven variables after the first expected variable'''
        if expected[0] in self.epsilons and len(expected)>1:
            follow = self.lookahead(expected[1:],lookahead)
            if follow: firsts+=follow

        '''if all the variables in expected and drive epsilon also add the original lookahead'''
        for e in expected:
            if not e in self.epsilons:
                reach_end=False
                break
        if reach_end: firsts+=lookahead

        return firsts

    def state_closures(self,state):
        '''calculate state's closure items'''

        '''state's header item'''
        header, check = state.header, []

        '''header's directly deriven closure items'''
        items = self.expected_lookahead(header.expected, header.lookahead)
        if items ==0: return state #no closure items
        for var,lookahead in items:
            for derivation in self.productions[var]:
                new_item = Item(var,[],derivation,lookahead)
                if not new_item in state.closure:
                    state.add_closure(new_item)
                    check.append(new_item)

        '''closure items driven from unchecked closure items in the state'''
        while check:
            item = check.pop(0)
            items = self.expected_lookahead(item.expected, item.lookahead)
            if items == 0: continue
            for var,lookahead in items:
                for derivation in self.productions[var]:
                    new_item = Item(var,[],derivation,lookahead)
                    if not new_item in state.closure:
                        state.add_closure(new_item)
                        check.append(new_item)

        return state

    def graph(self):
        '''parser's NFA'''
        automata ={}

        '''initial state: id0)ZERO-->[],START,EOF'''
        state0 = State(Item('ZERO', [], ['START'], ['EOF']))
        automata[state0.id]=state0

        '''create all possible states'''
        id=0
        while id<State.count:

            '''calculate closures based on the header'''
            state = self.state_closures(automata[id])

            '''claculate (edge,neighbor) based on first expected symbol'''
            edges= state.get_edges()

            for edge,item in edges:

                '''check if a state with similar header item doesn't already exist'''
                check,i=True,0
                for k, s in automata.items():
                    if s.header.equal(item):
                        check,i=False,s.id
                        break
                '''create new neighboring state'''
                if check:
                    state_new = State(item)
                    i= state_new.id
                    automata[i]=state_new

                '''add state's edges'''
                if (edge,i) not in state.edge: state.edge.append((edge,i))

            id+=1
        return automata

class Parser():
    def __init__(self):
        self.stack,self.output=[0],[]

    def table_automata(self,grammar):
        self.grammar = LR1(grammar)
        self.symbols,self.variables,self.automata,self.productions = self.grammar.components()
        self.table=self.parsing_table( self.symbols,self.variables,self.automata)
        return self.table_output()

    def parsing_table(self,symbols,variables,automata):
        table,self.ids = {},[]
        '''table[id][terminal-variable]'''
        for id, state in automata.items():
            self.ids.append(id)
            table[id]={}
            for symbol in symbols: table[id][symbol]=ERROR
            for edge, node in state.edge:
                if edge in variables: table[id][edge]=('G',node) #GoTo
                else: table[id][edge]=('S',node) #shiftTo
            items = [state.header]+state.closure
            for item in items:
                if not item.expected:  #leaf or epsilon derivation
                    var,der = item.variable, item.completed
                    if not item.completed: der = EPSILON
                    elif item.lookahead == [END] and item.completed==[START]:
                        for s in symbols: table[id][s]=ACCEPT
                        self.accept_id = id
                        continue
                    index = self.grammar.rule_index(var,der)
                    for terminal in item.lookahead:
                        table[id][terminal]=('R',index) #ReduceTo
        return table

    def table_output(self):
        txt = []
        txt.append('{:4s}'.format(' '))
        for symbol in self.symbols: txt.append('|{:11s}'.format(symbol))
        txt.append('\n'+'_'*(4+12*len(self.symbols))+'\n')
        for id in self.ids:
            txt.append('{:4.0f}'.format(id))
            for symbol in self.symbols:
                txt.append('|{:11s}'.format(str(self.table[id][symbol])))
            txt.append('\n'+'_'*(4+12*len(self.symbols))+'\n')
        return txt

    def automata_output(self):
        txt=[]
        for id, state in self.automata.items():
            txt.append(state.print_state())
        return txt

    def parse_tokens(self,tokens): #read tokens.txt
        if len(tokens)==0: return Error().parser(errors['prsr'],'Tokens Not Found')
        symbols=[]
        for token in tokens:
            code,const = token.split(EPSILON)
            symbols.append(const.strip())
        symbols.append(END)
        result = self.parse(symbols)
        if result==ERROR:
            self.output.append('<Production Rule Not Found>')
            Error().parser(errors['prsr'],'Invalid Command')
        else: Log().success('Program Was Parsed Successfully [LR1]')
        self.output.reverse()
        return self.output

    def get_rule(self,number,var,rule):
        der = ''
        for symbol in rule: der+= (symbol+' ')
        return number+') '+var+' -> '+der.strip()+'\n'

    def parse(self,tokens): #parse tokens according to the parse table
        while tokens:
            A ,a = self.stack[-1],tokens[0]
            if A in self.variables: #GoTo
                A2 = self.stack[-2]
                M = self.table[A2][A]
                if not type(M) is tuple or not M[0] =='G': return ERROR
                self.stack.append(M[1])
            elif A in self.ids:
                M = self.table[A][a]
                if M== ACCEPT: return self.output
                if not type(M) is tuple: return ERROR
                instruction, id = M
                if instruction not in ['R','S']: return ERROR
                if instruction=='S': #ShifTo
                    del tokens[0]
                    self.stack.append(a)
                    self.stack.append(id)
                else: #ReduceTo
                    variable, derivation = self.grammar.rule_derivation(id)
                    if derivation!=EPSILON:
                        for i in range(len(derivation)-1,-1,-1):
                            if self.stack[-2]!= derivation[i]: return ERROR
                            self.stack.pop(-1)
                            self.stack.pop(-1)
                    self.stack.append(variable)
                    self.output.append(self.get_rule(id,variable,derivation))
            else: return ERROR

def run(files,dir_grammar,dir_parsetable,dir_automata,dir_tokens,dir_rules):
    '''read grammar and calculate first and automata'''
    grammar = files.set_directory(dir_grammar,dir_parsetable)
    parser = Parser()
    '''calculate and print parsing table'''
    table = parser.table_automata(grammar)
    files.write(table)
    files.close()
    '''print automata'''
    f = files.set_directory(dir_grammar,dir_automata)
    files.write(parser.automata_output())
    files.close()
    '''read tokens and parse the program'''
    tokens = files.set_directory(dir_tokens,dir_rules)
    rules = parser.parse_tokens(tokens)
    '''print rules'''
    files.write(rules)
    files.close()

if __name__ == "__main__":
    from inout import File

    dir_program = 'program.txt'
    vis,outs = 'visualizations','outputs'
    files = File(dir_program,[vis,outs])

    dir_tokens, dir_grammar, dir_rules_lr1 = (outs+'/tokens.txt'),'grammar.txt', (outs+'/rules_lr1.txt')
    dir_parsetable_lr1, dir_automata =(vis+'/parsetable_lr1.txt'), (vis+'/automata.txt')

    run(files,dir_grammar,dir_parsetable_lr1,dir_automata,dir_tokens,dir_rules_lr1)


import os
import scipy.misc
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf

from model import DCGAN
from utils import pp, visualize, to_json, show_all_variables, inverse_transform
from time import gmtime, strftime

flags = tf.app.flags
flags.DEFINE_integer("epoch", 50, "Epoch to train [25]")
flags.DEFINE_float("learning_rate", 0.0002, "Learning rate of for adam [0.0002]")
flags.DEFINE_float("beta1", 0.5, "Momentum term of adam [0.5]")
flags.DEFINE_integer("train_size", 1000000000, "The size of train images [np.inf]")
flags.DEFINE_integer("batch_size", 64, "The size of batch images [64]")
flags.DEFINE_integer("input_height", 108, "The size of image to use (will be center cropped). [108]")
flags.DEFINE_integer("input_width", None, "The size of image to use (will be center cropped). If None, same value as input_height [None]")
flags.DEFINE_integer("output_height", 64, "The size of the output images to produce [64]")
flags.DEFINE_integer("output_width", None, "The size of the output images to produce. If None, same value as output_height [None]")
flags.DEFINE_string("data_dir", "./../data", "The name of dataset [celebA, mnist, lsun]")
flags.DEFINE_string("dataset", "celebA", "The name of dataset [celebA, mnist, lsun]")
flags.DEFINE_string("data_list", None, "The file list (if used, input_fname_pattern is ignored). Pathes should be relative to data_dir")
flags.DEFINE_string("input_fname_pattern", "*.png", "Glob pattern of filename of input images [*]")
flags.DEFINE_string("checkpoint_dir", "checkpoint", "Directory name to save the checkpoints [checkpoint]")
flags.DEFINE_string("sample_dir", "samples", "Directory name to save the image samples [samples]")
flags.DEFINE_boolean("train", False, "True for training, False for testing [False]")
flags.DEFINE_boolean("crop", False, "True for training, False for testing [False]")
flags.DEFINE_boolean("test", False, "True for visualizing, False for nothing [False]")
flags.DEFINE_boolean("generate", False, "True for generating set of images, False for training [False]")
flags.DEFINE_integer("generate_num", 10, "Number of images to generate [10]")
flags.DEFINE_float("generate_min_D_prob", None, "Min. discriminator output score for generated images [None]")
flags.DEFINE_string("generate_name_postfix", "" , "Postfix of generated images")
flags.DEFINE_string("generate_dir", "./../generated" , "Path to generated images")
flags.DEFINE_integer("test_mode", 7, "Visualizing mode: 0 - ; 1 - ; 2 - ;3 - ;4 - ;5 - ;6 - ;7 -")
flags.DEFINE_integer("z_dim", 100, "Latent space size")
flags.DEFINE_integer("y_dim", None, " space size")
flags.DEFINE_boolean("use_double_G_train", True, "use_double_G_train")
flags.DEFINE_boolean("add_dense", False, "Adding dense layer on top of generator [False]")
flags.DEFINE_boolean("generate_known_z", False, "Generate using predefined z values")
flags.DEFINE_string("known_z_path", "./../data", "File with mean Z values to use for generation")

FLAGS = flags.FLAGS

def main(_):
  pp(flags.FLAGS.__flags)

  if FLAGS.input_width is None:
    FLAGS.input_width = FLAGS.input_height
  if FLAGS.output_width is None:
    FLAGS.output_width = FLAGS.output_height

  if not os.path.exists(FLAGS.checkpoint_dir):
    os.makedirs(FLAGS.checkpoint_dir)
  if not os.path.exists(FLAGS.sample_dir):
    os.makedirs(FLAGS.sample_dir)

  #gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
  run_config = tf.ConfigProto()
  run_config.gpu_options.allow_growth=True

  with tf.Session(config=run_config) as sess:
    dcgan = DCGAN(
          sess,
          input_width=FLAGS.input_width,
          input_height=FLAGS.input_height,
          output_width=FLAGS.output_width,
          output_height=FLAGS.output_height,
          batch_size=FLAGS.batch_size,
          sample_num=FLAGS.batch_size,
          y_dim=FLAGS.y_dim,
          z_dim=FLAGS.z_dim,
          dataset_name=FLAGS.dataset,
          data_dir=FLAGS.data_dir,
          data_list=FLAGS.data_list,
          input_fname_pattern=FLAGS.input_fname_pattern,
          crop=FLAGS.crop,
          checkpoint_dir=FLAGS.checkpoint_dir,
          sample_dir=FLAGS.sample_dir,
          use_double_G_train=FLAGS.use_double_G_train,
          add_deconv=FLAGS.add_dense)

    show_all_variables()

    if FLAGS.train:
        train_dcgan(dcgan)

    if not FLAGS.train:
        if not dcgan.load(FLAGS.checkpoint_dir)[0]:
            raise Exception("[!] Train a model first, then run test mode")

    if FLAGS.test:
        test_dcgan(dcgan, sess)

    if FLAGS.generate:
        generated_file_pathes = generate_images(dcgan, sess)

    if FLAGS.generate_known_z:
        z_mean = np.loadtxt(FLAGS.known_z_path)
        generated_file_pathes = generate_images(dcgan, sess, z_mean)

def train_dcgan(dcgan):
    dcgan.train(FLAGS)

def test_dcgan(dcgan, sess):
    visualize(sess, dcgan, FLAGS, FLAGS.visualize_mode)

def generate_images(dcgan, sess, predefined_mean_z=None, z_sigma=0.05):
    print(FLAGS.generate_name_postfix)
    strLog = ""
    generation_data = []
    dir = FLAGS.generate_dir
    if not os.path.exists(dir):
        os.makedirs(dir)

    imgInd = 0
    while (imgInd < FLAGS.generate_num):
        if predefined_mean_z is None:
            z_sample = np.random.uniform(-1, 1, size=(FLAGS.batch_size, FLAGS.z_dim))
        else:
            z_sample = [[np.random.normal(mu, z_sigma) for mu in z_sigma] for _ in range(FLAGS.batch_size)]
            z_sample = np.array(z_sample)

        samples = sess.run(dcgan.sampler, feed_dict={dcgan.z: z_sample})

        ps = sess.run(dcgan.D, feed_dict={dcgan.inputs: samples})
        probs = ps[:, 0]

        images = inverse_transform(samples)
        for im, z, d_prob in zip(images, z_sample, probs):
            if imgInd >= FLAGS.generate_num:
                break

            if FLAGS.generate_min_D_prob is None or d_prob >= FLAGS.generate_min_D_prob:
                filename = 'gen_%d_dPr_%.2f.png' % (imgInd, d_prob)
                scipy.misc.imsave(os.path.join(dir, filename), np.squeeze(im))
                strLog += filename + ',' + FLAGS.generate_name_postfix + ',' + strftime("%Y%m%d%H%M%S", gmtime()) + '\n'
                generation_data.append((filename, str(d_prob),','.join([str(_) for _ in z])))
                imgInd += 1

            if imgInd % 50 == 0:
                print(imgInd)

    f = open(dir + "/generated_list.txt", "a")
    f.write(strLog)
    f.close()
    with open(dir + "/generated_info.csv", 'w') as ff:
        ff.write('file,D_prob,' + ','.join(['z' + str(_) for _ in range(0, FLAGS.z_dim)]) + '\n')
        for g_d in generation_data:
            ff.write(','.join(g_d) + '\n')
    # plt.hist(probs)
    # plt.show()
    return [os.path.join(dir, gd[0]) for gd in generation_data]

if __name__ == '__main__':
  tf.app.run()


from typing import List

class Solution:
    """灵感来自花花的longest substring of XXX的模板
    https://zxi.mytechroad.com/blog/hashtable/leetcode-1542-find-longest-awesome-substring/"""

    def maxSubArrayLen(self, nums: List[int], k: int) -> int:
        idx = dict()
        idx[0] = -1
        s = 0  # 初始状态, 表示前缀数组的和
        ans = 0
        for i, num in enumerate(nums):
            s += num
            if s - k in idx:
                ans = max(ans, i - idx[s - k])
            if s not in idx:
                idx[s] = i
        return ans

if __name__ == '__main__':
    nums = [1, -1, 5, -2, 3]
    k = 3
    s = Solution()
    assert s.maxSubArrayLen(nums, k) == 4

    nums = [-2, -1, 2, 1]
    k = 1
    assert s.maxSubArrayLen(nums, k) == 2


#!/usr/bin/env python3

def psfstudy(band_path, ztfname, filtercode, logger, args):
    import pathlib
    import gc
    import traceback

    import numpy as np
    from numpy.polynomial.polynomial import Polynomial
    import matplotlib.pyplot as plt
    import matplotlib
    from matplotlib.cm import ScalarMappable
    from matplotlib.colorbar import Colorbar
    from saunerie.plottools import binplot
    import pandas as pd
    from dask import delayed, compute

    from deppol_utils import quadrants_from_band_path
    from utils import get_header_from_quadrant_path, plot_ztf_focal_plan_values, ListTable, match_by_gaia_catalogs, quadrant_name_explode, match_to_gaia

    matplotlib.use('Agg')

    output_path = band_path.joinpath("psfstudy")
    output_path.mkdir(exist_ok=True)

    quadrant_paths = quadrants_from_band_path(band_path, logger)

    # Remove bad QC quadrants
    def _quadrant_good(quadrant_path):
        if not quadrant_path.joinpath("match_gaia.success").exists():
            return False

        if len(ListTable.from_filename(quadrant_path.joinpath("psfstars.list")).df) < 150:
            return False

        return True

    print("Found {} quadrants".format(len(quadrant_paths)))
    # quadrant_paths = list(filter(_quadrant_good, quadrant_paths))
    # print("After QC: {}".format(len(quadrant_paths)))

    min_mag, max_mag = 13., 19.
    low_bin, high_bin = 14, 17

    def _process_quadrant(quadrant_path, cat1, cat2):
        def _cat_to_name(cat):
            if cat == 'flux':
                return "PSF"
            else:
                return "Aperture-" + cat[-1]

        def _is_aperture(cat):
            return 'apfl' in cat

        def _get_cat_from_name(cat, refexp=None):
            path = quadrant_path
            if refexp:
                path = quadrant_path.with_name(refexp + quadrant_path.name[28:])

            if cat == 'flux':
                with pd.HDFStore(path.joinpath("matched_stars.hd5")) as hdfstore:
                    cat_df = hdfstore.get('/matched_stars')
            else:
                cat_df = ListTable.from_filename(path.joinpath("standalone_stars.list")).df

            return cat_df

        def _get_gaia_cat(refexp=None):
            path = quadrant_path
            if refexp:
                path = quadrant_path.with_name(refexp + quadrant_path.name[28:])

            with pd.HDFStore(path.joinpath("matched_stars.hd5")) as hdfstore:
                cat_gaia_df = hdfstore.get('/matched_gaia_stars')

            return cat_gaia_df

        if not _quadrant_good(quadrant_path):
            return

        try:
            refexp = None
            if cat1[:3] == 'ref':
                refexp = args.refexp
                cat1 = cat1[3:]

            quadrant_header = get_header_from_quadrant_path(quadrant_path)
            expid = int(quadrant_header['expid'])
            mjd = float(quadrant_header['obsmjd'])
            ccdid = int(quadrant_header['ccdid'])
            qid = int(quadrant_header['qid'])
            skylev = float(quadrant_header['sexsky'])
            moonillf = float(quadrant_header['moonillf'])
            seeing = float(quadrant_header['seeing'])
            airmass = float(quadrant_header['airmass'])

            year, month, day, _, _, ccdid, qid = quadrant_name_explode(quadrant_path.name)

            exposure_output_path = output_path.joinpath("{}-{}-{}-{}".format(year, month, day, expid))
            exposure_output_path.mkdir(exist_ok=True)

            # First load catalogs
            cat1_df = _get_cat_from_name(cat1, refexp=refexp)
            cat2_df = _get_cat_from_name(cat2)

            if refexp:
                cat1_gaia_df = _get_gaia_cat(refexp)
                cat2_gaia_df = _get_gaia_cat()
            else:
                cat1_gaia_df = _get_gaia_cat()
                cat2_gaia_df = cat1_gaia_df

            if _is_aperture(cat1):
                cat1_df['gaiaid'] = match_to_gaia(cat1_df, cat1_gaia_df)

            if _is_aperture(cat2):
                cat2_df['gaiaid'] = match_to_gaia(cat2_df, cat2_gaia_df)

            # Match catalogs to Gaia
            # i = match_pixel_space(cat_gaia_df[['x', 'y']].to_records(), cat2_df[['x', 'y']].to_records(), radius=1.)

            # cat2_df = cat2_df.iloc[i[i>=0]].reset_index(drop=True)
            # cat_gaia_df = cat_gaia_df.iloc[i[i>=0]].reset_index(drop=True)
            # cat1_df = cat1_df.iloc[i>=0].reset_index(drop=True)
            print(cat1_df)
            print(cat2_df)
            cat1_df, cat2_df, cat_gaia_df = match_by_gaia_catalogs(cat1_df, cat2_df, cat1_gaia_df, cat2_gaia_df)

            # Convert ADU flux into mag
            cat1_df['mag'] = -2.5*np.log10(cat1_df[cat1])
            cat2_df['mag'] = -2.5*np.log10(cat2_df[cat2])
            cat1_df['emag'] = -1.08*cat1_df['e'+cat1]/cat1_df[cat1]
            cat2_df['emag'] = -1.08*cat2_df['e'+cat2]/cat2_df[cat2]

            if _is_aperture(cat1):
                aper_size_1 = list(set(cat1_df['rad'+cat1[-1]]))[0]

            if _is_aperture(cat2):
                aper_size_2 = list(set(cat2_df['rad'+cat2[-1]]))[0]

            # Compute delta flux between PSF and aperture photometry
            deltamag = (cat1_df['mag'] - cat2_df['mag']).to_numpy()
            deltamag = deltamag - np.mean(deltamag)
            edeltamag = np.sqrt(cat1_df['emag']**2+cat2_df['emag']**2).to_numpy()

            # Bin delta flux
            plt.figure(figsize=(7., 4.))
            gmag_binned, deltamag_binned, edeltamag_binned = binplot(cat_gaia_df['gmag'].to_numpy(), deltamag, robust=True, data=False, scale=True, weights=1./edeltamag**2, bins=np.linspace(min_mag, max_mag, 8), color='black', zorder=15)

            bin_mask = (edeltamag_binned > 0.)
            gmag_binned = gmag_binned[bin_mask]
            deltamag_binned = np.array(deltamag_binned)[bin_mask]
            edeltamag_binned = edeltamag_binned[bin_mask]

            # Fit polynomials on delta mag bins
            poly0, ([poly0_chi2], _, _, _) = Polynomial.fit(gmag_binned, deltamag_binned, 0, w=1./edeltamag_binned, full=True)
            poly1, ([poly1_chi2], _, _, _) = Polynomial.fit(gmag_binned, deltamag_binned, 1, w=1./edeltamag_binned, full=True)
            poly2, ([poly2_chi2], _, _, _) = Polynomial.fit(gmag_binned, deltamag_binned, 2, w=1./edeltamag_binned, full=True)

            poly0_chi2 = poly0_chi2/(len(gmag_binned)-1)
            poly1_chi2 = poly1_chi2/(len(gmag_binned)-2)
            poly2_chi2 = poly2_chi2/(len(gmag_binned)-3)

            # Compute rough delta mag on delta mag bins
            low_bin_idx = np.argmin(np.abs(gmag_binned-low_bin))
            high_bin_idx = np.argmin(np.abs(gmag_binned-high_bin))
            low_bin_indices = list(set([max([0, low_bin_idx-1]), low_bin_idx, low_bin_idx+1]))
            high_bin_indices = list(set([high_bin_idx-1, high_bin_idx, min([high_bin_idx+1, len(gmag_binned)-1])]))
            bin_low = np.average(deltamag_binned[low_bin_indices], weights=1./edeltamag_binned[low_bin_indices])
            bin_high = np.average(deltamag_binned[high_bin_indices], weights=1./edeltamag_binned[high_bin_indices])
            deltamag_bin = bin_high-bin_low

            # Start plot
            name1 = _cat_to_name(cat1)
            name2 = _cat_to_name(cat2)

            if _is_aperture(cat1):
                name1 += " {:.2f} px".format(aper_size_1)

            if _is_aperture(cat2):
                name2 += " {:.2f} px".format(aper_size_2)

            plt.title("{} - {} [mag]\n({}-{}-{}) CCD={} Q={}\nSky level={:.2f}".format(name1, name2, year, month, day, ccdid, qid, skylev))
            plt.grid()
            gmag_space = np.linspace(min_mag, max_mag, 100)
            plt.plot(gmag_space, poly0(gmag_space), ls='-', lw=1.5, color='grey', label="Constant - $\chi^2_\\nu={:.3f}$".format(poly0_chi2), zorder=10)
            plt.plot(gmag_space, poly1(gmag_space), ls='dotted', lw=1.5, color='grey', label="Linear - $\chi^2_\\nu={:.3f}$".format(poly1_chi2), zorder=10)
            plt.plot(gmag_space, poly2(gmag_space), ls='--', lw=1.5, color='grey', label="Quadratic - $\chi^2_\\nu={:.3f}$".format(poly2_chi2), zorder=10)
            plt.errorbar(cat_gaia_df['gmag'].to_numpy(), deltamag, yerr=edeltamag, lw=0., marker='.', markersize=1.5, zorder=0, color='xkcd:light blue')
            plt.legend(fontsize='small', loc='upper left')
            plt.xlabel("$g_\mathrm{Gaia}$ [mag]")
            plt.ylabel("$\Delta_M$ [mag]")
            plt.text(14., -0.1, "$\Delta M_{{{}-{}}}={:.5f}$".format(low_bin, high_bin, deltamag_bin), bbox=dict(boxstyle="round", ec='grey', fc='xkcd:light blue'))
            plt.xlim(min_mag, max_mag)
            plt.ylim(-0.2, 0.2)

            plt.tight_layout()
            plt.savefig(exposure_output_path.joinpath(quadrant_path.name), dpi=200.)
            plt.close()

            print(".", end="", flush=True)

            return {'quadrant': quadrant_path.name,
                    'date': "{}-{}-{}".format(year, month, day),
                    'mjd': mjd,
                    'expid': expid,
                    'ccdid': ccdid,
                    'qid': qid,
                    'skylev': skylev,
                    'deltamag': deltamag_bin,
                    'poly0_chi2': poly0_chi2,
                    'poly1_chi2': poly1_chi2,
                    'poly2_chi2': poly2_chi2,
                    'seeing': seeing,
                    'airmass': airmass}

        except Exception as e:
            print(quadrant_path)
            print(e)
            print(traceback.format_exc())

    # d = [_process_quadrant(quadrant_path) for quadrant_path in quadrant_paths]

    tasks = [delayed(_process_quadrant)(quadrant_path, 'flux', 'apfl6') for quadrant_path in quadrant_paths]
    results = compute(tasks)

    # pairs = [('flux', 'apfl6'), ('refflux', 'flux')]
    pairs = [('refflux', 'flux'), ('flux', 'apfl6')]
    print("")
    print("Done.")
    results = list(filter(lambda x: x is not None, results[0]))
    df = pd.DataFrame(results)
    df.to_parquet(band_path.joinpath("myquadrants.parquet"))

    cmap = 'coolwarm'
    def _plot_focal_plane(exposure, vmin, vmax, val):
        cm = ScalarMappable(cmap=cmap)
        cm.set_clim(vmin=vmin, vmax=vmax)

        fig = plt.figure(figsize=(5., 6.), constrained_layout=False)
        f1, f2 = fig.subfigures(ncols=1, nrows=2, height_ratios=[12., 0.3], wspace=-1., hspace=-1.)
        f1.suptitle("\n{}-{}-{}\n$\Delta m$ [mag]".format(exposure[4:8], exposure[8:10], filtercode), fontsize='large')
        plot_ztf_focal_plan_values(f1, delta_mags[exposure], vmin=vmin, vmax=vmax, cmap=cmap, scalar=True)
        ax2 = f2.add_subplot(clip_on=False)
        f2.subplots_adjust(bottom=-5, top=5.)
        ax2.set_position((0.2, 0, 0.6, 1))
        cb = f1.colorbar(cm, cax=ax2, orientation='horizontal')
        ax2.set_clip_on(False)
        ax2.minorticks_on()
        ax2.tick_params(direction='inout')
        plt.savefig(output_path.joinpath("{}-{}_dmag_focal_plane.png".format(exposure, filtercode)), dpi=200., bbox_inches='tight', pad_inches=0.1)
        plt.close()

    fp_vals = {'PSF-Aper': {'limits': 0.2, 'val': lambda x: x.deltamag}}
    for exposure in exposures:
        _plot_focal_plane_dmag(exposure, -vext, vext)


import itertools
n=int(input())
my_dictionary = {"TOI": 26,"Hindu":20.5,"ET":34,"BM":10.5,"HT":18}
values = my_dictionary.values()
keys = my_dictionary.keys()
result = [seq for i in range(len(values), 0, -1)
              for seq in itertools.combinations(keys, 2)
              if sum(map(my_dictionary.get,seq)) <=n]
def unique(list1):
    unique_list = []
    for x in list1:
        if x not in unique_list:
            unique_list.append(x)
    # print list
    for x in unique_list:
        print(x)
unique(result)


#!/usr/bin/env python

# peak detection: https://qiita.com/wrblue_mica34/items/e174a71570abb710dcfb

import rospy
from std_msgs.msg import *
from geometry_msgs.msg import *
from scipy import signal
import numpy as np

sensitivity_ = 50

def callback(msg):
  global sensitivity_

  maxid_list = []

  y = np.array(msg.data)

  maxid = signal.argrelmax(y, order=1) #最大値
  minid = signal.argrelmin(y, order=1) #最小値

  maxid = maxid[0].tolist()
  minid = minid[0].tolist()

  for i in maxid:
    if y[i] > sensitivity_*255/100:
      maxid_list.append(i)

  max_id_msg = Int32MultiArray()
  min_id_msg = Int32MultiArray()

  max_id_msg.data = maxid_list
  min_id_msg.data = minid

  max_pub.publish(max_id_msg)
  min_pub.publish(min_id_msg)

def sensitivity_callback(msg):
#  print(msg)
  global sensitivity_
  sensitivity_ = msg.data

rospy.init_node("peak_detection")
max_pub = rospy.Publisher("~peak_max_index", Int32MultiArray, queue_size=10)
min_pub = rospy.Publisher("~peak_min_index", Int32MultiArray, queue_size=10)
sensitivity_state_pub = rospy.Publisher("/peak_detection/sensitivity_state", Float64, queue_size=10)
rospy.Subscriber("/edge_detection/edge_graph", Float32MultiArray, callback)
rospy.Subscriber("/edge_detection/edge_sensitivity", Float64, sensitivity_callback)

r = rospy.Rate(10)

while not rospy.is_shutdown():
  sensitivity_state_pub.publish(Float64(sensitivity_))
  print(sensitivity_)
  r.sleep()


# patricia.angie.kwarwp-mapa-indio.py
# SPDX-License-Identifier: GPL-3.0-or-later
""" Projeto Kwarwp

.. codeauthor:: Monica Novellino <monicanovellino@gmail.com>

Changelog
---------
.. versionadded::    20.07
        Montando o jogo usando mapa e o índio

"""
MAPA_INICIO = """
@....&
......
.....#
.#.^..
"""

MAPA_CERCA = """
++++++
.....&
......
...^..
++++++
"""
MAPA_ROCHA = """
+++++++
+..+..&
+..%..+
+^.+..+
+++++++
"""

class Indio():

    def __init__(self, imagem, x, y, cena):
        self.lado = lado = Kwarwp.LADO
        self.posicao = (x//lado, y//lado)  # XXX[3]XXX faltou definir posição
        self.indio = Kwarwp.VITOLLINO.a(imagem, w=lado, h=lado, x=x, y=y, cena=cena)

    def anda(self):
        """ Faz o índio caminhar na direção em que está olhando.
        """
        self.posicao = (self.posicao[0], self.posicao[1]-1)
        """Assumimos que o índio está olhando para cima, decrementamos a posição **y**"""
        self.indio.y = self.posicao[1]*self.lado
        self.indio.x = self.posicao[0]*self.lado

    def executa(self):
        """ Roteiro do índio. Conjunto de comandos para ele executar.
        """
        self.anda()

class Kwarwp():
    """ Jogo para ensino de programação.

        :param vitollino: Empacota o engenho de jogo Vitollino.
    """
    VITOLLINO = None
    """Referência estática para obter o engenho de jogo."""
    LADO = None
    """Referência estática para definir o lado do piso da casa."""

    #self.o_indio = None XXXX -> etava no lugar errado, mandei para XXX[1]XXXX
    """Instância do personagem principal, o índio, vai ser atribuído pela fábrica do índio"""

    def __init__(self, vitollino=None, mapa=MAPA_INICIO, medidas={}):
        Kwarwp.VITOLLINO = self.v = vitollino()
        """Cria um matriz com os elementos descritos em cada linha de texto"""
        self.mapa = mapa.split()
        """Largura da casa da arena dos desafios, número de colunas no mapa"""
        self.lado, self.col, self.lin = 100, len(self.mapa[0]), len(self.mapa)+1
        Kwarwp.LADO = self.lado
        #XXX[1]XXXX ---> aqui que tinha que estar

        self.o_indio = None
        """Instância do personagem principal, o índio, vai ser atribuído pela fábrica do índio"""
        w, h = self.col *self.lado, self.lin *self.lado
        self.taba = {}
        """Dicionário que a partir de coordenada (i,J) localiza um piso da taba"""
        medidas.update(width=w, height=f"{h}px")
        self.cena = self.cria(mapa=self.mapa) if vitollino else None

    def cria(self, mapa="  "):

        from collections import namedtuple as nt
        Fab = nt("Fab", "objeto imagem")

        IMGUR = "https://imgur.com/"
        IIMGUR = "https://i.imgur.com/"
        fabrica = {
        "&": Fab(self.coisa, f"{IIMGUR}dZQ8liT.jpg"), # OCA
        "^": Fab(self.indio, f"{IMGUR}8jMuupz.png"), # INDIO
        ".": Fab(self.coisa, f"{IIMGUR}npb9Oej.png"), # VAZIO
        "_": Fab(self.coisa, f"{IIMGUR}sGoKfvs.jpg"), # SOLO
        "#": Fab(self.coisa, f"{IMGUR}ldI7IbK.png"), # TORA
        "@": Fab(self.coisa, f"{IIMGUR}tLLVjfN.png"), # PICHE
        "~": Fab(self.coisa, f"{IMGUR}UAETaiP.gif"), # CEU
        "*": Fab(self.coisa, f"{IMGUR}PfodQmT.gif"), # SOL
        "+": Fab(self.coisa, f"{IMGUR}uwYPNlz.png"),  # CERCA
        "%": Fab(self.coisa, f"{IMGUR}Ry3Vmsn.png") #ROCHA
        }

        mapa = mapa if mapa != "" else self.mapa

        mapa = self.mapa
        lado = self.lado
        cena = self.v.c(fabrica["_"].imagem)
        ceu = self.v.a(fabrica["~"].imagem, w=lado*self.col, h=lado, x=0, y=0, cena=cena)
        sol = self.v.a(fabrica["*"].imagem, w=60, h=60, x=0, y=40, cena=cena)

        self.taba = {(i, j): fabrica[imagem].objeto(
            fabrica[imagem].imagem, x=i*lado, y=j*lado+lado, cena=cena)
            for j, linha in enumerate(mapa) for i, imagem in enumerate(linha)}

        ceu = self.v.a(fabrica["~"].imagem, w=lado*self.col, h=lado, x=0, y=0, cena=cena, vai= self.executa)
        """No argumento *vai*, associamos o clique no céu com o método **executa ()** desta classe"""

        cena.vai()
        return cena

    def coisa(self, imagem, x, y, cena):
        lado = self.lado
        return self.v.a(imagem, w=lado, h=lado, x=x, y=y, cena=cena)

    def indio(self, imagem, x, y, cena):
        self.o_indio = Indio(imagem, x=x, y=y, cena=cena)
        return self.o_indio

    def executa(self, *_):
        """ Ordena a execução do roteiro do índio.
        """
        self.o_indio.executa()

if __name__ == "__main__":
    from _spy.vitollino.main import Jogo, STYLE
    #XXX[2]XXX faltou estas duas linhas do STYLE para esticar o solo
    STYLE["width"] = 600
    STYLE["heigth"] = "500px"
    Kwarwp(Jogo, mapa=MAPA_INICIO)


import pandas as pd
from openpyxl import load_workbook

def write_to_excel(df, name):
	book = load_workbook('stock.xlsx')
	if name in book.sheetnames:
		book.remove(book[name])
		book.save('stock.xlsx')
	book = load_workbook('stock.xlsx')
	with pd.ExcelWriter('stock.xlsx') as writer:
		writer.book = book
		writer.sheets = dict((ws.title, ws) for ws in book.worksheets)
		df.to_excel(writer, sheet_name=name)
		writer.save()

def generate():
	xls = pd.ExcelFile('stock.xlsx')
	xls_split = pd.ExcelFile('stock_split.xlsx')
	df1 = pd.read_excel(xls, 'Stock')
	df2 = pd.read_excel(xls_split, 'Split')

	#Sym, Quanitiy Available, Avg price, Invested Amt, Current Amt, Abs return, % of Invested, % of current, % of change
	df1['Invested Amt'] = df1['Quantity Available'] * df1['Average Price']
	df1['Current Amt'] = df1['Quantity Available'] * df1['Previous Closing Price']
	df1['Abs return'] = round(((df1['Current Amt'] - df1['Invested Amt']) / df1['Invested Amt'])*100, 2)

	total_invested_price = df1['Invested Amt'].sum()
	total_current_price = df1['Current Amt'].sum()

	stock = {};
	c_name =''
	for (column_name, column_data) in df2.iteritems():
		if c_name != "" and c_name.lower() + ' percentage' == column_name.lower():
			stock[c_name] = dict(zip(stock[c_name], column_data.dropna().to_list()))
		else:
			stock[column_name] = column_data.dropna().to_list()
			c_name = column_name

	data=[]
	for (group, values) in stock.items():
		dfs = df1[df1['Symbol'].isin(values.keys())].copy()
		for key,val in values.items():
			if val < 100:
				d = dfs[ dfs['Symbol'] == key]
				d['Quantity Available'] = d['Quantity Available'] * (val/100)
				dfs[dfs['Symbol'] == key] = d

		dfs['Invested Amt'] = dfs['Quantity Available'] * dfs['Average Price']
		dfs['Current Amt'] = dfs['Quantity Available'] * dfs['Previous Closing Price']
		dfs['Abs return'] = round(((dfs['Current Amt'] - dfs['Invested Amt']) / dfs['Invested Amt']) * 100, 2)
		total_invested_price_dfs = dfs['Invested Amt'].sum()
		total_current_price_dfs = dfs['Current Amt'].sum()
		dfs['% of invested'] = round((dfs['Invested Amt']/total_invested_price_dfs)*100, 2)
		dfs['% of current'] = round((dfs['Current Amt']/total_current_price_dfs)*100, 2)
		dfs['% change'] = dfs['% of current'] - dfs['% of invested']
		data.append([group, total_invested_price_dfs, total_current_price_dfs])
		df = dfs[['Symbol', 'Quantity Available', 'Average Price', 'Invested Amt', 'Current Amt', 'Abs return', '% of invested', '% of current', '% change' ]]
		write_to_excel(df, group)

	df_tmp = pd.DataFrame(data, columns = ['Name', 'Invested Amt', 'Current Amt'])
	df_tmp['Abs return'] = round(((df_tmp['Current Amt'] - df_tmp['Invested Amt']) / df_tmp['Invested Amt']) * 100, 2)
	df_tmp['% of invested Amt'] = round((df_tmp['Invested Amt']/total_invested_price)*100, 2)
	df_tmp['% of current Amt'] = round((df_tmp['Current Amt']/total_current_price)*100, 2)
	df_tmp['% change'] = df_tmp['% of current Amt'] - df_tmp['% of invested Amt']
	write_to_excel(df_tmp, 'Output')

if __name__=="__main__":
	generate()


import time
from http import HTTPStatus

import uweb3

import invoices.login.model as login_model
from invoices.invoice.model import PRO_FORMA_PREFIX

API_VERSION = "/api/v1"

class PageMaker(
    uweb3.DebuggingPageMaker,
    uweb3.LoginMixin,
):
    """Holds all the request handlers for the application"""

    def __init__(self, *args, **kwds):
        super(PageMaker, self).__init__(*args, **kwds)

    def _PostInit(self):
        """Sets up all the default vars"""
        self.validatexsrf()
        self.parser.RegisterFunction("CentRound", lambda x: "%.2f" % x if x else None)
        self.parser.RegisterFunction("items", lambda x: x.items())
        self.parser.RegisterFunction("DateOnly", lambda x: str(x)[0:10])
        self.parser.RegisterFunction(
            "isProForma", lambda x: bool(str(x).startswith(PRO_FORMA_PREFIX))
        )
        self.parser.RegisterTag("year", time.strftime("%Y"))
        self.parser.RegisterTag(
            "header", self.parser.JITTag(lambda: self.parser.Parse("parts/header.html"))
        )
        self.parser.RegisterTag(
            "footer",
            self.parser.JITTag(
                lambda *args, **kwargs: self.parser.Parse(
                    "parts/footer.html", *args, **kwargs
                )
            ),
        )
        self.parser.RegisterTag("xsrf", self._Get_XSRF())
        self.parser.RegisterTag("user", self.user)

    def _PostRequest(self, response):
        response.headers.update(
            {
                "Access-Control-Allow-Origin": "*",
            }
        )
        return response

    def _ReadSession(self):
        """Attempts to read the session for this user from his session cookie"""
        try:
            user = login_model.Session(self.connection)
        except Exception:
            raise ValueError("Session cookie invalid")
        try:
            user = login_model.User.FromPrimary(self.connection, int(str(user)))
        except uweb3.model.NotExistError:
            return None
        if user["active"] != "true":
            raise ValueError("User not active, session invalid")
        return user

    def RequestInvalidcommand(self, command=None, error=None, httpcode=404):
        """Returns an error message"""
        uweb3.logging.warning(
            "Bad page %r requested with method %s", command, self.req.method
        )
        if command is None and error is None:
            command = "%s for method %s" % (self.req.path, self.req.method)
        page_data = self.parser.Parse("404.html", command=command, error=error)
        return uweb3.Response(content=page_data, httpcode=httpcode)

    @uweb3.decorators.ContentType("application/json")
    def FourOhFour(self, path):
        """The request could not be fulfilled, this returns a 404."""
        return uweb3.Response(
            {
                "error": True,
                "errors": ["Requested page not found"],
                "http_status": HTTPStatus.NOT_FOUND,
            },
            httpcode=HTTPStatus.NOT_FOUND,
        )

    def Error(self, error="", httpcode=500, link=None):
        """Returns a generic error page based on the given parameters."""
        uweb3.logging.error("Error page triggered: %r", error)
        page_data = self.parser.Parse("error.html", error=error, link=link)
        return uweb3.Response(content=page_data, httpcode=httpcode)


"""Tests for cell_img.image_grid.ts_index_lib."""

import copy

from absl.testing import absltest
from cell_img.image_grid import ts_index_lib
import numpy as np
import pandas as pd

class Single1DTest(absltest.TestCase):

  def test_get_index_individually(self):
    lookup = ts_index_lib._Single1D(
        ['zero', 'one', 'two', 'three', 'four', 'five'])
    self.assertEqual(lookup.get_index('five'), 5)
    self.assertEqual(lookup.get_index('one'), 1)
    self.assertEqual(lookup.get_index('four'), 4)
    self.assertEqual(lookup.get_index('three'), 3)
    self.assertEqual(lookup.get_index('two'), 2)
    self.assertEqual(lookup.get_index('zero'), 0)

  def test_get_index_multiple(self):
    lookup = ts_index_lib._Single1D(
        ['zero', 'one', 'two', 'three', 'four', 'five'])
    result = lookup.get_index(['two', 'five'])
    np.testing.assert_equal(result, [2, 5])

    result = lookup.get_index(np.array(['two', 'five']))
    np.testing.assert_equal(result, [2, 5])

    result = lookup.get_index(pd.Series(['two', 'five']))
    np.testing.assert_equal(result, [2, 5])

  def test_get_index_non_existent_raises(self):
    lookup = ts_index_lib._Single1D(
        ['zero', 'one', 'two', 'three', 'four', 'five'])
    with self.assertRaisesRegex(
        ValueError, 'Not all values of the query are in the index.'):
      lookup.get_index('six')
    with self.assertRaisesRegex(
        ValueError, 'Not all values of the query are in the index.'):
      lookup.get_index(['two', 'five', 'six'])

  def test_get_value_individually(self):
    lookup = ts_index_lib._Single1D(
        ['zero', 'one', 'two', 'three', 'four', 'five'])
    self.assertEqual(lookup.get_value(5), 'five')

  def test_get_value_multiple(self):
    lookup = ts_index_lib._Single1D(
        ['zero', 'one', 'two', 'three', 'four', 'five'])
    np.testing.assert_equal(lookup.get_value([2, 5]), ['two', 'five'])

  def test_bad_dimensions_raises(self):
    with self.assertRaisesRegex(ValueError, 'Expected 1D input but got.*'):
      ts_index_lib._Single1D([['A1', 'A2'], ['B1', 'B2']])

  def test_empty_raises(self):
    with self.assertRaisesRegex(ValueError, 'Input cannot be empty.'):
      ts_index_lib._Single1D([])

  def test_dim_length_one_works(self):
    lookup = ts_index_lib._Single1D(['one_value'])
    self.assertEqual(lookup.get_index('one_value'), 0)
    self.assertEqual(lookup.get_value(0), 'one_value')
    np.testing.assert_equal(
        lookup.get_index(['one_value', 'one_value']), [0, 0])
    np.testing.assert_equal(
        lookup.get_value([0, 0]), ['one_value', 'one_value'])

class Multi1DTest(absltest.TestCase):

  def get_three_level_lookup(self):
    coordinate_arrays = {
        'plate': ['plate0', 'plate1', 'plate2'],
        'well_col': ['1', '2', '3', '4', '5'],
        'site_col': ['site_col0', 'site_col1'],
    }
    coordinate_names = ['plate', 'well_col', 'site_col']

    lookup = ts_index_lib._Multi1D(coordinate_names, coordinate_arrays)
    return lookup

  def test_valid_missing_coords(self):
    valid_missing_coord_names = [['site_col',], ['site_col', 'well_col']]
    lookup = self.get_three_level_lookup()
    for coord_names in valid_missing_coord_names:
      # does not raise
      _ = lookup._check_valid_missing_coords(coord_names)

  def test_raises_when_not_smallest_coord(self):
    lookup = self.get_three_level_lookup()
    invalid_missing_coord_names = [['well_col',], ['plate']]
    for coord_names in invalid_missing_coord_names:
      with self.assertRaises(KeyError):
        lookup._check_valid_missing_coords(coord_names)

  def test_raises_when_all_missing(self):
    lookup = self.get_three_level_lookup()
    invalid_missing_coords = ['site_col', 'well_col', 'plate']
    with self.assertRaises(KeyError):
      lookup._check_valid_missing_coords(invalid_missing_coords)

  def test_raises_when_not_consecutive(self):
    lookup = self.get_three_level_lookup()
    invalid_missing_coords = ['site_col', 'plate']
    with self.assertRaises(KeyError):
      lookup._check_valid_missing_coords(invalid_missing_coords)

  def test_missing_coord_names_site(self):
    lookup = self.get_three_level_lookup()
    self.assertEqual(
        lookup._get_missing_coords({
            'plate': ['value1',],
            'well_col': ['value1',]
        }), ['site_col',])

  def test_missing_coord_names_site_well(self):
    lookup = self.get_three_level_lookup()
    self.assertEqual(
        lookup._get_missing_coords({
            'plate': ['value1',],
        }), ['well_col', 'site_col'])

  def test_get_well_level_smallest_block_size(self):
    lookup = self.get_three_level_lookup()
    well_level_labels = {'plate': ['plate0',], 'well_col': ['1',]}
    expected_block_size = 2  # 2 sites per well
    self.assertEqual(
        lookup.get_smallest_block_size(well_level_labels),
        expected_block_size)

  def test_get_plate_level_smallest_block_size(self):
    lookup = self.get_three_level_lookup()
    plate_level_labels = {'plate': ['plate0',]}
    expected_block_size = 2 * 5  # 2 sites per well, 5 wells per plate
    self.assertEqual(
        lookup.get_smallest_block_size(plate_level_labels),
        expected_block_size)

  def test_get_index_two_level(self):
    names = ['year', 'month']
    name_to_values = {
        'year': ['2019', '2020', '2021'],
        'month': [
            'jan', 'feb', 'mar', 'apr', 'may', 'jun', 'jul', 'aug', 'sep',
            'oct', 'nov', 'dec'
        ]
    }
    lookup = ts_index_lib._Multi1D(names, name_to_values)
    self.assertEqual(lookup.get_index({'year': '2019', 'month': 'jan'}), 0)
    self.assertEqual(lookup.get_index({'year': '2019', 'month': 'aug'}), 7)
    self.assertEqual(lookup.get_index({'year': '2020', 'month': 'aug'}), 19)
    self.assertEqual(lookup.get_index({'year': '2021', 'month': 'aug'}), 31)
    self.assertEqual(lookup.get_index({'year': '2021', 'month': 'dec'}), 35)

    np.testing.assert_equal(
        lookup.get_index({
            'year': ['2019', '2019', '2020', '2021', '2021'],
            'month': ['jan', 'aug', 'aug', 'aug', 'dec']
        }), [0, 7, 19, 31, 35])

  def test_get_index_non_existent_raises(self):
    names = ['year', 'month']
    name_to_values = {
        'year': ['2019', '2020', '2021'],
        'month': [
            'jan', 'feb', 'mar', 'apr', 'may', 'jun', 'jul', 'aug', 'sep',
            'oct', 'nov', 'dec'
        ]
    }
    lookup = ts_index_lib._Multi1D(names, name_to_values)
    with self.assertRaisesRegex(
        ValueError, 'Not all values of the query are in the index.'):
      lookup.get_index({'year': '2021', 'month': 'NON_EXISTENT'})
    with self.assertRaisesRegex(
        ValueError, 'Not all values of the query are in the index.'):
      lookup.get_index({
          'year': ['2020', '2021'],
          'month': ['feb', 'NON_EXISTENT']
      })

  def test_get_index_individually(self):
    names = ['x', 'y', 'z']
    name_to_values = {
        'x': ['x0', 'x1', 'x2'],
        'y': ['y0', 'y1', 'y2', 'y3'],
        'z': ['z0', 'z1']
    }
    lookup = ts_index_lib._Multi1D(names, name_to_values)
    self.assertEqual(lookup.get_index({'x': 'x0', 'y': 'y0', 'z': 'z0'}), 0)
    self.assertEqual(lookup.get_index({'x': 'x1', 'y': 'y2', 'z': 'z1'}), 13)
    self.assertEqual(lookup.get_index({'x': 'x2', 'y': 'y3', 'z': 'z1'}), 23)

  def test_get_index_multiple(self):
    names = ['x', 'y', 'z']
    name_to_values = {
        'x': ['x0', 'x1', 'x2'],
        'y': ['y0', 'y1', 'y2', 'y3'],
        'z': ['z0', 'z1']
    }
    lookup = ts_index_lib._Multi1D(names, name_to_values)
    results = lookup.get_index({
        'x': ['x0', 'x1'],
        'y': ['y0', 'y2'],
        'z': ['z0', 'z1']
    })
    np.testing.assert_equal(results, [0, 13])

  def test_get_index_exhaustive(self):
    names = ['x', 'y', 'z']
    name_to_values = {
        'x': ['x0', 'x1', 'x2'],
        'y': ['y0', 'y1', 'y2', 'y3'],
        'z': ['z0', 'z1']
    }
    lookup = ts_index_lib._Multi1D(names, name_to_values)

    i = 0
    for x in name_to_values['x']:
      for y in name_to_values['y']:
        for z in name_to_values['z']:
          self.assertEqual(lookup.get_index({'x': x, 'y': y, 'z': z}), i)
          i += 1

class TensorstoreIndexTest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    coordinate_arrays = {
        'batch': ['batch0', 'batch1'],
        'plate': ['plate0', 'plate1', 'plate2'],
        'well_row': ['A', 'B', 'C', 'D'],
        'well_col': ['1', '2', '3', '4', '5'],
        'site_row': ['site_row0', 'site_row1'],
        'site_col': ['site_col0', 'site_col1'],
        'stain': ['stain0', 'stain1']
    }
    axes = ['X', 'Y', 'stain']
    axes_wrap = {
        'X': ['plate', 'well_col', 'site_col'],
        'Y': ['batch', 'well_row', 'site_row']
    }
    image_shape = (480, 640)
    unit_sizes = {'X': image_shape[1], 'Y': image_shape[0]}

    self.lookup = ts_index_lib.TensorstoreIndex(axes, axes_wrap,
                                                coordinate_arrays, unit_sizes)

  def test_get_dimensions(self):
    self.assertEqual(self.lookup.get_dimensions(), [19200, 7680, 2])

  def test_get_coords_individually(self):
    query = {
        'batch': 'batch0',
        'plate': 'plate0',
        'well_row': 'A',
        'well_col': '1',
        'site_row': 'site_row0',
        'site_col': 'site_col0',
        'stain': 'stain0',
    }
    self.assertEqual(self.lookup.get_coords(query), [0, 0, 0])
    # Move along the stain coordinate by one.
    query['stain'] = 'stain1'
    self.assertEqual(self.lookup.get_coords(query), [0, 0, 1])
    # Move along the Y coordinate by one image by going to the second site_row.
    query['site_row'] = 'site_row1'
    self.assertEqual(self.lookup.get_coords(query), [0, 480, 1])
    # Move along the X coordinate by one image by going to the second site_col.
    query['site_col'] = 'site_col1'
    self.assertEqual(self.lookup.get_coords(query), [640, 480, 1])

  def test_get_coords_raises_when_missing_key(self):
    valid_query = {
        'batch': 'batch0',
        'plate': 'plate0',
        'well_row': 'A',
        'well_col': '1',
        'site_row': 'site_row0',
        'site_col': 'site_col0',
        'stain': 'stain0',
    }
    # valid_query should not raise.
    _ = self.lookup.get_coords(valid_query)
    # Assert that it raises when the query is missing a key
    invalid_missing_keys = ['batch', 'plate', 'well_row', 'well_col', 'stain']
    for k in invalid_missing_keys:
      query = copy.deepcopy(valid_query)
      del query[k]
      with self.assertRaises(KeyError):
        self.lookup.get_coords(query)

  def test_get_coords_valid_when_missing_key(self):
    valid_query = {
        'batch': 'batch0',
        'plate': 'plate0',
        'well_row': 'B',
        'well_col': '2',
        'site_row': 'site_row1',
        'site_col': 'site_col1',
        'stain': 'stain0',
    }
    # valid_query should not raise.
    valid_coords = np.array([640*3, 480*3, 0])
    self.assertEqual(self.lookup.get_coords(valid_query), list(valid_coords))
    # Assert that it doesn't raise when missing smallest key(s)
    # drop site_row
    query = copy.deepcopy(valid_query)
    del query['site_row']
    expected_coords = valid_coords - np.array([0, 480, 0])
    self.assertEqual(self.lookup.get_coords(query), list(expected_coords))

    # drop [well_col, site_ col]
    query = copy.deepcopy(valid_query)
    del query['site_col']
    del query['well_col']
    expected_coords = valid_coords - np.array([640*3, 0, 0])
    self.assertEqual(self.lookup.get_coords(query), list(expected_coords))

  def test_get_block_sizes(self):
    valid_query = {
        'batch': 'batch0',
        'plate': 'plate0',
        'well_row': 'B',
        'well_col': '2',
        'site_row': 'site_row1',
        'site_col': 'site_col1',
        'stain': 'stain0',
    }
    # All dimensions are present, size should be 1 * unit_sizes
    self.assertEqual(self.lookup.get_block_sizes(valid_query), [640, 480, 1])

  def test_get_block_sizes_missing_site(self):
    valid_query = {
        'batch': 'batch0',
        'plate': 'plate0',
        'well_row': 'B',
        'well_col': '2',
        'stain': 'stain0',
    }
    # Missing sites, so sizes should be num_sites_per_well * unit_sizes
    self.assertEqual(
        self.lookup.get_block_sizes(valid_query), [2 * 640, 2 * 480, 1])

  def test_get_block_sizes_missing_site_well(self):
    valid_query = {
        'batch': 'batch0',
        'plate': 'plate0',
        'stain': 'stain0',
    }
    # Missing sites and wells, so sizes should be
    # num_wells_per_plate * num_sites_per_well * unit_sizes
    self.assertEqual(
        self.lookup.get_block_sizes(valid_query), [10 * 640, 8 * 480, 1])

  def test_get_coords_with_dataframe(self):
    query_df = pd.DataFrame.from_dict({
        'batch': ['batch0', 'batch0'],
        'plate': ['plate0', 'plate0'],
        'well_row': ['A', 'A'],
        'well_col': ['1', '1'],
        'site_row': ['site_row0', 'site_row1'],
        'site_col': ['site_col0', 'site_col1'],
        'stain': ['stain0', 'stain1'],
    })
    coords_dict = self.lookup.get_coords_dict(query_df)
    self.assertSameElements(coords_dict.keys(), ['X', 'Y', 'stain'])
    np.testing.assert_equal(coords_dict['X'], [0, 640])
    np.testing.assert_equal(coords_dict['Y'], [0, 480])
    np.testing.assert_equal(coords_dict['stain'], [0, 1])

  def test_get_whole_image_slice(self):
    valid_query = {
        'batch': 'batch0',
        'plate': 'plate0',
        'well_row': 'A',
        'well_col': '1',
        'site_row': 'site_row0',
        'site_col': 'site_col0',
        'stain': 'stain0',
    }
    whole_img = self.lookup.get_whole_image_slice(valid_query)
    np.testing.assert_equal(whole_img[0], slice(0, 640, None))
    np.testing.assert_equal(whole_img[1], slice(0, 480, None))
    np.testing.assert_equal(whole_img[2], 0)

  def test_get_whole_image_slice_fails_on_multiple_labels(self):
    invalid_query = {
        'batch': ['batch0', 'batch0'],
        'plate': ['plate0', 'plate0'],
        'well_row': ['A', 'A'],
        'well_col': ['1', '1'],
        'site_row': ['site_row0', 'site_row1'],
        'site_col': ['site_col0', 'site_col1'],
        'stain': ['stain0', 'stain1'],
    }
    with self.assertRaises(TypeError):
      _ = self.lookup.get_whole_image_slice(invalid_query)

  def test_get_square_patch_slice(self):
    valid_query = {
        'batch': 'batch0',
        'plate': 'plate0',
        'well_row': 'A',
        'well_col': '1',
        'site_row': 'site_row0',
        'site_col': 'site_col0',
        'stain': 'stain0',
    }
    offsets = {'X': 16, 'Y': 16}
    patch_size = 32
    slice_img = self.lookup.get_square_patch_slice(
        valid_query, offsets, patch_size)
    np.testing.assert_equal(slice_img[0], slice(0, 32, None))
    np.testing.assert_equal(slice_img[1], slice(0, 32, None))
    np.testing.assert_equal(slice_img[2], 0)

  def test_get_square_patch_slice_fails_offset_not_used(self):
    valid_query = {
        'batch': 'batch0',
        'plate': 'plate0',
        'well_row': 'A',
        'well_col': '1',
        'site_row': 'site_row0',
        'site_col': 'site_col0',
        'stain': 'stain0',
    }
    # There is no 'Z' axis
    offsets = {'X': 16, 'Y': 16, 'Z': 16}
    patch_size = 32
    with self.assertRaises(ValueError):
      self.lookup.get_square_patch_slice(valid_query, offsets, patch_size)

  def test_get_square_patch_slice_fails_across_whole_images(self):
    valid_query = {
        'batch': 'batch0',
        'plate': 'plate0',
        'well_row': 'A',
        'well_col': '1',
        'site_row': 'site_row0',
        'site_col': 'site_col0',
        'stain': 'stain0',
    }
    # Make a patch before the first site
    offsets = {'X': 0, 'Y': 0}
    patch_size = 32
    with self.assertRaises(ValueError):
      self.lookup.get_square_patch_slice(valid_query, offsets, patch_size)
    # Make a patch across the first and second site
    offsets = {'X': 630, 'Y': 0}
    patch_size = 32
    with self.assertRaises(ValueError):
      self.lookup.get_square_patch_slice(valid_query, offsets, patch_size)
    # Make a patch across the first and second site in the Y direction
    offsets = {'X': 0, 'Y': 450}
    patch_size = 32
    with self.assertRaises(ValueError):
      self.lookup.get_square_patch_slice(valid_query, offsets, patch_size)
    # No exception if we turn off the require_within_single_image flag.
    offsets = {'X': 630, 'Y': 0}
    patch_size = 32
    self.lookup.get_square_patch_slice(valid_query, offsets, patch_size,
                                       require_within_single_image=False)

if __name__ == '__main__':
  absltest.main()