""" gnn_module.py

Hold basic GNN Types:
1. GGNN
2. NNConv

These classes should accept graphs and return featurizations at each node

The calling class should be responsible for pooling however is best
"""
import numpy as np

import torch
import torch.nn as nn

import dgl.function as fn
from dgllife.model import MPNNGNN
from dgl.backend import pytorch as dgl_F

class MoleculeGNN(nn.Module):
    """MoleculeGNN Module"""

    def __init__(
            self,
            hidden_size: int,
            num_step_message_passing: int = 4,
            gnn_node_feats: int = 74,
            gnn_edge_feats: int = 4,  # 12,
            mpnn_type: str = "NNConv",
            node_feat_symbol="h",
            set_transform_layers: int = 2,
            **kwargs):
        """__init__.
        Args:
            hidden_size (int): Hidden size
            num_mol_layers (int): Number of layers to encode for the molecule
        """
        super().__init__()
        self.hidden_size = hidden_size
        self.gnn_edge_feats = gnn_edge_feats
        self.gnn_node_feats = gnn_node_feats
        self.node_feat_symbol = node_feat_symbol

        self.mpnn_type = mpnn_type
        node_out_feats = self.hidden_size
        if self.mpnn_type == "NNConv":
            # https://github.com/awslabs/dgl-lifesci/blob/master/examples/property_prediction/csv_data_configuration/
            # Note: Edge hidden feats seems way too bloated in this model which
            # is driving params way up when node_out_feats is increased
            self.gnn = MPNNGNN(
                node_in_feats=self.gnn_node_feats,
                edge_in_feats=self.gnn_edge_feats,
                node_out_feats=self.hidden_size,
                edge_hidden_feats=self.hidden_size // 4,  #32,
                num_step_message_passing=num_step_message_passing,
            )
        elif self.mpnn_type == "GGNN":
            self.gnn = GGNN(
                hidden_size=self.hidden_size,
                edge_feats=self.gnn_edge_feats,
                node_feats=self.gnn_node_feats,
                num_step_message_passing=num_step_message_passing,
            )
        else:
            raise ValueError()

        # Keeping d_head only to 2x increase in size to avoid memory. Orig
        # transformer uses 4x
        self.set_transformer = SetTransformerEncoder(
            d_model=node_out_feats,
            n_heads=4,
            d_head=node_out_feats // 4,
            d_ff=hidden_size,
            n_layers=set_transform_layers,
        )

    def forward(self, g):
        """encode batch of molecule graph"""
        node_feats = g.ndata[self.node_feat_symbol]
        device = g.device
        if self.mpnn_type == "NNConv":
            edge_feats = g.edata.get(
                "e", torch.empty(0, self.gnn_edge_feats, device=device))
            output = self.gnn(g, node_feats, edge_feats)
        elif self.mpnn_type == "GGNN":
            # Get graph output
            output = self.gnn(g, self.node_feat_symbol, "e")
        else:
            raise NotImplementedError()

        output = self.set_transformer(g, output)
        return output

class GGNN(nn.Module):

    def __init__(self,
                 hidden_size=64,
                 edge_feats=4,
                 node_feats=74,
                 num_step_message_passing=4,
                 **kwargs):
        """GGNN.

        Define a gated graph neural network

        This is very similar to the NNConv models.

        Args:
            input_size (int): Size of edge features into the graph
            hidden_size (int): Hidden size
            edge_feats (int): Number of edge feats. Must be onehot!
            node_feats (int): Num of node feats (default 74)
            num_step_message_passing (int): Number of message passing steps
        """
        super().__init__()
        self.hidden_size = hidden_size
        self.num_step_message_passing = num_step_message_passing
        self.edge_feats = edge_feats
        self.node_feats = node_feats
        self.input_project = nn.Linear(self.node_feats, self.hidden_size)

        # Define LSTM Cell
        self.gru = nn.GRUCell(self.hidden_size, self.hidden_size)
        self.edge_transform_weights = torch.nn.Parameter(
            torch.randn(self.edge_feats, self.hidden_size, self.hidden_size))
        self.edge_transform_bias = torch.nn.Parameter(
            torch.randn(self.edge_feats, 1))

        ### Compute Output transformation
        self.gate_mlp = nn.Sequential(
            nn.Linear(self.hidden_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, 1),
            nn.Sigmoid(),
        )
        self.tan_mlp = nn.Sequential(
            nn.Linear(self.hidden_size * 2, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.hidden_size),
            nn.Tanh(),
        )

    def message_pass(self, edges):
        """Pass messages along the edges"""
        src_feat = edges.src["_h"]

        # Linear layer
        messages = (torch.einsum("nio,ni->no", edges.data["w"], src_feat) +
                    edges.data["b"])
        messages = nn.functional.relu(messages)
        return {"m": messages}

    def forward(self, graph, nfeat_name="h", efeat_name="e"):
        """forward.

        Args:
            graph (dgl graph): Graph object
            nfeat_name (str): Name of node feat data
            efeat_name (str): Name of e feat

        Return:
            h: Hidden state at each node

        """
        ndata = graph.ndata[nfeat_name]
        with graph.local_scope():

            # Set initial hidden
            h_init = self.input_project(ndata)
            graph.ndata.update({"_h": h_init})

            # Make these 1 dimensional, assuming they were onehot
            efeats = graph.edata[efeat_name].argmax(1)
            for layer in range(self.num_step_message_passing):
                h_in = graph.ndata["_h"]
                edge_transforms = self.edge_transform_weights[efeats]
                edge_biases = self.edge_transform_bias[efeats]
                graph.edata.update({"w": edge_transforms, "b": edge_biases})
                graph.update_all(self.message_pass, fn.sum("m", "_h"))
                h_out = graph.ndata["_h"]
                h_out = self.gru(h_out, h_in)
                graph.ndata.update({"_h": h_out})

            # Read out at the end
            # Because we have other layers, don't get this
            node_level_output = torch.cat([h_init, h_out], -1)
            # gate_weight = self.gate_mlp(h_out)
            # output_val = gate_weight * self.tan_mlp(node_level_output)
            output_val = self.tan_mlp(node_level_output)
        return output_val

# DGL Models
# https://docs.dgl.ai/en/0.6.x/_modules/dgl/nn/pytorch/glob.html#SetTransformerDecoder

class MultiHeadAttention(nn.Module):
    r"""Multi-Head Attention block, used in Transformer, Set Transformer and so on.

    Parameters
    ----------
    d_model : int
        The feature size (input and output) in Multi-Head Attention layer.
    num_heads : int
        The number of heads.
    d_head : int
        The hidden size per head.
    d_ff : int
        The inner hidden size in the Feed-Forward Neural Network.
    dropouth : float
        The dropout rate of each sublayer.
    dropouta : float
        The dropout rate of attention heads.

    Notes
    -----
    This module was used in SetTransformer layer.
    """

    def __init__(self,
                 d_model,
                 num_heads,
                 d_head,
                 d_ff,
                 dropouth=0.0,
                 dropouta=0.0):
        super(MultiHeadAttention, self).__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_head = d_head
        self.d_ff = d_ff
        self.proj_q = nn.Linear(d_model, num_heads * d_head, bias=False)
        self.proj_k = nn.Linear(d_model, num_heads * d_head, bias=False)
        self.proj_v = nn.Linear(d_model, num_heads * d_head, bias=False)
        self.proj_o = nn.Linear(num_heads * d_head, d_model, bias=False)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropouth),
            nn.Linear(d_ff, d_model),
        )
        self.droph = nn.Dropout(dropouth)
        self.dropa = nn.Dropout(dropouta)
        self.norm_in = nn.LayerNorm(d_model)
        self.norm_inter = nn.LayerNorm(d_model)
        self.reset_parameters()

    def reset_parameters(self):
        """Reinitialize learnable parameters."""
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def self_attention(self, x, mem, lengths_x, lengths_mem):
        batch_size = len(lengths_x)
        max_len_x = max(lengths_x)
        max_len_mem = max(lengths_mem)
        device = x.device
        lengths_x = torch.tensor(lengths_x, dtype=torch.int64, device=device)
        lengths_mem = torch.tensor(lengths_mem,
                                   dtype=torch.int64,
                                   device=device)

        queries = self.proj_q(x).view(-1, self.num_heads, self.d_head)
        keys = self.proj_k(mem).view(-1, self.num_heads, self.d_head)
        values = self.proj_v(mem).view(-1, self.num_heads, self.d_head)

        # padding to (B, max_len_x/mem, num_heads, d_head)
        queries = dgl_F.pad_packed_tensor(queries, lengths_x, 0)
        keys = dgl_F.pad_packed_tensor(keys, lengths_mem, 0)
        values = dgl_F.pad_packed_tensor(values, lengths_mem, 0)

        # attention score with shape (B, num_heads, max_len_x, max_len_mem)
        e = torch.einsum("bxhd,byhd->bhxy", queries, keys)
        # normalize
        e = e / np.sqrt(self.d_head)

        # generate mask
        mask = _gen_mask(lengths_x, lengths_mem, max_len_x, max_len_mem)
        e = e.masked_fill(mask == 0, -float("inf"))

        # apply softmax
        alpha = torch.softmax(e, dim=-1)
        # the following line addresses the NaN issue, see
        # https://github.com/dmlc/dgl/issues/2657
        alpha = alpha.masked_fill(mask == 0, 0.0)

        # sum of value weighted by alpha
        out = torch.einsum("bhxy,byhd->bxhd", alpha, values)
        # project to output
        out = self.proj_o(out.contiguous().view(batch_size, max_len_x,
                                                self.num_heads * self.d_head))
        # pack tensor
        out = dgl_F.pack_padded_tensor(out, lengths_x)
        return out

    def forward(self, x, mem, lengths_x, lengths_mem):
        """
        Compute multi-head self-attention.

        Parameters
        ----------
        x : torch.Tensor
            The input tensor used to compute queries.
        mem : torch.Tensor
            The memory tensor used to compute keys and values.
        lengths_x : list
            The array of node numbers, used to segment x.
        lengths_mem : list
            The array of node numbers, used to segment mem.
        """

        ### Following a _pre_ transformer

        # intra norm
        x = x + self.self_attention(self.norm_in(x), mem, lengths_x,
                                    lengths_mem)

        # inter norm
        x = x + self.ffn(self.norm_inter(x))

        ## intra norm
        # x = self.norm_in(x + out)

        ## inter norm
        # x = self.norm_inter(x + self.ffn(x))
        return x

