# -*- coding: utf-8 -*-

import sys
import os
import logging
import torch
import math
import numpy as np
import glob

def numparameters(model):
  npars = 0 #pars
  nbytes = 0 #bytes
  for name, param in model.named_parameters():
    if param.requires_grad: #learnable parameters only
      npars += param.numel()
      nbytes += param.numel() * param.data.element_size() #returns size of each parameter
      logging.debug("{} => {} = {} x {} bytes".format(name, list(param.data.size()), param.data.numel(), param.data.element_size()))

  name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
  if nbytes == 0:
    i = 0
  else:
    i = int(math.floor(math.log(nbytes, 1024)))
    p = math.pow(1024, i)
    nbytes /= p
  size = "{:.2f}{}".format(nbytes, name[i])

  return npars, size

def load_checkpoint(suffix, model, optimizer, device):
  step = 0
  files = sorted(glob.glob("{}.checkpoint_????????.pt".format(suffix))) ### I check if there is one model
  if len(files) == 0:
    logging.info('No model found')
    sys.exit()

  file = files[-1] ### last is the newest
  checkpoint = torch.load(file, map_location=device)
  step = checkpoint['step']
  model.load_state_dict(checkpoint['model'])
  optimizer.load_state_dict(checkpoint['optimizer'])
  logging.info('Loaded model/optimizer step={} from {}'.format(step,file))
  return step, model, optimizer ### this is for learning

def save_checkpoint(suffix, model, optimizer, step, keep_last_n):
  checkpoint = { 'step': step, 'model': model.state_dict(), 'optimizer': optimizer.state_dict() }
  torch.save(checkpoint, "{}.checkpoint_{:08d}.pt".format(suffix,step))
  logging.info('Saved {}.checkpoint_{:08d}.pt'.format(suffix,step))
  files = sorted(glob.glob(suffix + '.checkpoint_????????.pt'))
  while keep_last_n > 0 and len(files) > keep_last_n:
    f = files.pop(0)
    os.remove(f) ### first is the oldest
    logging.debug('Removed checkpoint {}'.format(f))

def load_model(suffix, model, device, fmodel=None):
  if fmodel is not None:
    if not os.path.isfile(fmodel):
      logging.error('No model found')
      sys.exit()
  else:
    files = sorted(glob.glob("{}.checkpoint_????????.pt".format(suffix))) ### I check if there is one model
    if len(files) == 0:
      logging.info('No checkpoint found')
      sys.exit()
    fmodel = files[-1] ### last is the newest
  checkpoint = torch.load(fmodel, map_location=device)
  step = checkpoint['step']
  model.load_state_dict(checkpoint['model'])
  logging.info('Loaded model step={} from {}'.format(step,fmodel))
  return step, model

def prepare_source(batch_src, idx_pad, device):
  src = [torch.tensor(seq) for seq in batch_src] #[bs, ls]
  src = torch.nn.utils.rnn.pad_sequence(src, batch_first=True, padding_value=idx_pad).to(device) #[bs,ls]
  msk_src = (src != idx_pad).unsqueeze(-2) #[bs,1,ls] (False where <pad> True otherwise)
  return src, msk_src

def prepare_source_cross(batch_src, idx_pad, device):
  src = [torch.tensor(seq[1:]) for seq in batch_src] #[bs, ls]
  src = torch.nn.utils.rnn.pad_sequence(src, batch_first=True, padding_value=idx_pad).to(device) #[bs,ls]
  msk_src = (src != idx_pad).unsqueeze(-2) #[bs,1,ls] (False where <pad> True otherwise)
  return src, msk_src

def prepare_prefix(batch_pre, idx_pad, device):
  pre = [torch.tensor(seq) for seq in batch_pre] #[bs, lp]
  pre = torch.nn.utils.rnn.pad_sequence(pre, batch_first=True, padding_value=idx_pad).to(device) #[bs,lp]
  msk_pre = (pre != idx_pad).unsqueeze(-2) #[bs,1,ls] (False where <pad> True otherwise)
  return pre, msk_pre

def prepare_target(batch_tgt, idx_pad, device):
  tgt = [torch.tensor(seq[:-1]) for seq in batch_tgt] #delete <eos>
  tgt = torch.nn.utils.rnn.pad_sequence(tgt, batch_first=True, padding_value=idx_pad).to(device)
  #
  ref = [torch.tensor(seq[1:])  for seq in batch_tgt] #delete <bos>
  ref = torch.nn.utils.rnn.pad_sequence(ref, batch_first=True, padding_value=idx_pad).to(device)
  #
  msk_tgt = (tgt != idx_pad).unsqueeze(-2) & (1 - torch.triu(torch.ones((1, tgt.size(1), tgt.size(1)), device=tgt.device), diagonal=1)).bool() #[bs,lt,lt]
  return tgt, ref, msk_tgt

##############################################################################################################
### Endcoder_Decoder #########################################################################################
##############################################################################################################
class Encoder_Decoder(torch.nn.Module):
  #https://www.linzehui.me/images/16005200579239.jpg
  def __init__(self, n_layers, ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout, share_embeddings, src_voc_size, tgt_voc_size, idx_pad):
    super(Encoder_Decoder, self).__init__()
    self.idx_pad = idx_pad
    self.src_emb = Embedding(src_voc_size, emb_dim, idx_pad)
    self.tgt_emb = Embedding(tgt_voc_size, emb_dim, idx_pad)
    if share_embeddings:
      self.tgt_emb.emb.weight = self.src_emb.emb.weight

    self.add_pos_enc = AddPositionalEncoding(emb_dim, dropout, max_len=5000)
    self.stacked_encoder = Stacked_Encoder(n_layers, ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout)
    self.stacked_decoder = Stacked_Decoder(n_layers, ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout)
    self.generator = Generator(emb_dim, tgt_voc_size)

  def type(self):
    return 's_sc'

  def forward(self, src, tgt, msk_src, msk_tgt):
    #src is [bs,ls]
    #tgt is [bs,lt]
    #msk_src is [bs,1,ls] (False where <pad> True otherwise)
    #mst_tgt is [bs,lt,lt]
    ### encoder #####
    src = self.add_pos_enc(self.src_emb(src)) #[bs,ls,ed]
    z_src = self.stacked_encoder(src, msk_src) #[bs,ls,ed]
    ### decoder #####
    tgt = self.add_pos_enc(self.tgt_emb(tgt)) #[bs,lt,ed]
    z_tgt = self.stacked_decoder(z_src, tgt, msk_src, msk_tgt) #[bs,lt,ed]
    ### generator ###
    y = self.generator(z_tgt) #[bs, lt, Vt]
    return y ### returns logits (for learning)

  def encode(self, src, msk_src):
    src = self.add_pos_enc(self.src_emb(src)) #[bs,ls,ed]
    z_src = self.stacked_encoder(src, msk_src) #[bs,ls,ed]
    return z_src

  def decode(self, z_src, tgt, msk_src, msk_tgt=None):
    assert z_src.shape[0] == tgt.shape[0] ### src/tgt batch_sizes must be equal
    #z_src are the embeddings of the source words (encoder) [bs, sl, ed]
    #tgt is the history (words already generated) for current step [bs, lt]
    tgt = self.add_pos_enc(self.tgt_emb(tgt)) #[bs,lt,ed]
    z_tgt = self.stacked_decoder(z_src, tgt, msk_src, msk_tgt) #[bs,lt,ed]
    y = self.generator(z_tgt) #[bs, lt, Vt]
    y = torch.nn.functional.log_softmax(y, dim=-1)
    return y ### returns log_probs (for inference)

##############################################################################################################
### PositionalEncoding #######################################################################################
##############################################################################################################
class AddPositionalEncoding(torch.nn.Module):
  def __init__(self, emb_dim, dropout, max_len=5000):
    super(AddPositionalEncoding, self).__init__()
    assert emb_dim%2 == 0, 'emb_dim must be pair'
    self.dropout = torch.nn.Dropout(dropout)
    pe = torch.zeros(max_len, emb_dim) #[max_len=5000, ed]
    position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) #[max_len, 1]
    div_term = torch.exp(torch.arange(0, emb_dim, 2).float() * (-math.log(10000.0) / emb_dim)) #[ed/2]
    pe[:, 0::2] = torch.sin(position*div_term) #[max_len, 1] * [1, ed/2] => [max_len, ed] (pairs of pe)
    pe[:, 1::2] = torch.cos(position*div_term) #[max_len, 1] * [1, ed/2] => [max_len, ed] (odds of pe)
    pe = pe.unsqueeze(0) #[1, max_len=5000, ed]
    self.register_buffer('pe', pe) #register_buffer is for params which are saved&restored in state_dict but not trained

  def forward(self, x):
    bs, l, ed = x.shape
    x = x + self.pe[:, :l] #[bs, l, ed] + [1, l, ed] => [bs, l, ed]
    return self.dropout(x)

##############################################################################################################
### Stacked_Encoder ##########################################################################################
##############################################################################################################
class Stacked_Encoder(torch.nn.Module):
  def __init__(self, n_layers, ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout):
    super(Stacked_Encoder, self).__init__()
    self.encoderlayers = torch.nn.ModuleList([Encoder(ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout) for _ in range(n_layers)])
    self.norm = torch.nn.LayerNorm(emb_dim, eps=1e-6)

  def forward(self, src, msk):
    for i,encoderlayer in enumerate(self.encoderlayers):
      src = encoderlayer(src, msk) #[bs, ls, ed]
    return self.norm(src)

