import argparse
import os
import time
import numpy as np
import tensorflow as tf
import transformer
import utils
import vae
import vaetransformer

if __name__ == '__main__':

    parser = argparse.ArgumentParser()

    parser.add_argument('--input_max_len', type=int, default=64)
    parser.add_argument('--target_max_len', type=int, default=64)
    parser.add_argument('--epochs', type=int, default=25)

    parser.add_argument('--kl_start', type=float, default=0.001)
    parser.add_argument('--kl_end', type=float, default=0.01)
    parser.add_argument('--kl_interval', type=int, default=5000)
    parser.add_argument('--kl_delta', type=float, default=0.0001)
    parser.add_argument('--kl_warmup', type=int, default=40000)

    parser.add_argument('--warmup', type=int, default=10000)

    parser.add_argument('--npy_num', type=int, default=10)
    parser.add_argument('--batch_size', type=int, default=64)
    parser.add_argument('--input_npy', type=str)
    parser.add_argument('--target_npy', type=str)

    parser.add_argument('--ckpt_path', type=str, default='ckpt')
    parser.add_argument('--step_count', type=int, default=0)
    parser.add_argument('--batch_log', type=str, default='./batch_log.csv')
    parser.add_argument('--epoch_log', type=str, default='./epoch_log.csv')
    args = parser.parse_args()
    print(args)

    learning_rate = utils.CustomLearningRateSchedule(d_model=256, warmup_steps=args.warmup)
    optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98,
                                         epsilon=1e-9)

    kl_loss = tf.keras.metrics.Mean(name='kl_loss')
    vae_ce_loss = tf.keras.metrics.Mean(name='vae_ce_loss')
    vae_accuracy = tf.keras.metrics.Mean(name='vae_acc')
    transformer_ce_loss = tf.keras.metrics.Mean(name='transformer_ce_loss')

    train_loss = tf.keras.metrics.Mean(name='train_loss')
    train_accuracy = tf.keras.metrics.Mean(name='train_accuracy')

    model = vaetransformer.VAETransformer(
        input_max_len=args.input_max_len, input_vocab_size=28,
        target_max_len=args.target_max_len, target_vocab_size=28,
        d_model=256, num_vae_layers=3, vae_kernel_size=3, latent_dim=64, pooling='max',
        num_transformer_layers=3, num_heads=4, dff=1024, dropout_rate=0.1)

    checkpoint_path = os.path.join('.', 'checkpoints', args.ckpt_path)
    os.makedirs(checkpoint_path, exist_ok=True)

    ckpt = tf.train.Checkpoint(model=model,
                               optimizer=optimizer)

    ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)

    # if a checkpoint exists, restore the latest checkpoint.
    if ckpt_manager.latest_checkpoint:
        print(ckpt_manager.latest_checkpoint)
        ckpt.restore(ckpt_manager.latest_checkpoint)
        print ('Latest checkpoint restored!!', ckpt_manager.latest_checkpoint)

    train_step_signature = [
        tf.TensorSpec(shape=(None, args.input_max_len), dtype=tf.int64),
        tf.TensorSpec(shape=(None, args.target_max_len), dtype=tf.int64),
        tf.TensorSpec(shape=(1, ), dtype=tf.float32)
    ]

    @tf.function(input_signature=train_step_signature)
    def train_step(inp, tar, gamma_kl):
        """Training function for one step.
        Args:
            inp: Input chain SMILES. shape = `[batch_size, input_length]`.
            tar: Output full SMILES. shape = `[batch_size, target_length]`.
            gamma_kl: Weight for KL-divergence loss.
        Returns:
            None.
        """
        tar_inp = tar[:, :-1]
        tar_real = tar[:, 1:]

        enc_padding_mask, combined_mask, dec_padding_mask = utils.create_masks(inp, tar_inp)

        with tf.GradientTape() as tape:
            z_mean, z_log_var, z, mid_output, final_output, attention_weights_dict = model(inp,
                                                                                tar_inp,
                                                                                enc_padding_mask,
                                                                                combined_mask,
                                                                                dec_padding_mask,
                                                                                training=True)

            kl_div = utils.kl_loss_function(z_mean, z_log_var)
            vae_loss = utils.crossentropy_loss_function_prob(inp, mid_output)
            transformer_loss = utils.crossentropy_loss_function_logits(tar_real, final_output)
            loss = gamma_kl * kl_div + vae_loss + transformer_loss

        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

        kl_loss(kl_div)
        vae_ce_loss(vae_loss)
        vae_accuracy(utils.accuracy_function(inp, mid_output))
        transformer_ce_loss(transformer_loss)
        train_loss(loss)
        train_accuracy(utils.accuracy_function(tar_real, final_output))

    gamma_kl_schedule = utils.CustomLossSchedule(
        args.kl_start,
        args.kl_end,
        args.kl_interval,
        args.kl_delta,
        args.kl_warmup)

    step_count = args.step_count

    train_batches = utils.CustomDataset(args.npy_num, args.batch_size, args.input_npy, args.target_npy)

    for epoch in range(args.epochs):
        begin = time.time()

        kl_loss.reset_states()
        vae_ce_loss.reset_states()
        vae_accuracy.reset_states()
        transformer_ce_loss.reset_states()
        train_loss.reset_states()
        train_accuracy.reset_states()

        for (batch_i, (inp, tar)) in enumerate(train_batches):
            gamma_kl = gamma_kl_schedule(step_count)
            train_step(inp, tar, gamma_kl)
            step_count += 1
            if (batch_i + 1) % 100 == 0:
                print(time.asctime( time.localtime(time.time()) ))
                print(f'[Epoch: {epoch + 1}] [Batch: {batch_i + 1}] [Step: {step_count}]\n'\
                      f'[KL-gamma: {gamma_kl.numpy()[0]:.4f}] [KL: {kl_loss.result():.4f}]\n'\
                      f'[VAE Accuracy: {vae_accuracy.result():.4f}]\n'\
                      f'[VAE Loss: {vae_ce_loss.result():.4f}] [Transformer Loss: {transformer_ce_loss.result():.4f}]\n'\
                      f'[Total Loss: {train_loss.result():.4f}] [Accuracy: {train_accuracy.result():.4f}]')
                print('-' * 32)
                with open(args.batch_log, 'a') as fbatchlog:
                    print(f'{epoch + 1},{step_count},'\
                          f'{gamma_kl.numpy()[0]:.4f},{kl_loss.result():.4f},'\
                          f'{vae_accuracy.result():.4f},'\
                          f'{vae_ce_loss.result():.4f},{transformer_ce_loss.result():.4f},'\
                          f'{train_loss.result():.4f},{train_accuracy.result():.4f}', file=fbatchlog)
        ckpt_save_path = ckpt_manager.save()

        print(f'[Epoch: {epoch + 1}] [Step: {step_count}]\n'\
              f'[KL-gamma: {gamma_kl.numpy()[0]:.4f}] [KL: {kl_loss.result():.4f}]\n'\
              f'[VAE Accuracy: {vae_accuracy.result():.4f}]\n'\
              f'[VAE Loss: {vae_ce_loss.result():.4f}] [Transformer Loss: {transformer_ce_loss.result():.4f}]\n'\
              f'[Total Loss: {train_loss.result():.4f}] [Accuracy: {train_accuracy.result():.4f}]')
        with open(args.epoch_log, 'a') as f_log:
            print(f'{epoch + 1},{step_count},'\
                  f'{gamma_kl.numpy()[0]:.4f},{kl_loss.result():.4f},'\
                  f'{vae_accuracy.result():.4f},'\
                  f'{vae_ce_loss.result():.4f},{transformer_ce_loss.result():.4f},'\
                  f'{train_loss.result():.4f},{train_accuracy.result():.4f}', file=f_log)
        print(f'Time taken for one epoch: {time.time() - begin:.2f} secs')
        print(f'Savedcheckpoint for epoch {epoch + 1} at {ckpt_save_path}\n')


# Library that gives us the functions to create
#   the map
import plotly.io as pio

# Reads csv files and sorts the states by their
#   fips codes
import pandas as pd

# Object that will be used to execute the
#   map creation features
import plotly.express as px

# Reads the map data from a website to help
#   generate the map
from urllib.request import urlopen
import json
with urlopen('https://raw.githubusercontent.com/plotly/datasets/master/geojson-counties-fips.json') as response:
    counties = json.load(response)

# To find the proper path to find the html file we want to display
import os

# Used to import and run the C++ program to generate CSV files
import subprocess

# Calculate the elapsed time for each check for a map or csv file
import timeit

class map_test() :
  # Constructor
  def __init__( self, interval, sDate, fDate, rates ) :
    # Stores the number of days the user chooses
    self.interval = interval

    # Stores the start date
    self.sDate = sDate

    # Stores the final date
    self.fDate = fDate

    # Stores the rates of change
    self.rates = rates

    # Make mapHash() a global variable
    self.mapHash = {}

  # Runs entire checks for generating map
  def main( self ) :
    # Reads in user input from the website
    options = [ self.interval, self.sDate,
                self.fDate, self.rates ]

    if int(options[3]) == int(8) :
      options[3] = 'ALL'

    # Stores the name of the temporary key
    tempKey = ''

    # Creates a key using the user input
    for i in range( len(options) ) :
      tempKey = tempKey + str( options[i] )

      # Separates each element in the key by a hyphen
      if i != len(options) - 1 :
        tempKey = tempKey + '-'

    # Reformats the key to be easier to read with the file names
    key = self.createKey( tempKey )

    # Each element is separated into
    keys = key.split( '-' )

    # Concatenates the strings that will be used for filename
    fName = ( keys[0] + '-CommID-Cases-' +
              keys[1] + '-' + keys[2] + '-' +
              keys[3] + 'vs' + keys[4] + '-' +
              keys[5] + '-' + keys[6] +
              '-intDays' + keys[7] )

    # Creates a subdirectory for the data set if it doesn't exist
    subName = ( 'Cases-'+ keys[1] + '-' +
                keys[2] + '-' + keys[3] +
                'vs' + keys[4] + '-' +
                keys[5] + '-' + keys[6] +
                '-intDays' + keys[7] )

    # Calculate to generate map hash table
    start_time = timeit.default_timer()

    # Initialize the hash tables for
    #   the maps files
    self.mapHash = self.initMapHash()

    # End timer
    elapsed = timeit.default_timer() - start_time

    # Calculate to generate map hash table
    start_time = timeit.default_timer()

    # Initialize the hash tables for
    #   the CSV files
    csvHash = self.initCSVHash()

    # End timer
    elapsed = timeit.default_timer() - start_time

    # Check if there is a map for this dataset
    if self.mapHash.get( key ) != None :
      #print( "map key: %s does exist" % key )
      fName = './v2/static/maps/' + subName + '/' + self.mapHash[key]
    else :
      if csvHash.get( key ) == None :
        subprocess.call(["g++", "-o", "hello", "./v2/maptest/hello.cpp"])
        subprocess.call(["./hello"])

        #print( "csv key: %s doesn't exist" % key )

      # Creates a file path for where to store the cases CSV files
      csvPath = os.getcwd() + '/v2/maptest/cases/'+ subName

      # Checks if a directory for the CSV data set exists
      # If not, then we create one
      if not os.path.isdir( csvPath ) :
        try :
          os.mkdir( csvPath )
        except OSError :
          print ("Command mkdir for directory %s has failed." % csvPath)
      else :
        print( 'this csv directory exists' )

      #print( "map key: %s doesn't exist" % key )

      # Creates a file path for where to store the map html files
      mapPath = os.getcwd() + '/v2/static/maps/'+ subName

      # Checks if a directory for the map data set exists
      # If not, then we create one
      if not os.path.isdir( mapPath ) :
        try :
          os.mkdir( mapPath )
        except OSError :
          print ("Command mkdir for directory %s has failed." % mapPath)

      fig = self.generateMap( subName, fName )

      # Add key and value to the HTML hash table
      with open( './v2/maptest/tables/html_table.txt', 'a' ) as fp :
        fp.write( key + ',' + fName + '.html\r\n' )

      fp.close()

      # Update the map hash table
      self.mapHash[key] = ( fName + '.html' )

    return key