class SetAttentionBlock(nn.Module):
    r"""SAB block introduced in Set-Transformer paper.

    Parameters
    ----------
    d_model : int
        The feature size (input and output) in Multi-Head Attention layer.
    num_heads : int
        The number of heads.
    d_head : int
        The hidden size per head.
    d_ff : int
        The inner hidden size in the Feed-Forward Neural Network.
    dropouth : float
        The dropout rate of each sublayer.
    dropouta : float
        The dropout rate of attention heads.

    Notes
    -----
    This module was used in SetTransformer layer.
    """

    def __init__(self,
                 d_model,
                 num_heads,
                 d_head,
                 d_ff,
                 dropouth=0.0,
                 dropouta=0.0):
        super(SetAttentionBlock, self).__init__()
        self.mha = MultiHeadAttention(d_model,
                                      num_heads,
                                      d_head,
                                      d_ff,
                                      dropouth=dropouth,
                                      dropouta=dropouta)

    def forward(self, feat, lengths):
        """
        Compute a Set Attention Block.

        Parameters
        ----------
        feat : torch.Tensor
            The input feature.
        lengths : list
            The array of node numbers, used to segment feat tensor.
        """
        return self.mha(feat, feat, lengths, lengths)

class SetTransformerEncoder(nn.Module):
    r"""

    Description
    -----------
    The Encoder module in `Set Transformer: A Framework for Attention-based
    Permutation-Invariant Neural Networks <https://arxiv.org/pdf/1810.00825.pdf>`__.

    Parameters
    ----------
    d_model : int
        The hidden size of the model.
    n_heads : int
        The number of heads.
    d_head : int
        The hidden size of each head.
    d_ff : int
        The kernel size in FFN (Positionwise Feed-Forward Network) layer.
    n_layers : int
        The number of layers.
    block_type : str
        Building block type: 'sab' (Set Attention Block) or 'isab' (Induced
        Set Attention Block).
    m : int or None
        The number of induced vectors in ISAB Block. Set to None if block type
        is 'sab'.
    dropouth : float
        The dropout rate of each sublayer.
    dropouta : float
        The dropout rate of attention heads.

    Examples
    --------
    >>> import dgl
    >>> import torch as th
    >>> from dgl.nn import SetTransformerEncoder
    >>>
    >>> g1 = dgl.rand_graph(3, 4)  # g1 is a random graph with 3 nodes and 4 edges
    >>> g1_node_feats = torch.rand(3, 5)  # feature size is 5
    >>> g1_node_feats
    tensor([[0.8948, 0.0699, 0.9137, 0.7567, 0.3637],
            [0.8137, 0.8938, 0.8377, 0.4249, 0.6118],
            [0.5197, 0.9030, 0.6825, 0.5725, 0.4755]])
    >>>
    >>> g2 = dgl.rand_graph(4, 6)  # g2 is a random graph with 4 nodes and 6 edges
    >>> g2_node_feats = torch.rand(4, 5)  # feature size is 5
    >>> g2_node_feats
    tensor([[0.2053, 0.2426, 0.4111, 0.9028, 0.5658],
            [0.5278, 0.6365, 0.9990, 0.2351, 0.8945],
            [0.3134, 0.0580, 0.4349, 0.7949, 0.3891],
            [0.0142, 0.2709, 0.3330, 0.8521, 0.6925]])
    >>>
    >>> set_trans_enc = SetTransformerEncoder(5, 4, 4, 20)  # create a settrans encoder.

    Case 1: Input a single graph

    >>> set_trans_enc(g1, g1_node_feats)
    tensor([[ 0.1262, -1.9081,  0.7287,  0.1678,  0.8854],
            [-0.0634, -1.1996,  0.6955, -0.9230,  1.4904],
            [-0.9972, -0.7924,  0.6907, -0.5221,  1.6211]],
           grad_fn=<NativeLayerNormBackward>)

    Case 2: Input a batch of graphs

    Build a batch of DGL graphs and concatenate all graphs' node features into one tensor.

    >>> batch_g = dgl.batch([g1, g2])
    >>> batch_f = torch.cat([g1_node_feats, g2_node_feats])
    >>>
    >>> set_trans_enc(batch_g, batch_f)
    tensor([[ 0.1262, -1.9081,  0.7287,  0.1678,  0.8854],
            [-0.0634, -1.1996,  0.6955, -0.9230,  1.4904],
            [-0.9972, -0.7924,  0.6907, -0.5221,  1.6211],
            [-0.7973, -1.3203,  0.0634,  0.5237,  1.5306],
            [-0.4497, -1.0920,  0.8470, -0.8030,  1.4977],
            [-0.4940, -1.6045,  0.2363,  0.4885,  1.3737],
            [-0.9840, -1.0913, -0.0099,  0.4653,  1.6199]],
           grad_fn=<NativeLayerNormBackward>)

    See Also
    --------
    SetTransformerDecoder

    Notes
    -----
    SetTransformerEncoder is not a readout layer, the tensor it returned is nodewise
    representation instead out graphwise representation, and the SetTransformerDecoder
    would return a graph readout tensor.
    """

    def __init__(
        self,
        d_model,
        n_heads,
        d_head,
        d_ff,
        n_layers=1,
        block_type="sab",
        m=None,
        dropouth=0.0,
        dropouta=0.0,
    ):
        super(SetTransformerEncoder, self).__init__()
        self.n_layers = n_layers
        self.block_type = block_type
        self.m = m
        layers = []
        if block_type == "isab" and m is None:
            raise KeyError(
                "The number of inducing points is not specified in ISAB block."
            )

        for _ in range(n_layers):
            if block_type == "sab":
                layers.append(
                    SetAttentionBlock(
                        d_model,
                        n_heads,
                        d_head,
                        d_ff,
                        dropouth=dropouth,
                        dropouta=dropouta,
                    ))
            elif block_type == "isab":
                # layers.append(
                #    InducedSetAttentionBlock(m, d_model, n_heads, d_head, d_ff,
                #                             dropouth=dropouth, dropouta=dropouta))
                raise NotImplementedError()
            else:
                raise KeyError(
                    "Unrecognized block type {}: we only support sab/isab")

        self.layers = nn.ModuleList(layers)

    def forward(self, graph, feat):
        """
        Compute the Encoder part of Set Transformer.

        Parameters
        ----------
        graph : DGLGraph
            The input graph.
        feat : torch.Tensor
            The input feature with shape :math:`(N, D)`, where :math:`N` is the
            number of nodes in the graph.

        Returns
        -------
        torch.Tensor
            The output feature with shape :math:`(N, D)`.
        """
        lengths = graph.batch_num_nodes()
        for layer in self.layers:
            feat = layer(feat, lengths)
        return feat

def _gen_mask(lengths_x, lengths_y, max_len_x, max_len_y):
    """Generate binary mask array for given x and y input pairs.

    Parameters
    ----------
    lengths_x : Tensor
        The int tensor indicates the segment information of x.
    lengths_y : Tensor
        The int tensor indicates the segment information of y.
    max_len_x : int
        The maximum element in lengths_x.
    max_len_y : int
        The maximum element in lengths_y.

    Returns
    -------
    Tensor
        the mask tensor with shape (batch_size, 1, max_len_x, max_len_y)
    """
    device = lengths_x.device
    # x_mask: (batch_size, max_len_x)
    x_mask = torch.arange(max_len_x,
                          device=device).unsqueeze(0) < lengths_x.unsqueeze(1)
    # y_mask: (batch_size, max_len_y)
    y_mask = torch.arange(max_len_y,
                          device=device).unsqueeze(0) < lengths_y.unsqueeze(1)
    # mask: (batch_size, 1, max_len_x, max_len_y)
    mask = (x_mask.unsqueeze(-1) & y_mask.unsqueeze(-2)).unsqueeze(1)
    return mask

def pad_packed_tensor(input, lengths, value):
    """pad_packed_tensor"""
    old_shape = input.shape
    device = input.device
    if not isinstance(lengths, torch.Tensor):
        lengths = torch.tensor(lengths, dtype=torch.int64, device=device)
    else:
        lengths = lengths.to(device)
    max_len = (lengths.max()).item()

    batch_size = len(lengths)
    x = input.new(batch_size * max_len, *old_shape[1:])
    x.fill_(value)

    # Initialize a tensor with an index for every value in the array
    index = torch.ones(len(input), dtype=torch.int64, device=device)

    # Row shifts
    row_shifts = torch.cumsum(max_len - lengths, 0)

    # Calculate shifts for second row, third row... nth row (not the n+1th row)
    # Expand this out to match the shape of all entries after the first row
    row_shifts_expanded = row_shifts[:-1].repeat_interleave(lengths[1:])

    # Add this to the list of inds _after_ the first row
    cumsum_inds = torch.cumsum(index, 0) - 1
    cumsum_inds[lengths[0]:] += row_shifts_expanded
    x[cumsum_inds] = input
    return x.view(batch_size, max_len, *old_shape[1:])


import json
from pathlib import Path

class JsonStorage:
    """
    This class handles Json Data.

    Arguments:
        project_root: the project root; this to make the class as generic as possible.
        json_directory_under_root: the location of the Json folder the file is in under the project root. Without '\'
                        can be either a path or string
    """
    def __init__(self, project_root, json_directory_under_root):
        self.project_root = str(project_root)
        self.folder_location_specified = str(json_directory_under_root)
        self.json_directory = Path(self.project_root, self.folder_location_specified)

    def read_json_file(self, file_name):
        """
        This function returns the data from a json file.

        Arguments:
            file_name: the name of the file to read.
        """
        file_path = Path(self.json_directory, file_name)

        try:
            with open(file_path) as json_file:
                data = json.load(json_file)
                return data
        except FileNotFoundError:
            print("File was not found!")
            print(f"File name: {file_name}")
            print(f"File file_path: {file_path}")
        except TypeError as _error:
            print(f"Type Error: {_error}")
            print(f"File name: {file_name}")

    def write_json_file(self, data, file_name):
        """
        This function writes to a json file.

        Arguments:
            file_name: the name of the file to write to.
            data: the data to write.
        """
        file_path = Path(self.json_directory, file_name)

        try:
            with open(file_path, "w") as file:
                json.dump(data, file, indent=4)
        except FileNotFoundError:
            print("File was not found!")
            print(f"File name: {file_name}")
            print(f"File file_path: {file_path}")
        except TypeError as _error:
            print(f"Type Error: {_error}")
            print(f"File name: {file_name}")

    def read_json_file_table(self, file_name, table: str):
        """
        This function returns the table data from a json file.

        Arguments:
            file_name: the name of the file to read.
            table: the name of the table to read.
        """
        try:
            return self.read_json_file(file_name)[table]
        except KeyError:
            print("Table wasn't found!")
        except TypeError as _error:
            print(f"Type Error: {_error}")
            print(f"File name: {file_name}, Table name: {table}")

    def write_json_file_table(self, file_name, table: str, input_data):
        """
        This function writes to table data from a json file.

        Arguments:
            file_name: the name of the file to write to.
            table: the name of the table to write to.
            input_data: the data to insert into the table.
        """
        data = {table: input_data}
        try:
            self.write_json_file(data, file_name)
        except KeyError:
            print("Table wasn't found!")
        except TypeError as _error:
            print(f"Type Error: {_error}")
            print(f"File name: {file_name}, Table name: {table}")


from UserDict import DictMixin
from xml.etree import ElementTree
from uuid import uuid4
from rdflib import Namespace, Graph, Literal, XSD
from hcsvlab_robochef.rdf.namespaces import *

from hcsvlab_robochef.annotations.annotation_names import NAMEMAP

