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

# Lint as: python3
"""Funnel Transformer model.

Funnel-Transformer: Filtering out Sequential Redundancy for Efficient
Language Processing https://arxiv.org/abs/2006.03236
"""
import functools

from trax import fastmath
from trax import layers as tl
from trax.fastmath import numpy as jnp
from trax.fastmath.ops import index_add
from trax.layers.assert_shape import assert_shape
from trax.layers.research.rel_attention import get_rel_att_inputs
from trax.layers.research.rel_attention import  RelativeAttentionWrapper
from trax.layers.research.resampling import _FeedForwardBlock
from trax.layers.research.resampling import AttentionResampling
from trax.layers.research.resampling import AveragePooling
from trax.layers.research.resampling import LinearPooling
from trax.layers.research.resampling import LinearUpsampling
from trax.models.reformer.reformer import DecoderBlock
from trax.models.research.configurable_transformer import  ApplyAttentionLayer
from trax.models.research.configurable_transformer import PositionalEncoder
from trax.models.transformer import _EncoderBlock

@assert_shape('bld->bSd')
def PoolLayer(pool_layer=tl.AvgPool,
              pool_size=(2,),
              strides=(2,),
              separate_cls=True):
  """Returns a pool layer for Funnel Transformer.

  Args:
    pool_layer: Type of pooling layer used for downsampling;
        should be `tl.AvgPool` or `tl.MaxPool`.
    pool_size: Shape of window that gets reduced to a single vector value.
        If the layer inputs are :math:`n`-dimensional arrays, then `pool_size`
        must be a tuple of length :math:`n-2`.
    strides: Offsets from the location of one window to the locations of
        neighboring windows along each axis. If specified, must be a tuple of
        the same length as `pool_size`. If None, then offsets of 1 along each
        window axis, :math:`(1, ..., 1)`, will be used.
    separate_cls: If `True`, pooling in funnel blocks is not applied to
          embeddings of the first token (`cls` from BERT paper).
  """
  if separate_cls:
    cls_selection = tl.Fn('select_cls_token', lambda x: x[:, :1, :])
    tokens_after_cls = tl.Fn('rest_tokens', lambda x: x[:, 1:, :])

    return tl.Serial(
        tl.Branch(
            cls_selection,
            tl.Serial(
                tokens_after_cls,
                pool_layer(pool_size, strides)
            )
        ),
        tl.Concatenate(axis=1)
    )
  else:
    return pool_layer(pool_size, strides)

@assert_shape('b11l->b11S')
def MaskPool(pool_size=(2,), strides=(2,), separate_cls=True):
  return tl.Serial(
      tl.Fn('reshape', lambda x: x.swapaxes(1, -1).squeeze(axis=-1)),
      PoolLayer(tl.MaxPool, pool_size, strides, separate_cls),
      tl.Fn('reshape_back', lambda x: x[..., None].swapaxes(1, -1))
  )

@assert_shape('bld->bd')
def SelectFirst():
  return tl.Fn('select_first', lambda x: x[:, 0, :])

def _Upsampler(total_pool_size, separate_cls):
  """Returns an upsampling layer for Funnel Transformer.

  Args:
    total_pool_size: The combined pool size of previously used funnel blocks.
    separate_cls: If `True`, pooling in funnel blocks is not applied to
          embeddings of the first token (`cls` from BERT paper).
  """

  def _Upsample(short, long):
    if separate_cls:
      upsampled_short = jnp.concatenate(
          (short[:, :1, :],
           short[:, 1:, :].repeat(total_pool_size, axis=1)),
          axis=1)
      return index_add(
          long,
          (slice(None),
           slice(None, upsampled_short.shape[1]),
           slice(None)),
          upsampled_short)
    else:
      upsampled_short = short.repeat(total_pool_size, axis=1)
      return long + upsampled_short

  return tl.Fn('Upsampler', _Upsample)

def _FunnelBlock(d_model, d_ff, n_heads,
                 dropout, dropout_shared_axes, mode, ff_activation,
                 pool_layer, pool_size, strides, separate_cls):
  """Internal funnel block. Returns a list of layers implementing it.

  The input is an activation tensor.

  Args:
    d_model: Final dimension of tensors at most points in the model, including
        the initial embedding output.
    d_ff: Size of special dense layer in the feed-forward part of each block.
    n_heads: Number of attention heads.
    dropout: Stochastic rate (probability) for dropping an activation value
        when applying dropout within a block.
    dropout_shared_axes: Tensor axes on which to share a dropout mask.
        Sharing along batch and sequence axes (`dropout_shared_axes=(0,1)`) is
        a useful way to save memory and apply consistent masks to activation
        vectors at different sequence positions.
    mode: If `'train'`, each block will include dropout; else, it will
        pass all values through unaltered.
    ff_activation: Type of activation function at the end of each block; must
        be an activation-type subclass of `Layer`.
    pool_layer: Type of pooling layer used for downsampling;
        should be `tl.AvgPool` or `tl.MaxPool`.
    pool_size: Shape of window that gets reduced to a single vector value.
        If the layer inputs are :math:`n`-dimensional arrays, then `pool_size`
        must be a tuple of length :math:`n-2`.
    strides: Offsets from the location of one window to the locations of
        neighboring windows along each axis. If specified, must be a tuple of
        the same length as `pool_size`. If None, then offsets of 1 along each
        window axis, :math:`(1, ..., 1)`, will be used.
    separate_cls: If `True`, pooling in funnel blocks is not applied to
          embeddings of the first token (`cls` from BERT paper).
  Returns:
      A list of layers that maps (activations, mask) to (activations', mask).
  """
  pooling = PoolLayer(pool_layer, pool_size, strides, separate_cls)
  mask_pooling = MaskPool(pool_size, strides, separate_cls)

  attention = tl.AttentionQKV(d_model, n_heads=n_heads, dropout=dropout,
                              mode=mode)
  hidden_dropout = tl.Dropout(
      rate=dropout, shared_axes=dropout_shared_axes, mode=mode)

  feed_forward = _FeedForwardBlock(
      d_model, d_ff, dropout, dropout_shared_axes, mode, ff_activation)

  return [                                          # h, mask
      tl.LayerNorm(),                               # h, mask
      tl.Branch(pooling, None),                     # h', h, mask
      tl.Residual(
          tl.Select([0, 1, 1, 2]),                  # h', h, h, mask
          attention,                                # attn, mask
          tl.Parallel(None, mask_pooling),          # attn, mask'
          hidden_dropout                            # attn, mask'
      ),                                            # funnel_activations, mask'
      tl.Residual(
          tl.LayerNorm(),
          feed_forward,
          hidden_dropout,
      )
  ]

def FunnelTransformerEncoder(vocab_size,
                             n_classes=10,
                             d_model=512,
                             d_ff=2048,
                             encoder_segment_lengths=(2, 2, 2),
                             n_heads=8,
                             max_len=2048,
                             dropout=0.1,
                             dropout_shared_axes=None,
                             mode='train',
                             ff_activation=tl.Relu,
                             pool_layer=tl.AvgPool,
                             pool_size=(2,),
                             strides=(2,),
                             separate_cls=True):
  """Returns a Funnel Encoder.

  This model performs text categorization:

    - input: rank 2 tensor representing a batch of text strings via token IDs
      plus padding markers; shape is (batch_size, sequence_length). The tensor
      elements are integers in `range(vocab_size)`, and `0` values mark padding
      positions.

    - output: rank 2 tensor representing a batch of log-probability
      distributions over N categories; shape is (batch_size, `n_classes`).

  Args:
    vocab_size: Input vocabulary size -- each element of the input tensor
        should be an integer in `range(vocab_size)`. These integers typically
        represent token IDs from a vocabulary-based tokenizer.
    n_classes: Final dimension of the output tensors, representing N-way
        classification.
    d_model: Final dimension of tensors at most points in the model, including
        the initial embedding output.
    d_ff: Size of special dense layer in the feed-forward part of each encoder
        block.
    encoder_segment_lengths: Tuple, where each element denotes the number of
        transformer encoder blocks preceding a funnel transformer block.
        There is no funnel block after the last sequence of encoder blocks,
        therefore the total number of blocks in the model is equal to
        `sum(encoder_segment_lengths) + len(encoder_segment_lengths) - 1`.
    n_heads: Number of attention heads.
    max_len: Maximum symbol length for positional encoding.
    dropout: Stochastic rate (probability) for dropping an activation value
        when applying dropout within an encoder block.
    dropout_shared_axes: Tensor axes on which to share a dropout mask.
        Sharing along batch and sequence axes (`dropout_shared_axes=(0,1)`) is
        a useful way to save memory and apply consistent masks to activation
        vectors at different sequence positions.
    mode: If `'train'`, each encoder block will include dropout; else, it will
        pass all values through unaltered.
    ff_activation: Type of activation function at the end of each encoder
        block; must be an activation-type subclass of `Layer`.
    pool_layer: Type of pooling layer used for downsampling in each of the
        funnel blocks; should be `tl.AvgPool` or `tl.MaxPool`.
    pool_size: Shape of window that gets reduced to a single vector value.
        If the layer inputs are :math:`n`-dimensional arrays, then `pool_size`
        must be a tuple of length :math:`n-2`.
    strides: Offsets from the location of one window to the locations of
        neighboring windows along each axis. If specified, must be a tuple of
        the same length as `pool_size`. If None, then offsets of 1 along each
        window axis, :math:`(1, ..., 1)`, will be used.
    separate_cls: If `True`, pooling in funnel blocks is not applied to
        embeddings of the first token (`cls` from BERT paper) and only final
        embedding of this token is used for categorization - the rest are
        discarded. If `False`, each token from the beginning is pooled and
        all embeddings are averaged and mapped to output categories like in
        original `TransformerEncoder` model.
  Returns:
    A Transformer model that maps strings (conveyed via token IDs) to
    probability-like activations over a range of output classes.
  """
  assert encoder_segment_lengths

  positional_encoder = [
      tl.Embedding(vocab_size, d_model),
      tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode),
      tl.PositionalEncoding(max_len=max_len)]

  encoder_blocks = []
  n_encoder_segments = len(encoder_segment_lengths)

  for i in range(n_encoder_segments):
    # Building i'th segment
    for _ in range(encoder_segment_lengths[i]):
      # Create segment_size encoder blocks
      encoder_blocks.append(
          _EncoderBlock(d_model, d_ff, n_heads, dropout, dropout_shared_axes,
                        mode, ff_activation))

    # If not last segment, add funnel block
    if i != n_encoder_segments - 1:
      encoder_blocks.append(
          _FunnelBlock(d_model, d_ff, n_heads, dropout,
                       dropout_shared_axes, mode,
                       ff_activation, pool_layer, pool_size,
                       strides, separate_cls))

  cls_pooling = SelectFirst() if separate_cls else tl.Mean(axis=1)

  # Assemble and return the model.
  return tl.Serial(                               # toks
      # Encode.
      tl.Branch(
          positional_encoder, tl.PaddingMask()),  # vecs masks
      encoder_blocks,                             # vecs masks
      tl.Select([0], n_in=2),                     # vecs
      tl.LayerNorm(),                             # vecs

      # Map to output categories.
      cls_pooling,                                # cls
      tl.Dense(n_classes),                        # cls
  )

