import subprocess, sys, os, argparse, yaml

# Initialize argparse
parser = argparse.ArgumentParser(description='Updater and installer for common packages used in many of my programs')
parser.add_argument('-i', type=str, help='Install package', metavar='package')
parser.add_argument('-u', action='store_true', help='Update all packages')
parser.add_argument('-r', type=str, help='Remove package', metavar='package')

args = parser.parse_args()
packageToInstall = args.i 	# str
doUpdate = bool(args.u)		# bool
packageToRemove = args.r 	# str

"""
# Required positional argument
parser.add_argument('pos_arg', type=int,
                    help='A required integer positional argument')

# Optional positional argument
parser.add_argument('opt_pos_arg', type=int, nargs='?',
                    help='An optional integer positional argument')

# Optional argument
parser.add_argument('--opt_arg', type=int,
                    help='An optional integer argument')

# Switch
parser.add_argument('--switch', action='store_true',
                    help='A boolean switch')
"""

# get os type
nt = (True if os.name == 'nt' else 'unix')

def fetchVersions():
	sys.stdout.flush()
	sys.stdout.write('Fetching versions file from GitHub...')
	# fetch versions file
	subprocess.call(
		['wget', 'https://raw.githubusercontent.com/Mase3206/py-common-modules/main/versions.yml'],
		stdout=subprocess.DEVNULL,
		stderr=subprocess.DEVNULL
	)
	sys.stdout.write('\rFetching versions file from GitHub... done.\n')
	sys.stdout.flush()

	# open versions file, store to dict, then close
	with open('versions.yml', 'r') as f2:
		versions = yaml.safe_load(f2)
	# remove versions file
	os.remove('versions.yml')

	try:
		with open('installed.yml', 'r') as f1:
			installedVersions = yaml.safe_load(f1)
	except:
		f1 = open('installed.yml', 'w+')
		f1.close()

	if installedVersions == None:
		installedVersions = {}

	return versions, installedVersions

def saveInstalledVersions(newInstalledVersions):
	with open('installed.yml', 'w') as f3:
		yaml.safe_dump(newInstalledVersions, f3)

def update(module='all', quiet=False):
	toUpdate = []

	versions, installedVersions = fetchVersions()

	if module == 'all':
		# make a list of all out-of-date modules
		for mod in versions:
			if mod in list(installedVersions.keys()):
				if installedVersions[mod] < versions[mod]:
					toUpdate.append(mod)

	else:
		# if a specific module is passed, add only that one to toUpdate
		mod = module
		if installedVersions[mod] < versions[mod]:
			toUpdate.append(mod)

	if toUpdate == []:
		print('All modules are already up to date.') if not quiet else None
	else:
		print('Modules to update:') if quiet else None

		for mod in toUpdate:
			print(f'{mod}: {installedVersions[mod]} => {versions[mod]}')

		proceed = input('Proceed with update? [Y/n] ') if not quiet else 'y'
		if proceed == '' or proceed.lower() == 'y':
			# delete old modules
			for mod in toUpdate:
				print(f'Removing {mod}...') if not quiet else None
				os.remove(mod + '.py')
			# fetch updated modules
			for mod in toUpdate:
				print(f'Downloading and installing {mod}...') if not quiet else None
				subprocess.call(
					['wget', 'https://raw.githubusercontent.com/Mase3206/py-common-modules/main/' + mod + '.py'],
					stdout=subprocess.DEVNULL,
					stderr=subprocess.DEVNULL
				)
				installedVersions[mod] = versions[mod]
			saveInstalledVersions(installedVersions)
			print('Update successful.') if not quiet else None
		else:
			print('Abort.') if not quiet else None

def install(module:str, quiet=False):
	versions, installedVersions = fetchVersions()
	installed = None

	# test if module is already installed. if so, set installed to True
	try:
		installedVersions[module]
		installed = True
	except:
		installed = False

	if installed == True:
		proceed = input('Module already installed. Would you like to try to update it? [Y/n] ') if not quiet else 'y'
		if proceed == '' or proceed.lower() == 'y':
			update(module)
		else:
			print(f'{module} not sent to updates. Exiting.') if not quiet else None
		del proceed
	else:
		try:
			proceed = input(f'Found {module} in GitHub repo. Would you like to install it? [Y/n] ') if not quiet else 'y'
			if proceed == '' or proceed.lower() == 'y':
				if not quiet:
					sys.stdout.flush()
					sys.stdout.write(f'Downloading {module}...')
				subprocess.call(
					['wget', 'https://raw.githubusercontent.com/Mase3206/py-common-modules/main/' + module + '.py'],
					stdout=subprocess.DEVNULL,
					stderr=subprocess.DEVNULL
				)
				del proceed
			else:
				exit('Abort.') if not quiet else None
				del proceed
		except:
			print(f'Fetching {module} failed. Verify that it exists.') if not quiet else None		#TODO: maybe list the installable modules?

		if not quiet:
			sys.stdout.write(f'\rDownloading {module}... done.\n')
			sys.stdout.write(f'Adding {module} to installed.yml...')
		installedVersions[module] = versions[module]
		saveInstalledVersions(installedVersions)
		if not quiet:
			sys.stdout.write(f'\rAdding {module} to installed.yml... done.\n')
			sys.stdout.flush()
			print(f'Successfully installed {module}.')

def remove(module:str):
	versions, installedVersions = fetchVersions()
	installed = None

	# test if module is already installed. if so, set installed to True
	try:
		installedVersions[module]
		installed = True
	except:
		installed = False

	if installed == False:
		print(f'{module} is not installed. Exiting.')
	elif installed == True:
		sys.stdout.flush()
		sys.stdout.write(f'Removing {module}...')
		try:
			os.remove(module + '.py')
			sys.stdout.write(f'\rRemoving {module}... done.\n')
		except:
			print(f'{module} not installed, but was found in installed.yml. Removing...')
		finally:
			sys.stdout.write(f'Removing {module} from installed.yml...')
			del installedVersions[module]
			saveInstalledVersions(installedVersions)
			sys.stdout.write(f'\rRemoving {module} from installed.yml... done.\n')
	else:
		raise RuntimeError('installed flag in remove() was not set to True or False.')

def installMultiple(modules:list, quiet=True):
	versions, installedVersions = fetchVersions()
	toInstall = []

	for module in modules:
		try:
			installedVersions[module]
			update(module)
		except:
			toInstall.append(module)

	# This function is meant to be called indirectly, so it defaults to quiet operation.
	# In fact, it can't even be loud, at the moment.

	for module in toInstall:
		subprocess.call(
			['wget', 'https://raw.githubusercontent.com/Mase3206/py-common-modules/main/' + module + '.py'],
			stdout=subprocess.DEVNULL,
			stderr=subprocess.DEVNULL
		)
		installedVersions[module] = versions[module]
		saveInstalledVersions(installedVersions)

if doUpdate == True:
	update()

if packageToInstall != None:
	install(packageToInstall)

if packageToRemove != None:
	remove(packageToRemove)

while True:
	try:
		os.remove('wget-log*')
	except:
		break


#!/usr/bin/python3

# Python code to encrypt, compress, and base64 encode shellcode from a file

import argparse
import gzip
from base64 import b64encode
import os

def xor(data, key):
    # xor encrypt shellcode
    output_str = ""
    for i in range(len(data)):
        current = data[i]
        current_key = key[i%len(key)]
        output_str += chr((current) ^ ord(current_key))

    return bytearray(output_str, 'latin-1')