def _subnode_(node, name, text, attr={}):
    """Utility fn to create a child node called 'name'
    with text content 'text' as a child of 'node'"""

    sn = ElementTree.Element(name)
    sn.text = str(text)
    sn.attrib = attr
    node.append(sn)

class AnnotationCollection:
    """All the annotations on an item"""

    def __init__(self, annotationList, itemid, metadataMap):

        self.annotations = annotationList
        self.metaMap = metadataMap
        self.itemid = itemid
        self.id = str(uuid4())   # random unique identifier

    def uri(self):
        """Return an identifier URI for this annotation collection"""

        return corpus_annotation_uri(self.metaMap.corpusID, self.id)

    def to_rdf(self):
        """Add RDF for all of the annotations in the collection
        in an RDF graph"""

        graph = Graph()
        graph = bind_graph(graph)

        for a in self.annotations:
            a.to_rdf(graph, self.metaMap, self.uri())

        graph.add((self.uri(), RDF.type, DADA.AnnotationCollection))
        graph.add((self.uri(), DADA.annotates, self.metaMap.item_uri(self.itemid)))

        return graph

class Annotation(DictMixin):
  """
  An annotation on a region of a source document with the position
  of the source region defined by start and end offsets.
  By default (this class) these are interpreted as UTF8 character
  offsets from the start of the file. Subclasses define alternate
  interpretations, eg. second times in an audio file.
  """

  uniqueid = 0

  def __init__(self, tipe, val, start, end, id=None, properties=None):
    # generate an id unless we're given one
    if id:
        self.id = str(id)
    else:
        self.id = str(Annotation.uniqueid)
        Annotation.uniqueid = Annotation.uniqueid + 1

    # here we map the string type onto a URI

    if NAMEMAP.has_key(tipe):
       self.tipe = NAMEMAP[tipe]
    else:
        raise Exception("Annotation name '%s' not in the list in annotation_names.py" % tipe )

    self.start = start
    self.end = end

    # dictionary for properties
    if properties:
        self.properties = properties
    else:
        self.properties = dict()
    # val is a special property
    if val:
        self.properties['val'] = val
    else:
        self.properties['val'] = ""

  # define the dictionary interface
  def __getitem__(self, key):
      return self.properties[key]

  def __setitem__(self, key, value):
      self.properties[key] = value

  def __delitem__(self, key):
      self.properties.__delitem__(key)

  def keys(self):
      return self.properties.keys()

  def __repr__(self):
    return (self.tipe + ': ' + self['val'] + ' ' + str(self.start) + ' -> ' + str(self.end))

  def __cmp__(self, other):
    if (cmp(self.tipe, other.tipe) == 0):
      if (cmp(self["val"], other["val"]) == 0):
        if (cmp(self.start, other.start) == 0):
          return cmp(self.end, other.end)
        else:
          return cmp(self.start, other.start)
      else:
        return cmp(self["val"], other["val"])
    else:
      return cmp(self.tipe, other.tipe)

  def grow(self, by):
    self.end += by

  def shrink(self, by):
    self.end -= by

  def forward(self, by):
    self.start += by
    self.end   += by

  def backwards(self, by):
    self.start -= by
    self.end   -= by

  def to_rdf(self, g, metaMap, collectionUri):
      """Add triples to this rdf graph to represent this
      annotation. Nodes go into the given namespace
      and are part of the collectionUri"""

      # some identifiers
      property_namespace = corpus_property_namespace(metaMap.corpusID)
      annoturi = corpus_annotation_uri(metaMap.corpusID, self.id )
      locatoruri = corpus_annotation_uri(metaMap.corpusID, self.id + "L")

      # annotation
      g.add((annoturi, RDF.type, DADA.Annotation))
      g.add((annoturi, DADA.partof, collectionUri))

      # locator info depends on the type of annotation
      locatoruri = self.locator_rdf(locatoruri, g)
      g.add((annoturi, DADA.targets, locatoruri))

      g.add((annoturi, DADA.type, self.tipe))

      for key in self.keys():
          if self[key] != '':
              if key == "speakerid":
                  g.add((annoturi, AUSNC.speakerid, metaMap.speaker_uri(self[key])))
              elif key == 'val':
                  g.add((annoturi, DADA.label, Literal(unicode(self[key]))))
              else:
                  g.add((annoturi, property_namespace[key], Literal(unicode(self[key]))))

  def locator_rdf(self, locatoruri, graph):
        """Add RDF triples to the graph to represent the locator information
        for this annotation"""

        graph.add((locatoruri, RDF.type, DADA.UTF8Region))
        graph.add((locatoruri, DADA.start, Literal(int(self.start), datatype=XSD.integer)))
        graph.add((locatoruri, DADA.end, Literal(int(self.end), datatype=XSD.integer)))

        return locatoruri

class SecondAnnotation(Annotation):
    """An annotation on a audio/video document with endpoints defined by offsets in seconds,
    defines the serialisation of the locator"""

    def locator_rdf(self, locatoruri, graph):
        """Add RDF triples to the graph to represent the locator information
        for this annotation"""

        graph.add((locatoruri, RDF.type, DADA.SecondRegion))
        graph.add((locatoruri, DADA.start, Literal(float(self.start), datatype=XSD.float)))
        graph.add((locatoruri, DADA.end, Literal(float(self.end), datatype=XSD.float)))
        return locatoruri

class HMSAnnotation(Annotation):
    """An annotation on a audio/video document with endpoints defined by offsets in HH:MM:SS
    defines the serialisation of the locator"""

    def locator_rdf(self, locatoruri, graph):
        """Add RDF triples to the graph to represent the locator information
        for this annotation"""

        graph.add((locatoruri, RDF.type, DADA.HMSRegion))
        graph.add((locatoruri, DADA.start, Literal(self.start)))
        graph.add((locatoruri, DADA.end, Literal(self.end)))
        return locatoruri


import os

from setuptools import setup, find_packages

def read(fname):
    return open(os.path.join(os.path.dirname(__file__), fname)).read()

setup(
    name = "splitfdr",
    version = "0.0.1",
    author = "Hangyu LIN",
    author_email = "linhy960303@gmail.com",
    description = ("The package of SplitKnockoffs and SplitModelX"),
    license = "BSD",
    keywords = "SplitKnockoffs and SplitModelX",
    packages=find_packages("src"),
    package_dir={"":"src"},
    long_description=read('README.md'),
    classifiers=[
        "Development Status :: 3 - Alpha",
        "Topic :: Utilities",
        "License :: OSI Approved :: BSD License",
    ],
)


import boto3
import time
import os

def redshift_unload_partitioned(cluster_id, sql_query, s3_path, role, partition_by_column, client = None):
    # Create the Boto3 client
    client = client or boto3.client('redshift-data')
    # Launch the SQL query with Redshift Data API
    execution_id = client.execute_statement(
        ClusterIdentifier=cluster_id,
        Database='dev',
        DbUser='awsuser',
        Sql=f"unload ('{sql_query}') TO '{s3_path}' iam_role '{role}' partition by ({partition_by_column}) csv header parallel off;",
    )['Id']

    # Wait for completion
    status = client.describe_statement(Id=execution_id)['Status']
    print(status)
    while status not in ['FINISHED', 'FAILED']:
        time.sleep(15)
        status = client.describe_statement(Id=execution_id)['Status']
        print(status)
        if status == 'FAILED':
            print(client.describe_statement(Id=execution_id)['Error'])
            return False
        else:
            return True

def lambda_handler(event, context):
    # Read the values from the event - event is the dict 'inputs' of the LambdaStep
    cluster_id = event['CLUSTER_ID']
    sql_query = event['SQL_QUERY']
    s3_path = event['S3_PATH']
    role = event['REDSHIFT_ROLE']
    partition_by_column = event['PARTITION_BY_COLUMN']
    # Launch the unload by partition
    status = redshift_unload_partitioned(sql_query, s3_path, role, partition_by_column)
    return {
        'status': status,
        's3_path': s3_path
    }

if __name__=='__main__':
    # Set-up the path and roles
    bucket = 'YOUR_BUCKET_HERE'
    key_prefix = 'some/s3/prefix/where/to/store/data'
    role = 'YOUR_ROLE_ARN_HERE'
    cluster_id = 'YOUR_CLUSTER_ID_HERE'
    sql_query = 'SELECT * from my_table'
    partition_by_column = 'YOUR_COLUMN_TO_PARTITION_BY'

    # Define the s3 path
    s3_path = f's3://{bucket}/{key_prefix}/'

    # Launch the unload by partition
    status = redshift_unload_partitioned(cluster_id, sql_query, s3_path, role, partition_by_column)
    print(status)


import random
import config as cf
import strat as st
import matplotlib.pyplot as plt
import numpy as np
import math

def welcome_msg():
    print('Welcome to Blackjack \n')
    mode = input('Play manually? (y/n): ')
    if mode.lower() == 'y':
        cf.mode = 'man'
        print('Guide:\nh = hit\ns = stand\nd = double\nv = split\nx = surrender\nq = quit\n')
    elif mode.lower() == 'n':
        cf.mode = 'sim'
        print('Entering simulation mode')
        # cf.sim_hands = int(input('Enter number of hands to be played: '))
    elif mode.lower() == 't':
        cf.mode = 'test'
        print('Entering testing mode.')
    else:
        print('Invalid input, entering sim mode')
    cf.bankroll = cf.initial = int(input("Enter buy-in amount: "))

# converts card to numeric value
def numeric(card):
    if card == 'A':
        return 11
    elif card == 'T' or card == 'J' or card == 'Q' or card == 'K':
        return 10
    else:
        return int(card)

# returns sum of cards
def get_total(cards):
    total = 0
    for i in range(0, len(cards)):
        total += numeric(cards[i])
    return total

# generate new deck of cards and shuffle
def gen_cards():
    cards = ['A', '2', '3', '4', '5', '6', '7', '8', '9', 'T', 'J', 'Q', 'K']
    shoe = []
    for i in range(0, 4 * cf.numDecks):
        shoe.extend(cards)
    random.shuffle(shoe)

    # init counts
    cf.rc = 0
    cf.tc = 0

    return shoe

def deal_cards(deck):
    if len(deck) < cf.cut_card:
        deck = gen_cards()  # reshuffle

    # uncomment to test predetermined deck
    # deck = ['T', '9', '6', '5', 'T', 'T', '8', 'T', '6', 'A', '2', '5', 'T', '4']

    # dealer gets 2 cards
    c1 = deck.pop(0)
    c2 = deck.pop(0)
    d_cards = [c1, c2]
    adjust_counts(c1, deck)
    adjust_counts(c2, deck)

    # player gets 2 cards
    c1 = deck.pop(0)
    c2 = deck.pop(0)
    p_cards = [[c1, c2]]
    adjust_counts(c1, deck)
    adjust_counts(c2, deck)
    return d_cards, p_cards, deck

def reset_globals():
    cf.total_hands += 1
    cf.hi = 0
    cf.num_splits = 0
    cf.handled_insurance = False
    cf.insurance_bet = 0
    cf.bought_insurance = False
    for i in range(0, cf.num_splits + 1):
        if cf.doubled_down[i]:
            cf.bet_size[i] /= 2
            cf.doubled_down[i] = False