#-------------------------------------------------------------

  # Takes a filename and generates the map requested by the user
  def generateMap( self, dirName, fName ) :
    """
    TODO: Create a way to parse a folder from filename

    """
    path = ( './v2/maptest/cases/'+ dirName +
             '/' + fName + '.csv' )

    # Reads in the info from our COVID data set
    df = pd.read_csv(path, dtype={"FIPS": str}, encoding='UTF-8')

    # We manually insert the CommunityID column in the DataFrame
    #   object so it is compatible with the cholopleth parameter
    if fName[:3] != 'ALL' :
      # The color parameter in the choropleth function
      #   only takes in int data types
      df.loc[:,'CommunityID'] = int( fName[0] )

    df['text'] = '' + df['Area_Name'] + ', ' + df['State'] + '' + '
' + \
                 'Pop. Density per Square Mile: ' + str(df['DensityPerSquaremileOfLandarea-Population']) + '
' + \
                 'Median Household Income in 2018: ' + str(df['Median_Household_Income_2018']) + '
' + \
                 '% Adults with High School Diplomas: ' + str(df['Percent_of_adults_with_a_high_school_diploma_only_2014-18'])

    text = df['text'].tolist()

    # Creates the map object and fills it with
    #   information
    fig = px.choropleth(df, geojson=counties, locations='FIPS', color='CommunityID',
                               range_color=[1,7],
                               # color_continuous_scale=px.colors.diverging.RdYlGn[::-1],
                               # color_continuous_scale=px.colors.sequential.OrRd,
                               #color_continuous_scale=([(0.000, 'rgb(175, 45, 36)'), (0.142,'rgb(175, 45, 36)'), # dark red
                               #                          (0.142, 'rgb(193, 78, 79)'), (0.285,'rgb(193, 78, 79)'), # light red
                               #                          (0.285, 'rgb(202, 122, 92)'), (0.428,'rgb(202, 122, 92)'), # nude
                               #                          (0.428, 'rgb(211, 166, 88)'), (0.571,'rgb(211, 166, 88)'), # mustard
                               #                          (0.571, 'rgb(183, 192, 85)'), (0.714,'rgb(183, 192, 85)'), # olive
                               #                          (0.714, 'rgb(143, 174, 98)'), (0.857,'rgb(143, 174, 98)'), # light green
                               #                          (0.857, 'rgb(119, 163, 111)'), (1.000,'rgb(119, 163, 111)')]), # green apple
                               color_continuous_scale=(['rgb(119, 163, 111)', # green apple
                                                        'rgb(143, 174, 98)', # light green
                                                        'rgb(183, 192, 85)', # olive
                                                        'rgb(211, 166, 88)', # mustard
                                                        'rgb(202, 122, 92)', # nude
                                                        'rgb(193, 78, 79)', # light red
                                                        'rgb(175, 45, 36)']), # dark red
                               scope="usa",
                               labels={'State':'State',
                                       'Area_Name':'County',
                                       'DensityPerSquaremileOfLandarea-Population':'Pop. Density per Square Mile of Land Area',
                                       'Median_Household_Income_2018':'Median Household Income in 2018',
                                       'Percent_of_adults_with_a_high_school_diploma_only_2014-18':'% Adults with High School Diplomas',
                                       'FIPS':'',
                                       'CommunityID':''},
                               hover_data=["Area_Name",
                                           "DensityPerSquaremileOfLandarea-Population",
                                           "Median_Household_Income_2018",
                                           "Percent_of_adults_with_a_high_school_diploma_only_2014-18"])
                               #hover_data=["text"])

    # Fixes and displays the legend to the
    #   right of the map
    fig.update_layout(margin={"r":0,"t":0,"l":0,"b":0},
                      coloraxis_colorbar = dict(
                        title = dict(
                          text = "Percentage Change in
the Number of New Cases",
                          font = dict(
                            size = 15)),
                        ticks = 'inside',
                        tickvals = [1, 2, 3, 4, 5, 6, 7],
                        ticktext = ['-100%: BIG DIP', '(-100%, 50%]: DOWNTICK',
                                    '(-50%, 0%): DECREASE', '0%: FLAT',
                                    '(0%, 50%): INCREASE', '[50%, 100%): UPTICK',
                                    'Greater than 100%: SPIKE'],
                        tickfont = dict(
                          size = 10),
                        thickness = 30,
                        xpad = 5,
                        len = 0.75))

    # Displays the figure when program is executed
    #fig.show()

    # Path to the specific directory the file is stored in
    nDir  = ( './v2/static/maps/' + dirName + '/' )

    # Path to where to store the file for the hash table
    nFile = fName + '.html'

    # Creates an html file of the map
    pio.write_html(fig, file= nDir + nFile, auto_open=False)

    return fig
#-------------------------------------------------------------

  # Initializes the CSV Hash Table
  def initCSVHash( self ) :
    # The dictionary that stores all the keys to
    #   the csv file we want to generate maps from
    csvHash = {}

    # Empties out the file before we store new values
    open( './v2/maptest/tables/csv_table.txt', 'w+' ).close()

    # Walks through each subdirectory in /cases/ and
    #   stores the names of the CSVs in csv_table.txt
    for path, subd, files in os.walk( './cases/' ) :
      for names in files:
        keys = names.replace( '.', '-' ).split( '-' )

        # Only the keys and CSV files will be stored
        #   in csv_table.txt
        if keys[-1] == 'csv' :
          # Creates the key for each value in the
          #   CSV hash table
          key = ( keys[0] + '-' + keys[3] +
                  '-' + keys[4] + '-' + keys[5][:2] +
                  '-' + keys[5][-2:] + '-' + keys[6] +
                  '-' + keys[7] + '-' + keys[-2][7:] )

          # Writes the values to csv_table.txt
          with open( './v2/maptest/tables/csv_table.txt', 'a+' ) as fp :
            fp.write( key + ',' + names + '\r\n' )

          fp.close()

    with open( './v2/maptest/tables/csv_table.txt', 'r' ) as fp :
      lines = fp.read().replace( '\r', '' ).split( '\n' )

    # Checks if the file is not empty
    if lines[0] != '' :
      for line in lines :
        line = line.split( ',' )

        key = line[0]

        # Indicates an end of file
        if key == '' :
          val = ''
        else :
          val = line[1]

        csvHash[key] = val

    return csvHash
#-------------------------------------------------------------

  # Initializes the Map Hash Table
  def initMapHash( self ) :
    # The dictionary that stores all the keys to
    #   the html file we want to display to the
    #   home page iframe
    mHash = {}

    with open( './v2/maptest/tables/html_table.txt', 'r' ) as fp :
      lines = fp.read().replace( '\r', '' ).split( '\n' )

    fp.close()

    # Checks if the file is not empty
    if lines[0] != '' :
      for line in lines :
        line = line.split( ',' )

        key = line[0]

        # Indicates an end of file
        if key == '' :
          val = ''
        else :
          val = line[1]

        mHash[key] = val

    return mHash
#-------------------------------------------------------------

  # Formats the key of the hash table to be similar
  #   the names of all the .csv and .html files
  def createKey( self, tempKey ) :
    key = ''
    front = ''
    middle = ''
    end = ''

    if( tempKey.endswith('ALL') ) :
      # Holds the ALL part of the string to be moved
      #   to the front of a key
      front = tempKey[-3:]

      if tempKey[1] == '-' :
        # Rearranges the string for key in order for
        #   it to fit the format of each files that
        #   will need to be read in
        middle = tempKey[1:-3]
        end    = tempKey[:1]
      else :
        middle = tempKey[2:-3]
        end = tempKey[:2]
    else :
      # Otherwise, last element in the key will be
      #   a single digit
      front = tempKey[-1:]

      if tempKey[1] == '-' :
        middle = tempKey[1:-1]
        end = tempKey[:1]
      else :
        middle = tempKey[2:-1]
        end = tempKey[:2]

    key = front + middle + end

    return key
#-------------------------------------------------------------

  def get_maphash( self ) :
    return self.mapHash
#-------------------------------------------------------------

# This is how you'll call the class
# test = map_test( 7, '04-JUL-20', '04-AUG-20', 2 )
# test = map_test( 30, '04-JUL-20', '04-AUG-20', 8 )
# key = test.main()


# This is where we test on a real chain (Rinkeby in this case)

from time import time
import pytest
from brownie import network
from scripts.deploy_lottery import deploy_lottery
from scripts.helpful_scripts import (
    LOCAL_BLOCKCHAIN_ENVIRONMENTS,
    fund_with_link,
    get_account,
)

def test_can_pick_winner():
    if network.show_active() in LOCAL_BLOCKCHAIN_ENVIRONMENTS:
        pytest.skip()
    lottery = deploy_lottery()
    account = get_account()
    lottery.startLottery({"from": account})
    lottery.enter(
        {"from": account, "value": lottery.getEntranceFee()}
    )  # We could add + 1000 Wei to make sure we match the Entrance Fee
    lottery.enter({"from": account, "value": lottery.getEntranceFee()})
    fund_with_link(lottery)
    lottery.endLottery({"from": account})
    time.sleep(60)
    assert lottery.recentWinner() == account
    assert lottery.balance() == 0


import typer
from pathlib import Path
from typing import Optional, List

from src.pipelines.transcribe import transcribe_audios_in_folder
from src.pipelines.export import export_corpus_dataset

from src.clients.scp_transfer import FileTransfer
from src.clients.database import Database

from src.models.file import AudioFormat

app = typer.Typer(
    no_args_is_help=True,
    help="Run data pipelines for diarization, transcription, metadata extraction for feeding the BrazSpeechData platform.",
)

DATA_PATH = Path("./data/")

@app.command(name="export")
def export(
    corpus_id: int = typer.Option(..., help="Corpus ID"),
    output_folder: Path = typer.Option(DATA_PATH / "export", help="Output folder"),
    csv: bool = typer.Option(False, help="Export to CSV"),
    json_metadata: bool = typer.Option(False, help="Export to JSON metadata"),
    textgrid: bool = typer.Option(False, help="Export to TextGrid"),
    continuous_text: bool = typer.Option(False, help="Export to continuous text"),
    speakers_text: bool = typer.Option(False, help="Export to speakers text"),
    original_audios: bool = typer.Option(False, help="Export original audios"),
    export_audio_to_formats: List[AudioFormat] = typer.Option(
        [AudioFormat.WAV, AudioFormat.MP3], help="Export audio to formats"
    ),
    google_drive_folder_ids: Optional[List[str]] = typer.Option(
        None, help="Google Drive folder IDs"
    ),
    filter_format: Optional[AudioFormat] = typer.Option(
        None, help="Filter audios by format"
    ),
    sample_rate: int = typer.Option(48000, help="Sample rate"),
    all: bool = typer.Option(False, help="Export all"),
    debug: bool = typer.Option(False, help="Debug mode"),
):
    if all:
        csv = True
        textgrid = True
        continuous_text = True
        speakers_text = True
        original_audios = True
        json_metadata = True

    with Database() as db:
        export_corpus_dataset(
            corpus_id=corpus_id,
            output_folder=output_folder,
            db=db,
            export_original_audios=original_audios,
            export_audio_to_formats=export_audio_to_formats,
            google_drive_folder_ids=google_drive_folder_ids,
            sample_rate=sample_rate,
            filter_format=filter_format,
            export_concanated_text=continuous_text,
            export_json_metadata=json_metadata,
            export_speakers_text=speakers_text,
            export_text_grid=textgrid,
            export_to_csv=csv,
            debug=debug,
        )