##############################################################################################################
### Encoder ##################################################################################################
##############################################################################################################
class Encoder(torch.nn.Module):
  def __init__(self, ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout):
    super(Encoder, self).__init__()
    self.multihead_attn = MultiHead_Attn(n_heads, emb_dim, qk_dim, v_dim, dropout)
    self.feedforward = FeedForward(emb_dim, ff_dim, dropout)
    self.norm_att = torch.nn.LayerNorm(emb_dim, eps=1e-6)
    self.norm_ff = torch.nn.LayerNorm(emb_dim, eps=1e-6)

  def forward(self, src, msk):
    #NORM
    tmp1 = self.norm_att(src)
    #ATTN over source words
    tmp2 = self.multihead_attn(q=tmp1, k=tmp1, v=tmp1, msk=msk) #[bs, ls, ed] contains dropout
    #ADD
    tmp = tmp2 + src

    #NORM
    tmp1 = self.norm_ff(tmp)
    #FF
    tmp2 = self.feedforward(tmp1) #[bs, ls, ed] contains dropout
    #ADD
    z = tmp2 + tmp
    return z

##############################################################################################################
### Stacked_Decoder ##########################################################################################
##############################################################################################################
class Stacked_Decoder(torch.nn.Module):
  def __init__(self, n_layers, ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout):
    super(Stacked_Decoder, self).__init__()
    self.decoderlayers = torch.nn.ModuleList([Decoder(ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout) for _ in range(n_layers)])
    self.norm = torch.nn.LayerNorm(emb_dim, eps=1e-6)

  def forward(self, z_src, tgt, msk_src, msk_tgt):
    for i,decoderlayer in enumerate(self.decoderlayers):
      tgt = decoderlayer(z_src, tgt, msk_src, msk_tgt)
    return self.norm(tgt)

##############################################################################################################
### Decoder ##################################################################################################
##############################################################################################################
class Decoder(torch.nn.Module):
  def __init__(self, ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout):
    super(Decoder, self).__init__()
    self.multihead_attn_self = MultiHead_Attn(n_heads, emb_dim, qk_dim, v_dim, dropout)
    self.multihead_attn_enc = MultiHead_Attn(n_heads, emb_dim, qk_dim, v_dim, dropout)
    self.feedforward = FeedForward(emb_dim, ff_dim, dropout)
    self.norm_att_self = torch.nn.LayerNorm(emb_dim, eps=1e-6)
    self.norm_att_enc = torch.nn.LayerNorm(emb_dim, eps=1e-6)
    self.norm_ff = torch.nn.LayerNorm(emb_dim, eps=1e-6)

  def forward(self, z_src, tgt, msk_src, msk_tgt):
    #NORM
    tmp1 = self.norm_att_self(tgt)
    #ATTN over tgt (previous) words : q, k, v are tgt words
    tmp2 = self.multihead_attn_self(q=tmp1, k=tmp1, v=tmp1, msk=msk_tgt) #[bs, lt, ed] contains dropout
    #ADD
    tmp = tmp2 + tgt

    #NORM
    tmp1 = self.norm_att_enc(tmp)
    #ATTN over src words : q are tgt words, k, v are src words
    tmp2 = self.multihead_attn_enc(q=tmp1, k=z_src, v=z_src, msk=msk_src) #[bs, lt, ed] contains dropout
    #ADD
    tmp = tmp2 + tmp

    #NORM
    tmp1 = self.norm_ff(tmp)
    #FF
    tmp2 = self.feedforward(tmp1) #[bs, lt, ed] contains dropout
    #ADD
    z = tmp2 + tmp
    return z

##############################################################################################################
### Stacked_CrossAdapter #####################################################################################
##############################################################################################################
class Stacked_CrossAdapter(torch.nn.Module):
  def __init__(self, n_layers, ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout):
    super(Stacked_CrossAdapter, self).__init__()
    self.adapterlayers = torch.nn.ModuleList([CrossAdapter(ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout) for _ in range(n_layers)])
    self.norm = torch.nn.LayerNorm(emb_dim, eps=1e-6)

  def forward(self, z_src, tgt, msk_src, msk_tgt):
    for i,adapterlayer in enumerate(self.adapterlayers):
      tgt = adapterlayer(z_src, tgt, msk_src, msk_tgt)
    return self.norm(tgt)

##############################################################################################################
### CrossAdapter #############################################################################################
##############################################################################################################
class CrossAdapter(torch.nn.Module):
  def __init__(self, ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout):
    super(CrossAdapter, self).__init__()
    self.multihead_attn = MultiHead_Attn(n_heads, emb_dim, qk_dim, v_dim, dropout)
    self.norm_att_cross = torch.nn.LayerNorm(emb_dim, eps=1e-6)

  def forward(self, z_src, tgt, msk_src, msk_tgt):
    #NORM
    tmp_norm = self.norm_att_cross(tgt)
    #ATTN over src words : q are tgt words, k, v are src words
    tmp = self.multihead_attn(q=tmp_norm, k=z_src, v=z_src, msk=msk_src) #[bs, lt, ed] contains dropout
    #ADD
    tmp = tmp + tgt
    return tmp

##############################################################################################################
### Stacked_Encoder_scc ######################################################################################
##############################################################################################################
class Stacked_Encoder_scc(torch.nn.Module):
  def __init__(self, n_layers, ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout):
    super(Stacked_Encoder_scc, self).__init__()
    self.encoderlayers = torch.nn.ModuleList([Encoder_scc(ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout) for _ in range(n_layers)])
    self.norm = torch.nn.LayerNorm(emb_dim, eps=1e-6)

  def forward(self, z_src, z_xsrc, xtgt, msk_src, msk_xsrc, msk_xtgt):
    for i,encoderlayer in enumerate(self.encoderlayers):
      xtgt = encoderlayer(z_src, z_xsrc, xtgt, msk_src, msk_xsrc, msk_xtgt) #[bs, ls, ed]
    return self.norm(xtgt)

##############################################################################################################
### Encoder_scc ##############################################################################################
##############################################################################################################
class Encoder_scc(torch.nn.Module):
  def __init__(self, ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout):
    super(Encoder_scc, self).__init__()
    self.multihead_attn_self = MultiHead_Attn(n_heads, emb_dim, qk_dim, v_dim, dropout)
    self.multihead_attn_cross1 = MultiHead_Attn(n_heads, emb_dim, qk_dim, v_dim, dropout)
    self.multihead_attn_cross2 = MultiHead_Attn(n_heads, emb_dim, qk_dim, v_dim, dropout)
    self.feedforward = FeedForward(emb_dim, ff_dim, dropout)
    self.norm_att_self = torch.nn.LayerNorm(emb_dim, eps=1e-6)
    self.norm_att_cross1 = torch.nn.LayerNorm(emb_dim, eps=1e-6)
    self.norm_att_cross2 = torch.nn.LayerNorm(emb_dim, eps=1e-6)
    self.norm_ff = torch.nn.LayerNorm(emb_dim, eps=1e-6)

  def forward(self, z_src, z_xsrc, xtgt, msk_src, msk_xsrc, msk_xtgt):
    #NORM
    tmp_norm = self.norm_att_self(xtgt)
    #Self ATTN over xtgt words
    tmp2 = self.multihead_attn_self(q=tmp_norm, k=tmp_norm, v=tmp_norm, msk=msk_xtgt) #[bs, ls, ed] contains dropout
    #ADD
    tmp = tmp2 + xtgt

    #NORM
    tmp_norm = self.norm_att_cross1(tmp)
    #Cross ATTN over xsrc words : q are xtgt words, k, v are xsrc words
    tmp2 = self.multihead_attn_cross1(q=tmp_norm, k=z_xsrc, v=z_xsrc, msk=msk_xsrc) #[bs, ls, ed] contains dropout
    #ADD
    tmp = tmp2 + tmp

    #NORM
    tmp_norm = self.norm_att_cross2(tmp)
    #Cross ATTN over src words : q are xtgt words, k, v are src words
    tmp2 = self.multihead_attn_cross2(q=tmp_norm, k=z_src, v=z_src, msk=msk_src) #[bs, ls, ed] contains dropout
    #ADD
    tmp = tmp2 + tmp

    #NORM
    tmp_norm = self.norm_ff(tmp)
    #FF
    tmp2 = self.feedforward(tmp_norm) #[bs, ls, ed] contains dropout
    #ADD
    z = tmp2 + tmp
    return z

##############################################################################################################
### Stacked_Encoder_sc ######################################################################################
##############################################################################################################
class Stacked_Encoder_sc(torch.nn.Module):
  def __init__(self, n_layers, ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout):
    super(Stacked_Encoder_sc, self).__init__()
    self.encoderlayers = torch.nn.ModuleList([Encoder_sc(ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout) for _ in range(n_layers)])
    self.norm = torch.nn.LayerNorm(emb_dim, eps=1e-6)

  def forward(self, z_src, xtgt, msk_src, msk_xtgt):
    for i,encoderlayer in enumerate(self.encoderlayers):
      xtgt = encoderlayer(z_src, xtgt, msk_src, msk_xtgt) #[bs, ls, ed]
    return self.norm(xtgt)