def adjust_counts(card, deck):

    if not cf.enable_count:
        cf.tc = 0
        return
    if numeric(card) <= 6:
        cf.rc += 1
    elif numeric(card) > 9:
        cf.rc -= 1

    num_decks = math.ceil(len(deck) / 52)
    cf.tc = math.floor(cf.rc / num_decks)

def print_cards(d_cards, p_cards, show_hole):
    # dealer hole-card is only shown after player acts
    if show_hole:
        print(f'Dealer Cards: {d_cards}')
        print(f'Your cards: {p_cards}')
    else:
        print(f'Dealer Up-card: {d_cards[0]}')
        print(f'Your cards: {p_cards}')

def check_bj(cards):
    if get_total(cards) == 21 and len(cards) == 2:
        return True
    else:
        return False

def ask_insurance():
    if cf.mode == 'man':
        ans = input('Would you like to buy insurance? (y/n): ')
    else:
        if cf.tc >= 3:
            cf.bought_insurance = True
            cf.insurance_bet = 0.5 * cf.bet_size[0]
            return True
        else:
            return False

    if ans.lower() == 'y':
        cf.bought_insurance = True
        cf.insurance_bet = 0.5 * cf.bet_size[0]
        return True
    else:
        return False

# checks to see if game is over due to blackjack
def handle_bj(p_cards, d_cards, is_split):
    p_blackjack = check_bj(p_cards)
    d_blackjack = check_bj(d_cards)
    ins = False

    if d_cards[0] == 'A':
        if cf.mode != 'sim':
            print_cards(d_cards, p_cards, False)
        ins = ask_insurance()

    if d_blackjack and ins and not cf.handled_insurance:
        if cf.mode != 'sim':
            print('You win insurance bet.')
        cf.bankroll += cf.insurance_bet
        cf.handled_insurance = True
    elif not d_blackjack and ins and not cf.handled_insurance:
        if cf.mode != 'sim':
            print('You lose insurance bet.')
        cf.bankroll -= cf.insurance_bet
        cf.handled_insurance = True

    if p_blackjack and not d_blackjack:
        if cf.mode == 'man' or cf.mode == 'test':
            print_cards(d_cards, p_cards, True)
            print('Blackjack! You win\n')
        cf.sess_wins += 1
        if is_split:
            cf.bankroll += cf.bet_size[cf.hi]  # can't get natural blackjack on split hands
        else:
            cf.bankroll += cf.bj_payout * cf.bet_size[cf.hi]
        return True
    elif d_blackjack and not p_blackjack:
        if cf.mode == 'man' or cf.mode == 'test':
            print_cards(d_cards, p_cards, True)
            print('Dealer has blackjack and wins.\n')
        cf.sess_losses += 1
        cf.bankroll -= cf.bet_size[cf.hi]
        return True
    elif p_blackjack and d_blackjack:
        if cf.mode == 'man' or cf.mode == 'test':
            print_cards(d_cards, p_cards, True)
            print('Push.\n')
        cf.sess_pushes += 1
        return True
    else:
        return False

# Possible combos: using all aces as 1's
# or using one of the aces as an 11 while the rest are 1's.
# Returns boolean for if busted and total that should be used
def handle_aces(cards):
    num_aces = cards.count('A')  # number of aces in hand
    hard_total = 0
    for i in range(0, len(cards)):
        if cards[i] != 'A':
            hard_total += numeric(cards[i])

    # try using one ace as 11
    new_total = hard_total + 11 + (num_aces - 1)
    if new_total > 21:
        new_total = hard_total + num_aces  # using all aces as 1s
        if new_total > 21:
            return False, new_total
        else:
            return True, new_total
    else:
        return True, new_total

# returns boolean for if busted (True = no bust, continue game)
def check_for_bust(cards):

    total = get_total(cards)
    if 'A' in cards:
        run, total = handle_aces(cards)  # need to do special check for aces
        return run
    else:
        if total > 21:
            return False
        else:
            return True

# check for a valid split
def verify_split(cards):
    if len(cards[cf.hi]) == 2:
        if numeric(cards[cf.hi][0]) == numeric(cards[cf.hi][1]) and cf.num_splits <= cf.max_splits:
            return True
    return False

def verify_double(cards):
    if cf.doubled_down[cf.hi] or len(cards[cf.hi]) > 2:
        return False
    else:
        return True

def handle_split(cards, deck):
    aces = False
    if not verify_split(cards):
        if cf.mode == 'man' or cf.mode == 'test':
            print("Cannot split here.")
        return cards, False, aces

    if cards[cf.hi][0] == cards[cf.hi][1] == 'A':
        aces = True

    # add new hand to player hand list
    cards.append([0, 0])
    cards[cf.num_splits][0] = cards[cf.hi][1]
    cards[cf.hi][1] = deck.pop(0)
    cards[cf.num_splits][1] = deck.pop(0)

    adjust_counts(cards[cf.hi][1], deck)
    adjust_counts(cards[cf.num_splits][1], deck)

    return cards, True, aces

# determines winner when neither player busts
def find_winner(p_cards, d_cards, index):

    p_total = get_total(p_cards)
    while p_total > 21:
        p_total -= 10  # must be using an ace as a 1

    d_total = get_total(d_cards)
    while d_total > 21:
        d_total -= 10  # must be using an ace as a 1

    if cf.mode == 'man' or cf.mode == 'test':
        print(f'Your total is {p_total}')
        print(f'Dealer total is {d_total}')

    if p_total > d_total:
        if cf.mode == 'man' or cf.mode == 'test':
            print('You win!\n')
        cf.sess_wins += 1
        cf.bankroll += cf.bet_size[index]
    elif d_total > p_total:
        if cf.mode == 'man' or cf.mode == 'test':
            print('Dealer wins.\n')
        cf.sess_losses += 1
        cf.bankroll -= cf.bet_size[index]
    else:
        if cf.mode == 'man' or cf.mode == 'test':
            print('Push\n')
        cf.sess_pushes += 1

def random_choice():
    choices = ['h', 'd', 's', 'x', 'v']
    random.shuffle(choices)
    return choices.pop()

# returns correct action based on basic strategy along with table used
def basic_strategy(p_cards, d_card):

    col = numeric(d_card) - 1
    total = get_total(p_cards)
    while total > 21:
        total -= 10  # must be using an ace as a 1

    hard_total = 0
    for i in range(0, len(p_cards)):
        if p_cards[i] != 'A':
            hard_total += numeric(p_cards[i])

    if len(p_cards) == 2 and numeric(p_cards[0]) == numeric(p_cards[1]) and cf.num_splits < cf.max_splits:
        # use split table
        row = 12 - numeric(p_cards[0])
        return st.split[row][col], 'split'
    elif hard_total <= 9 and 'A' in p_cards and cf.num_splits < cf.max_splits:
        # use soft table
        row = 10 - hard_total
        return st.soft[row][col], 'soft'
    # else use hard table
    row = 22 - total
    return st.hard[row][col], 'hard'

# adjust bet size based on true count
def adjust_bet_size():

    if cf.tc <= -2:
        return 0
    if cf.tc <= 0:
        return cf.bet_spread[0]  # TC <= 0
    elif 1 <= cf.tc <= 5:
        return cf.bet_spread[cf.tc]
    else:
        return cf.bet_spread[5]  # TC > 5

def print_stats():

    margin = cf.bankroll - cf.initial
    pct_gain = (margin / cf.initial) * 100
    ev_per_hand = pct_gain / cf.total_hands
    winnings_per_hand = margin / cf.total_hands
    winnings_per_hour = winnings_per_hand * cf.hph

    print(f'\nTotal rounds played: {cf.total_hands}')
    print(f'# of wins: {cf.sess_wins}')
    print(f'# of losses: {cf.sess_losses}')
    print(f'# of pushes: {cf.sess_pushes}')
    print(f'% of hands won: {"{:.2f}".format((cf.sess_wins / cf.total_hands) * 100)}')
    print(f'Money won/lost: {cf.bankroll - cf.initial}')
    print(f'% money gain: {"{:.2f}".format(pct_gain)}')
    print(f'% EV/hand: {"{:.6f}".format(ev_per_hand)}')
    print(f'$/hand: ${"{:.2f}".format(winnings_per_hand)}')
    print(f'$/hr: ${"{:.2f}".format(winnings_per_hour)}')
    plot_winnings()

def plot_winnings():

    x = np.linspace(0, len(cf.winnings_history), len(cf.winnings_history))
    baseline = np.full((len(cf.winnings_history), ), cf.initial)
    y = np.array(cf.winnings_history)
    m = np.polyfit(x, y, 1)
    best_fit = m[0]*x + m[1]

    plt.plot(x, y, 'r--', x, baseline, 'k--', x, best_fit, 'b--')
    plt.title("Winnings over time")
    plt.xlabel("Hands")
    plt.ylabel("Total Money")
    plt.show()

def hit(cards, deck, double):
    c1 = deck.pop(0)
    cards.append(c1)  # hit
    adjust_counts(c1, deck)
    run = check_for_bust(cards)
    busted = not run
    if double:
        run = False  # can only get one card after doubling
    if cf.mode == 'man' or cf.mode == 'test':
        print(f'{c1} is drawn')
    return cards, busted, run

def player_turn(cards, d_cards, deck):

    # flags
    run = True
    busted = False
    end = False
    sur = False
    has_split = False
    bj = False
    aces = False

    while run:
        if cf.num_splits > 0 and (cf.mode == 'man' or cf.mode == 'test'):
            print(f'Hand #{cf.hi + 1}:')
        if handle_bj(cards[cf.hi], d_cards, True):
            bj = True
            break

        if cf.mode == 'man':
            print_cards(d_cards[0], cards[cf.hi], False)
            action = input('Action: ')
        else:
            action, table = basic_strategy(cards[cf.hi], d_cards[0])
            # print(f'Action is {action}')
            if cf.mode == 'test':
                print_cards(d_cards[0], cards[cf.hi], False)
                print(f'Action is {action} from {table} table')
        action = action.lower()

        # switch on action
        if action == 'h':
            cards[cf.hi], busted, run = hit(cards[cf.hi], deck, False)
        elif action == 's':
            run = False  # stand
        elif action == 'd':
            if verify_double(cards):
                cf.bet_size[cf.hi] *= 2
                cf.doubled_down[cf.hi] = True
                cards[cf.hi], busted, run = hit(cards[cf.hi], deck, True)
            else:
                if cf.mode == 'man' or cf.mode == 'test':
                    print("Invalid double, hitting instead")
                cards[cf.hi], busted, run = hit(cards[cf.hi], deck, False)  # just hit
        elif action == 'v':
            if verify_split(cards) or aces:
                has_split = True
                cf.num_splits += 1
                cards, valid_split, aces = handle_split(cards, deck)
                run = False
            else:
                if cf.mode == 'man' or cf.mode == 'test':
                    print("Invalid split")
        elif action == 'x':
            if len(cards[cf.hi]) <= 2 and cf.allow_sur:
                sur = True
                run = False
            else:
                if cf.mode == 'man' or cf.mode == 'test':
                    print("Invalid surrender")
                cards[cf.hi], busted, run = hit(cards[cf.hi], deck, False)  # just hit
        elif action == 'q':
            end = True
            run = False
        else:
            if cf.mode == 'man' or cf.mode == 'test':
                print('Not a valid input.')
    return busted, end, sur, has_split, bj, aces

