"""Autoregressive transformer.

Modeled after https://github.com/openai/gpt-2/blob/master/src/model.py.

Positional embeddings are not used, since in our context, we don't care about
sequence element order.
"""

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

#### Masking
#
# Assume we have (x0, x1, x2) and order=natural in all cases below.
#
# Scheme 0 (base) --> OK
#         Input  [SOS, x0, x1]
#         Hidden [p(x0), p(x1|x0), p(x2|x0,x1)]
#      (We use 1 block so hidden == output).
#
#      mask(3) is used:
#         mask (row=destination):
#         tensor([[1, 0, 0],
#                 [1, 1, 0],
#                 [1, 1, 1]], dtype=torch.uint8)
# This scheme only works for natural ordering.
#
# Scheme 1 (default) --> OK
#    # TODO: document
#
##### Some failed schemes
#
# Scheme F1 --> FAILS
# [column 0 depends on self; because Query(p(x0)) is produced from x0]
#    Input  [SOS, x0, x1, x2]
#    Hidden [SOS, p(x0), p(x1|x0), p(x2|x0,x1)]
#     mask (row=destination):
#     tensor([[1., 0., 0., 0.],
#             [1., 0., 0., 0.],
#             [1., 1., 0., 0.],
#             [1., 1., 1., 0.]], dtype=torch.float64)
#
# Scheme F2 --> FAILS
# [column 0 depends on self; because Query(p(x0)) is produced from x0]
#
#    Input  [x0, x1, x2, EOS]
#    Hidden [p(x0), p(x1|x0), p(x2|x0,x1), EOS]
#    mask (row=destination):
#    tensor([[1., 0., 0., 1.],
#           [1., 1., 0., 1.],
#           [1., 1., 1., 1.],
#           [0., 0., 0., 1.]], dtype=torch.float64)

DEFAULT_MASK_SCHEME = 0  # 0: original, only works for natural order
DEFAULT_MASK_SCHEME = 1
DEBUG_PRINT = True
DEBUG_PRINT = False

## Notes on unk_embeddings & pos_embeddings
#
#  pos_embeddings index: *position* of input sequence
#    - aganostic to what column index it is (or even if it is SOS)
#
#  unk_embeddings index: *natural_idx* of the column it's supposed to mask
#    - thus, SOS does not have an unk_embeddings
#
#  How they interact: potentially dropout first, then potentially add pos emb.

def pprint(*args):
    do_print = DEBUG_PRINT
    if do_print:
        print(*args)

def mask(n):
    # tensor([[1, 0, 0],
    #         [1, 1, 0],
    #         [1, 1, 1]], dtype=torch.uint8)
    # for ns = nd = 3.
    ns = n
    nd = n
    i = torch.arange(nd)[:, None]
    j = torch.arange(ns)
    m = i >= j - ns + nd
    m.requires_grad = False
    return m

# def order_respecting_mask(ncols, ordering, fill_diagonal=False):
#     """Construct appropriate mask for attention.

#     Assuming o=(2,0,1):
#      - set inputs = [ SOS=0, x0,         x1,         x2 ]
#      - so outputs = [ 0,   h(x0|x2), h(x1|x0,x2), h(x2) ]

#     There's an added column at the left end.

#     Desired mask (row=destination):
#         [[1, 0, 0, 0],
#          [1, 0, 0, 1],
#          [1, 1, 0, 1],
#          [1, 0, 0, 0]]

#     Mask after the first attention == above + "each col sees directly below":
#      - inputs =  [ 0, h(x0|x2), h(x1|x0,x2), h(x2) ]
#      - outputs = [ 0, h(x0|x2), h(x1|x0,x2), h(x2) ]

#         [[1, 0, 0, 0],
#          [1, 1, 0, 1],
#          [1, 1, 1, 1],
#          [1, 0, 0, 1]]
#       (just fill diagnoal with 1's)
#     """
#     # natural idx -> position
#     nat_to_pos = [None] * ncols
#     for natural_idx in range(ncols):
#         nat_to_pos[ordering[natural_idx]] = natural_idx

#     mask = np.zeros((ncols + 1, ncols + 1))
#     mask[:, 0] = 1  # First column is SOS -- everyone can see.
#     for pos_src in range(ncols):
#         src_nat_idx = ordering[pos_src]
#         for pos_dst in range(pos_src + 1, ncols):
#             # Variable at pos_dst should see pos_src.
#             dst_nat_idx = ordering[pos_dst]
#             mask[dst_nat_idx + 1, src_nat_idx + 1] = 1
#     if fill_diagonal:
#         np.fill_diagonal(mask, 1)

#     mask = torch.as_tensor(mask)
#     # mask = torch.as_tensor(mask, dtype=torch.uint8)
#     mask.requires_grad = False
#     return mask

# def order_respecting_mask2(ncols, ordering, fill_diagonal=False):
#     """Construct appropriate mask for attention.

#     Assuming o=(2,0,1):
#      - set inputs = [  x0,         x1,         x2 , EOS]
#      - so outputs = [ h(x0|x2), h(x1|x0,x2), h(x2), EOS]

#     There's an added column at the right end.

#     Desired mask (row=destination):
#         tensor([[0., 0., 1., 1.],
#                 [1., 0., 1., 1.],
#                 [0., 0., 0., 1.],
#                 [0., 0., 0., 1.]], dtype=torch.float64)

#     For natural order (0,1,2), the desired mask is:
#         tensor([[0., 0., 0., 1.],
#                 [1., 0., 0., 1.],
#                 [1., 1., 0., 1.],
#                 [0., 0., 0., 1.]], dtype=torch.float64)

#     Mask after the first attention == above + "each col sees directly below".
#     """
#     # natural idx -> position
#     nat_to_pos = [None] * ncols
#     for natural_idx in range(ncols):
#         nat_to_pos[ordering[natural_idx]] = natural_idx

#     mask = np.zeros((ncols + 1, ncols + 1))
#     mask[:, -1] = 1  # Last column is SOS -- everyone can see.
#     for pos_src in range(ncols):
#         src_nat_idx = ordering[pos_src]
#         for pos_dst in range(pos_src + 1, ncols):
#             # Variable at pos_dst should see pos_src.
#             dst_nat_idx = ordering[pos_dst]
#             mask[dst_nat_idx, src_nat_idx] = 1
#     if fill_diagonal:
#         np.fill_diagonal(mask, 1)

#     mask = torch.as_tensor(mask)
#     # mask = torch.as_tensor(mask, dtype=torch.uint8)
#     mask.requires_grad = False
#     return mask

def order_respecting_mask3(ncols, ordering, input_layer=True):
    """Construct appropriate mask for attention.

    Assuming o=(2,0,1):
     - set inputs = [ SOS=0,          x0,    x1,     x2 ]
     - so outputs = [ h(x0|x2), h(x1|x0,x2), h(x2), EOS ]

    No one connects to EOS.  SOS connects to everyone.

    Desired mask (row=destination):
        [[1, 0, 0, 1],
         [1, 1, 0, 1],
         [1, 0, 0, 0],
         [0, 0, 0, 0]]

    Mask after the first attention + see self (diagonal)
    Basically == shift above to the left 1 column, then fill diagonal
     - inputs  = [ h(x0|x2), h(x1|x0,x2), h(x2), EOS ]
     - outputs = [ h(x0|x2), h(x1|x0,x2), h(x2), EOS ]
        [[1, 0, 1, 0],
         [1, 1, 1, 0],
         [0, 0, 1, 0],
         [0, 0, 0, 0]]
    """
    mask = np.zeros((ncols + 1, ncols + 1))

    if input_layer:
        mask[:, 0] = 1  # First column is SOS -- everyone can see.
        mask[-1, :] = 0  # No one connects to EOS
        for pos_src in range(ncols):
            src_nat_idx = ordering[pos_src]
            for pos_dst in range(pos_src + 1, ncols):
                # Variable at pos_dst should see pos_src.
                dst_nat_idx = ordering[pos_dst]
                mask[dst_nat_idx, src_nat_idx + 1] = 1
    else:
        for pos_src in range(ncols):
            src_nat_idx = ordering[pos_src]
            for pos_dst in range(pos_src, ncols):
                dst_nat_idx = ordering[pos_dst]
                mask[dst_nat_idx, src_nat_idx] = 1

    mask = torch.as_tensor(mask, dtype=torch.float32)
    mask.requires_grad = False
    return mask

class LayerNorm(nn.Module):
    """Norm to 0-mean 1-std , then do a learned diagonal affine transform."""

    def __init__(self, features, eps=1e-5):
        super(LayerNorm, self).__init__()
        self.scale = nn.Parameter(torch.ones(features))
        self.shift = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        s = (x - mean).pow(2).mean(-1, keepdim=True)
        x = (x - mean) * torch.rsqrt(s + self.eps)
        return self.scale * x + self.shift

class Conv1d(nn.Module):
    """Linear (with bias).  Weights ~ N(std), bias ~ 0."""

    def __init__(self, d_in, d_out, w_init_std=0.02):
        super(Conv1d, self).__init__()

        self.w = nn.Parameter(torch.zeros(d_in, d_out))
        self.b = nn.Parameter(torch.zeros(d_out))
        nn.init.normal_(self.w, std=w_init_std)
        nn.init.zeros_(self.b)
        self.d_in = d_in
        self.d_out = d_out

    def forward(self, x):
        *start, d_in = x.size()
        out = torch.matmul(x.view(-1, d_in), self.w) + self.b
        return out.view(start + [self.d_out])