def getShellcode(fileName):
    # Open shellcode file and read bytes
    try:
        with open(fileName, 'rb') as shellcodeFileHandle:
            shellcodeBytes = shellcodeFileHandle.read()
            shellcodeFileHandle.close()
        print("--- Shellcode file -{}- successfully loaded".format(fileName))
    except IOError:
        print("!!! Could not open and read -{}-".format(fileName))
        quit()

    print("--- Shellcode size: {} bytes".format(len(shellcodeBytes)))

    return shellcodeBytes

def encryptShellcode(shellcodeBytes, key):
    # XOR Encrypt Shellcode
    encryptedShellcode = xor(shellcodeBytes, key)
    print("--- Encrypted shellcode size: {} bytes".format(len(encryptedShellcode)))
    try:
        f = open("reverseencrypted.raw", "wb+")
        f.write(encryptedShellcode)
        f.close()
    except:
        print("!!! Unable to write encrypted shellcode to file.")
        quit()

    return encryptedShellcode

def compressShellcode(encryptedShellcode):
    # Use gzip to compress shellcode
    try:
        compressedShellcode = gzip.compress(encryptedShellcode)
    except:
        print("!!! Unable to compress shellcode.")
        quit()

    return compressedShellcode

def writeEncodedToFile(compressedShellcode, outfile):
    # Base64 encode the shellcode and write out to file
    try:
        encodedShellcode = b64encode(bytearray(compressedShellcode))
    except:
        print("!!! Unable to encode shellcode.")
        quit()

    try:
        f = open(outfile, "wb+")
        numChars = f.write(encodedShellcode)
        f.close()
        print("--- Encoded shellcode written to {}.".format(outfile))
    except IOError:
        print("!!! Could not write encoded shellcode to file.")

if __name__ == '__main__':
    # Parse arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("shellcodeFileName", help="File name of the raw shellcode to be encrypted/compressed/encoded")
    parser.add_argument("key", help="Key to XOR encrypt the shellcode")
    parser.add_argument("outfile", help="Name of file where encrypted/compressed/encoded shellcode will be output")
    args = parser.parse_args()

    shellcodeBytes = getShellcode(args.shellcodeFileName)
    encryptedShellcode = encryptShellcode(shellcodeBytes, args.key)
    compressedShellcode = compressShellcode(encryptedShellcode)
    writeEncodedToFile(compressedShellcode, args.outfile)


import torch
from torch import nn
from torch.nn import functional as F
import zounds
from config.experiment import Experiment
from modules.ddsp import overlap_add
from modules.dilated import DilatedStack
from modules.fft import fft_convolve
from modules.linear import LinearOutputStack
from modules.matchingpursuit import SparseCodingLoss
from modules.normalization import max_norm
from modules.reverb import ReverbGenerator
from modules.softmax import hard_softmax, sparse_softmax
from modules.sparse import sparsify, sparsify_vectors

from modules.stft import morlet_filter_bank
from modules.transfer import ImpulseGenerator, PosEncodedImpulseGenerator
from time_distance import optimizer
from train.experiment_runner import BaseExperimentRunner
from util import device
from util.readmedocs import readme

exp = Experiment(
    samplerate=zounds.SR22050(),
    n_samples=2**15,
    weight_init=0.1,
    model_dim=128,
    kernel_size=512)

class Impulse(nn.Module):
    def __init__(self, n_frames, n_samples, n_impulses):
        super().__init__()
        self.n_frames = n_frames
        self.n_samples = n_samples
        self.n_impulses = n_impulses

        self.impulses = nn.Parameter(torch.zeros(
            n_impulses, n_frames).uniform_(0, 1))

    def forward(self, x):
        x = x.view(-1, self.n_impulses)
        x = x @ (self.impulses ** 2)
        x = x.view(-1, 1, self.n_frames)
        x = F.interpolate(x, size=self.n_samples, mode='linear')
        noise = torch.zeros_like(x).uniform_(-1, 1)
        x = x * noise
        return x