def dealer_turn(cards, p_cards, deck):

    if cf.mode == 'man' or cf.mode == 'test':
        print_cards(cards, p_cards, deck)
    total = get_total(cards)
    run = True
    while run:

        # hit
        if total < 17:
            c1 = deck.pop(0)
            cards.append(c1)
            adjust_counts(c1, deck)
            total += numeric(cards[len(cards) - 1])
            if cf.mode == 'man' or cf.mode == 'test':
                print(f'Dealer draws a {cards[len(cards) - 1]}')
        else:
            if 'A' in cards:
                run, total = handle_aces(cards)  # special check to check hand with aces
                if total > 16:
                    run = False
            else:
                run = False

    # check if busted
    if total > 21:
        return True
    else:
        return False


from fastapi import APIRouter
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import List

import requests
import json

from call_chatgpt import call_chatgpt, call_chatgpt_w_function
from logger import logger

router = APIRouter()

class Message(BaseModel):
    role: str
    content: str

@router.post("/booksearch")
def booksearch(messages: List[Message]):

    # messagesの最後から10要素のみにする
    messages = messages[-1:]

    messages_for_query = messages.copy()
    messages_for_summary = messages.copy()

    # System Prompt
    system_prompt = """あなたは書籍検索クエリ作成マシーンです。ユーザからのリクエストに対して、書籍に含まれそうな単語を2個から5個、以下のフォーマットに従って回答してください。なお、順序は関連度順です。

出力フォーマット）
["単語1","単語2","単語3","単語4","単語5"]
例えば、
ユーザ）エベレストで遭難する話
あなた）["エベレスト","遭難","登山"]
"""
    # messagesの先頭に要素を追加
    messages_for_query.insert(0, Message(role="system", content=system_prompt))

    logger.debug(messages_for_query)

    # ChatGPTを呼び出す
    answer_query = call_chatgpt([m.dict() for m in messages_for_query])

    # array配列を+で連結した文字列を作成する
    querystring = "+".join(eval(answer_query))

    googlebooksapi = "https://www.googleapis.com/books/v1/volumes?q="

    option = "&maxResults=5&printType=books&langRestrict=ja&country=JP"

    # GoogleBooksAPIのURLを生成する(googlebooksapi + querystring + test)
    url = googlebooksapi + querystring + option

    logger.debug(url)
    # urlにGETリクエストを送り結果をJSONで取得する
    result = requests.get(url).json()

    total_num = result["totalItems"]

    # bookresult用の配列を用意する
    bookresult = []

    # result.itemsでループする。要素が無い場合は空を入れる
    for item in result.get("items", []):
        # itemのvolumeInfoを取得する
        volumeInfo = item.get("volumeInfo", {})
        # volumeInfoのtitleを取得する
        title = volumeInfo.get("title", "")
        # vlumeInfoのinfolinkを取得する
        infolink = volumeInfo.get("infoLink", "")
        # volumeInfoのauthorsを取得する
        authors = volumeInfo.get("authors", [])
        # volumeInfoのdescriptionを取得する
        description = volumeInfo.get("description", "")
        # volumeInfoのimageLinksのthumbnailを取得する
        imageLinks = volumeInfo.get("imageLinks", {})
        thumbnail = imageLinks.get("smallThumbnail", "")

        # bookresultに辞書型で追加する
        bookresult.append({"title": title, "authors": ",".join(authors), "infolink":infolink,"description": description, "smallThumbnail": thumbnail})

    # bookresultの各要素について、descriptionが空、もしくは20文字以下の場合は、リストから除外する
    bookresult  = [book for book in bookresult if len(book["description"]) > 20]

    # bookresultの要素を先頭から3つに絞る
    bookresult = bookresult[:3]

    # bookresultのtitleとauthorsとdescriptionを改行で連結した文字列を作成する
    bookinfo = "\n".join([f"タイトル:{book['title']}\n著者:{book['authors']}\nあらすじ:{book['description']}\n" for book in bookresult])

    # System Prompt
    system_prompt = f"""あなたは本のコンシェルジュです。以下の書籍情報のみを使って、ユーザからのリクエストに回答してください。
ユーザからのリクエストにどうマッチしているか、また、ユーザが読みたくなるような推薦文書を簡潔に書いてください。
**1冊のみおススメしてください**
**書籍情報以外の情報は使用してはなりません。**

### 書籍情報
{bookinfo}
"""
    logger.debug(system_prompt)

    # messagesの先頭に要素を追加
    messages_for_summary.insert(0, Message(role="system", content=system_prompt))

    logger.debug(messages_for_summary)

    # ChatGPTを呼び出す
    answer = call_chatgpt([m.dict() for m in messages_for_summary])

    data = {
        "message": answer,
        "bookresult": bookresult,
        "totalnum":total_num,
        "query":answer_query
    }
    # JSONResponseに辞書型のデータを渡して返す
    return JSONResponse(content=data)

@router.post("/booksearch/function")
def booksearch_function(messages: List[Message]):

    system_prompt = """\
あなたは書籍検索コンシェルジュです。ユーザの要望をもとに書籍検索を行い、その結果からお薦めの書籍を1冊だけ紹介します。
"""

    functions= [
        {
            "name": "search_book",
            "type": "function",
            "description": "クエリをもとに書籍を検索する",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "array",
                        "description": "書籍検索クエリ",
                        "items": {
                            "type": "string"
                        }
                    }
                },
                "required": ["query"],
            },
        }
    ]
    # messagesの最後から10要素のみにする
    messages = messages[-1:]

    messages_for_query = messages.copy()
    messages_for_summary = messages.copy()

    # messagesの先頭に要素を追加
    messages_for_query.insert(0, Message(role="system", content=system_prompt))

    logger.debug(messages_for_query)

    first_answer = call_chatgpt_w_function([m.dict() for m in messages_for_query],functions)

    function_call = first_answer.get("function_call")

    # first_answerにfunction_name属性があれば
    if function_call:
        logger.debug("Function call detected")
        # function_nameを取得する
        function_name = function_call.get("name")
        # function_nameがsearch_bookの場合
        if function_name == "search_book":

            # queryを取得する
            function_args = json.loads(function_call["arguments"])
            query = function_args.get("query")

            # queryを+で連結した文字列を作成する
            querystring = "+".join(query)

            googlebooksapi = "https://www.googleapis.com/books/v1/volumes?q="

            option = "&maxResults=5&printType=books&langRestrict=ja&country=JP"

            # GoogleBooksAPIのURLを生成する(googlebooksapi + querystring + test)
            url = googlebooksapi + querystring + option

            # urlにGETリクエストを送り結果をJSONで取得する
            logger.debug(url)
            result = requests.get(url).json()

            total_num = result["totalItems"]

            # bookresult用の配列を用意する
            bookresult = []

            # result.itemsでループする。要素が無い場合は空を入れる
            for item in result.get("items", []):
                # itemのvolumeInfoを取得する
                volumeInfo = item.get("volumeInfo", {})
                # volumeInfoのtitleを取得する
                title = volumeInfo.get("title", "")
                # vlumeInfoのinfolinkを取得する
                infolink = volumeInfo.get("infoLink", "")
                # volumeInfoのauthorsを取得する
                authors = volumeInfo.get("authors", [])
                # volumeInfoのdescriptionを取得する
                description = volumeInfo.get("description", "")
                # volumeInfoのimageLinksのthumbnailを取得する
                imageLinks = volumeInfo.get("imageLinks", {})
                thumbnail = imageLinks.get("smallThumbnail", "")

                # bookresultに辞書型で追加する
                bookresult.append({"title": title, "authors": ",".join(authors), "infolink":infolink,"description": description, "smallThumbnail": thumbnail})

            # bookresultの各要素について、descriptionが空、もしくは20文字以下の場合は、リストから除外する
            bookresult  = [book for book in bookresult if len(book["description"]) > 20]

            # bookresultの要素を先頭から3つに絞る
            bookresult = bookresult[:3]

            # bookresultのtitleとauthorsとdescriptionを改行で連結した文字列を作成する
            bookinfo = "\n".join([f"タイトル:{book['title']}\n著者:{book['authors']}\nあらすじ:{book['description']}\n" for book in bookresult])

            temp_message_for_query = [m.dict() for m in messages_for_query]
            temp_message_for_query.append({
                "role":"function",
                "name":"search_book",
                "content":bookinfo
            })

            second_answer = call_chatgpt_w_function(temp_message_for_query,functions)

            data = {
                "message": second_answer.content,
                "bookresult": bookresult,
                "totalnum":total_num,
                "query":query
            }
            # JSONResponseに辞書型のデータを渡して返す
            logger.debug(data)
            return JSONResponse(content=data)
    else:
        data = {
            "message": first_answer.content,
            "bookresult": [],
            "totalnum":0,
            "query":[]
        }

        # JSONResponseに辞書型のデータを渡して返す
        return JSONResponse(content=data)


import unittest

from anndata import AnnData
import pandas as pd
import numpy as np
from pathlib import Path
from scipy.io import mmread
import zarr

from vitessce.data_utils import (
    optimize_arr,
    optimize_adata,
    sort_var_axis,
    to_uint8,
    adata_to_multivec_zarr,
)

data_path = Path('tests/data')