def FunnelTransformer(vocab_size,
                      d_model=512,
                      d_ff=2048,
                      encoder_segment_lengths=(2, 2, 2),
                      n_decoder_blocks=2,
                      n_heads=8,
                      max_len=2048,
                      dropout=0.1,
                      dropout_shared_axes=None,
                      mode='train',
                      ff_activation=tl.Relu,
                      pool_layer=tl.AvgPool,
                      pool_size=(2,),
                      separate_cls=True):
  """Returns a Full Funnel Transformer, that can be used for example for BERT.

  This model outputs token-level categorical distributions over all vocab:

    - input: rank 2 tensor representing a batch of text strings via token IDs
      plus padding markers; shape is (batch_size, sequence_length). The tensor
      elements are integers in `range(vocab_size)`, and `0` values mark padding
      positions.

    - output: rank 3 tensor representing a batch of log-probability
      distributions over `vocab_size` categories for each token; shape is
      (batch_size, sequence_length, vocab_size).

  Args:
    vocab_size: Input vocabulary size -- each element of the input tensor
        should be an integer in `range(vocab_size)`. These integers typically
        represent token IDs from a vocabulary-based tokenizer.
    d_model: Final dimension of tensors at most points in the model, including
        the initial embedding output.
    d_ff: Size of special dense layer in the feed-forward part of each encoder
        block.
    encoder_segment_lengths: Tuple, where each element denotes the number of
        transformer encoder blocks preceding a funnel transformer block.
        There is no funnel block after the last sequence of encoder blocks,
        therefore the total number of blocks in the model is equal to
        `sum(encoder_segment_lengths) + len(encoder_segment_lengths) - 1`.
    n_decoder_blocks: Number of transformer blocks in the upsampling decoder.
    n_heads: Number of attention heads.
    max_len: Maximum symbol length for positional encoding.
    dropout: Stochastic rate (probability) for dropping an activation value
        when applying dropout within an encoder block.
    dropout_shared_axes: Tensor axes on which to share a dropout mask.
        Sharing along batch and sequence axes (`dropout_shared_axes=(0,1)`) is
        a useful way to save memory and apply consistent masks to activation
        vectors at different sequence positions.
    mode: If `'train'`, each encoder block will include dropout; else, it will
        pass all values through unaltered.
    ff_activation: Type of activation function at the end of each encoder
        block; must be an activation-type subclass of `Layer`.
    pool_layer: Type of pooling layer used for downsampling in each of the
        funnel blocks; should be `tl.AvgPool` or `tl.MaxPool`.
    pool_size: Shape of window that gets reduced to a single vector value.
        If the layer inputs are :math:`n`-dimensional arrays, then `pool_size`
        must be a tuple of length :math:`n-2`.
    separate_cls: If `True`, pooling in funnel blocks is not applied to
        embeddings of the first token (`cls` from BERT paper) and only final
        embedding of this token is used for categorization - the rest are
        discarded. If `False`, each token from the beginning is pooled and
        all embeddings are averaged and mapped to output categories like in
        original `TransformerEncoder` model.
  """
  assert encoder_segment_lengths

  positional_encoder = [
      tl.Embedding(vocab_size, d_model),
      tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode),
      tl.PositionalEncoding(max_len=max_len)]

  n_encoder_segments = len(encoder_segment_lengths)

  encoder_blocks_before_first_pooling = [
      _EncoderBlock(d_model, d_ff, n_heads, dropout,
                    dropout_shared_axes, mode, ff_activation)
      for _ in range(encoder_segment_lengths[0])]
  encoder_blocks_from_first_pooling = []

  for i in range(1, n_encoder_segments):
    # Building i'th segment

    # Add funnel block between segments
    encoder_blocks_from_first_pooling.append(
        _FunnelBlock(d_model, d_ff, n_heads, dropout,
                     dropout_shared_axes, mode,
                     ff_activation, pool_layer,
                     pool_size=pool_size, strides=pool_size,
                     separate_cls=separate_cls))

    for _ in range(encoder_segment_lengths[i]):
      # Create segment_size encoder blocks
      encoder_blocks_from_first_pooling.append(
          _EncoderBlock(d_model, d_ff, n_heads, dropout,
                        dropout_shared_axes, mode, ff_activation))

  decoder_blocks = [_EncoderBlock(d_model, d_ff, n_heads, dropout,
                                  dropout_shared_axes, mode, ff_activation)
                    for _ in range(n_decoder_blocks)]

  total_pool_size = pool_size[0] ** (len(encoder_segment_lengths) - 1)

  # Assemble and return the model.
  return tl.Serial(                               # toks
      tl.Branch(
          positional_encoder, tl.PaddingMask()),  # vecs masks
      encoder_blocks_before_first_pooling,        # vecs masks
      tl.Select([0, 1, 0, 1]),
      # vecs masks residual = vecs old_masks
      encoder_blocks_from_first_pooling,          # vecs masks residual masks
      tl.Select([0, 2, 3]),                       # vecs residual masks
      tl.Parallel(
          # residual from first segment is taken before
          # normalization, so apply it now
          None, tl.LayerNorm(), None),            # vecs norm(residual) masks
      _Upsampler(total_pool_size, separate_cls),  # vecs masks
      decoder_blocks,
      tl.Select([0], n_in=2),                     # vecs
      tl.LayerNorm(),
      tl.Dense(vocab_size),
  )

def _RelativeDecoderBlock(attention_type, d_model,
                          d_ff, n_heads, dropout, dropout_shared_axes,
                          mode, ff_activation, context_bias_layer,
                          location_bias_layer, total_pooling):
  """Returns a list of layers that implements a Transformer encoder block.

  The input to the block is a pair, (activations, mask), where the mask was
  created from the original source tokens to prevent attending to the padding
  part of the input.

  Args:
    attention_type: attention type.
    d_model: Final dimension of tensors at most points in the model, including
      the initial embedding output.
    d_ff: Size of special dense layer in the feed-forward part of each block.
    n_heads: Number of attention heads.
    dropout: Stochastic rate (probability) for dropping an activation value
      when applying dropout within a block.
    dropout_shared_axes: Tensor axes on which to share a dropout mask.
      Sharing along batch and sequence axes (`dropout_shared_axes=(0,1)`) is
      a useful way to save memory and apply consistent masks to activation
      vectors at different sequence positions.
    mode: If `'train'`, each block will include dropout; else, it will
      pass all values through unaltered.
    ff_activation: Type of activation function at the end of each block; must
      be an activation-type subclass of `Layer`.
    context_bias_layer: context bias layer.
    location_bias_layer: location bias layer.
    total_pooling: The combined pool size of previously used funnel blocks.
  Returns:
    A list of layers that maps (activations, att_vecs, mask) to
                               (activations, att_vecs, mask).
  """
  if attention_type == RelativeAttentionWrapper:
    attention = RelativeAttentionWrapper(
        d_model,
        n_heads,
        dropout,
        mode=mode,
        context_bias_layer=context_bias_layer,
        location_bias_layer=location_bias_layer,
        total_pooling=total_pooling
    )
  else:
    attention = ApplyAttentionLayer(
        attention_type,
        d_model,
        n_heads,
        d_model // n_heads,
        d_model // n_heads,
        causal=True,
        masked=False,
        attention_dropout=dropout,
        output_dropout=dropout,
        attention_chunk_size=0,  # Disables tl.Chunk in ApplyAttentionLayer.
        mode=mode,
    )

  feed_forward = _FeedForwardBlock(
      d_model, d_ff, dropout, dropout_shared_axes, mode, ff_activation)

  def _Dropout():
    return tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode)

  return [
      tl.Residual(               # vecs
          tl.LayerNorm(),
          attention,
          _Dropout(),
      ),                         # vecs
      tl.Residual(
          tl.LayerNorm(),
          feed_forward,
          _Dropout(),
      ),                         # vecs
  ]

def RelformerLM(vocab_size,
                d_model=512,
                d_ff=2048,
                vanilla_layers=(0, 1),
                shorten_factors=(3,),
                n_funnel_blocks=(6,),
                n_heads=8,
                dropout=0.1,
                dropout_shared_axes=None,
                mode='train',
                ff_activation=tl.FastGelu,
                vanilla_attn_type=RelativeAttentionWrapper,
                middle_attn_type=RelativeAttentionWrapper,
                downsampling_fn=AttentionResampling,
                upsampling_fn=AttentionResampling,
                attention_downsampling_fn=AveragePooling,
                attention_upsampling_fn=LinearUpsampling):
  """Returns a Transformer language model.

  This model performs autoregressive language modeling:

    - input: rank 2 tensor representing a batch of text strings via token IDs
      plus padding markers; shape is (batch_size, sequence_length). The tensor
      elements are integers in `range(vocab_size)`, and `0` values mark padding
      positions.

    - output: rank 3 tensor representing a batch of log-probability
      distributions for each sequence position over possible token IDs;
      shape is (batch_size, sequence_length, `vocab_size`).

  This model uses only the decoder part of the overall Transformer.

  Args:
    vocab_size: Input vocabulary size -- each element of the input tensor
        should be an integer in `range(vocab_size)`. These integers typically
        represent token IDs from a vocabulary-based tokenizer.
    d_model: Final dimension of tensors at most points in the model, including
        the initial embedding output.
    d_ff: Size of special dense layer in the feed-forward part of each encoder
        block.
    vanilla_layers: (pre_layers, post_layers) tuple - number of full token-level
        Transformer decoder layers before and after shortening.
    shorten_factors: by how much to shorten at each step - tuple of arbitrary
        length denoting by how much shorten at each pooling stage.
    n_funnel_blocks: number of Transformer decoder blocks after each stage of
        pooling - tuple of the same length as `shorten_factors`.
    n_heads: Number of attention heads.
    dropout: Stochastic rate (probability) for dropping an activation value
        when applying dropout within an encoder block.
    dropout_shared_axes: Tensor axes on which to share a dropout mask.
        Sharing along batch and sequence axes (`dropout_shared_axes=(0,1)`) is
        a useful way to save memory and apply consistent masks to activation
        vectors at different sequence positions.
    mode: str: 'train' or 'eval'.
    ff_activation: Type of activation function at the end of each encoder
        block; must be an activation-type subclass of `Layer`.
    vanilla_attn_type: class: attention class such as SelfAttention to use in
        the layers before and after shortening (vanilla layers).
    middle_attn_type: class: attention class to use in the middle layers
        (these operating on the shortened sequence).
    downsampling_fn: function that takes full token-level vectors of
        length `l` and transforms them into `l` / `k` vectors, where `k`
        denotes `shorten_factor` parameter.
    upsampling_fn: function that takes shortened representations of a sequence,
        consisting of `l` / `k` vectors and transforms them into full
        token-level representations of length `l`.
    attention_downsampling_fn: Downsampling function that transforms token-level
        vectors into query vectors with reduced length. Necessary only when
        AttentionResampling is used as `downsampling_fn`.
    attention_upsampling_fn: Upsampling function for AttentionResampling.
        Valid only when AttentionResampling is used as a `upsampling_fn`.

  Returns:
    A Transformer language model as a layer that maps from a tensor of tokens
    to activations over a vocab set.
  """
  assert mode != 'predict'  # For now, 'predict' mode is unsupported.
  assert len(n_funnel_blocks) == len(shorten_factors)

  token_encoder = [
      tl.Embedding(vocab_size, d_model),
      tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode)]

  context_bias_layer, location_bias_layer = get_rel_att_inputs(d_model,
                                                               n_heads)

  n_pre_decoder_blocks, n_post_decoder_blocks = vanilla_layers

  def create_decoder_blocks(n_layers, total_pooling, attention_type):  # pylint: disable=invalid-name
    decoder_blocks = [
        # pylint: disable=g-complex-comprehension
        _RelativeDecoderBlock(attention_type, d_model, d_ff, n_heads, dropout,
                              dropout_shared_axes, mode, ff_activation,
                              context_bias_layer, location_bias_layer,
                              total_pooling)
        for _ in range(n_layers)]
    return decoder_blocks + [tl.LayerNorm()]

  total_pooling_acc = 1
  pre_decoder_blocks = create_decoder_blocks(n_pre_decoder_blocks,
                                             total_pooling_acc,
                                             vanilla_attn_type)

  funnel_blocks = []

  for shorten_factor, block_len in zip(shorten_factors, n_funnel_blocks):
    funnel_blocks = funnel_blocks + [
        downsampling_fn(shorten_factor, d_model,
                        is_upsampling=False, d_ff=d_ff, n_heads=n_heads,
                        dropout=dropout,
                        dropout_shared_axes=dropout_shared_axes, mode=mode,
                        ff_activation=ff_activation,
                        context_bias_layer=context_bias_layer,
                        location_bias_layer=location_bias_layer,
                        total_pooling=total_pooling_acc,
                        resampling_fn=attention_downsampling_fn)]
    total_pooling_acc *= shorten_factor
    funnel_blocks = funnel_blocks + create_decoder_blocks(
        block_len,
        total_pooling_acc,
        middle_attn_type)

  upsampling_layer = upsampling_fn(shorten_factor=total_pooling_acc,
                                   d_model=d_model, is_upsampling=True,
                                   d_ff=d_ff, n_heads=n_heads,
                                   dropout=dropout,
                                   dropout_shared_axes=dropout_shared_axes,
                                   mode=mode, ff_activation=ff_activation,
                                   context_bias_layer=context_bias_layer,
                                   location_bias_layer=location_bias_layer,
                                   total_pooling=total_pooling_acc,
                                   resampling_fn=attention_upsampling_fn)

  post_decoder_blocks = create_decoder_blocks(n_post_decoder_blocks, 1,
                                              vanilla_attn_type)

  # Assemble and return the model.
  return tl.Serial(              # tokens (or chunked tuple of tokens)
      tl.ShiftRight(mode=mode),  # toks
      token_encoder,             # vecs
      pre_decoder_blocks,        # vecs
      tl.Dup(),                  # vecs
      tl.ShiftRight(n_positions=total_pooling_acc - 1),  # shifted_vecs, vecs
      funnel_blocks,             # shifted_vecs, vecs
      upsampling_layer,          # vecs, vecs
      tl.LayerNorm(),            # vecs, vecs
      tl.Add(),                  # vecs
      post_decoder_blocks,       # vecs
      tl.Dense(vocab_size),      # vecs
  )