class MultiHeadSelfAttention(nn.Module):
    """Params:

      d_model: last dim of input and output of this module.
      num_heads: number of parallel heads.

    Internally, queries, keys, and values are all produced from the input
    (hence "self"), and all of them are (d_model/num_heads)-dimensional.
    """

    def __init__(self, d_model, num_heads):
        super(MultiHeadSelfAttention, self).__init__()

        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_state = d_model // num_heads

        self.qkv_linear = Conv1d(d_model, self.d_state * 3 * num_heads)
        self.linear = Conv1d(num_heads * self.d_state, d_model)

        self.attn_mask = None  # Will be set by caller.

    def _split_heads(self, x):
        # Each input has shape [bs, num cols, d_state * num_heads].
        *start, m = x.size()
        x = x.view(start + [self.num_heads, m // self.num_heads])
        return x.permute(0, 2, 1, 3)

    def _do_attention(self, query, key, value, mask):
        """Accepts Q,K,V each shaped [bs, num heads, num cols, d_state].

        Returns transformed [bs, num_heads, num cols, d_state].
        """
        d_k = query.size()[-1]
        scores = torch.matmul(query, key.transpose(-1, -2)) / np.sqrt(d_k)
        pprint('qk scores', scores)
        mask = mask.to(scores.dtype)
        scores = scores * mask - (1 - mask) * 1e10
        pprint('masked scores', scores)
        attn_weights = F.softmax(scores, dim=-1)
        pprint('attn weights', attn_weights)
        pprint('values', value)

        out = torch.matmul(attn_weights, value)
        pprint('read', out)
        return out

    def forward(self, x, query_input=None):
        """x: [bs, num cols, d_model].  Output has the same shape."""
        assert x.dim() == 3, x.size()
        bs, ncols, _ = x.size()
        pprint('In MultiHead', x.size())

        # [bs, num cols, d_state * 3 * num_heads]
        qkv = self.qkv_linear(x)
        # [bs, num heads, num cols, d_state] each
        qs, ks, vs = map(self._split_heads, torch.chunk(qkv, 3, dim=-1))

        if query_input is not None:
            # TODO: obviously can avoid redundant calc...
            qkv = self.qkv_linear(query_input)
            qs, _, _ = map(self._split_heads, torch.chunk(qkv, 3, dim=-1))

        # [bs, num heads, num cols, d_state]
        x = self._do_attention(qs, ks, vs, mask=self.attn_mask.to(x.device))

        # [bs, num cols, num heads, d_state]
        x = x.transpose(1, 2)
        # Concat all heads' outputs: [bs, num cols, num heads * d_state]
        x = x.contiguous().view(bs, ncols, -1)
        # Then do a transform: [bs, num cols, d_model].
        x = self.linear(x)
        return x

class GeLU(nn.Module):

    def __init__(self):
        super().__init__()

    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(
            np.sqrt(2 / np.pi) * (x + 0.044715 * torch.pow(x, 3))))

class Block(nn.Module):
    """Params:

      d_model: last dim of input and output of this module.
      d_ff: the hidden dim inside the 2-layer MLP.
      num_heads: number of parallel heads.
    """

    def __init__(self,
                 d_model,
                 d_ff,
                 num_heads,
                 activation='relu',
                 do_residual=False):
        super(Block, self).__init__()

        self.mlp = nn.Sequential(
            Conv1d(d_model, d_ff),
            nn.ReLU() if activation == 'relu' else GeLU(),
            Conv1d(d_ff, d_model),
        )
        self.norm1 = LayerNorm(features=d_model)
        self.norm2 = LayerNorm(features=d_model)
        self.attn = MultiHeadSelfAttention(d_model, num_heads)
        self.do_residual = do_residual

    def set_attn_mask(self, mask):
        self.attn.attn_mask = mask

    def forward(self, x, query_input=None):
        residual = x
        x = self.norm1(x)
        x = self.attn(x, query_input=query_input)
        if self.do_residual:
            x += residual

        residual = x
        x = self.norm2(x)
        x = self.mlp(x)
        if self.do_residual:
            x += residual

        return x

class Transformer(nn.Module):
    """An autoregressive Transformer (decoder only)."""

    def __init__(self,
                 num_blocks,
                 d_model,
                 d_ff,
                 num_heads,
                 nin,
                 input_bins,
                 use_positional_embs=False,
                 activation='relu',
                 dropout=False,
                 fixed_ordering=None,
                 draw_dropout_per_col=False,
                 seed=None,
                 first_query_shared=False,
                 prefix_dropout=False,
                 mask_scheme=DEFAULT_MASK_SCHEME):
        super().__init__()

        print('self.mask_scheme', mask_scheme)
        self.mask_scheme = mask_scheme

        # Common attributes below.
        self.nin = nin
        self.input_bins = input_bins
        encoded_bins = [d_model] * nin
        pprint('input_bins', self.input_bins)
        pprint('encoded', encoded_bins)
        self.logit_indices = np.cumsum(encoded_bins)
        self.nout = self.logit_indices[-1]
        self.num_blocks = num_blocks
        self.d_model = d_model
        self.d_ff = d_ff
        self.num_heads = num_heads
        self.embed_size = d_model
        self.emb_dim = d_model
        self.use_positional_embs = use_positional_embs
        assert activation in ['relu', 'gelu']
        self.activation = activation
        self.fixed_ordering = fixed_ordering
        if fixed_ordering is None:
            natural = np.arange(nin)
            if seed is None:
                self.fixed_ordering = natural
            else:
                self.fixed_ordering = np.random.RandomState(seed).permutation(
                    natural)
        print('ordering', self.fixed_ordering)
        self.draw_dropout_per_col = draw_dropout_per_col
        self.first_query_shared = first_query_shared
        print('first_query_shared', first_query_shared)
        if first_query_shared:
            assert self.mask_scheme == 1, self.mask_scheme
        self.prefix_dropout = prefix_dropout

        # Build.
        self.blocks = nn.Sequential(*[
            Block(d_model,
                  d_ff,
                  num_heads,
                  activation,
                  do_residual=(self.mask_scheme == 0 or i > 0))
            for i in range(num_blocks)
        ])
        # Set masks.

        orig_mask = None
        if self.mask_scheme == 0:
            orig_mask = mask(nin)
        elif self.mask_scheme == 1:
            init_attn_mask = order_respecting_mask3(nin, self.fixed_ordering)
            attn_mask = order_respecting_mask3(nin,
                                               self.fixed_ordering,
                                               input_layer=False)
        else:
            assert False, self.mask_scheme

        if orig_mask is not None:
            print('using orig mask\n', orig_mask)
            for b in self.blocks:
                b.set_attn_mask(orig_mask)
        else:
            print('init_attn_mask\n', init_attn_mask)
            print('after 1st layer attn_mask\n', attn_mask)
            self.blocks[0].set_attn_mask(init_attn_mask)
            for b in self.blocks[1:]:
                b.set_attn_mask(attn_mask)

        self.norm = LayerNorm(d_model)

        # TODO: scale by sqrt(d_model)?
        self.embeddings = nn.ModuleList()
        if len(set(self.input_bins)) == 1:
            print("Detected tied columns")
            tied_embed = nn.Embedding(self.input_bins[0], d_model)
            for i in range(nin):
                self.embeddings.append(tied_embed)
        else:
            for i in range(nin):
                self.embeddings.append(nn.Embedding(self.input_bins[i], d_model))
        for e in self.embeddings:
            nn.init.normal_(e.weight, std=0.02)

        if use_positional_embs:
            if self.mask_scheme == 1:
                self.pos_embeddings = nn.Embedding(self.nin + 1, d_model)
            else:
                self.pos_embeddings = nn.Embedding(self.nin, d_model)
            nn.init.normal_(self.pos_embeddings.weight, std=0.01)

        print('Transformer ctor, dropout {}'.format(dropout))
        self.dropout = dropout
        if dropout or prefix_dropout:
            self.unk_embeddings = nn.ParameterList()
            for i, dist_size in enumerate(self.input_bins):
                self.unk_embeddings.append(nn.Parameter(torch.zeros(d_model)))

        # Interface required by ProgressiveSampling.
        self.input_bins_encoded_cumsum = np.cumsum(encoded_bins)
        self.orderings = [self.fixed_ordering]

    def name(self):
        n = 'transformer'
        n += '-blocks' + str(self.num_blocks)
        n += '-model' + str(self.d_model)
        n += '-ff' + str(self.d_ff)
        n += '-heads' + str(self.num_heads)
        if self.use_positional_embs:
            n += '-posEmb'
        n += '-' + self.activation
        if self.dropout:
            n += '-dropout'
        if self.mask_scheme == 1:
            n += '-scheme1'
        if self.first_query_shared:
            n += '-1stqshared'
        return n

    def EncodeInput(self, x, natural_col=None, out=None, return_pos_embs=False, skip_prefix=[]):
        """Right shift by one token.

        Suppose we want to model x=(x0,x1,x2).
        Set model inputs = [ SOS=0, x0, x1 ]
            (SOS = start of sequence)
        outputs =          [ p(x0); p(x1|x0); p(x2|x0,x1) ].
            (because output i depends on inputs <= i).

        If self.fixed_ordering is supplied and non-natural,
        we set inputs = [ SOS=0, x_o(0), x_o(1) ]
        so    outputs = [ p(x_o(0)), p(x_o(1) | x_o(0)), p(x_o(2) | x_o(0..1)) ]

        This (1) requires when calculating the loss, seq [x_o(0), ..., x_o(2)]
        is passed, (2) assumes we don't change the diagonal attention mask.

        Alternatively (assuming o=(2,0,1)):
          - change diagonal mask to respect ordering o
          - set inputs = [ SOS=0, x_o(0)=x2, x_o(1)=x0 ]
          - so outputs = [ p(x0|x2), p(x1|x0,x2), p(x2) ]
          - doesn't require feeding targets under order o
        """
        if natural_col is not None:
            assert not return_pos_embs
            return self.EncodeInputInference(x, natural_col, out)

        if x.dtype != torch.long:
            x = x.long()
        bs = x.size()[0]

        if self.mask_scheme == 0:
            # SOS = start of sequence symbol, just zeros.
            y_embed = [torch.zeros(bs, self.embed_size, device=x.device)]
            for nat_idx in range(self.nin - 1):
                y_embed.append(self.embeddings[nat_idx](x[:, nat_idx]))
        elif self.mask_scheme == 1:
            y_embed = [torch.zeros(bs, self.embed_size, device=x.device)]
            for nat_idx in range(self.nin):
                y_embed.append(self.embeddings[nat_idx](x[:, nat_idx]))
        else:
            assert False, self.mask_scheme

        # [batch size, num cols (+ 1), d_model].  +1 or not depends on scheme.
        inp = torch.stack(y_embed, 1)

        # pprint('after stacking embs', inp, 'shape', inp.shape)
        # pprint(self.embeddings[0].weight, self.embeddings[1].weight,
        #        self.embeddings[2].weight)

        inp_seq_len = inp.shape[1]

        if self.dropout:
            if self.draw_dropout_per_col:
                vecs = []
                for _ in range(inp_seq_len):
                    vecs.append(
                        torch.dropout(
                            torch.ones(bs, 1, 1, device=x.device),  # ncol=1
                            p=1.0 -
                            np.random.randint(1, self.nin + 1) * 1. / self.nin,
                            train=self.training))
                dropout_vec = torch.cat(vecs, dim=1)
            else:
                dropout_vec = torch.dropout(
                    torch.ones(bs, inp_seq_len, 1, device=x.device),
                    p=1.0 - np.random.randint(1, self.nin + 1) * 1. / self.nin,
                    train=self.training)
            # During training, non-dropped 1's are scaled by 1/(1-p), so we
            # clamp back to 1.  Shaped [bs, num cols, 1].
            batch_mask = torch.clamp(dropout_vec, 0, 1)
            # Shaped [1, num cols, d_model].
            dropped_repr = self.get_dropped_repr()
            inp = batch_mask * inp + (1. - batch_mask) * dropped_repr

        if self.prefix_dropout and self.training:
            prefix_mask = (
                torch.arange(1, inp_seq_len + 1).repeat(bs).reshape([bs, -1])
                > (torch.rand(bs) * inp_seq_len).unsqueeze(1)
            ).int().unsqueeze(2).to(x.device)
            dropped_repr = self.get_dropped_repr()
            inp = (prefix_mask * inp) + (1. - prefix_mask) * dropped_repr
        elif len(skip_prefix) > 0 and any(skip != 0 for skip in skip_prefix):
            dropped_repr = self.get_dropped_repr()
            pos_arr = torch.arange(inp_seq_len).repeat(bs).reshape([bs, -1])
            prefix_mask = torch.stack(
                [pos >= (off + 1) for pos, off in zip(pos_arr, skip_prefix)]
            ).int().unsqueeze(2).to(x.device)
            inp = (prefix_mask * inp) + (1. - prefix_mask) * dropped_repr

        if self.use_positional_embs:
            # [1, inp_seq_len, d_model]
            # NOTE: indexes into pos embs == positions \in [0, inp_seq_len).
            pos_embs = self.pos_embeddings(
                torch.arange(inp_seq_len, device=x.device)).unsqueeze(0)
            inp += pos_embs
            if return_pos_embs:
                return inp, pos_embs
            return inp

        assert not return_pos_embs
        return inp

    def get_dropped_repr(self):
        dropped_repr = torch.stack(tuple(self.unk_embeddings)).unsqueeze(0)
        if self.mask_scheme == 0:
            # Namely, [0, unk(0), unk(1)] for ncols=3.  This means:
            #   (1) SOS is never dropped.
            #   (2) indexing into unk_embeddings is based on natural_idx.
            dropped_repr = torch.cat((torch.zeros_like(
                dropped_repr[:, 0:1, :]), dropped_repr[:, :-1, :]),
                                     dim=1)
        else:
            dropped_repr = torch.cat(
                (torch.zeros_like(dropped_repr[:, 0:1, :]), dropped_repr),
                dim=1)
        return dropped_repr

    def EncodeInputInference(self, x, natural_col, out):
        """Special inference path.

        Args:
          x: [batch size, 1].  Just the data for column 'natural_col'.
          natural_col (int): [0, num cols).
          out: shaped [batch size, d_model].  To hold the encoded data.
        """
        if natural_col < 0:
            # Potentially handling SOS.
            if self.use_positional_embs:
                # Let's also add E_pos=0 to SOS (if enabled).
                out.copy_(
                    self.pos_embeddings(torch.as_tensor(
                        0,
                        device=x.device)).unsqueeze(0).expand(x.size()[0], -1))
            return

        skip = x[0, 0] < 0
        if skip:
            # [bs, d_model]
            embs = self.unk_embeddings[natural_col].unsqueeze(0).expand(
                x.shape[0], -1)
        else:
            # [bs, d_model]
            embs = self.embeddings[natural_col](x).squeeze(1)

        if self.use_positional_embs:
            # NOTE: this is tricky.  Under self.mask_scheme=0 or 1, E_pos=0 is added
            # to SOS, E_pos=1 is added to x0, etc.  So we need to take this into
            # account.
            pos = self.pos_embeddings(
                torch.as_tensor(natural_col + 1, device=x.device)).unsqueeze(0)
            embs = embs + pos

        out.copy_(embs)

    def forward(self, x, skip_prefix=[]):
        """Outputs logits for (x0, x1|x0, x2|x0,x1, ...)."""
        # [bs, ncols] -> [bs, ncols, d_model].  Right-shifted.
        if self.mask_scheme == 1:

            if self.first_query_shared:
                x = self.EncodeInput(x)
                # So that each position gets the same query vec, Q(0).
                pos_embs = torch.zeros_like(x, device=x.device)
            else:
                assert self.use_positional_embs, 'Need pos_embs for 1st layer query vecs'
                x, pos_embs = self.EncodeInput(x, return_pos_embs=True)

            x = self.blocks[0](x, query_input=pos_embs)
            for b in self.blocks[1:]:
                x = b(x)
        else:
            x = self.EncodeInput(x, skip_prefix=skip_prefix)
            pprint('going through {} blocks'.format(len(self.blocks)))
            x = self.blocks(x)

        # pprint('done blocks, out: ', x)
        x = self.norm(x)
        return x

    def forward_with_encoded_input(self, x):
        # [batch size, num cols * d_model] -> [bs, num cols, d_model]
        x = x.view(x.shape[0], -1, self.d_model)

        if self.mask_scheme == 1:
            inp_seq_len = x.shape[1]

            if self.first_query_shared:
                # So that each position gets the same query vec, Q(0).
                pos_embs = torch.zeros_like(x, device=x.device)
            else:
                assert self.use_positional_embs, 'Need pos_embs for 1st layer query vecs'
                pos_embs = self.pos_embeddings(
                    torch.arange(inp_seq_len, device=x.device)).unsqueeze(0)

            x = self.blocks[0](x, query_input=pos_embs)
            for b in self.blocks[1:]:
                x = b(x)
        else:
            x = self.blocks(x)

        x = self.norm(x)
        return x

    def nll(self, logits, data):
        """Calculates -log p(data), given logits (the conditionals).

        Args:
          logits: [batch size, ncols+1, d_model].
          data: [batch size, ncols].

        Returns:
          nll: [batch size].
        """
        if data.dtype != torch.long:
            data = data.long()
        nll = torch.zeros(logits.size()[0], device=logits.device)
        for i in range(self.nin):
            logits_i = self.logits_for_col(i, logits)
            ce = F.cross_entropy(logits_i, data[:, i], reduction='none')
            nll += ce
        return nll

    def logits_for_col(self, idx, logits):
        """Returns the logits (vector) corresponding to log p(x_i | x_(<i)).

        Args:
          idx: int, in natural (table) ordering.
          logits: [batch size, ncols+1, d_model] .

        Returns:
          logits_for_col: [batch size, domain size for column idx].
        """
        embed = self.embeddings[idx]
        return torch.matmul(logits[:, idx, :], embed.weight.t())

if __name__ == '__main__':
    # ncols = 4
    # d_model = 3
    # bs = 1
    # rng = np.random.RandomState(134)
    # attn = MultiHeadSelfAttention(d_model=d_model, num_heads=1)
    # attn.attn_mask = mask(ncols)

    # # Test i depends only on <i.
    # for i in range(ncols):
    #     inp = torch.tensor(rng.rand(bs, ncols, d_model).astype(np.float32),
    #                        requires_grad=True)

    #     out = attn(inp)
    #     l = out[:, i, :].view(-1, )
    #     pprint(l)

    #     out[:, i, :].view(-1, )[0].backward()

    #     pprint('inp.grad', inp.grad)
    #     depends = (inp.grad[0].numpy() != 0).astype(
    #         np.uint8)  # is there a gradient on the input for this k
    #     # pprint('inp.grad[0]', inp.grad[0])
    #     pprint('col', i, 'inp.grads', inp.grad, '\ndepends', depends)

    #     # Shape of depends: [ncols, d_model]
    #     assert np.all(
    #         depends[i + 1:, :].reshape(-1, ) == 0), 'i={} depends={}'.format(
    #             i, depends)

    # print('[MultiHeadSelfAttention] Passes autoregressive-ness check!')

    # block = Block(d_model=16, d_ff=64, num_heads=4)
    # block.set_attn_mask(mask(5))
    # out = block(torch.randn(1, 5, 16))

    ######

    num_cols = 3
    vocab = 1
    bs = 1
    num_cols = 11
    vocab = 5
    bs = 3
    orderings = [
        np.arange(num_cols),
        # [2, 0, 1],
        np.arange(num_cols)[::-1],
        np.random.permutation(np.arange(num_cols)),
    ]
    for ordering in orderings:
        print('Testing ordering', ordering)
        model = Transformer(
            num_blocks=2,  #2,
            # d_model=4,
            d_model=16,
            d_ff=64,
            num_heads=4,
            nin=num_cols,
            input_bins=[
                vocab,
            ] * num_cols,
            use_positional_embs=True,
            activation='gelu',
            fixed_ordering=ordering)
        print('attn_mask for blk 0', model.blocks[0].attn.attn_mask)

        for i in range(num_cols):
            nat_idx = ordering[i]
            print('\nchecking output column {} nat_idx {}...'.format(
                i, nat_idx))
            inp = torch.randint(vocab, (bs, num_cols))
            # [bs, num cols, d_model], the logits
            out = model(inp)
            pprint('inp:', inp)

            out[:, nat_idx, :].contiguous().view(-1,)[0].backward()
            pprint('grads:')
            ok = True
            for n, p in model.named_parameters():
                if 'embed' in n:
                    if p.grad is None:
                        print(n, p.grad)
                        continue
                    dep = (p.grad.reshape(-1) != 0).numpy().any()
                    pprint(
                        n,
                        'gradients?:',
                        dep,
                        'grads',  #p.grad,
                    )
                    # for j in range(i + 1, num_cols):
                    for j in range(i + 1, len(ordering)):
                        nat_idx_j = ordering[j]
                        # i.e., p corresponds to nat_idx j
                        if n == 'embeddings.{}.weight'.format(nat_idx_j):
                            ok &= (not dep)
            assert ok

        print('[Transformer] Passes autoregressive-ness check!')

        # inp = torch.randint(vocab, (bs, num_cols))
        # logits = model(inp)
        # print(model.nll(logits, inp))

        # from torchviz import make_dot, make_dot_from_trace
        # print(
        #     make_dot(logits, params=dict(model.named_parameters())))


import argparse
import csv
import sys
from datetime import datetime
from io import StringIO
from dgh import CsvDGH

_DEBUG = True

class _Table:

    def __init__(self, pt_path: str, dgh_paths: dict):

        """
        Instantiates a table and the specified Domain Generalization Hierarchies from the
        corresponding files.

        :param pt_path:             Path to the table to anonymize.
        :param dgh_paths:           Dictionary whose values are paths to DGH files and whose keys
                                    are the corresponding attribute names.
        :raises IOError:            If a file cannot be read.
        :raises FileNotFoundError:  If a file cannot be found.
        """

        self.table = None
        """
        Reference to the table file.
        """
        self.attributes = dict()
        """
        Dictionary whose keys are the table attributes names and whose values are the corresponding
        column indices.
        """
        self._init_table(pt_path)
        """
        Reference to the table file.
        """
        self.dghs = dict()
        """
        Dictionary whose values are DGH instances and whose keys are the corresponding attribute
        names.
        """
        for attribute in dgh_paths:
            self._add_dgh(dgh_paths[attribute], attribute)

    def __del__(self):

        """
        Closes the table file.
        """

        self.table.close()

    def anonymize(self, qi_names: list, sd_names: list, k: int, output_path: str, v=True):

        """
        Writes a k-anonymous representation of this table on a new file. The maximum number of
        suppressed rows is k.

        :param qi_names:    List of names of the Quasi Identifiers attributes to consider during
                            k-anonymization.
        :param sd_names:    List of names of the Sensitive Data attributes.
        :param k:           Level of anonymity.
        :param output_path: Path to the output file.
        :param v:           If True prints some logging.
        :raises KeyError:   If a QI attribute name is not valid.
        :raises IOError:    If the output file cannot be written.
        """
        """"
        Algo Pseudocode:
            1. freq ← a frequency list contains distinct sequences of values of PT[QI],along with the number of occurrences of each sequence.
            2. while there exists sequences in freq occurring less than k times that account for more than k tuples do
                2.1. let A j be attribute in freq having the most number of distinct values
                2.2. freq ← generalize the values of A j in freq
            3. freq ← suppress sequences in freq occurring less than k times.
            4. freq ← enforce k requirement on suppressed tuples in freq.
            5. Return MGT ← construct table from freq
        """
        global _DEBUG

        if v:
            _DEBUG = False

        self._debug("[DEBUG] Creating the output file...", _DEBUG)
        try:
            output = open(output_path, 'w')
        except IOError:
            raise
        self._log("[LOG] Created output file.", endl=True, enabled=v)

        # Start reading the table file from the top:
        self.table.seek(0)

        self._debug("[DEBUG] Instantiating the QI frequency dictionary...", _DEBUG)
        # Dictionary whose keys are sequences of values for the Quasi Identifiers and whose values
        # are couples (n, s) where n is the number of occurrences of a sequence and s is a set
        # containing the indices of the rows in the original table file with those QI values:
        qi_frequency = dict()

        self._debug("[DEBUG] Instantiating the attributes domains dictionary...", _DEBUG)
        # Dictionary whose keys are the indices in the QI attribute names list, and whose values are
        # sets containing the corresponding domain elements:
        domains = dict()
        for i, attribute in enumerate(qi_names):
            domains[i] = set()

        # Dictionary whose keys are the indices in the QI attribute names list, and whose values are
        # the current levels of generalization, from 0 (not generalized):
        gen_levels = dict()
        for i, attribute in enumerate(qi_names):
            gen_levels[i] = 0
        # 1. freq ← a frequency list contains distinct sequences of values of PT[QI],
        # along with the number of occurrences of each sequence.
        for i, row in enumerate(self.table):

            qi_sequence = self._get_values(row, qi_names, i)

            # Skip if this row must be ignored:
            if qi_sequence is None:
                self._debug("[DEBUG] Ignoring row %d with values '%s'..." % (i, row.strip()),
                            _DEBUG)
                continue
            else:
                qi_sequence = tuple(qi_sequence)

            if qi_sequence in qi_frequency:
                occurrences = qi_frequency[qi_sequence][0] + 1
                rows_set = qi_frequency[qi_sequence][1].union([i])
                qi_frequency[qi_sequence] = (occurrences, rows_set)
            else:
                # Initialize number of occurrences and set of row indices:
                qi_frequency[qi_sequence] = (1, set())
                qi_frequency[qi_sequence][1].add(i)

                # Update domain set for each attribute in this sequence:
                for j, value in enumerate(qi_sequence):
                    domains[j].add(value)

            self._log("[LOG] Read line %d from the input file." % i, endl=False, enabled=v)

        self._log('[LOG] Starting to anonymize table.', endl=True, enabled=v)

        # 2. while there exists sequences in freq occurring less than k times that account for more than k tuples do
        while True:
            to_suppress = self.__check_k_req(qi_frequency, k, v)
            if to_suppress is not None:
                break
            aj = self.__get_aj(domains)
            qi_frequency = self.__generalize_qi_frequency_key_field(qi_frequency, aj, qi_names[aj], gen_levels[aj])
            self.__generalize_domain_set(qi_frequency, domains, aj)
            gen_levels[aj] += 1

        # 3.freq ← suppress sequences in freq occurring less than k times
        for key in to_suppress:
            self._log("[DEBUG] Supressing row with values '%s'." % key, endl=True, enabled=v)
            qi_frequency[key] = (0, set()) #imposto empty set

        # todo 4. ??????

        # 5. Return MGT ← construct table from freq
        self._log('[LOG] Initializing sensitive data list.', endl=True, enabled=v)
        self.table.seek(0)
        sd_dict = dict()
        for i, row in enumerate(self.table):
            sd_sequence = self._get_values(row, sd_names, i)
            if sd_sequence is None:
                self._debug("[DEBUG] Ignoring row %d with values '%s'..." % (i, row.strip()),_DEBUG)
                continue
            else:
                sd_dict[i] = tuple(sd_sequence)

        self._log('[LOG] Start printing the anonymized table.', endl=True, enabled=v)
        row_index = 1
        for key, val in qi_frequency.items():
            for line in val[1]:
                output.write("%d," % row_index)
                # print qi_id_data
                for item in key:
                    output.write("%s," % str(item))
                j = 0
                while j < len(sd_dict[line]) - 1:
                    output.write("%s," % sd_dict[line][j])
                    j += 1
                output.write("%s\n" % sd_dict[line][j])
                row_index += 1

        output.close()
        self._log("[LOG] All done.", endl=True, enabled=v)
        return

    @staticmethod
    def _log(content, enabled=True, endl=True):

        """
        Prints a log message.

        :param content: Content of the message.
        :param enabled: If False the message is not printed.
        """

        if enabled:
            if endl:
                print(content)
            else:
                sys.stdout.write('\r' + content)

    @staticmethod
    def _debug(content, enabled=False):

        """
        Prints a debug message.

        :param content: Content of the message.
        :param enabled: If False the message is not printed.
        """

        if enabled:
            print(content)

    def _init_table(self, pt_path: str):

        """
        Gets a reference to the table file and instantiates the attribute dictionary.

        :param pt_path:             Path to the table file.
        :raises IOError:            If the file cannot be read.
        :raises FileNotFoundError:  If the file cannot be found.
        """

        try:
            self.table = open(pt_path, 'r')
        except FileNotFoundError:
            raise

    def _get_values(self, row: str, attributes: list, row_index=None):

        """
        Gets the values corresponding to the given attributes from a row.

        :param row:         Line of the table file.
        :param attributes:  Names of the attributes to get the data of.
        :param row_index:   Index of the row in the table file.
        :return:            List of corresponding values if valid, None if this row must be ignored.
        :raises KeyError:   If an attribute name is not valid.
        :raises IOError:    If the row cannot be read.
        """

        # Ignore empty lines:
        if row.strip() == '':
            return None

    def _set_values(self, row, values, attributes: list) -> str:

        """
        Sets the values of a row for the given attributes and returns the row as a formatted string.

        :param row:         List of values of the row.
        :param values:      Values to set.
        :param attributes:  Names of the attributes to set.
        :return:            The new row as a formatted string.
        """

        pass

    def _add_dgh(self, dgh_path: str, attribute: str):

        """
        Adds a Domain Generalization Hierarchy to this table DGH collection, from its file.

        :param dgh_path:            Path to the DGH file.
        :param attribute:           Name of the attribute with this DGH.
        :raises IOError:            If the file cannot be read.
        :raises FileNotFoundError:  If the file cannot be found.
        """

        pass

    def __check_k_req(self, d: dict, k: int, v):
        sum_occ = 0
        l = []
        min_occ = sys.maxsize
        max_occ = 0
        for key, val in d.items():
            if val[0] < k:
                sum_occ += val[0]
                l.append(key)

        if sum_occ <= k:
            #calcolo il valore k ottenuto
            for key, val in d.items():
                if val[0] < min_occ:
                    min_occ = val[0]
                else:
                    if val[0] > max_occ:
                        max_occ = val[0]
            self._log("[LOG] Obtained table with q.i. max occurrences:%d and min occurrences(actual k):%d ." % (max_occ, min_occ), endl=True, enabled=v)
            return l
        else:
            return None

    @staticmethod
    def __get_aj(dom: dict):
        aj = 0
        s_max_size = len(dom[aj])
        for key, val in dom.items():
            if len(val) > s_max_size:
                s_max_size = len(val)
                aj = key
        return aj

    def __generalize_qi_frequency_key_field(self, d: dict, aj: int, qi_name: str, current_gen_level: int):
        dgh = self.dghs[qi_name]
        qi_freq_new = dict()
        self._log('[LOG] Generalizing field %s on level %d.' % (qi_name, current_gen_level), endl=True)
        for key, val in d.items():
            new_key_l = list(key)
            new_key_l[aj] = dgh.generalize(key[aj], current_gen_level)
            new_key_t = tuple(new_key_l)
            if new_key_t in qi_freq_new:
                t = list(qi_freq_new[new_key_t])
                t[0] += val[0]
                t[1].update(val[1])
                qi_freq_new[new_key_t] = tuple(t)
            else:
                qi_freq_new[new_key_t] = val

        return qi_freq_new

    @staticmethod
    def __generalize_domain_set(qi_frequency: dict, domains: dict, aj: int):
        domains[aj] = set()
        for key in qi_frequency.keys():
            domains[aj].update(key[aj])
        return

class CsvTable(_Table):

    def __init__(self, pt_path: str, dgh_paths: dict):

        super().__init__(pt_path, dgh_paths)

    def __del__(self):

        super().__del__()

    def anonymize(self, qi_names, sd_names, k, output_path, v=False):

        super().anonymize(qi_names, sd_names, k, output_path, v)

    def _init_table(self, pt_path):

        super()._init_table(pt_path)

        try:
            # Try to read the first line (which contains the attribute names):
            csv_reader = csv.reader(StringIO(next(self.table)))
        except IOError:
            raise

        # Initialize the dictionary of table attributes:
        for i, attribute in enumerate(next(csv_reader)):
            self.attributes[attribute] = i

    def _get_values(self, row: str, attributes: list, row_index=None):

        super()._get_values(row, attributes, row_index)

        # Ignore the first line (which contains the attribute names):
        if row_index is not None and row_index == 0:
            return None

        # Try to parse the row:
        try:
            csv_reader = csv.reader(StringIO(row))
        except IOError:
            raise
        parsed_row = next(csv_reader)

        values = list()
        for attribute in attributes:
            if attribute in self.attributes:
                values.append(parsed_row[self.attributes[attribute]])
            else:
                raise KeyError(attribute)

        return values

    def _set_values(self, row: list, values, attributes: list):

        for i, attribute in enumerate(attributes):
            row[self.attributes[attribute]] = values[i]

        values = StringIO()
        csv_writer = csv.writer(values)
        csv_writer.writerow(row)

        return values.getvalue()

    def _add_dgh(self, dgh_path, attribute):

        try:
            self.dghs[attribute] = CsvDGH(dgh_path)
        except FileNotFoundError:
            raise
        except IOError:
            raise

if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description="Python implementation of the Datafly algorithm. Finds a k-anonymous "
                    "representation of a table.")
    parser.add_argument("--private_table", "-pt", required=True,
                        type=str, help="Path to the CSV table to K-anonymize.")
    parser.add_argument("--quasi_identifier", "-qi", required=True,
                        type=str, help="Names of the attributes which are Quasi Identifiers.",
                        nargs='+')
    parser.add_argument("--sensitive_data_identifier", "-sdi", required=True,
                        type=str, help="Names of the attributes which are Sensitive Data.",
                        nargs='+')
    parser.add_argument("--domain_gen_hierarchies", "-dgh", required=True,
                        type=str, help="Paths to the generalization files (must have same order as "
                                       "the QI name list.",
                        nargs='+')
    parser.add_argument("-k", required=True,
                        type=int, help="Value of K.")
    parser.add_argument("--output", "-o", required=True,
                        type=str, help="Path to the output file.")
    args = parser.parse_args()

    try:

        start = datetime.now()

        dgh_paths = dict()
        for i, qi_name in enumerate(args.quasi_identifier):
            dgh_paths[qi_name] = args.domain_gen_hierarchies[i]
        table = CsvTable(args.private_table, dgh_paths)
        try:
            table.anonymize(args.quasi_identifier, args.sensitive_data_identifier, args.k, args.output, v=True)
        except KeyError as error:
            if len(error.args) > 0:
                _Table._log("[ERROR] Quasi Identifier '%s' is not valid." % error.args[0],
                            endl=True, enabled=True)
            else:
                _Table._log("[ERROR] A Quasi Identifier is not valid.", endl=True, enabled=True)

        end = (datetime.now() - start).total_seconds()
        _Table._log("[LOG] Done in %.2f seconds (%.3f minutes (%.2f hours))" %
                    (end, end / 60, end / 60 / 60), endl=True, enabled=True)

    except FileNotFoundError as error:
        _Table._log("[ERROR] File '%s' has not been found." % error.filename,
                    endl=True, enabled=True)
    except IOError as error:
        _Table._log("[ERROR] There has been an error with reading file '%s'." % error.filename,
                    endl=True, enabled=True)