@app.command(name="transcribe")
def transcribe(
    corpus_id: int = typer.Option(..., help="Corpus ID"),
    folder_ids: List[str] = typer.Option(
        ..., help="Google Drive folder IDs that contain the audios to transcribe"
    ),
    output_folder: Path = typer.Option(DATA_PATH, help="Output folder"),
    storage_output_folder_id: str = typer.Option(
        None,
        help="Google Drive folder ID to save the transcriptions. If none is provided, the transcriptions will be saved in the same folder as the audios.",
    ),
    format_filter: Optional[AudioFormat] = typer.Option(
        None, help="Filter audios by format"
    ),
    save_to_db: bool = typer.Option(False, help="Save transcriptions to database"),
    save_to_drive: bool = typer.Option(
        False, help="Save transcriptions to Google Drive"
    ),
    transfer_to_server: bool = typer.Option(
        False, help="Transfer transcriptions to server"
    ),
):
    get_db_search_key = lambda x: x

    # Handle MUPE special case:
    if corpus_id == 1:
        format_filter = AudioFormat.MP4
        get_db_search_key = lambda x: "_".join(x.split("_")[:3])
    elif corpus_id == 2:
        format_filter = AudioFormat.WAV

    with FileTransfer() as ft:
        with Database() as db:
            transcribe_audios_in_folder(
                corpus_id=corpus_id,
                folder_ids=folder_ids,
                output_folder=output_folder,
                db=db if save_to_db else None,
                storage_output_folder_id=storage_output_folder_id,
                format_filter=format_filter,
                file_transfer_client=ft if transfer_to_server else None,
                save_to_drive=save_to_drive,
                get_db_search_key=get_db_search_key,
            )

if __name__ == "__main__":
    app()


# import pdb
import argparse
import os
import subprocess

from defaults import template_str_profile as template_str
from jinja2 import Template
from itertools import product
from rsparse.utils import params_to_filename, params_dict_to_str

from defaults import var_dict, CKPT_DIR, DATA_DIR, LOG_DIR

# all possible combinations for the python script's parameters

var_lists = list()

for k in var_dict.keys():
    var_lists.append([(k, v) for v in var_dict[k]])

def main(vars_, script_dir, overwrite, **args):
    script = os.path.join(script_dir, 'profile.py')
    template = Template(template_str)
    ckpt_dir = os.path.join(script_dir, CKPT_DIR)
    data_dir = os.path.join(script_dir, DATA_DIR)
    log_dir = os.path.join(script_dir, LOG_DIR)

    if not os.path.exists(ckpt_dir):
        os.mkdir(ckpt_dir)
    if not os.path.exists(data_dir):
        os.mkdir(data_dir)
    if not os.path.exists(log_dir):
        os.mkdir(log_dir)

    for i, x in enumerate(product(*var_lists)):
        x = dict(x)

        # exclude product + proximal
        if x['prox'] and x['reg'] == 'linfprod':
            continue
        if x['prox'] and x['reg'] == 'linfpath':
            args['time'] = '7:00:00'
        else:
            args['time'] = '1:00:00'

        if x['reg'] == 'linfproj':
            if x['prox'] and x['lambda_'] < 1e-2:
                continue
        else:
            if x['lambda_'] > 1e2:
                continue

        filename_ = params_to_filename(**x)
        ckpt = os.path.join(ckpt_dir, filename_ + '.pt')

        if os.path.isfile(ckpt) and not overwrite:
            continue

        x['ckpt_dir'] = ckpt_dir
        x['data_dir'] = data_dir
        x['overwrite'] = overwrite
        params = params_dict_to_str(x)

        with open('job_script.sh', 'w') as f:
            f.write(template.render(
                params=params, script=script, log_dir=log_dir, **args))

        subprocess.run('sbatch job_script.sh', shell=True)

    if os.path.isfile('job_script.sh'):
        os.remove('job_script.sh')

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    # python interpreter and where the project folder is located
    parser.add_argument(
            '--interpreter', type=str,
            default='/home/latorre/miniconda3/envs/pytorch/bin/python')
    parser.add_argument(
            '--script_dir', type=str,
            default='/home/latorre/robust_sparsity')

    # SLURM parameters
    parser.add_argument('--job_name', type=str, default='robsparse')
    parser.add_argument('--time', type=str, default='2:30:00')
    parser.add_argument('--mem', type=str, default='4GB')
    # parser.add_argument('--partition', type=str, default='gpu')
    parser.add_argument('--nodes', type=str, default='1')
    parser.add_argument('--overwrite', action='store_true')
    args = parser.parse_args()
    main(vars_=var_dict, **vars(args))


import os
from collections import Counter

class Utilities:

    def is_card_in_hand(self, guessing_hand, target_hand, guess, deck_list, deck, current_player, correct_guess):
        indices = [i for i, x in enumerate(target_hand) if x.card_number == guess]

        # Draw card if guess doesn't match
        if not indices:
            print('Go Fish!')
            hand, deck_list = deck.draw_cards(deck_list, guessing_hand, 1)
            correct_guess = False
        # Remove all matching cards from hand
        else:
            print('You guessed it!')
            for index in sorted(indices, reverse=True):
                guessing_hand.append(target_hand.pop(index))
                correct_guess = True

        return guessing_hand, target_hand, deck_list, correct_guess

    def scoring_cycle(self, hand, player_name, player_pts):
        counts = Counter(hand)
        match = False
        for ele in counts:
            if (counts[ele] == 4):
                print("Theres 4 of this card:")
                print(ele)
                # Remove element from hand and increment players score by 1.
                hand.remove(ele)
                print(player_name + str(player_pts))
                match = True
        if match:
            player_pts += 1

        return hand, player_pts

    def print_hand(self, hand, player_name):
        print(player_name + ":")
        hand.sort(key=lambda x: x.card_number)
        print(hand)

    def end_player_cycle(self):
        input("Press enter when ready...")
        os.system('clear')
        input("Please pass the game to the next player...")


from tools.general import load_input_list

def resolve_bsp(bsp_code, low_char, high_char):

    lower = 0
    upper = 2 ** len(bsp_code) - 1

    for c in bsp_code:
        mid = (lower + upper) // 2
        if c == high_char:
            lower = mid + 1
        elif c == low_char:
            upper = mid
        else:
            raise ValueError(f"Invalid character '{c}'")

    return lower

def seat_id(bp_code):

    row = resolve_bsp(bp_code[:7], 'F', 'B')
    col = resolve_bsp(bp_code[7:], 'L', 'R')

    return (8 * row) + col

taken_seats = sorted(seat_id(s) for s in load_input_list("day5.txt").split('\n'))

print(f"Part 1 => {taken_seats[-1]}")

prev = taken_seats[0]
for s in taken_seats[1:]:
    if s - prev == 2:
        print(f"Part 2 => {s - 1}")
        break
    prev = s


import io
import os
import platform
import struct
import threading

from fleetspeak.src.client.channel.proto.fleetspeak_channel import channel_pb2
from fleetspeak.src.common.proto.fleetspeak import common_pb2

_WINDOWS = platform.system() == "Windows"
if _WINDOWS:
  import msvcrt  # pylint: disable=g-import-not-at-top

class ProtocolError(Exception):
  """Raised when we do not understand the data received from Fleetspeak."""

# Constants to match behavior of channel.go.
_MAGIC = 0xF1EE1001

# We recommend that messages be ~1MB or smaller, and daemonservice has has 2MB
# hardcoded maximum.
MAX_SIZE = 2 * 1024 * 1024

# Format for the struct module to pack/unpack a 32 bit unsigned integer to/from
# a little endian byte sequence.
_STRUCT_FMT = "<I"

# The number of bytes required/produced when using _STRUCT_FMT.
_STRUCT_LEN = 4

# Environment variables, used to find the filedescriptors left open for
# us when started by Fleetspeak.
_INFD_VAR = "FLEETSPEAK_COMMS_CHANNEL_INFD"
_OUTFD_VAR = "FLEETSPEAK_COMMS_CHANNEL_OUTFD"

def _EnvOpen(var, mode):
  """Open a file descriptor identified by an environment variable."""
  value = os.getenv(var)
  if value is None:
    raise ValueError("%s is not set" % var)

  fd = int(value)

  # If running on Windows, convert the file handle to a C file descriptor; see:
  # https://groups.google.com/forum/#!topic/dev-python/GeN5bFJWfJ4
  if _WINDOWS:
    fd = msvcrt.open_osfhandle(fd, 0)

  return io.open(fd, mode)

class FleetspeakConnection(object):
  """A connection to the Fleetspeak system.

  It's safe to call methods of this class in parallel.
  """

  def __init__(self, version=None, read_file=None, write_file=None):
    """Connect to Fleetspeak.

    Connects to and begins an initial exchange of magic numbers with the
    Fleetspeak process. In normal use, the arguments are not required and will
    be created using the environment variables set by daemonservice.

    Args:
      version: A string identifying the version of the service being run. Will
        be included in resource reports for this service.
      read_file: A python file object, or similar, used to read bytes from
        Fleetspeak. If None, will be created based on the execution environment
        provided by daemonservice.
      write_file: A python file object, or similar, used to write bytes to
        Fleetspeak. If None, will be created based on the execution environment
        provided by daemonservice.

    Raises:
      ValueError: If read_file and write_file are not provided, and the
        corresponding environment variables are not set.
      ProtocolError: If we receive unexpected data from Fleetspeak.
    """
    self._read_file = read_file
    if not self._read_file:
      self._read_file = _EnvOpen(_INFD_VAR, "rb")

    self._read_lock = threading.Lock()

    self._write_file = write_file
    if not self._write_file:
      self._write_file = _EnvOpen(_OUTFD_VAR, "wb")

    self._write_lock = threading.Lock()

    # It is safer to send the magic number before reading it, in case the other
    # end does the same. Also, we'll be killed as unresponsive if we don't
    # write the magic number quickly enough. (Currently though, the other end is
    # the go implementation, which reads and writes in parallel.)
    self._WriteMagic()

    self._WriteStartupData(version)
    self._ReadMagic()

  def Send(self, message):
    """Send a message through Fleetspeak.

    Args:
      message: A message protocol buffer.

    Returns:
      Size of the message in bytes.
    Raises:
      ValueError: If message is not a common_pb2.Message.
    """
    if not isinstance(message, common_pb2.Message):
      raise ValueError("Send requires a fleetspeak.Message")

    if message.destination.service_name == "system":
      raise ValueError(
          "Only predefined messages can have destination.service_name =="
          ' "system"'
      )

    return self._SendImpl(message)

  def _SendImpl(self, message):
    if not isinstance(message, common_pb2.Message):
      raise ValueError("Send requires a fleetspeak.Message")

    buf = message.SerializeToString()
    if len(buf) > MAX_SIZE:
      raise ValueError(
          "Serialized message too large, size must be at most %d, got %d"
          % (MAX_SIZE, len(buf))
      )

    with self._write_lock:
      self._write_file.write(struct.pack(_STRUCT_FMT, len(buf)))
      self._write_file.write(buf)
      self._WriteMagic()

    return len(buf)

  def Recv(self):
    """Accept a message from Fleetspeak.

    Returns:
      A tuple (common_pb2.Message, size of the message in bytes).
    Raises:
      ProtocolError: If we receive unexpected data from Fleetspeak.
    """
    size = struct.unpack(_STRUCT_FMT, self._ReadN(_STRUCT_LEN))[0]
    if size > MAX_SIZE:
      raise ProtocolError(
          "Expected size to be at most %d, got %d" % (MAX_SIZE, size)
      )
    with self._read_lock:
      buf = self._ReadN(size)
      self._ReadMagic()

    res = common_pb2.Message()
    res.ParseFromString(buf)

    return res, len(buf)

  def Heartbeat(self):
    """Sends a heartbeat to the Fleetspeak client.

    If this daemonservice is configured to use heartbeats, clients that don't
    call this method often enough are considered faulty and are restarted by
    Fleetspeak.
    """
    heartbeat_msg = common_pb2.Message(
        message_type="Heartbeat",
        destination=common_pb2.Address(service_name="system"),
    )
    self._SendImpl(heartbeat_msg)

  def _ReadMagic(self):
    got = struct.unpack(_STRUCT_FMT, self._ReadN(_STRUCT_LEN))[0]
    if got != _MAGIC:
      raise ProtocolError(
          "Expected to read magic number {}, got {}.".format(_MAGIC, got)
      )

  def _WriteMagic(self):
    buf = struct.pack(_STRUCT_FMT, _MAGIC)
    self._write_file.write(buf)
    self._write_file.flush()

  def _WriteStartupData(self, version):
    startup_msg = common_pb2.Message(
        message_type="StartupData",
        destination=common_pb2.Address(service_name="system"),
    )
    startup_msg.data.Pack(
        channel_pb2.StartupData(pid=os.getpid(), version=version)
    )
    self._SendImpl(startup_msg)

  def _ReadN(self, n):
    """Reads n characters from the input stream, or until EOF.

    This is equivalent to the current CPython implementation of read(n), but
    not guaranteed by the docs.

    Args:
      n: int

    Returns:
      string
    """
    ret = b""
    while True:
      chunk = self._read_file.read(n - len(ret))
      ret += chunk

      if len(ret) == n or not chunk:
        return ret