class Bandpass(nn.Module):
    def __init__(self, filter_size, n_filters):
        super().__init__()
        self.filter_size = filter_size
        self.n_filters = n_filters
        self.filters = nn.Parameter(torch.zeros(self.n_filters, self.filter_size).uniform_(-1, 1))

    def forward(self, x, signal):
        batch = x.shape[0]
        signal = signal.view(1, batch, signal.shape[-1])
        x = x.view(-1, self.n_filters)
        x = x @ self.filters
        x = x.view(batch, 1, self.filter_size)
        x = x * torch.hamming_window(self.filter_size, device=x.device)[None, None, :]
        x = F.conv1d(signal, x, padding=self.filter_size // 2, groups=batch)
        x = x.view(batch, 1, -1)
        return x

class BandLimitedNoise(nn.Module):
    def __init__(self, n_frames, n_samples, n_impulses, filter_size, n_filters):
        super().__init__()
        self.n_samples = n_samples
        self.impulses = Impulse(n_frames, n_samples, n_impulses)
        self.bandpass = Bandpass(filter_size, n_filters)

    def forward(self, impulse_choice, filter_choice):
        impulse = self.impulses.forward(impulse_choice)
        filtered = self.bandpass.forward(filter_choice, impulse)
        return filtered[..., :self.n_samples]

class Resonance(nn.Module):
    def __init__(self, n_samples, n_frequencies, samplerate: zounds.SampleRate):
        super().__init__()
        self.n_samples = n_samples
        self.n_frequencies = n_frequencies

        band = zounds.FrequencyBand(20, samplerate.nyquist)
        scale = zounds.MelScale(band, n_frequencies)
        bank = morlet_filter_bank(
            samplerate, n_samples, scale, 0.01, normalize=False)

        bank = torch.from_numpy(bank.real).float().view(n_frequencies, n_samples)
        bank = max_norm(bank, dim=-1)

        self.register_buffer('bank', bank)

        # self.bank = nn.Parameter(torch.zeros(n_frequencies, n_samples).uniform_(-1, 1))

    def forward(self, balance: torch.Tensor, decay: torch.Tensor):
        balance = balance.view(-1, self.n_frequencies)
        decay = decay.view(-1, 1)
        decay = torch.clamp(decay, 0, 1)

        x = balance @ self.bank
        x = x.view(-1, 1, self.n_samples)
        x = F.pad(x, (0, 256))
        windowed = \
            x.unfold(-1, size=512, step=256) * torch.hamming_window(512,
                                                                   device=x.device)[None, None, None, :]
        windowed = max_norm(windowed, dim=-1)

        decay = decay.repeat(1, windowed.shape[-2])
        decay = torch.exp(torch.cumsum(torch.log(decay + 1e-8), dim=-1))

        x = windowed * decay[:, None, :, None]

        x = overlap_add(x, apply_window=False)
        x = x[..., :self.n_samples]
        return x

class Event(nn.Module):
    def __init__(
            self,
            n_samples,
            n_frequencies,
            samplerate,
            n_frames,
            impulse_samples,
            n_impulses,
            filter_size,
            n_filters):

        super().__init__()
        self.n_samples = n_samples
        self.impulse_samples = impulse_samples
        self.resonance = Resonance(
            n_samples, n_frequencies, samplerate)
        self.impulse = BandLimitedNoise(
            n_frames, impulse_samples, n_impulses, filter_size, n_filters)

    def forward(self, balance, decay, impulse_choice, filter_choice, amp):
        batch = balance.shape[0]

        res = self.resonance.forward(balance, decay)
        impulse = self.impulse.forward(impulse_choice, filter_choice)

        res = res.view(1, batch, self.n_samples)
        impulse = impulse.view(batch, 1, self.impulse_samples)

        x = F.conv1d(res, impulse, padding=self.impulse_samples // 2, groups=batch)[..., :self.n_samples]
        x = x.view(-1, 1, self.n_samples)
        imp = F.pad(impulse.reshape(x.shape[0], 1, -1), (0, self.n_samples - self.impulse_samples))
        imp = imp.view(-1, 1, self.n_samples)
        x = x + imp
        x = x * amp.view(-1, 1, 1)
        x = x.view(-1, 1, self.n_samples)
        return x

class Model(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.channels = channels
        self.encoder = DilatedStack(
            channels,
            [1, 3, 9, 27, 81, 1],
            dropout=0.1,
            soft_sparsity=False,
            internally_sparse=False,
            sparsity_amt=1)

        self.event = Event(
            exp.n_samples,
            n_frequencies=512,
            samplerate=exp.samplerate,
            n_frames=128,
            impulse_samples=2048,
            n_impulses=512,
            filter_size=32,
            n_filters=512)

        self.to_balance = LinearOutputStack(channels, 3, out_channels=512)
        self.to_decay = LinearOutputStack(channels, 3, out_channels=1)
        self.to_impulse = LinearOutputStack(channels, 3, out_channels=512)
        self.to_filter = LinearOutputStack(channels, 3, out_channels=512)
        self.to_amp = LinearOutputStack(channels, 3, out_channels=1)

        self.to_time = LinearOutputStack(channels, 3, out_channels=33)

        self.to_pos = PosEncodedImpulseGenerator(
            exp.n_samples, exp.n_samples, softmax=lambda x: hard_softmax(x, invert=True, tau=0.01))
        # self.to_pos = ImpulseGenerator(exp.n_samples, softmax=lambda x: hard_softmax(x, invert=True, tau=0.01))

        self.attn = nn.Conv1d(channels, 1, 1, 1, 0)

        self.verb = ReverbGenerator(channels, 3, exp.samplerate, exp.n_samples)
        self.apply(lambda p: exp.init_weights(p))

    def forward(self, x, debug=False):
        batch = x.shape[0]
        x = exp.pooled_filter_bank(x)
        x = x[..., :exp.n_samples]
        encoded = self.encoder(x)

        context = torch.mean(encoded, dim=-1)
        attn = self.attn.forward(encoded)
        events, indices = sparsify_vectors(encoded, attn, n_to_keep=16)

        b = torch.relu(self.to_balance.forward(events))
        d = torch.sigmoid(self.to_decay.forward(events))
        i = hard_softmax(self.to_impulse.forward(events), invert=True, tau=0.01)
        f = hard_softmax(self.to_filter.forward(events), invert=True, tau=0.01)
        a = self.to_amp.forward(events) ** 2
        t = self.to_time.forward(events)

        events = self.event.forward(
            b.view(-1, 512),
            d.view(-1, 1),
            i.view(-1, 512),
            f.view(-1, 512),
            a.view(-1, 1))

        impulses, _ = self.to_pos.forward(t.view(batch * 16, -1))

        impulses = impulses.view(batch, 16, exp.n_samples)
        events = events.view(batch, 16, exp.n_samples)
        final = fft_convolve(events, impulses)
        final = torch.sum(final, dim=1, keepdim=True)

        final = self.verb.forward(context, final)
        return final

model = Model(exp.model_dim).to(device)

optim = optimizer(model, lr=1e-4)

loss_module = SparseCodingLoss(
    512,
    512,
    128,
    approx=None,
    learning_steps=256,
    device=device,
    pooling=(512, 256))

def train(batch, i):
    optim.zero_grad()
    recon = model.forward(batch, debug=False)
    loss = loss_module.loss(recon, batch)
    # loss = exp.perceptual_loss(recon, batch)
    loss.backward()
    optim.step()
    return loss, recon

@readme
class PhysicalModel(BaseExperimentRunner):
    def __init__(self, stream):
        super().__init__(stream, train, exp)
        self.real = None
        self.fake = None


import logging
from typing import Optional

from src.core.entities import Profile, Shelf, ShelfItem
from src.core.exceptions import ProfileNotFoundError
from src.core.gateways import IDataGateway
from src.core.repositories import IDataRepository

logger = logging.getLogger(__name__)

class RefreshShelfItemsUseCase:
    def __init__(self, gateway: IDataGateway, repository: IDataRepository) -> None:
        self._repository = repository
        self._gateway = gateway

    async def execute(self, shelf: Shelf) -> None:
        logger.info(f"Refreshing items on shelf {shelf.name}")

        async with self._repository.context():
            profile: Optional[Profile] = await self._repository.profile.read(
                uuid=shelf.profile_uuid
            )

            if not profile:
                raise ProfileNotFoundError(f"Profile {shelf.profile_uuid} not found.")

        async with self._gateway.context():
            remote_shelf_items: list[ShelfItem] = []
            fetch_results = self._gateway.shelf_item.fetch_for_profile_and_shelf(
                profile=profile,
                shelf=shelf,
            )
            async for shelf_item in fetch_results:
                if not shelf_item.book_editions:
                    book_editions = await self._gateway.book_edition.fetch_for_shelf_item(
                        shelf_item=shelf_item
                    )
                    if book_editions:
                        shelf_item.book_editions = book_editions

                if shelf_item.book_editions:
                    logger.info(
                        f"Found {len(shelf_item.book_editions)} editions"
                        f" for '{shelf_item.title}' by {shelf_item.author}"
                    )

                remote_shelf_items.append(shelf_item)

        if not remote_shelf_items:
            logger.info(f"No shelf items found on shelf {shelf.name}")
            return

        remote_shelf_items_len = len(remote_shelf_items)
        logger.info(f"Found {remote_shelf_items_len} shelf items on shelf {shelf.name}")

        remote_book_editions_len = sum(
            len(shelf_item.book_editions) for shelf_item in remote_shelf_items
        )
        logger.info(f"Found {remote_book_editions_len} book editions on shelf {shelf.name}")

        async with self._repository.context():
            await self._repository.shelf_item.sync_on_shelf(
                shelf=shelf,
                shelf_items=remote_shelf_items,
            )

            remote_book_editions = [
                book_edition
                for shelf_item in remote_shelf_items
                for book_edition in shelf_item.book_editions
            ]
            await self._repository.book_edition.sync_on_shelf(
                shelf=shelf,
                book_editions=remote_book_editions,
            )

        logger.info(f"Synchronized items on shelf {shelf.name}")


import pandas as pd
import numpy as np
import json

def read_from_file(filename, is_json=False):
    with open(filename, "r") as f:
        if is_json:
            return json.load(f)
        else:
            return f.read()

def write_to_file(filename, data, is_json=False):
    with open(filename, "w") as f:
        if is_json:
            json.dump(data, f)
        else:
            f.write(str(data))

df = pd.read_csv("../data/Timeouts-3.csv")
pbp = pd.read_csv("../data/New_Play_by_Play_data_2.csv")

temp = []
i = 0
for index, row in df.iterrows():
    if i >= len(temp):
        game_id = row["Game_id"]
        team_id = row["Team_id"]
        index1 = row["index"]

        cur_point_diff = abs(int(pbp[(pbp["index"] == index1)]["Point_diff"]))

        largest_point_diff = abs(pbp[pbp["Game_id"] == game_id]["Point_diff"].max())

        temp.append(cur_point_diff - largest_point_diff)

    i += 1
    print(len(temp))

df["Lead_diff"] = temp

df.to_csv("../data/Timeouts-3.csv")


from pokemons import pokemons
from time import sleep
from class_pokemon import Pokemon
from random import choice
import inquirer

collors = {
    'X': '\33[32m',
    'Y': '\33[4;35m',
    'Z': '\33[31m',
    'N': '\33[7;31m',
    'C': '\33[m'
}

def select_poke():
    pokemon = [
                inquirer.List('choice',
                                message="Select:",
                                choices=[pokemons[0].name, pokemons[1].name, pokemons[2].name]
                                ),
                ]
    user_choice = inquirer.prompt(pokemon)

    if user_choice['choice'] == "Blastoise":
        return pokemons[0]
    elif user_choice['choice'] == "Charizard":
        return pokemons[1]
    elif user_choice['choice'] == "Venusaur":
        return pokemons[2]

def play():
    print(f"{collors['N']}Welcome to Pokemon Game{collors['C']}")
    print()
    poke_user = select_poke()
    sleep(2)
    print("You choose your Pokemon, now it's pc's turn")
    poke_pc = choice(pokemons)
    sleep(3)
    print()
    print("Both chose, now it's battle time")
    print(f'''
    {collors['X']}Your pokemon traits:{collors['C']}
    Name = {poke_user.name}
    Life = {poke_user.life}
    Attack = {poke_user.attack}
    Speed = {poke_user.speed}''')
    print()
    print(f'''
    {collors['Z']}Pc's pokemon traits:{collors['C']}
    Name = {poke_pc.name}
    Life = {poke_pc.life}
    Attack = {poke_pc.attack}
    Speed = {poke_pc.speed}''')
    sleep(2)
    print()

    Pokemon.attacking(poke_user, poke_pc)


from torch.optim import Adam
import torch
from u_net import Unet
from train import train_model
from prepare_data import BitmogiDataset
from torchvision.transforms import Compose
from torchvision import transforms
from torch.utils.data import DataLoader
from pathlib import Path

#will go to config
image_size = 128
channels   = 1
epochs     = 1
data_path  = "/Users/sbhar/Riju/TorchDemoCode/TrinaCode/bitmojis/"
batch_size = 6
results_folder = Path("./results")
results_folder.mkdir(exist_ok = True)

def main():
    if not torch.backends.mps.is_available():
        if not torch.backends.mps.is_built():
            print("MPS not available because the current PyTorch install was not "
              "built with MPS enabled.")
            device = torch.device("cpu")
        else:
            print("MPS not available because the current MacOS version is not 12.3+ "
              "and/or you do not have an MPS-enabled device on this machine.")
            device = torch.device("cpu")

    else:
        device = torch.device("mps")

    model = Unet(
        dim=image_size,
        channels=channels,
        dim_mults=(1, 2, 4,)
    )

    model.to(device)
    optimizer = Adam(model.parameters(), lr=1e-3)

    # define image transformations (e.g. using torchvision)
    transform_ = Compose([
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Lambda(lambda t: (t * 2) - 1)
        ])

    bitmoji_dataset = BitmogiDataset(data_dir=data_path,transform=transform_)
    dataloader = DataLoader(bitmoji_dataset, batch_size=batch_size, shuffle=True)

    train_model(epochs=epochs,model=model,optimizer=optimizer,device=device,results_folder=results_folder,dataloader=dataloader)

if __name__ == "__main__":
    main()


import os
import pytest
from unittest import mock

def get_filenames_in_directory(directory):
    fileNames = []
    for root, dirs, files in os.walk(directory):
        for name in files:
            fileNames.append(name)
    return fileNames

# https://docs.pytest.org/en/latest/parametrize.html
# https://docs.python.org/3/library/unittest.mock.html
@mock.patch('os.walk')
@pytest.mark.parametrize('files,expected_count', [
  (['a.txt', 'b.txt', 'c.png'], 3),
  (['a.txt', 'c.png'], 2),
  ([], 0),
  (['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'], 10)
])
def test_with_mocked_logic_and_test_cases(mockwalk, files, expected_count):
    folders = ['folder1', 'folder2']
    mockwalk.return_value = [('/root', folders, files)]

    result = get_filenames_in_directory('path')
    assert len(result) == expected_count
    assert result == files
    mockwalk.assert_called_once_with('path')


import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm

from ..utils.learning import adjust_learning_rate
from ..utils.log import logger
from ..base.module import Module

from .config import DEVICE, DEFAULT_CONFIG, Config, MODELS
from .tool import rl_tool, ENTITY, RELATION
from .utils.get_neg_batch import get_neg_batch

seed = 2019
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)

class KRL(Module):
    """
    """

    def __init__(self):
        self._model = None
        self._entity_vocab = None
        self._relation_vocab = None

    def train(self, train_path, model_type=DEFAULT_CONFIG['model_type'], save_path=DEFAULT_CONFIG['save_path'], dev_path=None, **kwargs):
        train_dataset = rl_tool.get_dataset(train_path)
        if dev_path:
            dev_dataset = rl_tool.get_dataset(dev_path)
            entity_vocab, relation_vocab = rl_tool.get_vocab(train_dataset, dev_dataset)
        else:
            entity_vocab, relation_vocab = rl_tool.get_vocab(train_dataset)
        self._entity_vocab = entity_vocab
        self._relation_vocab = relation_vocab
        train_iter = rl_tool.get_iterator(train_dataset, batch_size=DEFAULT_CONFIG['batch_size'])
        config = Config(entity_vocab, relation_vocab, save_path=save_path, **kwargs)
        if model_type in MODELS:
            model = MODELS[model_type](config)
        else:
            raise Exception('there is no model named {}! please check the name carefully'.format(model_type))
        # print(self._model)
        self._model = model
        optim = torch.optim.Adam(self._model.parameters(), lr=config.lr)
        # criterion = nn.MarginRankingLoss(config.loss_margin, reduction='sum')
        criterion = nn.MarginRankingLoss(config.loss_margin, reduction='mean')
        for epoch in range(config.epoch):
            self._model.train()
            acc_loss = 0
            for item in tqdm(train_iter):
                optim.zero_grad()
                # item.head = item.head.reshape(item.head.size(0), -1)
                # item.rel = item.rel.reshape(item.rel.size(0), -1)
                # item.tail = item.tail.reshape(item.tail.size(0), -1)
                neg_head, neg_tail = get_neg_batch(item.head, item.tail, self._model.entity_num)
                pos_score = self._model(item.head, item.rel, item.tail)
                neg_score = self._model(neg_head, item.rel, neg_tail)
                item_loss = criterion(neg_score, pos_score, torch.ones_like(pos_score))
                acc_loss += item_loss.item()
                item_loss.backward()
                optim.step()
            logger.info('epoch: {}, acc_loss: {}'.format(epoch, acc_loss))
            if dev_path:
                dev_score = self._validate(dev_dataset)
                logger.info('dev score:{}'.format(dev_score))
            adjust_learning_rate(optim, config.lr / (1 + (epoch + 1) * config.lr_decay))
        config.save()
        self._model.save()

    def predict(self, head: str, rel: str, tail:str):
        self._model.eval()
        if head not in self._entity_vocab.stoi:
            print('there is not entity named:{} in database!'.format(head))
            return None
        if rel not in self._relation_vocab.stoi:
            print('there is not relation named:{} in database!'.format(rel))
            return None
        if tail not in self._entity_vocab.stoi:
            print('there is not entity named:{} in database!'.format(tail))
            return None
        head_id = torch.tensor([self._entity_vocab.stoi[head]]).view((-1, 1)).to(DEVICE)
        rel_id = torch.tensor([self._relation_vocab.stoi[rel]]).view((-1, 1)).to(DEVICE)
        tail_id = torch.tensor([self._entity_vocab.stoi[tail]]).view((-1, 1)).to(DEVICE)
        score = torch.exp(- self._model(head_id, rel_id, tail_id))[0].item()
        return score

    def predict_head(self, rel: str, tail: str, topk=3):
        self._model.eval()
        if rel not in self._relation_vocab.stoi:
            print('there is not relation named:{} in database!'.format(rel))
            return None
        if tail not in self._entity_vocab.stoi:
            print('there is not entity named:{} in database!'.format(tail))
            return None
        rel_id = torch.tensor([self._relation_vocab.stoi[rel]] * self._model.entity_num).view((-1, 1)).to(DEVICE)
        tail_id = torch.tensor([self._entity_vocab.stoi[tail]] * self._model.entity_num).view((-1, 1)).to(DEVICE)
        candidate_heads = torch.arange(self._model.entity_num).view((-1, 1)).to(DEVICE)
        head_scores = torch.exp(- self._model(candidate_heads, rel_id, tail_id))
        topk_scores, topk_index = torch.topk(head_scores, topk)
        topk_scores = topk_scores.tolist()
        topk_entities = [self._entity_vocab.itos[i] for i in topk_index]
        return list(zip(topk_entities, topk_scores))

    def predict_rel(self, head, tail, topk=3):
        self._model.eval()
        if head not in self._entity_vocab.stoi:
            print('there is not entity named:{} in database!'.format(head))
            return None
        if tail not in self._entity_vocab.stoi:
            print('there is not entity named:{} in database!'.format(tail))
            return None
        head_id = torch.tensor([self._entity_vocab.stoi[head]] * self._model.rel_num).view((-1, 1)).to(DEVICE)
        tail_id = torch.tensor([self._entity_vocab.stoi[tail]] * self._model.rel_num).view((-1, 1)).to(DEVICE)
        candidate_rels = torch.arange(self._model.rel_num).view((-1, 1)).to(DEVICE)
        rel_scores = torch.exp(- self._model(head_id, candidate_rels, tail_id))
        topk_scores, topk_index = torch.topk(rel_scores, topk)
        topk_scores = topk_scores.tolist()
        topk_rels = [self._relation_vocab.itos[i] for i in topk_index.tolist()]
        return list(zip(topk_rels, topk_scores))

    def predict_tail(self, head, rel, topk=3):
        self._model.eval()
        if head not in self._entity_vocab.stoi:
            print('there is not entity named:{} in database!'.format(head))
            return None
        if rel not in self._relation_vocab.stoi:
            print('there is not relation named:{} in database!'.format(rel))
            return None
        head_id = torch.tensor([self._entity_vocab.stoi[head]] * self._model.entity_num).view((-1, 1)).to(DEVICE)
        rel_id = torch.tensor([self._relation_vocab.stoi[rel]] * self._model.entity_num).view((-1, 1)).to(DEVICE)
        candidate_tails = torch.arange(self._model.entity_num).view((-1, 1)).to(DEVICE)
        head_scores = torch.exp(- self._model(head_id, rel_id, candidate_tails))
        topk_scores, topk_index = torch.topk(head_scores, topk)
        topk_scores = topk_scores.tolist()
        topk_entities = [self._entity_vocab.itos[i] for i in topk_index]
        return list(zip(topk_entities, topk_scores))

    def load(self, save_path=DEFAULT_CONFIG['save_path'], model_type=DEFAULT_CONFIG['model_type']):
        config = Config.load(save_path)
        if model_type in MODELS:
            self._model = MODELS[model_type](config)
        else:
            raise Exception('there is no model named {}! please check the name carefully'.format(model_type))
        self._model.load()
        self._model = self._model
        self._entity_vocab = config.entity_vocab
        self._relation_vocab = config.rel_vocab
        self._check_vocab()

    def test(self, test_path):
        test_dataset = rl_tool.get_dataset(test_path)
        test_score = self._validate(test_dataset)
        logger.info('test score:{}'.format(test_score))

    def _validate(self, dev_dataset, batch_size=DEFAULT_CONFIG['batch_size']):
        self._model.eval()
        dev_score_list = []
        dev_iter = rl_tool.get_iterator(dev_dataset, batch_size=batch_size)
        for dev_item in tqdm(dev_iter):
            item_score = torch.exp(- self._model(dev_item.head, dev_item.rel, dev_item.tail))
            # item_score = self._model(dev_item.head, dev_item.rel, dev_item.tail)
            dev_score_list.extend(item_score.cpu().tolist())
        return sum(dev_score_list) / len(dev_score_list)

    def _check_vocab(self):
        if not hasattr(ENTITY, 'vocab'):
            ENTITY.vocab = self._entity_vocab
        if not hasattr(RELATION, 'vocab'):
            RELATION.vocab = self._relation_vocab


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

# Imports
from torch.utils.data.dataset import Dataset
import urllib
import os
import zipfile
import json
import codecs

# TIRA Author identification data set
class TIRAAuthorIdentificationDataset(Dataset):
    """
    TIRA Author identification data set
    """

    # Constructor
    def __init__(self, root, problem_name, transform, train, encoding):
        """
        Constructor
        :param root:
        :param problem_name:
        """
        # Properties
        self.root = root
        self.problem_name = problem_name
        self.transform = transform
        self.authors = list()
        self.last_text = ""
        self.n_authors = 0
        self.train = train
        self.encoding = encoding

        # List of text
        self.texts = list()

        # Generate data set
        self._load()
    # end if

    ##############################################
    # Static
    ##############################################

    # Get collection info
    @staticmethod
    def collection_infos(root):
        """
        Get collection info
        :return:
        """
        return json.load(open(os.path.join(root, "collection-info.json"), 'r'))
    # end colelction_infos

    ##############################################
    # OVERRIDE
    ##############################################

    # Length
    def __len__(self):
        """
        Length
        :return:
        """
        return len(self.texts)
    # end __len__

    # Get item
    def __getitem__(self, item):
        """
        Get item
        :param item:
        :return:
        """
        # Current file
        text_path, author_name = self.texts[item]
        self.last_text = text_path[text_path.rfind('/')+1:]

        # Read text
        text_content = codecs.open(text_path, 'r', encoding=self.encoding).read()

        # Transformed
        transformed = self.transform(text_content)

        # Unsqueeze
        transformed = transformed.squeeze(0)

        # Transform
        if self.transform is not None:
            return transformed, author_name
        else:
            return text_content, author_name
        # end if
    # end __getitem__

    ################################################
    # PRIVATE
    ################################################

    # Load dataset
    def _load(self):
        """
        Load the dataset
        :return:
        """
        # Load problem info
        problem_info = json.load(open(os.path.join(self.root, self.problem_name, "problem-info.json")))

        # Info
        unknown_folder = problem_info['unknown-folder']
        candidate_authors = problem_info['candidate-authors']

        # Load texts
        if self.train:
            self.n_authors = len(candidate_authors)

            # Load candidate texts
            for candidate_author in candidate_authors:
                # Author name
                author_name = candidate_author['author-name']

                # Add  to authors
                if author_name not in self.authors:
                    self.authors.append(author_name)
                # end if

                # For each text
                for file_path in os.listdir(os.path.join(self.root, self.problem_name, author_name)):
                    # Add with author name
                    self.texts.append((os.path.join(self.root, self.problem_name, author_name, file_path), author_name))
                # end for
            # end for
        else:
            # For each test files
            for unknown_file in os.listdir(os.path.join(self.root, self.problem_name, unknown_folder)):
                # Add with author name
                self.texts.append((os.path.join(self.root, self.problem_name, unknown_folder, unknown_file), ""))
            # end for
        # end if
    # end _load

# end TIRAAuthorIdentificationDataset


# -*- coding: utf-8 -*-
"""
Created on Fri Jan  3 15:50:48 2020

@author: rramn
"""

#Enabling Intellisense
#%config IPCompleter.greedy=True
#Importing Relevant packages
import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
import matplotlib.pyplot as mp
import seaborn as sns
from sklearn.ensemble import RandomForestClassifier
from sklearn.feature_selection import SelectFromModel
from sklearn_pandas import CategoricalImputer
from sklearn.preprocessing import LabelEncoder
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn import metrics
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.neural_network import MLPRegressor
from sklearn.linear_model import LogisticRegressionCV
from scipy.stats import chi2_contingency
from sklearn import linear_model
from sklearn import cluster
from scipy.spatial.distance import cdist
from sklearn.preprocessing import MinMaxScaler

# Data extraction and exploratory Analysis
rawdat=pd.read_csv(r'C:\Users\rramn\Documents\Kaggle_projects\Telecom_churn\cell2celltrain.csv')

df_vis=pd.DataFrame(rawdat)
df_vis['MonthlyRevenue'].max()
type(df_vis)
#print(np.dtype(t.iloc[0]))
nrow,ncol=df_vis.shape
colnames=list(df_vis)
mp.figure(num=None,figsize=(6*8,8*ncol),dpi=80)
for i in range(ncol):
    mp.subplot(ncol,8,i+1)
    t=df_vis.iloc[:,i]
    if(np.issubdtype(type(t.iloc[0]),np.number)):
        #t.hist()
        sns.boxplot(t)
    else:
        df_vis.iloc[:,i]=df_vis.iloc[:,i].astype('category')
        valcount=t.value_counts()
        valcount.plot.bar()
    mp.ylabel('counts')
    mp.title(f'{colnames[i]} (column {i})')
    mp.xticks(rotation=90)
    #mp.title()
mp.tight_layout(pad = 1.0, w_pad = 1.0, h_pad = 1.0)
mp.show()

#Bivariate analysis
mp.figure(figsize=(40,30))
crdf=df_vis.select_dtypes(exclude='category')
cor=crdf.corr(method='pearson')
sns.heatmap(cor)
mp.show()

#Data cleaning
df_vis.info()
sns.heatmap(df_vis.isnull(),cbar=False)
df_vis.isnull().sum()
df_vis.isna().sum()

df_vis['HandsetPrice']=df_vis['HandsetPrice'].replace('Unknown',-1)
df_vis['HandsetPrice']
df_vis['HandsetPrice']=df_vis['HandsetPrice'].astype('int')
df_vis['HandsetPrice']=df_vis['HandsetPrice'].replace(-1,np.NaN)
df_vis['HandsetPrice']
df_vis.Churn.value_counts()

#Missing values imputation
temp=df_vis
temp=temp.fillna(temp.mean())
imputer=CategoricalImputer(missing_values='NaN', strategy='most_frequent')
imputer.fit_transform(temp['ServiceArea'])
temp=temp.apply(lambda x: x.fillna(x.value_counts().index[0]))
temp.isna().sum()

#Statistical analysis
nr,nc=temp.shape
for j in range(nc):
    if((temp.iloc[:,j].dtype!=np.int64) & (temp.iloc[:,j].dtype!=np.number)):
        xx=temp.iloc[:,j]
        yy=temp['Churn']
        ct=pd.crosstab(xx,yy)
        ch=chi2_contingency(ct)
        print('Chisquare result for',temp.columns[j],'is  ',ch)
    else:
        print('Not a category')

#Predicting Churn index
#Label encoding categorical variables
temp.dtypes
cl_df=temp
cl_df.isna().sum()
cl_df=cl_df.drop(['CustomerID'],axis=1)

#cl_df['Churn']=cl_df['Churn'].cat.codes
cl_df['Churn']=LabelEncoder().fit_transform(cl_df.Churn)
cl_df.Churn.value_counts()
for n,v in cl_df.iteritems():
    if((cl_df[n].dtype!=np.int64) & (cl_df[n].dtype!=np.number) & (n!='Churn')):
        mean_encode=cl_df.groupby(n)['Churn'].mean()
        print(mean_encode)
        cl_df.loc[:,n]=cl_df[n].map(mean_encode)

cl_df.dtypes
x_cl=pd.DataFrame(cl_df.iloc[:,1:])
y_cl=cl_df.iloc[:,0]

#Logistic Regression significance test

lrr=LogisticRegression(penalty="l1")
lrr.fit(lrdf,y_cl)

#cat_columns=x_cl.select_dtypes(['category']).columns
#cat_columns
#x_cl[cat_columns]=x_cl[cat_columns].apply(lambda x: x.cat.codes)
#y_cl=y_cl.cat.codes

#Feature selection and importance
#
#rf=RandomForestClassifier(min_samples_split=2)
#rf.fit(xcl_train,ycl_train)
#predcl=rf.predict(xcl_test)
#metrics.accuracy_score(ycl_test,predcl)

# Feature selection Lasso regression
lr=SelectFromModel(LogisticRegression(penalty="l1"),max_features=40)
lr.fit(x_cl,y_cl)

lr_fea=lr.get_support()
lr_list=x_cl.loc[:,lr_fea].columns.tolist()
lr_list
#Feature importance using Random forests
rf=RandomForestClassifier(n_estimators=100,min_samples_split=4)
rf.fit(x_cl,y_cl)
imp_feat=pd.Series(rf.feature_importances_, index=x_cl.columns)
imp_feat.nlargest(20).plot(kind='barh')
ind=np.argsort(imp)
rf.feature_importances_

# Stacking model
xcl_train,xcl_test,ycl_train,ycl_test=train_test_split(x_cl,y_cl,test_size=0.3)
#Support vector classifier
#svm=SVC(C=5, probability=True,gamma='auto')
#svm.fit(xcl_train,ycl_train)
#
#
lr=LogisticRegressionCV(cv=10)
lr.fit(x_cl,y_cl)
metrics.accuracy_score(y_cl,lr.predict(x_cl))
metrics.roc_auc_score(y_cl,lr.predict(x_cl))
#nn = MLPClassifier((80, 10), early_stopping=False, random_state=SEED)
gb = GradientBoostingClassifier(n_estimators=100)
gb.fit(xcl_train,ycl_train)
metrics.accuracy_score(ycl_test,gb.predict(xcl_test))
#rf = RandomForestClassifier(n_estimators=10, max_features=3, random_state=SEED)

#Predicting Monthly Revenue
reg_df=temp
cl_df['Churn']
reg_df['Revenue_loss']=cl_df.Churn * reg_df['MonthlyRevenue']
reg_df['Revenue_loss']
reg_df=reg_df.drop(['MonthlyRevenue'],axis=1)
#Casting categorical types
for n,v in reg_df.iteritems():
    if((reg_df[n].dtype!=np.int64) & (reg_df[n].dtype!=np.number)):
        mean_encode=reg_df.groupby(n)['Revenue_loss'].mean()
        print(mean_encode)
        reg_df.loc[:,n]=reg_df[n].map(mean_encode)

reg_df.dtypes
y_reg=reg_df['Revenue_loss']
x_reg=pd.DataFrame(reg_df.drop(['CustomerID','Revenue_loss'],axis=1))
x_reg.shape
#scaler=StandardScaler()
#x_reg=scaler.fit_transform(x_reg)

min_max_scaler = MinMaxScaler()
x_scaled = min_max_scaler.fit_transform(x_reg)
x_reg = pd.DataFrame(x_scaled)
#Stochastic gradient Descent and Neural Networks

sgdr= linear_model.LinearRegression()
sgdr.fit(x_reg,y_reg)
np.sqrt(metrics.mean_squared_error(y_reg,sgdr.predict(x_reg)))

nnr=MLPRegressor(hidden_layer_sizes=(20,),activation='relu')
nnr.fit(x_reg,y_reg)
metrics.mean_squared_error(y_reg,nnr.predict(x_reg))

#Customer segmentation
clus_df=reg_df
scaler=StandardScaler()
scale_df=scaler.fit_transform(clus_df)

#Choosing the best k
cluster_range=[2,3,4,5,6,7,8,9,10]
distortions=[]
for k in cluster_range:
    clus=cluster.KMeans(n_clusters=k,max_iter=100)
    clus.fit(scale_df)
    distortions.append(sum(np.min(cdist(scale_df,clus.cluster_centers_,'euclidean'), axis=1))/scale_df.shape[0])

mp.plot(cluster_range,distortions,'bx-')
mp.xlabel('k')
mp.ylabel('Distortion')
mp.title('Elbow method showing optimal k')
mp.show()

#Finding silhouette measure with best k
clusk=cluster.KMeans(n_clusters=5)
clusk.fit(scale_df)
metrics.silhouette_score(scale_df,clusk.labels_)
pred=clusk.transform(scale_df)
clust=pd.DataFrame(pred)
clust.shape
sns.barplot(x=clust,y=df_vis['MonthlyMinutes'])


from django.contrib import admin
from django.urls import path, include
from . import views
urlpatterns = [
path('', views.homepage, name="homepage"),
path('rest_prem',views.rest_prem,name="rest_prem"),
path('signup', views.signup, name="signup"),
path('signout', views.signout, name="signout"),
path('signin', views.signin, name="signin"),
path('homepage', views.homepage, name="homepage"),
path('team',views.rest_prem,name="team"),
path('my_team_stats',views.my_team_stats, name="my_team_stats"),
path('my_fixtures',views.my_fixtures, name="my_fixtures"),
path('my_minutes',views.my_minutes, name="my_minutes"),
path('my_shooting',views.my_shooting, name="my_shooting"),
path('fixtures',views.fixtures,name="fixtures"),
path('squad_stats',views.squad_stats,name="squad_stats"),
path('get_pred',views.get_pred,name="get_pred"),
path('go_to_pred',views.go_to_pred,name="go_to_pred"),
]


# README: moved here because of module cyclic dependencies problems
from typing import List, Set

class ShelleyPyError(Exception):
    DUPLICATED_METHOD = "Duplicated method!"
    MISSING_RETURN = "Missing return!"
    CASE_MISSING_RETURN = "Missing return for case!"
    MATCH_CALL_TYPE = "Match call type mismatch. Accepted types are: Call, Await!"
    MATCH_CASE_VALUE_TYPE = "Cases values must be strings or list of strings!"
    RETURN_PARSE_ERROR = "Could not parse return. Expecting str|list[,*]"
    ALL_BRANCH_RETURN_INSIDE_LOOP = (
        "All branches return which makes loop useless. Please remove the loop."
    )

    # DECORATOR_PARSE_ERROR = "Could not parse decorator!"

    def __init__(self, lineno: int, msg: str):
        self.lineno = lineno
        self.msg = msg
        super().__init__(self.msg)

class ReturnMatchesNext(ShelleyPyError):
    def __init__(self, lineno: int, return_name: List[str]):
        self.lineno = lineno
        super().__init__(
            lineno, f"Return names {return_name} are not listed in the next operations!"
        )

class NextOpsNotInReturn(ShelleyPyError):
    def __init__(self, lineno: int, missing_returns: Set[str]):
        self.lineno = lineno
        super().__init__(
            lineno,
            f"Next operations {missing_returns} do not have a corresponding return!",
        )

class CompilationError(Exception):
    pass


from celery.result import AsyncResult
import math
import re
import pickle
import collections
import json
import celery
from pymongo import MongoClient, DESCENDING
from rest_framework.reverse import reverse
from collections import OrderedDict
from datetime import datetime

from celery.task.control import inspect
from api import config

class celeryConfig:
    BROKER_URL = config.BROKER_URL
    BROKER_USE_SSL = config.BROKER_USE_SSL
    CELERY_SEND_EVENTS = True
    CELERY_TASK_RESULT_EXPIRES = None
    CELERY_RESULT_BACKEND = config.CELERY_RESULT_BACKEND
    CELERY_MONGODB_BACKEND_SETTINGS = config.CELERY_MONGODB_BACKEND_SETTINGS

def check_memcache(host=config.MEMCACHE_HOST, port=config.MEMCACHE_PORT):
    """ Check if memcache is running on server """
    import socket

    s = socket.socket()
    try:
        s.connect((host, port))
        return True
    except:
        return False

class QueueTask():
    def __init__(self, database=config.MONGO_DB,
                 log_collection=config.MONGO_LOG_COLLECTION,
                 tomb_collection=config.MONGO_TOMBSTONE_COLLECTION,
                 memcache=check_memcache()):
        self.db = MongoClient(config.DATA_STORE_MONGO_URI)
        self.database = database
        self.collection = log_collection
        self.tomb_collection = tomb_collection
        self.memcache=memcache
        celery.Celery().config_from_object(celeryConfig)
        self.i = inspect()

    def run(self, task, task_args, task_kwargs, task_queue, user_info, tags):
        """
        Submit task to celerey async tasks
        """
        # print(celeryConfig)
        #app = Celery().config_from_object(celeryConfig)
        print(user_info)
        # Submit task
        task_obj = celery.current_app.send_task(
            task, args=task_args, kwargs=task_kwargs, queue=task_queue, track_started=True,headers={"authenticated_user":user_info})
        task_log = {
            'task_id': task_obj.task_id,
            'user': user_info['username'],
            'task_name': task,
            'args': task_args,
            'kwargs': task_kwargs,
            'queue': task_queue,
            'timestamp': datetime.now(),
            'tags': tags
        }
        self.db[self.database][self.collection].insert(task_log)

        return {"task_id": task_obj.task_id}

    def list_tasks(self):
        """ List available tasks """
        REGISTERED_TASKS, AVAILABLE_QUEUES = self.update_tasks()
        # print REGISTERED_TASKS, AVAILABLE_QUEUES
        REGISTERED_TASKS = [task for task in list(
            REGISTERED_TASKS) if task[0:6] != "celery"]
        AVAILABLE_QUEUES = list(AVAILABLE_QUEUES)
        REGISTERED_TASKS.sort()
        AVAILABLE_QUEUES.sort()
        return {"available_tasks": REGISTERED_TASKS, "available_queues": AVAILABLE_QUEUES}
        #return list_tasks()

    def status(self, task_id=None):
        """ Return a task's status """
        col = self.db[self.database][self.tomb_collection]
        if task_id:
            result = [item for item in col.find({'_id': task_id})]
            if len(result) == 0:
                try:
                    res = celery.result.AsyncResult(task_id)
                    return {"status": "%s" % (res.status), "task_id": "%s" % (task_id)}
                except:
                    raise
                    # return {"status": "ERROR IN OBTAINING STATUS", "task_id": "%s" % (task_id)}
            else:
                return {"status": result[0]['status']}
        else:
            raise Exception("Not a valid task_id")

    def result(self, task_id=None, redirect=True):
        """ Get the result of a task  """
        col = self.db[self.database][self.tomb_collection]
        if task_id:
            result = [item for item in col.find({'_id': task_id})]
            try:
                result = pickle.loads(result[0]['result'])
                return result
            except:
                raise Exception("Not a valid task_id")
        else:
            raise Exception("Not a valid task_id")

    def task(self, task_id=None):
        """Return task log and task results"""
        doc = self.db[self.database][self.collection].find_one(
            {'task_id': task_id}, {'_id': False})
        col = self.db[self.database][self.tomb_collection]
        if doc:
            result = col.find_one({'_id': task_id}, {'_id': False})
            if result:
                result = self.unpickle_result(result)
            else:
                result = self.status(task_id=task_id)
            doc['result'] = result
            return doc
        else:
            result = col.find_one({'_id': task_id}, {'_id': False})
            if result:
                result = self.unpickle_result(result)
                return result
            else:
                result = self.status(task_id=task_id)
            return result

    def unpickle_result(self, result):
        if 'traceback' in result:
            if type(result['traceback']) == bytes:
                result['traceback'] = pickle.loads(result['traceback'])
            try:
                result['traceback'] =json.loads(result['traceback'])
            except:
                pass
        if 'children' in result:
            if type(result['children']) == bytes:
                result['children'] = pickle.loads(result['children'])
            try:
                result['children'] =json.loads(result['children'])
            except:
                pass

        if 'result' in result:
            if type(result['result']) == bytes:
                result['result'] = pickle.loads(result['result'])
            if isinstance(result['result'], Exception):
                result['result'] = "ERROR: {0}".format(str(result['result']))
            try:
                result['result'] =json.loads(result['result'])
            except:
                pass
        return result

    def reset_tasklist(self, user="guest"):
        """
        Delete and reload memcached record of available tasks, useful for development
        when tasks are being frequently reloaded.
        """
        if self.memcache:
            import memcache
            mc = memcache.Client(
                ['%s:%s' % (config.MEMCACHE_HOST, config.MEMCACHE_PORT)])
            tasks = "REGISTERED_TASKS_%s" % user
            queues = "AVAILABLE_QUEUES_%s" % user
            mc.delete(tasks)
            mc.delete(queues)
        return self.list_tasks()

    def history(self, user, task_name=None, page=1, limit=0, request=None):
        """ Show a history of tasks """
        if page < 1:
            page = 1
        limit = int(limit)
        col = self.db[self.database][self.collection]
        result = {'count': 0, 'next': None, 'previous': None, 'results': []}
        history = []
        if task_name:
            tasks_list=task_name.split(',')
            result['count'] = col.find(
                {'task_name': {"$in":tasks_list}, 'user': user}).count()
            data = col.find({'task_name': {"$in":tasks_list}, 'user': user}, {'_id': False}, skip=(page - 1) * limit,
                            limit=limit).sort('timestamp', DESCENDING)
        else:
            data = col.find({'user': user}, {'_id': False}, skip=(page - 1) * limit, limit=limit).sort('timestamp',
                                                                                DESCENDING)
            result['count'] = col.find({'user': user}).count()
        if result['count'] <= page*limit:
            if page != 1:
                result['previous'] = "%s?page=%d&page_size=%d" % (
                    reverse('queue-user-tasks', request=request), page-1, limit)
        if result['count'] >= page*limit:
            if result['count'] != page*limit:
                result['next'] = "%s?page=%d&page_size=%d" % (
                    reverse('queue-user-tasks', request=request), page+1, limit)
            if page > 1:
                result['previous'] = "%s?page=%d&page_size=%d" % (
                    reverse('queue-user-tasks', request=request), page-1, limit)
        result['meta'] = {'page': page, 'page_size': limit,
                          'pages': math.ceil(float(result['count'])/float(limit))}
        for item in data:
            if type(item['kwargs']) is dict:
                for i, v in item['kwargs'].items():
                    try:
                        item['kwargs'][i] = json.loads(v)
                    except:
                        pass
            try:
                item['result'] = reverse(
                    'queue-task-result', kwargs={'task_id': item['task_id']}, request=request)
            except:
                item['result'] = ""
            history.append(item)
        result['results'] = history
        try:
            od = collections.OrderedDict(sorted(result.items()))
        except:
            # older python versions < 2.7
            od = OrderedDict(sorted(result.items()))
        return od

    def task_docstring(self, taskname):
        """
        Get task docstring of a registered tasks from celery.
        """
        data = self.i.registered('__doc__')

        for x, v in data.items():
            for task in v:
                name, doc = self.get_taskname_doc(task, ']')
                if name.strip() == taskname.strip():
                    doc=doc.replace('.\n',' ')
                    doc=doc.replace('\n','. ')
                    return doc
        return None

    def get_taskname_doc(self, thestring, ending):
        temp = thestring.split('[__doc__=')
        if len(temp) > 1:
            if temp[1].endswith(ending):
                return temp[0], re.sub(' +', ' ', temp[1][:-len(ending)])
            return temp[0], re.sub(' +', ' ', temp[1])
        return temp[0], ""

    def update_tasks(self,timeout=60000, user="guest"):
        """
        Get list of registered tasks from celery, store in memcache for
            `timeout` period if set (default to 6000s) if available
        """
        try:
            if self.memcache:
                import memcache
                mc = memcache.Client(
                    ['%s:%s' % (config.MEMCACHE_HOST, config.MEMCACHE_PORT)])
                tasks = "REGISTERED_TASKS_%s" % user
                queues = "AVAILABLE_QUEUES_%s" % user
                REGISTERED_TASKS = mc.get(tasks)
                AVAILABLE_QUEUES = mc.get(queues)
                if not REGISTERED_TASKS:
                    REGISTERED_TASKS = set()
                    for item in self.i.registered().values():
                        REGISTERED_TASKS.update(item)
                    mc.set(tasks, REGISTERED_TASKS, timeout)
                    REGISTERED_TASKS = mc.get(tasks)
                if not AVAILABLE_QUEUES:
                    mc.set(queues, set([item[0]["exchange"]["name"]
                                        for item in self.i.active_queues().values()]), timeout)
                    AVAILABLE_QUEUES = mc.get(queues)
            else:
                REGISTERED_TASKS = set()
                for item in self.i.registered().values():
                    REGISTERED_TASKS.update(item)
                AVAILABLE_QUEUES = set([item[0]["exchange"]["name"]
                                        for item in self.i.active_queues().values()])
        except:
            REGISTERED_TASKS = set()
            AVAILABLE_QUEUES = set()
        return (REGISTERED_TASKS, AVAILABLE_QUEUES)