##############################################################################################################
### Encoder_sc ###############################################################################################
##############################################################################################################
class Encoder_sc(torch.nn.Module):
  def __init__(self, ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout):
    super(Encoder_sc, self).__init__()
    self.multihead_attn_self = MultiHead_Attn(n_heads, emb_dim, qk_dim, v_dim, dropout)
    self.multihead_attn_cross = MultiHead_Attn(n_heads, emb_dim, qk_dim, v_dim, dropout)
    self.feedforward = FeedForward(emb_dim, ff_dim, dropout)
    self.norm_att_self = torch.nn.LayerNorm(emb_dim, eps=1e-6)
    self.norm_att_cross = torch.nn.LayerNorm(emb_dim, eps=1e-6)
    self.norm_ff = torch.nn.LayerNorm(emb_dim, eps=1e-6)

  def forward(self, z_src, xtgt, msk_src, msk_xtgt):
    #NORM
    tmp_norm = self.norm_att_self(xtgt)
    #Self ATTN over xtgt words
    tmp2 = self.multihead_attn_self(q=tmp_norm, k=tmp_norm, v=tmp_norm, msk=msk_xtgt) #[bs, ls, ed] contains dropout
    #ADD
    tmp = tmp2 + xtgt

    #NORM
    tmp_norm = self.norm_att_cross(tmp)
    #Cross ATTN over xsrc words : q are xtgt words, k, v are xsrc words
    tmp2 = self.multihead_attn_cross(q=tmp_norm, k=z_src, v=z_src, msk=msk_src) #[bs, ls, ed] contains dropout
    #ADD
    tmp = tmp2 + tmp

    #NORM
    tmp_norm = self.norm_ff(tmp)
    #FF
    tmp2 = self.feedforward(tmp_norm) #[bs, ls, ed] contains dropout
    #ADD
    z = tmp2 + tmp
    return z

##############################################################################################################
### Stacked_Decoder_scc ######################################################################################
##############################################################################################################
class Stacked_Decoder_scc(torch.nn.Module):
  def __init__(self, n_layers, ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout):
    super(Stacked_Decoder_scc, self).__init__()
    self.decoderlayers = torch.nn.ModuleList([Decoder_scc(ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout) for _ in range(n_layers)])
    self.norm = torch.nn.LayerNorm(emb_dim, eps=1e-6)

  def forward(self, z_src, z_xtgt, tgt, msk_src, msk_xtgt, msk_tgt):
    for i,decoderlayer in enumerate(self.decoderlayers):
      tgt = decoderlayer(z_src, z_xtgt, tgt, msk_src, msk_xtgt, msk_tgt)
    return self.norm(tgt)

##############################################################################################################
### Decoder_scc ##############################################################################################
##############################################################################################################
class Decoder_scc(torch.nn.Module):
  def __init__(self, ff_dim, n_heads, emb_dim, qk_dim, v_dim, dropout):
    super(Decoder_scc, self).__init__()
    self.multihead_attn_self = MultiHead_Attn(n_heads, emb_dim, qk_dim, v_dim, dropout)
    self.multihead_attn_cross1 = MultiHead_Attn(n_heads, emb_dim, qk_dim, v_dim, dropout)
    self.multihead_attn_cross2 = MultiHead_Attn(n_heads, emb_dim, qk_dim, v_dim, dropout)
    self.feedforward = FeedForward(emb_dim, ff_dim, dropout)
    self.norm_att_self = torch.nn.LayerNorm(emb_dim, eps=1e-6)
    self.norm_att_cross1 = torch.nn.LayerNorm(emb_dim, eps=1e-6)
    self.norm_att_cross2 = torch.nn.LayerNorm(emb_dim, eps=1e-6)
    self.norm_ff = torch.nn.LayerNorm(emb_dim, eps=1e-6)

  def forward(self, z_src, z_xtgt, tgt, msk_src, msk_xtgt, msk_tgt):
    #NORM
    tmp_norm = self.norm_att_self(tgt)
    #Self ATTN over tgt (previous) words : q, k, v are tgt words
    tmp2 = self.multihead_attn_self(q=tmp_norm, k=tmp_norm, v=tmp_norm, msk=msk_tgt) #[bs, lt, ed] contains dropout
    #ADD
    tmp = tmp2 + tgt

    #NORM
    tmp_norm = self.norm_att_cross1(tmp)
    #Cross ATTN over xtgt words : q are tgt words, k, v are xtgt words
    tmp2 = self.multihead_attn_cross1(q=tmp_norm, k=z_xtgt, v=z_xtgt, msk=msk_xtgt) #[bs, lt, ed] contains dropout
    #ADD
    tmp = tmp2 + tmp

    #NORM
    tmp_norm = self.norm_att_cross2(tmp)
    #Cross ATTN over src words : q are tgt words, k, v are src words
    tmp2 = self.multihead_attn_cross2(q=tmp_norm, k=z_src, v=z_src, msk=msk_src) #[bs, lt, ed] contains dropout
    #ADD
    tmp = tmp2 + tmp

    #NORM
    tmp_norm = self.norm_ff(tmp)
    #FF
    tmp2 = self.feedforward(tmp_norm) #[bs, lt, ed] contains dropout
    #ADD
    z = tmp2 + tmp
    return z

##############################################################################################################
### MultiHead_Attn ###########################################################################################
##############################################################################################################
class MultiHead_Attn(torch.nn.Module):
  def __init__(self, n_heads, emb_dim, qk_dim, v_dim, dropout):
    super(MultiHead_Attn, self).__init__()
    self.nh = n_heads
    self.ed = emb_dim
    self.qd = qk_dim
    self.kd = qk_dim
    self.vd = v_dim
    self.WQ = torch.nn.Linear(emb_dim, qk_dim*n_heads)
    self.WK = torch.nn.Linear(emb_dim, qk_dim*n_heads)
    self.WV = torch.nn.Linear(emb_dim, v_dim*n_heads)
    self.WO = torch.nn.Linear(v_dim*n_heads, emb_dim)
    self.dropout = torch.nn.Dropout(dropout)

  def forward(self, q, k, v, msk=None):
    #q is [bs, lq, ed]
    #k is [bs, lk, ed]
    #v is [bs, lv, ed]
    #msk is [bs, 1, ls] or [bs, lt, lt]
    #logging.info('q = {}'.format(q.shape))
    #logging.info('k = {}'.format(k.shape))
    #logging.info('v = {}'.format(v.shape))

    if msk is not None:
      msk = msk.unsqueeze(1) #[bs, 1, 1, ls] or [bs, 1, lt, lt]

    #logging.info('msk = {}'.format(msk.shape))

    bs = q.shape[0]
    lq = q.shape[1] ### sequence length of q vectors (length of target sentences)
    lk = k.shape[1] ### sequence length of k vectors (may be length of source/target sentences)
    lv = v.shape[1] ### sequence length of v vectors (may be length of source/target sentences)
    ed = q.shape[2]
    assert self.ed == q.shape[2] == k.shape[2] == v.shape[2]
    assert lk == lv #when applied in decoder both refer the source-side (lq refers the target-side)
    Q = self.WQ(q).contiguous().view([bs,lq,self.nh,self.qd]).permute(0,2,1,3) #=> [bs,lq,nh*qd] => [bs,lq,nh,qd] => [bs,nh,lq,qd]
    K = self.WK(k).contiguous().view([bs,lk,self.nh,self.kd]).permute(0,2,1,3) #=> [bs,lk,nh*kd] => [bs,lk,nh,kd] => [bs,nh,lk,kd]
    V = self.WV(v).contiguous().view([bs,lv,self.nh,self.vd]).permute(0,2,1,3) #=> [bs,lv,nh*vd] => [bs,lv,nh,vd] => [bs,nh,lv,vd]
    #Scaled dot-product Attn from multiple Q, K, V vectors (bs*nh*l vectors)
    Q = Q / math.sqrt(self.kd)
    s = torch.matmul(Q, K.transpose(2, 3)) #[bs,nh,lq,qd] x [bs,nh,kd,lk] = [bs,nh,lq,lk] # thanks to qd==kd #in decoder lq are target words and lk are source words

    #logging.info('s = {}'.format(s.shape))

    if msk is not None:
      s = s.masked_fill(msk == 0, float('-inf')) #score=-Inf to masked tokens
    w = torch.nn.functional.softmax(s, dim=-1) #[bs,nh,lq,lk] (these are the attention weights)
    #### Minh uses relu instead of softmax: w = torch.nn.functional.relu(s)
    w = self.dropout(w) #[bs,nh,lq,lk]

    z = torch.matmul(w,V) #[bs,nh,lq,lk] x [bs,nh,lv,vd] = [bs,nh,lq,vd] #thanks to lk==lv
    z = z.transpose(1,2).contiguous().view([bs,lq,self.nh*self.vd]) #=> [bs,lq,nh,vd] => [bs,lq,nh*vd]
    z = self.WO(z) #[bs,lq,ed]
    return self.dropout(z)