import os
import random
import tempfile

import pytest

import teek

SMILEY_PATH = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), 'data', 'smiley.gif')

def test_width_and_height():
    image = teek.Image(file=SMILEY_PATH)
    assert image.width == 32
    assert image.height == 32

# tests that use this so that it returns True (for loops run) are marked slow
def slow_content_eq_check(image1, image2):
    # magic ftw, the double for loop is slow with big images which is why the
    # test images are small
    return (
        image1.width == image2.width and
        image1.height == image2.height and
        all(
            image1.get(x, y) == image2.get(x, y)
            for x in range(image1.width)
            for y in range(image1.height)
        )
    )

@pytest.mark.slow
def test_images_contain_same_data_util():
    image1 = teek.Image(file=SMILEY_PATH)
    image2 = teek.Image(file=SMILEY_PATH)
    assert slow_content_eq_check(image1, image1)
    assert slow_content_eq_check(image1, image2)

    image3 = teek.Image(width=1, height=2)
    assert not slow_content_eq_check(image1, image3)

def test_config(check_config_types):
    check_config_types(teek.Image(file=SMILEY_PATH).config, 'Image')

@pytest.mark.slow
def test_data_base64():
    with open(SMILEY_PATH, 'rb') as file:
        binary_data = file.read()

    assert slow_content_eq_check(teek.Image(file=SMILEY_PATH),
                                 teek.Image(data=binary_data))

@pytest.mark.slow
def test_from_to_tcl():
    image1 = teek.Image(file=SMILEY_PATH)
    teek.tcl_eval(None, 'proc returnArg {arg} {return $arg}')
    image2 = teek.tcl_call(teek.Image, 'returnArg', image1)
    assert image1.to_tcl() == image2.to_tcl()   # no implicit copying
    assert image1 is not image2                 # no implicit cache
    assert slow_content_eq_check(image1, image2)

def test_eq_hash():
    image1 = teek.Image(file=SMILEY_PATH)
    image2 = teek.Image.from_tcl(image1.to_tcl())
    image3 = teek.Image(file=SMILEY_PATH)

    assert image1 == image2
    assert {image1: 'woot'}[image2] == 'woot'
    assert image1 != image3
    assert image2 != image3
    assert image2 != 'woot'

def test_delete_and_all_images():
    old_images = set(teek.Image.get_all_images())
    new_image = teek.Image()
    assert set(teek.Image.get_all_images()) == old_images | {new_image}
    new_image.delete()
    assert set(teek.Image.get_all_images()) == old_images

def test_repr():
    image = teek.Image(file=SMILEY_PATH)
    assert repr(image) == '<Image: from %s, 32x32>' % repr(SMILEY_PATH)
    image.delete()
    assert repr(image) == '<Image: from %s, deleted>' % repr(SMILEY_PATH)

    image2 = teek.Image(width=123, height=456)
    assert repr(image2) == '<Image: 123x456>'
    image2.delete()
    assert repr(image2) == '<Image: deleted>'

def test_in_use():
    image = teek.Image(file=SMILEY_PATH)
    assert image.in_use() is False
    window = teek.Window()
    assert image.in_use() is False

    widget = teek.Label(window, image=image)
    assert image.in_use() is True
    widget2 = teek.Label(window, image=image)
    assert image.in_use() is True
    widget.destroy()
    assert image.in_use() is True
    widget2.destroy()
    assert image.in_use() is False

def test_blank():
    image1 = teek.Image(file=SMILEY_PATH)
    image2 = image1.copy()
    image2.blank()
    assert not slow_content_eq_check(image1, image2)      # not slow

@pytest.mark.slow
def test_read():
    image1 = teek.Image(file=SMILEY_PATH)
    assert image1.width > 0 and image1.height > 0

    image2 = teek.Image()
    assert image2.width == image2.height == 0
    image2.read(SMILEY_PATH)
    assert (image2.width, image2.height) == (image1.width, image1.height)
    assert slow_content_eq_check(image1, image2)

def test_redither():
    # no idea what this should do, so let's test that it does something...
    # except that it doesn't seem to actually do anything to the smiley, so
    # just test that it calls the Tcl redither command
    called = []
    fake_image = teek.create_command(called.append, [str])
    teek.Image.from_tcl(fake_image).redither()
    assert called == ['redither']

def test_transparency():
    image = teek.Image(file=SMILEY_PATH)
    x = random.randint(0, image.width - 1)
    y = random.randint(0, image.height - 1)
    assert image.transparency_get(x, y) is False
    image.transparency_set(x, y, True)
    assert image.transparency_get(x, y) is True
    image.transparency_set(x, y, False)
    assert image.transparency_get(x, y) is False

@pytest.mark.slow
def test_write():
    image1 = teek.Image(file=SMILEY_PATH)
    with tempfile.TemporaryDirectory() as tmpdir:
        asd = os.path.join(tmpdir, 'asd.gif')
        image1.write(asd, format='gif')
        image2 = teek.Image(file=asd)
    assert slow_content_eq_check(image1, image2)

def test_get_bytes():
    image = teek.Image(file=SMILEY_PATH)

    with tempfile.TemporaryDirectory() as tmpdir:
        asd = os.path.join(tmpdir, 'asd.gif')
        image.write(asd, format='gif')
        with open(asd, 'rb') as file:
            actual_gif_bytes = file.read()

    assert image.get_bytes('gif') == actual_gif_bytes


import os
from PIL import Image
import requests
from io import BytesIO
from flask import Flask, flash, request, redirect, url_for, render_template, send_from_directory
import tweepy
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address

app = Flask(__name__)
limiter = Limiter(
    app,
    key_func=get_remote_address,
    default_limits=["200 per day", "50 per hour"]
)
from werkzeug.utils import secure_filename

SECRET_KEY = "2*S28~NQ{p~c.aT"
UPLOAD_FOLDER = "C:/Users/tensor/Desktop/leaflet/static/uploads"
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
MAX_IMAGE_FILESIZE = 0.5 * 1024 * 1024

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['SECRET_KEY'] = SECRET_KEY

import base64

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route('/', methods=["GET", "POST"])
def index():
    return render_template("index.html")

@app.route('/test', methods=["GET", "POST"])
def test():
    return render_template("test.html")

@app.route('/uploads/<filename>')
def uploaded_file(filename):
    return send_from_directory(app.config['UPLOAD_FOLDER'],
                               filename)

@app.route("/twitter", methods=['GET', 'POST'])
@limiter.limit("100/day")
@limiter.limit("15/hour")
def twitter():
    if request.method == 'POST':
        username = request.form.get('text')
        auth = tweepy.OAuthHandler("VkrreD50glR0MKTM60ds2twHd", "TfVrWJumsxTcK9Aj8UAbk9jeGxnoZGikNFY9beQEP4DxOShc87")
        auth.set_access_token("878974576570617856-mDJZap5Jy3OeS9TxoLcgII9z6TM7SH9", "4BgyiRsxxkORNe0FApchpFgVLUsjAb6HMPPojnbCJa4gq")
        api = tweepy.API(auth)
        user = None
        try:
                user = api.get_user(screen_name=username)
                print(user)

        except tweepy.errors.TweepyException as e:
                flash("User not found")
                print(e)

                return render_template('index.html')
        else:
                global user_profile_image
                user_profile_image = user.profile_image_url.replace('normal', '400x400')

        data1 = BytesIO()
        data2 = BytesIO()
        data3 = BytesIO()
        data4 = BytesIO()
        data5 = BytesIO()
        data6 = BytesIO()
        data7 = BytesIO()
        data8 = BytesIO()
        data9 = BytesIO()
        data10 = BytesIO()

        data11 = BytesIO()
        data12 = BytesIO()
        data13 = BytesIO()

        response = requests.get(user_profile_image)

        userimage1 = Image.open(BytesIO(response.content))
        #userimage1 = userimage1.resize((280,280),Image.ANTIALIAS)
        image1 = Image.open('static/img/nullframes/are_you_hired_son.png')
        userimage1.paste(image1, (0,0), image1)
        userimage1.save(data1, 'PNG', quality=100)
        img1 = base64.b64encode(data1.getvalue())

        userimage2 = Image.open(BytesIO(response.content))
        #image2 = userimage2.resize((280,280),Image.ANTIALIAS)
        image2 = Image.open('static/img/nullframes/blood_red.png')
        userimage2.paste(image2, (0,0), image2)
        userimage2.save(data2, 'PNG', quality=100)
        img2 = base64.b64encode(data2.getvalue())

        userimage3 = Image.open(BytesIO(response.content))
        #image3 = userimage3.resize((280,280),Image.ANTIALIAS)
        image3 = Image.open('static/img/nullframes/chroma.png')
        userimage3.paste(image3, (0,0), image3)
        userimage3.save(data3, 'PNG', quality=100)
        img3 = base64.b64encode(data3.getvalue())

        userimage4 = Image.open(BytesIO(response.content))
        #image4 = userimage4.resize((280,280),Image.ANTIALIAS)
        image4 = Image.open('static/img/nullframes/finedog.png')
        userimage4.paste(image4, (0,0), image4)
        userimage4.save(data4, 'PNG', quality=100)
        img4 = base64.b64encode(data4.getvalue())

        userimage5 = Image.open(BytesIO(response.content))
        #image5 = userimage5.resize((280,280),Image.ANTIALIAS)
        image5 = Image.open('static/img/nullframes/mini_world.png')
        userimage5.paste(image5, (0,0), image5)
        userimage5.save(data5, 'PNG', quality=100)
        img5 = base64.b64encode(data5.getvalue())

        userimage6 = Image.open(BytesIO(response.content))
        #image6 = userimage6.resize((280,280),Image.ANTIALIAS)
        image6 = Image.open('static/img/nullframes/party_hard.png')
        userimage6.paste(image6, (0,0), image6)
        userimage6.save(data6, 'PNG', quality=100)
        img6 = base64.b64encode(data6.getvalue())

        userimage7 = Image.open(BytesIO(response.content))
        #image7 = userimage7.resize((280,280),Image.ANTIALIAS)
        image7 = Image.open('static/img/nullframes/finedog.png')
        userimage7.paste(image7, (0,0), image7)
        userimage7.save(data7, 'PNG', quality=100)
        img7 = base64.b64encode(data7.getvalue())

        userimage8 = Image.open(BytesIO(response.content))
        #image8 = userimage8.resize((280,280),Image.ANTIALIAS)
        image8 = Image.open('static/img/nullframes/shootme.png')
        userimage8.paste(image8, (0,0), image8)
        userimage8.save(data8, 'PNG', quality=100)
        img8 = base64.b64encode(data8.getvalue())

        userimage9 = Image.open(BytesIO(response.content))
        #image9 = userimage9.resize((280,280),Image.ANTIALIAS)
        image9 = Image.open('static/img/nullframes/sunger.png')
        userimage9.paste(image9, (0,0), image9)
        userimage9.save(data9, 'PNG', quality=100)
        img9 = base64.b64encode(data9.getvalue())

        userimage10 = Image.open(BytesIO(response.content))
        #image10 = userimage10.resize((280,280),Image.ANTIALIAS)
        image10 = Image.open('static/img/nullframes/natural.png')
        userimage10.paste(image10, (0,0), image10)
        userimage10.save(data10, 'PNG', quality=100)
        img10 = base64.b64encode(data10.getvalue())

        userimage11 = Image.open(BytesIO(response.content))
        #image11 = userimage11.resize((280,280),Image.ANTIALIAS)
        image11 = Image.open('static/img/nullframes/finedog_s.png')
        userimage11.paste(image11, (0,0), image11)
        userimage11.save(data11, 'PNG', quality=100)
        img11 = base64.b64encode(data11.getvalue())

        userimage12 = Image.open(BytesIO(response.content))
        #image8 = userimage12.resize((280,280),Image.ANTIALIAS)
        image12 = Image.open('static/img/nullframes/shootme_s.png')
        userimage12.paste(image12, (0,0), image12)
        userimage12.save(data12, 'PNG', quality=100)
        img12 = base64.b64encode(data12.getvalue())

        userimage13 = Image.open(BytesIO(response.content))
        #image9 = userimage13.resize((280,280),Image.ANTIALIAS)
        image13 = Image.open('static/img/nullframes/sunger_s.png')
        userimage13.paste(image13, (0,0), image13)
        userimage13.save(data13, 'PNG', quality=100)
        img13 = base64.b64encode(data13.getvalue())

    return render_template('processing.html' , filename=user_profile_image, are_you_hired_son=img1.decode('utf-8'), blood_red=img2.decode('utf-8'), chroma=img3.decode('utf-8'), mini_world=img5.decode('utf-8'), nature=img10.decode('utf-8'), party_hard= img6.decode('utf-8'), finedog=img7.decode('utf-8'), shootme=img8.decode('utf-8'), sunger=img9.decode('utf-8'), finedog_s=img11.decode('utf-8'), shootme_s=img12.decode('utf-8'), sunger_s=img13.decode('utf-8'))