class TestAnnDataUtils(unittest.TestCase):

    def setUp(self):
        X = np.array([
            [5.1, 4.0, 5.0],
            [5.0, 1.0, 5.1],
            [2.1, 1.0, 2.0],
            [4.0, 1.4, 4.1],
        ])
        obs = pd.DataFrame(data=[
            {"cell_id": "cell_1", "leiden": "Cluster 2", "bad": "_"},
            {"cell_id": "cell_2", "leiden": "Cluster 2", "bad": "+"},
            {"cell_id": "cell_3", "leiden": "Cluster 1", "bad": "%"},
            {"cell_id": "cell_4", "leiden": "Cluster 5", "bad": "$"}
        ]).set_index("cell_id")
        var = pd.DataFrame(data=[
            {"gene_id": "SOX10"},
            {"gene_id": "LAMP5"},
            {"gene_id": "RORB"}
        ]).set_index("gene_id")
        obsm = {
            "X_umap": np.array([
                [0.5, 1.0],
                [1.0, 2.0],
                [3.1, 2.2],
                [4.0, 1.0]
            ]),
            "X_ints": np.array([
                [1.0, 2.0],
                [1.0, 2.0],
                [3.0, 2.0],
                [4.0, 1.0]
            ])
        }
        self.adata = AnnData(X=X, obs=obs, var=var, obsm=obsm)

    def test_optimize_arr(self):
        int_arr = self.adata.obsm["X_ints"]
        assert optimize_arr(int_arr).dtype.name == 'uint8'

        float_arr = self.adata.obsm["X_umap"]
        assert optimize_arr(float_arr).dtype.name == 'float32'

    def test_optimize_adata(self):
        adata = self.adata
        adata_optim = optimize_adata(adata, obs_cols=['leiden'], obsm_keys=['X_umap'])
        assert list(adata_optim.obsm.keys()) == ['X_umap']
        assert adata_optim.obsm['X_umap'].dtype.name == 'float32'
        assert adata_optim.obs.columns.values.tolist() == ['leiden']

    def test_sort_var_axis(self):
        adata = self.adata
        leaf_list = sort_var_axis(adata.X, adata.var.index.values)
        adata_sorted = adata[:, leaf_list].copy()
        assert adata.var.index.values.tolist() == ['SOX10', 'LAMP5', 'RORB']
        assert adata_sorted.var.index.values.tolist() == ['LAMP5', 'SOX10', 'RORB']

    def test_to_uint8(self):
        adata = self.adata
        norm_X = to_uint8(adata.X)
        assert norm_X.tolist() == [[5, 4, 5], [5, 1, 5], [2, 1, 2], [4, 1, 4]]

    def test_to_uint8_global_norm(self):
        adata = self.adata
        norm_X = to_uint8(adata.X, norm_along="global")
        assert norm_X.tolist() == [[255, 200, 250], [250, 50, 255], [104, 50, 100], [200, 70, 205]]

    def test_to_uint8_gene_norm(self):
        adata = self.adata
        norm_X = to_uint8(adata.X, norm_along="var")
        assert norm_X.tolist() == [[255, 255, 246], [246, 0, 254], [0, 0, 0], [161, 33, 172]]

    def test_to_uint8_cell_norm(self):
        adata = self.adata
        norm_X = to_uint8(adata.X, norm_along="obs")
        assert norm_X.tolist() == [[255, 0, 231], [248, 0, 255], [255, 0, 231], [245, 0, 255]]

    def test_multivec_zarr(self):
        mtx = mmread(data_path / 'test.snap.mtx').toarray()
        bins_df = pd.read_csv(
            data_path / 'test.snap.bins.txt', header=None, names=["interval"])
        clusters_df = pd.read_csv(
            data_path / 'test.snap.clusters.csv', index_col=0)

        zarr_filepath = data_path / 'test_out.snap.multivec.zarr'

        # The genome assembly is GRCh38 but the chromosome names in the bin names do not start with the "chr" prefix.
        # This is incompatible with the chromosome names from `negspy`, so we need to append the prefix.
        bins_df["interval"] = bins_df["interval"].apply(lambda x: "chr" + x)

        obs = clusters_df[["cluster"]]
        obs["cluster"] = obs["cluster"].astype(str)
        obsm = {"X_umap": clusters_df[["umap.1", "umap.2"]].values}
        adata = AnnData(X=mtx, obs=obs, var=bins_df, obsm=obsm)

        # Sort cluster IDs
        cluster_ids = obs["cluster"].unique().tolist()
        cluster_ids.sort(key=int)
        # Save genomic profiles to multivec-zarr format.
        adata_to_multivec_zarr(adata, zarr_filepath, obs_set_col="cluster", obs_set_name="Cluster", obs_set_vals=cluster_ids)

        z = zarr.open(zarr_filepath, mode='r')

        self.assertEqual(z['chromosomes/chr1/5000'].shape, (4, 49792))
        self.assertEqual(z['chromosomes/chr1/5000'][:, 0].sum(), 0)
        self.assertEqual(z['chromosomes/chr1/5000'][:, 1].sum(), 0)
        self.assertEqual(z['chromosomes/chr1/5000'][:, 2].sum(), 7)
        self.assertEqual(z['chromosomes/chr1/5000'][:, 3].sum(), 7)
        self.assertEqual(z['chromosomes/chr1/5000'][:, 4].sum(), 0)
        self.assertEqual(z['chromosomes/chr1/5000'][0, :].sum(), 17)
        self.assertEqual(z['chromosomes/chr1/10000'][0, :].sum(), 17)
        self.assertEqual(z['chromosomes/chr1/5000'][:, 2].sum(), 7)
        self.assertEqual(z['chromosomes/chr2/5000'][:, 0].sum(), 0)
        self.assertEqual(z['chromosomes/chr2/5000'][:, 1].sum(), 0)
        self.assertEqual(z['chromosomes/chr2/10000'][:, 0].sum(), 0)
        self.assertEqual(z['chromosomes/chr2/5000'][:, 2].sum(), 4)
        self.assertEqual(z['chromosomes/chr2/5000'][:, 3].sum(), 9)
        self.assertEqual(z['chromosomes/chr2/10000'][:, 1].sum(), 13)
        self.assertEqual(z['chromosomes/chr3/5000'][:, 3].sum(), 9)
        self.assertEqual(z['chromosomes/chr3/5000'][:].sum(), 9)
        self.assertEqual(z['chromosomes/chr18/5000'][:].sum(), 8)


import os
import types
import numpy as np

from moviepy.video.VideoClip import VideoClip, ImageClip, TextClip
from moviepy.video.compositing.CompositeVideoClip import CompositeVideoClip
from moviepy.editor import concatenate_videoclips
from moviepy.video.fx.resize import resize
from moviepy.video.fx.scroll import scroll
import moviepy.video.compositing.transitions as transfx
from moviepy.video.fx.fadein import fadein

from ..SVGImageClip import SVGImageClip

def zoom(clip, screensize, show_full_height=False):
    """Zooms preferably image clip for clip duration a little

    To make slideshow more movable

    Parameters
    ---------
    clip
        ImageClip on which to work with duration
    screensize
        Wanted (width, height) tuple
    show_full_height
        Should this image be shown in full height. This is usefull when 4:3
        images are shown in 16:9 video and need to be shown in full.
        Otherwise they are shown in full width and top and bottom is cut off.

    Returns
    ------
    VideoClip in desired size
    """
    #We need to resize high imageč differently
    if clip.h > clip.w or show_full_height:
        clip_resized = (clip.fx(resize, width=screensize[0]*2)
                    .fx(resize, lambda t : 1+0.02*t)
                    .set_position(('center', 'center'))
                    )
        clip_composited = CompositeVideoClip([clip_resized]) \
                .fx(resize, height=screensize[1])
    else:
        clip_resized = (clip.fx(resize, height=screensize[1]*2)
                    .fx(resize, lambda t : 1+0.02*t)
                    .set_position(('center', 'center'))
                    )
        clip_composited = CompositeVideoClip([clip_resized]) \
                .fx(resize, width=screensize[0])

    vid = CompositeVideoClip([clip_composited.set_position(('center', 'center'))],
                            size=screensize)
    return vid

def panorama(clip, screensize, duration=None, speed=None, freeze_duration=1.5):
    """Makes image panorama of large images

    It pans image from left to right since they are very wide
    It works in two ways if duration is given it calculates speed to pan whole
    panorama in given number of seconds. If speed is given it sets duration so
    that whole pan will happen at given speed

    Parameters
    ---------
    clip
        ImageClip (not resized)
    screensize
        Wanted (width, height) tuple
    duration
        If given wanted duration of pan
    speed
        Number of pixel of move per time unit (duration is calculated)
    freeze_duration : float
        Number of seconds to freeze before scroll and at the end
    """

    #TODO: pause image at start and/or end for some time
    print("DURATION", duration, "SPEED:", speed)
    #FIXME: add workable assert
    #assert duration is None or speed is None
    clip = clip.fx(resize, height=screensize[1])

    width = screensize[0]

    if speed is not None:
        clip = clip.set_duration((clip.w-width)/speed)
    else:
        clip = clip.set_duration(duration)
        speed = (clip.w-width)/duration

    scrolled = scroll(clip, w=width, x_speed=speed)
    return concatenate_videoclips([scrolled.to_ImageClip(0).set_duration(freeze_duration),
        scrolled,
        #FIXME: why is +1 actually needed without it there is no freeze at end
        scrolled.to_ImageClip(scrolled.duration).set_duration(freeze_duration+1)])

def image_effect(clip, screensize, duration=None, speed=None, show_full_height=False):
    if clip.w/clip.h > 2:
        return panorama(clip, screensize, duration, speed)
    return zoom(clip, screensize, show_full_height)

def make_clip(clip, text, height, width, font, font_color, fontsize):
    if font is None:
        return clip
    #TODO: make text caption style configurable
    #Swimming changes text color to light blue, moves text up and resized
    #it
    #It also adds black stroke that's why resizing is needed
    #And adds icon of swimming on the bottom
    if text.startswith("KOP:"):
        text = text[4:]
        font_color_c="#00aaff"
        pos = "top"
        swimm = True
        stroke_color="black"
        factor=2
        mult=1.4*factor
    else:
        font_color_c = font_color
        pos = "bottom"
        swimm = False
        stroke_color = None
        mult=1
        factor=1
    if width is None:
        text_space_size = None
    else:
        text_space_size = (width*factor, 40*factor)
    text_caption = TextClip(text, size=text_space_size, method="caption",
            align="center",
            color=font_color_c, fontsize=fontsize*mult, font=font,
            stroke_color=stroke_color)
    if swimm:
        text_caption = text_caption.fx(resize, height=text_caption.h/2)
    text_caption = text_caption.set_pos(("center", "center"))
    #print (clip.size, tc.size)
    #TODO: makes transparent bar size same size (based on highest caption)
    #adds 60% transparent bar under the caption
    #Bar starts 5 pixels above caption and ends at the bottom
    if not swimm:
        text_caption_bar = text_caption.on_color(size=(clip.w, text_caption.h+5),
                color=(0,0,0), pos=('center'), col_opacity=0.6)
        text_caption_with_bar = text_caption_bar.set_pos(('center',
            'bottom'))
    else:
        text_caption_with_bar = text_caption.set_pos(('center', 20))
    clips = [clip, text_caption_with_bar]
    if swimm:
        swimm_size = fontsize
        #TODO: make location of this configurable
        swim_clip = \
        SVGImageClip("/home/mabu/programiranje/overlay/projects/glein/images/swimming-15.svg",
                width=swimm_size, height=swimm_size).set_pos(('center',
            clip.h-20-swimm_size))
        clips.append(swim_clip)
    return (CompositeVideoClip(clips) \
            .set_duration(clip.duration))