##############################################################################################################
### Embedding ################################################################################################
##############################################################################################################
class Embedding(torch.nn.Module):
  def __init__(self, vocab_size, emb_dim, idx_pad):
    super(Embedding, self).__init__()
    self.emb = torch.nn.Embedding(vocab_size, emb_dim, padding_idx=idx_pad)
    self.sqrt_emb_dim = math.sqrt(emb_dim)

  def forward(self, x):
    return self.emb(x) * self.sqrt_emb_dim

##############################################################################################################
### FeedForward ##############################################################################################
##############################################################################################################
class FeedForward(torch.nn.Module):
  def __init__(self, emb_dim, ff_dim, dropout):
    super(FeedForward, self).__init__()
    self.FF_in = torch.nn.Linear(emb_dim, ff_dim)
    self.FF_out = torch.nn.Linear(ff_dim, emb_dim)
    self.dropout = torch.nn.Dropout(dropout) #this regularization is not used in the original model

  def forward(self, x):
    tmp = self.FF_in(x)
    tmp = torch.nn.functional.relu(tmp)
    tmp = self.dropout(tmp)
    tmp = self.FF_out(tmp)
    tmp = self.dropout(tmp)
    return tmp

##############################################################################################################
### Generator ################################################################################################
##############################################################################################################
class Generator(torch.nn.Module):
  def __init__(self, emb_dim, voc_size):
    super(Generator, self).__init__()
    self.proj = torch.nn.Linear(emb_dim, voc_size) #[bs, Vt]

  def forward(self, x):
    y = self.proj(x)
    return y


from typing import Union
from typing import Optional
from typing import List
from typing import Sequence
from typing import IO
from typing import Any
from typing import cast
from typing import overload
from typing import Literal

import os
import io
import tempfile
import pathlib

from rdkit.Chem import PandasTools
from rdkit.Chem import rdmolfiles  # type: ignore

import pandas as pd
import fsspec
import fsspec.utils

import datamol as dm

from .types import Mol

def read_csv(
    urlpath: Union[str, os.PathLike, IO],
    smiles_column: Optional[str] = None,
    mol_column: str = "mol",
    **kwargs: Any,
) -> pd.DataFrame:
    """Read a CSV file.

    Args:
        urlpath: Path to a file or a file-like object. Path can be remote or local.
        smiles_column: Use this column to build a mol column.
        mol_column: Name to give to the mol column. If not None a mol column will be build.
            Avoid when loading a very large file.
        **kwargs: Arguments to pass to `pd.read_csv()`.

    Returns:
        df: a `pandas.DataFrame`
    """

    df = pd.read_csv(urlpath, **kwargs)
    df = cast(pd.DataFrame, df)

    if smiles_column is not None:
        PandasTools.AddMoleculeColumnToFrame(df, smiles_column, mol_column)

    return df

def read_excel(
    urlpath: Union[str, os.PathLike, IO],
    sheet_name: Optional[Union[str, int, list]] = 0,
    smiles_column: Optional[str] = None,
    mol_column: str = "mol",
    **kwargs: Any,
) -> pd.DataFrame:
    """Read an excel file.

    Args:
        urlpath: Path to a file or a file-like object. Path can be remote or local.
        sheet_name: see `pandas.read_excel()` doc.
        mol_column: Name to give to the mol column. If not None a mol column will be build.
            Avoid when loading a very large file.
        mol_column: name to give to the mol column.
        **kwargs: Arguments to pass to `pd.read_excel()`.

    Returns:
        df: a `pandas.DataFrame`
    """

    df = pd.read_excel(urlpath, sheet_name=sheet_name, **kwargs)
    df = cast(pd.DataFrame, df)

    if smiles_column is not None:
        PandasTools.AddMoleculeColumnToFrame(df, smiles_column, mol_column)

    return df

def _get_supplier_mols(
    supplier: rdmolfiles.ForwardSDMolSupplier,
    max_num_mols: Optional[int],
) -> List[dm.Mol]:
    """Given a supplier, read the molecules until we reach the `max_num_mols` limit.
    Useful when reading SDF files.
    """
    if max_num_mols is None:
        mols = list(supplier)
    else:
        mols = []
        for _ in range(max_num_mols):
            try:
                mols.append(next(supplier))
            except StopIteration:
                break
    return mols

@overload
def read_sdf(
    urlpath: Union[str, os.PathLike, IO],
    sanitize: bool = ...,
    as_df: Literal[False] = ...,
    smiles_column: Optional[str] = ...,
    mol_column: Optional[str] = ...,
    include_private: bool = ...,
    include_computed: bool = ...,
    strict_parsing: bool = ...,
    remove_hs: bool = ...,
    max_num_mols: Optional[int] = ...,
    discard_invalid: bool = ...,
    n_jobs: Optional[int] = ...,
) -> List[Mol]:
    ...

@overload
def read_sdf(
    urlpath: Union[str, os.PathLike, IO],
    sanitize: bool = ...,
    as_df: Literal[True] = ...,
    smiles_column: Optional[str] = ...,
    mol_column: Optional[str] = ...,
    include_private: bool = ...,
    include_computed: bool = ...,
    strict_parsing: bool = ...,
    remove_hs: bool = ...,
    max_num_mols: Optional[int] = ...,
    discard_invalid: bool = ...,
    n_jobs: Optional[int] = ...,
) -> pd.DataFrame:
    ...

@overload
def read_sdf(
    urlpath: Union[str, os.PathLike, IO],
    sanitize: bool = ...,
    as_df: bool = ...,
    smiles_column: Optional[str] = ...,
    mol_column: Optional[str] = ...,
    include_private: bool = ...,
    include_computed: bool = ...,
    strict_parsing: bool = ...,
    remove_hs: bool = ...,
    max_num_mols: Optional[int] = ...,
    discard_invalid: bool = ...,
    n_jobs: Optional[int] = ...,
) -> Union[List[Mol], pd.DataFrame]:
    ...

def read_sdf(
    urlpath: Union[str, os.PathLike, IO],
    sanitize: bool = True,
    as_df: bool = False,
    smiles_column: Optional[str] = "smiles",
    mol_column: Optional[str] = None,
    include_private: bool = False,
    include_computed: bool = False,
    strict_parsing: bool = True,
    remove_hs: bool = True,
    max_num_mols: Optional[int] = None,
    discard_invalid: bool = True,
    n_jobs: Optional[int] = 1,
) -> Union[List[Mol], pd.DataFrame]:
    """Read an SDF file.

    Note: This function is meant to be used with dataset that fit _in-memory_.
    For a more advanced usage we suggest you to use directly `Chem.ForwardSDMolSupplier`.

    Args:
        urlpath: Path to a file or a file-like object. Path can be remote or local.
        sanitize: Whether to sanitize the molecules.
        as_df: Whether to return a list mol or a pandas DataFrame.
        smiles_column: Name of the SMILES column. Only relevant if `as_df` is True.
        mol_column: Name of the mol column. Only relevant if `as_df` is True.
        include_private: Include private properties in the columns.  Only relevant if
            `as_df` is True.
        include_computed: Include computed properties in the columns.  Only relevant if
            `as_df` is True.
        strict_parsing: If set to false, the parser is more lax about correctness of the contents.
        remove_hs: Whether to remove the existing hydrogens in the SDF files.
        max_num_mols: Maximum number of molecules to read from the SDF file. Read all by default when set
            to `None`.
        discard_invalid: Discard the molecules that failed to be read correctly. Otherwise,
            invalid molecules will be loaded as `None`.
        n_jobs: Optional number of jobs for parallelization of `to_df`. Leave to 1 for no
            parallelization. Set to -1 to use all available cores. Only relevant is `as_df` is True
    """

    # File-like object
    if isinstance(urlpath, io.IOBase):
        supplier = rdmolfiles.ForwardSDMolSupplier(
            urlpath,
            sanitize=sanitize,
            strictParsing=strict_parsing,
            removeHs=remove_hs,
        )
        mols = _get_supplier_mols(supplier, max_num_mols)

    # Regular local or remote paths
    else:
        with fsspec.open(urlpath, compression="infer") as f:
            supplier = rdmolfiles.ForwardSDMolSupplier(
                f,
                sanitize=sanitize,
                strictParsing=strict_parsing,
                removeHs=remove_hs,
            )
            mols = _get_supplier_mols(supplier, max_num_mols)

    # Discard None values
    if discard_invalid:
        mols = [mol for mol in mols if mol is not None]

    # Convert to dataframe
    if as_df:
        return dm.to_df(
            mols,
            smiles_column=smiles_column,
            mol_column=mol_column,
            include_private=include_private,
            include_computed=include_computed,
            n_jobs=n_jobs,
        )  # type: ignore

    return mols