@app.route("/processing", methods=['GET', 'POST'])
def processing():
    return render_template('processing.html')

@app.errorhandler(404)
def page_not_found(e):
    return render_template('404.html'), 404


from PIL import Image
import os
from tqdm import tqdm
import os.path
import numpy as np
import pickle
import sys
from typing import Any, Callable, Optional, Tuple

from torch.utils.data import Dataset
from torch.autograd import Variable

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torchvision.datasets import VisionDataset
import torchvision.models as models

from dataset.VanillaBackprop import VanillaBackprop, convert_to_grayscale
import models.resnet as RN
from train_utils.cutout import Cutout

class CIFAR10(VisionDataset):
    base_folder = 'cifar-10-batches-py'
    url = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
    filename = "cifar-10-python.tar.gz"
    tgz_md5 = 'c58f30108f718f92721af3b95e74349a'
    train_list = [
        ['data_batch_1', 'c99cafc152244af753f735de768cd75f'],
        ['data_batch_2', 'd4bba439e000b95fd0a9bffe97cbabec'],
        ['data_batch_3', '54ebc095f3ab1f0389bbae665268c751'],
        ['data_batch_4', '634d18415352ddfa80567beed471001a'],
        ['data_batch_5', '482c414d41f54cd18b22e5b47cb7c3cb'],
    ]

    test_list = [
        ['test_batch', '40351d587109b95175f43aff81a1287e'],
    ]
    meta = {
        'filename': 'batches.meta',
        'key': 'label_names',
        'md5': '5ff9c542aee3614f3951f8cda6e48888',
    }

    def __init__(self, root, train=True, transform=None, target_transform=None, download=False, args= None):
        self.train = train  # training set or test set
        self.args = args
        self.root = root
        if self.train:
            downloaded_list = self.train_list
        else:
            downloaded_list = self.test_list

        self.data = []
        self.targets = []
        self.transform=transform

        # now load the picked numpy arrays
        for file_name, checksum in downloaded_list:
            file_path = os.path.join(self.root, self.base_folder, file_name)
            with open(file_path, 'rb') as f:
                if sys.version_info[0] == 2:
                    entry = pickle.load(f)
                else:
                    entry = pickle.load(f, encoding='latin1')
                self.data.append(entry['data'])
                if 'labels' in entry:
                    self.targets.extend(entry['labels'])
                else:
                    self.targets.extend(entry['fine_labels'])

        self.data = np.vstack(self.data).reshape(-1, 3, 32, 32)
        self.data = self.data.transpose((0, 2, 3, 1))  # convert to HWC

    def __getitem__(self, index):
        target = self.targets[index]
        image = self.data[index]
        image = image.astype(np.uint8)
        image = Image.fromarray(image)
        if self.transform is not None:
            image = self.transform(image)

        image = torch.FloatTensor(image)

        return image, target

    def __len__(self):
        return len(self.data)

class CIFAR100(CIFAR10):
    base_folder = 'cifar-100-python'
    url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz"
    filename = "cifar-100-python.tar.gz"
    tgz_md5 = 'eb9058c3a382ffc7106e4002c42a8d85'
    train_list = [
        ['train', '16019d7e3df5f24257cddd939b257f8d'],
    ]

    test_list = [
        ['test', 'f0ef6b0ae62326f3e7ffdfab6717acfc'],
    ]
    meta = {
        'filename': 'meta',
        'key': 'fine_label_names',
        'md5': '7973b15100ade9c7d40fb424638fde48',
    }

def get_CIFAR_dataloader(args):
    n_holes = 1
    if args.dataset == 'cifar100':
        length = 8
    elif args.dataset == 'cifar10':
        length = 16
    args.lr = 0.25
    args.momentum = 0.9
    args.weight_decay = 1e-4
    args.batch_size = 64
    args.epochs = 300
    normalize = transforms.Normalize(mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],
                                        std=[x / 255.0 for x in [63.0, 62.1, 66.7]])

    if args.augmentation == 'baseline':
        print("Augmentation Apply Base")
        transform_train = transforms.Compose([
            transforms.ToTensor(),
            normalize,
        ])
    elif args.augmentation == 'cutout':
        print("Augmentation Apply cutout")
        transform_train = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
            Cutout(n_holes=n_holes, length=length)
        ])
    else:
        print("Augmentation Apply augmentation")
        transform_train = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ])

    transform_test = transforms.Compose([
        transforms.ToTensor(),
        normalize
    ])

    if args.dataset == 'cifar100':
        train_loader = torch.utils.data.DataLoader(
            CIFAR100('../data', train=True, download=True, transform=transform_train, args=args),
            batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=False)
        val_loader = torch.utils.data.DataLoader(
            CIFAR100('../data', train=False, download=True, transform=transform_test),
            batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=False)
        numberofclass = 100
    elif args.dataset == 'cifar10':
        train_loader = torch.utils.data.DataLoader(
            CIFAR10('../data', train=True, download=True, transform=transform_train, args=args),
            batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=False)
        val_loader = torch.utils.data.DataLoader(
            CIFAR10('../data', train=False, download=True, transform=transform_test),
            batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=False)
        numberofclass = 10
    else:
        raise Exception('unknown dataset: {}'.format(args.dataset))

    return train_loader, val_loader, numberofclass


# -*- coding: utf-8 -*-
"""
Created on Tue Oct  5 09:01:40 2021

@author: staar
"""

#%%
import matplotlib as mpl
help(mpl.figure) # too big -- this is sub-module !
dir(mpl.figure)

help(mpl.figure.Figure)  # it's still big!

#%% see first simpler pyplot constructor

import matplotlib.pyplot as plt
help(plt.figure)

#%%
"""
figure(num=None, figsize=None, dpi=None, facecolor=None, edgecolor=None, frameon=True,
       FigureClass=<class 'matplotlib.figure.Figure'>, clear=False, **kwargs)

    Create a new figure, or activate an existing figure.

    Parameters
    ----------
    num : int or str, optional
        A unique identifier for the figure.

        If a figure with that identifier already exists, this figure is made
        active and returned.
        An integer refers to the ``Figure.number``
        attribute, a string refers to the figure label.

        If there is no figure with the identifier or *num* is not given, a new
        figure is created, made active and returned.
        If *num* is an int, it
        will be used for the ``Figure.number`` attribute, otherwise, an
        auto-generated integer value is used (starting at 1 and incremented
        for each new figure). If *num* is a string, the figure label and the
        window title is set to this value.

    figsize : (float, float), default: :rc:`figure.figsize`
        Width, height in inches.

    dpi : float, default: :rc:`figure.dpi`
        The resolution of the figure in dots-per-inch.

    facecolor : color, default: :rc:`figure.facecolor`
        The background color.

    edgecolor : color, default: :rc:`figure.edgecolor`
        The border color.

    frameon : bool, default: True
        If False, suppress drawing the figure frame.

    FigureClass : subclass of `~matplotlib.figure.Figure`
        Optionally use a custom `.Figure` instance.

    clear : bool, default: False
        If True and the figure already exists, then it is cleared.

    tight_layout : bool or dict, default: :rc:`figure.autolayout`
        If ``False`` use *subplotpars*. If ``True`` adjust subplot
        parameters using `.tight_layout` with default padding.
        When providing a dict containing the keys ``pad``, ``w_pad``,
        ``h_pad``, and ``rect``, the default `.tight_layout` paddings
        will be overridden.

    constrained_layout : bool, default: :rc:`figure.constrained_layout.use`
        If ``True`` use constrained layout to adjust positioning of plot
        elements.  Like ``tight_layout``, but designed to be more
        flexible.  See
        :doc:`/tutorials/intermediate/constrainedlayout_guide`
        for examples.  (Note: does not work with `add_subplot` or
        `~.pyplot.subplot2grid`.)

    **kwargs : optional
        See `~.matplotlib.figure.Figure` for other possible arguments.

    Returns
    -------
    `~matplotlib.figure.Figure`
        The `.Figure` instance returned will also be passed to
        new_figure_manager in the backends, which allows to hook custom
        `.Figure` classes into the pyplot interface. Additional kwargs will be
        passed to the `.Figure` init function.

    Notes
    -----
    If you are creating many figures, make sure you explicitly call
    `.pyplot.close` on the figures you are not using, because this will
    enable pyplot to properly clean up the memory.

    `~matplotlib.rcParams` defines the default values, which can be modified
    in the matplotlibrc file.
"""
#%%

plt.figure(1, clear=True)  # clearing figure 1 (if exists)
plt.figure(1, tight_layout=True)

#%%
fig = mpl.figure.Figure()
ax = fig.add_subplot()
ax.plot(range(5), [3, 2, 5, 4, 1])
# nothing...

fig.draw(...)     # ???

fig.show()  # AttributeError: Figure.show works only for figures managed by pyplot,
            # normally created by pyplot.figure()

#%%
#%%
import matplotlib as mpl
help(mpl.figure) # too big -- this is sub-module !
dir(mpl.figure)