class SlideshowImagesClip(VideoClip):

    def __init__(self, sequence, titles, height=None, width=None, image_duration=4,
        transition_duration=1, fontsize=30, font="M+-1p-medium",
        font_color="white", zoom_images=False, test=False):
        """Function makes slideshow VideoClip from sequence of images with their
        captions

        Parameters
        ---------
        sequence
            List of paths to images
        titles : list
            List of image captions
        height
            Height of wanted VideoClip
        width
            Width of wanted VideoClip
        image_duration
            How long is one image visible
        transition_duration
            How long is fade transition between images
        fontsize
            Font point size
        font
            Name of the font to use. See ``TextClip.list('font')`` for
        the list of fonts you can use on your computer.
        font_color
            Color of the text. See ``TextClip.list('color')`` for a
        list of acceptable names.
        test : bool
            If true shows each image with caption in preview

        """

        VideoClip.__init__(self, ismask=False)
        image_duration+=transition_duration*2
        #FIXME: duration is wrong if there is panorama in images
        tt = np.cumsum([0] + list(np.repeat(image_duration, len(sequence))))
        #print (tt)
        self.images_starts = np.maximum(0, tt + (-1*transition_duration)*np.arange(len(tt)))
        self.duration = self.images_starts[-1]
        self.end = self.images_starts[-1]
        self.sequence = sequence
        #print(self.images_starts)
        #print ("DUR:", self.duration)
        def find_image_index(t):
            return max([i for i in range(len(self.sequence))
                              if self.images_starts[i]<=t])
        self.lastindex = None
        self.lastimage = None

        self.previndex = None
        self.previmage = None

        def load_clip(index):
            image = self.sequence[index]
            text = titles[index]
            if text.startswith("W:"):
                text = text[2:]
                show_full_height = True
            else:
                show_full_height = False
            if height is None and width is None:
                clip = ImageClip(image, duration=image_duration)
            else:
                if zoom_images:
                    clip = ImageClip(image, duration=image_duration) \
                            .fx(image_effect, screensize=(width, height), \
                            duration=20, show_full_height=show_full_height)
                elif show_full_height:
                    clip = ImageClip(image, duration=image_duration) \
                            .fx(resize, height=height).set_position('center',
                                    'center')
                    clip = CompositeVideoClip([clip],
                                            size=(width, height))

                else:
                    clip = ImageClip(image, duration=image_duration) \
                            .fx(resize, height=height, width=width)
            #Adds text label etc. on clip
            clip = make_clip(clip, text, height, width, font, font_color,
            fontsize)
            return clip

        def make_frame(t):

            index = find_image_index(t)
            fade = False
            #print ("INDEX:", index)
            #print ("TIME:", t, t-self.images_starts[index])
            clip_time = t- self.images_starts[index]
            #At the start of the clip we need to fade to the next
            if clip_time < transition_duration:
                fade = True
                #print ("FADE IN")

                if self.lastindex == index-1 and self.previndex != index-1:
                    self.previmage = self.lastimage
                    self.previndex = self.lastindex

                #print ("PREV IDX:", self.previndex)
            #When we are on last image we delete previous one from memory
            #Instead of 2 images in memory we have only one
            #Bu this can only happen on last image
            elif self.previndex and index == len(self.sequence)-1:
                #print ("LAST INDEX")
                del self.previmage
                self.previmage = None
                self.previndex = None

            #If we need to load new image we load it add the mask and add fade
            #if needed
            if index != self.lastindex:
                clip = load_clip(index)
                #print ("MASK:", clip.mask)
                if clip.mask:
                    #TODO: is mask actually needed outside fading?
                    clip.mask.duration = clip.duration
                    newclip = clip.copy()
                    newclip.mask = clip.mask.fx(fadein, transition_duration)
                else:
                    newclip = clip

                self.lastimage = newclip
                self.lastindex = index

            #Fading between two images
            if fade and self.previmage:
                image = self.lastimage.blit_on(
                        self.previmage.get_frame(t-self.images_starts[self.previndex]),
                        t-self.images_starts[self.lastindex])
            else:
                ##Fade at the start where we don't have previous image
                ##TODO: is this actually needed
                #if fade:
                    #self.set_mask(
                            #self.lastimage.mask)
                    ##print ("Fade MASK")
                #else:
                    #self.mask = None
                self.mask = None
                #Normal image without fading
                image = self.lastimage.get_frame(t-self.images_starts[index])
            #print ("DIFF:", self.duration-t, 1/25)
            #Deletes image from memory if we are closer to end of the clip as
            #last frame in 25 FPS
            #If we are wrong we will just load the image again
            if index == len(self.sequence)-1 and self.duration-t < 2/25:
                #print ("LAST FRAME")
                del self.lastimage
                self.lastimage = None
                self.lastindex = None
            return image

        self.make_frame = make_frame
        if width is not None and height is not None:
            self.size = (width, height)
        else:
            self.size = make_frame(0).shape[:2][::-1]

#@profile
def make_image_slideshow(sequence, titles, height=None, width=None, image_duration=4,
        transition_duration=1, fontsize=30, font="M+-1p-medium",
        font_color="white", zoom_images=False, test=False):
    """Function makes slideshow VideoClip from sequence of images with their
    captions

    Parameters
    ---------
    sequence
        List of paths to images
    titles : list
        List of image captions
    height
        Height of wanted VideoClip
    width
        Width of wanted VideoClip
    image_duration
        How long is one image visible
    transition_duration
        How long is fade transition between images
    fontsize
        Font point size
    font
        Name of the font to use. See ``TextClip.list('font')`` for
      the list of fonts you can use on your computer.
    font_color
        Color of the text. See ``TextClip.list('color')`` for a
      list of acceptable names.
    test : bool
        If true shows each image with caption in preview

    """
    #TODO: list of clip effects (so that panoramas can be panorama or specific
    #panorama)
    #TODO: support folder as a sequence
    images = sequence
    assert len(images) == len(titles), "Number of images and titles needs \
    to be the same"
    return SlideshowImagesClip(images, titles, height, width, image_duration,
            transition_duration, fontsize, font, font_color, zoom_images, test)

    def dump(self):
        ret_string = ["<slideshow>"]
        ret_string.append("size:{}x{} transition:{}s fontsize:{} "
                "font:{} font_color:{} zoom_images:{}".format(width, height,
                    transition_duration, fontsize, font, font_color,
                    zoom_images))
        ret_string.append("Duration:{}s".format(final_clip.duration))
        for image, title, clip in zip(sequence, titles, slided_clips):
            ret_string.append("{2: 6.2f}s {0} {1}".format(
                os.path.basename(image), title, clip.duration))
        ret_string.append("</slideshow>")
        return "\n".join(ret_string)

    final_clip.dump = types.MethodType(dump, final_clip)
    return final_clip


import os
from typing import List

from vkbottle import ABCRule, BaseStateGroup, VKAPIError
from vkbottle.bot import Bot, BotLabeler, Message

# A simple rule to demonstrate labeler
# setup for custom rules later
class SpamRule(ABCRule[Message]):
    def __init__(self, chars: List[str]):
        self.chars = "".join(chars)

    async def check(self, event: Message):
        return len(event.text) and event.text.strip(self.chars)

# Create a bot, or a single labeler:
# from vkbottle.bot import BotLabeler
# labeler = BotLabeler()

# Load token from system environment variable
# https://12factor.net/config
bot = Bot(os.environ["TOKEN"])

# Labeler can be accessed with bot.labeler
# or with bot.on (.on is property which returns
# .labeler, this shortcut is cute legacy from
# vkbottle 2.x

# This is first shortcut for VBMLRule from custom_rules
# <vbml_ignore_case = True> makes get_vbml_rule to add
# re.IGNORECASE to flags
bot.labeler.vbml_ignore_case = True  # type: ignore
# You can add default flags if ignore case is False
# <bot.labeler.default_flags = ...>

# We can add rule to custom_rules and it will be accessible
# in handlers in any place but is it of importance that
# labeler is always local (shortcuts work only for a local
# instance, for eg Bot, Blueprint, or pure Labeler)
bot.labeler.custom_rules["spam"] = SpamRule

# BotLabeler has fixed views. If you want to add yours you need
# to implement custom labeler, take it in account that labeler
# views are GLOBAL(!)
bot.labeler.views()  # {"message": MessageView, "raw": RawEventView}

bot.labeler.load(BotLabeler())  # Labeler can be loaded in another labeler

# Patcher for vbml rule shortcut can be set:
# <bot.labeler.patcher = ...>

# We will add some states
# The comments for states are skipped because
# we have another topic of the example
class SpamState(BaseStateGroup):
    GOOD = "good"
    BAD = "bad"

# Lets add some handlers
@bot.on.chat_message(spam=["!", ".", "?", "$", "#", "@", "%"])
async def spam_handler(message: Message):
    state_peer = await bot.state_dispenser.get(message.from_id)

    if state_peer and state_peer.state == SpamState.BAD:
        try:
            await bot.api.messages.remove_chat_user(message.chat_id, message.from_id)
            return "Как можно игнорировать мои просьбы"
        except VKAPIError[15]:
            return "Где мои права администратора?"

    await message.answer("Пожалуйста перестаньте спамить")
    await bot.state_dispenser.set(message.from_id, SpamState.BAD)

@bot.on.message(text="прости меня")
async def forgive_handler(message: Message):
    await bot.state_dispenser.set(message.from_id, SpamState.GOOD)
    return "Ладно, извинения приняты"

bot.run_forever()


import socket
import subprocess
import struct
import json

phone=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
# phone.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,1)
phone.bind(('127.0.0.1',9909)) #0-65535:0-1024给操作系统使用
phone.listen(5)

print('starting...')
while True: # 链接循环
    conn,client_addr=phone.accept()
    print(client_addr)

    while True: #通信循环
        try:
            #1、收命令
            cmd=conn.recv(8096)
            if not cmd:break #适用于linux操作系统

            #2、执行命令，拿到结果
            obj = subprocess.Popen(cmd.decode('utf-8'), shell=True,
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE)

            stdout=obj.stdout.read()
            stderr=obj.stderr.read()

            #3、把命令的结果返回给客户端
            #第一步：制作固定长度的报头
            header_dic={
                'filename':'a.txt',
                'md5':'xxdxxx',
                'total_size': len(stdout) + len(stderr)
            }

            header_json=json.dumps(header_dic)

            header_bytes=header_json.encode('utf-8')

            #第二步：先发送报头的长度
            conn.send(struct.pack('i',len(header_bytes)))

            #第三步：再发报头
            conn.send(header_bytes)

            #第四步：再发送真实的数据
            conn.send(stdout)
            conn.send(stderr)

        except ConnectionResetError: #适用于windows操作系统
            break
    conn.close()

phone.close()

3面向对象编程/面向对象&网络编程基础/网络编程入门&FTP服务开发实战/第二章网络编程/04 解决粘包问题/终极版/服务端.py


from django.contrib.auth.models import User
from django.test import TestCase, override_settings

from analysis.models import Analysis, CohortNode, CohortNodeZygosityFiltersCollection, FilterNode, FilterNodeItem, \
    SampleNode, GeneListNode, IntersectionNode, PhenotypeNode, PopulationNode, TagNode, VariantTag, \
    AlleleFrequencyNode, NodeAlleleFrequencyFilter, NodeVCFFilter, TrioNode, MOINode, MergeNode, PedigreeNode
from annotation.fake_annotation import get_fake_annotation_version, create_fake_variants
from annotation.models import VariantGeneOverlap, AnnotationRun
from annotation.tests.test_data_fake_genes import create_fake_transcript_version
from genes.models import GeneList, GeneListGeneSymbol
from ontology.models import OntologyTerm
from patients.models_enums import GnomADPopulation
from pedigree.models import PedigreeInheritance
from snpdb.models import GenomeBuild, ImportStatus, Variant, GenomicInterval, Tag
from snpdb.tests.test_data import create_fake_trio, create_fake_pedigree