def to_sdf(
    mols: Union[Mol, Sequence[Mol], pd.DataFrame],
    urlpath: Union[str, os.PathLike, IO],
    smiles_column: Optional[str] = "smiles",
    mol_column: Optional[str] = None,
):
    """Write molecules to a file.

    Args:
        mols: a dataframe, a molecule or a list of molecule.
        urlpath: Path to a file or a file-like object. Path can be remote or local.
        smiles_column: Column name to extract the molecule.
        mol_column: Column name to extract the molecule. It takes
            precedence over `smiles_column`.
    """

    if isinstance(mols, pd.DataFrame):
        mols = dm.from_df(mols, smiles_column=smiles_column, mol_column=mol_column)

    elif isinstance(mols, Mol):
        mols = [mols]

    # Filter out None values
    mols = [mol for mol in mols if mol is not None]

    # File-like object
    if isinstance(urlpath, io.IOBase):
        writer = rdmolfiles.SDWriter(urlpath)
        for mol in mols:
            writer.write(mol)
        writer.close()

    # Regular local or remote paths
    else:
        with fsspec.open(urlpath, mode="w") as f:
            writer = rdmolfiles.SDWriter(f)
            for mol in mols:
                writer.write(mol)
            writer.close()

def read_mol2file(
    urlpath: Union[str, os.PathLike, IO],
    sanitize: bool = True,
    cleanup_substructures: bool = True,
    remove_hs: bool = True,
    fail_if_invalid: bool = False,
) -> List[Mol]:
    """Read a Mol2 File

    Args:
        urlpath: Path to a file or a file-like object. Path can be remote or local.
        sanitize: Whether to sanitize the molecules.
        remove_hs: Whether to remove the existing hydrogens in the SDF files.
        cleanup_substructures: Whether to clean up substructure in the Mol2 Files.
        fail_if_invalid: If set to true, the parser will raise an exception if the molecule is invalid
            instead of returning None.
    """

    block = []
    mols = []
    with fsspec.open(urlpath, compression="infer") as f:
        f = cast(IO, f)
        fReadLines = f.readlines()
        # reversing due to ambiguous end line for mol2 files
        fReadLines.reverse()
        for line in fReadLines:
            # ignores any header info
            if b"#" not in line:
                block.append(str(line, "utf-8"))
            # since reversed, this is the 'end' a mol2
            if b"@<TRIPOS>MOLECULE" in line:
                block.reverse()
                mol2block = ",".join(block).replace(",", "")
                mol = rdmolfiles.MolFromMol2Block(
                    mol2block,
                    sanitize=sanitize,
                    removeHs=remove_hs,
                    cleanupSubstructures=cleanup_substructures,
                )
                if mol is None and fail_if_invalid:
                    raise ValueError(f"Invalid molecule: {mol2block}")
                mols.append(mol)
                block = []

    mols.reverse()
    return mols

def read_molblock(
    molblock: str,
    sanitize: bool = True,
    strict_parsing: bool = True,
    remove_hs: bool = True,
    fail_if_invalid: bool = False,
) -> Optional[dm.Mol]:
    """Read a Mol block.

    Note that potential molecule properties are **not** read.

    Args:
        molblock: String containing the Mol block.
        sanitize: Whether to sanitize the molecules.
        strict_parsing: If set to false, the parser is more lax about correctness of the contents.
        remove_hs: Whether to remove the existing hydrogens in the SDF files.
        fail_if_invalid: If set to true, the parser will raise an exception if the molecule is invalid
            instead of returning None.
    """

    mol = rdmolfiles.MolFromMolBlock(
        molblock,
        sanitize=sanitize,
        removeHs=remove_hs,
        strictParsing=strict_parsing,
    )

    if mol is None and fail_if_invalid:
        raise ValueError(f"Invalid molecule: {molblock}")

    return mol

def to_molblock(
    mol: Mol,
    include_stereo: bool = True,
    conf_id: int = -1,
    kekulize: bool = True,
    force_V3000: bool = False,
):
    """Convert a molecule to a mol block string.

    Note that any molecule properties are lost.

    Args:
        mol: A molecule.
        include_stereo: Toggles inclusion of stereochemical information in the output.
        conf_id: Selects which conformation to output.
        kekulize: Triggers kekulization of the molecule before it's written,
            as suggested by the MDL spec.
        force_V3000: Force generation a V3000 mol block (happens automatically
            with more than 999 atoms or bonds).
    """

    molblock = rdmolfiles.MolToMolBlock(
        mol,
        includeStereo=include_stereo,
        confId=conf_id,
        kekulize=kekulize,
        forceV3000=force_V3000,
    )

    return molblock

def read_pdbblock(
    molblock: str,
    sanitize: bool = True,
    remove_hs: bool = True,
    flavor: int = 0,
    proximity_bonding: bool = True,
) -> dm.Mol:
    """Read a PDB string block.

    Args:
        molblock: String containing the Mol block.
        sanitize: Whether to sanitize the molecules.
        remove_hs: Whether to remove the existing hydrogens in the SDF files.
        flavor: RDKit flavor options.
        proximity_bonding: Whether to toggles automatic proximity bonding.
    """

    mol = rdmolfiles.MolFromPDBBlock(
        molblock,
        sanitize=sanitize,
        removeHs=remove_hs,
        flavor=flavor,
        proximityBonding=proximity_bonding,
    )
    return mol

def to_pdbblock(mol: Mol, conf_id: int = -1) -> str:
    """Convert a molecule to a PDB string block.

    Args:
        mol: A molecule.
        conf_id: Selects which conformation to use.
    """
    molblock = rdmolfiles.MolToPDBBlock(mol, confId=conf_id)
    return molblock

def read_pdbfile(
    urlpath: Union[str, os.PathLike],
    sanitize: bool = True,
    remove_hs: bool = True,
    flavor: int = 0,
    proximity_bonding: bool = True,
) -> Mol:
    """Read a PDB file.

    Args:
        urlpath: Path to a file or a file-like object. Path can be remote or local.
        sanitize: Whether to sanitize the molecules.
        remove_hs: Whether to remove the existing hydrogens in the SDF files.
        flavor: RDKit flavor options.
        proximity_bonding: Whether to toggles automatic proximity bonding.

    Returns:
        mol: a molecule
    """

    with fsspec.open(urlpath, "r") as f:
        f = cast(IO, f)
        mol = read_pdbblock(
            f.read(),
            sanitize=sanitize,
            remove_hs=remove_hs,
            flavor=flavor,
            proximity_bonding=proximity_bonding,
        )
    return mol

def to_pdbfile(
    mol: Mol,
    urlpath: Union[str, os.PathLike],
    conf_id: int = -1,
):
    """Save a molecule to a PDB file.

    Args:
        mol: A molecule.
        urlpath: Path to a file or a file-like object. Path can be remote or local.
        conf_id: Selects which conformation to use.
    """
    molblock = to_pdbblock(mol, conf_id=conf_id)
    with fsspec.open(urlpath, "w") as f:
        f = cast(IO, f)
        f.write(molblock)

def to_smi(
    mols: Sequence[Mol],
    urlpath: Union[str, os.PathLike, IO],
    error_if_empty: bool = False,
):
    """Save a list of molecules in an `.smi` file.

    Note: We **strongly** recommend you to use `dm.to_csv` instead
    of `dm.to_smi` since `.smi` files are CSV-like format. The only difference are the
    default settings which changes:

    - The default separator is a space ` ` instead of a comma `,`.
    - The headers of the column are not included.

    By modifying the args of `dm.to_csv()`, you will be able to save a SMI compatible file.

    Args:
        mols: a list of molecules.
        urlpath: Path to a file or a file-like object. Path can be remote or local.
        error_if_empty: whether to raise and error if the input list is empty.
    """

    if len(mols) == 0 and error_if_empty:
        raise ValueError("The list of mols/smiles provided is empty.")

    # Filter out None values
    mols = [mol for mol in mols if mol is not None]

    # File-like object
    if isinstance(urlpath, io.IOBase):
        writer = rdmolfiles.SmilesWriter(urlpath, includeHeader=False, nameHeader="")
        for mol in mols:
            writer.write(mol)
        writer.close()

    # Regular local or remote paths
    else:
        with fsspec.open(urlpath, "w") as f:
            writer = rdmolfiles.SmilesWriter(f, includeHeader=False, nameHeader="")
            for mol in mols:
                writer.write(mol)
            writer.close()

def read_smi(
    urlpath: Union[str, pathlib.Path, io.IOBase, fsspec.core.OpenFile],
) -> Sequence[Mol]:
    """Read a list of smiles from am `.smi` file.

    Note: We **strongly** recommend you to use `dm.read_csv` or `pandas.read_csv` instead
    of `dm.read_smi` since `.smi` files are CSV-like format. The only difference are the
    default settings which changes:

    - The default separator is a space ` ` instead of a comma `,`.
    - The headers of the column are not included.

    By modifying the args of `dm.read_csv()`, you will be able to read an `.smi` files.

    Args:
        urlpath: Path to a file or a file-like object. Path can be remote or local.
    """

    active_path = urlpath

    # NOTE(hadim): the temporary local file copy
    # is because `SmilesMolSupplier` does not support
    # using file-like object, only path.

    # Copy to a local temporary path if the path is a remote one.
    if not fsspec.utils.can_be_local(str(urlpath)):
        active_path = pathlib.Path(tempfile.mkstemp()[1])
        dm.utils.fs.copy_file(urlpath, active_path, force=True)

    # Read the molecules
    supplier = rdmolfiles.SmilesMolSupplier(str(active_path), titleLine=0)
    mols = [mol for mol in supplier if mol is not None]

    # Delete the local temporary path
    if not fsspec.utils.can_be_local(str(urlpath)):
        pathlib.Path(str(active_path)).unlink()

    return mols