class RelformerCacher(tl.Layer):
  """Cache for Relformer.

  A class for caching tokens going through model to provide fast inference
  for Relformer model.
  """

  def __init__(self,
               total_kv_pooling,
               n_raw_tokens_generated=1,
               max_inference_length=64 * 64 * 3,
               shift=0,
               sliding=False,
               mode='train'):
    super().__init__(n_in=1, n_out=1)
    self._total_kv_pooling = total_kv_pooling
    self._n_raw_tokens_generated = n_raw_tokens_generated
    self._max_len = max_inference_length
    self._shift = shift
    self._sliding = sliding
    self._mode = mode

  def forward(self, inputs):
    if self._mode != 'predict':
      return inputs
    return self.update_state(inputs=inputs)

  def init_weights_and_state(self, input_signature):
    if self._mode == 'predict':
      shape, dtype = input_signature.as_tuple()
      batch_size, _, d_feature = shape
      cache = jnp.zeros((batch_size, 2 * self._total_kv_pooling, d_feature),
                        dtype=dtype)
      self.state = cache, jnp.array(0)

  def update_state(self, inputs):
    cache, idx = self.state
    cache = fastmath.dynamic_update_slice_in_dim(
        cache,
        inputs, (idx + self._shift) % (2 * self._total_kv_pooling),
        axis=1)

    if self._sliding:
      cache = fastmath.dynamic_update_slice_in_dim(
          cache,
          inputs,
          (idx + self._total_kv_pooling * 2 - 1) % (2 * self._total_kv_pooling),
          axis=1)

    if self._sliding:
      left_index = idx % self._total_kv_pooling
    else:
      left_index = (idx -
                    (idx % self._total_kv_pooling)) % (2 *
                                                       self._total_kv_pooling)

    output = fastmath.dynamic_slice(
        cache, [0, left_index, 0],
        [cache.shape[0], self._total_kv_pooling, cache.shape[2]])

    self.state = cache, idx + self._n_raw_tokens_generated
    return output

class RelformerPicker(tl.Layer):
  """Relformer Picker.

  A class for picking tokens going through model to provide fast inference
  for Relformer model.
  """

  def __init__(self, total_kv_pooling, n_raw_tokens_generated=1, mode='train'):
    super().__init__(n_in=1, n_out=1)
    self._total_kv_pooling = total_kv_pooling
    self._n_raw_tokens_generated = n_raw_tokens_generated
    self._mode = mode

  def forward(self, inputs):
    if self._mode != 'predict':
      return inputs

    output = fastmath.dynamic_slice(
        inputs, [0, self.state, 0],
        [inputs.shape[0], self._n_raw_tokens_generated, inputs.shape[2]])
    self.state = (self.state +
                  self._n_raw_tokens_generated) % self._total_kv_pooling
    return output

  def init_weights_and_state(self, input_signature):
    if self._mode == 'predict':
      self.state = jnp.array(0)

def PickLastTokenInPredict(mode='train'):
  """Picks the last token logits.

  Self-descriptive layer for picking the last token logits in predict mode
  for fast inference.

  Args:
    mode: the model mode (train, predict, ...)

  Returns:
    The last token logits.
  """

  def last_token(x):  # pylint: disable=invalid-name
    if mode == 'predict':
      return x[:, -1:, :]
    return x

  return tl.Fn('Pick last token in predict', last_token)