@override_settings(ANALYSIS_NODE_CACHE_Q=False)
class TestCloneAnalysisNodes(TestCase):
    """ It's only worth testing AnalysisNodes with related objects, as normal fields are just copied OK """

    @classmethod
    def setUpTestData(cls):
        super().setUpTestData()

        user = User.objects.get_or_create(username='testuser')[0]
        cls.grch37 = GenomeBuild.get_name_or_alias("GRCh37")
        cls.annotation_version_grch37 = get_fake_annotation_version(cls.grch37)
        gene_annotation_release = cls.annotation_version_grch37.gene_annotation_version.gene_annotation_release
        cls.transcript_version = create_fake_transcript_version(cls.grch37,
                                                                release=gene_annotation_release)
        cls.gene_symbol = cls.transcript_version.gene_version.gene_symbol
        cls.trio = create_fake_trio(user, cls.grch37)
        cls.pedigree = create_fake_pedigree(user, cls.grch37)

        cls.analysis = Analysis(genome_build=cls.grch37)
        cls.analysis.set_defaults_and_save(user)

        cls.sample = cls.trio.get_samples()[0]

        # Gene List
        cls.gene_list = GeneList.objects.get_or_create(name="fake list",
                                                       user=cls.analysis.user,
                                                       import_status=ImportStatus.SUCCESS)[0]
        GeneListGeneSymbol.objects.get_or_create(gene_list=cls.gene_list, gene_symbol=cls.gene_symbol)

        # Need some overlapping variants so gene list will work
        create_fake_variants(cls.grch37)
        # Note: Variant probably doesn't overlap with gene, just want a random one
        variant = Variant.objects.filter(Variant.get_no_reference_q()).first()
        annotation_run = AnnotationRun.objects.create()
        VariantGeneOverlap.objects.create(version=cls.annotation_version_grch37.variant_annotation_version,
                                          annotation_run=annotation_run,
                                          gene=cls.transcript_version.gene,
                                          variant=variant)

        # Tag that variant
        cls.tag = Tag.objects.get_or_create(pk="foo")[0]
        VariantTag.objects.create(genome_build=cls.grch37, analysis=cls.analysis,
                                  variant=variant, tag=cls.tag, user=user)

    def _test_clone_node(self, node):
        clone = node.save_clone()
        node_arg_q_dict = str(node._get_node_arg_q_dict())
        clone_arg_q_dict = str(clone._get_node_arg_q_dict())
        class_name = node.get_class_name()
        msg = f"Clone of {class_name} has same get_q()"
        print(f"Type '{class_name}' Node: {node_arg_q_dict} <=> {clone_arg_q_dict}")
        self.assertEqual(node_arg_q_dict, clone_arg_q_dict, msg)

    @staticmethod
    def _set_af_for_node(node):
        naff = NodeAlleleFrequencyFilter.objects.get_or_create(node=node)[0]
        naff.nodeallelefrequencyrange_set.get_or_create(min=15, max=100)

    @staticmethod
    def _set_vcf_filter_for_node(node, vcf):
        vcf_filter = vcf.vcffilter_set.first()
        NodeVCFFilter.objects.create(node=node, vcf_filter=vcf_filter)

    def test_clone_allele_frequency_node(self):
        af_node = AlleleFrequencyNode.objects.create(analysis=self.analysis, sample=self.sample)
        self._set_af_for_node(af_node)
        self._test_clone_node(af_node)

    def _get_cohort_node(self):
        cohort = self.trio.cohort
        cohort_node = CohortNode.objects.create(analysis=self.analysis, cohort=cohort,
                                                accordion_panel=CohortNode.PER_SAMPLE_ZYGOSITY)

        self._set_af_for_node(cohort_node)
        self._set_vcf_filter_for_node(cohort_node, cohort.vcf)
        return cohort_node

    def _get_sample_node(self):
        return SampleNode.objects.create(analysis=self.analysis, sample=self.sample)

    def test_clone_cohort_node(self):
        cohort_node = self._get_cohort_node()
        fc = CohortNodeZygosityFiltersCollection.get_for_node_and_cohort(cohort_node, cohort_node.cohort)
        zf1 = fc.cohortnodezygosityfilter_set.order_by("pk").first()
        zf1.zygosity_hom = False
        zf1.save()

        zf2 = fc.cohortnodezygosityfilter_set.order_by("pk").last()
        zf2.zygosity_het = False
        zf2.save()

        self._test_clone_node(cohort_node)

    def test_clone_filter_node(self):
        filter_node = FilterNode.objects.create(analysis=self.analysis)
        FilterNodeItem.objects.create(filter_node=filter_node, sort_order=1,
                                      operation="eq", field="id", data="42")

        self._test_clone_node(filter_node)

    def test_clone_gene_list_node(self):
        gene_list_node = GeneListNode.objects.create(analysis=self.analysis)
        gene_list_node.genelistnodegenelist_set.create(gene_list_node=gene_list_node, gene_list=self.gene_list)
        # TODO: This isn't actually working properly
        self._test_clone_node(gene_list_node)

    def test_clone_intersection_node(self):
        # Doesn't have any related objects, but does need to make its own copy of GenomicInterval so test that
        genomic_interval = GenomicInterval.objects.create(chrom="1", start=10000, end=20000)
        intersection_node = IntersectionNode.objects.create(analysis=self.analysis, genomic_interval=genomic_interval)
        # Save the PK as the actual object gets changed (known issue - ok if we don't save the node)
        original_genomic_interval_id = genomic_interval.pk
        clone = intersection_node.save_clone()
        self.assertNotEqual(original_genomic_interval_id, clone.genomic_interval.pk,
                            "Intersection node made copy of genomic interval")

    def test_clone_merge_node(self):
        _merge_node = MergeNode.objects.create(analysis=self.analysis)
        # TODO: This doesn't run the test, as it calls _get_node_q which MergeNode doesn't implement
        # self._test_clone_node(merge_node)

    def test_clone_moi_node(self):
        # TODO: Need to also test patient setup w/ontology etc
        moi_node = MOINode.objects.create(analysis=self.analysis, sample=self.sample)

        ontology_term = OntologyTerm.objects.first()
        moi_node.moinodeontologyterm_set.create(ontology_term=ontology_term)
        self._test_clone_node(moi_node)

    def test_clone_pedigree_node(self):
        pedigree_node = PedigreeNode.objects.create(analysis=self.analysis, pedigree=self.pedigree,
                                                    inheritance_model=PedigreeInheritance.AUTOSOMAL_DOMINANT)
        self._set_af_for_node(pedigree_node)
        self._set_vcf_filter_for_node(pedigree_node, self.trio.cohort.vcf)
        self._test_clone_node(pedigree_node)

    def test_clone_phenotype_node(self):
        # TODO: Need to insert genes against the ontology term so that it actually does filtering
        phenotype_node = PhenotypeNode.objects.create(analysis=self.analysis)
        ontology_term = OntologyTerm.objects.first()
        phenotype_node.phenotypenodeontologyterm_set.create(ontology_term=ontology_term)
        self._test_clone_node(phenotype_node)

    def test_clone_population_node(self):
        population_node = PopulationNode.objects.create(analysis=self.analysis)
        population_node.populationnodegnomadpopulation_set.create(population=GnomADPopulation.NON_FINNISH_EUROPEAN)
        self._test_clone_node(population_node)

    def test_clone_sample_node(self):
        # Test allele frequency, VCF filters
        sample_node = self._get_sample_node()
        self._set_af_for_node(sample_node)
        self._set_vcf_filter_for_node(sample_node, self.sample.vcf)
        self._test_clone_node(sample_node)

    def test_clone_tag_node(self):
        tag_node = TagNode.objects.create(analysis=self.analysis)
        tag_node.tagnodetag_set.create(tag=self.tag)
        self._test_clone_node(tag_node)

    def test_clone_trio_node(self):
        trio_node = TrioNode.objects.create(analysis=self.analysis, trio=self.trio)
        self._set_af_for_node(trio_node)
        self._set_vcf_filter_for_node(trio_node, self.trio.cohort.vcf)
        self._test_clone_node(trio_node)


from scrapy.spider import Spider
from scrapy.selector import Selector

from fa_scraper.items import FreeAgent
import re

class FaSpider(Spider):
    name = "fa"
    allowed_domains = ["http://en.wikipedia.org"]
    start_urls = [
        "http://en.wikipedia.org/wiki/List_of_2006-07_NBA_season_transactions",
        "http://en.wikipedia.org/wiki/List_of_2007-08_NBA_season_transactions",
        "http://en.wikipedia.org/wiki/List_of_2008-09_NBA_season_transactions",
        "http://en.wikipedia.org/wiki/List_of_2009-10_NBA_season_transactions",
        "http://en.wikipedia.org/wiki/List_of_2010-11_NBA_season_transactions",
        "http://en.wikipedia.org/wiki/List_of_2011-12_NBA_season_transactions",
        "http://en.wikipedia.org/wiki/List_of_2012-13_NBA_season_transactions"
    ]

    title_re = re.compile(r'^List of 20(\d\d).*$')

    def parse(self, response):
        sel = Selector(response)
        title = sel.xpath('//h1/span[1]/text()').extract()[0]
        year = self.title_re.match(title).groups()[0]
        season = "{year1}-{year2}".format(year1=year
                ,year2=str(int(year)+1).zfill(2))
        if season == "06-07":
            names = (sel.xpath(
                '//h2/span[@id="Free_Agency"]/'
                '../../table[@style="text-align:center"][1]/tr/td[1]'
                '/a'))
        elif season == "07-08":
            names = (sel.xpath(
                '//h2/span[@id="Free_Agency"]/'
                '../../table[@style="text-align:left"][1]/tr/td'
                '/span/span/a'))
        elif season == "08-09":
            names = (sel.xpath(
                '//h3/span[@id="Free_agency"]/'
                '../../table[@style="text-align:center"][2]/tr/td[@align="left"]'
                '/a'))
        elif season == "09-10":
            names = (sel.xpath(
                '//h3/span[@id="Signed_from_free_agency"]/'
                '../../table[@style="text-align: center"][1]/tr/td[1]/a'))
        elif season == "10-11":
            names = (sel.xpath(
                '//h3/span[@id="Signed_from_free_agency"]/'
                '../../table[@style="text-align: center; width:80%"][1]/tr/td[1]/a'))
        elif season == "11-12" or season == "12-13":
            names = (sel.xpath(
                '//h3/span[@id="Free_agency"]/'
                '../../table[@style="text-align:center"][2]/tr/td/'
                'span/span/a'))
        for name in names:
            fa = FreeAgent()
            fa['Player'] = name.xpath('text()').extract()
            fa['Season'] = season
            team = name.xpath('../../td/a/text()')
            if team:
                team = team[1].extract()
            else:
                team = name.xpath('../../../../td/a/text()')
                if team:
                    team = team[0].extract()
            fa['Team'] =  team
            yield fa