def to_xlsx(
    mols: Union[Mol, Sequence[Mol], pd.DataFrame],
    urlpath: Union[str, os.PathLike],
    smiles_column: Optional[str] = "smiles",
    mol_column: str = "mol",
    mol_size: List[int] = [300, 300],
):
    """Write molecules to an Excel file with a molecule column as an RDKit rendered
    image.

    Args:
        mols: a dataframe, a molecule or a list of molecule.
        urlpath: Path to a file or a file-like object. Path can be remote or local.
        smiles_column: Column name to extract the molecule.
        mol_column: Column name to extract the molecule. It takes
            precedence over `smiles_column`.
            Column name to write the RDKit rendered image. If none,
            the molecule images are not written.
    """

    if isinstance(mols, Mol):
        mols = [mols]

    if isinstance(mols, Sequence):
        mols = [mol for mol in mols if mol is not None]
        mols = dm.to_df(mols, smiles_column=smiles_column, mol_column=mol_column)

    if mols is None or mols.empty:  # type: ignore
        raise ValueError("No molecules to write")

    with fsspec.open(urlpath, mode="wb") as f:
        PandasTools.SaveXlsxFromFrame(mols, f, molCol=mol_column, size=mol_size)

EXTENSIONS_DICT = {
    "csv": [
        ".csv",
        ".csv.gz",
        ".csv.bz2",
        ".csv.zip",
        ".csv.xz",
        ".csv.zst",
        ".csv.tar",
        ".csv.tar.gz",
        ".csv.tar.xz",
        ".csv.tar.bz2",
    ],
    "excel": [".xlsx"],
    "parquet": [".parquet"],
    "json": [
        ".json",
        ".json.gz",
        ".json.bz2",
        ".json.zip",
        ".json.xz",
        ".json.zst",
        ".json.tar",
        ".json.tar.gz",
        ".json.tar.xz",
        ".json.tar.bz2",
    ],
    "sdf": [
        ".sdf",
        ".sdf.gz",
    ],
}

def _guess_filetype(path: str):
    """Return a filetype given an input path. Filetypes returned can be from
    `csv, excel, parquet, json, sdf`.
    """
    for name, extensions in EXTENSIONS_DICT.items():
        for ext in extensions:
            if path.endswith(ext):
                return name

def open_df(path: str, **kwargs: Any) -> pd.DataFrame:
    """Open a dataframe file whatever its filetype from
    `csv, excel, parquet, json, sdf`.

    Args:
        path: path to the file.
        **kwargs: keyword arguments to pass to the underlying reader.
    """

    filetype = _guess_filetype(path)

    data = None
    if filetype == "csv":
        data = pd.read_csv(path, **kwargs)
    elif filetype == "excel":
        data = pd.read_excel(path, **kwargs)
    elif filetype == "parquet":
        data = pd.read_parquet(path, **kwargs)
    elif filetype == "json":
        data = pd.read_json(path, **kwargs)
    elif filetype == "sdf":
        kwargs.setdefault("as_df", True)
        data = dm.read_sdf(path, **kwargs)
    else:
        raise ValueError(f"The file type of {path} is not supported.")

    data = cast(pd.DataFrame, data)

    return data

def save_df(
    data: pd.DataFrame,
    path: str,
    **kwargs: Any,
):
    """Save a dataframe file whatever its filetype from
    `csv, excel, parquet, json, sdf`.

    Args:
        data: dataframe to save.
        path: path to save the file.
        **kwargs: additional arguments to pass that are specific to the file save type.
    """

    filetype = _guess_filetype(path)

    if filetype == "csv":
        kwargs.setdefault("index", False)
        data.to_csv(path, **kwargs)
    elif filetype == "excel":
        kwargs.setdefault("index", False)
        data.to_excel(path, **kwargs)
    elif filetype == "parquet":
        data.to_parquet(path, **kwargs)
    elif filetype == "json":
        data.to_json(path, **kwargs)
    elif filetype == "sdf":
        dm.to_sdf(data, path, **kwargs)
    else:
        raise ValueError(f"The file type of {path} is not supported.")


from __future__ import unicode_literals
import json
import logging
from flask import Flask, request
import requests
import pandas as pd

import os
os.environ["GOOGLE_APPLICATION_CREDENTIALS"]="marusya-made-12e4432bff10.json"

app = Flask(__name__)
app.debug = True

def detect_intent_texts(project_id, session_id, texts, language_code):
    """Returns the result of detect intent with texts as inputs.

    Using the same `session_id` between requests allows continuation
    of the conversation."""
    import dialogflow_v2 as dialogflow
    session_client = dialogflow.SessionsClient()

    session = session_client.session_path(project_id, session_id)
    print('Session path: {}\n'.format(session))

    for text in texts:
        text_input = dialogflow.types.TextInput(
            text=text, language_code=language_code)

        query_input = dialogflow.types.QueryInput(text=text_input)

        response = session_client.detect_intent(
            session=session, query_input=query_input)

        print('=' * 20)
        print('Query text: {}'.format(response.query_result.query_text))
        print('Detected intent: {} (confidence: {})\n'.format(
            response.query_result.intent.display_name,
            response.query_result.intent_detection_confidence))
        print('Fulfillment text: {}\n'.format(
            response.query_result.fulfillment_text))
        return {'Query text': response.query_result.query_text,
                'Detected intent': response.query_result.intent.display_name,
                'Fulfillment text': response.query_result.fulfillment_text}, response

def get_intent(text):
    answer, r = detect_intent_texts('marusya-made', 7777777, [text], 'ru')
    return answer['Detected intent'], r

@app.route('/')
def index():
    return '
На главной ничего нет

вебхук для запросов: https://what-about-the-virus.herokuapp.com/mrs
'

@app.route('/mrs', methods=['POST'])
def main():

    answ, r = get_intent(request.json['request']['command'])

    card = {}
    buttons = []
    end_session = False
    text = ''

    if request.json['session']['new']:
        text = 'Привет! Меня зовут Алиса. Я могу дать текущую сводку по \
                коронавирусу в разрезе стран и мира. \
                Вы можете ознакомиться с примерами команд или начать спрашивать)'
        # Потом раскоментить картинку!!!!!!!!
        # card = {
        #         'type': 'BigImage',
        #         'image_id': 457239017,
        #         'title': '',
        #         'description': ''
        #         }

        buttons = [
                    {'title': 'Примеры'},
                    {'title': 'Не надо примеров'}
                  ]

    elif request.json['request']['command'] == 'примеры':
        text = '1. Какая ситуация в мире? Что в мире?\n \
                2. Сводка по России. Ситуация в нашей стране\n \
                3. Статистика в Германии. Ситуация в Монголии.\n \
                4. Какие симптомы? Симптомы\n \
                5. Где найти советы по вирусу? Советы.'

    elif request.json['request']['command'] == 'не надо примеров':
        text = 'Хорошо, тогда слушаю вас)) \n (＾▽＾)'

    elif answ == 'Default Welcome Intent':
        text = r.query_result.fulfillment_text

    elif answ == 'world request':
        try:
            res = requests.get('https://covid-api.com/api/reports/total')
            data = res.json()['data']

            text = 'СВОДКА ПО МИРУ \n\
                    Актуально на {d}\n\
                    Заражения: {z}\n\
                    Смерти: {de}\n\
                    Смертность: {s}%\n\
                    Выздоровления: {v}'.format(d=data['date'],
                                               z='{0:,}'.format(data['confirmed']).replace(',', ' '),
                                               de='{0:,}'.format(data['deaths']).replace(',', ' '),
                                               s=round(data['fatality_rate']*100,2),
                                               v='{0:,}'.format(data['recovered']).replace(',', ' '))
        except:
            text = 'Кажется, не только вам сейчас интересна статистика по миру. \
                    Не удается загрузить данные.\n\
                    Простите, может сработают другие команды.'

    elif answ == 'russia':
        try:
            res = requests.get('https://covid19-api.org/api/status/ru')
            data = res.json()

            text = 'СВОДКА ПО РОССИИ \n\
                    Актуально на {d}\n\
                    Заражения: {z}\n\
                    Смерти: {de}\n\
                    Смертность: {s}%\n\
                    Выздоровления: {v}'.format(d=data['last_update'][0:10],
                                               z='{0:,}'.format(data['cases']).replace(',', ' '),
                                               de='{0:,}'.format(data['deaths']).replace(',', ' '),
                                               s=round(data['deaths']/data['cases']*100,2),
                                               v='{0:,}'.format(data['recovered']).replace(',', ' '))
        except:
            text = 'Кажется, не только вам сейчас интересна статистика в России. \
                    Не удается загрузить данные.\n\
                    Простите, может сработают другие команды.'

    elif answ == 'countries':
        df = pd.read_csv('countries_codes.csv')
        country = r.query_result.parameters.values()[0]
        code = df.loc[df['Страна']==country]['Код'].values[0]
        try:
            res = requests.get('https://covid19-api.org/api/status/'+code)
            data = res.json()

            text = 'СВОДКА - {c}\n\
            Актуально на {d}\n\
            Заражения: {z}\n\
            Смерти: {de}\n\
            Смертность: {s}%\n\
            Выздоровления: {v}'.format(c=country.upper(),
                                       d=data['last_update'][0:10],
                                       z='{0:,}'.format(data['cases']).replace(',', ' '),
                                       de='{0:,}'.format(data['deaths']).replace(',', ' '),
                                       s=round(data['deaths']/data['cases']*100,2),
                                       v='{0:,}'.format(data['recovered']).replace(',', ' '))
        except:
            text = 'Кажется, не только вам сейчас интересна статистика по странам. \
                    Не удается загрузить данные.\n\
                    Простите, может сработают другие команды.'

    elif answ == 'symptoms':
        text = 'К наиболее распространенным симптомам COVID-19 относятся:\
                \n\nлихорадка;\nсухой кашель;\nутомляемость.\n\n\
                К другим, менее распространенным симптомам, которые встречаются у ряда пациентов, относятся:\
                \n\nутрата обоняния или вкусовых ощущений;\nзаложенность носа;\n\
                конъюнктивит (или покраснение глаз);\nболь в горле;\nголовная боль;\
                \nболь в мышцах или суставах;\nразличные виды высыпаний на коже;\
                \nтошнота или рвота;\nдиарея;\nозноб или головокружение.'

    elif answ == 'advices':
        text = 'Исчерпывающую информацию можно получить на сайте ВОЗ\
                \nhttps://www.who.int/ru/emergencies/diseases/novel-coronavirus-2019/advice-for-public\n\
                или минздрава\
                \nhttps://covid19.rosminzdrav.ru/#prevention'

    elif request.json['request']['command'] == 'on_interrupt':
        text = 'Всего доброго! Берегите себя и близких!'

    else:
        text = r.query_result.fulfillment_text

    response = {
    'version': request.json['version'],
    'session': request.json['session'],
    'response': {
        'end_session': end_session,
        'text': text,
        'card': card,
        'buttons': buttons
        }
    }

    logging.info('Response: %r', response)

    return json.dumps(response, ensure_ascii=False, indent=2)

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