help(mpl.figure.Figure)

# https://matplotlib.org/stable/api/figure_api.html#module-matplotlib.figure
"""
Parameters:

figsize : 2-tuple of floats, default: rcParams["figure.figsize"] (default: [6.4, 4.8])
    Figure dimension (width, height) in inches.

dpi : float, default: rcParams["figure.dpi"] (default: 100.0)
    Dots per inch.

facecolor : default: rcParams["figure.facecolor"] (default: 'white')
    The figure patch facecolor.

edgecolor : default: rcParams["figure.edgecolor"] (default: 'white')
    The figure patch edge color.

linewidth : float
    The linewidth of the frame (i.e. the edge linewidth of the figure patch).

frameonbool : default: rcParams["figure.frameon"] (default: True)
    If False, suppress drawing the figure background patch.

subplotpars : SubplotParams
    Subplot parameters. If not given, the default subplot parameters rcParams["figure.subplot.*"] are used.

tight_layout : bool or dict, default: rcParams["figure.autolayout"] (default: False)
    If False use subplotpars.
    If True adjust subplot parameters using tight_layout with default padding.
    When providing a dict containing the keys pad, w_pad, h_pad, and rect,
    the default tight_layout paddings will be overridden.

constrained_layout : bool, default: rcParams["figure.constrained_layout.use"] (default: False)
    If True use constrained layout to adjust positioning of plot elements.
    Like tight_layout, but designed to be more flexible. See Constrained Layout Guide for examples. (Note: does not work with add_subplot or subplot2grid.)

"""
#%%


from operator import itemgetter

from pytz import common_timezones, common_timezones_set
from wtforms.fields import BooleanField, EmailField, IntegerField, SelectField, StringField
from wtforms.validators import DataRequired, Email, NumberRange, ValidationError

from indico.core.auth import multipass
from indico.core.config import config
from indico.modules.auth.forms import LocalRegistrationForm, _check_existing_email
from indico.modules.users import User
from indico.modules.users.models.emails import UserEmail
from indico.modules.users.models.users import NameFormat
from indico.util.i18n import _, get_all_locales
from indico.web.forms.base import IndicoForm
from indico.web.forms.fields import (IndicoEnumSelectField, IndicoSelectMultipleCheckboxField, MultiStringField,
                                     PrincipalField, PrincipalListField)
from indico.web.forms.util import inject_validators
from indico.web.forms.validators import HiddenUnless
from indico.web.forms.widgets import SwitchWidget

class UserPreferencesForm(IndicoForm):
    lang = SelectField(_('Language'))

    force_language = BooleanField(
        _('Use my language'),
        widget=SwitchWidget(),
        description=_("Always use my preferred language instead of an event's supported language."
                      ' This may result in a mix of two languages in some areas of an event.'))

    timezone = SelectField(_('Timezone'))

    force_timezone = BooleanField(
        _('Use my timezone'),
        widget=SwitchWidget(),
        description=_("Always use my current timezone instead of an event's timezone."))

    show_future_events = BooleanField(
        _('Show future events'),
        widget=SwitchWidget(),
        description=_('Show future events by default.'))

    show_past_events = BooleanField(
        _('Show past events'),
        widget=SwitchWidget(),
        description=_('Show past events by default.'))

    name_format = IndicoEnumSelectField(_('Name format'), enum=NameFormat,
                                        description=_('Default format in which names are displayed'))

    use_previewer_pdf = BooleanField(
        _('Use previewer for PDF files'),
        widget=SwitchWidget(),
        description=_('The previewer is used by default for image and text files, but not for PDF files.'))

    add_ical_alerts = BooleanField(
        _('Add alerts to iCal'),
        widget=SwitchWidget(),
        description=_('Add an event reminder to exported iCal files/URLs.'))

    add_ical_alerts_mins = IntegerField(
        _('iCal notification time'),
        [HiddenUnless('add_ical_alerts'), NumberRange(min=0)],
        description=_('Number of minutes to notify before an event.'))

    use_markdown_for_minutes = BooleanField(
        _('Markdown editor for minutes'),
        widget=SwitchWidget(),
        description=_('Use Markdown editor instead of HTML editor when editing the minutes of a meeting.'))

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        locales = [(code, f'{name} ({territory})' if territory else name)
                   for code, (name, territory, __) in get_all_locales().items()]
        self.lang.choices = sorted(locales, key=itemgetter(1))
        self.timezone.choices = list(zip(common_timezones, common_timezones))
        if self.timezone.object_data and self.timezone.object_data not in common_timezones_set:
            self.timezone.choices.append((self.timezone.object_data, self.timezone.object_data))

class UserEmailsForm(IndicoForm):
    email = EmailField(_('Add new email address'), [DataRequired(), Email()], filters=[lambda x: x.lower() if x else x])

    def validate_email(self, field):
        conflict = (UserEmail.query
                    .filter(~User.is_pending,
                            ~UserEmail.is_user_deleted,
                            UserEmail.email == field.data)
                    .join(User)
                    .has_rows())
        if conflict:
            raise ValidationError(_('This email address is already in use.'))

class SearchForm(IndicoForm):
    last_name = StringField(_('Family name'))
    first_name = StringField(_('First name'))
    email = StringField(_('Email'), filters=[lambda x: x.lower() if x else x])
    affiliation = StringField(_('Affiliation'))
    exact = BooleanField(_('Exact match'))
    include_deleted = BooleanField(_('Include deleted'))
    include_pending = BooleanField(_('Include pending'))
    external = BooleanField(_('External'))

class MergeForm(IndicoForm):
    source_user = PrincipalField(_('Source user'), [DataRequired()],
                                 description=_('The user that will be merged into the target one'))
    target_user = PrincipalField(_('Target user'), [DataRequired()],
                                 description=_('The user that will remain active in the end'))

class AdminUserSettingsForm(IndicoForm):
    notify_account_creation = BooleanField(_('Registration notifications'), widget=SwitchWidget(),
                                           description=_('Send an email to all administrators whenever someone '
                                                         'registers a new local account.'))
    email_blacklist = MultiStringField(_('Email blacklist'), field=('email_blacklist', _('email')),
                                       unique=True, flat=True,
                                       description=_('Prevent users from creating Indico accounts with these email '
                                                     'addresses. Supports wildcards, e.g. *@gmail.com'))
    allow_personal_tokens = BooleanField(_('Personal API tokens'), widget=SwitchWidget(),
                                         description=_('Whether users are allowed to generate personal API tokens. '
                                                       'If disabled, only admins can create them, but users will '
                                                       'still be able to regenerate the tokens assigned to them.'))
    mandatory_fields_account_request = IndicoSelectMultipleCheckboxField(
        _('Mandatory fields in account request'),
        choices=[('affiliation', _('Affiliation')), ('comment', _('Comment'))],
        description=_('Fields a new user has to fill in when requesting an account')
    )

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if not multipass.has_moderated_providers:
            del self.mandatory_fields_account_request

class AdminAccountRegistrationForm(LocalRegistrationForm):
    email = EmailField(_('Email address'), [DataRequired(), Email(), _check_existing_email],
                       filters=[lambda s: s.lower() if s else s])
    create_identity = BooleanField(_('Set login details'), widget=SwitchWidget(), default=True)

    def __init__(self, *args, **kwargs):
        if config.LOCAL_IDENTITIES:
            for field in ('username', 'password', 'confirm_password'):
                inject_validators(self, field, [HiddenUnless('create_identity')], early=True)
        super().__init__(*args, **kwargs)
        if not config.LOCAL_IDENTITIES:
            del self.username
            del self.password
            del self.confirm_password
            del self.create_identity

class AdminsForm(IndicoForm):
    admins = PrincipalListField(_('Admins'), [DataRequired()])


from django.urls import re_path, include
from rest_framework import routers
from .viewsets import (
    FoodViewSet,
    IngredientViewSet,
    RestrictionTagViewSet,
    RecipeViewSet
)

# Routers provide an easy way of automatically determining the URL conf.
router = routers.DefaultRouter()
router.register(r'food', FoodViewSet)
router.register(r'ingredients', IngredientViewSet)
router.register(r'restriction_tags', RestrictionTagViewSet)
router.register(r'recipes', RecipeViewSet)

# Additionally, we include login URLs for the browsable API.
urlpatterns = [
    re_path(r'^', include(router.urls)),
    re_path(r'^api-auth/', include('rest_framework.urls', namespace='rest_framework'))
]


from merlinpyspark.config import BigQuerySourceConfig
from merlinpyspark.source import create_source
from merlinpyspark.spec.prediction_job_pb2 import BigQuerySource
import pytest

@pytest.mark.ci
def test_bq_source(spark_session):
    bq_src_proto = BigQuerySource(
        table="bigquery-public-data:samples.shakespeare",
        features=["word", "word_count"],
        options={"filter": "word_count > 100"}
    )
    src_cfg = BigQuerySourceConfig(bq_src_proto)

    bq_source = create_source(spark_session, src_cfg)
    df = bq_source.load()
    assert len(df.columns) == 2
    assert df.columns[0] == "word"
    assert df.columns[1] == "word_count"
    assert df.sort("word_count").take(1)[0]["word_count"] > 100


"""Reads in python dictionary files."""

import io

from celestine.stream import (
    SECTION_BREAK,
    module_open,
)
from celestine.typed import (
    GS,
    S,
)
from celestine.unicode import (
    LINE_FEED,
    REVERSE_SOLIDUS,
)
from celestine.unicode.alphabet2 import UNICODE

def fix_line_split(*path: S) -> GS:
    """"""
    skip = False

    document = module_open(*path)
    for string in document:
        for character in string:
            if character not in UNICODE:
                continue

            if character == REVERSE_SOLIDUS:
                skip = True
                continue

            if skip:
                skip = False
                continue

            yield from character

def make_dictionary(document):
    dictionary = {}

    lines = document.split(LINE_FEED)
    for line in lines:
        if line[0:3] == '"""':
            continue

        if not line:
            continue
        split = line.split("=")
        key = split[0].strip()
        value = split[-1].strip()[1:-1]
        dictionary[key] = value
    return dictionary

def open_language(*path):
    """Convert a dictionary to a string and save it to a file."""
    text = fix_line_split(*path)
    lines = read_new_lines(text)
    # normal = normalize(lines)
    normal = lines

    string = io.StringIO()
    for line in normal:
        string.write(line)

    value = string.getvalue()

    split = value.split(SECTION_BREAK)

    head = make_dictionary(split[0])
    body = make_dictionary(split[-1])

    return (head, body)

def read_new_lines(string: S) -> GS:
    """"""

    buffer = io.StringIO()
    count = 0

    for character in string:
        count += buffer.write(character)

        if character == LINE_FEED:
            buffer.seek(0, io.SEEK_SET)
            line = buffer.read(count)
            buffer.seek(0, io.SEEK_SET)
            count = 0
            yield from line


import torch
from sentence_transformers import SentenceTransformer
import numpy as np
from tqdm import tqdm
import json
import argparse
from PIL import Image
from lavis.models import load_model_and_preprocess
import torch

def sbert_encoder(utterances, directory, split):
    """
    Encodes `utterances` with sentence_transformers library and saves embeddings to .npy file.

    Params
    ------
        utterances: list[str], all utterances from dataset
        path: str, where to save .npy file
    """
    model = SentenceTransformer('sentence-transformers/all-distilroberta-v1').to('cuda')
    embeddings = model.encode(utterances, show_progress_bar=True)
    np.save(f'{directory}/{split}_embeddings', embeddings)

