import torch
import torch.nn as nn
from dreamerv2.utils.rssm import RSSMUtils, RSSMContState, RSSMDiscState

class RSSM(nn.Module, RSSMUtils):
    def __init__(
        self,
        action_size,
        rssm_node_size,
        embedding_size,
        device,
        rssm_type,
        info,
        act_fn=nn.ELU,
    ):
        nn.Module.__init__(self)
        RSSMUtils.__init__(self, rssm_type=rssm_type, info=info)
        self.device = device
        self.action_size = action_size
        self.node_size = rssm_node_size
        self.embedding_size = embedding_size
        self.act_fn = act_fn
        self.rnn = nn.GRUCell(self.deter_size, self.deter_size)
        self.fc_embed_state_action = self._build_embed_state_action()
        self.fc_prior = self._build_temporal_prior()
        self.fc_posterior = self._build_temporal_posterior()

    def _build_embed_state_action(self):
        """
        model is supposed to take in previous stochastic state and previous action
        and embed it to deter size for rnn input
        """
        fc_embed_state_action = [nn.Linear(self.stoch_size + self.action_size, self.deter_size)]
        fc_embed_state_action += [self.act_fn()]
        return nn.Sequential(*fc_embed_state_action)

    def _build_temporal_prior(self):
        """
        model is supposed to take in latest deterministic state
        and output prior over stochastic state
        """
        temporal_prior = [nn.Linear(self.deter_size, self.node_size)]
        temporal_prior += [self.act_fn()]
        if self.rssm_type == 'discrete':
            temporal_prior += [nn.Linear(self.node_size, self.stoch_size)]
        elif self.rssm_type == 'continuous':
             temporal_prior += [nn.Linear(self.node_size, 2 * self.stoch_size)]
        return nn.Sequential(*temporal_prior)

    def _build_temporal_posterior(self):
        """
        model is supposed to take in latest embedded observation and deterministic state
        and output posterior over stochastic states
        """
        temporal_posterior = [nn.Linear(self.deter_size + self.embedding_size, self.node_size)]
        temporal_posterior += [self.act_fn()]
        if self.rssm_type == 'discrete':
            temporal_posterior += [nn.Linear(self.node_size, self.stoch_size)]
        elif self.rssm_type == 'continuous':
            temporal_posterior += [nn.Linear(self.node_size, 2 * self.stoch_size)]
        return nn.Sequential(*temporal_posterior)

    def rssm_imagine(self, prev_action, prev_rssm_state, nonterms=True):
        state_action_embed = self.fc_embed_state_action(torch.cat([prev_rssm_state.stoch*nonterms, prev_action],dim=-1))
        deter_state = self.rnn(state_action_embed, prev_rssm_state.deter*nonterms)
        if self.rssm_type == 'discrete':
            prior_logit = self.fc_prior(deter_state)
            stats = {'logit':prior_logit}
            prior_stoch_state = self.get_stoch_state(stats)
            prior_rssm_state = RSSMDiscState(prior_logit, prior_stoch_state, deter_state)

        elif self.rssm_type == 'continuous':
            prior_mean, prior_std = torch.chunk(self.fc_prior(deter_state), 2, dim=-1)
            stats = {'mean':prior_mean, 'std':prior_std}
            prior_stoch_state, std = self.get_stoch_state(stats)
            prior_rssm_state = RSSMContState(prior_mean, std, prior_stoch_state, deter_state)
        return prior_rssm_state

    def rollout_imagination(self, horizon:int, actor:nn.Module, prev_rssm_state):
        rssm_state = prev_rssm_state
        next_rssm_states = []
        action_entropy = []
        imag_log_probs = []
        for t in range(horizon):
            action, action_dist = actor((self.get_model_state(rssm_state)).detach())
            rssm_state = self.rssm_imagine(action, rssm_state)
            next_rssm_states.append(rssm_state)
            action_entropy.append(action_dist.entropy())
            imag_log_probs.append(action_dist.log_prob(torch.round(action.detach())))

        next_rssm_states = self.rssm_stack_states(next_rssm_states, dim=0)
        action_entropy = torch.stack(action_entropy, dim=0)
        imag_log_probs = torch.stack(imag_log_probs, dim=0)
        return next_rssm_states, imag_log_probs, action_entropy

    def rssm_observe(self, obs_embed, prev_action, prev_nonterm, prev_rssm_state):
        prior_rssm_state = self.rssm_imagine(prev_action, prev_rssm_state, prev_nonterm)
        deter_state = prior_rssm_state.deter
        x = torch.cat([deter_state, obs_embed], dim=-1)
        if self.rssm_type == 'discrete':
            posterior_logit = self.fc_posterior(x)
            stats = {'logit':posterior_logit}
            posterior_stoch_state = self.get_stoch_state(stats)
            posterior_rssm_state = RSSMDiscState(posterior_logit, posterior_stoch_state, deter_state)

        elif self.rssm_type == 'continuous':
            posterior_mean, posterior_std = torch.chunk(self.fc_posterior(x), 2, dim=-1)
            stats = {'mean':posterior_mean, 'std':posterior_std}
            posterior_stoch_state, std = self.get_stoch_state(stats)
            posterior_rssm_state = RSSMContState(posterior_mean, std, posterior_stoch_state, deter_state)
        return prior_rssm_state, posterior_rssm_state

    def rollout_observation(self, seq_len:int, obs_embed: torch.Tensor, action: torch.Tensor, nonterms: torch.Tensor, prev_rssm_state):
        priors = []
        posteriors = []
        for t in range(seq_len):
            prev_action = action[t]*nonterms[t]
            prior_rssm_state, posterior_rssm_state = self.rssm_observe(obs_embed[t], prev_action, nonterms[t], prev_rssm_state)
            priors.append(prior_rssm_state)
            posteriors.append(posterior_rssm_state)
            prev_rssm_state = posterior_rssm_state
        prior = self.rssm_stack_states(priors, dim=0)
        post = self.rssm_stack_states(posteriors, dim=0)
        return prior, post


import json
import base64

from xcalar.compute.coretypes.LibApisCommon.ttypes import *
from xcalar.compute.coretypes.LibApisCommon.constants import *
from xcalar.compute.coretypes.LibApisEnums.ttypes import *
from xcalar.compute.coretypes.LibApisConstants.ttypes import *
from xcalar.compute.coretypes.DagTypes.ttypes import *
from xcalar.compute.coretypes.DagTypes.constants import *
from xcalar.compute.coretypes.XcalarApiVersionSignature.ttypes import XcalarApiVersionT
from xcalar.compute.coretypes.UdfTypes.ttypes import *
from xcalar.compute.coretypes.DataTargetTypes.ttypes import *
from xcalar.compute.coretypes.DataTargetTypes.constants import *
from xcalar.compute.coretypes.DataTargetEnums.constants import *
from xcalar.compute.coretypes.DataFormatEnums.ttypes import *
from xcalar.compute.coretypes.DataFormatEnums.constants import *
from xcalar.compute.coretypes.OrderingEnums.constants import *
from xcalar.compute.coretypes.JoinOpEnums.constants import *
from xcalar.compute.coretypes.UnionOpEnums.constants import *

#
# This file defines convenience classes for constructing Xcalar API work
# items.
#

# Base class for all work items. Defines a few basic things.
class WorkItem(object):
    def __init__(self, userName=None, userIdUnique=None, sessionName=None):
        self.workItem = XcalarApiWorkItemT()
        self.workItem.apiVersionSignature = XcalarApiVersionT.XcalarApiVersionSignature
        if sessionName:
            self.workItem.sessionName = sessionName
        if userIdUnique:
            self.workItem.userIdUnique = userIdUnique
        if userName:
            self.workItem.userId = userName

    # Gives the work item a chance to define how output should look to callers.
    def output(self, executeOutput):
        return executeOutput