# @app.route('/mrs', methods=['POST'])
# def main():
#
#     card = {}
#     buttons = []
#     end_session = False
#     text = ''
#
#     if request.json['session']['new']:
#         text = 'Привет! Это еще тестовый навык.'
#     elif request.json['request']['command'].lower() == 'привет':
#         text = 'Привет!!!'
#     elif request.json['request']['command'].lower() == 'хватит':
#         text = 'Хорошо, до встречи!'
#         end_session = True
#     elif request.json['request']['command'] == 'on_interrupt':
#         text = 'Приходи еще.'
#     elif request.json['request']['command'].lower() == 'картинка':
#         text = 'Вот такая есть картинка:'
#         card = {
#             'type': 'BigImage',
#             'image_id': 457239018,
#             'title': 'Это Коржик — офисный пёс',
#             'description': 'текст'
#                 }
#     elif request.json['request']['command'].lower() == 'кнопки':
#         text = 'Вот такие есть кнопки:'
#         buttons = [
#                     {'title': 'кнопка 1'},
#                     {'title': 'кнопка 2'}
#                   ]
#     elif request.json['request']['command'].lower() == 'карусель':
#         text = 'Карусель:'
#         card = {
#             "type": "ItemsList",
#             # "title": "Две картинки",
#             # "description": "текст",
#             "items": [{"image_id": 457239018}, {"image_id": 457239017}, {"image_id": 457239019}]
#                 }
#     else:
#         text = request.json['request']['command'].lower()
#
#     response = {
#     'version': request.json['version'],
#     'session': request.json['session'],
#     'response': {
#         'end_session': end_session,
#         'text': text,
#         'card': card,
#         'buttons': buttons
#         }
#     }
#
#     logging.info('Response: %r', response)
#
#     return json.dumps(response, ensure_ascii=False, indent=2)


import socketserver
import json

class TcpServer(socketserver.BaseRequestHandler):
    def put(self,*args):
        """接收客户端文件"""
        cmd_data = args[0]
        filename = cmd_data['filename']
        filesize = cmd_data['size']

        # 保存文件
        print(cmd_data)
        print('save file')
        print(filename)
        print(filesize)
        f = open(filename, 'wb')
        self.request.send(b'200 ok')

        received_size = 0
        while True:
            if received_size < filesize:
                data = self.request.recv(1024)
                f.write(data)
                received_size += len(data)
            else:
                print('file [%s] has uploaded...' % filename)

                f.close()
                break

        pass

    def handle(self):
        while True:
            self.data = self.request.recv(1024).strip()
            print('{} wrote.'.format(self.client_address))
            print(self.data)
            if not self.data:
                break

            cmd_dic = json.loads(self.data.decode())
            action = cmd_dic['action']
            if hasattr(self,action):
                func = getattr(self,action)
                func(cmd_dic)
        pass

if __name__ == '__main__':
    host = ('localhost',9999)
    server = socketserver.ThreadingTCPServer(host,TcpServer)
    server.serve_forever()


#!/usr/bin/env python3
import os
import sys
import pathlib
import patoolib
import argparse

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[0m'
    UNDERLINE = '\033[4m'

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-d', '--directory', action='store',
            dest='root_dir', help='The directory to parse \
                    and zip all subdirectories', default='',
                    type=str)
    parser.add_argument('-o', '--output-directory', action='store',
            dest='output_dir', help='Where should the zipped files be stored')
    parser.add_argument('-z', '--zip-format', action='store', default='zip',
            dest='zip_format', help='The format to use for archiving')
    parser.add_argument('-v', '--verbosity', action='count', \
            help='Make the program talk more', dest='verbose', default=0)
    parser.add_argument('-c', '--checksum', action='store_true',
            dest='checksum', help='Should a checksum file be generated')
    parser.add_argument('-t', '--test', action='store_true', \
            default='store_false', dest='test')
    parser.set_defaults(root_dir = os.path.dirname(os.path.realpath(__file__)),
            output_dir = os.path.dirname(os.path.realpath(__file__)),
                verbosity=0, test=False)
    return parser.parse_args()

if __name__ == '__main__':
    args = parse_args()
    if args.verbose > 0:
        if args.test:
            print( bcolors.OKBLUE + ' ===> ' + 'running as ' + \
                    bcolors.WARNING + 'test' + bcolors.ENDC)
    if args.verbose > 2:
        print( bcolors.OKBLUE + ' ===> ' + bcolors.OKGREEN + \
                ' retrieving direcotries within: ' + \
                bcolors.ENDC  + ''.join(args.root_dir))
    dirs = os.listdir(args.root_dir)
    if args.verbose > 2:
        print(bcolors.OKBLUE + ' ===> ' + bcolors.OKGREEN + \
                'Directories found:\n\t' + bcolors.ENDC + \
                '\n\t'.join(dirs))
    for d in dirs:
        if args.verbose > 0:
            print(bcolors.OKBLUE + ' ====> ' + bcolors.ENDC + \
                    'creating file %s' % d)
        archive_out = os.path.join(args.output_dir, d) + '.' + args.zip_format
        if args.verbose > 1:
            print(bcolors.OKBLUE + ' ====>  ' + bcolors.ENDC + \
                    '  Creating archive %s' % archive_out)
        outpath = os.path.join(args.root_dir, d)
        if not args.test:
            if args.verbose > 0:
                print(bcolors.OKGREEN + ' ====> ' + bcolors.ENDC +  archive_out)
                print(bcolors.OKGREEN + ' ====> ' + bcolors.ENDC + outpath)
            patoolib.create_archive("%s" % archive_out, \
                    (outpath + '/'))


import sc2
from sc2.player import Bot, Computer, BotAI
from sc2 import Result
import random
import cv2
import numpy as np
from ComputerBot import ComputerBot