def RelformerChunkedLM(vocab_size,
                       d_model=512,
                       d_ff=2048,
                       vanilla_layers=(1, 1),
                       shorten_factor=3,
                       n_rel_layers=6,
                       rel_chunk_len=None,
                       vanilla_chunk_len=None,
                       n_heads=8,
                       dropout=0.1,
                       dropout_shared_axes=None,
                       vanilla_attn_type=tl.LSHSelfAttention,
                       pos_type='fixed-base',
                       max_len=3072,
                       n_raw_tokens_generated=1,
                       mode='train',
                       ff_activation=tl.FastGelu):
  """Returns a Transformer language model.

  This model performs autoregressive language modeling:

    - input: rank 2 tensor representing a batch of text strings via token IDs
      plus padding markers; shape is (batch_size, sequence_length). The tensor
      elements are integers in `range(vocab_size)`, and `0` values mark padding
      positions.

    - output: rank 3 tensor representing a batch of log-probability
      distributions for each sequence position over possible token IDs;
      shape is (batch_size, sequence_length, `vocab_size`).

  This model uses only the decoder part of the overall Transformer.

  Args:
    vocab_size: Input vocabulary size -- each element of the input tensor
        should be an integer in `range(vocab_size)`. These integers typically
        represent token IDs from a vocabulary-based tokenizer.
    d_model: Final dimension of tensors at most points in the model, including
        the initial embedding output.
    d_ff: Size of special dense layer in the feed-forward part of each encoder
        block.
    vanilla_layers: (pre_layers, post_layers) tuple - number of full token-level
        Transformer decoder layers before and after shortening.
    shorten_factor: by how much to shorten
    n_rel_layers: number of Transformer blocks after the pooling. These blocks
        use relative attention.
    rel_chunk_len (optional): Number of tokens per chunk. Setting this option
        will enable chunked relative attention.
    vanilla_chunk_len (optional): If set, enables chunked relative attention
        also in layers before and after shortening.
    n_heads: Number of attention heads.
    dropout: Stochastic rate (probability) for dropping an activation value
        when applying dropout within an encoder block.
    dropout_shared_axes: Tensor axes on which to share a dropout mask.
        Sharing along batch and sequence axes (`dropout_shared_axes=(0,1)`) is
        a useful way to save memory and apply consistent masks to activation
        vectors at different sequence positions.
    vanilla_attn_type: class: attention class such as SelfAttention to use in
        the layers before and after shortening (vanilla layers).
    pos_type: string, the type of positional embeddings to use.
    max_len: int: maximum symbol length both for positional encoding and it is
      also the maximum length of the possible inference in 'predict' mode
    n_raw_tokens_generated: int: number of tokens generated with every pass
      through model in 'predict' mode. Number of tokens should be smaller and
      divisible by the first shorten factor we are using in the model.
      It cannot be larger than one if we use vanilla layers because we would
      lose autoregressive property of the model.
    mode: str: 'train' or 'eval' or 'predict'.
    ff_activation: Type of activation function at the end of each encoder
        block; must be an activation-type subclass of `Layer`.

  Returns:
    A Transformer language model as a layer that maps from a tensor of tokens
    to activations over a vocab set.
  """

  token_encoder = [
      tl.Embedding(vocab_size, d_model),
      tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode)]

  if vanilla_chunk_len is None:
    positional_encoder = PositionalEncoder(mode, dropout, max_len, pos_type)
  else:
    positional_encoder = []

  n_pre_decoder_blocks, n_post_decoder_blocks = vanilla_layers

  def create_reformer_blocks(  # pylint: disable=invalid-name
      n_layers,
      total_kv_pooling=1,
      layer_chunk_len=None,
      force_relative=False,
      dense=True):
    if n_layers == 0:
      return [tl.LayerNorm()]

    def determine_attn_type(layer_number):  # pylint: disable=invalid-name
      if layer_chunk_len is None and not force_relative:
        return vanilla_attn_type

      if layer_chunk_len is not None:
        chunk_offset = (layer_number % 2) * (layer_chunk_len // 2)
      else:
        chunk_offset = None

      return functools.partial(
          RelativeAttentionWrapper,
          n_raw_tokens_generated=n_raw_tokens_generated,
          max_inference_length=max_len,
          total_kv_pooling=total_kv_pooling,
          chunk_len=layer_chunk_len,
          chunk_offset=chunk_offset)

    d_per_head = d_model // n_heads

    decoder_blocks = []
    for i in range(n_layers):
      layer_attn_type = determine_attn_type(i)

      decoder_blocks.append(
          DecoderBlock(
              d_model,
              d_ff,
              d_per_head,
              d_per_head,
              n_heads,
              layer_attn_type,
              dropout,
              ff_activation,
              dropout,
              ff_use_sru=0,
              ff_chunk_size=0,
              ff_sparsity=0,
              attention_chunk_size=0,
              mode=mode))

    return [
        tl.Dup(),
        tl.ReversibleSerial(decoder_blocks),
        tl.Concatenate(),
        tl.LayerNorm(),
        tl.Dense(d_model) if dense else [],
    ]

  pre_decoder_blocks = create_reformer_blocks(
      n_pre_decoder_blocks, layer_chunk_len=vanilla_chunk_len)

  relative_decoder_blocks = create_reformer_blocks(
      n_rel_layers,
      total_kv_pooling=shorten_factor,
      layer_chunk_len=rel_chunk_len,
      force_relative=True)

  conv_layer = tl.Serial(
      tl.CausalConv(d_model, shorten_factor),
      ff_activation()
  )

  post_decoder_blocks = create_reformer_blocks(
      n_post_decoder_blocks, layer_chunk_len=vanilla_chunk_len, dense=False)

  cacher = RelformerCacher(
      total_kv_pooling=shorten_factor,
      n_raw_tokens_generated=n_raw_tokens_generated,
      max_inference_length=max_len,
      shift=shorten_factor - 1,
      mode=mode)

  picker = RelformerPicker(
      total_kv_pooling=shorten_factor,
      n_raw_tokens_generated=n_raw_tokens_generated,
      mode=mode)

  cacher_conv = RelformerCacher(
      total_kv_pooling=shorten_factor,
      n_raw_tokens_generated=n_raw_tokens_generated,
      max_inference_length=max_len,
      shift=shorten_factor - 1,
      sliding=True,
      mode=mode)

  picker_conv = PickLastTokenInPredict(mode=mode)

  # Assemble and return the model.
  return tl.Serial(  # tokens (or chunked tuple of tokens)
      tl.ShiftRight(mode=mode),  # toks
      token_encoder,  # vecs
      positional_encoder,
      pre_decoder_blocks,  # vecs
      tl.Dup(),
      cacher,
      tl.ShiftRight(n_positions=shorten_factor - 1, mode=mode),
      LinearPooling(shorten_factor, d_model),
      relative_decoder_blocks,
      tl.Dropout(rate=dropout, shared_axes=[-2], mode=mode),
      LinearUpsampling(shorten_factor, d_model),
      tl.LayerNorm(),
      picker,
      tl.Concatenate(),
      cacher_conv,
      conv_layer,
      picker_conv,
      post_decoder_blocks,
      tl.Dense(vocab_size),  # vecs
  )


from django.core.wsgi import get_wsgi_application
import wsgiserver

def start_server( address="0.0.0.0", port=8080):

    """
    server = wsgiserver.WSGIServer(
        (address, port),  # Use '127.0.0.1' to only bind to the localhost
        get_wsgi_application()
    )
    """

    server = wsgiserver.WSGIServer(get_wsgi_application(), host=address, port=port)

    try:
        server.start()
    except KeyboardInterrupt:
        print('Stopping server')
        server.stop()


from splinter import Browser
from bs4 import BeautifulSoup as bs
import pandas as pd

def scrape():
    executable_path = {'executable_path': 'chromedriver.exe'}
    browser = Browser('chrome', **executable_path, headless=False)
    news_title, news_p= mars_news(browser)
    results = {
        "title": news_title,
        "paragraph": news_p,
        "image_url": image(browser),
        "weather": mars_weather(browser),
        "facts": mars_facts(),
        "hemispheres": mars_hemisphere(browser),

    }
    browser.quit()
    print(results)
    return results

def mars_news(browser):
    url = 'https://mars.nasa.gov/news/'
    browser.visit(url)
    html = browser.html
    mars_news_soup = bs(html, 'html.parser')

    # Scrape the first article title and teaser paragraph text; return them
    title = mars_news_soup.find('div', {"class":"content_title"})
    news_title = title.text.strip()
    #print(news_title)
    paragraph = mars_news_soup.find('div', {'class':'article_teaser_body'})
    news_p = paragraph.string.strip()
    #print(news_p)
    return news_title, news_p

def image(browser):
    img_url = 'https://www.jpl.nasa.gov/spaceimages/'
    browser.visit(img_url)
    browser.click_link_by_partial_text('FULL IMAGE')
    expand = browser.find_by_css('a.fancybox-expand')
    expand.click()
    html = browser.html
    soup = bs(html, 'html.parser')
    img = soup.find('img', class_='fancybox-image')['src']
    # pass the path so i can see the actual image
    featured_image_url = f'https://www.jpl.nasa.gov{img}'
    return featured_image_url

def mars_weather(browser):
    url = 'https://twitter.com/marswxreport?lang=en'
    browser.visit(url)
    html = browser.html
    tweet_soup = bs(html, 'html.parser')

    # Scrape the tweet info and return
    first_tweet = tweet_soup.find('p', class_='TweetTextSize').text
    return first_tweet

def mars_facts():
    url = 'https://space-facts.com/mars/'
    tables = pd.read_html(url)
    df = tables[0]
    #df.columns = ['Element', 'Facts']
     # Convert to HTML table string and return
    return df.to_html()

def mars_hemisphere(browser):
    hemisphere_url = 'https://astrogeology.usgs.gov/search/results?q=hemisphere+enhanced&k1=target&v1=Mars'
    hem_dict = []
    for i in range(1,9,2):
        hemi_dict = {}
        browser.visit(hemisphere_url)
        html = browser.html
        soup = bs(html, 'html.parser')
        links = soup.findAll('a', class_='product-item')
        hemi_name = links[i].text
        #print(hemi_name)
        link_detail = browser.find_by_css('a.product-item')
        link_detail[i].click()
        browser.find_link_by_text('Sample').first.click()
        browser.windows.current = browser.windows[-1]
        html2 = browser.html
        browser.windows.current = browser.windows[0]
        browser.windows[-1].close()
        img_soup = bs(html2, 'html.parser')
        img_path = img_soup.find('img')['src']

        hemi_dict['title'] = hemi_name.strip('Enhanced')
        hemi_dict['img_url'] = img_path

        hem_dict.append(hemi_dict)
        #print(results)
    return hem_dict
scrape()


import torch
import torchvision
from torch import nn
import torch.nn.functional as F
from torch.autograd import Variable

IMAGE_SIZE = 784
IMAGE_WIDTH = IMAGE_HEIGHT = 28

class autoencoder(nn.Module):
    def __init__(self, input_dim=28*28, code_dim=32):
        super(autoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.ReLU(True),
            nn.Linear(128, 64),
            nn.ReLU(True),
            nn.Linear(64, code_dim),
            nn.Tanh()
        )
        self.decoder = nn.Sequential(
            nn.Linear(code_dim, 64),
            nn.ReLU(True),
            nn.Linear(64, 128),
            nn.ReLU(True),
            nn.Linear(128, input_dim),
            nn.Tanh()
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

class conv_autoencoder_with_tanh(nn.Module):

    def __init__(self, code_size):
        super(conv_autoencoder_with_tanh, self).__init__()
        self.code_size = code_size

        # Encoder specification
        self.enc_cnn_1 = nn.Conv2d(1, 10, kernel_size=5)
        self.enc_cnn_2 = nn.Conv2d(10, 20, kernel_size=5)
        self.enc_linear_1 = nn.Linear(4 * 4 * 20, 50)
        self.enc_linear_2 = nn.Linear(50, self.code_size)

        # Decoder specification
        self.dec_linear_1 = nn.Linear(self.code_size, 160)
        self.dec_linear_2 = nn.Linear(160, IMAGE_SIZE)

    def forward(self, images):
        code = self.encode(images)
        out = self.decode(code)
        return out, code

    def encode(self, images):
        code = self.enc_cnn_1(images)
        code = F.relu(F.max_pool2d(code, 2))

        code = self.enc_cnn_2(code)
        code = F.relu(F.max_pool2d(code, 2))

        code = code.view([images.size(0), -1])
        code = F.relu(self.enc_linear_1(code))
        code = self.enc_linear_2(code)
        code = nn.Tanh()(code)
        return code

    def decode(self, code):
        out = F.relu(self.dec_linear_1(code))
        out = nn.Tanh()(self.dec_linear_2(out))
        out = out.view([code.size(0), 1, IMAGE_WIDTH, IMAGE_HEIGHT])
        return out

class conv_autoencoder_with_sigmoid(nn.Module):

    def __init__(self, code_size):
        super(conv_autoencoder_with_sigmoid, self).__init__()
        self.code_size = code_size

        # Encoder specification
        self.enc_cnn_1 = nn.Conv2d(1, 10, kernel_size=5)
        self.enc_cnn_2 = nn.Conv2d(10, 20, kernel_size=5)
        self.enc_linear_1 = nn.Linear(4 * 4 * 20, 50)
        self.enc_linear_2 = nn.Linear(50, self.code_size)

        # Decoder specification
        self.dec_linear_1 = nn.Linear(self.code_size, 160)
        self.dec_linear_2 = nn.Linear(160, IMAGE_SIZE)

    def forward(self, images):
        code = self.encode(images)
        out = self.decode(code)
        return out, code

    def encode(self, images):
        code = self.enc_cnn_1(images)
        code = F.selu(F.max_pool2d(code, 2))

        code = self.enc_cnn_2(code)
        code = F.selu(F.max_pool2d(code, 2))

        code = code.view([images.size(0), -1])
        code = F.selu(self.enc_linear_1(code))
        code = self.enc_linear_2(code)
        return code

    def decode(self, code):
        out = F.selu(self.dec_linear_1(code))
        out = F.sigmoid(self.dec_linear_2(out))
        out = out.view([code.size(0), 1, IMAGE_WIDTH, IMAGE_HEIGHT])
        return out

class autoencoder_GMMN(nn.Module):
    def __init__(self, input_dim=28*28, code_dim=32):
        super(autoencoder_GMMN, self).__init__()
        self.encoder_fc1 = nn.Linear(input_dim, 1024)
        self.encoder_fc2 = nn.Linear(1024, code_dim)

        self.decoder_fc1 = nn.Linear(code_dim, 1024)
        self.decoder_fc2 = nn.Linear(1024, input_dim)

    def forward(self, x):
        e = self.encode(x)
        d = self.decode(e)
        return e, d

    def encode(self, x):
        e = F.sigmoid(self.encoder_fc1(x))
        e = F.sigmoid(self.encoder_fc2(e))
        return e

    def decode(self, x):
        d = F.sigmoid(self.decoder_fc1(x))
        d = F.sigmoid(self.decoder_fc2(d))
        return d

class autoencoder_cifar10(nn.Module):
    def __init__(self, code_dim=32):
        super(autoencoder_cifar10, self).__init__()
        self.code_dim = code_dim
        # Input size: [batch, 3, 32, 32]
        # Output size: [batch, 3, 32, 32]
        self.encoder_cnn = nn.Sequential(
            nn.Conv2d(3, 12, 4, stride=2, padding=1),            # [batch, 12, 16, 16]
            nn.ReLU(),
            nn.Conv2d(12, 24, 4, stride=2, padding=1),           # [batch, 24, 8, 8]
            nn.ReLU(),
			#nn.Conv2d(12, 24, 4, stride=2, padding=1),           # [batch, 48, 4, 4]
            #nn.ReLU(),
 			nn.Conv2d(24, 48, 4, stride=2, padding=1),           # [batch, 96, 2, 2]
            nn.ReLU(),
            #nn.Conv2d(48, 96, 4, stride=2, padding=1),           # [batch, 96, 2, 2]
            #nn.ReLU(),
        )
        #self.encoder_linear = nn.Linear(48 * 4 * 4, self.code_dim)
        #self.decoder_linear = nn.Linear(self.code_dim, 48 * 4 * 4)
        self.decoder_cnn = nn.Sequential(
            #nn.ConvTranspose2d(96, 48, 4, stride=2, padding=1),  # [batch, 48, 4, 4]
            #nn.ReLU(),
            nn.ConvTranspose2d(48, 24, 4, stride=2, padding=1),  # [batch, 48, 4, 4]
            nn.ReLU(),
			#nn.ConvTranspose2d(24, 12, 4, stride=2, padding=1),  # [batch, 24, 8, 8]
            #nn.ReLU(),
			nn.ConvTranspose2d(24, 12, 4, stride=2, padding=1),  # [batch, 12, 16, 16]
            nn.ReLU(),
            nn.ConvTranspose2d(12, 3, 4, stride=2, padding=1),   # [batch, 3, 32, 32]
            nn.Sigmoid(),
            #nn.Tanh()
        )

    def encoder(self, images):
        code = self.encoder_cnn(images)
        code = code.view(images.size(0), -1)
        #code = self.encoder_linear(code)
        code = nn.Sigmoid()(code)
        #code = nn.Tanh()(code)
        return code

    def decoder(self, code):
        out = code.view(code.size(0), 48, 4, 4)
        #out = nn.ReLU()(self.decoder_linear(code))
        #out = out.view(code.size(0), 48, 4, 4)
        out = self.decoder_cnn(out)
        return out

    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return encoded, decoded


import http.client
import json

conn = http.client.HTTPSConnection("api.razorpay.com")
payload = json.dumps({
  "amount": 680000,
  "currency": "INR",
  "accept_partial": True,
  "first_min_partial_amount": 100,
  "expire_by": 55791273360000,
  "reference_id": "#9937",
  "description": "cloud computing",
  "customer": {
    "name": "Gaurav Kumar",
    "contact": "+919999999999",
    "email": "gaurav.kumar@example.com"
  },
  "notify": {
    "sms": True,
    "email": True
  },
  "reminder_enable": True,
  "notes": {
    "policy_name": "Jeevan Bima"
  },
  "callback_url": "https://example-callback-url.com/",
  "callback_method": "get"
})
headers = {
  'Content-Type': 'application/json',
  'Authorization': 'Basic cnpwX3Rlc3RfaDlFSG1HdU5HcGJsM2Y6aXdYN2ZhcWljaE9ReWprcVAxS2ZCbFJo'
}
conn.request("POST", "/v1/payment_links", payload, headers)
res = conn.getresponse()
data = res.read()
print(data.decode("utf-8"))


from django.template.loader import render_to_string
from django.utils.text import slugify
from quantaly.models import User

def create_user(strategy, details, backend, user=None, *args, **kwargs):
    if user:
        return {'is_new': False}

    fields = dict((name, kwargs.get(name, details.get(name)))
                  for name in backend.setting('USER_FIELDS', ['email', 'username']))
    if not fields:
        return

    if fields['email'] is None or fields['email'] == '':
        fields['email'] = kwargs.get('response')['id'] + "@" + backend.name + ".com"

    auth_user = strategy.create_user(**fields)
    auth_user.is_verify = True
    auth_user.save()

    message = render_to_string('quantaly/mail/welcome.html', {
        'user': auth_user,
    })
    auth_user.email_user('Welcome to Quantaly.', message)

    return {
        'is_new': True,
        'user': auth_user
    }


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

from collections import defaultdict
import tempfile

from enum import Enum
import errno
import os
import semantic_version
import shutil
import yaml

from solar import errors
from solar import utils

RES_TYPE = Enum("Resource Types", 'Normal Composer')

class RepositoryException(errors.SolarError):
    pass

class ResourceNotFound(RepositoryException):

    def __init__(self, spec):
        self.message = 'Resource definition %r not found' % spec

    def __str__(self):
        return str(self.message)

def read_meta(base_path):
    base_meta_file = os.path.join(base_path, 'meta.yaml')

    metadata = utils.yaml_load(base_meta_file)
    metadata.setdefault('version', '1.0.0')
    # NOTE(jnowak): when `version: 0.1` then it's float
    metadata['version'] = str(metadata['version'])
    metadata['base_path'] = os.path.abspath(base_path)
    actions_path = os.path.join(metadata['base_path'], 'actions')
    metadata['actions_path'] = actions_path
    metadata['base_name'] = os.path.split(metadata['base_path'])[-1]

    return metadata

class RepositoryExists(RepositoryException):
    pass

class Repository(object):

    db_obj = None
    _REPOS_LOCATION = '/var/lib/solar/repositories'
    _TMP_DIRNAME = '.tmp'

    def __init__(self, name):
        self.name = name
        # TODO: (jnowak) sanitize name
        self.fpath = self.repo_path(self.name)

    @classmethod
    def _list_source_contents(cls, source):
        if source.endswith('.yaml'):
            # single CR
            pth = os.path.split(source)[-1][:-5]
            return ((RES_TYPE.Composer, pth, source), )
        elif os.path.isdir(source):
            meta_path = os.path.join(source, 'meta.yaml')
            if os.path.exists(meta_path):
                # single normal
                pth = os.path.split(source)[-1]
                valid = semantic_version.validate(pth)
                if not valid:
                    name = pth
                else:
                    # if it was semver then single_path may look like
                    # /a/b/name/version
                    # and source may look like /a/b/name
                    # we can extract name from this path then
                    name = source.split(os.path.sep)[-2]
                return ((RES_TYPE.Normal, name, source), )
            return tuple(cls._list_source_contents_from_multidir(source))

    @classmethod
    def _list_source_contents_from_multidir(cls, source):
        for pth in os.listdir(source):
            single_path = os.path.join(source, pth)
            if pth.endswith('.yaml'):
                pth = pth[:-5]
                yield RES_TYPE.Composer, pth, single_path
            elif os.path.exists(os.path.join(single_path, 'meta.yaml')):
                valid = semantic_version.validate(pth)
                if not valid:
                    name = pth
                else:
                    # if it was semver then single_path may look like
                    # /a/b/name/version
                    # and source may look like /a/b/name
                    # we can extract name from this path then
                    name = os.path.split(source)[-1]
                yield RES_TYPE.Normal, name, single_path
            else:
                maybe_cr = os.path.join(single_path,
                                        "{}.yaml".format(
                                            os.path.split(source)[-1]))
                if os.path.exists(maybe_cr):
                    name = os.path.split(source)[-1]
                    yield RES_TYPE.Composer, name, maybe_cr
                    continue
                if not os.path.isdir(single_path):
                    continue
                for single in os.listdir(single_path):
                    valid = semantic_version.validate(single)
                    if not valid:
                        fp = os.path.join(single_path, single)
                        raise RepositoryException(
                            "Unexpected repository content "
                            "structure: {} ""Expected directory "
                            "with version number".format(single_path))
                    else:
                        fp = os.path.join(single_path, single)
                        if os.path.exists(os.path.join(fp, 'meta.yaml')):
                            yield RES_TYPE.Normal, pth, fp
                        elif os.path.exists(
                                os.path.join(fp, '{}.yaml'.format(pth))):
                            cr = os.path.join(fp, '{}.yaml'.format(pth))
                            yield RES_TYPE.Composer, pth, cr

    @classmethod
    def repo_path(cls, repo_name):
        return os.path.join(cls._REPOS_LOCATION, repo_name)

    def create(self, source=None, link_only=False):
        if source is None:
            os.mkdir(self.fpath)
            return
        if not link_only:
            if os.path.isdir(self.fpath):
                raise RepositoryExists("Repository %s "
                                       "already exists" % self.name)

            if not os.path.isdir(self.tmp_dir):
                os.makedirs(self.tmp_dir)

            old_fpath = self.fpath
            self.fpath = tempfile.mkdtemp(dir=self.tmp_dir)
            try:
                self._add_contents(source)
                os.rename(self.fpath, old_fpath)
            except Exception as e:
                shutil.rmtree(self.fpath)
                raise
            finally:
                self.fpath = old_fpath
        else:
            try:
                os.symlink(source, self.fpath)
            except OSError as e:
                if e.errno == errno.EEXIST:
                    raise RepositoryExists("Repository %s "
                                           "already exists" % self.name)
                else:
                    raise

    def update(self, source, overwrite=False):
        self._add_contents(source, overwrite)

    def _add_contents(self, source, overwrite=False):
        cnts = self._list_source_contents(source)
        for res_type, single_name, single_path in cnts:
            if res_type is RES_TYPE.Normal:
                self.add_single_normal(single_name, single_path, overwrite)
            else:
                self.add_single_cr(single_name, single_path, overwrite)

    def add_single(self, name, source, overwrite=False):
        if os.path.isfile(source):
            name = name.replace('.yaml', '')
            return self.add_single_cr(name, source, overwrite)
        return self.add_single_normal(name, source, overwrite)

    def add_single_normal(self, name, source, overwrite=False):
        try:
            metadata = read_meta(source)
        except IOError as e:
            if e.errno == errno.ENOENT:
                raise RepositoryException(
                    "meta.yaml not found: %s" % e.filename)
            raise
        version = metadata['version']
        valid = semantic_version.validate(version)
        if not valid:
            raise RepositoryException("Invalid version "
                                      "%s for %s" % (version,
                                                     source))
        # TODO: (jnowak) sanitize version
        target_path = os.path.join(self.fpath, name, version)
        try:
            shutil.copytree(source, target_path, symlinks=True)
        except OSError as e:
            if e.errno != errno.EEXIST:
                raise
            if not overwrite:
                raise
            shutil.rmtree(target_path)
            shutil.copytree(source, target_path, symlinks=True)

    def add_single_cr(self, name, source, overwrite=False):
        with open(source, 'rb') as f:
            parsed = yaml.safe_load(f.read())
        version = parsed.get('version', '1.0.0')
        target_dir = os.path.join(self.fpath, name, version)
        target_path = os.path.join(target_dir, "{}.yaml".format(name))
        if not os.path.exists(target_dir):
            os.makedirs(target_dir)
        try:
            shutil.copy(source, target_path)
        except OSError as e:
            if e.errno != errno.EEXIST:
                raise
            if not overwrite:
                raise
            shutil.rm(target_dir)
            shutil.copy(source, target_path)

    def remove(self):
        shutil.rmtree(self.fpath)

    def remove_single(self, spec):
        spec = self._parse_spec(spec)
        if spec['version_sign'] != '==':
            raise RepositoryException("Removal possible only with `==` sign")
        path = self._make_version_path(spec)
        shutil.rmtree(path)
        return True

    def iter_contents(self, resource_name=None):

        def _single(single_path):
            try:
                for version in os.listdir(os.path.join(self.fpath,
                                                       single_path)):
                    yield {"name": single_path,
                           'version': version}
            except OSError:
                return

        if resource_name is None:
            for single in os.listdir(self.fpath):
                for gen in _single(single):
                    yield gen
        else:
            for gen in _single(resource_name):
                yield gen

    def get_contents(self, resource_name=None):
        out = defaultdict(list)
        cnt = self.iter_contents(resource_name)
        for curr in cnt:
            out[curr['name']].append(curr['version'])
        return out

    @classmethod
    def _parse_spec(cls, spec):
        if isinstance(spec, dict):
            return spec
        if ':' in spec:
            repos, version = spec.split(':', 1)
        else:
            repos = spec
            version = None
        if '/' in repos:
            repo_name, resource_name = repos.split('/', 1)
        else:
            repo_name = 'resources'
            resource_name = repos
        if version is None:
            version_sign = ">="
        elif '>=' in version or '<=' in version or '==' in version:
            version_sign = version[:2]
            version = version[2:]
        elif '>' in version or '<' in version:
            version_sign = version[:1]
            version = version[1:]
        else:
            version_sign = '=='
        return {'repo': repo_name,
                'resource_name': resource_name,
                'version': version,
                'version_sign': version_sign}

    def _get_version(self, spec):
        spec = self._parse_spec(spec)
        version = spec['version']
        version_sign = spec['version_sign']
        resource_name = spec['resource_name']
        if version_sign == '==':
            return os.path.join(self.fpath, spec['resource_name'], version)
        found = self.iter_contents(resource_name)
        if version is None:
            sc = semantic_version.compare
            sorted_vers = sorted(found,
                                 cmp=lambda a, b: sc(a['version'],
                                                     b['version']),
                                 reverse=True)
            if not sorted_vers:
                raise ResourceNotFound(spec)
            version = sorted_vers[0]['version']
        else:
            version = '{}{}'.format(version_sign, version)
            matched = filter(lambda x: semantic_version.match(version,
                                                              x['version']),
                             found)
            sorted_vers = sorted(matched,
                                 cmp=lambda a, b: semantic_version.compare(
                                     a['version'],
                                     b['version']),
                                 reverse=True)
            version = next((x['version'] for x in sorted_vers
                            if semantic_version.match(version, x['version'])),
                           None)
        if version is None:
            raise ResourceNotFound(spec)
        return version

    def _make_version_path(self, spec, version=None):
        spec = self._parse_spec(spec)
        if version is None:
            version = self._get_version(spec)
        return os.path.join(self.fpath, spec['resource_name'], version)

    def read_meta(self, spec):
        path = self.get_path(spec)
        return read_meta(path)

    def get_path(self, spec):
        spec = self._parse_spec(spec)
        return self._make_version_path(spec)

    @property
    def tmp_dir(self):
        return os.path.join(self._REPOS_LOCATION, self._TMP_DIRNAME)

    @classmethod
    def get_metadata(cls, spec):
        spec = cls._parse_spec(spec)
        repo = Repository(spec['repo'])
        return repo.read_meta(spec)

    @classmethod
    def contains(cls, spec):
        repo, spec = cls.parse(spec)
        try:
            version = repo._get_version(spec)
            path = repo._make_version_path(spec, version=version)
        except ResourceNotFound:
            return False
        return os.path.exists(path)

    @classmethod
    def what_version(cls, spec):
        repo, spec = cls.parse(spec)
        try:
            version = repo._get_version(spec)
            path = repo._make_version_path(spec, version=version)
        except ResourceNotFound:
            return False
        if not os.path.exists(path):
            return False
        return version

    @classmethod
    def list_repos(cls):
        return filter(
            lambda x: (os.path.isdir(os.path.join(cls._REPOS_LOCATION, x))
                       and x != cls._TMP_DIRNAME),
            os.listdir(cls._REPOS_LOCATION)
        )

    @classmethod
    def parse(cls, spec):
        spec = cls._parse_spec(spec)
        return Repository(spec['repo']), spec

    def is_composer_file(self, spec):
        return os.path.exists(self.get_composer_file_path(spec))

    def get_composer_file_path(self, spec):
        spec = self._parse_spec(spec)
        p = self.get_path(spec)
        return os.path.join(p, "{}.yaml".format(spec['resource_name']))


import sys

import json
from requests_toolbelt import MultipartEncoder
import requests
import mimetypes

def main(argv=None):
	if argv is None:
		argv = sys.argv
		if len(argv) is 1:
			argv.append(input("Type in URL of the media you are trying to transcribe:\n"))
		elif len(argv) > 2:
			print("Too many arguments!\n")

		media_url = argv[1]
		process_speech_to_text(media_url)

###############
##
###############
def process_speech_to_text(media_url):
	# CREDENTIALS
	with open("credentials.json") as cred_file:
		credentials = json.load(cred_file)

	audio = open(media_url, "rb")
	content_type = mimetypes.guess_type(media_url, strict=True)[0]

	s2t_endpoint = credentials["speech-to-text"]["url"] + "/v1/recognize"
	auth = (credentials["speech-to-text"]["username"], credentials["speech-to-text"]["password"])
	headers = {"Content-Type": content_type}
	params = {"timestamps": "true", "profanity_filter":"false"}

	r = requests.post(s2t_endpoint, data=audio, headers=headers, auth=auth, params=params)

	print(r.json())

if __name__ == "__main__":
    sys.exit(main())


#(c) 2016 by Authors
#This file is a part of ABruijn program.
#Released under the BSD license (see LICENSE file)

"""
Runs assemble binary
"""

import subprocess
import logging
import os

from abruijn.utils import which

ASSEMBLE_BIN = "abruijn-assemble"
logger = logging.getLogger()

class AssembleException(Exception):
    pass

def check_binaries():
    if not which(ASSEMBLE_BIN):
        raise AssembleException("Assemble binary was not found. "
                                "Did you run 'make'?")
    try:
        devnull = open(os.devnull, "w")
        subprocess.check_call([ASSEMBLE_BIN, "-h"], stderr=devnull)
    except (subprocess.CalledProcessError, OSError) as e:
        if e.returncode == -9:
            logger.error("Looks like the system ran out of memory")
        raise AssembleException(str(e))

def assemble(args, out_file, log_file):
    logger.info("Assembling reads")
    logger.debug("-----Begin assembly log------")
    cmdline = [ASSEMBLE_BIN, "-k", str(args.kmer_size), "-l", log_file,
               "-t", str(args.threads), "-v", str(args.min_overlap)]
    if args.debug:
        cmdline.append("-d")
    if args.min_kmer_count is not None:
        cmdline.extend(["-m", str(args.min_kmer_count)])
    if args.max_kmer_count is not None:
        cmdline.extend(["-x", str(args.max_kmer_count)])
    cmdline.extend([args.reads, out_file, str(args.coverage)])

    try:
        subprocess.check_call(cmdline)
    except (subprocess.CalledProcessError, OSError) as e:
        if e.returncode == -9:
            logger.error("Looks like the system ran out of memory")
        raise AssembleException(str(e))


import geopandas
from osgeo import gdal, gdal_array
import numpy as np

traceResult={}
polluting= geopandas.read_file('../Data/geoJson/pollution.geojson')

raster="../resultData/computingResult/TN_experience_formula.tif"
srcArray = gdal_array.LoadFile(raster)

srcArray[np.isnan(srcArray)]=-999

n=10
maxten=[]
maxtenposition=[]
for i in range(n):
    bm = np.nanmax(srcArray)
    pos = np.unravel_index(np.argmax(srcArray), srcArray.shape)
    maxten.append(bm)
    maxtenposition.append(pos)
    srcArray[pos[0],pos[1]]=-999

dataset = gdal.Open(raster)

adfGeoTransform = dataset.GetGeoTransform()

nXSize = dataset.RasterXSize #列数
nYSize = dataset.RasterYSize #行数

arrSlope = [] # 用于存储每个像素的（X，Y）坐标

i=0
polluting_list=[]
for i in range(polluting.shape[0]):
    polluting_list.append([])
    polluting_list[-1].append(i+1)
    polluting_list[-1].append(polluting.iat[i,1])
    polluting_list[-1].append(polluting.iat[i,-1].x)
    polluting_list[-1].append(polluting.iat[i,-1].y)

i=0
for i in range(len(maxtenposition)):
    num=maxtenposition[i]
    x=num[0]
    y=num[1]
    px = adfGeoTransform[0] + y * adfGeoTransform[1] + x * adfGeoTransform[2]
    py = adfGeoTransform[3] + y * adfGeoTransform[4] + x * adfGeoTransform[5]

    distance_list=[]
    for m in polluting_list:
        dis=pow((m[2]-px),2)+pow((m[3]-py),2)
        distance_list.append(dis)

    id=distance_list.index(min(distance_list))+1
    traceResult[id]=0

print(traceResult.keys())


import matplotlib.pyplot as plt
import random

"""
Script for drawing lines demo
"""

# create data
y = []
bins = []

for i in range(500):
    bins.append(i)
    y.append(random.randrange(50))

plt.hist(y, bins)
plt.show()


from unittest import TestCase, mock

import pytest
from django.contrib.auth import get_user_model
from rest_framework.authtoken.models import Token

from custom_auth.oauth.serializers import GoogleLoginSerializer, OAuthLoginSerializer, GithubLoginSerializer

UserModel = get_user_model()

@pytest.mark.django_db
class OAuthLoginSerializerTestCase(TestCase):
    def setUp(self) -> None:
        self.test_email = "testytester@example.com"
        self.test_first_name = "testy"
        self.test_last_name = "tester"
        self.test_id = "test-id"
        self.mock_user_data = {
            "email": self.test_email,
            "first_name": self.test_first_name,
            "last_name": self.test_last_name,
            "google_user_id": self.test_id
        }

    @mock.patch("custom_auth.oauth.serializers.get_user_info")
    def test_create(self, mock_get_user_info):
        # Given
        access_token = "access-token"
        data = {
            "access_token": access_token
        }
        serializer = OAuthLoginSerializer(data=data)

        # monkey patch the get_user_info method to return the mock user data
        serializer.get_user_info = lambda: self.mock_user_data

        # When
        serializer.is_valid()
        response = serializer.save()

        # Then
        assert UserModel.objects.filter(email=self.test_email).exists()
        assert isinstance(response, Token)
        assert response.user.email == self.test_email

class GoogleLoginSerializerTestCase(TestCase):
    @mock.patch("custom_auth.oauth.serializers.get_user_info")
    def test_get_user_info(self, mock_get_user_info):
        # Given
        access_token = "some-access-token"
        serializer = GoogleLoginSerializer(data={"access_token": access_token})

        # When
        serializer.is_valid()
        serializer.get_user_info()

        # Then
        mock_get_user_info.assert_called_with(access_token)

class GithubLoginSerializerTestCase(TestCase):
    @mock.patch("custom_auth.oauth.serializers.GithubUser")
    def test_get_user_info(self, MockGithubUser):
        # Given
        access_token = "some-access-token"
        serializer = GithubLoginSerializer(data={"access_token": access_token})

        mock_github_user = mock.MagicMock()
        MockGithubUser.return_value = mock_github_user

        # When
        serializer.is_valid()
        serializer.get_user_info()

        # Then
        MockGithubUser.assert_called_with(code=access_token)
        mock_github_user.get_user_info.assert_called()


# -*- encoding:ascii -*-
from mako import runtime, filters, cache
UNDEFINED = runtime.UNDEFINED
__M_dict_builtin = dict
__M_locals_builtin = locals
_magic_number = 6
_modified_time = 1399385602.205022
_template_filename='templates/webapps/galaxy/dataset/embed.mako'
_template_uri='dataset/embed.mako'
_template_cache=cache.Cache(__name__, _modified_time)
_source_encoding='ascii'
_exports = ['render_item_links', 'render_summary_content']

# SOURCE LINE 2

from galaxy.web.framework.helpers import iff

def _mako_get_namespace(context, name):
    try:
        return context.namespaces[(__name__, name)]
    except KeyError:
        _mako_generate_namespaces(context)
        return context.namespaces[(__name__, name)]
def _mako_generate_namespaces(context):
    pass
def _mako_inherit(template, context):
    _mako_generate_namespaces(context)
    return runtime._inherit_from(context, u'/embed_base.mako', _template_uri)
def render_body(context,**pageargs):
    context.caller_stack._push_frame()
    try:
        __M_locals = __M_dict_builtin(pageargs=pageargs)
        __M_writer = context.writer()
        # SOURCE LINE 1
        __M_writer(u'\n')
        # SOURCE LINE 4
        __M_writer(u'\n\n')
        # SOURCE LINE 12
        __M_writer(u'\n\n')
        # SOURCE LINE 19
        __M_writer(u'\n')
        return ''
    finally:
        context.caller_stack._pop_frame()

def render_render_item_links(context,dataset):
    context.caller_stack._push_frame()
    try:
        h = context.get('h', UNDEFINED)
        trans = context.get('trans', UNDEFINED)
        item = context.get('item', UNDEFINED)
        __M_writer = context.writer()
        # SOURCE LINE 6
        __M_writer(u'\n    \n')
        # SOURCE LINE 9
        __M_writer(u'    \n    \n    \n')
        return ''
    finally:
        context.caller_stack._pop_frame()

def render_render_summary_content(context,dataset,data):
    context.caller_stack._push_frame()
    try:
        __M_writer = context.writer()
        # SOURCE LINE 14
        __M_writer(u'\n')
        return ''
    finally:
        context.caller_stack._pop_frame()


from django.shortcuts import render
from django.http import HttpResponse
from django.shortcuts import redirect
from .models import Friend
from .forms import FriendForm

def index(request):
    data = Friend.objects.all()
    params = {
            'title': 'Hello',
            'data': data,
        }
    return render(request, 'hello/index.html' , params)

#create model
def create(request):
    if (request.method == 'POST'):
        obj = Friend()
        friend = FriendForm(request.POST, instance=obj)
        friend.save()
        return redirect(to='/hello')
    params = {
        'title': 'Hello',
        'form': FriendForm(),
    }

    return render(request, 'hello/create.html' , params)


import setuptools
import scyjava_config
from os import path

here = path.abspath(path.dirname(__file__))

with open(path.join(here, 'README.md')) as f:
    scyjava_long_description = f.read()

setuptools.setup(
    name='scyjava',
    python_requires='>=3',
    packages=['scyjava', 'scyjava.convert'],
    py_modules=['scyjava_config'],
    version=scyjava_config.version,
    author='Philipp Hanslovsky, Curtis Rueden',
    author_email='hanslovskyp@janelia.hhmi.org',
    description='scyjava',
    long_description=scyjava_long_description,
    long_description_content_type='text/markdown',
    license='Public domain',
    url='https://github.com/scijava/scyjava',
    install_requires=['pyjnius', 'jgo'],
)


from config import get_config
from Learner import Learner
from PatchLearner import PatchLearner, PatchLearnerMult
import argparse
import torch
from functools import partial
from torch.nn import MSELoss
from Pearson import pearson_corr_loss, ncl_loss

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='for CBIS-DDSM')
    parser.add_argument("-e", "--epochs", help="training epochs", default=20, type=int)
    parser.add_argument('-lr', '--lr', help='learning rate', default=1e-3, type=float)
    parser.add_argument("-b", "--batch_size", help="batch size", default=96, type=int)
    parser.add_argument("-w", "--num_workers", help="number of workers", default=3, type=int)
    parser.add_argument("-d", "--data_mode", help="databases: [crop_data, whole_data]", default='crop_data', type=str)
    parser.add_argument("-s", "--epoch_per_save", help="save_every s epochs", default=1, type=int)
    parser.add_argument("-net", "--net_mode", help="choose net", default='resnet50', type=str)

    parser.add_argument("-n", "--n_models", help="how many duplicate nets to use. 1 leads to basic training, "
                                                 "making -a and -p flags redundant", default=1, type=int)
    parser.add_argument("-patch", "--n_patch", help="how many patches per image", default=2, type=int)
    parser.add_argument("-bkg", "--bkg_prob", help="how many patches per image", default=.5, type=float)
    parser.add_argument("-mul", "--mul", help="use mult mode?", default=0, type=int)
    parser.add_argument("-rank", "--local_rank", help="rank for mul", default=0, type=int)
    parser.add_argument("-pre", "--pre_train", help="use a pretrain net?", default=0, type=int)
    parser.add_argument("-pre_layers", "--pre_layers", help="layer steps to use?", default=[], type=int, nargs='*')
    parser.add_argument("-pre_step", "--pre_steps", help="what steps to use?", default=[3, 10, 37], type=int, nargs='*')
    parser.add_argument("-ngpu", "--ngpu", help="how many gpu's to use?", default=1, type=int)
    parser.add_argument("-half", "--half", help="use half precisions?", default=0, type=int)
    parser.add_argument("-no_bkg", "--no_bkg", help="4 class mode", default=0, type=int)
    parser.add_argument("-cancer", "--cancer_only", help="limit label to benign/malignant", default=0, type=int)
    parser.add_argument("-type", "--type_only", help="limit label to mass/calcification", default=0, type=int)
    parser.add_argument("-single_type", "--single_type", help="limit labels to mass=1 or calcification=2", default=0, type=int)

    # TODO maybe add option to specify a network mix instead of duplicates
    parser.add_argument("-m", "--milestones", help="fractions of where lr will be tuned", default=[], type=int, nargs='*')
    parser.add_argument("-a", "--alpha", help="balancing parameter", default=0, type=float)
    parser.add_argument("-t", "--sig_thresh", help="thresholding of the most correct class", default=0.9, type=float)
    parser.add_argument("-p", "--pearson", help="using pearson loss", default=False, type=bool)
    parser.add_argument("-ncl", "--ncl", help="using Negative Correlation Loss", default=False, type=bool)
    parser.add_argument("-mean", "--joint_mean", help="using mean loss", default=False, type=bool)
    parser.add_argument("-morph_dir", "--morph_dir", help="use a morph directory", default='', type=str)
    parser.add_argument("-morph_a", "--morph_alpha", help="balance parameter", default=10., type=float)

    parser.add_argument("-logdir", "--logdir", help="extend log/saves to a folder for group experiment", default='', type=str)

    parser.add_argument("-c", "--cpu_mode", help="force cpu mode", default=0, type=int)

    args = parser.parse_args()
    conf = get_config(logext=args.logdir)

    # training param
    assert not (args.cancer_only and args.type_only)  # choose at most one
    conf.single_type = args.single_type
    conf.cancer_only = args.cancer_only
    conf.type_only = args.type_only
    conf.no_bkg = args.no_bkg
    conf.half = args.half
    conf.ngpu = args.ngpu
    conf.pre_layers = args.pre_layers
    conf.pre_steps = args.pre_steps
    conf.pre_train = args.pre_train
    conf.local_rank = args.local_rank
    conf.n_patch = args.n_patch
    conf.bkg_prob = args.bkg_prob if not args.no_bkg else 0
    conf.net_mode = args.net_mode
    conf.evaluate_every = 3  # TODO see if relevant
    conf.epoch_per_save = args.epoch_per_save
    conf.data_mode = args.data_mode
    conf.cpu_mode = args.cpu_mode
    conf.device = torch.device("cuda" if (torch.cuda.is_available() and not conf.cpu_mode) else "cpu")
    conf.lr = args.lr
    conf.batch_size = args.batch_size
    conf.num_workers = args.num_workers
    conf.epochs = args.epochs
    conf.milestones = args.milestones

    # pearson param
    conf.alpha = args.alpha
    conf.sig_thresh = args.sig_thresh
    conf.n_models = args.n_models
    conf.pearson = args.pearson
    conf.joint_mean = args.joint_mean
    conf.ncl = args.ncl

    # morph param
    conf.morph_alpha = args.morph_alpha
    conf.morph_dir = args.morph_dir

    # loss funcs
    conf.pearson_loss = partial(pearson_corr_loss, threshold=conf.sig_thresh)
    conf.ncl_loss = partial(ncl_loss)
    conf.morph_loss = MSELoss()

    # create learner and go
    param_desc = '_'.join([
        str(conf.net_mode), 'lr='+str(conf.lr), 'm='+'_'.join([str(m) for m in conf.milestones]),
        ('a='+str(conf.alpha) if conf.n_models>1 else ''),
        str(conf.batch_size), str(conf.n_patch), 'type_only' if conf.type_only else '',
                              'cancer_only' if conf.cancer_only else ''] +
        ([] if not conf.pre_train else
         ['pre', 'pre_layers='+'_'.join([str(m) for m in conf.pre_layers]),
          'pre_steps='+'_'.join([str(m) for m in conf.pre_steps])])
        )
    conf.log_path = str(conf.log_path) + '_' + param_desc
    conf.save_path = str(conf.save_path) + '_' + param_desc

    learner = PatchLearnerMult(conf) if args.mul else PatchLearner(conf)
    # face_learner(conf) if conf.n_models == 1 else face_learner_corr(conf)

    if conf.pre_train:
        learner.pretrain(conf)
    else:
        learner.train(conf, conf.epochs)


# -*- coding: utf-8 -*-
from itertools import chain

from django.utils.html import conditional_escape
from django.utils.encoding import force_unicode
from django.utils.safestring import mark_safe
from django import forms
from django.forms.widgets import CheckboxSelectMultiple, CheckboxInput
from shop.models import OrderExtraInfo

from models import Trailer

class OrderExtraInfoForm(forms.ModelForm):
    class Meta:
        model = OrderExtraInfo
        fields = ('text',)

class CustomCheckboxSelectMultiple(CheckboxSelectMultiple):
    def render(self, name, value, attrs=None, choices=()):
        if value is None: value = []
        has_id = attrs and 'id' in attrs
        final_attrs = self.build_attrs(attrs, name=name)
        output = [u'']
        # Normalize to strings
        str_values = set([force_unicode(v) for v in value])
        for i, (option_value, option_label) in enumerate(
        chain(self.choices, choices)):
            # If an ID attribute was given, add a numeric index as a suffix,
            # so that the checkboxes don't all have the same ID attribute.
            if has_id:
                final_attrs = dict(final_attrs, id='%s_%s' % (attrs['id'], i))
                label_for = u' for="%s"' % final_attrs['id']
            else:
                label_for = ''

            cb = CheckboxInput(final_attrs,
                check_test=lambda value: value in str_values)
            option_value = force_unicode(option_value)
            rendered_cb = cb.render(name, option_value)
            option_label = conditional_escape(force_unicode(option_label))
            output.append(u'<label%s>%s %s</label>' % (
                label_for, rendered_cb, option_label))
        return mark_safe(u'\n'.join(output))

class TrailerSearchForm(forms.Form):
    SHORT = u'1'
    LONG = u'2'
    LENGTH_CHOICES = (
        (SHORT, 'от 1,5 до 3 метров'),
        (LONG, 'от 3 до 6 метров'),
    )

    SMALL = u'1'
    BIG = u'2'
    CAPACITY_CHOICES = (
        (SMALL, 'до 8 тонн'),
        (BIG, 'от 8 тонн'),
    )

    NO_BRAKES = u'0'
    WITH_BRAKES = u'1'
    BRAKES_CHOICES = (
        (NO_BRAKES, 'без тормозов'),
        (WITH_BRAKES, 'с тормозами'),
    )

    number_axis = forms.MultipleChoiceField(Trailer.NUMBER_AXIS_CHOICES, False,
        CustomCheckboxSelectMultiple, 'Количество осей')
    suspension = forms.MultipleChoiceField(Trailer.SUSPENSION_CHOICES, False,
        CustomCheckboxSelectMultiple, 'Подвеска')
    length = forms.MultipleChoiceField(LENGTH_CHOICES, False,
        CustomCheckboxSelectMultiple, 'Длина платформы')
    capacity = forms.MultipleChoiceField(CAPACITY_CHOICES, False,
        CustomCheckboxSelectMultiple, 'Грузоподъемность')
    availability_of_brakes = forms.MultipleChoiceField(BRAKES_CHOICES, False,
        CustomCheckboxSelectMultiple, 'Наличие тормозов')


import unittest
from collections import namedtuple

from p1844 import class_solution

class MyTestCase(unittest.TestCase):
    def setUp(self) -> None:
        self.solution = class_solution.Solution()
        self.Case = namedtuple('Case', 's expect')
        # example cases
        self.cases_example = [
            self.Case(s="a1c1e1", expect="abcdef"),
            self.Case(s="a1b2c3d4e", expect="abbdcfdhe"),
        ]

    # replaceDigits1
    def test_examples1(self):
        for i, Case in enumerate(self.cases_example):
            with self.subTest(f'replaceDigits1: case: {i}'):
                act = self.solution.replaceDigits1(Case.s)
                self.assertEqual(act, Case.expect)

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


import numpy as np
import pandas as pd
import copy, os
os.chdir('/Users/eloncha/Documents/GitHub/Elon/python')

### Kmeans
def loadDataSet(fileName):
    dataMat = []
    fr = open(fileName)
    for line in fr.readlines():
        curLine = line.strip().split('\t')
        floatLine = map(float,curLine)
        dataMat.append(list(floatLine))
    return dataMat
def loadDataSetInPandas(fileName, by = '\t', is_header = None):
    dataMat = pd.read_table(fileName, sep = by, header = is_header)
    colNum = dataMat.shape[1]
    newColName = []
    for i in range(0,colNum):
        newColName.append("X{0}".format(str(i+1)))
    dataMat.columns = newColName
    return dataMat
def distMeasure(vecA, vecB, order = 2):
    dist = np.sqrt(np.sum(np.abs(np.power(vecA - vecB, order))))
    return dist #np.linalg.norm(vecA-vecB)
def proximityMatrix(dataMat, order=2):
    n = dataMat.shape[0]
    Mat = np.zeros((n,n))
    for i in range(n):
        for j in range(n):
            Mat[i,j] = distMeasure(dataMat.iloc[i,:], dataMat.iloc[j,:])
    return Mat
def assignCent(dataMat, k):
    n = dataMat.shape[1]
    centroids = np.zeros((k,n))
    for i in range(0,n):
        xMax = np.max(dataMat.iloc[:,i])
        xMin = np.min(dataMat.iloc[:,i])
        xRand = np.random.uniform(xMin, xMax, k)
        for j in range(0,k):
            centroids[j,i] = xRand[j]
    return centroids
def kMeans(dataMat, k):
    m = dataMat.shape[0]
    p = dataMat.shape[1]
    centroids = assignCent(dataMat, k)
    clusterChanged = True
    iteTime = 0
    print('Random Centroids assigned.')
    print(centroids)
    clusterDoc = copy.copy(dataMat)
    clusterDoc['minIndex'] = -1
    clusterDoc['minDist'] = np.Inf
    while clusterChanged:
        iteTime += 1
        clusterChanged = False
        for i in range(0, m): # loop through each data point to assign it to the nearest centroid
            minDist, minIndex = np.Inf, -1
            distList = np.zeros((1, k))
            for j in range(0, k):
                curDist = distMeasure(centroids[j], dataMat.iloc[i, :], order=2)
                distList[0, j] = curDist
                minIndex = np.argmin(distList)
                minDist = np.min(distList)
            if minDist < clusterDoc.iloc[i, p+1]:
                clusterChanged = True
                clusterDoc.iloc[i, p+1] = minDist
                clusterDoc.iloc[i, p] = minIndex
        for cent in range(0, k): # loop through each clusters to update centroids
            centroids[cent] = np.mean(clusterDoc[clusterDoc['minIndex'] == cent].iloc[:, 0:p], axis=0)
        print('Cluster Assignment Changed.')
        print('Iteration {0}'.format(str(iteTime)))
    return centroids, clusterDoc
dataMat = loadDataSetInPandas('testSet_kmeans.txt')
centroids, cluster = kMeans(dataMat,6)

### Bi-Kmeans
def rankSSE(dataMat1, dataMat2):
    SSE1 = np.sum(dataMat1['minDist'])
    SSE2 = np.sum(dataMat2['minDist'])
    if SSE1 > SSE2:
        print('Cluster 1 has higher SSE')
        return dataMat1, dataMat2
    elif SSE1 <= SSE2:
        print('Cluster 2 has higher SSE')
        return dataMat2, dataMat1
def biKmeans(dataMat, k):
    curK = 0
    m = dataMat.shape[0]
    p = dataMat.shape[1]
    cList = np.zeros((k,p))
    dList = pd.DataFrame(columns = kMeans(dataMat, 2)[1].columns)
    while curK < k:
        centroids, cluster = kMeans(dataMat.iloc[:,0:p], 2)
        dataMat1, dataMat2 = cluster[cluster['minIndex'] == 0], cluster[cluster['minIndex'] == 1]
        deadMat, surviveMat = rankSSE(dataMat1, dataMat2)
        cdf = centroids[surviveMat['minIndex'].values[0]]
        cList[curK] = cdf
        surviveMat.iloc[:,p] = curK
        dList = surviveMat.merge(dList, how = 'outer')
        curK += 1
        dataMat = copy.copy(deadMat)
    return cList, dList
cList, dList = biKmeans(dataMat, 6)

### Kmeans in scikit.learn
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn import datasets

np.random.seed(6)
kmeans = KMeans(n_clusters = 6, init = 'k-means++', n_init = 1).fit(dataMat)
kmeans.cluster_centers_
kmeans.labels_

plt.scatter(kmeans.cluster_centers_[0:6,0], kmeans.cluster_centers_[0:6,1])
plt.show()

### Vector Quantization (LVQ)
import scipy as sp
from sklearn import cluster

from scipy.misc import face
face = face(gray=True)
plt.figure(1, figsize = (3,2.2))
plt.imshow(face, cmap = plt.cm.gray, vmin = 0, vmax = 250)
plt.show()

X = face.reshape((-1,1))
n_clusters = 3
kmeans = KMeans(n_clusters = n_clusters, n_init = 1).fit(X)
values = kmeans.cluster_centers_.squeeze()
labels = kmeans.labels_
face_compressed = np.choose(labels, values)
face_compressed.shape = face.shape

plt.figure(2, figsize = (3,2.2))
plt.imshow(face_compressed, cmap = plt.cm.gray, vmin = 0, vmax = 250)
plt.show()


from PIL import Image
import matplotlib.pylab as plt
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D
from keras.layers import Flatten, Dense
from keras import backend as K
from keras.callbacks import ModelCheckpoint
import keras

def image_size(img_path):

    """

    :param img_path: Path of one image to determine the height and the width
    :return: height and width of the image
    """
    img_width, img_height = img_path.size
    return img_width, img_height

def model_build(img_width, img_height, num_classes):

    """

    :param img_width: Width of the images
    :param img_height: Height of the image
    :param num_classes: Number of classes. In this case 10
    :return: Custom built CNN model
    """

    if K.image_data_format() == 'channels_first':
        input_shape = (1, img_width, img_height)
    else:
        input_shape = (img_width, img_height, 1)

    model = Sequential()
    model.add(Conv2D(16, kernel_size=(5, 5), strides=(1, 1), padding = 'same',
                     activation='relu',
                     input_shape=input_shape))
    model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))
    model.add(Conv2D(32, (5, 5), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Flatten())
    model.add(Dense(128, activation='relu'))
    model.add(Dense(num_classes, activation='softmax'))

    model.compile(loss=keras.losses.categorical_crossentropy,
                  optimizer=keras.optimizers.Adam(),
                  metrics=['accuracy'])

    checkpointer = ModelCheckpoint(filepath='./mnist_weights.hdf5',
                                   verbose=1, save_best_only=True)

    return model, checkpointer

def datagenerator():

    """
    :return: datagen required for data augmentation to prevent overfitting
    """

    # this is the augmentation configuration we will use for training
    train_datagen = ImageDataGenerator(
        rescale=1. / 255,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True)

    # this is the augmentation configuration we will use for testing:
    # only rescaling
    test_datagen = ImageDataGenerator(rescale=1. / 255)

    return train_datagen, test_datagen

def train_validation_split(train_datagen, test_datagen, img_height, img_width, batch_size):

    """

    :param train_datagen: training data generator
    :param test_datagen: testing data generator
    :param img_height: height of the image
    :param img_width: width of the image
    :param batch_size: batch size depending on the number of samples
    :return: training data and the testing data
    """

    train_generator = train_datagen.flow_from_directory(
        train_data_dir,
        target_size=(img_width, img_height),
        color_mode="grayscale",
        batch_size=batch_size,
        class_mode='categorical')

    validation_generator = test_datagen.flow_from_directory(
        validation_data_dir,
        target_size=(img_width, img_height),
        color_mode="grayscale",
        batch_size=batch_size,
        class_mode='categorical')

    return train_generator, validation_generator

def fit(train_generator, validation_generator, checkpointer, model, epochs,
        n_train_samples, n_test_samples):

    """

    :param train_generator: training data
    :param validation_generator: testing data
    :param checkpointer: checkpoint to save the weights of the best model
    :param model: custom built CNN model
    :param epochs: number of epochs
    :param n_train_samples: number of train samples in the dataset
    :param n_test_samples: number of test samples in the dataset
    :return:
    """

    history = model.fit_generator(
        train_generator,
        steps_per_epoch=n_train_samples // batch_size,
        epochs=epochs,
        validation_data=validation_generator,
        validation_steps=n_test_samples // batch_size,
        callbacks=[checkpointer])

    model.save_weights('mnist_model.h5')

    score = model.evaluate_generator(validation_generator, verbose=0, callbacks=[history])

    return history, score

def run(img_path, batch_size, num_classes, epochs, n_train_samples, n_test_samples):

    """

    :param img_path: Path of the image to determine the height and the width
    :param batch_size: batch size depending on the number of samples
    :param num_classes: number of classes in this case 10
    :param epochs: number of epochs
    :param n_train_samples: number of training samples
    :param n_test_samples: number of testing samples
    :return:
    """

    img_width, img_height = image_size(img_path)

    train_datagen, test_datagen = datagenerator()

    train_generator, validation_generator = train_validation_split(train_datagen,
                                                             test_datagen,
                                                             img_height,
                                                             img_width,
                                                             batch_size)

    model, checkpointer = model_build(img_width, img_height, num_classes)

    history, score = fit(train_generator,
                         validation_generator,
                         checkpointer,
                         model,
                         epochs,
                         n_train_samples,
                         n_test_samples)

    return history, score

def plots(History):

    """

    :param History: Metrics of the model.
    :return: Accuracy and loss plots
    """

    plt.plot(History.history['loss'])
    plt.plot(History.history['val_loss'])
    plt.title('Model Loss')
    plt.ylabel('Loss')
    plt.xlabel('Epochs')
    plt.legend(['train', 'test'])
    plt.show()

    plt.plot(History.history['accuracy'])
    plt.plot(History.history['val_accuracy'])
    plt.title('Model Accuracy')
    plt.ylabel('Accuracy')
    plt.xlabel('Epochs')
    plt.legend(['train', 'test'])
    plt.show()

if __name__ == '__main__':

    train_data_dir = '../../data/raw/data/mnist/training'
    validation_data_dir = '../../data/raw/data/mnist/testing'

    n_train_samples = 60000
    n_test_samples = 10000

    batch_size = 512
    num_classes = 10
    epochs = 50

    img_path = Image.open('../../data/raw/data/mnist/training/0/1.png')

    history, score = run(img_path, batch_size, num_classes, epochs, n_train_samples, n_test_samples)

    plots(history)

    print("Training Accuracy: {}".format(history.history['accuracy']))
    print("Testing Accuracy: {}".format(history.history['val_accuracy']))

    """
    The validation accuracy was 98.81% at the 23rd Epoch.
    loss: 0.0479 - accuracy: 0.9847 - val_loss: 0.0229 - val_accuracy: 0.9881

    """


import os
import sys
from io import BytesIO, IOBase

_str = str
str = lambda x=b"": x if type(x) is bytes else _str(x).encode()

BUFSIZE = 8192

class FastIO(IOBase):
    newlines = 0

    def __init__(self, file):
        self._buffer = BytesIO()
        self._fd = file.fileno()
        self._writable = "x" in file.mode or "r" not in file.mode
        self.write = self._buffer.write if self._writable else None

    def read(self):
        if self._buffer.tell():
            return self._buffer.read()
        return os.read(self._fd, os.fstat(self._fd).st_size)

    def readline(self):
        while self.newlines == 0:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            self.newlines = b.count(b"\n") + (not b)
            ptr = self._buffer.tell()
            self._buffer.seek(0, 2), self._buffer.write(b), self._buffer.seek(ptr)
        self.newlines -= 1
        return self._buffer.readline()

    def flush(self):
        if self._writable:
            os.write(self._fd, self._buffer.getvalue())
            self._buffer.truncate(0), self._buffer.seek(0)

def print(*args, sep=b" ", end=b"\n", file=sys.stdout, flush=False):
    at_start = True
    for x in args:
        if not at_start:
            file.write(sep)
        file.write(str(x))
        at_start = False
    file.write(end)
    if flush:
        file.flush()

sys.stdin, sys.stdout = FastIO(sys.stdin), FastIO(sys.stdout)
input = lambda: sys.stdin.readline().rstrip(b"\r\n")