def clip_encoder(utterances, directory, split, image_batch_size=8, text_batch_size=32):
    """
    Encodes `utterances` with sentence_transformers library and saves embeddings to .npy file.

    Params
    ------
        utterances: list[str], all utterances from dataset
        path: str, where to save .npy file
    """

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model, vis_processors, txt_processors = load_model_and_preprocess(
        name="clip_feature_extractor",
        model_type="ViT-B-16",
        is_eval=True,
        device=device
    )

    # divide images, texts and last utterances
    images = []
    texts = []
    last_utters = []
    is_image = []
    for i, ut in enumerate(utterances):
        if isinstance(ut, list):
            is_image.append(True)
            img, lut = ut
            images.append(Image.open(img).convert('RGB'))
            last_utters.append(lut)
        else:
            is_image.append(False)
            texts.append(ut)
    is_image = np.array(is_image)

    pbar = tqdm(total=len(images) * 2 + len(texts))

    # encode them all
    image_embeddings = []
    for i in range(0, len(images), image_batch_size):
        inp = images[i:i+image_batch_size]
        preproc = torch.stack(list(map(vis_processors['eval'], inp)))
        embs = model.extract_features({'image': preproc.to(device)})
        image_embeddings.append(embs.detach().cpu().numpy())
        pbar.update(len(inp))
    image_embeddings = np.concatenate(image_embeddings, axis=0)

    text_embeddings = []
    for i in range(0, len(texts), text_batch_size):
        inp = texts[i:i+text_batch_size]
        preproc = txt_processors['eval'](inp)
        embs = model.extract_features({'text_input': preproc})
        text_embeddings.append(embs.detach().cpu().numpy())
        pbar.update(len(inp))
    text_embeddings = np.concatenate(text_embeddings, axis=0)

    last_utters_embeddings = []
    for i in range(0, len(texts), text_batch_size):
        inp = last_utters[i:i+text_batch_size]
        preproc = txt_processors['eval'](inp)
        embs = model.extract_features({'text_input': preproc})
        last_utters_embeddings.append(embs.detach().cpu().numpy())
        pbar.update(len(inp))
    last_utters_embeddings = np.concatenate(last_utters_embeddings, axis=0)

    # these will be used in dgac_clustering.py
    embeddings = np.empty((len(utterances), text_embeddings.shape[1]))
    embeddings[is_image] = image_embeddings
    embeddings[~is_image] = text_embeddings
    np.save(f'{directory}/{split}_embeddings', embeddings)

    # these will be used in response_selection.py and my_dff.py
    embeddings[is_image] = last_utters_embeddings
    np.save(f'{directory}/{split}_response_embeddings', embeddings)

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument('--task', dest='task', required=True)
    args = ap.parse_args()

    directory = f'{args.task}_stuff'

    # load json files
    train_utterances = json.load(open(f'{directory}/train_utterances.json', 'r'))
    test_utterances = json.load(open(f'{directory}/test_utterances.json', 'r'))
    val_utterances = json.load(open(f'{directory}/val_utterances.json', 'r'))

    if args.task == 'persona_chat':
        encoder = sbert_encoder
    elif args.task == 'imad':
        encoder = clip_encoder
    else:
        raise ValueError(f'not supported: {args.task}')

    encoder(train_utterances, directory, 'train')
    encoder(test_utterances, directory, 'test')
    encoder(val_utterances, directory, 'val')


from dataclasses import dataclass
from typing import List

from action.argmapping.analytics import (
    CollFormArgs, CTFreqFormArgs, FreqFormArgs)
from action.control import http_action
from action.errors import ImmediateRedirectException, UserReadableException
from action.krequest import KRequest
from action.model.concordance import ConcActionModel
from action.response import KResponse
from conclib.calc import require_existing_conc
from conclib.errors import ConcNotFoundException
from conclib.pyconc import PyConc
from dataclasses_json import dataclass_json
from main_menu import MainMenu
from sanic import Blueprint

bp = Blueprint('dispersion', url_prefix='dispersion')

@dataclass_json
@dataclass
class FreqDispersionBin:
    start: float
    position: float
    end: float
    freq: int

def _get_freq_dispersion(conc: PyConc, resolution: int) -> List[FreqDispersionBin]:
    conc_begs, values = conc.xdistribution([0] * resolution, 101)

    abs_freq = []
    last_valid_item = None
    for beg in reversed(conc_begs):
        # if beg is 0, it means there are no concordances in the bin
        if beg > 0:
            if last_valid_item is None:
                abs_freq.append(int(conc.size()) - beg)
            else:
                # `last_valid_item - beg` is number of concordances
                # between beginnig of last non empty bin and the beginning of current bin
                # (for cycle is going backwards)
                abs_freq.append(last_valid_item - beg)
            last_valid_item = beg
        else:
            abs_freq.append(0)

    freq_dispersion = [
        FreqDispersionBin(
            100 * i / len(conc_begs),
            100 * (i + 0.5) / len(conc_begs),
            100 * (i + 1) / len(conc_begs),
            freq,
        ) for i, freq in enumerate(reversed(abs_freq))
    ]

    return freq_dispersion

@bp.route('/ajax_get_freq_dispersion')
@http_action(action_model=ConcActionModel, return_type='json')
async def ajax_get_freq_dispersion(amodel: ConcActionModel, req: KRequest, resp: KResponse) -> List[FreqDispersionBin]:
    conc = await require_existing_conc(amodel.corp, amodel.args.q, amodel.args.cutoff)
    resolution = int(req.args.get('resolution', 100))
    if 0 < resolution <= 1000:
        return _get_freq_dispersion(conc, resolution)
    raise UserReadableException('Invalid dispersion resolution. Acceptable values [1, 1000].')

@bp.route('/index')
@http_action(action_model=ConcActionModel, page_model='dispersion', template='dispersion.html')
async def index(amodel: ConcActionModel, req: KRequest, response: KResponse):
    amodel.disabled_menu_items = (
        MainMenu.CONCORDANCE('query-save-as'),
        MainMenu.VIEW('kwic-sent-switch'),
        MainMenu.CONCORDANCE('query-overview'))

    try:
        conc = await require_existing_conc(amodel.corp, amodel.args.q, amodel.args.cutoff)
    except ConcNotFoundException:
        args = list(req.args.items()) + [('next', 'dispersion')]
        raise ImmediateRedirectException(req.create_url('restore_conc', args))

    resolution = int(req.args.get('resolution', 100))
    if resolution > 1000:
        resolution = 1000
    elif resolution < 1:
        resolution = 1

    result = {
        'coll_form_args': CollFormArgs().update(amodel.args).to_dict(),
        'freq_form_args': FreqFormArgs().update(amodel.args).to_dict(),
        'ctfreq_form_args': CTFreqFormArgs().update(amodel.args).to_dict(),
        'text_types_data': await amodel.tt.export_with_norms(ret_nums=True),
        'dispersion_resolution': resolution,
        'initial_data': _get_freq_dispersion(conc, resolution),
    }
    await amodel.export_query_forms(result)
    return result


from datetime import datetime, timedelta

def get_birthdays_per_week(users):
    users = sorted(users,key=lambda x: x["birthday"].replace(year=1970))

    curr = datetime.now()
    start = datetime(curr.year, curr.month, curr.day) - timedelta(curr.weekday()) + timedelta(5)
    end = start + timedelta(7)

    days = [[], [], [], [], []]

    for user in users:
        name = user["name"]
        birth = user["birthday"]

        if birth.replace(year=start.year) < start:
            continue
        if birth >= end:
            continue;

        if birth.weekday() == 0 or birth.weekday() > 4:
            days[0].append(name);
        else:
            days[birth.weekday()].append(name);

    for i, day in enumerate(days):
        if i == 0:
            print("Monday: ", end='')
        elif i == 1:
            print("Tuesday: ", end='')
        elif i == 2:
            print("Wednesday: ", end='')
        elif i == 3:
            print("Thursday: ", end='')
        elif i == 4:
            print("Friday: ", end='')

        print(*day, sep=', ')

users = [
    { "name": "Bill", "birthday": datetime.fromisoformat("1970-02-10 10:10:10") },
    { "name": "Ash", "birthday": datetime.fromisoformat("1970-02-10 10:10:10") },
    { "name": "Vanya", "birthday": datetime.fromisoformat("2022-02-05 10:10:10") },
    { "name": "Ak", "birthday": datetime.fromisoformat("2022-02-10 10:10:10") },
    { "name": "Zuk", "birthday": datetime.fromisoformat("2023-02-11 10:10:10") },
]

get_birthdays_per_week(users)


import argparse
import os
import glob
import pickle
import numpy as np
import pandas as pd
from sklearn.cluster import KMeans

import torch

from src.ebirl import EBIRL
from src.data import pad_collate
from src.train_utils import train, Logger

def parse_args():
    bool_ = lambda x: x if isinstance(x, bool) else x == "True"

    parser = argparse.ArgumentParser()
    # data args
    parser.add_argument("--data_path", type=str, default="../../data", help="data path, default=../../data")
    parser.add_argument("--exp_path", type=str, default="../../exp", help="experiment path, default=../exp")
    # agent args
    parser.add_argument("--state_dim", type=int, default=2, help="agent state dimension, default=2")
    parser.add_argument("--horizon", type=int, default=30, help="agent max plan horizon, default=30")
    # prior args
    parser.add_argument("--prior_cov", type=str, choices=["diag", "full"], default="full", help="prior covariance type, default=full")
    parser.add_argument("--bc_penalty", type=float, default=1., help="prior behavior cloning penalty, default=1.")
    parser.add_argument("--obs_penalty", type=float, default=1., help="observation likelihood penalty, default=1.")
    parser.add_argument("--kl_penalty", type=float, default=1., help="prior kl penalty, default=1.")
    # train args
    parser.add_argument("--cp_path", type=str, default="none", help="checkpoint path, default=none")
    parser.add_argument("--lr", type=float, default=0.01)
    parser.add_argument("--decay", type=float, default=0.)
    parser.add_argument("--epochs", type=int, default=100)
    parser.add_argument("--cp_every", type=int, default=20)
    parser.add_argument("--verbose", type=int, default=10)
    parser.add_argument("--save", type=bool_, default=True)
    parser.add_argument("--seed", type=int, default=0)
    arglist = parser.parse_args()
    arglist = vars(arglist)
    return arglist