class WorkItemGetNumNodes(WorkItem):
    def __init__(self, userName=None, userIdUnique=None):
        super(WorkItemGetNumNodes, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiGetNumNodes

    def output(self, executeOutput):
        return executeOutput.output.outputResult.getNumNodesOutput

class WorkItemTop(WorkItem):
    def __init__(self,
                 measureIntervalInMs=None,
                 cacheValidityInMs=None,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemTop, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiTop

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.topInput = XcalarApiTopInputT()
        assert (measureIntervalInMs is not None)
        assert (cacheValidityInMs is not None)
        self.workItem.input.topInput.measureIntervalInMs = measureIntervalInMs
        self.workItem.input.topInput.cacheValidityInMs = cacheValidityInMs

    def output(self, executeOutput):
        return executeOutput.output.outputResult.topOutput

    def updateMeasureIntervalInMs(self, measureIntervalInMs=None):
        assert (measureIntervalInMs is not None)
        self.workItem.input.topInput.measureIntervalInMs = measureIntervalInMs

    def updateCacheValidity(self, cacheValidityInMs=None):
        assert (cacheValidityInMs is not None)
        self.workItem.input.topInput.cacheValidityInMs = \
            cacheValidityInMs

class WorkItemListDataset(WorkItem):
    def __init__(self, userName=None, userIdUnique=None):
        # XXX super() in python 3.
        super(WorkItemListDataset, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiListDatasets

    def output(self, executeOutput):
        return executeOutput.output.outputResult.listDatasetsOutput

class WorkItemGetDatasetsInfo(WorkItem):
    def __init__(self, datasetsNamePattern, userName=None, userIdUnique=None):
        super(WorkItemGetDatasetsInfo, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiGetDatasetsInfo

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.getDatasetsInfoInput = XcalarApiGetDatasetsInfoInputT(
        )
        self.workItem.input.getDatasetsInfoInput.datasetsNamePattern = datasetsNamePattern

    def output(self, executeOutput):
        return executeOutput.output.outputResult

class WorkItemListDagInfo(WorkItem):
    def __init__(self, namePattern, srcType, userName=None, userIdUnique=None):
        super(WorkItemListDagInfo, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiListDagNodeInfo

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.listDagNodesInput = XcalarApiDagNodeNamePatternInputT(
        )
        self.workItem.input.listDagNodesInput.srcType = srcType
        self.workItem.input.listDagNodesInput.namePattern = namePattern

    def output(self, executeOutput):
        return executeOutput.output.outputResult.listNodesOutput

class WorkItemQueryState(WorkItem):
    def __init__(self,
                 queryName,
                 detailedStats=True,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemQueryState, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiQueryState

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.queryStateInput = XcalarApiQueryNameInputT()
        self.workItem.input.queryStateInput.queryName = queryName
        self.workItem.input.queryStateInput.detailedStats = detailedStats

    def output(self, executeOutput):
        return executeOutput.output.outputResult.queryStateOutput

class WorkItemQueryCancel(WorkItem):
    def __init__(self, queryName, userName=None, userIdUnique=None):
        super(WorkItemQueryCancel, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiQueryCancel

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.queryStateInput = XcalarApiQueryNameInputT()
        self.workItem.input.queryStateInput.queryName = queryName

class WorkItemQueryDelete(WorkItem):
    def __init__(self, queryName="", userName=None, userIdUnique=None):
        super(WorkItemQueryDelete, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiQueryDelete

        self.workItem.input = XcalarApiInputT()
        # Not a typo. Query delete uses queryStateInput for some reason...
        self.workItem.input.queryStateInput = XcalarApiQueryNameInputT()
        self.workItem.input.queryStateInput.queryName = queryName

class WorkItemCancelOp(WorkItem):
    def __init__(self, name):
        super(WorkItemCancelOp, self).__init__()
        self.workItem.api = XcalarApisT.XcalarApiCancelOp

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.dagTableNameInput = XcalarApiDagTableNameInputT()
        self.workItem.input.dagTableNameInput.tableInput = name

class WorkItemGetOpStats(WorkItem):
    def __init__(self, name, userName=None, userIdUnique=None):
        super(WorkItemGetOpStats, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiGetOpStats

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.dagTableNameInput = XcalarApiDagTableNameInputT()
        self.workItem.input.dagTableNameInput.tableInput = name

    def output(self, executeOutput):
        return executeOutput.output.outputResult.opStatsOutput

class WorkItemDeleteDagNode(WorkItem):
    def __init__(self,
                 namePattern,
                 srcType,
                 userName=None,
                 userIdUnique=None,
                 sessionName=None,
                 deleteCompletely=False):
        super(WorkItemDeleteDagNode, self).__init__(userName, userIdUnique,
                                                    sessionName)
        self.workItem.api = XcalarApisT.XcalarApiDeleteObjects

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.deleteDagNodeInput = XcalarApiDagNodeNamePatternInputT(
        )
        self.workItem.input.deleteDagNodeInput.namePattern = namePattern
        self.workItem.input.deleteDagNodeInput.srcType = srcType
        self.workItem.input.deleteDagNodeInput.deleteCompletely = deleteCompletely

    def output(self, executeOutput):
        return executeOutput.output.outputResult.deleteDagNodesOutput

class WorkItemRenameNode(WorkItem):
    def __init__(self, oldName, newName, userName=None, userIdUnique=None):
        super().__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiRenameNode

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.renameNodeInput = XcalarApiRenameNodeInputT()
        self.workItem.input.renameNodeInput.oldName = oldName
        self.workItem.input.renameNodeInput.newName = newName

class WorkItemArchiveTables(WorkItem):
    def __init__(self,
                 archive,
                 tableNames,
                 allTables=False,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemArchiveTables, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiArchiveTables

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.archiveTablesInput = XcalarApiArchiveTablesInputT()
        self.workItem.input.archiveTablesInput.archive = archive
        self.workItem.input.archiveTablesInput.allTables = allTables
        self.workItem.input.archiveTablesInput.tableNames = tableNames

    def output(self, executeOutput):
        return executeOutput.output.outputResult.archiveTablesOutput

class WorkItemSessionNew(WorkItem):
    def __init__(self,
                 name,
                 fork,
                 forkedSession=None,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemSessionNew, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiSessionNew

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.sessionNewInput = XcalarApiSessionNewInputT()
        self.workItem.input.sessionNewInput.sessionName = name
        self.workItem.input.sessionNewInput.fork = fork
        self.workItem.input.sessionNewInput.forkedSessionName = forkedSession

class WorkItemSessionDelete(WorkItem):
    def __init__(self, namePattern, userName=None, userIdUnique=None):
        super(WorkItemSessionDelete, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiSessionDelete

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.sessionDeleteInput = XcalarApiSessionDeleteInputT()
        self.workItem.input.sessionDeleteInput.sessionName = namePattern
        self.workItem.input.sessionDeleteInput.noCleanup = False

class WorkItemSessionActivate(WorkItem):
    def __init__(self, name, userName=None, userIdUnique=None):
        super(WorkItemSessionActivate, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiSessionActivate

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.sessionActivateInput = XcalarApiSessionActivateInputT(
        )
        self.workItem.input.sessionActivateInput.sessionName = name

class WorkItemSessionInact(WorkItem):
    def __init__(self, name, userName=None, userIdUnique=None):
        super(WorkItemSessionInact, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiSessionInact

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.sessionDeleteInput = XcalarApiSessionDeleteInputT()
        self.workItem.input.sessionDeleteInput.sessionName = name
        self.workItem.input.sessionDeleteInput.noCleanup = False

class WorkItemSessionList(WorkItem):
    def __init__(self, pattern, userName=None, userIdUnique=None):
        super(WorkItemSessionList, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiSessionList

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.sessionListInput = XcalarApiSessionListArrayInputT(
        )
        self.workItem.input.sessionListInput.sesListInput = pattern

    def output(self, executeOutput):
        return executeOutput.output.outputResult.sessionListOutput

class WorkItemSessionRename(WorkItem):
    def __init__(self, newName, oldName, userName=None, userIdUnique=None):
        super().__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiSessionRename

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.sessionRenameInput = XcalarApiSessionRenameInputT()
        self.workItem.input.sessionRenameInput.sessionName = newName
        self.workItem.input.sessionRenameInput.origSessionName = oldName

class WorkItemSessionPersist(WorkItem):
    def __init__(self, name, userName=None, userIdUnique=None):
        super().__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiSessionPersist

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.sessionDeleteInput = XcalarApiSessionDeleteInputT()
        self.workItem.input.sessionDeleteInput.sessionName = name

class WorkItemSessionDownload(WorkItem):
    def __init__(self,
                 name,
                 pathToAdditionalFiles,
                 userName=None,
                 userIdUnique=None):
        super().__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiSessionDownload

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.sessionDownloadInput = XcalarApiSessionDownloadInputT(
        )
        self.workItem.input.sessionDownloadInput.sessionName = name
        self.workItem.input.sessionDownloadInput.pathToAdditionalFiles = pathToAdditionalFiles

    def output(self, executeOutput):
        downloadStr = executeOutput.output.outputResult.sessionDownloadOutput.sessionContent
        return base64.b64decode(downloadStr)

class WorkItemSessionUpload(WorkItem):
    def __init__(self,
                 name,
                 sessionContent,
                 pathToAdditionalFiles,
                 userName=None,
                 userIdUnique=None):
        super().__init__(userName, userIdUnique)

        encodedSessionContent = base64.b64encode(sessionContent).decode(
            "ascii")

        self.workItem.api = XcalarApisT.XcalarApiSessionUpload

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.sessionUploadInput = XcalarApiSessionUploadInputT()
        self.workItem.input.sessionUploadInput.sessionName = name
        self.workItem.input.sessionUploadInput.sessionContent = encodedSessionContent
        self.workItem.input.sessionUploadInput.sessionContentCount = len(
            encodedSessionContent)
        self.workItem.input.sessionUploadInput.pathToAdditionalFiles = pathToAdditionalFiles

    def output(self, executeOutput):
        return executeOutput.output.outputResult.sessionNewOutput

class WorkItemGetTableMeta(WorkItem):
    def __init__(self,
                 isTable,
                 name,
                 isPrecise,
                 userName=None,
                 userIdUnique=None,
                 sessionName=None):
        super(WorkItemGetTableMeta, self).__init__(userName, userIdUnique,
                                                   sessionName)
        self.workItem.api = XcalarApisT.XcalarApiGetTableMeta

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.getTableMetaInput = XcalarApiGetTableMetaInputT()
        self.workItem.input.getTableMetaInput.tableNameInput = XcalarApiNamedInputT(
        )
        self.workItem.input.getTableMetaInput.tableNameInput.isTable = isTable
        self.workItem.input.getTableMetaInput.tableNameInput.name = name
        self.workItem.input.getTableMetaInput.tableNameInput.xid = XcalarApiXidInvalidT
        self.workItem.input.getTableMetaInput.isPrecise = isPrecise

    def output(self, executeOutput):
        return executeOutput.output.outputResult.getTableMetaOutput

class WorkItemPreview(WorkItem):
    def __init__(self,
                 sourceArgs,
                 numBytes,
                 offset,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemPreview, self).__init__(userName, userIdUnique)
        inputObj = {
            "func": "preview",
            "sourceArgs": {
                "targetName": sourceArgs.targetName,
                "path": sourceArgs.path,
                "fileNamePattern": sourceArgs.fileNamePattern,
                "recursive": sourceArgs.recursive,
            },
            "offset": offset,
            "bytesRequested": numBytes
        }
        self.workItem.api = XcalarApisT.XcalarApiPreview
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.previewInput = XcalarApiPreviewInputT()
        self.workItem.input.previewInput.inputJson = json.dumps(inputObj)

    def output(self, executeOutput):
        return executeOutput.output.outputResult.previewOutput

class WorkItemAddTarget2(WorkItem):
    def __init__(self,
                 targetTypeId,
                 targetName,
                 targetParams,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemAddTarget2, self).__init__(userName, userIdUnique)
        inputObj = {
            "func": "addTarget",
            "targetTypeId": targetTypeId,
            "targetName": targetName,
            "targetParams": targetParams
        }
        self.workItem.api = XcalarApisT.XcalarApiTarget
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.targetInput = XcalarApiTargetInputT()
        self.workItem.input.targetInput.inputJson = json.dumps(inputObj)

    def output(self, executeOutput):
        return json.loads(
            executeOutput.output.outputResult.targetOutput.outputJson)

class WorkItemDeleteTarget2(WorkItem):
    def __init__(self, targetName, userName=None, userIdUnique=None):
        super(WorkItemDeleteTarget2, self).__init__(userName, userIdUnique)
        inputObj = {"func": "deleteTarget", "targetName": targetName}
        self.workItem.api = XcalarApisT.XcalarApiTarget
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.targetInput = XcalarApiTargetInputT()
        self.workItem.input.targetInput.inputJson = json.dumps(inputObj)

    def output(self, executeOutput):
        return json.loads(
            executeOutput.output.outputResult.targetOutput.outputJson)

class WorkItemListTargets2(WorkItem):
    def __init__(self, userName=None, userIdUnique=None):
        super(WorkItemListTargets2, self).__init__(userName, userIdUnique)
        inputObj = {"func": "listTargets"}
        self.workItem.api = XcalarApisT.XcalarApiTarget
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.targetInput = XcalarApiTargetInputT()
        self.workItem.input.targetInput.inputJson = json.dumps(inputObj)

    def output(self, executeOutput):
        return json.loads(
            executeOutput.output.outputResult.targetOutput.outputJson)

class WorkItemListTargetTypes2(WorkItem):
    def __init__(self, userName=None, userIdUnique=None):
        super(WorkItemListTargetTypes2, self).__init__(userName, userIdUnique)
        inputObj = {"func": "listTypes"}
        self.workItem.api = XcalarApisT.XcalarApiTarget
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.targetInput = XcalarApiTargetInputT()
        self.workItem.input.targetInput.inputJson = json.dumps(inputObj)

    def output(self, executeOutput):
        return json.loads(
            executeOutput.output.outputResult.targetOutput.outputJson)

class WorkItemDatasetCreate(WorkItem):
    def __init__(self,
                 name,
                 loadArgs,
                 userName=None,
                 userIdUnique=None,
                 sessionName=None):
        super(WorkItemDatasetCreate, self).__init__(userName, userIdUnique,
                                                    sessionName)
        self.workItem.api = XcalarApisT.XcalarApiDatasetCreate
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.datasetCreateInput = XcalarApiDatasetCreateInputT()
        self.workItem.input.datasetCreateInput.dest = name
        self.workItem.input.datasetCreateInput.loadArgs = loadArgs

class WorkItemLoad(WorkItem):
    def __init__(self,
                 name,
                 loadArgs,
                 userName=None,
                 userIdUnique=None,
                 sessionName=None):
        super(WorkItemLoad, self).__init__(userName, userIdUnique, sessionName)
        self.workItem.api = XcalarApisT.XcalarApiBulkLoad
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.loadInput = XcalarApiBulkLoadInputT()
        self.workItem.input.loadInput.dest = name
        self.workItem.input.loadInput.loadArgs = loadArgs

    def output(self, executeOutput):
        return executeOutput.output.outputResult.loadOutput

class WorkItemDatasetGetMeta(WorkItem):
    def __init__(self,
                 name,
                 userName=None,
                 userIdUnique=None,
                 sessionName=None):
        super(WorkItemDatasetGetMeta, self).__init__(userName, userIdUnique,
                                                     sessionName)
        self.workItem.api = XcalarApisT.XcalarApiDatasetGetMeta
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.datasetGetMetaInput = XcalarApiDatasetGetMetaInputT(
        )
        self.workItem.input.datasetGetMetaInput.datasetName = name

    def output(self, executeOutput):
        return executeOutput.output.outputResult.datasetGetMetaOutput

class WorkItemDatasetUnload(WorkItem):
    def __init__(self,
                 namePattern,
                 userName=None,
                 userIdUnique=None,
                 sessionName=None):
        super(WorkItemDatasetUnload, self).__init__(userName, userIdUnique,
                                                    sessionName)
        self.workItem.api = XcalarApisT.XcalarApiDatasetUnload
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.datasetUnloadInput = XcalarApiDatasetUnloadInputT()
        self.workItem.input.datasetUnloadInput.datasetNamePattern = namePattern

    def output(self, executeOutput):
        return executeOutput.output.outputResult.datasetUnloadOutput

class WorkItemDatasetDelete(WorkItem):
    def __init__(self,
                 name,
                 userName=None,
                 userIdUnique=None,
                 sessionName=None):
        super(WorkItemDatasetDelete, self).__init__(userName, userIdUnique,
                                                    sessionName)
        self.workItem.api = XcalarApisT.XcalarApiDatasetDelete
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.datasetDeleteInput = XcalarApiDatasetDeleteInputT()
        self.workItem.input.datasetDeleteInput.datasetName = name

class WorkItemIndex(WorkItem):
    def __init__(self,
                 datasetName,
                 dstTableName,
                 keyNames,
                 dhtName,
                 ordering,
                 isTable,
                 fatptrPrefixName="",
                 delaySort=False,
                 keyFieldNames=None):
        super(WorkItemIndex, self).__init__()
        self.workItem.api = XcalarApisT.XcalarApiIndex
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.indexInput = XcalarApiIndexInputT()

        self.workItem.input.indexInput.source = datasetName
        self.workItem.input.indexInput.dest = dstTableName
        self.workItem.input.indexInput.key = []

        if (isinstance(keyNames, list)):
            for ii in range(len(keyNames)):
                key = XcalarApiKeyT()
                key.name = keyNames[ii]
                key.type = "DfUnknown"

                if keyFieldNames:
                    key.keyFieldName = keyFieldNames[ii]

                if (isinstance(ordering, list)):
                    key.ordering = XcalarOrderingTStr[ordering[ii]]
                else:
                    key.ordering = XcalarOrderingTStr[ordering]

                self.workItem.input.indexInput.key.append(key)
        else:
            key = XcalarApiKeyT()
            key.name = keyNames
            if keyFieldNames:
                key.keyFieldName = keyFieldNames

            key.type = "DfUnknown"
            key.ordering = XcalarOrderingTStr[ordering]

            self.workItem.input.indexInput.key.append(key)

        self.workItem.input.indexInput.dhtName = dhtName
        self.workItem.input.indexInput.prefix = fatptrPrefixName
        self.workItem.input.indexInput.delaySort = delaySort

    def output(self, executeOutput):
        return executeOutput.output.outputResult.indexOutput.tableName

class WorkItemGetRowNum(WorkItem):
    def __init__(self, srcTable, dstTable, newFieldName):
        super(WorkItemGetRowNum, self).__init__()
        self.workItem.api = XcalarApisT.XcalarApiGetRowNum
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.getRowNumInput = XcalarApiGetRowNumInputT()
        self.workItem.input.getRowNumInput.source = srcTable
        self.workItem.input.getRowNumInput.dest = dstTable
        self.workItem.input.getRowNumInput.newField = newFieldName

    def output(self, executeOutput):
        return executeOutput.output.outputResult.getRowNumOutput.tableName

class WorkItemMap(WorkItem):
    def __init__(self, srcTable, dstTable, evalStrs, newFieldNames):
        super(WorkItemMap, self).__init__()
        self.workItem.api = XcalarApisT.XcalarApiMap
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.mapInput = XcalarApiMapInputT()
        self.workItem.input.mapInput.source = srcTable
        self.workItem.input.mapInput.dest = dstTable

        self.workItem.input.mapInput.eval = []

        if (isinstance(evalStrs, list)):
            for ii in range(0, len(evalStrs)):
                eval = XcalarApiEvalT()
                eval.evalString = evalStrs[ii]
                eval.newField = newFieldNames[ii]

                self.workItem.input.mapInput.eval.append(eval)
        else:
            eval = XcalarApiEvalT()
            eval.evalString = evalStrs
            eval.newField = newFieldNames

            self.workItem.input.mapInput.eval.append(eval)

    def output(self, executeOutput):
        return executeOutput.output.outputResult.mapOutput.tableName

class WorkItemFilter(WorkItem):
    def __init__(self, srcTable, dstTable, filterStr):
        super(WorkItemFilter, self).__init__()
        self.workItem.api = XcalarApisT.XcalarApiFilter
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.filterInput = XcalarApiFilterInputT()
        self.workItem.input.filterInput.source = srcTable
        self.workItem.input.filterInput.dest = dstTable
        self.workItem.input.filterInput.eval = []
        eval = XcalarApiEvalT()
        eval.evalString = filterStr
        self.workItem.input.filterInput.eval.append(eval)

    def output(self, executeOutput):
        return executeOutput.output.outputResult.filterOutput.tableName

class WorkItemJoin(WorkItem):
    def __init__(self,
                 leftTable,
                 rightTable,
                 joinTable,
                 joinType,
                 leftColumns,
                 rightColumns,
                 filterString,
                 keepAllColumns=True):
        super(WorkItemJoin, self).__init__()
        self.workItem.api = XcalarApisT.XcalarApiJoin
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.joinInput = XcalarApiJoinInputT()

        self.workItem.input.joinInput.source = [leftTable, rightTable]
        self.workItem.input.joinInput.dest = joinTable

        self.workItem.input.joinInput.joinType = JoinOperatorTStr[joinType]
        self.workItem.input.joinInput.columns = [leftColumns, rightColumns]
        self.workItem.input.joinInput.evalString = filterString
        self.workItem.input.joinInput.keepAllColumns = keepAllColumns

    def output(self, executeOutput):
        return executeOutput.output.outputResult.joinOutput.tableName

class WorkItemUnion(WorkItem):
    def __init__(self, sources, dest, columns, dedup, unionType):
        super(WorkItemUnion, self).__init__()
        self.workItem.api = XcalarApisT.XcalarApiUnion
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.unionInput = XcalarApiUnionInputT()

        self.workItem.input.unionInput.source = sources
        self.workItem.input.unionInput.dest = dest
        self.workItem.input.unionInput.columns = columns
        self.workItem.input.unionInput.dedup = dedup
        self.workItem.input.unionInput.unionType = UnionOperatorTStr[unionType]

    def output(self, executeOutput):
        return executeOutput.output.outputResult.unionOutput.tableName

class WorkItemAggregate(WorkItem):
    def __init__(self,
                 srcTableName,
                 dstTableName,
                 evalStr,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemAggregate, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiAggregate
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.aggregateInput = XcalarApiAggregateInputT()

        self.workItem.input.aggregateInput.source = srcTableName
        self.workItem.input.aggregateInput.dest = dstTableName
        self.workItem.input.aggregateInput.eval = []
        eval = XcalarApiEvalT()
        eval.evalString = evalStr
        self.workItem.input.aggregateInput.eval.append(eval)

    def output(self, executeOutput):
        return json.loads(executeOutput.output.outputResult.aggregateOutput.
                          jsonAnswer)['Value']

class WorkItemGroupBy(WorkItem):
    def __init__(self,
                 srcTable,
                 dstTable,
                 evalStrs,
                 newFieldNames,
                 includeSrcTableSample=False,
                 groupAll=False):
        super(WorkItemGroupBy, self).__init__()
        self.workItem.api = XcalarApisT.XcalarApiGroupBy
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.groupByInput = XcalarApiGroupByInputT()
        self.workItem.input.groupByInput.source = srcTable
        self.workItem.input.groupByInput.dest = dstTable
        self.workItem.input.groupByInput.eval = []

        if (isinstance(evalStrs, list)):
            for ii in range(0, len(evalStrs)):
                eval = XcalarApiEvalT()
                eval.evalString = evalStrs[ii]
                eval.newField = newFieldNames[ii]

                self.workItem.input.groupByInput.eval.append(eval)
        else:
            eval = XcalarApiEvalT()
            eval.evalString = evalStrs
            eval.newField = newFieldNames

            self.workItem.input.groupByInput.eval.append(eval)

        self.workItem.input.groupByInput.includeSample = includeSrcTableSample
        self.workItem.input.groupByInput.icv = False
        self.workItem.input.groupByInput.newKeyField = ""
        self.workItem.input.groupByInput.groupAll = groupAll

    def output(self, executeOutput):
        return executeOutput.output.outputResult.groupByOutput.tableName

class WorkItemProject(WorkItem):
    def __init__(self, srcTable, dstTable, columns):
        super(WorkItemProject, self).__init__()
        self.workItem.api = XcalarApisT.XcalarApiProject
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.projectInput = XcalarApiProjectInputT()
        self.workItem.input.projectInput.source = srcTable
        self.workItem.input.projectInput.dest = dstTable
        self.workItem.input.projectInput.columns = columns

    def output(self, executeOutput):
        return executeOutput.output.outputResult.projectOutput.tableName

class WorkItemShutdown(WorkItem):
    def __init__(self):
        super(WorkItemShutdown, self).__init__()
        self.workItem.api = XcalarApisT.XcalarApiShutdown

class WorkItemUdfAdd(WorkItem):
    def __init__(self, moduleName, source, userName=None, userIdUnique=None):
        super(WorkItemUdfAdd, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiUdfAdd
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.udfAddUpdateInput = UdfModuleSrcT()
        self.workItem.input.udfAddUpdateInput.type = UdfTypeT.UdfTypePython
        self.workItem.input.udfAddUpdateInput.moduleName = moduleName
        self.workItem.input.udfAddUpdateInput.source = source

    def output(self, executeOutput):
        return executeOutput.output.outputResult.udfAddUpdateOutput

class WorkItemUdfUpdate(WorkItem):
    def __init__(self, moduleName, source):
        super(WorkItemUdfUpdate, self).__init__()
        self.workItem.api = XcalarApisT.XcalarApiUdfUpdate
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.udfAddUpdateInput = UdfModuleSrcT()
        self.workItem.input.udfAddUpdateInput.type = UdfTypeT.UdfTypePython
        self.workItem.input.udfAddUpdateInput.moduleName = moduleName
        self.workItem.input.udfAddUpdateInput.source = source

    def output(self, executeOutput):
        return executeOutput.output.outputResult.udfAddUpdateOutput

class WorkItemUdfDelete(WorkItem):
    def __init__(self, moduleName):
        super(WorkItemUdfDelete, self).__init__()
        self.workItem.api = XcalarApisT.XcalarApiUdfDelete
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.udfDeleteInput = XcalarApiUdfDeleteInputT()
        self.workItem.input.udfDeleteInput.moduleName = moduleName

class WorkItemUdfGet(WorkItem):
    def __init__(self, moduleName, userName=None, userIdUnique=None):
        super(WorkItemUdfGet, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiUdfGet
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.udfGetInput = XcalarApiUdfGetInputT()
        self.workItem.input.udfGetInput.moduleName = moduleName

    def output(self, getOutput):
        return getOutput.output.outputResult.udfGetOutput

class WorkItemListXdfs(WorkItem):
    def __init__(self,
                 fnNamePattern,
                 categoryPattern,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemListXdfs, self).__init__(userName, userIdUnique)

        self.workItem.api = XcalarApisT.XcalarApiListXdfs
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.listXdfsInput = XcalarApiListXdfsInputT()
        self.workItem.input.listXdfsInput.fnNamePattern = fnNamePattern
        self.workItem.input.listXdfsInput.categoryPattern = categoryPattern

    def output(self, getOutput):
        return getOutput.output.outputResult.listXdfsOutput

class WorkItemGetStats(WorkItem):
    def __init__(self, nodeId=0, userName=None, userIdUnique=None):
        super(WorkItemGetStats, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiGetStat
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.statInput = XcalarApiStatInputT()
        self.workItem.input.statInput.nodeId = nodeId

    def output(self, rsOutput):
        return rsOutput.output.outputResult.statOutput.stats

class WorkItemGetStatGroupIdMap(WorkItem):
    def __init__(self, nodeId=0, userName=None, userIdUnique=None):
        super(WorkItemGetStatGroupIdMap, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiGetStatGroupIdMap
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.statInput = XcalarApiStatInputT()
        self.workItem.input.statInput.nodeId = nodeId

    def output(self, rsOutput):
        return rsOutput.output.outputResult.statGroupIdMapOutput

class WorkItemExport(WorkItem):
    def __init__(self,
                 tableName,
                 driverName,
                 driverParams,
                 columns,
                 destDagNodeName="",
                 userName=None,
                 userIdUnique=None):
        super(WorkItemExport, self).__init__(userName, userIdUnique)
        self.workItem.input = XcalarApiInputT()
        self.workItem.api = XcalarApisT.XcalarApiExport
        self.workItem.input.exportInput = XcalarApiExportInputT()
        self.workItem.input.exportInput.source = tableName
        self.workItem.input.exportInput.driverName = driverName
        self.workItem.input.exportInput.driverParams = driverParams
        self.workItem.input.exportInput.columns = columns
        self.workItem.input.exportInput.dest = destDagNodeName

class WorkItemImportRetina(WorkItem):
    def __init__(self,
                 retinaName,
                 overwrite,
                 retinaTar=None,
                 retinaJsonStr=None,
                 userName=None,
                 userIdUnique=None,
                 udfUserName=None,
                 udfSessionName=None):
        super(WorkItemImportRetina, self).__init__(userName, userIdUnique)

        self.workItem.input = XcalarApiInputT()
        self.workItem.api = XcalarApisT.XcalarApiImportRetina

        self.workItem.input.importRetinaInput = XcalarApiImportRetinaInputT()
        self.workItem.input.importRetinaInput.retinaName = retinaName
        self.workItem.input.importRetinaInput.overwriteExistingUdf = overwrite

        if retinaTar:
            encodedRetina = base64.b64encode(retinaTar).decode("ascii")
            self.workItem.input.importRetinaInput.retinaCount = len(
                encodedRetina)
            self.workItem.input.importRetinaInput.retina = encodedRetina
            self.workItem.input.importRetinaInput.loadRetinaJson = False

        if retinaJsonStr:
            self.workItem.input.importRetinaInput.loadRetinaJson = True
            self.workItem.input.importRetinaInput.retinaJson = retinaJsonStr

        self.workItem.input.importRetinaInput.udfUserName = udfUserName
        self.workItem.input.importRetinaInput.udfSessionName = udfSessionName

    def output(self, executeOutput):
        return executeOutput.output.outputResult.importRetinaOutput

class WorkItemExportRetina(WorkItem):
    def __init__(self, retinaName, userName=None, userIdUnique=None):
        super(WorkItemExportRetina, self).__init__(userName, userIdUnique)

        self.workItem.input = XcalarApiInputT()
        self.workItem.api = XcalarApisT.XcalarApiExportRetina

        self.workItem.input.exportRetinaInput = XcalarApiExportRetinaInputT()
        self.workItem.input.exportRetinaInput.retinaName = retinaName

    def output(self, executeOutput):
        retinaStr = executeOutput.output.outputResult.exportRetinaOutput.retina

        return base64.b64decode(retinaStr)

class WorkItemMakeRetina(WorkItem):
    def __init__(self,
                 retinaName,
                 dstTables,
                 srcTables,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemMakeRetina, self).__init__(userName, userIdUnique)

        self.workItem.input = XcalarApiInputT()
        self.workItem.api = XcalarApisT.XcalarApiMakeRetina

        self.workItem.input.makeRetinaInput = XcalarApiMakeRetinaInputT()
        self.workItem.input.makeRetinaInput.retinaName = retinaName
        self.workItem.input.makeRetinaInput.numTables = len(dstTables)
        self.workItem.input.makeRetinaInput.tableArray = dstTables
        self.workItem.input.makeRetinaInput.numSrcTables = len(srcTables)
        self.workItem.input.makeRetinaInput.srcTables = srcTables

class WorkItemListRetina(WorkItem):
    def __init__(self, userName=None, userIdUnique=None):
        super(WorkItemListRetina, self).__init__(userName, userIdUnique)

        self.workItem.input = XcalarApiInputT()
        self.workItem.api = XcalarApisT.XcalarApiListRetinas
        self.workItem.input.listRetinasInput = XcalarApiListRetinasInputT()
        self.workItem.input.listRetinasInput.namePattern = "*"

    def output(self, executeOutput):
        return executeOutput.output.outputResult.listRetinasOutput

class WorkItemDeleteRetina(WorkItem):
    def __init__(self, retinaName, userName=None, userIdUnique=None):
        super(WorkItemDeleteRetina, self).__init__(userName, userIdUnique)

        self.workItem.input = XcalarApiInputT()
        self.workItem.api = XcalarApisT.XcalarApiDeleteRetina

        self.workItem.input.deleteRetinaInput = XcalarApiDeleteRetinaInputT()
        self.workItem.input.deleteRetinaInput.delRetInput = retinaName

class WorkItemAppSet(WorkItem):
    def __init__(self,
                 name,
                 hostType,
                 execStr,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemAppSet, self).__init__(userName, userIdUnique)
        self.workItem.input = XcalarApiInputT()
        self.workItem.api = XcalarApisT.XcalarApiAppSet

        self.workItem.input.appSetInput = XcalarApiAppSetInputT()
        self.workItem.input.appSetInput.name = name
        self.workItem.input.appSetInput.hostType = hostType
        self.workItem.input.appSetInput.execStr = execStr

class WorkItemAppRun(WorkItem):
    def __init__(self, name, isGlobal, inStr, userName=None,
                 userIdUnique=None):
        super(WorkItemAppRun, self).__init__(userName, userIdUnique)
        self.workItem.input = XcalarApiInputT()
        self.workItem.api = XcalarApisT.XcalarApiAppRun

        self.workItem.input.appRunInput = XcalarApiAppRunInputT()
        self.workItem.input.appRunInput.name = name
        self.workItem.input.appRunInput.isGlobal = isGlobal
        self.workItem.input.appRunInput.inStr = inStr

class WorkItemAppReap(WorkItem):
    def __init__(self, groupId, cancel=False, userName=None,
                 userIdUnique=None):
        super(WorkItemAppReap, self).__init__(userName, userIdUnique)
        self.workItem.input = XcalarApiInputT()
        self.workItem.api = XcalarApisT.XcalarApiAppReap

        self.workItem.input.appReapInput = XcalarApiAppReapInputT()
        self.workItem.input.appReapInput.appGroupId = groupId
        self.workItem.input.appReapInput.cancel = cancel

class WorkItemRuntimeSetParam(WorkItem):
    def __init__(self, runtimeSetParams, userName=None, userIdUnique=None):
        super(WorkItemRuntimeSetParam, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiRuntimeSetParam
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.runtimeSetParamInput = XcalarApiRuntimeSetParamInputT(
        )
        self.workItem.input.runtimeSetParamInput = runtimeSetParams

class WorkItemRuntimeGetParam(WorkItem):
    def __init__(self, userName=None, userIdUnique=None):
        super(WorkItemRuntimeGetParam, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiRuntimeGetParam

    def output(self, executeOutput):
        return executeOutput.output.outputResult.runtimeGetParamOutput

class WorkItemSetConfigParam(WorkItem):
    def __init__(self, paramName, paramValue, userName=None,
                 userIdUnique=None):
        super(WorkItemSetConfigParam, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiSetConfigParam
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.setConfigParamInput = XcalarApiSetConfigParamInputT(
        )
        self.workItem.input.setConfigParamInput.paramName = paramName
        self.workItem.input.setConfigParamInput.paramValue = paramValue

class WorkItemGetConfigParams(WorkItem):
    def __init__(self, userName=None, userIdUnique=None):
        super(WorkItemGetConfigParams, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiGetConfigParams

    def output(self, executeOutput):
        return executeOutput.output.outputResult.getConfigParamsOutput

class WorkItemGetRetina(WorkItem):
    def __init__(self, retinaName, userName=None, userIdUnique=None):
        super(WorkItemGetRetina, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiGetRetina
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.getRetinaInput = XcalarApiGetRetinaInputT()
        self.workItem.input.getRetinaInput.retInput = retinaName

    def output(self, executeOutput):
        return executeOutput.output.outputResult.getRetinaOutput

class WorkItemGetRetinaJson(WorkItem):
    def __init__(self, retinaName, userName=None, userIdUnique=None):
        super(WorkItemGetRetinaJson, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiGetRetinaJson
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.getRetinaJsonInput = XcalarApiGetRetinaJsonInputT()
        self.workItem.input.getRetinaJsonInput.retinaName = retinaName

    def output(self, executeOutput):
        return executeOutput.output.outputResult.getRetinaJsonOutput

class WorkItemUpdateRetina(WorkItem):
    def __init__(self,
                 retinaName,
                 retinaJson,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemUpdateRetina, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiUpdateRetina

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.updateRetinaInput = XcalarApiUpdateRetinaInputT()
        self.workItem.input.updateRetinaInput.retinaName = retinaName
        self.workItem.input.updateRetinaInput.retinaJson = retinaJson

class WorkItemUpdateTable(WorkItem):
    def __init__(self,
                 srcTableNames,
                 dstTableNames,
                 times=[],
                 dropSrc=False,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemUpdateTable, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiUpdate

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.updateInput = XcalarApiUpdateInputT()
        self.workItem.input.updateInput.updates = []

        if (isinstance(srcTableNames, list)):
            for ii in range(0, len(srcTableNames)):
                update = XcalarApiUpdateTableInputT()
                update.source = srcTableNames[ii]
                update.dest = dstTableNames[ii]

                if len(times) > 0:
                    update.unixTS = times[ii]
                else:
                    update.unixTS = 0
                update.dropSrc = dropSrc
                self.workItem.input.updateInput.updates.append(update)
        else:
            update = XcalarApiUpdateTableInputT()
            update.source = srcTableNames
            update.dest = dstTableNames

            if isinstance(times, list):
                update.unixTS = 0
            else:
                update.unixTS = times
            update.dropSrc = dropSrc

            self.workItem.input.updateInput.updates.append(update)

        def output(self, executeOutput):
            return executeOutput.output.outputResult.updateOutput.batchId

class WorkItemPublishTable(WorkItem):
    def __init__(self,
                 srcTableName,
                 dstTableName,
                 time=0,
                 dropSrc=False,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemPublishTable, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiPublish

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.publishInput = XcalarApiPublishInputT()
        self.workItem.input.publishInput.source = srcTableName
        self.workItem.input.publishInput.dest = dstTableName
        self.workItem.input.publishInput.unixTS = time
        self.workItem.input.publishInput.dropSrc = dropSrc

class WorkItemSelectTable(WorkItem):
    def __init__(self,
                 srcTableName,
                 dstTableName,
                 maxBatchId=-1,
                 minBatchId=-1,
                 filterString="",
                 columns=[],
                 limitRows=0,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemSelectTable, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiSelect

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.selectInput = XcalarApiSelectInputT()
        self.workItem.input.selectInput.source = srcTableName
        self.workItem.input.selectInput.dest = dstTableName
        self.workItem.input.selectInput.minBatchId = minBatchId
        self.workItem.input.selectInput.maxBatchId = maxBatchId
        self.workItem.input.selectInput.filterString = filterString
        self.workItem.input.selectInput.columns = columns
        self.workItem.input.selectInput.limitRows = limitRows

    def output(self, executeOutput):
        return executeOutput.output.outputResult.selectOutput.tableName

class WorkItemUnpublishTable(WorkItem):
    def __init__(self,
                 srcTableName,
                 inactivateOnly=False,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemUnpublishTable, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiUnpublish

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.unpublishInput = XcalarApiUnpublishInputT()
        self.workItem.input.unpublishInput.source = srcTableName
        self.workItem.input.unpublishInput.inactivateOnly = inactivateOnly

class WorkItemRestoreTable(WorkItem):
    def __init__(self, srcTableName, userName=None, userIdUnique=None):
        super(WorkItemRestoreTable, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiRestoreTable

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.restoreTableInput = XcalarApiRestoreTableInputT()
        self.workItem.input.restoreTableInput.publishedTableName = srcTableName

    def output(self, executeOutput):
        return executeOutput.output.outputResult.restoreTableOutput

class WorkItemCoalesceTable(WorkItem):
    def __init__(self, srcTableName, userName=None, userIdUnique=None):
        super(WorkItemCoalesceTable, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiCoalesce

        self.workItem.input = XcalarApiInputT()
        self.workItem.input.coalesceInput = XcalarApiCoalesceInputT()
        self.workItem.input.coalesceInput.source = srcTableName

class WorkItemSynthesize(WorkItem):
    def __init__(self,
                 src,
                 dstTable,
                 columns,
                 userName=None,
                 userIdUnique=None):
        super(WorkItemSynthesize, self).__init__(userName, userIdUnique)
        self.workItem.api = XcalarApisT.XcalarApiSynthesize
        self.workItem.input = XcalarApiInputT()
        self.workItem.input.synthesizeInput = XcalarApiSynthesizeInputT()
        self.workItem.input.synthesizeInput.source = src
        self.workItem.input.synthesizeInput.dest = dstTable
        self.workItem.input.synthesizeInput.columns = columns
        self.workItem.input.synthesizeInput.sameSession = True

    def output(self, executeOutput):
        return executeOutput.output.outputResult.synthesizeOutput


# -*- coding: utf-8 -*-
import logging
import werkzeug

from openerp import http
from openerp.http import request
from openerp.addons.payment_mollie_official.controllers.main import (
    MollieController,
)

_logger = logging.getLogger(__name__)

class ProductSubscriptionMollieController(MollieController):
    @http.route(
        ["/payment/mollie/redirect"], type="http", auth="none", methods=["GET"]
    )
    def mollie_redirect(self, **post):
        route = "/shop/payment/validate"
        pay_tx_obj = request.env["payment.transaction"]
        pay_tx_obj.sudo().form_feedback(post, "mollie")
        orderid = post["reference"]
        tx = pay_tx_obj.sudo().search([("reference", "=", orderid)])
        if tx and tx.product_subscription_request_id:
            _logger.info("mollie redirect : tx status is %s" % tx.state)
            if tx.state == "done":
                route = "/render/online_payment_success"
            elif tx.state == "cancel":
                route = "/render/online_payment_cancel"
            elif tx.state == "error":
                route = "/render/online_payment_error"
        elif tx:
            _logger.info("mollie redirect : tx status is %s" % tx.state)
        else:
            _logger.info(
                "mollie redirect : no tx found for reference %s" % orderid
            )
        return werkzeug.utils.redirect(route)


import os
import os.path
import sys
import datetime
import dateutil.parser
import requests
import radix
import gzip
from io import StringIO
import bz2
import math
import ipaddress
import urllib.request

##
# Configuration
##

ripe_base_url	= 'https://ftp.ripe.net/pub/stats'
afrinic_sub	= 'afrinic'
arin_sub	= 'arin'
apnic_sub	= 'apnic'
lacnic_sub	= 'lacnic'
ripencc_sub	= 'ripencc'

tmpdir		= '/tmp'

# Retrieve NRO stats for the AfriNIC region for the
# specified date
def fetch_afrinic_nrostats(day):
	# Where do the archives at the RIPE NCC
	start_date = datetime.date(2005,3,3)

	if day < start_date:
		# Write an empty file if the day requested
		# is not in the archives
		tmpfile = '{}/afrinic-empty'.format(tmpdir)

		fd = open(tmpfile, 'w')
		fd.close()

		return tmpfile

	# The archive at the RIPE NCC misses data for
	# certain dates which we replace with newer
	# data
	missing_dateranges = []
	missing_dateranges.append((datetime.date(2011,1,1), datetime.date(2011,5,15), datetime.date(2011,5,16)))
	missing_dateranges.append((datetime.date(2014,12,31), datetime.date(2015,1,6), datetime.date(2015,1,7)))
	missing_dateranges.append((datetime.date(2015,12,31), datetime.date(2016,1,4), datetime.date(2016,1,5)))
	missing_dateranges.append((datetime.date(2016,12,31), datetime.date(2016,12,31), datetime.date(2017,1,1)))
	missing_dateranges.append((datetime.date(2017,12,31), datetime.date(2018,1,3), datetime.date(2018,1,4)))

	for daterange in missing_dateranges:
		if day >= daterange[0] and day <= daterange[1]:
			print('{} is missing for AfriNIC, replacing it by {}'.format(day, daterange[2]))
			day = daterange[2]

	# Some days ended up in different years, note
	# these exceptions
	exception_dates = []
	exception_dates.append((datetime.date(2012,12,31), 2013))

	year = day.year

	for exception in exception_dates:
		if day == exception[0]:
			print('{} is in another year directory for AfriNIC, adjusting year directory to {}'.format(day, exception[1]))
			year = exception[1]

	afrinic_url = '{}/{}/{}/delegated-afrinic-{:04d}{:02d}{:02d}'.format(ripe_base_url, afrinic_sub, year, day.year, day.month, day.day)

	sys.stdout.write('Fetching AfriNIC data from {} ... '.format(afrinic_url))
	sys.stdout.flush()

	tmpfile = '{}/afrinic-{:04d}{:02d}{:02d}'.format(tmpdir, day.year, day.month, day.day)

	try:
		urllib.request.urlretrieve(afrinic_url, tmpfile)
	except:
		raise Exception('Failed to download {} to {}'.format(afrinic_url, tmpfile))

	print('OK ({})'.format(tmpfile))

	return tmpfile

# Retrieve NRO stats for the ARIN region for
# the specified date
def fetch_arin_nrostats(day):
	# Where do the archives at the RIPE NCC
	start_date = datetime.date(2003,11,20)

	if day < start_date:
		# Write an empty file if the day requested
		# is not in the archives
		tmpfile = '{}/arin-empty'.format(tmpdir)

		fd = open(tmpfile, 'w')
		fd.close()

		return tmpfile

	# The ARIN stats are a bit of a mess unfortunately
	# so determining the right URL requires some black
	# magic. We assume that ARIN always has a full year
	# of data in the current dir + any data from the
	# running year, and that everything else is under
	# the "archive" directory. Furthermore, the file
	# format and file naming scheme change as of
	# March 5, 2013. To top it of, until September 30,
	# 2007, the files are gzipped, after that, they're
	# just plain ASCII text.
    #
    # 2020-04-30:
    # Turns out that from 2017 onward the ARIN files
    # do reside in a single directory and not in the
    # "archive" directory... go figure...
	subdir = None
	filename_prefix = None
	filename_postfix = None

	today = datetime.date.today()

	if day.year < 2017:
		subdir = 'archive/{}'.format(day.year)

	# The archive at the RIPE NCC misses data for
	# certain dates which we replace with newer
	# data
	missing_dateranges = []
	missing_dateranges.append((datetime.date(2019,8,25), datetime.date(2019,8,25), datetime.date(2019,8,26)))

	for daterange in missing_dateranges:
		if day >= daterange[0] and day <= daterange[1]:
			print('{} is missing for ARIN, replacing it by {}'.format(day, daterange[2]))
			day = daterange[2]

	if day >= datetime.date(2013,3,5):
		filename_prefix = 'delegated-arin-extended'
	else:
		filename_prefix = 'delegated-arin'

	if day <= datetime.date(2007,9,30):
		filename_postfix = '.gz'

	arin_url = '{}/{}'.format(ripe_base_url, arin_sub)

	if subdir is not None:
		arin_url += '/{}'.format(subdir)

	arin_url += '/{}-{:04d}{:02d}{:02d}'.format(filename_prefix, day.year, day.month, day.day)

	if filename_postfix is not None:
		arin_url += filename_postfix

	sys.stdout.write('Fetching ARIN data from {} ... '.format(arin_url))
	sys.stdout.flush()

	tmpfile = '{}/arin-{:04d}{:02d}{:02d}'.format(tmpdir, day.year, day.month, day.day)

	if filename_postfix is not None:
		tmpfile += filename_postfix

	try:
		urllib.request.urlretrieve(arin_url, tmpfile)
	except:
		raise Exception('Failed to download {} to {}'.format(arin_url, tmpfile))

	print('OK ({})'.format(tmpfile))

	return tmpfile

# Retrieve NRO stats for the APNIC region for
# the specified date
def fetch_apnic_nrostats(day):
	# Where do the archives at the RIPE NCC
	start_date = datetime.date(2001,5,1)

	if day < start_date:
		# Write an empty file if the day requested
		# is not in the archives
		tmpfile = '{}/apnic-empty'.format(tmpdir)

		fd = open(tmpfile, 'w')
		fd.close()

		return tmpfile

	# The archive at the RIPE NCC misses data for
	# certain dates which we replace with newer
	# data
	missing_dateranges = []
	missing_dateranges.append((datetime.date(2001,5,2), datetime.date(2001,5,31), datetime.date(2001,6,1)))
	missing_dateranges.append((datetime.date(2001,6,2), datetime.date(2001,8,31), datetime.date(2001,9,1)))
	missing_dateranges.append((datetime.date(2001,9,2), datetime.date(2001,9,30), datetime.date(2001,10,1)))
	missing_dateranges.append((datetime.date(2001,10,2), datetime.date(2001,10,31), datetime.date(2001,11,1)))
	missing_dateranges.append((datetime.date(2001,11,2), datetime.date(2001,11,30), datetime.date(2001,12,1)))
	missing_dateranges.append((datetime.date(2001,12,2), datetime.date(2001,12,31), datetime.date(2002,1,1)))
	missing_dateranges.append((datetime.date(2002,1,2), datetime.date(2002,1,31), datetime.date(2002,2,1)))
	missing_dateranges.append((datetime.date(2002,2,2), datetime.date(2002,2,28), datetime.date(2002,3,1)))
	missing_dateranges.append((datetime.date(2002,3,2), datetime.date(2002,3,31), datetime.date(2002,4,1)))
	missing_dateranges.append((datetime.date(2002,4,2), datetime.date(2002,4,30), datetime.date(2002,5,1)))
	missing_dateranges.append((datetime.date(2002,5,2), datetime.date(2002,5,31), datetime.date(2002,6,1)))
	missing_dateranges.append((datetime.date(2002,6,2), datetime.date(2002,6,30), datetime.date(2002,7,1)))
	missing_dateranges.append((datetime.date(2002,7,2), datetime.date(2002,7,31), datetime.date(2002,8,1)))
	missing_dateranges.append((datetime.date(2002,8,2), datetime.date(2002,8,31), datetime.date(2002,9,1)))
	missing_dateranges.append((datetime.date(2002,9,2), datetime.date(2002,9,30), datetime.date(2002,10,1)))
	missing_dateranges.append((datetime.date(2002,10,2), datetime.date(2002,10,31), datetime.date(2002,11,1)))
	missing_dateranges.append((datetime.date(2002,11,2), datetime.date(2002,11,30), datetime.date(2002,12,1)))
	missing_dateranges.append((datetime.date(2002,12,2), datetime.date(2002,12,31), datetime.date(2003,1,1)))
	missing_dateranges.append((datetime.date(2003,1,2), datetime.date(2003,1,31), datetime.date(2003,2,1)))
	missing_dateranges.append((datetime.date(2003,2,2), datetime.date(2003,2,28), datetime.date(2003,3,1)))
	missing_dateranges.append((datetime.date(2003,3,2), datetime.date(2003,3,31), datetime.date(2003,4,1)))
	missing_dateranges.append((datetime.date(2003,4,2), datetime.date(2003,4,30), datetime.date(2003,5,1)))
	missing_dateranges.append((datetime.date(2003,5,2), datetime.date(2003,5,7), datetime.date(2003,5,8)))

	for daterange in missing_dateranges:
		if day >= daterange[0] and day <= daterange[1]:
			print('{} is missing for APNIC, replacing it by {}'.format(day, daterange[2]))
			day = daterange[2]

	# Some days ended up in different years, note
	# these exceptions
	exception_dates = []
	exception_dates.append((datetime.date(2010,12,31), 2011))
	exception_dates.append((datetime.date(2011,12,31), 2012))
	exception_dates.append((datetime.date(2012,12,31), 2013))
	exception_dates.append((datetime.date(2013,12,31), 2014))
	exception_dates.append((datetime.date(2014,12,31), 2015))
	exception_dates.append((datetime.date(2015,12,31), 2016))
	exception_dates.append((datetime.date(2016,12,31), 2017))
	exception_dates.append((datetime.date(2017,12,31), 2018))
	exception_dates.append((datetime.date(2018,12,31), 2019))

	year = day.year

	for exception in exception_dates:
		if day == exception[0]:
			print('{} is in another year directory for APNIC, adjusting year directory to {}'.format(day, exception[1]))
			year = exception[1]

	# The filename for APNIC is different before
	# October 9, 2003
	filename = None

	if day <= datetime.date(2003,10,8):
		filename = 'apnic-{:04d}-{:02d}-{:02d}.gz'.format(day.year, day.month, day.day)
	else:
		filename = 'delegated-apnic-{:04d}{:02d}{:02d}.gz'.format(day.year, day.month, day.day)

	apnic_url = '{}/{}/{}/{}'.format(ripe_base_url, apnic_sub, year, filename)

	sys.stdout.write('Fetching APNIC data from {} ... '.format(apnic_url))
	sys.stdout.flush()

	tmpfile = '{}/apnic-{:04d}{:02d}{:02d}.gz'.format(tmpdir, day.year, day.month, day.day)

	try:
		urllib.request.urlretrieve(apnic_url, tmpfile)
	except:
		raise Exception('Failed to download {} to {}'.format(apnic_url, tmpfile))

	print('OK ({})'.format(tmpfile))

	return tmpfile

# Retrieve NRO stats for the LACNIC region for
# the specified date
def fetch_lacnic_nrostats(day):
	# Where do the archives at the RIPE NCC
	start_date = datetime.date(2004,1,1)

	if day < start_date:
		# Write an empty file if the day requested
		# is not in the archives
		tmpfile = '{}/lacnic-empty'.format(tmpdir)

		fd = open(tmpfile, 'w')
		fd.close()

		return tmpfile

	# The archive at the RIPE NCC misses data for
	# certain dates which we replace with newer
	# data
	missing_dateranges = []
	missing_dateranges.append((datetime.date(2018,9,26), datetime.date(2018,9,26), datetime.date(2018,9,27)))
	missing_dateranges.append((datetime.date(2018,11,10), datetime.date(2018,11,10), datetime.date(2018,11,11)))
	missing_dateranges.append((datetime.date(2019,12,21), datetime.date(2019,12,21), datetime.date(2019,12,22)))
	missing_dateranges.append((datetime.date(2020,4,22), datetime.date(2020,4,22), datetime.date(2020,4,23)))

	for daterange in missing_dateranges:
		if day >= daterange[0] and day <= daterange[1]:
			print('{} is missing for LACNIC, replacing it by {}'.format(day, daterange[2]))
			day = daterange[2]

	# The LACNIC data should all be in a single
	# directory, without any exceptions
	lacnic_url = '{}/{}/delegated-lacnic-{:04d}{:02d}{:02d}'.format(ripe_base_url, lacnic_sub, day.year, day.month, day.day)

	sys.stdout.write('Fetching LACNIC data from {} ... '.format(lacnic_url))
	sys.stdout.flush()

	tmpfile = '{}/lacnic-{:04d}{:02d}{:02d}'.format(tmpdir, day.year, day.month, day.day)

	try:
		urllib.request.urlretrieve(lacnic_url, tmpfile)
	except:
		raise Exception('Failed to download {} to {}'.format(lacnic_url, tmpfile))

	print('OK ({})'.format(tmpfile))

	return tmpfile

# Retrieve NRO stats for the RIPE region for
# the specified date
def fetch_ripencc_nrostats(day):
	# Where do the archives at the RIPE NCC
	start_date = datetime.date(2003,11,26)

	if day < start_date:
		# Write an empty file if the day requested
		# is not in the archives
		tmpfile = '{}/ripencc-empty'.format(tmpdir)

		fd = open(tmpfile, 'w')
		fd.close()

		return tmpfile

	# The RIPE NCC's archives are split up
	# by year with, annoyingly, a single
	# exception for Jan. 1st, 2004
	year = day.year

	if day == datetime.date(2004,1,1):
		year = 2003
		print('RIPE data for {} lives in year {}'.format(day, year))

	ripencc_url = '{}/{}/{}/delegated-ripencc-{:04d}{:02d}{:02d}.bz2'.format(ripe_base_url, ripencc_sub, year, day.year, day.month, day.day)

	sys.stdout.write('Fetching RIPE data from {} ... '.format(ripencc_url))
	sys.stdout.flush()

	tmpfile = '{}/ripencc-{:04d}{:02d}{:02d}.bz2'.format(tmpdir, day.year, day.month, day.day)

	try:
		urllib.request.urlretrieve(ripencc_url, tmpfile)
	except:
		raise Exception('Failed to download {} to {}'.format(ripencc_url, tmpfile))

	print('OK ({})'.format(tmpfile))

	return tmpfile

# Added the specified prefix to the radix tree and check if
# it might already be present
def add_prefix_to_radix(radix, prefix, prefix_info):
	# Find out if the prefix is already there
	rnode = radix.search_exact(prefix)

	if rnode is None:
		rnode = radix.add(prefix)
	else:
		print('Warning: {} already in radix tree'.format(prefix))

	info_arr = rnode.data.get('info', [])
	info_arr.append(prefix_info)
	rnode.data['info'] = info_arr

# Adding an IPv4 address to the radix tree is tricky, since the
# number of IP addresses in the block in the NRO stats is not
# necessarily a power of two. We first test if it is a power of
# 2, and if it is not, we need to factorize it to get the
# subprefixes
def add_v4_block_to_radix(v4radix, prefix_addr, addrcount, prefix_info):
	prefix_addr = ipaddress.ip_address(prefix_addr)

	while addrcount > 0:
		prefix_size = 32 - int(math.log(addrcount, 2))

		prefix_found = False

		while not prefix_found:
			try:
				test_nw = ipaddress.IPv4Network('{}/{}'.format(prefix_addr, prefix_size))
				prefix_found = True
			except ValueError as v:
				prefix_size += 1

		prefix = '{}/{}'.format(prefix_addr, prefix_size)

		used_addrcount = int(math.pow(2, 32 - prefix_size))
		add_prefix_to_radix(v4radix, prefix, prefix_info + (used_addrcount,))

		addrcount -= used_addrcount
		prefix_addr += used_addrcount

# This is a very crude parser that will only process NRO stat lines
# that are for IPv4 or IPv6 prefixes and the contain exactly 7 fields
# separated by a '|' sign.
def parse_nro_stats(fd, v4radix, v6radix):
	for line in fd:
		if type(line) is bytes:
			line = line.decode('utf8')
		line = line.strip('\r').strip('\n')
		fields = line.split('|')

		if len(fields) < 7 or (fields[2] != 'ipv4' and fields[2] != 'ipv6'):
			continue

		# We put the following information about the
		# prefix in the tuple:
		#  - RIR name
		#  - country
		#  - date of modification/assignment
		#  - prefix status
		#  - original starting block
		#  - original block IP count
		prefix_info = (fields[0], fields[1], fields[5], fields[6], fields[3], int(fields[4]))

		prefix_size = 0
		prefix_addr = fields[3]

		if fields[2] == 'ipv4':
			addrcount = int(fields[4])
			add_v4_block_to_radix(v4radix, prefix_addr, addrcount, prefix_info)
		elif fields[2] == 'ipv6':
			# IPv6 entries just specify the prefix
			# size
			prefix_size = int(fields[4])
			prefix = '{}/{}'.format(prefix_addr, prefix_size)
			add_prefix_to_radix(v6radix, prefix, prefix_info)

def open_nro_stats_file(filename):
	if filename.endswith('.gz'):
		return gzip.open(filename, 'r')
	elif filename.endswith('.bz2'):
		return bz2.BZ2File(filename, 'r')
	else:
		return open(filename, 'r')

##
# Main entry point
##

def main():
	if len(sys.argv) != 3:
		print('Usage: merge-nrostats.py <date> <output directory>')
		return

	day = dateutil.parser.parse(sys.argv[1]).date()
	outdir = sys.argv[2]

	v4radix = radix.Radix()
	v6radix = radix.Radix()

	afrinic_file	= None
	arin_file	= None
	apnic_file	= None
	lacnic_file	= None
	ripencc_file	= None

	try:
		afrinic_file	= fetch_afrinic_nrostats(day)
		arin_file	= fetch_arin_nrostats(day)
		apnic_file	= fetch_apnic_nrostats(day)
		lacnic_file	= fetch_lacnic_nrostats(day)
		ripencc_file	= fetch_ripencc_nrostats(day)

		# AfriNIC
		sys.stdout.write('Parsing AfriNIC NRO statistics for {} ... '.format(day))
		sys.stdout.flush()

		afrinic_fd = open_nro_stats_file(afrinic_file)

		parse_nro_stats(afrinic_fd, v4radix, v6radix)

		afrinic_fd.close()

		print('OK')

		# ARIN
		sys.stdout.write('Parsing ARIN NRO statistics for {} ... '.format(day))
		sys.stdout.flush()

		arin_fd = open_nro_stats_file(arin_file)

		parse_nro_stats(arin_fd, v4radix, v6radix)

		arin_fd.close()

		print('OK')

		# APNIC
		sys.stdout.write('Parsing APNIC NRO statistics for {} ... '.format(day))
		sys.stdout.flush()

		apnic_fd = open_nro_stats_file(apnic_file)

		parse_nro_stats(apnic_fd, v4radix, v6radix)

		apnic_fd.close()

		print('OK')

		# LACNIC
		sys.stdout.write('Parsing LACNIC NRO statistics for {} ... '.format(day))
		sys.stdout.flush()

		lacnic_fd = open_nro_stats_file(lacnic_file)

		parse_nro_stats(lacnic_fd, v4radix, v6radix)

		lacnic_fd.close()

		print('OK')

		# RIPE NCC
		sys.stdout.write('Parsing RIPE NCC NRO statistics for {} ... '.format(day))
		sys.stdout.flush()

		ripencc_fd = open_nro_stats_file(ripencc_file)

		parse_nro_stats(ripencc_fd, v4radix, v6radix)

		ripencc_fd.close()

		print('OK')

		# Write the consolidated NRO statistics to file
		v4_file = '{}/nrostats-{:04d}{:02d}{:02d}-v4.csv'.format(outdir, day.year, day.month, day.day)
		v6_file = '{}/nrostats-{:04d}{:02d}{:02d}-v6.csv'.format(outdir, day.year, day.month, day.day)

		v4_fd = open(v4_file, 'w')
		v6_fd = open(v6_file, 'w')

		v4_fd.write('prefix,rir,date,country_code,status,block_start,block_ip_count\n')
		v6_fd.write('prefix,rir,date,country_code,status\n')

		sys.stdout.write('Writing IPv4 NRO statistics to {} ... '.format(v4_file))
		sys.stdout.flush()

		for rnode in v4radix:
			prefix = rnode.prefix

			for info in rnode.data['info']:
				rir 		= info[0]
				date 		= info[1]
				country 	= info[2]
				prefix_status	= info[3]
				block_start	= info[4]
				block_ip_count	= info[5]

				v4_fd.write('{},{},{},{},{},{},{}\n'.format(prefix, rir, country, date, prefix_status, block_start, block_ip_count))

		print('OK')

		sys.stdout.write('Writing IPv6 NRO statistics to {} ... '.format(v6_file))
		sys.stdout.flush()

		for rnode in v6radix:
			prefix = rnode.prefix

			for info in rnode.data['info']:
				rir 		= info[0]
				date 		= info[1]
				country 	= info[2]
				prefix_status	= info[3]

				v6_fd.write('{},{},{},{},{}\n'.format(prefix, rir, country, date, prefix_status))

		print('OK')

		v4_fd.close()
		v6_fd.close()
	except Exception as e:
		print('Failed to retrieve and merge NRO statistics ({})'.format(e))
	finally:
		if afrinic_file is not None:
			os.unlink(afrinic_file)
		if arin_file is not None:
			os.unlink(arin_file)
		if apnic_file is not None:
			os.unlink(apnic_file)
		if lacnic_file is not None:
			os.unlink(lacnic_file)
		if ripencc_file is not None:
			os.unlink(ripencc_file)

	return

if __name__ == "__main__":
	main()


import cv2
import numpy as np
import matplotlib.pyplot as plt

def process_state_image(state):
    state = cv2.cvtColor(state, cv2.COLOR_BGR2GRAY)
    state = state.astype(float)
    state /= 255.0
    return state

def generate_state_frame_stack_from_queue(deque):
    frame_stack = np.array(deque)
    # Move stack dimension to the channel dimension (stack, x, y) -> (x, y, stack)
    return np.transpose(frame_stack, (1, 2, 0))

def plotLearning(x, scores, epsilons, filename, lines=None):
    fig = plt.figure()
    ax = fig.add_subplot(111, label="1")
    ax2 = fig.add_subplot(111, label="2", frame_on=False)

    ax.plot(x, epsilons, color="C0")
    ax.set_xlabel("Game", color="C0")
    ax.set_ylabel("Epsilon", color="C0")
    ax.tick_params(axis='x', colors="C0")
    ax.tick_params(axis='y', colors="C0")

    N = len(scores)
    running_avg = np.empty(N)
    for t in range(N):
	    running_avg[t] = np.mean(scores[max(0, t-20):(t+1)])

    ax2.scatter(x, running_avg, color="C1")
    ax2.axes.get_xaxis().set_visible(False)
    ax2.yaxis.tick_right()
    ax2.set_ylabel('Score', color="C1")
    ax2.yaxis.set_label_position('right')
    ax2.tick_params(axis='y', colors="C1")

    if lines is not None:
        for line in lines:
            plt.axvline(x=line)

    plt.savefig(filename)


from random import choice, sample

from django.contrib.auth.decorators import login_required
from django.core.files import File
from django.db.models import QuerySet
from django.shortcuts import render, redirect

# Create your views here.
from Licencjat.settings import STATICFILES_DIRS
from Models.models import PhotoModel, UserModel, SelectionModel, AlgorithmModel
from Pages import algorithms
from .apps import get_client_ip
from numpy import zeros

def homepageView(request, *args, **kwargs):
    user_ip = get_client_ip(request)
    context = {}
    if not UserModel.objects.filter(ip__exact=user_ip):
        user_model_instance = UserModel()
        user_model_instance.ip = user_ip
        user_model_instance.save()
        photos = sample(list(PhotoModel.objects.all()), 100)
        for photo in photos:
            SelectionModel.objects.create(user=user_model_instance, photo=photo)
    else:
        user_model_instance = UserModel.objects.filter(ip__exact=user_ip)[
            0]
    if not user_model_instance.has_finished:
        selection = \
            SelectionModel.objects.filter(user__exact=user_model_instance,
                                          selection=None)[0]
        if request.method == 'POST':
            if 'likebutton' in request.POST:
                selection.selection = True
            elif 'dislikebutton' in request.POST:
                selection.selection = False
            user_model_instance.selection_count += 1
            user_model_instance.save()
            selection.save()
            if user_model_instance.selection_count == 100:
                user_model_instance.has_finished = True
                user_model_instance.save()
                return redirect("results")
            selection = SelectionModel.objects.filter(
                user__exact=user_model_instance,
                selection=None)[0]
        context["photo_loc"] = "img/" + selection.photo.file
        context["selections_left"] = "Zostało Ci " + str(
            100 - int(
                user_model_instance.selection_count)) + " zdjęć do wybrania."
        return render(request, "index.html", context)
    else:
        return redirect("results")

# TODO: REFACTOR
def resultsView(request, *args, **kwargs):
    user_ip = get_client_ip(request)
    context = {}
    if not UserModel.objects.filter(ip__exact=user_ip):
        redirect('home')
    else:
        user_model_instance = UserModel.objects.filter(ip__exact=user_ip)[
            0]
        if not user_model_instance.has_finished:
            redirect('home')
    if not AlgorithmModel.objects.filter(user__exact=user_model_instance):
        for x in AlgorithmModel.TYPE_CHOICES:
            AlgorithmModel.objects.create(user=user_model_instance, type=x[0])
    if not user_model_instance.has_finished_results:
        algorithm = \
            AlgorithmModel.objects.filter(user__exact=user_model_instance,
                                          selection=None)[
                0]
        if algorithm.type == '1' and algorithm.photo is None:
            algorithms.linear_adding_algorithm(user_model_instance, algorithm)
        elif algorithm.type == '2' and algorithm.photo is None:
            algorithms.exponential_adding_algorithm(user_model_instance,
                                                    algorithm)
        elif algorithm.type == '3' and algorithm.photo is None:
            algorithms.selection_average_algorithm(user_model_instance,
                                                   algorithm)
        elif algorithm.type == '4' and algorithm.photo is None:
            algorithms.average_human_offset_algorithm(user_model_instance,
                                                      algorithm)
        elif algorithm.type == '5' and algorithm.photo is None:
            algorithms.attribute_weights_on_test_sample_algorithm(
                user_model_instance,
                algorithm)
        elif algorithm.type == '6' and algorithm.photo is None:
            algorithms.exponential_attribute_weights_on_test_sample_algorithm(
                user_model_instance, algorithm)
        elif algorithm.type == '7' and algorithm.photo is None:
            algorithms.average_landmarks(
                user_model_instance, algorithm)
        elif algorithm.type == '8' and algorithm.photo is None:
            algorithms.eawots_average_landmarks(
                user_model_instance, algorithm)
        if request.method == 'POST':
            if 'likebutton' in request.POST:
                algorithm.selection = True
            elif 'dislikebutton' in request.POST:
                algorithm.selection = False
            algorithm.save()
            if not AlgorithmModel.objects.filter(
                    user__exact=user_model_instance):
                user_model_instance.has_finished_results = True
                user_model_instance.save()
                user_photos = AlgorithmModel.objects.filter(
                    user__exact=user_model_instance)

                user_photos = AlgorithmModel.objects.filter(
                    user__exact=user_model_instance)
                user_photos = list(user_photos)
                for i, alg in enumerate(user_photos):
                    user_photos[i] = "img/" + alg.photo.file
                context["liked_list"] = user_photos
                return render(request, "thankyou.html", context)
            try:
                algorithm = \
                    AlgorithmModel.objects.filter(
                        user__exact=user_model_instance,
                        selection=None)[0]
                if algorithm.type == '1' and algorithm.photo is None:
                    algorithms.linear_adding_algorithm(user_model_instance,
                                                       algorithm)
                elif algorithm.type == '2' and algorithm.photo is None:
                    algorithms.exponential_adding_algorithm(user_model_instance,
                                                            algorithm)
                elif algorithm.type == '3' and algorithm.photo is None:
                    algorithms.selection_average_algorithm(user_model_instance,
                                                           algorithm)
                elif algorithm.type == '4' and algorithm.photo is None:
                    algorithms.average_human_offset_algorithm(
                        user_model_instance,
                        algorithm)
                elif algorithm.type == '5' and algorithm.photo is None:
                    algorithms.attribute_weights_on_test_sample_algorithm(
                        user_model_instance,
                        algorithm)
                elif algorithm.type == '6' and algorithm.photo is None:
                    algorithms.exponential_attribute_weights_on_test_sample_algorithm(
                        user_model_instance, algorithm)
                elif algorithm.type == '7' and algorithm.photo is None:
                    algorithms.average_landmarks(
                        user_model_instance, algorithm)
                elif algorithm.type == '8' and algorithm.photo is None:
                    algorithms.eawots_average_landmarks(
                        user_model_instance, algorithm)
            except:
                user_model_instance.has_finished_results = True
                user_model_instance.save()
        context["photo_loc"] = "img/" + algorithm.photo.file
        return render(request, "results.html", context)
    else:
        user_photos = AlgorithmModel.objects.filter(
            user__exact=user_model_instance)
        user_photos = list(user_photos)
        for i, alg in enumerate(user_photos):
            user_photos[i] = "img/" + alg.photo.file
        context["liked_list"] = user_photos
        return render(request, "thankyou.html", context)

def mealsView(request, *args, **kwargs):
    land_loc = STATICFILES_DIRS[0] + "/landmarks.txt"
    land_list = []
    attr_loc = STATICFILES_DIRS[0] + "/attributes.txt"
    attr_list = []
    with open(land_loc, encoding='utf-8') as f:
        for line in f.readlines():
            temp = list(line.rstrip().split()[1:])
            temp = [int(x) for x in temp]
            land_list.append(temp)
    with open(attr_loc, encoding='utf-8') as f:
        for line in f.readlines():
            attr_list.append(list(line.rstrip().split()))
    for x in range(202599):  # 202599
        f_name = attr_list[x].pop(0)
        meal_data = PhotoModel()
        meal_data.file = f_name
        meal_data.set_attributes([int(x) for x in attr_list[x]])
        meal_data.set_landmarks(land_list[x])
        meal_data.save()
    return redirect("home")

@login_required
def chartView(request, *args, **kwargs):
    labels = []
    users = UserModel.objects.all()
    for name in AlgorithmModel.TYPE_CHOICES:
        labels.append(name)
    all_to_add = []
    for user in users:
        if user.has_finished_results:
            algorithms = list(AlgorithmModel.objects.filter(user__exact=user))
            all_to_add.append(algorithms)
    data = list(zeros(len(labels)))
    for user_sel_set in all_to_add:
        for i, alg in enumerate(user_sel_set):
            if alg.selection:
                data[i] += 1
    data_raw = data
    data = [(x / len(all_to_add)) * 100 for x in data]
    return render(request, 'charts.html', {
        'UniqueValidUsers': len(all_to_add),
        'DataRaw': data_raw,
        'labels': labels,
        'data': data,
    })

@login_required
def fixView(request, *args, **kwargs):
    users = UserModel.objects.filter(has_finished=True, selection_count=100,
                                     has_finished_results=False)
    users = list(users)
    for user in users:
        try:
            AlgorithmModel.objects.filter(user__exact=user, selection=None)[0]
        except:
            user.has_finished_results = True
            user.save()
    return redirect("home")


import os
import time

from dotenv import load_dotenv

import telegram

load_dotenv()

CHAT_ID = os.getenv('CHAT_ID')
TELEGRAM_TOKEN = os.getenv('TELEGRAM_TOKEN')

tg_bot = telegram.Bot(token=TELEGRAM_TOKEN)

def send_message(text_massge):
    return tg_bot.send_message(chat_id=CHAT_ID, text=text_massge)

def main():
    mesege_number = 1
    while True:
        try:
            text = f"Messege # {mesege_number}. + test autodeploy."
            send_message(text)
            time.sleep(300)
            mesege_number += 1

        except Exception as e:
            print(f'Бот упал с ошибкой: {e}')
            time.sleep(300)
            continue

if __name__ == '__main__':
    main()


import pandas as pd
import numpy as np
from scipy.interpolate import interp1d
from scipy import signal
import matplotlib.pyplot as plt
import os

# 扩展到45
def up_sample(data):
    length = len(data)
    x = np.arange(0, length)

    # 采用一维插值进行上采样
    fc = interp1d(x, data, kind='cubic', axis=0)
    # fc2 = interp1d(x, data, kind='nearest')
    # fc3 = interp1d(x, data, kind='linear')

    xint = np.linspace(x.min(), x.max(), 45)
    yint = fc(xint)
    # yint2 = fc2(xint)
    # yint3 = fc3(xint)

    # plt.figure(dpi=100)
    # plt.plot(yint, color='red', label='cubic')
    # plt.legend()
    #
    # plt.figure(dpi=100)
    # plt.plot(yint2, color='green', label='nearest')
    # plt.legend()
    #
    # plt.figure(dpi=100)
    # plt.plot(yint3, color='blue', label='linear')  # 最牛逼的
    # plt.legend()
    #
    # plt.figure(dpi=100)
    # plt.plot(data, color='red', label='raw')
    #
    # plt.legend()
    # plt.show()

    return yint

# 下采样到45
def down_sample(data):
    length = len(data)
    x = np.arange(0, length)
    # print(x)

    # 利用插值法下采样
    fc = interp1d(x, data, kind='cubic', axis=0)
    # fc2 = interp1d(x, data, kind='nearest')
    # fc3 = interp1d(x, data, kind='linear')

    # 控制为45个采样点
    xint = np.linspace(x.min(), x.max(), 45)
    # print(xint)

    # 进行插值downsample
    yint = fc(xint)
    # yint2 = fc2(xint)
    # yint3 = fc3(xint)
    # print(yint)

    # plt.figure(dpi=100)
    # plt.plot(yint, color='red', label='one')
    #
    # plt.figure(dpi=100)
    # plt.plot(yint2, color='green', label='nearest')
    #
    # plt.figure(dpi=100)
    # plt.plot(yint3, color='blue', label='linear')
    #
    # plt.figure(dpi=100)
    # plt.plot(data, color='red', label='two')
    #
    # plt.legend()
    # plt.show()

    return yint

if __name__ == '__main__':
    # one = pd.read_csv('../data/dynamic/data/a_sample_0_0.txt', header=None, delim_whitespace=True)
    # kkk = pd.read_csv('../data/dynamic/data/k_sample_0_46.txt', header=None, delim_whitespace=True)
    # print(len(one), len(kkk))
    # flex1_one = one.iloc[:, 0]
    # flex1_k = kkk.iloc[:, 0]
    # flex2 = one.iloc[:, 1]
    # flex3 = one.iloc[:, 2]
    # flex4 = one.iloc[:, 3]
    # flex5 = one.iloc[:, 4]

    # print(flex1)
    # up_sample(flex1_one)
    # down_sample(flex1_k)

    # 批量进行数据长度固定
    dir_path = '../data/dynamic/data/'
    target_path = '../data/dynamic_unify/'
    files = os.listdir(dir_path)
    # print(files)
    for i in files:
        print(i)
        data = pd.read_csv(dir_path + i, header=None, delim_whitespace=True)
        print(data)
        sam_len = len(data)
        if sam_len >= 45:
            data_unify = pd.DataFrame(down_sample(data))
        elif sam_len < 45:
            data_unify = pd.DataFrame(up_sample(data))
        else:
            data_unify = data

        if not os.path.exists(target_path):
            os.mkdir(target_path)
        else:
            print('文件夹已经存在')
        name = i[:-3]
        data_unify.to_csv(target_path + name + 'csv', header=False, index_label=None, index=False)


'''
This script runs a full install of the dagger-related codes
B.G.
'''

import subprocess as sub
import sys

class bcolors:
	HEADER = '\033[95m'
	OKBLUE = '\033[94m'
	OKCYAN = '\033[96m'
	OKGREEN = '\033[92m'
	WARNING = '\033[93m'
	FAIL = '\033[91m'
	ENDC = '\033[0m'
	BOLD = '\033[1m'
	UNDERLINE = '\033[4m'
	YELLOW = '\033[33;49m'

# none = True if (("-g" in sys.argv) == False and ("-d" in sys.argv) == False and ("-p" in sys.argv) == False ) else True
# dagger = True if ("-d" in sys.argv or none) else False
# graphflood = True if ("-g" in sys.argv or none) else False
# popscape = True if ("-p" in sys.argv or none) else False

print(f"\n\n{bcolors.YELLOW}~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~={bcolors.ENDC}")
print(f"{bcolors.HEADER}I am going to install/update dagger suite codes from source in python.{bcolors.ENDC}\n")
print(f"{bcolors.BOLD}DAGGER core:{bcolors.ENDC} {bcolors.OKGREEN }{'included'}{bcolors.ENDC} ")
print(f"{bcolors.BOLD}graphflood core:{bcolors.ENDC} {bcolors.OKGREEN }{'included'}{bcolors.ENDC} ")
print(f"{bcolors.BOLD}popscape core:{bcolors.ENDC} {bcolors.OKGREEN}{'included'}{bcolors.ENDC} ")

can_run_test = False
try:
	import pytest
	can_run_test = True
except:
	pass

print("\n")

if("-v" in sys.argv):
	stttd = None
else:
	print(f"{bcolors.WARNING}use option -v to get the whole install log{bcolors.ENDC}\n")
	stttd = sub.DEVNULL

print(f"{bcolors.BOLD}Installing ... {bcolors.ENDC}", end = "", flush = True)
run_test = False
try:
	sub.run("pip uninstall -y dagger", shell = True, check = False, stdout = stttd, stderr = stttd)
	sub.run("pip uninstall -y DAGGER", shell = True, check = False, stdout = stttd, stderr = stttd)
	sub.run("pip uninstall -y daggerpy", shell = True, check = False, stdout = stttd, stderr = stttd)
	sub.run("pip uninstall -y popscape-dagger", shell = True, check = False, stdout = stttd, stderr = stttd)
	sub.run("pip uninstall -y fastflood-dagger", shell = True, check = False, stdout = stttd, stderr = stttd)
	sub.run("rm -r dist", cwd = "./", shell = True, check = False, stdout=stttd, stderr=stttd)
	sub.run("rm -r build", cwd = "./", shell = True, check = False, stdout=stttd, stderr=stttd)
	sub.run("python setup.py install", cwd = "./", shell = True, check = True, stdout=stttd, stderr=stttd)
	print(f"{bcolors.OKGREEN}OK!{bcolors.ENDC}\n")
except:
	print(f"{bcolors.FAIL}Failed...{bcolors.ENDC}\n")
	run_test = False

if(run_test):
	print(f"{bcolors.BOLD}Running tests ... {bcolors.ENDC}", end = "\n", flush = True)

	if(can_run_test):
		try:
			# sub.run("./tests/pytest", shell = True, check = False, stdout = stttd, stderr = stttd)
			pytest.main()
			print(f"{bcolors.OKGREEN}\n\n\n\n\nTest ran. See report above for results{bcolors.ENDC}\n")

		except:
			print(f"{bcolors.FAIL}Failed... code compiles but some tests produce erroneous results. Rerun with -v for more info.{bcolors.ENDC}\n")

	else:
		print(f"{bcolors.YELLOW}Cannot run tests - you need to install pytest (pip install pytest){bcolors.ENDC}\n")

print(f"{bcolors.BOLD}Done!{bcolors.ENDC}")
print(f"\n{bcolors.YELLOW}~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~={bcolors.ENDC}")


import time
import psutil
import os
import csv
from multiprocessing import Process, Manager

from map import Map
from algorithms import BFS_checker, Astar, PRM, DQN

def run_test(test_map, algorithm, results, index):
    start_time = time.time()
    status, path = algorithm(test_map.map, test_map.start, test_map.end)
    end_time = time.time()

    results[index] = {
        'algorithm': algorithm.__name__,
        'size': test_map.width,  # assuming square map so width = height = size
        'time': end_time - start_time,
        'cpu_percent': psutil.cpu_percent(),
        'memory_percent': psutil.virtual_memory().percent,
        'path_length': len(path) if status else None
    }

def main():
    sizes = [50, 100, 200]
    algorithms = [BFS_checker, Astar, PRM, DQN]

    # Warm up
    print("Warming Up Baby")
    #test_map = Map(10000, 10000)
    print("Warming Up D0n€ :)")

    # How many tests for average value
    average = 200

    results = []

    for size in sizes:
        for _ in range(average):  # Run 10 times
            test_map = Map(size, size)
            manager = Manager()
            algorithm_results = manager.dict()

            # Run each algorithm on a separate process
            processes = []
            for i, algorithm in enumerate(algorithms):
                process = Process(target=run_test, args=(test_map, algorithm, algorithm_results, i))
                process.start()
                processes.append(process)

            for process in processes:
                process.join()

            results += list(algorithm_results.values())

    # Save results to CSV
    with open('results.csv', 'w') as f:
        writer = csv.DictWriter(f, fieldnames=['algorithm', 'size', 'time', 'cpu_percent', 'memory_percent', 'path_length'])
        writer.writeheader()
        writer.writerows(results)

if __name__ == '__main__':
    main()


import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--name', type=str, default='svhn')
parser.add_argument('--mode', type=str, default='train')
parser.add_argument('--root', type=str, default='datasets')
args = parser.parse_args()

import os
labels = os.listdir(os.path.join(args.root, args.name, args.mode))

with open(os.path.join('sourcefiles', f'{args.name}_{args.mode}.txt'), 'w') as f:
    for l in labels:
        for img in os.listdir(os.path.join(args.root, args.name, args.mode, l)):
            f.write(f'{l}/{img} {l}\n')