class Game:
    max_char_per_name = 15
    max_char_per_rank = 2
    max_char_per_line = max_char_per_name + max_char_per_rank + 3
    disable_char = 'X'
    cell_size = (75, 150)

    def __init__(self, side, player1_dict=None, player2_dict=None, realtime=True):
        self.realtime = realtime
        if player1_dict is not None:
            self.player1 = player1_dict["player"]
            self.rank1 = player1_dict["rank"]
            self.enable1 = player1_dict["enable"]

        else:
            self.player1 = None
            self.rank1 = None
            self.enable1 = None

        self.player1_color = (255, 255, 255)

        if player2_dict is not None:
            self.player2 = player2_dict["player"]
            self.rank2 = player2_dict["rank"]
            self.enable2 = player2_dict["enable"]

        else:
            self.player2 = None
            self.rank2 = None
            self.enable2 = None

        self.player2_color = (255, 255, 255)

        self.side = side
        assert self.side == "left" or self.side == "right"

    def run(self, debug=False):
        assert self.player1 is not None or self.player2 is not None

        if self.player1 is None:
            return [False, True]
        elif self.player2 is None:
            return [True, False]

        if debug:
            result = [False, False]
            rn_idx = int(random.random() > 0.5)
            result[rn_idx] = True
            # self.rank1 = self.rank1 if result[0] else self.disable_char
            # self.rank2 = self.rank2 if result[1] else self.disable_char
            # print(result, (self.rank1, self.rank2))
            self.player1_color = (0, 255, 0) if result[0] else (0, 0, 255)
            self.player2_color = (0, 255, 0) if result[1] else (0, 0, 255)
            return result

        game_result = self.run_game()

        print("game_result", type(game_result), game_result)
        print("game_result[0]", type(game_result[0]), game_result[0])
        result = [Result.Victory == game_result[0], Result.Victory == game_result[1]]
        if not result[0] and not result[1]:
            rn_idx = int(random.random() > 0.5)
            result[rn_idx] = True
        # print("run return: ", result)
        # self.rank1 = self.rank1 if result[0] else self.disable_char
        # self.rank2 = self.rank2 if result[1] else self.disable_char

        self.player1_color = (0, 255, 0) if result[0] else (0, 0, 255)
        self.player2_color = (0, 255, 0) if result[1] else (0, 0, 255)
        return result

    def run_game(self):

        try:
            player1_is_bot = issubclass(self.player1, BotAI)
        except TypeError:
            player1_is_bot = False

        try:
            player1_is_computer = issubclass(self.player1, ComputerBot)
        except TypeError:
            player1_is_computer = isinstance(self.player1, ComputerBot)

        try:
            player2_is_bot = issubclass(self.player2, BotAI)
        except TypeError:
            player2_is_bot = False

        try:
            player2_is_computer = issubclass(self.player2, ComputerBot)
        except TypeError:
            player2_is_computer = isinstance(self.player2, ComputerBot)

        # print(player1_is_bot, player1_is_computer, player2_is_bot, player2_is_computer)
        assert (player1_is_bot or player1_is_computer) and (player2_is_bot or player2_is_computer)

        if player1_is_computer and player2_is_bot:
            raw_game_result = sc2.run_game(sc2.maps.get("AbyssalReefLE"), [
                Bot(self.player2.BOTRACE, self.player2(), name=self.player2.BOTNAME),
                Computer(race=self.player1.BOTRACE, difficulty=self.player1.DIFFICULTY)
            ], realtime=self.realtime)
            game_result = [Result.Defeat if raw_game_result == Result.Victory else Result.Victory, raw_game_result]
        elif player1_is_bot and player2_is_computer:
            raw_game_result = sc2.run_game(sc2.maps.get("AbyssalReefLE"), [
                Bot(self.player1.BOTRACE, self.player1(), name=self.player1.BOTNAME),
                Computer(self.player2.BOTRACE, self.player2.DIFFICULTY)
            ], realtime=self.realtime)
            game_result = [raw_game_result, Result.Defeat if raw_game_result == Result.Victory else Result.Victory]
        elif player1_is_computer and player2_is_computer:
            game_result = sc2.run_game(sc2.maps.get("AbyssalReefLE"), [
                Computer(self.player1.BOTRACE, self.player1.DIFFICULTY),
                Computer(self.player2.BOTRACE, self.player2.DIFFICULTY)
            ], realtime=self.realtime)
        elif player1_is_bot and player2_is_bot:
            game_result = sc2.run_game(sc2.maps.get("AbyssalReefLE"), [
                Bot(self.player1.BOTRACE, self.player1(), name=self.player1.BOTNAME),
                Bot(self.player2.BOTRACE, self.player2(), name=self.player2.BOTNAME)
            ], realtime=self.realtime)
        else:
            raise ValueError("players must be ComputerBot or BatAI instance")

        return game_result

    def empty_str(self):
        out = '-' * self.max_char_per_line + '\n'
        if self.side == "left":
            out += '|' + " " * self.max_char_per_rank + '|' + " " * self.max_char_per_name + '|' + '\n'
            out += '|' + " " * self.max_char_per_rank + '|' + " " * self.max_char_per_name + '|' + '\n'
        else:
            out += '|' + " " * self.max_char_per_name + '|' + " " * self.max_char_per_rank + '|' '\n'
            out += '|' + " " * self.max_char_per_name + '|' + " " * self.max_char_per_rank + '|' '\n'

        out += '-' * self.max_char_per_line + '\n'
        return out

    def __str__(self):
        if self.player1 is None and self.player2 is None:
            return self.empty_str()

        out = '-' * self.max_char_per_line + '\n'
        out += self.get_line1(for_str=True)
        out += self.get_line2(for_str=True)
        out += '-' * self.max_char_per_line + '\n'
        return out

    def get_line1(self, for_str=False):
        out = ""

        rank1 = self.rank1

        if for_str and not self.enable1:
            rank1 = self.disable_char

        if self.player1 is None:
            rank1_str = ' ' * self.max_char_per_rank
        elif len(str(rank1)) > self.max_char_per_rank:
            rank1_str = '.' * max(3, self.max_char_per_rank)
        else:
            rank1_str = str(rank1) + " " * (self.max_char_per_rank - len(str(rank1)))

        if self.player1 is None:
            player1_str = ' ' * self.max_char_per_name
        elif len(self.player1.BOTNAME) >= self.max_char_per_name:
            player1_str = ''.join([self.player1.BOTNAME[i] for i in range(self.max_char_per_name)])
        else:
            player1_str = self.player1.BOTNAME + " " * (self.max_char_per_name - len(self.player1.BOTNAME))

        if self.side == "left":
            out += '|' + rank1_str + '|' + player1_str + '|' + '\n'
        else:
            out += '|' + player1_str + '|' + rank1_str + '|' '\n'
        return out

    def get_line2(self, for_str=False):
        out = ""

        rank2 = self.rank2

        if for_str and not self.enable2:
            rank2 = self.disable_char

        if self.player2 is None:
            rank2_str = ' ' * self.max_char_per_rank
        elif len(str(rank2)) > self.max_char_per_rank:
            rank2_str = '.' * max(3, self.max_char_per_rank)
        else:
            rank2_str = str(rank2) + " " * (self.max_char_per_rank - len(str(rank2)))

        if self.player2 is None:
            player2_str = ' ' * self.max_char_per_name
        elif len(self.player2.BOTNAME) >= self.max_char_per_name:
            player2_str = ''.join([self.player2.BOTNAME[i] for i in range(self.max_char_per_name)])
        else:
            player2_str = self.player2.BOTNAME + " " * (self.max_char_per_name - len(self.player2.BOTNAME))

        if self.side == "left":
            out += '|' + rank2_str + '|' + player2_str + '|' + '\n'
        else:
            out += '|' + player2_str + '|' + rank2_str + '|' '\n'
        return out

    def get_img_cell(self):
        img = np.zeros((self.cell_size[0], self.cell_size[1], 3))
        y0, dy = 15, self.cell_size[0]//2 - 2
        font_size = (self.cell_size[0] * 0.30) / 50
        lines = [self.get_line1(), self.get_line2()]
        for i, line in enumerate(lines):
            y = y0 + i*dy
            line = line.replace('\n', '')
            line = line.replace('|', ' ')
            p_color = self.player1_color if i == 0 else self.player2_color
            cv2.putText(img, line, (0, y), cv2.FONT_HERSHEY_SIMPLEX, font_size, p_color, thickness=1)
        cv2.rectangle(img, (0, 0), (self.cell_size[1] - 1, self.cell_size[0] - 1), color=(255, 0, 0), thickness=1)
        return img

if __name__ == '__main__':
    from MarineOverflow import MarineOverflow
    from JarexTerran import JarexTerran
    from JarexProtoss import JarexProtoss

    class Dummy1(MarineOverflow):
        BOTNAME = "Dummy1"

    lgames = list()
    rgames = list()
    game = Game("left", {"player": MarineOverflow, "rank": 1, "enable": True},
                {"player": JarexProtoss, "rank": 12, "enable": True}, realtime=False)
    print(game)

    while True:
        img = game.get_img_cell()
        # img = cv2.resize(img, (500, 250))
        cv2.imshow('test cell', img)
        if cv2.waitKey(25) & 0xFF == ord('q'):
            cv2.destroyAllWindows()
            break
    print(game.run())

    print(game)

    while True:
        img = game.get_img_cell()
        # img = cv2.resize(img, (500, 250))
        cv2.imshow('test cell', img)
        if cv2.waitKey(25) & 0xFF == ord('q'):
            cv2.destroyAllWindows()
            break

    rgame = Game("right", {"player": JarexTerran, "rank": 1, "enable": True},
                 {"player": Dummy1, "rank": 'X', "enable": False})
    rgame.player1_color = (0, 255, 0)
    print(rgame)

    while True:
        img = rgame.get_img_cell()
        # img = cv2.resize(img, (500, 250))
        cv2.imshow('test cell', img)
        if cv2.waitKey(25) & 0xFF == ord('q'):
            cv2.destroyAllWindows()
            break

    empty_game = Game("left")
    print(empty_game)

    # for i in range(1, 6):
    #     lgames.append(Game({"player": MarineOverflow, "rank": i}, {"player": MarineOverflow, "rank": i+5}, "left"))
    # for i in range(6, 11):
    #     rgames.append(Game({"player": MarineOverflow, "rank": i}, {"player": MarineOverflow, "rank": i+5}, "right"))
    #
    # for i in range(len(rgames)):
    #     print(str(lgames[i]) + ' '*10 + str(rgames[i]) + '\n')