import sys
import json
import argparse
from os import path

sys.path.append(path.dirname(path.dirname(path.abspath(__file__))))

from filters.topk import TopkFilter

def main(config, input_workers):

    with open(config) as f:
        config_data = json.load(f)

    replicator = TopkFilter(input_workers, config_data)

    replicator.run()

if __name__ == '__main__':

    parser = argparse.ArgumentParser(
                    description='Top K Filter NBA',
                    formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument(
            '--config',
            help='The Net Topology configuration file'
    )
    parser.add_argument(
            '--iworkers',
            type=int,
            default=2,
            help='The number of input workers'
    )

    args = parser.parse_args()

    main(args.config, args.iworkers)


import pytest

@pytest.mark.django_db
def test_welcome_django(client):
    """
    Asserts whether the base URL 404s

    404 means that the HTTP request could not reach a page. We don't have a
    page yet, so 404 is the expected return.
    """
    response = client.get("http://localhost:8000/")
    assert response.status_code == 404


import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from torch.autograd import Variable
import numpy as np
from skipthoughts import BayesianUniSkip
from torchvision import models

epsilon = 1e-16

class Attention(nn.Module):
	def __init__(self,):
		super(Attention,self).__init__()
		self.pool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
		self.conv1 = nn.Conv2d(512,64,kernel_size=1,stride=1,padding=0,bias=True)
		self.conv2 = nn.Conv2d(64,128,kernel_size=3,stride=1,padding=1,bias=True)
		self.pool2 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
		self.conv3 = nn.Conv2d(128,64,kernel_size=1,stride=1,padding=0,bias=True)
		self.conv4 = nn.Conv2d(64,128,kernel_size=3,stride=1,padding=1,bias=True)
		self.conv5 = nn.Conv2d(128,1,kernel_size=1,stride=1,padding=0,bias=True)

	def forward(self,x):
		x = self.pool1(x)
		x = F.relu(self.conv1(x))
		x = F.relu(self.conv2(x))
		x = self.pool2(x)
		x = F.relu(self.conv3(x))
		x = F.relu(self.conv4(x))
		x = F.relu(self.conv5(x))
		x = F.interpolate(x,scale_factor=4)
		return torch.sigmoid(x)

class dcn_vgg(nn.Module):
	def __init__(self,):
		super(dcn_vgg,self).__init__()
		self.backbone = nn.Sequential(*list(models.vgg16(pretrained=True).features.children())[:-1])
		self.spatial_attention = Attention()

	def forward(self,x):
		x = self.backbone(x)
		att = self.spatial_attention(x)
		x = x + torch.mul(x,att.expand_as(x))
		return x

class Conv_LSTM_multi(nn.Module):
	def __init__(self,embed_size=256):
		super(Conv_LSTM_multi,self).__init__()
		self.input_x = nn.Conv2d(512,embed_size,kernel_size=3,stride=1,padding=1,bias=True)
		self.forget_x = nn.Conv2d(512,embed_size,kernel_size=3,stride=1,padding=1,bias=True)
		self.output_x = nn.Conv2d(512,embed_size,kernel_size=3,stride=1,padding=1,bias=True)
		self.memory_x = nn.Conv2d(512,embed_size,kernel_size=3,stride=1,padding=1,bias=True)
		self.input_h = nn.Conv2d(embed_size,embed_size,kernel_size=3,stride=1,padding=1,bias=True)
		self.forget_h = nn.Conv2d(embed_size,embed_size,kernel_size=3,stride=1,padding=1,bias=True)
		self.output_h = nn.Conv2d(embed_size,embed_size,kernel_size=3,stride=1,padding=1,bias=True)
		self.memory_h = nn.Conv2d(embed_size,embed_size,kernel_size=3,stride=1,padding=1,bias=True)
		self.input_c = nn.Conv2d(embed_size,embed_size,kernel_size=3,stride=1,padding=1,bias=True)
		self.forget_c = nn.Conv2d(embed_size,embed_size,kernel_size=3,stride=1,padding=1,bias=True)
		self.output_c = nn.Conv2d(embed_size,embed_size,kernel_size=3,stride=1,padding=1,bias=True)

		self.input_pos = nn.Linear(embed_size,embed_size)
		self.forget_pos = nn.Linear(embed_size,embed_size)
		self.output_pos = nn.Linear(embed_size,embed_size)
		self.input_neg = nn.Linear(embed_size,embed_size)
		self.forget_neg = nn.Linear(embed_size,embed_size)
		self.output_neg = nn.Linear(embed_size,embed_size)

	def forward(self,x,state,semantic_pos,semantic_neg):
		h, c = state
		i = torch.sigmoid(self.input_x(x) + self.input_h(h) + self.input_c(c) + self.input_pos(semantic_pos).unsqueeze(-1).unsqueeze(-1).expand_as(h)+self.input_neg(semantic_neg).unsqueeze(-1).unsqueeze(-1).expand_as(h))
		f = torch.sigmoid(self.forget_x(x) + self.forget_h(h) + self.forget_c(c) + self.forget_pos(semantic_pos).unsqueeze(-1).unsqueeze(-1).expand_as(h)+ self.forget_neg(semantic_neg).unsqueeze(-1).unsqueeze(-1).expand_as(h))
		g = torch.tanh(self.memory_x(x) + self.memory_h(h))

		next_c = torch.mul(f,c) + torch.mul(i,g)
		o = torch.sigmoid(self.output_x(x) + self.output_h(h) + self.output_c(next_c) + self.output_pos(semantic_pos).unsqueeze(-1).unsqueeze(-1).expand_as(h) + self.output_neg(semantic_neg).unsqueeze(-1).unsqueeze(-1).expand_as(h))
		h = torch.mul(o,next_c)
		state = (h,next_c)

		return state

class semantic_att(nn.Module):
	def __init__(self,embedding_size=512):
		super(semantic_att,self).__init__()
		self.semantic_q = nn.Linear(embedding_size,embedding_size)
		self.semantic_pool = nn.Linear(embedding_size,embedding_size)
		self.semantic_cur = nn.Linear(embedding_size,embedding_size)
		self.semantic_att = nn.Linear(embedding_size,1)

	def forward(self,semantic,q,cur_v):
		q = F.relu(self.semantic_q(q)).unsqueeze(1).expand_as(semantic)
		v_pool = F.relu(self.semantic_pool(semantic))
		v_cur = F.relu(self.semantic_cur(cur_v)).unsqueeze(1).expand_as(semantic)
		semantic_att = F.softmax(self.semantic_att(q+v_pool+v_cur),dim=1).expand_as(semantic)
		semantic = torch.mul(semantic_att,semantic).sum(1)

		return semantic

class SWM_agg(nn.Module):
	def __init__(self,embed_size=256,vocab=None):
		super(SWM_agg,self).__init__()
		self.backbone = dcn_vgg()
		self.embed_size = embed_size
		self.rnn = Conv_LSTM_multi(embed_size)
		self.q_model = BayesianUniSkip(dir_st='/srv/chenshi/data/Skip-thoughts-pretrained' ,vocab=vocab) # using pretrained skip-thought
		self.q_embed = nn.Linear(2400,embed_size) # skip-thought vectors have pre-defined size 2400
		self.sal_pred = nn.Conv2d(embed_size,2,kernel_size=1,stride=1,padding=0,bias=True)
		self.semantic_att = semantic_att(embed_size)
		self.semantic_embed = nn.Linear(512,embed_size)

		# weight for combining the two maps
		self.weight_q = nn.Linear(embed_size,embed_size)
		self.weight_v = nn.Linear(embed_size*2,embed_size)
		self.weight = nn.Linear(embed_size,2)

	def init_hidden_state(self,batch,height,width,embedding_size=256):
		init_h = torch.zeros(batch,embedding_size,height,width).cuda()
		init_m = torch.zeros(batch,embedding_size,height,width).cuda()
		return (init_h,init_m)

	def get_semantic(self,sal_map,semantic_feature):
		semantic_feature_pos = sal_map[:,0].unsqueeze(1).expand_as(semantic_feature) * semantic_feature
		semantic_feature_neg = sal_map[:,1].unsqueeze(1).expand_as(semantic_feature) * semantic_feature
		semantic_feature_pos = semantic_feature_pos.view(semantic_feature.size(0),semantic_feature.size(1),-1).mean(-1)  # for sigmoid output
		semantic_feature_neg = semantic_feature_neg.view(semantic_feature.size(0),semantic_feature.size(1),-1).mean(-1)  # for sigmoid output

		return semantic_feature_pos, semantic_feature_neg

	def forward(self,x,que):
		if len(que) == 1:
			q = self.q_model(que,[len(que[0])])
		else:
			q = self.q_model(que)
		q = F.relu(self.q_embed(q)) # for semantic attention with hadamard product

		batch, seq, c, h, w = x.size()
		state = self.init_hidden_state(batch,int(h/16),int(w/16),self.embed_size)
		semantic_mem_pos = torch.zeros(batch,self.embed_size).cuda()
		semantic_mem_neg = torch.zeros(batch,self.embed_size).cuda()

		output = []
		semantic_pool_pos = []
		semantic_pool_neg = []
		for i in range(seq):
			cur_v = x[:,i].contiguous()
			cur_v = self.backbone(cur_v)

			state = self.rnn(cur_v,state,semantic_mem_pos,semantic_mem_neg)

			cur_sal = state[0] + q.unsqueeze(-1).unsqueeze(-1).expand_as(state[0])
			cur_sal = self.sal_pred(cur_sal)
			# semantic memory module
			rand_prob = np.random.rand(1,)[0]
			sal_map = torch.sigmoid(cur_sal)
			semantic_feature_pos, semantic_feature_neg = self.get_semantic(sal_map,cur_v)
			semantic_feature_pos = F.relu(self.semantic_embed(semantic_feature_pos))
			semantic_feature_neg = F.relu(self.semantic_embed(semantic_feature_neg))

			semantic_pool_pos.append(semantic_feature_pos)
			semantic_pool_neg.append(semantic_feature_neg)
			semantic_mem_pos = self.semantic_att(torch.cat([ _.unsqueeze(1) for _ in semantic_pool_pos], 1),q,semantic_feature_pos)
			semantic_mem_neg = self.semantic_att(torch.cat([ _.unsqueeze(1) for _ in semantic_pool_neg], 1),q,semantic_feature_neg)

			cur_sal = torch.sigmoid(cur_sal)
			# obtaining the aggregated saliency map
			weight_q = F.relu(self.weight_q(q))
			weight_v = F.relu(self.weight_v(torch.cat((semantic_mem_pos,semantic_mem_neg),dim=-1)))
			cur_weight = torch.sigmoid(self.weight(weight_q+weight_v))
			cur_sal = torch.mul(cur_sal,cur_weight.unsqueeze(-1).unsqueeze(-1).expand_as(cur_sal))
			cur_sal = cur_sal[:,0] + cur_sal[:,1]
			batch_, h_, w_ = cur_sal.size()
			cur_sal = cur_sal.view(batch_,h_*w_)
			cur_sal = torch.div(cur_sal,cur_sal.max(-1,keepdim=True)[0].expand_as(cur_sal)+epsilon)
			cur_sal = cur_sal.view(batch_,h_,w_)
			output.append(cur_sal.unsqueeze(1))

		output = torch.cat([_ for _ in output],dim=1)

		return output


import urllib.request
import re
from bs4 import BeautifulSoup
import pickle
import sys; sys.setrecursionlimit(90000)

class section:
    def __init__(self):
        self.class_list = None

class individual_class:
    def __init__(self):
        self.name = None
        self.course_number = None
        self.credit_hours = None
        self.price = None
        self.prerequisites = None
        self.professors = None
        self.dates = None
        self.class_number = None

def get_data_from_site(site):
    response = urllib.request.urlopen(site)
    html = response.read()
    text = html.decode('cp1252')
    return text

def store_data_locally(data, dest):
    with open(dest, 'w') as f:
        f.write(data)

def table_check(tab, line):
    if tab:
        if re.search('</tr>', line):
            return False
        else:
            return True
    else:
        if re.search('<tr>', line):
            return True
        else:
            return False

def get_links(line):
    return re.findall('(?:.*?href.+?\#)(.+?)(?=\")', line)

def parse_ufv_links(data):
    sub_names = []
    in_table = False
    for line in data.split():
        in_table = table_check(in_table, line)
        if in_table:
            sub_names.append(get_links(line))
    return [x[0] for x in sub_names if x] # only returns lists that have data

def split_sections(html_soup, links):
    """returns the raw html of each block of subject data in the soup"""
    links = links[:]
    subject_data = []
    buffer = []
    # html = html_soup.prettify().split('\n')
    with open('ufv_class_data.html', 'r') as f: # TODO REMOVE
        html = f.readlines()
    for line in html:
        if 'name' in line:
            for link in links:
                if re.search(r'a\s*name\s*=\s*\"{}'.format(link), line):
                    if buffer:
                        subject_data.append(buffer)
                        buffer = []
                    buffer.append(line)
                    links.remove(link)
            else:
                if buffer:
                    buffer.append(line)
        elif buffer:
            buffer.append(line)
    else:
        subject_data.append(buffer)
    return subject_data

def split_sections_naive(html_soup):
    """returns the raw html of each block of subject data in the soup"""
    subject_data = []
    buffer = []
    # html = html_soup.prettify().split('\n')
    with open('ufv_class_data.html', 'r') as f:  # TODO REMOVE
        html = f.readlines()
    for line in html:
        if 'a name=\"' in line:
            if buffer:
                subject_data.append(buffer)
            buffer = [line]
        elif buffer:
            buffer.append(line)
    subject_data.append(buffer)
    return subject_data

def parse_ufv_data(html_soup):
    # links = parse_ufv_links(html_soup.prettify())
    with open('links.pk', 'rb') as f:  # TODO REMOVE
        links = pickle.load(f)
    subject_data = split_sections_naive(html_soup)
    return links

def remove_struck_out_classes(text_data):
    all_data_buf = []
    single_class_buf = []
    in_class = False
    last_newline = True
    for line in text_data:
        if last_newline:
            if re.search('\s+\d+'):
                pass
        if '\n' in line:
            last_newline = True


#-*-coding: utf-8 -*-
"""
/dms/folder/views_navigation.py

.. enthaelt den View zum Aendern des linken Navigationsbereichs
         Django content Management System

Hans Rauch
hans.rauch@gmx.net

Die Programme des dms-Systems koennen frei genutzt und den spezifischen
Beduerfnissen entsprechend angepasst werden.

0.01  22.01.2007  Beginn der Arbeit
"""

from django.http        import HttpResponse, HttpResponseRedirect
from django.shortcuts   import render_to_response
from django             import newforms as forms

from django.utils.translation import ugettext as _

from dms.queries        import get_site_url
from dms.queries        import save_navigation
from dms.queries        import get_menu_by_name_left
from dms.queries        import get_all_menus_left

from dms.roles          import require_permission
from dms.roles          import UserEditPerms
from dms.utils          import get_breadcrumb
from dms.utils          import get_tabbed_form
from dms.utils          import get_footer_email
from dms.utils          import get_folderish_actions

from dms.utils_navigation import get_navmenu_choices_left
from dms.folder.help_form import help_form

from dms_ext.extension    import * # dms-Funktionen ueberschreiben

# -----------------------------------------------------
@require_permission('perm_manage_site')
def folder_navigation(request, item_container):
  """ Zuordnung zu einem Navigationspunkt aendern """

  if request.GET.has_key('name'):
    name = request.GET.copy()['name']
    nav_menu = get_menu_by_name_left(name)
    item_container.container.menu_left_id = abs(nav_menu.menu_id)
    item_container.container.nav_name_left = '|'
    item_container.container.save()

  class dms_itemForm(forms.Form):
    # --- vorhandenes Menue aendern
    navigation_left = forms.ChoiceField(
                            choices=get_navmenu_choices_left(item_container.container.menu_left_id),
                            widget=forms.RadioSelect({})
                            )
  app_name = u'folder'
  my_title = _(u'Navigationspunkt auswählen')

  data_init = { 'navigation_left': item_container.container.nav_name_left, }
  if request.method == 'POST' :
    data = request.POST.copy()
  else :
    data = data_init
  f = dms_itemForm(data)

  # --- Reihenfolge, Ueberschriften, Hilfetexte
  tabs = [('tab_navigation', ['navigation_left',]),]
  content = get_tabbed_form(tabs, help_form, app_name, f, has_tabs=False)

  if request.method == 'POST' and not f.errors:
    save_navigation(item_container, data_init, f.data)
    return HttpResponseRedirect(get_site_url(item_container, 'index.html'))
  else :
    user_perms = UserEditPerms(request.user.username, request.path)
    nav_menus = get_all_menus_left(mode=1)
    vars={'content_div_style': 'frame-main-manage',
          'site'             : item_container.container.site,
          'my_name'          : item_container.item.name,
          'action'           : get_folderish_actions(request,user_perms, item_container, app_name, False),
          'breadcrumb'       : get_breadcrumb(item_container),
          'errors'           : f.errors,
          'content'          : content,
          'menus_left'       : nav_menus,
          'title'            : my_title,
          'sub_title'        : item_container.item.title,
          'footer_email'     : get_footer_email(item_container.item),
          'last_modified'    : item_container.get_last_modified(),
          'submit'           : my_title,
        }
    return render_to_response ( 'app/folder/menu_left.html', vars )


class Solution5677(object):
    def countHomogenous(self, s):
        """
        :type s: str
        :rtype: int
        """
        array = []
        cur = '+'
        number = 0
        for i in range(len(s)):
            cur2 = s[i]
            if cur2 == cur:
                number += 1
            else:
                array.append(number)
                number = 1
                cur = cur2
        array.append(number)
        count = 0
        x = 1000000007
        for i in range(len(array)):
            mm = array[i]
            count += mm * (mm + 1) / 2
        return count % x

s = Solution5677()
print(s.countHomogenous("abbcccaa"))

competition/First Half 2021/Feb 14. 2021/Solution5677.py


import csv
import json

# date format mm/dd/yyyy
def toJson(path,mostRecentDate):
  with open(path, mode='r') as csv_file:
      csv_reader = csv.DictReader(csv_file)
      arr = []
      for row in csv_reader:
          if row["modeldate"] == mostRecentDate:
            obj = {
              "stateId":row["state"],
              "trumpChance":row['winstate_inc'],
              "bidenChance":row['winstate_chal']
            }
            print(obj)
            arr.append(obj)
  return arr

data = toJson('./election-forecasts-2020/presidential_state_toplines_2020.csv','10/17/2020')

with open('data.json', 'w') as outfile:
    json.dump(data, outfile)


import multiprocessing

from celery.utils.log import get_task_logger
from dynaconf import settings

from opulence import App
from opulence.common.job import Result, StatusCode
from opulence.common.plugins import PluginManager

logger = get_task_logger(__name__)

manager = multiprocessing.Manager()
available_collectors = manager.dict()

@App.task(name="collectors:reload_collectors", ignore_result=True)
def reload_collectors(flush=False):
    global available_collectors
    logger.info("Reload collectors")
    if flush:
        available_collectors.clear()
    for path in settings.COLLECTORS_MODULES:
        PluginManager().discover(path)
        for plugin in PluginManager().get_plugins(path):
            if plugin.plugin_name not in available_collectors:
                available_collectors[plugin.plugin_name] = plugin

@App.task(name="collectors:list_collectors")
def list_collectors():
    global available_collectors
    logger.info("List collectors")
    return [c.get_info() for _, c in available_collectors.items()]

@App.task(name="collectors:execute_collector_by_name")
def execute_collector_by_name(collector_name, fact_or_composite):
    global available_collectors
    logger.info(
        "Execute collector {} with {}".format(collector_name, type(fact_or_composite))
    )
    result = Result(input=fact_or_composite, status=StatusCode.empty)
    if collector_name in available_collectors:
        return available_collectors[collector_name].run(result)
    result.status = (
        StatusCode.error,
        "Could not find collector {}".format(collector_name),
    )
    return result

# Reload collectors at startup
reload_collectors(flush=True)
[
    print(
        "{} ({})".format(
            i["plugin_data"]["name"],
            i["plugin_data"]["error"] if i["plugin_data"]["error"] else "OK",
        )
    )
    for i in list_collectors()
]


import geopandas as gpd
import pandas as pd
import geojsonio
from shapely.geometry import Point, Polygon, shape, LineString
from shapely.ops import cascaded_union
from scipy.spatial import cKDTree,distance_matrix
import numpy as np
import missingno as msn
import seaborn as sns
import matplotlib.pyplot as plt
import pycpt
from matplotlib.patches import Patch
from matplotlib.lines import Line2D
from matplotlib.collections import LineCollection
from sklearn.decomposition import PCA
import os, argparse, fiona, re, mds, commons
import distribution as Dist
from PIL import Image
import core, commons
import imageBasedInputPostProcessing as IBIPP
import elevation
import pycpt
from osgeo import gdal

# dataset = osgeo.gdal.Open("MergedDEM")
# gt = dataset.GetGeoTransform()
def hillshade(array, azimuth, angle_altitude):
    # Source: http://geoexamples.blogspot.com.br/2014/03/shaded-relief-images-using-gdal-python.html
    x, y = np.gradient(array)
    slope = np.pi/2. - np.arctan(np.sqrt(x*x + y*y))
    aspect = np.arctan2(-x, y)
    azimuthrad = azimuth*np.pi / 180.
    altituderad = angle_altitude*np.pi / 180.

    shaded = np.sin(altituderad) * np.sin(slope) \
     + np.cos(altituderad) * np.cos(slope) \
     * np.cos(azimuthrad - aspect)
    return 255*(shaded + 1)/2

def downloadAndSaveElevationData(cityBoundary, cityName):
    maxx,maxy,minx,miny = cityBoundary.bounds['maxx'][0],cityBoundary.bounds['maxy'][0],cityBoundary.bounds['minx'][0],cityBoundary.bounds['miny'][0]
    outputfName = '{}-DEM-geojson.tif'.format(cityName)

    # Query and download elevation data
    os.system('eio clip -o {} --bounds {} {} {} {}'.format('temp.tif', minx, miny, maxx, maxy))
    os.system('scp temp.tif {}'.format(commons.getElevationfName(cityName)))
    os.system('rm temp.tif')
    elevation.clean()

def saveElevationOverlayZoning(cityName, validBlocks, overlayName,cityBoundary):
    try:
        topocmap = pycpt.load.cmap_from_cptcity_url('wkp/schwarzwald/wiki-schwarzwald-cont.cpt')
    except:
        topocmap = 'Spectral_r'
    if not os.path.exists(commons.getElevationfName(cityName)):
        downloadAndSaveElevationData(cityBoundary, cityName)
    gdal_data = gdal.Open(commons.getElevationfName(cityName))
    gdal_band = gdal_data.GetRasterBand(1)
    nodataval = gdal_band.GetNoDataValue()
    # convert to a numpy array
    data_array = gdal_data.ReadAsArray().astype(np.float)

    # replace missing values if necessary
    if np.any(data_array == nodataval):
        data_array[data_array == nodataval] = np.nan

    vmin = 180
    vmax = 575
    maxx,maxy,minx,miny = cityBoundary.bounds['maxx'][0],cityBoundary.bounds['maxy'][0],cityBoundary.bounds['minx'][0],cityBoundary.bounds['miny'][0]
    extent = minx, maxx, miny, maxy
    myDpi = 192
    imH, imW = commons.getOutputImageHeightWidth()
    fig = plt.figure(figsize=(imW, imH))
    ax = fig.add_subplot(111)
    ax.matshow(data_array, extent=extent, cmap='Greys', alpha=0.6, zorder=10) #fignum=1
    cityBase = validBlocks.plot(ax=ax,figsize=(imW,imH), color = commons.getColorHexPalette()['residential'], alpha=1.0, antialiased=False)
    plt.axis('off')
    fig.savefig('{}{}_elevationHillshade-validBlock-Overlay-{}.png'.format(commons.getOutputPath(cityName), cityName, overlayName),bbox_inches='tight', interpolation = 'nearest')

def saveElevationOverlayMultipleZoning(cityName, residentialBlocks, parkBlocks, industrialBlocks, commercialBlocks, cityBoundary ):
    try:
        import pycpt
        topocmap = pycpt.load.cmap_from_cptcity_url('wkp/schwarzwald/wiki-schwarzwald-cont.cpt')
    except:
        topocmap = 'Spectral_r'
    if not os.path.exists(commons.getElevationfName(cityName)):
        downloadAndSaveElevationData(cityBoundary, cityName)
    gdal_data = gdal.Open(commons.getElevationfName(cityName))
    gdal_band = gdal_data.GetRasterBand(1)
    nodataval = gdal_band.GetNoDataValue()
    # convert to a numpy array
    data_array = gdal_data.ReadAsArray().astype(np.float)

    # replace missing values if necessary
    if np.any(data_array == nodataval):
        data_array[data_array == nodataval] = np.nan

    vmin = 180
    vmax = 575
    maxx,maxy,minx,miny = cityBoundary.bounds['maxx'][0],cityBoundary.bounds['maxy'][0],cityBoundary.bounds['minx'][0],cityBoundary.bounds['miny'][0]
    extent = minx, maxx, miny, maxy
    myDpi = 192
    imH, imW = commons.getOutputImageHeightWidth()
    fig = plt.figure(figsize=(imW, imH))
    ax = fig.add_subplot(111)
    ax.matshow(data_array, extent=extent, cmap=topocmap, alpha=1.0, zorder=1) #fignum=1
    blockAlpha = 1.0
    cityRes = residentialBlocks.plot(ax=ax,figsize=(imW,imH), color = commons.getColorHexPalette()['residential'], alpha=blockAlpha, antialiased=False, zorder=2, linewidth=2)
    cityResCom = commercialBlocks.plot(ax=cityRes,figsize=(imW,imH), color = commons.getColorHexPalette()['commercial'], alpha=blockAlpha, antialiased=False,zorder=2,linewidth=2)
    cityResComInd = industrialBlocks.plot(ax=cityResCom,figsize=(imW,imH), color = commons.getColorHexPalette()['industrial'], alpha=blockAlpha, antialiased=False, zorder=2,linewidth=2)
    cityResComIndPark = parkBlocks.plot(ax=cityResComInd,figsize=(imW,imH), color = commons.getColorHexPalette()['park'], alpha=blockAlpha, antialiased=False, zorder=2,linewidth=2)
    plt.axis('off')
    fig.savefig('{}{}_elevationHillshade-validBlock-Overlay-{}.png'.format(commons.getOutputPath(cityName), cityName, 'allZoning'),bbox_inches='tight', interpolation = 'nearest')

    plt.clf()
    fig = plt.figure(figsize=(imW, imH))
    ax = fig.add_subplot(111)
    x, y = np.gradient(data_array)
    elevationGradient = np.sqrt(x*x + y*y)
    ax.matshow(elevationGradient, extent=extent, cmap='Greys', alpha=0.8, zorder=3) #fignum=1
    cityRes = residentialBlocks.plot(ax=ax,figsize=(imW,imH), color = commons.getColorHexPalette()['residential'], alpha=blockAlpha, antialiased=False, zorder=2, linewidth=2)
    cityResCom = commercialBlocks.plot(ax=cityRes,figsize=(imW,imH), color = commons.getColorHexPalette()['commercial'], alpha=blockAlpha, antialiased=False,zorder=2,linewidth=2)
    cityResComInd = industrialBlocks.plot(ax=cityResCom,figsize=(imW,imH), color = commons.getColorHexPalette()['industrial'], alpha=blockAlpha, antialiased=False, zorder=2,linewidth=2)
    cityResComIndPark = parkBlocks.plot(ax=cityResComInd,figsize=(imW,imH), color = commons.getColorHexPalette()['park'], alpha=blockAlpha, antialiased=False, zorder=2,linewidth=2)
    plt.axis('off')
    fig.savefig('{}{}_elevationHillshade-validBlock-Overlay-gradient-{}.png'.format(commons.getOutputPath(cityName), cityName, 'allZoning'),bbox_inches='tight', interpolation = 'nearest')

def getCityElevationData(cityName, cityBoundary, normalizeOn=True):
    if not os.path.exists(commons.getElevationfName(cityName)):
        downloadAndSaveElevationData(cityBoundary, cityName)
    gdal_data = gdal.Open(commons.getElevationfName(cityName))
    gdal_band = gdal_data.GetRasterBand(1)
    nodataval = gdal_band.GetNoDataValue()

    # convert to a numpy array
    data_array = gdal_data.ReadAsArray().astype(np.float)

    # replace missing values if necessary
    if np.any(data_array == nodataval):
        data_array[data_array == nodataval] = np.nan

    data_array = np.array(data_array)
    if normalizeOn:
        # normalize data
        print("normalizing data")
        upper_quartile = np.percentile(data_array, 95)
        lower_quartile = np.percentile(data_array, 5)

        data_array[np.where((data_array < lower_quartile))] = lower_quartile
        data_array[np.where((data_array > upper_quartile))] = upper_quartile
    return data_array

def saveElevationGraph(cityName):
    cityBoundary = gpd.read_file(commons.getBoundaryfName(cityName))
    maxx,maxy,minx,miny = cityBoundary.bounds['maxx'][0],cityBoundary.bounds['maxy'][0],cityBoundary.bounds['minx'][0],cityBoundary.bounds['miny'][0]
    vmin = 0
    vmax = 1000
    extent = minx, maxx, miny, maxy

    # convert to a numpy array
    data_array = np.array(getCityElevationData(cityName, cityBoundary))
    minElev = np.min(np.array(data_array))
    maxElev = np.max(np.array(data_array))

    data_array = (data_array - minElev) * 1000.0 / (maxElev - minElev)
    minElev = np.min(np.array(data_array))
    maxElev = np.max(np.array(data_array))
    print( "current min %.4f" % np.min(np.array(data_array)))
    print( "current max %.4f" % np.max(np.array(data_array)))
    #Plot out data with Matplotlib's 'contour'
    imH, imW = commons.getOutputImageHeightWidth()

    fig, ax = plt.subplots(figsize=(imW,imH))
    ax.axis('off')
    ax.matshow(data_array, extent=extent, cmap='Greys', alpha=0.0, zorder=10)
    plot = ax.contour(data_array, origin='image', antialiased=False, extent=extent, cmap = "terrain", levels = list(range(int(minElev), int(maxElev)+3, 200)))
    # ax.clabel(plot, inline=1, fontsize=10)
    # cbar = plt.colorbar()
    # plt.gca().set_aspect('equal', adjustable='box')

    plt.savefig('{}{}_elevationContour.png'.format(commons.getOutputPath(cityName), cityName), bbox_inches='tight', interpolation = 'nearest')

    #Hillsahde
    # topoName = 'wkp/schwarzwald/wiki-schwarzwald-cont'
    topoName = 'meyers'
    topocmap = pycpt.load.cmap_from_cptcity_url('wkp/encyclopedia/{}.cpt'.format(topoName))
    # topocmap = pycpt.load.cmap_from_cptcity_url('wkp/schwarzwald/nordisk-familjebok.cpt')

    plt.clf()
    fig, ax = plt.subplots(figsize=(imW,imH))
    ax.axis('off')
    ax.matshow(data_array, extent=extent, cmap='Greys', alpha=0.0, zorder=10)
    # plt.matshow(hillshade(data_array, 30, 30), extent=extent, cmap='Greys', alpha=1.0, zorder=10, fignum=1)
    cax = ax.contourf(data_array, 10,extent=extent,
                      cmap='terrain', vmin=vmin, vmax=vmax, origin='image', antialiased=False,legend=False)
    ax.axis('off')
    # fig.colorbar(cax, ax=ax)
    fig.savefig('{}{}_elevationContourAxContourf.png'.format(commons.getOutputPath(cityName), cityName),bbox_inches='tight', interpolation = 'nearest')

    #Clean Version For Input
    plt.clf()
    fig, ax = plt.subplots(figsize=(imW,imH))
    ax.matshow(data_array, extent=extent, cmap='Greys', alpha=1.0, zorder=10) #cmap='Greys' 'topocmap'
    ax.axis('off')
    elevationfName = '{}{}_elevation'.format(commons.getOutputPath(cityName), cityName)
    fig.savefig(commons.appendPng(elevationfName), bbox_inches='tight', interpolation = 'nearest')
    imGrey = core.openImgNameAsRGB(commons.appendPng(elevationfName))
    core.write_ppm(commons.appendPpm(elevationfName), np.asarray(imGrey), 255)
    ax.matshow(data_array, extent=extent, cmap=topocmap, alpha=1.0, zorder=10)
    fig.savefig('{}{}_elevationColored.png'.format(commons.getOutputPath(cityName), cityName),bbox_inches='tight', interpolation = 'nearest')

    ax.matshow(hillshade(data_array, 1, 1), extent=extent, cmap=topocmap, alpha=1.0, zorder=10)
    cax = ax.contourf(data_array, np.arange(vmin, vmax, 10),extent=extent,
                      cmap=topocmap, vmin=vmin, vmax=vmax, origin='image', antialiased=False)
    fig.savefig('{}{}_elevationColored{}.png'.format(commons.getOutputPath(cityName), cityName, topoName),bbox_inches='tight', interpolation = 'nearest')

    #Clean Version For Input
    plt.clf()
    fig, ax = plt.subplots(figsize=(imW,imH))
    ax.matshow(data_array, extent=extent, cmap='Greys', alpha=1.0, zorder=10) #cmap='Greys' 'topocmap'
    ax.axis('off')
    x, y = np.gradient(data_array)
    elevationGradient = np.sqrt(x*x + y*y)
    elevationGradientfName = '{}{}_elevationGradient'.format(commons.getOutputPath(cityName), cityName)
    ax.matshow(elevationGradient, extent=extent, cmap='Greys', alpha=1.0, zorder=10)
    fig.savefig(commons.appendPng(elevationGradientfName), bbox_inches='tight', interpolation = 'nearest')

    imGradientGrey = core.openImgNameAsRGB(commons.appendPng(elevationGradientfName))
    core.write_ppm(commons.appendPpm(elevationfName), np.asarray(imGradientGrey), 255)
    ax.matshow(elevationGradient, extent=extent, cmap=topocmap, alpha=1.0, zorder=10)
    fig.savefig('{}{}_elevationGradientColored.png'.format(commons.getOutputPath(cityName), cityName),bbox_inches='tight', interpolation = 'nearest')

    mergeElevationImgName = '{}{}_ElevationMerge'.format(commons.getOutputPath(cityName), cityName)
    mergeElevationImg = core.compositeRGBAndSave(imGrey, imGradientGrey, imGradientGrey, mergeElevationImgName)

def main2():
    saveElevationGraph("Pittsburgh")

def main():
    zoning, nonormZoning, cityName, outputPath, cityBoundary, water, waterNonNorm = core.getInputCityAndZoning()
    cityBoundary = gpd.read_file(commons.getBoundaryfName(cityName))
    maxx,maxy,minx,miny = cityBoundary.bounds['maxx'][0],cityBoundary.bounds['maxy'][0],cityBoundary.bounds['minx'][0],cityBoundary.bounds['miny'][0]
    # convert to a numpy array
    data_array = np.array(getCityElevationData(cityName, cityBoundary, False))
    minElev = np.min(np.array(data_array))
    maxElev = np.max(np.array(data_array))

    data_array = (data_array - minElev) * 1000.0/ (maxElev - minElev)
    #Hillsahde
    # topoName = 'wkp/schwarzwald/wiki-schwarzwald-cont'
    topoName = 'terrain'
    # topoName = 'wiki-ice-greenland'

    # topocmap = pycpt.load.cmap_from_cptcity_url('wkp/encyclopedia/{}.cpt'.format(topoName))
    # topocmap = pycpt.load.cmap_from_cptcity_url('wkp/schwarzwald/nordisk-familjebok.cpt')
    topocmap = topoName

    vmin = 0
    vmax = 1000
    extent = minx, maxx, miny, maxy
    imH, imW = commons.getOutputImageHeightWidth()

    #Clean Version For Input
    # 192 pixels/ inch
    #
    # 2000 pixels / (192 pixels/inch) = 10.41 inch
    # 1200 pixels / (192 pixels/inch) = 6.25 inch
    # 4105= 2449
    # 21.38 × 12.75inch
    print((imW,imH))
    plt.clf()
    fig, ax = plt.subplots(figsize=(imW,imH))
    ax.axis('off')
    elevationfName = '{}{}_elevation'.format(commons.getOutputPath(cityName), cityName)
    cax = ax.contourf(data_array, np.arange(vmin, vmax, 200),extent=extent,
                      cmap=topocmap, vmin=vmin, vmax=vmax, origin='image', antialiased=True)
    fig.savefig('{}{}_elevationContourColored{}.png'.format(commons.getOutputPath(cityName), cityName, topoName),bbox_inches='tight', interpolation = 'nearest')

if __name__ == "__main__":
    main2()


import argparse
import collections

from torch.utils.data.dataloader import DataLoader
from evaluator import Evaluator
# torch
import torch
import torch.optim as optim
from torchvision import transforms
# retinanet
from retinanet.model import create_retinanet
from retinanet.dataloader import AspectRatioBasedSampler, IL_dataset, collater
from retinanet.dataloader import Resizer, Augmenter, Normalizer
# preprocessing
from preprocessing.params import Params
# train
import time
import numpy as np
# Global Setting

ROOT_DIR = "/home/deeplab307/Documents/Anaconda/Shiang/IL/"
PRINT_INFO = True # whether print some information about continual learning on the scrren
DEFAULT_ALPHA = 0.25
DEFAULT_GAMMA = 2.0
DEFAULT_BATCH_SIZE = 5

def fast_zero_grad(model):
    for param in model.parameters():
        param.grad = None

def str2bool(v):
    if isinstance(v, bool):
        return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

def create_tools(params:Params):
    start_state = params['start_state']
    start_epoch = params['start_epoch']
    # Training dataloader
    dataset_train = IL_dataset(params,
                               transform=transforms.Compose([Normalizer(), Augmenter(), Resizer()]),
                               start_state=start_state,
                               use_all_class=True)

    # Create the model
    if start_epoch == 1 and start_state != 0:
        retinanet = create_retinanet(params['depth'], params.states[start_state - 1]['num_knowing_class'])
    else:
        retinanet = create_retinanet(params['depth'], params.states[start_state]['num_knowing_class'])
    retinanet = retinanet.cuda()
    retinanet.training = True
    optimizer = optim.Adam(retinanet.parameters(), lr=params['lr'])

    # Read checkpoint
    if start_state != 0 or start_epoch != 1:
        if start_epoch == 1:
            params.load_model(start_state - 1,-1,retinanet)
        else:
            params.load_model(start_state,start_epoch - 1,retinanet)

    # IL_Trainer

    sampler = AspectRatioBasedSampler(dataset_train, batch_size = params['batch_size'], drop_last=False)
    dataloader_train = DataLoader(dataset_train, num_workers=2, collate_fn=collater, batch_sampler=sampler)
    return retinanet, optimizer, dataloader_train

def get_parser(args=None):
    parser = argparse.ArgumentParser()
    # must set params
    parser.add_argument('--root_dir', help='the root dir for training', default=ROOT_DIR)
    parser.add_argument('--dataset', help='Dataset name, must contain name and years, for instance: voc2007,voc2012', default='voc2007')
    parser.add_argument('--start_epoch', type=int)
    parser.add_argument('--end_epoch', help='Number of epochs', type=int)
    parser.add_argument('--start_state', type=int)
    parser.add_argument('--end_state', type=int)
    # retinanet params
    parser.add_argument('--alpha', type=float, default=DEFAULT_ALPHA)
    parser.add_argument('--gamma', type=float, default=DEFAULT_GAMMA)
    # Other params
    parser.add_argument('--record', help='whether record training with tensorboard default=True', type=str2bool, default=True)
    parser.add_argument('--debug', help='whether debug in Train process, default = False', type=str2bool, default=False)
    parser.add_argument('--val', help='whether do validation after training', type=str2bool, default=False)

    ###############################
    # Incremental learning params #
    ###############################

    # Warm up
    parser.add_argument('--warm_stage', help='the number of warm-up stage, 0 mean not warm up, default = 0', type=int, default=0)
    parser.add_argument('--warm_epoch', help='the number of epoch for each warm-up stage, use " "(space) to split epoch for different stage', type=int, nargs='*', default=[10,10])
    parser.add_argument('--warm_layers', help='the layers which will be warmed up, must be "output", "resnet", "fpn", and split each stage by space " "', nargs='*', default=['output','resnet'])

    # IL params
    parser.add_argument('--scenario', help='the scenario of states, must be "20", "19 1", "10 10", "15 1", "15 1 1 1 1"', nargs="+", default=[20])
    parser.add_argument('--shuffle_class', help='whether shuffle the class, default = False',type=str2bool , default=False)
    parser.add_argument('--distill', help='whether add distillation loss, default = False',type=str2bool , default=False)
    parser.add_argument('--enhance_on_new', type=str2bool, default=False)
    parser.add_argument('--ignore_past_class', help='when calculating the focal loss, whether ignore past class), default = False',type=str2bool , default=False)
    parser.add_argument('--decrease_positive', help="the upper score of the new class in incremental state, default=1.0",type=float , default=1.0)
    # Record
    parser.add_argument('--description', help="description for this experiment", default="None")

    # learning rate
    parser.add_argument('--lr', help="learning rate, default=1e-5", type=float, default=1e-5)
    parser.add_argument('--scheduler_milestone', type=int, nargs="+", default=[40])
    parser.add_argument('--scheduler_decay',  help="learning rate decay for scheduler, default=0.1", type=float, default=0.1)

    # always default paras
    parser.add_argument('--depth', help='Resnet depth, must be one of 18, 34, 50, 101, 152', type=int, default=50)
    parser.add_argument('--batch_size', help='batch_size', type=int, default=DEFAULT_BATCH_SIZE)

    parser = vars(parser.parse_args(args))
    return parser

def main(args=None):
    parser = get_parser(args)
    params = Params(parser)

    model, optimizer, dataloader_train = create_tools(params)

    model.train()
    model.freeze_layers(['classificationModel.output','regressionModel.output'])
    loss_hist = collections.deque(maxlen=500)
    print("Total IterNum:",len(dataloader_train))
    for epoch in range(params['start_epoch'], params['end_epoch'] + 1):
        for iter_num, data in enumerate(dataloader_train):
            start = time.time()
            fast_zero_grad(model)

            with torch.cuda.device(0):
                try:
                    cls_loss, reg_loss = model.cal_simple_focal_loss(data['img'].float().cuda(),
                                                data['annot'].cuda(),
                                                params)

                    # if cls_loss < 0.005:
                    #     continue

                    loss = cls_loss + reg_loss
                    loss.backward()
                    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1)
                    optimizer.step()
                    loss_hist.append(float(loss))
                    end = time.time()
                    print("Epoch: {} | Iter: {} | Cls_loss: {:.3f} | Reg_loss: {:.3f} | Total_loss: {:.3f} | Running_loss: {:.3f} | Time: {:.2f}s".format(epoch,
                                                                                iter_num,
                                                                                float(cls_loss),
                                                                                float(reg_loss),
                                                                                float(loss),
                                                                                np.mean(loss_hist),
                                                                                float(end - start)))
                except Exception as e:
                    print(e)
                    return None
        params.save_checkpoint(params['start_state'],epoch, model)
        if epoch % 5 == 0:
            params.auto_delete(params['start_state'],epoch)

if __name__ == '__main__':
    assert torch.__version__.split('.')[0] == '1'
    if not torch.cuda.is_available():
        print("CUDA isn't abailable")
        exit(-1)
    main()