def main(arglist):
    np.random.seed(arglist["seed"])
    torch.manual_seed(arglist["seed"])
    print(f"learning with args: {arglist}\n")

    # load data
    with open(os.path.join(arglist["data_path"], "data.p"), "rb") as f:
        data = pickle.load(f)

    # parse data
    data = [
        {
            k: torch.from_numpy(v).to(torch.float32)
            for k, v in d.items() if k not in ["drive_id", "t"]
        } for d in data
    ]
    data, mask = pad_collate(data)

    # process data
    nan_mask = mask.clone()
    nan_mask[mask == 0] = torch.nan
    obs_np = (nan_mask.unsqueeze(-1) * data["obs"]).numpy()

    obs_mean = np.nanmean(obs_np, axis=(0, 1))
    obs_std = np.nanstd(obs_np, axis=(0, 1))

    obs = (data["obs"] - obs_mean) / obs_std
    act = data["act"].long()

    data = (obs, act, mask)

    # compute kmeans obs initialization
    kmeans_data = obs.flatten(0, 1).numpy()
    kmeans_data = kmeans_data[mask.flatten() == 1]
    kmeans = KMeans(n_clusters=arglist["state_dim"]).fit(kmeans_data)
    kmeans_centers = torch.from_numpy(kmeans.cluster_centers_).to(torch.float32)

    # init model
    batch_size = obs.shape[1]
    act_dim = len(torch.unique(act))
    obs_dim = obs.shape[-1]
    model = EBIRL(
        arglist["state_dim"], act_dim, obs_dim, arglist["horizon"],
        prior_cov=arglist["prior_cov"], bc_penalty=arglist["bc_penalty"],
        obs_penalty=arglist["obs_penalty"], kl_penalty=arglist["kl_penalty"]
    )
    model.init_params(kmeans_centers)
    model.init_q(batch_size, freeze_prior=False)
    print(model)

    # init loss function
    loss_fn = lambda m, o, a, mask: m.compute_loss(o, a, mask)
    optimizer = torch.optim.Adam(model.parameters(), lr=arglist["lr"], weight_decay=arglist["decay"])

    # load pretrained prior
    cp_history = None
    if arglist["cp_path"] != "none":
        cp_path = os.path.join(arglist["exp_path"], arglist["cp_path"])
        cp_model_path = glob.glob(os.path.join(cp_path, "models/*.pt"))
        cp_model_path.sort(key=lambda x: int(os.path.basename(x).replace(".pt", "").split("_")[-1]))

        state_dict = torch.load(cp_model_path[-1], map_location="cpu")
        model.load_state_dict(state_dict["model_state_dict"], strict=False)
        optimizer.load_state_dict(state_dict["optimizer_state_dict"])

        # load history
        cp_history = pd.read_csv(os.path.join(cp_path, "history.csv"))
        print(f"loaded checkpoint from {cp_path}")

    callback = None
    if arglist["save"]:
        if not os.path.exists(arglist["exp_path"]):
            os.mkdir(arglist["exp_path"])
        callback = Logger(arglist, model.plot_keys, cp_history)

    # train loop
    model, optimizer, df_history, callback = train(
        model, data, loss_fn, optimizer, arglist["epochs"],
        verbose=arglist["verbose"], callback=callback, grad_check=False
    )

    if arglist["save"]:
        callback.save_checkpoint(model, optimizer)
        callback.save_history(df_history)
    return model, df_history

if __name__ == "__main__":
    arglist = parse_args()
    main(arglist)


## PyTorch
import torch
import torch.nn as nn
import torch.optim as optim

import numpy as np
import random
import math
import json
from functools import partial
import torch.nn.functional as F

# PyTorch Lightning
import pytorch_lightning as pl
from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint

class CosineWarmupScheduler(optim.lr_scheduler._LRScheduler):

    def __init__(self, optimizer, warmup, max_iters):
        self.warmup = warmup
        self.max_num_iters = max_iters
        super().__init__(optimizer)

    def get_lr(self):
        lr_factor = self.get_lr_factor(epoch=self.last_epoch)
        return [base_lr * lr_factor for base_lr in self.base_lrs]

    def get_lr_factor(self, epoch):
        lr_factor = 0.5 * (1 + np.cos(np.pi * epoch / self.max_num_iters) )
        if epoch <= self.warmup:
            lr_factor *= epoch * 1.0 / self.warmup
        return lr_factor

def scaled_dot_product(q, k, v, mask=None):
    d_k = q.size()[-1]
    attn_logits = torch.matmul(q, k.transpose(-2, -1))
    attn_logits = attn_logits / math.sqrt(d_k)
    if mask is not None:
        nHeads=attn_logits.shape[-3]
        mask=mask.unsqueeze(dim=-3)
        mask=torch.repeat_interleave(mask,nHeads,dim=-3)
        attn_logits = attn_logits.masked_fill(mask == 0, -9e15)
    attention = F.softmax(attn_logits, dim=-1)
    values = torch.matmul(attention, v)
    return values, attention

class MultiheadAttention(nn.Module):
    """
    Applying a Multi-Head Attention layer in a neural network, where we don’t have an arbitrary query, key, and value vector.
       A simple but effective implementation is to set the current feature map in a NN, , as , and ( being the batch size,
       the sequence length, the hidden dimensionality of ). The consecutive weight matrices ,
       and can transform to the corresponding feature vectors that represent the queries, keys, and values of the input.
    """
    def __init__(self, input_dim, embed_dim, num_heads):
        super().__init__()
        assert embed_dim % num_heads == 0, "Embedding dimension must be 0 modulo number of heads. " +str(embed_dim)+"%"+str(num_heads)

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        # Stack all weight matrices 1...h together for efficiency
        # Note that in many implementations you see "bias=False" which is optional
        self.qkv_proj = nn.Linear(input_dim, 3*embed_dim)
        self.o_proj = nn.Linear(embed_dim, embed_dim)

        self._reset_parameters()

    def _reset_parameters(self):
        # Original Transformer initialization, see PyTorch documentation
        nn.init.xavier_uniform_(self.qkv_proj.weight)
        self.qkv_proj.bias.data.fill_(0)
        nn.init.xavier_uniform_(self.o_proj.weight)
        self.o_proj.bias.data.fill_(0)

    def forward(self, x, mask=None, return_attention=False):
        batch_size, seq_length, _ = x.size()
        qkv = self.qkv_proj(x)

        # Separate Q, K, V from linear output
        qkv = qkv.reshape(batch_size, seq_length, self.num_heads, 3*self.head_dim)
        qkv = qkv.permute(0, 2, 1, 3) # [ Batch, Head, SeqLen, Dims ]
        q, k, v = qkv.chunk(3, dim=-1)

        # Determine value outputs
        values, attention = scaled_dot_product(q, k, v, mask=mask)
        values = values.permute(0, 2, 1, 3) # [Batch, SeqLen, Head, Dims]
        values = values.reshape(batch_size, seq_length, self.embed_dim)
        o = self.o_proj(values)

        if return_attention:
            return o, attention
        else:
            return o

class EncoderBlock(nn.Module):

    def __init__(self, input_dim, num_heads, dim_feedforward, dropout=0.0):
        """
        Inputs:
            input_dim - Dimensionality of the input
            num_heads - Number of heads to use in the attention block
            dim_feedforward - Dimensionality of the hidden layer in the MLP
            dropout - Dropout probability to use in the dropout layers
        """
        super().__init__()

        # Attention layer
        self.self_attn = MultiheadAttention(input_dim, input_dim, num_heads)

        # Two-layer MLP
        self.linear_net = nn.Sequential(
            nn.Linear(input_dim, dim_feedforward),
            nn.Dropout(dropout),
            nn.ReLU(inplace=True),
            nn.Linear(dim_feedforward, input_dim)
        )

        # Layers to apply in between the main layers
        self.norm1 = nn.LayerNorm(input_dim)
        self.norm2 = nn.LayerNorm(input_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Attention part
        attn_out = self.self_attn(x, mask=mask)
        x = x + self.dropout(attn_out)
        x = self.norm1(x)

        # MLP part
        linear_out = self.linear_net(x)
        x = x + self.dropout(linear_out)
        x = self.norm2(x)

        return x

class TransformerEncoder(nn.Module):
    def __init__(self, num_layers, **block_args):
        super().__init__()
        self.layers = nn.ModuleList([EncoderBlock(**block_args) for _ in range(num_layers)])

    def forward(self, x, mask=None):
        for l in self.layers:
            x = l(x, mask=mask)
        return x

    def get_attention_maps(self, x, mask=None):
        attention_maps = []
        for l in self.layers:
            _, attn_map = l.self_attn(x, mask=mask, return_attention=True)
            attention_maps.append(attn_map)
            x = l(x)
        return attention_maps

class PositionalEncoding(nn.Module):

    def __init__(self, d_model, max_len=5000):
        """
        Inputs
            d_model - Hidden dimensionality of the input.
            max_len - Maximum length of a sequence to expect.
        """
        super().__init__()

        # Create matrix of [SeqLen, HiddenDim] representing the positional encoding for max_len inputs
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)

        # register_buffer => Tensor which is not a parameter, but should be part of the modules state.
        # Used for tensors that need to be on the same device as the module.
        # persistent=False tells PyTorch to not add the buffer to the state dict (e.g. when we save the model)
        self.register_buffer('pe', pe, persistent=False)

    def forward(self, x):
        x = x + self.pe[:, :x.size(1)]
        return x

class TransformerPredictor(pl.LightningModule):

    def __init__(self, input_dim, model_dim, num_classes, num_heads,
                 num_layers, lr, warmup, max_iters, dropout=0.0, input_dropout=0.0):
        """
        Inputs:
            input_dim - Hidden dimensionality of the input
            model_dim - Hidden dimensionality to use inside the Transformer
            num_classes - Number of classes to predict per sequence element
            num_heads - Number of heads to use in the Multi-Head Attention blocks
            num_layers - Number of encoder blocks to use.
            lr - Learning rate in the optimizer
            warmup - Number of warmup steps. Usually between 50 and 500
            max_iters - Number of maximum iterations the model is trained for. This is needed for the CosineWarmup scheduler
            dropout - Dropout to apply inside the model
            input_dropout - Dropout to apply on the input features
        """
        super().__init__()
        self.save_hyperparameters()
        self._create_model()

    def _create_model(self):
        # Input dim -> Model dim
        self.input_net = nn.Sequential(
            nn.Dropout(self.hparams.input_dropout),
            nn.Linear(self.hparams.input_dim, self.hparams.model_dim)
        )
        # Positional encoding for sequences
        self.positional_encoding = PositionalEncoding(d_model=self.hparams.model_dim)
        # Transformer
        self.transformer = TransformerEncoder(num_layers=self.hparams.num_layers,
                                              input_dim=self.hparams.model_dim,
                                              dim_feedforward=2*self.hparams.model_dim,
                                              num_heads=self.hparams.num_heads,
                                              dropout=self.hparams.dropout)
        # Output classifier per sequence element
        self.output_net = nn.Sequential(
            nn.Linear(self.hparams.model_dim, self.hparams.model_dim),
            nn.LayerNorm(self.hparams.model_dim),
            nn.ReLU(inplace=True),
            nn.Dropout(self.hparams.dropout),
            nn.Linear(self.hparams.model_dim, self.hparams.num_classes)
        )

    def forward(self, x, mask=None, add_positional_encoding=True):
        """
        Inputs:
            x - Input features of shape [Batch, SeqLen, input_dim]
            mask - Mask to apply on the attention outputs (optional)
            add_positional_encoding - If True, we add the positional encoding to the input.
                                      Might not be desired for some tasks.
        """
        x = self.input_net(x)
        if add_positional_encoding:
            x = self.positional_encoding(x)
        x = self.transformer(x, mask=mask)
        x = self.output_net(x)
        return x

    @torch.no_grad()
    def get_attention_maps(self, x, mask=None, add_positional_encoding=True):
        """
        Function for extracting the attention matrices of the whole Transformer for a single batch.
        Input arguments same as the forward pass.
        """
        x = self.input_net(x)
        if add_positional_encoding:
            x = self.positional_encoding(x)
        attention_maps = self.transformer.get_attention_maps(x, mask=mask)
        return attention_maps

    def configure_optimizers(self):
        optimizer = optim.Adam(self.parameters(), lr=self.hparams.lr)

        # Apply lr scheduler per step
        lr_scheduler = CosineWarmupScheduler(optimizer,
                                             warmup=self.hparams.warmup,
                                             max_iters=self.hparams.max_iters)
        return [optimizer], [{'scheduler': lr_scheduler, 'interval': 'step'}]

    def training_step(self, batch, batch_idx):
        raise NotImplementedError

    def validation_step(self, batch, batch_idx):
        raise NotImplementedError

    def test_step(self, batch, batch_idx):
        raise NotImplementedError


from flask import Flask, redirect, url_for, request,render_template
import pandas as pd
import json
app = Flask(__name__)
@app.route('/GetDataPandaCSV',methods = ['GET'])
def GetCSVDataPandas():
       data = request.args.get('data')
       print(data)
       lst = []
       test = pd.read_csv('word_search (3).csv')
       res = test[test['name'].str.match(data,na=False)]
       for index, row in res.iterrows():
           d = {}
           d[data] = row['frequency']
           lst.append(d)
       return json.dumps(lst)
if __name__ == '__main__':
   app.run("127.0.0.1","5005",True)