"""Module with core functionality for a single pipeline stage """

import pathlib
import os
import sys
from textwrap import dedent
import shutil
import cProfile
import pdb
import datetime

from abc import abstractmethod
from . import errors
from .monitor import MemoryMonitor
from .config import StageParameter, StageConfig, cast_to_streamable

SERIAL = "serial"
MPI_PARALLEL = "mpi"
DASK_PARALLEL = "dask"

IN_PROGRESS_PREFIX = "inprogress_"

class PipelineStage:
    """A PipelineStage implements a single calculation step within a wider pipeline.

    Each different type of analysis stage is represented by a subclass of this
    base class.  The base class handles the connection between different pipeline
    stages, and the execution of the stages within a workflow system (parsl),
    potentially in parallel (MPI).

    An instance of one of these classes represents an actual run of the stage,
    with the required inputs, outputs, and configuration specified.

    See documentation pages for more details.

    """

    parallel = True
    dask_parallel = False
    config_options = {}
    doc = ""
    allow_reload = False

    def __init__(self, args, comm=None):
        """Construct a pipeline stage, specifying the inputs, outputs, and configuration for it.

        The constructor needs a dict or namespace. It should include:
        - input paths (required)
        - config path (required)
        - output paths (optional but usual)
        - additional configuration (required if not specified elsewhere)

        Input and output paths should map tags to paths.
        Tags are strings, and the first elements in each item in the subclass's
        "inputs" and "output" attributes.
        e.g. for a subclass with:
            inputs = [('eggs', TextFile)]
            outputs = [('spam', TextFile)]
        the args could contain:
            {'eggs': 'inputs/eggs.txt',
             'spam': 'outputs/spam.txt' }
        If spam is not specified it will default to "./spam.txt"

        }

        The config should map "config" to a path where a YAML config file
        is located, e.g. {'config':'/path/to/config.yml'}

        Any config variables that are specified in the class's config attribute
        will be searched for first in args, then in the config file, and then
        by looking at any default value they have been given.
        If they have no default value (and just a type, like int, is listed), then
        it's an error if they are not specified somewhere.

        The execute method can instantiate and run the class together, with added bonuses
        like profiling and debugging tools.

        Parameters
        ----------
        args: dict or namespace
            Specification of input and output paths and any missing config options
        comm: MPI communicator
            (default is None) An MPI comm object to use in preference to COMM_WORLD
        """
        self._configs = StageConfig(**self.config_options)
        self._inputs = None
        self._outputs = None
        self._parallel = SERIAL
        self._comm = None
        self._size = 1
        self._rank = 0
        self._io_checked = False
        self.dask_client = None

        self.load_configs(args)
        if comm is not None:
            self.setup_mpi(comm)

    @classmethod
    def make_stage(cls, **kwargs):
        """Make a stage of a particular type"""
        kwcopy = kwargs.copy()
        kwcopy.setdefault("config", None)
        comm = kwcopy.pop("comm", None)
        name = kwcopy.get("name", None)
        aliases = kwcopy.pop("aliases", {})
        for input_ in cls.inputs:
            kwcopy.setdefault(input_[0], "None")
        if name is not None:
            for output_ in cls.outputs:  # pylint: disable=no-member
                outtag = output_[0]
                aliases[outtag] = f"{outtag}_{name}"
        kwcopy["aliases"] = aliases
        return cls(kwcopy, comm=comm)

    def get_aliases(self):
        """Returns the dictionary of aliases used to remap inputs and outputs
        in the case that we want to have multiple instance of this class in the pipeline"""
        return self.config.get("aliases", None)

    def get_aliased_tag(self, tag):
        """Returns the possibly remapped value for an input or output tag

        Parameter
        ---------
        tag : `str`
            The input or output tag we are checking

        Returns
        -------
        aliased_tag : `str`
            The aliases version of the tag
        """
        aliases = self.get_aliases()
        if aliases is None:
            return tag
        return aliases.get(tag, tag)

    @abstractmethod
    def run(self):  # pragma: no cover
        """Run the stage and return the execution status"""
        raise NotImplementedError("run")

    def load_configs(self, args):
        """
        Load the configuraiton

        Parameters
        ----------
        args: dict or namespace
            Specification of input and output paths and any missing config options
        """
        if not isinstance(args, dict):
            args = vars(args)

        # We alwys assume the config arg exists, whether it is in input_tags or not
        if "config" not in args:  # pragma: no cover
            raise ValueError("The argument --config was missing on the command line.")

        _name = args.get("name")
        if _name is not None:
            self._configs.name = _name

        # First, we extract configuration information from a combination of
        # command line arguments and optional 'config' file
        self._inputs = dict(config=args["config"])
        try:
            self.read_config(args)
        except Exception as error:
            error_class = type(error)
            msg = str(error)
            raise error_class(f"Error configuring {self.instance_name}: {msg}")
        self.check_io(args)

    def check_io(self, args=None):
        """
        Check the inputs and outputs.
        This function is seperate so that when Stages are configured interactively after
        construction then can invove this

        Parameters
        ----------
        args: dict or namespace
            Specification of input and output paths and any missing config options
        """

        # We first check for missing input files, that's a show stopper
        if self._io_checked:  # pragma: no cover
            return
        if args is None:  # pragma: no cover
            args = self.config

        missing_inputs = []
        for x in self.input_tags():
            val = args.get(x)
            aliased_tag = self.get_aliased_tag(x)

            if val is None:
                val = args.get(aliased_tag)

            if val is None:  # pragma: no cover
                missing_inputs.append(f"--{x}")
            else:
                self._inputs[aliased_tag] = val
        if missing_inputs:  # pragma: no cover
            missing_inputs = "  ".join(missing_inputs)
            raise ValueError(
                f"""

{self.instance_name} Missing these names on the command line:
    Input names: {missing_inputs}"""
            )

        # We prefer to receive explicit filenames for the outputs but will
        # tolerate missing output filenames and will default to tag name in
        # current folder (this is for CWL compliance)
        self._outputs = {}
        for i, x in enumerate(self.output_tags()):
            aliased_tag = self.get_aliased_tag(x)
            if args.get(x) is None:
                ftype = self.outputs[i][1]  # pylint: disable=no-member
                self._outputs[aliased_tag] = ftype.make_name(aliased_tag)
            else:
                self._outputs[aliased_tag] = args[x]
        self._io_checked = True

    def setup_mpi(self, comm=None):
        """
        Setup the MPI interface

        Parameters
        ----------
        comm: MPI communicator
            (default is None) An MPI comm object to use in preference to COMM_WORLD
        """
        mpi = self.config.get("mpi", False)

        if mpi:  # pragma: no cover
            try:
                # This isn't a ceci dependency, so give a sensible error message if not installed.
                import mpi4py.MPI
            except ImportError:
                print("ERROR: Using --mpi option requires mpi4py to be installed.")
                raise

        # For scripting and testing we allow an MPI communicator or anything
        # with the same API to be passed in directly, overriding the --mpi
        # flag.
        if comm is not None:
            self._parallel = MPI_PARALLEL
            self._comm = comm
            self._size = self._comm.Get_size()
            self._rank = self._comm.Get_rank()
        elif mpi:  # pragma: no cover
            self._parallel = MPI_PARALLEL
            self._comm = mpi4py.MPI.COMM_WORLD
            self._size = self._comm.Get_size()
            self._rank = self._comm.Get_rank()
        else:
            self._parallel = SERIAL
            self._comm = None
            self._size = 1
            self._rank = 0

        # If we are running under MPI but this subclass has enabled dask
        # then we note that here. It stops various MPI-specific things happening
        # later
        if (self._parallel == MPI_PARALLEL) and self.dask_parallel:
            self._parallel = DASK_PARALLEL

    pipeline_stages = {}
    incomplete_pipeline_stages = {}

    def __init_subclass__(cls, **kwargs):
        """
        Python 3.6+ provides a facility to automatically
        call a method (this one) whenever a new subclass
        is defined.  In this case we use that feature to keep
        track of all available pipeline stages, each of which is
        defined by a class.

        """
        super().__init_subclass__(**kwargs)

        # This is a hacky way of finding the file
        # where our stage was defined
        filename = sys.modules[cls.__module__].__file__

        stage_is_complete = (
            hasattr(cls, "inputs")
            and hasattr(cls, "outputs")
            and not getattr(cls.run, "__isabstractmethod__", False)
        )

        # If there isn't an explicit name already then set it here.
        # by default use the class name.
        if not hasattr(cls, "name"):  # pragma: no cover
            cls.name = cls.__name__
        if cls.name is None:  # pragma: no cover
            cls.name = cls.__name__

        if stage_is_complete:
            # Deal with duplicated class names
            if cls.name in cls.pipeline_stages and not cls.allow_reload:
                other = cls.pipeline_stages[cls.name][1]
                raise errors.DuplicateStageName(
                    "You created two pipeline stages with the"
                    f"name {cls.name}.\nOne was in {filename}\nand the "
                    f"other in {other}\nYou can either change the class "
                    "name or explicitly put a variable 'name' in the top"
                    "level of the class."
                )

            # Check for "config" in the inputs list - this is implicit
            for name, _ in cls.inputs:
                if name == "config":
                    raise errors.ReservedNameError(
                        "An input called 'config' is implicit in each pipeline "
                        "stage and should not be added explicitly.  Please update "
                        f"your pipeline stage called {cls.name} to remove/rename "
                        "the input called 'config'."
                    )

        # Check if user has over-written the config variable.
        # Quite a common error I make myself.
        if not isinstance(cls.config, property):
            raise errors.ReservedNameError(
                "You have a class variable called 'config', which "
                "is reserved in ceci for its own configuration. "
                "You may have meant to specify config_options?"
            )
        # Find the absolute path to the class defining the file
        path = pathlib.Path(filename).resolve()

        # Add a description of the parameters to the end of the docstring
        if stage_is_complete:
            config_text = cls._describe_configuration_text()
            if cls.__doc__ is None:
                cls.__doc__ = f"Stage {cls.name}\n\nConfiguration Parameters:\n{config_text}"
            else:
                # strip any existing configuration text from parent classes that is at the end of the doctring
                cls.__doc__ = cls.__doc__.split("Configuration Parameters:")[0]
                cls.__doc__ += f"\n\nConfiguration Parameters:\n{config_text}"

        # Register the class
        if stage_is_complete:
            cls.pipeline_stages[cls.name] = (cls, path)
        else:
            cls.incomplete_pipeline_stages[cls.__name__] = (cls, path)

    #############################################
    # Life cycle-related methods and properties.
    #############################################

    @classmethod
    def get_stage(cls, name, module_name=None):
        """
        Return the PipelineStage subclass with the given name.

        This is used so that we do not need a new entry point __main__ function
        for each new stage - instead we can just use a single one which can query
        which class it should be using based on the name.

        If module_name is provided, this will import that module
        in order to load the required class.

        Returns
        -------
        cls: class
            The corresponding subclass
        """
        stage = cls.pipeline_stages.get(name)
        if stage is None:
            if module_name:
                __import__(module_name)
            stage = cls.pipeline_stages.get(name)

        # If not found, then check for incomplete stages
        if stage is None:
            if name in cls.incomplete_pipeline_stages:
                raise errors.IncompleteStage(
                    f"The stage {name} is not completely written. "
                    "Stages must specify 'inputs', 'outputs' as class variables "
                    f"and a 'run' method.\n{name} might be unfinished, or it might "
                    "be intended as a base for other classes and not to be run."
                )
            raise errors.StageNotFound(f"Unknown stage '{name}'")
        return stage[0]

    @classmethod
    def get_module(cls):
        """
        Return the path to the python package containing the current sub-class

        If we have a PipelineStage subclass defined in a module called "bar", in
        a package called "foo" e.g.:
        /path/to/foo/bar.py  <--   contains subclass "Baz"

        Then calling Baz.get_module() will return "foo.bar".

        We use this later to construct command lines like "python -m foo Baz"

        Returns
        -------
        module: str
            The module containing this class.
        """
        return cls.pipeline_stages[cls.name][0].__module__

    @classmethod
    def describe_configuration(cls):
        print(cls._describe_configuration_text())

    @classmethod
    def _describe_configuration_text(cls):
        s = []
        for name, val in cls.config_options.items():
            if isinstance(val, StageParameter):
                if val.required:
                    if val.dtype is None:
                        txt = f"[type not specified]: {val._help} (required)"
                    else:
                        txt = f"[{val.dtype.__name__}]: {val._help}  (required)"
                else:
                    if val.dtype is None:
                        txt = f"[type not specified]: {val._help} (default={val.default})"
                    else:
                        txt = f"[{val.dtype.__name__}]: {val._help} (default={val.default})"
            elif isinstance(val, type):
                txt = f"[{val.__name__}]: (required)"
            else:
                txt = f"[{type(val).__name__}]: (default={val})"
            s.append(f"{name} {txt} ")
        return '\n'.join(s)

    @classmethod
    def usage(cls):  # pragma: no cover
        """
        Print a usage message.
        """
        names = []
        docs = []
        for name, (stage, _) in cls.pipeline_stages.items():
            # find the first non-empty doc line, if there is one.
            try:
                doc_lines = [s.strip() for s in stage.__doc__.split("\n")]
                doc_lines = [d for d in doc_lines if d]
                doc = doc_lines[0]
            except (AttributeError, IndexError):
                doc = ""
            # cut off any very long lines
            if len(doc) > 100:
                doc = doc[:100] + " ..."
            # print the text
            names.append(name)
            docs.append(doc)

        # Make it look like a nice table by finding the maximum
        # length of the names, so that all the docs line up
        n = max(len(name) for name in names) + 1
        stage_texts = [f"- {name:{n}} - {d}" for name, d in zip(names, docs)]
        stage_text = "\n".join(stage_texts)

        try:
            module = cls.get_module().split(".")[0]
        except:  # pylint: disable=bare-except
            module = "<module_name>"
        sys.stderr.write(
            f"""
Usage: python -m {module} <stage_name> <stage_arguments>

If no stage_arguments are given then usage information
for the chosen stage will be given.

I currently know about these stages:

{stage_text}
"""
        )

    @classmethod
    def main(cls):
        """
        Create an instance of this stage and execute it with
        inputs and outputs taken from the command line
        """
        try:
            stage_name = sys.argv[1]
        except IndexError:  # pragma: no cover
            cls.usage()
            return 1
        if stage_name in ["--help", "-h"] and len(sys.argv) == 2:  # pragma: no cover
            cls.usage()
            return 1
        if stage_name.find(".") >= 0:
            tokens = stage_name.split(".")
            module_name = ".".join(tokens[:-1])
            stage_name = tokens[-1]
        else:
            module_name = None

        stage = cls.get_stage(stage_name, module_name)
        args = stage.parse_command_line()
        stage.execute(args)
        return 0

    @classmethod
    def parse_command_line(cls, cmd=None):
        """Set up and argument parser and parse the command line

        Parameters
        ----------
        cmd : str or None
            The command line to part (if None this will use the system arguments)

        Returns
        -------
        args : Namespace
            The resulting Mapping of arguement to values
        """
        import argparse

        parser = argparse.ArgumentParser(description=f"Run pipeline stage {cls.name}")
        parser.add_argument("stage_name")
        for conf, def_val in cls.config_options.items():
            if isinstance(def_val, StageParameter):
                opt_type = def_val.dtype
                def_val = def_val.default
            else:
                opt_type = def_val if isinstance(def_val, type) else type(def_val)
            if opt_type == bool:
                parser.add_argument(f"--{conf}", action="store_const", const=True)
                parser.add_argument(
                    f"--no-{conf}", dest=conf, action="store_const", const=False
                )
            elif opt_type == list:
                if not def_val:
                    out_type = str
                else:
                    out_type = (
                        def_val[0] if isinstance(def_val[0], type) else type(def_val[0])
                    )
                if out_type is str:  # pragma: no cover
                    parser.add_argument(
                        f"--{conf}", type=lambda string: string.split(",")
                    )
                elif out_type is int:  # pragma: no cover
                    parser.add_argument(
                        f"--{conf}",
                        type=lambda string: [int(i) for i in string.split(",")],
                    )
                elif out_type is float:
                    parser.add_argument(
                        f"--{conf}",
                        type=lambda string: [float(i) for i in string.split(",")],
                    )
                else:  # pragma: no cover
                    raise NotImplementedError(
                        "Only handles str, int and float list arguments"
                    )
            else:  # pragma: no cover
                parser.add_argument(f"--{conf}", type=opt_type)
        for inp in cls.input_tags():
            parser.add_argument(f"--{inp}")
        for out in cls.output_tags():
            parser.add_argument(f"--{out}")
        parser.add_argument(
            "--name",
            action="store",
            default=cls.name,
            type=str,
            help="Rename the stage",
        )
        parser.add_argument("--config")

        if cls.parallel:
            parser.add_argument(
                "--mpi", action="store_true", help="Set up MPI parallelism"
            )
        parser.add_argument(
            "--pdb", action="store_true", help="Run under the python debugger"
        )
        parser.add_argument(
            "--cprofile",
            action="store",
            default="",
            type=str,
            help="Profile the stage using the python cProfile tool",
        )
        parser.add_argument(
            "--memmon",
            type=int,
            default=0,
            help="Report memory use. Argument gives interval in seconds between reports",
        )

        # Error message we will return if --mpi used on a non-supported
        # stage.
        mpi_err = (
            "Error: you used the --mpi flag (or set MPI parallelism options) "
            f"for the stage {cls.name}, but that stage cannot be run in parallel."
        )

        if cmd is None:
            if ("--mpi" in sys.argv) and not cls.parallel:
                raise ValueError(mpi_err)
            ret_args = parser.parse_args()
        else:
            if ("--mpi" in cmd) and not cls.parallel:
                raise ValueError(mpi_err)
            ret_args = parser.parse_args(cmd)

        return ret_args

    @classmethod
    def execute(cls, args, comm=None):
        """
        Create an instance of this stage and run it
        with the specified inputs and outputs.

        This is calld by the main method.

        Parameters
        ----------
        args: namespace
            The argparse namespace for this subclass.
        """

        # Create the stage instance.  Running under dask this only
        # actually needs to happen for one process, but it's not a major
        # overhead and lets us do a whole bunch of other setup above
        stage = cls(args)
        stage.setup_mpi(comm)

        # This happens before dask is initialized
        start_time = datetime.datetime.now()
        if stage.rank == 0:
            start_time_text = start_time.isoformat(" ")
            print(f"Executing stage: {cls.name} @ {start_time_text}")

        if stage.is_dask():
            is_client = stage.start_dask()
            # worker and scheduler stages do not execute the
            # run method under dask
            if not is_client:
                return

        if args.cprofile:  # pragma: no cover
            profile = cProfile.Profile()
            profile.enable()

        if args.memmon:  # pragma: no cover
            monitor = MemoryMonitor.start_in_thread(interval=args.memmon)

        try:
            stage.run()
        except Exception as error:  # pragma: no cover
            if args.pdb:
                print(
                    "There was an exception - starting python debugger because you ran with --pdb"
                )
                print(error)
                pdb.post_mortem()
            else:
                if stage.rank == 0:
                    end_time = datetime.datetime.now()
                    end_time_text = end_time.isoformat(" ")
                    minutes = (end_time - start_time).total_seconds() / 60
                    print(
                        f"Stage failed: {cls.name} @ {end_time_text} after {minutes:.2f} minutes"
                    )
                raise
        finally:
            if args.memmon:  # pragma: no cover
                monitor.stop()
            if stage.is_dask():
                stage.stop_dask()

        # The default finalization renames any output files to their
        # final location, but subclasses can override to do other things too
        try:
            stage.finalize()
        except Exception as error:  # pragma: no cover
            if args.pdb:
                print(
                    "There was an exception in the finalization - starting python debugger because you ran with --pdb"
                )
                print(error)
                pdb.post_mortem()
            else:
                raise
        if args.cprofile:  # pragma: no cover
            profile.disable()
            profile.dump_stats(args.cprofile)
            profile.print_stats("cumtime")

        # Under dask the
        # the root process has gone off to become the scheduler,
        # and process 1 becomes the client which runs this code
        # and gets to this point
        if stage.rank == 0 or stage.is_dask():
            end_time = datetime.datetime.now()
            end_time_text = end_time.isoformat(" ")
            minutes = (end_time - start_time).total_seconds() / 60
            print(
                f"Stage complete: {cls.name} @ {end_time_text} took {minutes:.2f} minutes"
            )

    def finalize(self):
        """Finalize the stage, moving all its outputs to their final locations."""
        # Synchronize files so that everything is closed
        if self.is_mpi():  # pragma: no cover
            self.comm.Barrier()

        # Move files to their final path
        # Only the root process moves things, except under dask it is
        # process 1, which is the only process that reaches this point
        # (as noted above)
        if (self.rank == 0) or self.is_dask():
            for tag in self.output_tags():
                # find the old and new names
                self._finalize_tag(tag)

    def _finalize_tag(self, tag):
        """Finalize the data for a particular tag.

        This can be overridden by sub-classes for more complicated behavior
        """
        aliased_tag = self.get_aliased_tag(tag)
        temp_name = self.get_output(aliased_tag)
        final_name = self.get_output(aliased_tag, final_name=True)

        # it's not an error here if the path does not exist,
        # because that will be handled later.
        if pathlib.Path(temp_name).exists():
            # replace directories, rather than nesting more results
            if pathlib.Path(final_name).is_dir():  # pragma: no cover
                shutil.rmtree(final_name)
            shutil.move(temp_name, final_name)
        else:  # pragma: no cover
            sys.stderr.write(
                f"NOTE/WARNING: Expected output file {final_name} was not generated.\n"
            )
        return final_name

    #############################################
    # Parallelism-related methods and properties.
    #############################################
    @property
    def rank(self):
        """The rank of this process under MPI (0 if not running under MPI)"""
        return self._rank

    @property
    def size(self):
        """The number or processes under MPI (1 if not running under MPI)"""
        return self._size

    @property
    def comm(self):
        """The MPI communicator object (None if not running under MPI)"""
        return self._comm

    def is_parallel(self):
        """
        Returns True if the code is being run in parallel.
        Right now is_parallel() will return the same value as is_mpi(),
        but that may change in future if we implement other forms of
        parallelization.
        """
        return self._parallel != SERIAL

    def is_mpi(self):
        """
        Returns True if the stage is being run under MPI.
        """
        return self._parallel == MPI_PARALLEL

    def is_dask(self):
        """
        Returns True if the stage is being run in parallel with Dask.
        """
        return self._parallel == DASK_PARALLEL

    def start_dask(self):
        """
        Prepare dask to run under MPI. After calling this method
        only a single process, MPI rank 1 will continue to exeute code
        """

        # using the programmatic dask configuration system
        # does not seem to work. Presumably the loggers have already
        # been created by the time we modify the config. Doing it with
        # env vars seems to work. If the user has already set this then
        # we use that value. Otherwise we only want error logs
        key = "DASK_LOGGING__DISTRIBUTED"
        os.environ[key] = os.environ.get(key, "error")
        try:
            import dask
            import dask_mpi
            import dask.distributed
        except ImportError:  # pragma: no cover
            print(
                "ERROR: Using --mpi option on stages that use dask requires "
                "dask[distributed] and dask_mpi to be installed."
            )
            raise

        if self.size < 3:  # pragma: no cover
            raise ValueError(
                "Dask requires at least three processes. One becomes a scheduler "
                "process, one is a client that runs the code, and more are required "
                "as worker processes."
            )

        # This requires my fork until/unless they merge the PR, to allow
        # us to pass in these two arguments. In vanilla dask-mpi sys.exit
        # is called at the end of the event loop without returning to us.
        # After this point only a single process, MPI rank 1,
        # should continue to exeute code. The others enter an event
        # loop and return with is_client=False, which we return here
        # to tell the caller that they should not run everything.
        is_client = dask_mpi.initialize(comm=self.comm, exit=False)

        if is_client:
            # Connect this local process to remote workers.
            self.dask_client = dask.distributed.Client()
            # I don't yet know how to see this dashboard link at nersc
            print(f"Started dask. Diagnostics at {self.dask_client.dashboard_link}")

        return is_client

    @staticmethod
    def stop_dask():
        """
        End the dask event loop
        """
        from dask_mpi import send_close_signal

        send_close_signal()

    def split_tasks_by_rank(self, tasks):
        """Iterate through a list of items, yielding ones this process is responsible for/

        Tasks are allocated in a round-robin way.

        Parameters
        ----------
        tasks: iterable
            Tasks to split up

        """
        for i, task in enumerate(tasks):
            if i % self.size == self.rank:
                yield task

    def map_tasks_by_rank(self, function, inputs, allgather=False):
        """Run a function over a series of inputs, in parallel

        This mirrors the map function, and returns the equivalent of
        [function(input) for input in inputs], but executes in parallel.

        Parameters
        ----------
        function: Callable
            Function to be run on each item in inputs

        inputs: Iterable
            Any sequence of inputs, which should be the same
            on all processes. Or at least the same length:
            inputs not assigned to this process are ignored so
            you could get away with a dummy input for them.

        allgather: bool
            Whether to give all ranks the results (True) or just the
            root process (False). Default = False.

        Returns
        -------
        results: list
            A list of the results of calling the function on each input,
            in the same order as the input tasks
        """
        results = []
        # We keep track of the number of inputs manually rather
        # than calling len(inputs) because this allows inputs to
        # be an iterator.
        n = 0
        for i, inp in enumerate(inputs):
            n += 1
            if i % self.size == self.rank:
                results.append(function(inp))

        # If this is running in serial then the above just functions
        # like a basic map or list comprehension.
        if self.comm is not None:
            # Collate result as a list-of-lists, one sub-list for
            # each process
            if allgather:
                collected_results = self.comm.allgather(results)
            else:
                collected_results = self.comm.gather(results)
                if self.rank != 0:
                    return
            # convert the list-of-lists back into a single list
            # of results, returning to the original ordering.
            # The round-robin way we allocated them in the first
            # place is reversed by this.
            results = []
            for i in range(n):
                j = i % self.size
                k = i // self.size
                results.append(collected_results[j][k])

        return results

    def data_ranges_by_rank(self, n_rows, chunk_rows, parallel=True):
        """Split a number of rows by process.

        Given a total number of rows to read and a chunk size, yield
        the ranges within them that this process should handle.

        Parameters
        ----------
        n_rows: int
            Total number of rows to split up

        chunk_rows: int
            Size of each chunk to be read.

        Parallel: bool
            Whether to split data by rank or just give all procs all data.
            Default=True
        """
        n_chunks = n_rows // chunk_rows
        if n_chunks * chunk_rows < n_rows:  # pragma: no cover
            n_chunks += 1
        if parallel:
            it = self.split_tasks_by_rank(range(n_chunks))
        else:
            it = range(n_chunks)
        for i in it:
            start = i * chunk_rows
            end = min((i + 1) * chunk_rows, n_rows)
            yield start, end

    ##################################################
    # Input and output-related methods and properties.
    ##################################################

    def get_input(self, tag):
        """
        Return the path of an input file with the given tag,
        which can be aliased.
        """
        tag = self.get_aliased_tag(tag)
        return self._inputs[tag]

    def get_output(self, tag, final_name=False):
        """
        Return the path of an output file with the given tag,
        which can be aliased already.

        If final_name is False then use a temporary name - file will
        be moved to its final name at the end
        """

        tag = self.get_aliased_tag(tag)
        path = self._outputs[tag]

        # If not the final version, add a tag at the start of the filename
        if not final_name:
            p = pathlib.Path(path)
            p = p.parent / (IN_PROGRESS_PREFIX + p.name)
            path = str(p)
        return path

    def open_input(self, tag, wrapper=False, **kwargs):
        """
        Find and open an input file with the given tag, in read-only mode.

        For general files this will simply return a standard
        python file object.

        For specialized file types like FITS or HDF5 it will return
        a more specific object - see the types.py file for more info.

        """
        path = self.get_input(tag)
        input_class = self.get_input_type(tag)
        obj = input_class(path, "r", **kwargs)

        if wrapper:  # pragma: no cover
            return obj
        return obj.file

    def open_output(
        self, tag, wrapper=False, final_name=False, **kwargs
    ):  # pragma: no cover
        """
        Find and open an output file with the given tag, in write mode.

        If final_name is True then they will be opened using their final
        target output name.  Otherwise we will prepend "inprogress_" to their
        file name. This means we know that if the final file exists then it
        is completed.

        If wrapper is True this will return an instance of the class
        of the file as specified in the cls.outputs.  Otherwise it will
        return an open file object (standard python one or something more
        specialized).

        Parameters
        ----------

        tag: str
            Tag as listed in self.outputs

        wrapper: bool
            Default=False.  Whether to return a wrapped file

        final_name: bool
            Default=False. Whether to save to

        **kwargs:
            Extra args are passed on to the file's class constructor.

        """
        path = self.get_output(tag, final_name=final_name)
        output_class = self.get_output_type(tag)

        # HDF files can be opened for parallel writing
        # under MPI.  This checks if:
        # - we have been told to open in parallel
        # - we are actually running under MPI
        # and adds the flags required if all these are true
        run_parallel = kwargs.pop("parallel", False) and self.is_mpi()
        if run_parallel:
            kwargs["driver"] = "mpio"
            kwargs["comm"] = self.comm

            # XXX: This is also not a dependency, but it should be.
            #      Or even better would be to make it a dependency of descformats where it
            #      is actually used.
            import h5py

            if not h5py.get_config().mpi:
                print(
                    dedent(
                        """\
                Your h5py installation is not MPI-enabled.
                Options include:
                  1) Set nprocess to 1 for all stages
                  2) Upgrade h5py to use mpi.  See instructions here:
                     http://docs.h5py.org/en/latest/build.html#custom-installation
                Note: If using conda, the most straightforward way is to enable it is
                    conda install -c spectraldns h5py-parallel
                """
                    )
                )
                raise RuntimeError("h5py module is not MPI-enabled.")

        # Return an opened object representing the file
        obj = output_class(path, "w", **kwargs)
        if wrapper:
            return obj
        return obj.file

    @classmethod
    def inputs_(cls):
        """
        Return the dict of inputs
        """
        return cls.inputs  # pylint: disable=no-member

    @classmethod
    def outputs_(cls):
        """
        Return the dict of inputs
        """
        return cls.outputs  # pylint: disable=no-member

    @classmethod
    def output_tags(cls):
        """
        Return the list of output tags required by this stage
        """
        return [tag for tag, _ in cls.outputs_()]

    @classmethod
    def input_tags(cls):
        """
        Return the list of input tags required by this stage
        """
        return [tag for tag, _ in cls.inputs_()]

    def get_input_type(self, tag):
        """Return the file type class of an input file with the given tag."""
        tag = self.get_aliased_tag(tag)
        for t, dt in self.inputs_():
            t = self.get_aliased_tag(t)
            if t == tag:
                return dt
        raise ValueError(f"Tag {tag} is not a known input")  # pragma: no cover

    def get_output_type(self, tag):
        """Return the file type class of an output file with the given tag."""
        tag = self.get_aliased_tag(tag)
        for t, dt in self.outputs_():
            t = self.get_aliased_tag(t)
            if t == tag:
                return dt
        raise ValueError(f"Tag {tag} is not a known output")  # pragma: no cover

    ##################################################
    # Configuration-related methods and properties.
    ##################################################

    @property
    def instance_name(self):
        """Return the name associated to this particular instance of this stage"""
        return self._configs.get("name", self.name)

    @property
    def config(self):
        """
        Returns the configuration dictionary for this stage, aggregating command
        line options and optional configuration file.
        """
        return self._configs

    def read_config(self, args):
        """
        This function looks for the arguments of the pipeline stage using a
        combination of default values, command line options and separate
        configuration file.

        The order for resolving config options is first looking for a default
        value, then looking for a

        In case a mandatory argument (argument with no default) is missing,
        an exception is raised.

        Note that we recognize arguments with no default as the ones where
        self.config_options holds a type instead of a value.
        """
        # Try to load configuration file if provided
        import yaml

        config_file = self.get_input("config")

        # This is all the config information in the file, including
        # things for other stages
        if config_file is not None:
            with open(config_file) as _config_file:
                overall_config = yaml.safe_load(_config_file)
        else:
            overall_config = {}

        # The user can define global options that are inherited by
        # all the other sections if not already specified there.
        input_config = overall_config.get("global", {})

        # This is just the config info in the file for this stage.
        # It may be incomplete - there may be things specified on the
        # command line instead, or just using their default values
        stage_config = overall_config.get(self.instance_name, {})
        input_config.update(stage_config)

        self._configs.set_config(input_config, args)

    def get_config_dict(self, ignore=None, reduce_config=False):
        """Write the current configuration to a dict

        Parameters
        ----------
        ignore : dict or None
            Global parameters not to write
        reduce_config : bool
            If true, reduce the configuration by parsing out the inputs, outputs and global params

        Returns
        -------
        out_dict : dict
            The configuration
        """
        out_dict = {}
        if reduce_config:
            ignore_keys = self.input_tags() + self.output_tags() + ["config"]
        else:
            ignore_keys = []
        ignore = ignore or {}
        for key, val in self.config.items():
            if reduce_config:
                if key in ignore:
                    if ignore[key] == val:
                        continue
                if key in ignore_keys:
                    continue
            out_dict[key] = cast_to_streamable(val)
        return out_dict

    def find_inputs(self, pipeline_files):
        """Find and retrun all the inputs associated to this stage in the FileManager

        These are returned as a dictionary of tag : path pairs
        """
        ret_dict = {}
        for tag, _ in self.inputs_():
            aliased_tag = self.get_aliased_tag(tag)
            ret_dict[aliased_tag] = pipeline_files[aliased_tag]
        return ret_dict

    def find_outputs(self, outdir):
        """Find and retrun all the outputs associated to this stage

        These are returned as a dictionary of tag : path pairs
        """
        ret_dict = {}
        for tag, ftype in self.outputs_():
            aliased_tag = self.get_aliased_tag(tag)
            if not aliased_tag in self._outputs.keys():  # pragma: no cover
                self._outputs[aliased_tag] = ftype.make_name(aliased_tag)
            ret_dict[aliased_tag] = f"{outdir}/{self._outputs[aliased_tag]}"
        return ret_dict

    def print_io(self, stream=sys.stdout):
        """Print out the tags, paths and types for all the inputs and outputs of this stage"""
        stream.write("Inputs--------\n")
        for tag, ftype in self.inputs_():
            aliased_tag = self.get_aliased_tag(tag)
            stream.write(
                f"{tag:20} : {aliased_tag:20} :{str(ftype):20} : {self._inputs[tag]}\n"
            )
        stream.write("Outputs--------\n")
        for tag, ftype in self.outputs_():
            aliased_tag = self.get_aliased_tag(tag)
            stream.write(
                f"{tag:20} : {aliased_tag:20} :{str(ftype):20} : {self._outputs[aliased_tag]}\n"
            )

    def should_skip(self, run_config):
        """Return true if we should skip a stage b/c it's outputs already exist and we are in resume mode"""
        outputs = self.find_outputs(run_config["output_dir"]).values()
        already_run_stage = all(os.path.exists(output) for output in outputs)
        return already_run_stage and run_config["resume"]

    def already_finished(self):
        """Print a warning that a stage is being skipped"""
        print(f"Skipping stage {self.instance_name} because its outputs exist already")

    def iterate_fits(
        self, tag, hdunum, cols, chunk_rows, parallel=True
    ):  # pragma: no cover
        """
        Loop through chunks of the input data from a FITS file with the given tag

        TODO: add ceci tests of this functions
        Parameters
        ----------
        tag: str
            The tag from the inputs list to use

        hdunum: int
            The extension number to read

        cols: list
            The columns to read

        chunk_rows: int
            Number of columns to read and return at once

        parallel: bool
            Whether to split up data among processes (parallel=True) or give
            all processes all data (parallel=False).  Default = True.

        Returns
        -------
        it: iterator
            Iterator yielding (int, int, array) tuples of (start, end, data)
            data is a structured array.
        """
        fits = self.open_input(tag)
        ext = fits[hdunum]
        n = ext.get_nrows()
        for start, end in self.data_ranges_by_rank(n, chunk_rows, parallel=parallel):
            data = ext.read_columns(cols, rows=range(start, end))
            yield start, end, data

    def iterate_hdf(
        self, tag, group_name, cols, chunk_rows, parallel=True, longest=False
    ):
        """
        Loop through chunks of the input data from an HDF5 file with the given tag.

        All the selected columns must have the same length.

        Parameters
        ----------
        tag: str
            The tag from the inputs list to use

        group: str
            The group within the HDF5 file to use, looked up as
            file[group]

        cols: list
            The columns to read

        chunk_rows: int
            Number of columns to read and return at once

        parallel: bool
            Whether to split up data among processes (parallel=True) or give
            all processes all data (parallel=False).  Default = True.

        longest: bool
            Whether to allow mixed length arrays and keep going until the longest
            array is completed, returning empty arrays for shorter ones

        Returns
        -------
        it: iterator
            Iterator yielding (int, int, dict) tuples of (start, end, data)
        """
        import numpy as np

        hdf = self.open_input(tag)
        group = hdf[group_name]

        # Check all the columns are the same length
        N = [len(group[col]) for col in cols]
        n = max(N)
        if not longest:
            if not np.equal(N, n).all():
                raise ValueError(
                    f"Different columns among {cols} in file {tag} group {group_name}"
                    "are different sizes - if this is acceptable set longest=True"
                )

        # Iterate through the data providing chunks
        for start, end in self.data_ranges_by_rank(n, chunk_rows, parallel=parallel):
            data = {col: group[col][start:end] for col in cols}
            yield start, end, data

    ################################
    # Pipeline-related methods
    ################################

    @classmethod
    def generate_command(
        cls, inputs, config, outputs, aliases=None, instance_name=None
    ):
        """
        Generate a command line that will run the stage
        """
        module = cls.get_module()
        module = module.split(".")[0]

        if sys.modules[module].__file__:
            # Regular module, stage will be imported with module
            flags = [f"{cls.name}"]
        else:
            # Namescape module, use 'ceci' to the get main
            # and specify the full path
            flags = [f"{cls.get_module()}.{cls.name}"]
            module = "ceci"

        aliases = aliases or {}

        for tag, _ in cls.inputs_():
            aliased_tag = aliases.get(tag, tag)
            try:
                fpath = inputs[aliased_tag]
            except KeyError as msg:  # pragma: no cover
                raise ValueError(
                    f"Missing input location {aliased_tag} {str(inputs)}"
                ) from msg
            flags.append(f"--{tag}={fpath}")

        if instance_name is not None and instance_name != cls.name:
            flags.append(f"--name={instance_name}")

        flags.append(f"--config={config}")

        for tag, _ in cls.outputs_():
            aliased_tag = aliases.get(tag, tag)
            try:
                fpath = outputs[aliased_tag]
            except KeyError as msg:  # pragma: no cover
                raise ValueError(
                    f"Missing output location {aliased_tag} {str(outputs)}"
                ) from msg
            flags.append(f"--{tag}={fpath}")

        flags = "   ".join(flags)

        # We just return this, instead of wrapping it in a
        # parsl job
        cmd = f"python3 -m {module} {flags}"
        return cmd

    @classmethod
    def generate_cwl(cls, log_dir=None):
        """
        Produces a CWL App object which can then be exported to yaml
        """
        import cwl_utils.parser.cwl_v1_0 as cwlgen

        module = cls.get_module()
        module = module.split(".")[0]

        # Basic definition of the tool
        cwl_tool = cwlgen.CommandLineTool(
            [],
            [],
            id=cls.name,
            label=cls.name,
            baseCommand="python3",
            cwlVersion="v1.0",
            doc=cls.__doc__,
            arguments=[],
        )
        if log_dir is not None:
            cwl_tool.stdout = f"{cls.name}.out"
            cwl_tool.stderr = f"{cls.name}.err"

        # Adds the first input binding with the name of the module and pipeline stage
        input_arg = cwlgen.CommandLineBinding(position=-1, valueFrom=f"-m{module}")
        cwl_tool.arguments.append(input_arg)
        input_arg = cwlgen.CommandLineBinding(position=0, valueFrom=f"{cls.name}")
        cwl_tool.arguments.append(input_arg)

        type_dict = {int: "int", float: "float", str: "string", bool: "boolean"}
        # Adds the parameters of the tool
        for opt, def_val in cls.config_options.items():

            # Handles special case of lists:
            if isinstance(def_val, list):
                v = def_val[0]
                param_type = {
                    "type": "array",
                    "items": type_dict[v]
                    if isinstance(v, type)
                    else type_dict[type(v)],
                }
                default = def_val if not isinstance(v, type) else None
                input_binding = cwlgen.CommandLineBinding(
                    prefix=f"--{opt}=", itemSeparator=",", separate=False
                )
            else:
                param_type = (
                    type_dict[def_val]
                    if isinstance(def_val, type)
                    else type_dict[type(def_val)]
                )
                default = def_val if not isinstance(def_val, type) else None
                if param_type == "boolean":
                    input_binding = cwlgen.CommandLineBinding(prefix=f"--{opt}")
                else:  # pragma: no cover
                    input_binding = cwlgen.CommandLineBinding(
                        prefix=f"--{opt}=", separate=False
                    )

            input_param = cwlgen.CommandInputParameter(
                opt,
                label=opt,
                type=param_type,
                inputBinding=input_binding,
                default=default,
                doc="Some documentation about this parameter",
            )

            # We are bypassing the cwlgen builtin type check for the special case
            # of arrays until that gets added to the standard
            if isinstance(def_val, list):
                input_param.type = param_type

            cwl_tool.inputs.append(input_param)

        # Add the inputs of the tool
        for i, inp in enumerate(cls.input_tags()):
            input_binding = cwlgen.CommandLineBinding(prefix=f"--{inp}")
            input_param = cwlgen.CommandInputParameter(
                inp,
                label=inp,
                type="File",
                format=cls.inputs[i][1].format,  # pylint: disable=no-member
                inputBinding=input_binding,
                doc="Some documentation about the input",
            )
            cwl_tool.inputs.append(input_param)

        # Adds the overall configuration file
        input_binding = cwlgen.CommandLineBinding(prefix="--config")
        input_param = cwlgen.CommandInputParameter(
            "config",
            label="config",
            type="File",
            format="http://edamontology.org/format_3750",
            inputBinding=input_binding,
            doc="Configuration file",
        )
        cwl_tool.inputs.append(input_param)

        # Add the definition of the outputs
        for i, out in enumerate(cls.output_tags()):
            output_name = cls.outputs[i][1].make_name(out)  # pylint: disable=no-member
            output_binding = cwlgen.CommandOutputBinding(glob=output_name)
            output = cwlgen.CommandOutputParameter(
                out,
                label=out,
                type="File",
                outputBinding=output_binding,
                format=cls.outputs[i][1].format,  # pylint: disable=no-member
                doc="Some results produced by the pipeline element",
            )
            cwl_tool.outputs.append(output)

        if log_dir is not None:
            output = cwlgen.CommandOutputParameter(
                f"{cls.name}@stdout",
                label="stdout",
                type="stdout",
                doc="Pipeline elements standard output",
            )
            cwl_tool.outputs.append(output)
            error = cwlgen.CommandOutputParameter(
                f"{cls.name}@stderr",
                label="stderr",
                type="stderr",
                doc="Pipeline elements standard output",
            )
            cwl_tool.outputs.append(error)

        # Potentially add more metadata
        # This requires a schema however...
        # metadata = {'name': cls.name,
        #         'about': 'Some additional info',
        #         'publication': [{'id': 'one_doi'}, {'id': 'another_doi'}],
        #         'license': ['MIT']}
        # cwl_tool.metadata = cwlgen.Metadata(**metadata)

        return cwl_tool


from .AstroDataset import Minibatch

import torch

from lhotse.dataset import OnTheFlyFeatures
from lhotse.dataset.input_strategies import BatchIO, PrecomputedFeatures
from lhotse import Fbank, FbankConfig
from lhotse import CutSet
from lhotse.workarounds import Hdf5MemoryIssueFix
from lhotse.dataset.speech_recognition import validate_for_asr
from lhotse.utils import ifnone
from lhotse.dataset.input_strategies import AudioSamples

class Wav2Vec2Dataset(torch.utils.data.Dataset):
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        self.hdf5_fix = Hdf5MemoryIssueFix(reset_interval=100)
        self.input_method = AudioSamples()

    def __getitem__(self, cuts: CutSet) -> Minibatch:
        validate_for_asr(cuts)
        self.hdf5_fix.update()
        cuts = cuts.sort_by_duration(ascending=False)

        inputs, _ = self.input_method(cuts)
        supervision_intervals = self.input_method.supervision_intervals(cuts)
        tokens, token_lens = self.tokenizer(cuts)
        metadata = {
            'input_lens': supervision_intervals['num_samples'],
            'target_lens': token_lens,
            'utt_ids': [s.recording_id for cut in cuts for s in cut.supervisions],
            'text': [s.text for cut in cuts for s in cut.supervisions],
            'ids': [s.id for cut in cuts for s in cut.supervisions],
            'intervals': supervision_intervals,
        }
        targets = [t for t in tokens]
        return Minibatch(
            {
                'input': inputs,
                'targets': targets,
                'metadata': metadata,
            }
        )


import sys
import  socket

objetivo = socket.gethostbyname(input("Inserte la dirección IP: "))

# Acá ingresamos (por ej) la dirección ip obtenida anteriormente
print("Escaneando objetivo: " + objetivo)

try: # Recorremos uno a uno por lo que la búsqueda va a ser lenta
    for port in range (1, 150):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        socket.setdefaulttimeout(1)
        resultado = s.connect_ex((objetivo, port))
        if resultado == 0:
            print(f"El puerto {port} está abierto")
        s.close()
except:
    print("\n Saliendo...")
    sys.exit(0)


from openpyxl import load_workbook
import xlsxwriter

setting = {
    # 用于转换一些中文的字段名转换成英文，数据在输出前会修改字典返回的键名,默认转换是从第一个字段开始
    'CHANGE_FIELD': ('sys_sku','mum_id',)
}
# 操作xlsx文档某一页的类
class XlsxSheet:
    # 初始化
    # @ workBookPath str 工作薄的地址，绝对位置或者相对位置
    # @ sheet  第几个工作薄，以数字从0开始，
    def __init__(self, load_workbook, workBookPath, sheetNum):
        self.workBookPath = workBookPath
        self.sheetNum = sheetNum
        self.wb = load_workbook(workBookPath)
        self.sheets = self.wb.worksheets
        self.sheet = self.sheets[sheetNum]
        # 定义一个字典用来存放一行数据
        self.data = {}
        self.field = self.get_field()
        # 设置文件
        #

    # 根据行数查询一行记录
    def get_row_by_id(self, rowNum = 2):
        data = self.sheet[rowNum]
        lis = []
        for row in data:
            lis.append(row.value)

        self.data = self.format_data(self.field,lis)
        return self.data

    # 取得行数据中的cell 对象的value值返回一个list
    # @param data Tuple 元祖对象
    def get_cell_values(self,data):
        lis = []
        for row in data:
            lis.append(row.value)
        return lis

    # 根据条件查询多行
    # @param dic dict 包含条件的字典数据
    # return list or null 符合参数dic 键名 = 值的数据会被以list的形式返回
    def get_many_rows(self, dic):
        datas = []
        for row in self.sheet:
            d = self.get_cell_values(row)
            data = self.format_data(self.field, d)
            boo = 0
            for i in dic.items():
                if(i[1] != data[i[0]]):
                    boo = 1
            if boo != 1:
                # print(data)
                datas.append(data)
        return datas

    # 添加一行到文本尾部
    # @param list 添加的一行从A1开始按照list顺序添加
    def append_row(self,data):
        # 需要打开文档的可编辑权限，要不调用该函数就会报错
        # 建文件及sheet.
        # 获取当前活跃的sheet，默认是第一个sheet
        ws = self.wb.active
        # ws['A1'] = 'class'
        # ws['B1'] = 'name'
        # ws['C1'].value = 'score'
        # row1 = ['class1', 'zhangsan', 90]
        # row2 = ['class2', 'lisi', 88]
        # ws.append(row1)
        ws.append(data)
        self.wb.save(self.workBookPath)

    # 取字段数据，默认取当前sheet的第一行数据，一般为字段名
    # @ num 取第几行元素默认为1
    # @return list 返回一个包含第一行内容的list
    def get_field(self, num = 1):
        li = []
        data = self.sheet[num]
        for row in data:
            li.append(row.value)

        # 转换字段名
        for i in range(len(setting['CHANGE_FIELD'])):
            li[i] = setting['CHANGE_FIELD'][i]

        return li

    # 格式化数据,合并两个list 转换成一个字典
    # @param keys list 作为键名的list
    # @param data list 作为值的list
    # @return dict
    def format_data(self, keys, data):
        dic = {}
        for i in range(len(keys)):
            dic[keys[i]] = data[i]
        return dic

    # 更新数据
    # @ param
    def update_row(self):
        pass

    # 删除一行数据
    def delete_row(self):
        pass

def main():
    xlsxSheet = XlsxSheet(load_workbook, '图片资料和供应商(1).xlsx', 0)
    # 取表头字段
    # print(xlsxSheet.get_field())

    #取数据
    # print(xlsxSheet.get_row_by_id(4))

    # 查询符合条件的数据
    lis = xlsxSheet.get_many_rows({'sys_sku':'3179347'})
    print(lis)
    xlsxSheet.append_row(['1222111','4556554456'])
    # sublime 不支持repr 输出
    # repr(lis)
if __name__ == '__main__':
    main()


#!/usr/bin/python3
import socket
import threading
import subprocess

bind_ip = "0.0.0.0"
bind_port = 9999

server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
server.bind((bind_ip, bind_port))
server.listen(5)

print("* Listening on %s:%d" % (bind_ip, bind_port))

def handle_client(client_socket):
	while True:
		client_socket.send(b"Input (a) name(s) of file/directory => ")
		request = client_socket.recv(1024).decode().strip()
		print("* Received: %s" % request)
		if (request == "end"):
			break

		command = "ls " + request
		print(command)
		subprocess.run(command, shell=True, stdout=client_socket, stderr=client_socket, text=True)
	client_socket.close()

while True:
	client, addr = server.accept()
	print("* Accepted connection from: %s:%d" % (addr[0], addr[1]))

	client_handler = threading.Thread(target=handle_client, args=(client,))
	client_handler.start()


# -*- coding:utf-8 -*-
from typing import Optional
from copy import deepcopy
from ._fieldinfo import FieldInfo, _not_set, is_set
from .processors import Const
from ._packetbase import FieldBase

__all__ = ['Field']

class Field(FieldBase):
    __instances_created = 0

    @property
    def info(self):
        return self._info

    def __new__(cls, *args, **kwargs):
        instance = object.__new__(cls)
        cls.__instances_created += 1
        instance.__number = cls.__instances_created
        return instance

    def __init__(self, processor=_not_set, name=_not_set, default=_not_set, required=_not_set, override=_not_set):
        """Packet field constructor

        Args:
            processor (FieldProcessor, optional): field type processor. Defaults to _not_set.
            name (str, optional): serialized field name. Defaults to _not_set.
            default (Any, optional): default field value. Defaults to _not_set.
            required (bool, optional): flag if the field is required. Defaults to _not_set.
            override (bool, optional): flag if the field is overloaded. Defaults to _not_set.
        """
        self._info = FieldInfo(processor, name, default, required, override)

    @property
    def name(self, default=None) -> Optional[str]:
        return self.info.name if is_set(self.info.name) else default

    def on_packet_class_create(self, parent_field, field_name):
        """Callback to set field name on packet creation

        Args:
            parent_field (Field): parent field
            field_name (str): in-python field name
        """

        if parent_field is not None:
            if not self._info.override == True:
                raise TypeError(f'Repeated field {field_name}')
            else:
                new_info = parent_field._info.copy()
                new_info.update(self._info)
                self._info = new_info
        self._info.set_defaults(None, False, False)
        self._info.update_name(field_name)

    def raw_to_py(self, raw_value, strict):
        if raw_value is None:
            if self._info.py_default is None:
                py_value = None
            else:
                py_value = deepcopy(self._info.py_default)
        else:
            self._info.processor.check_raw(raw_value)
            py_value = self._info.processor.raw_to_py(raw_value, strict)

        if self._info.required and py_value is None:
            if not strict:
                py_value = self._info.processor.zero_value
            else:
                raise ValueError(f'Field required {self}')
        return py_value

    def py_to_raw(self, py_value):
        if py_value is None:
            if self._info.py_default is None:
                raw_value = None
            else:
                raw_value = deepcopy(self._info.default)
        else:
            self._info.processor.check_py(py_value)
            raw_value = self._info.processor.py_to_raw(py_value)

        if self._info.required and raw_value is None:
            raise ValueError(f'Field required {self}')
        return raw_value

    def clone(self, **kwargs):
        new_info = self._info.copy()
        new_info.update_params(**kwargs)
        field = self.__class__()
        field._info = new_info
        return field

    def frozen_clone(self, value):
        field = self.__class__(
            Const(value),
            name=self._info.name,
            default=value,
            required=self._info.required,
            override=True
        )
        field.on_packet_class_create(self, self._info.py_name)
        return field

    def dump_partial(self, value):
        return self._info.processor.dump_partial(value)

    def __str__(self):
        return f'<{self.__class__.__name__}("{self._info.py_name}", "{self._info.name}")>'

    def __cmp__(self, other):
        return ((self.__number > other.__number) - (self.__number < other.__number))


from .util import get_data_file_path
import logging
from collections import OrderedDict, namedtuple
from operator import itemgetter
import re
import string
import traceback
from vt.object import Object

__all__ = ("SampleInfo",
           "LabeledSample",
           "AvLabels",)

SampleInfo = namedtuple('SampleInfo',
                        ['md5', 'sha1', 'sha256', 'labels'])

LabeledSample = namedtuple('LabeledSample',
                           ['md5', 'sha1', 'sha256', 'labels', 'family', 'tokens', 'groundtruth', 'is_pup'])

def log(message):
    print(message)

class AvLabels(object):

    def __init__(self, gen_file=None, alias_file=None, av_file=None):
        """Class to operate on AV labels, such as extracting the most likely family name.

        :param gen_file:
        :param alias_file:
        :param av_file:
        """
        # Read generic token set from file
        try:
            if gen_file is None:
                gen_file = get_data_file_path("default.generics")
            self.gen_set = self.read_generics(gen_file)
        except:
            logging.warning(traceback.format_exc())
            self.gen_set = set()

        # Read aliases map from file
        try:
            if alias_file is None:
                alias_file = get_data_file_path("default.aliases")
            self.aliases_map = self.read_aliases(alias_file)
        except:
            logging.warning(traceback.format_exc())
            self.aliases_map = {}

        # Read AV engine set from file
        self.avs = self.read_avs(av_file) if av_file else None

    @staticmethod
    def read_aliases(alfile):
        """Read aliases map from given file"""
        if alfile is None:
            return {}
        almap = {}
        with open(alfile, 'r') as fd:
            for line in fd:
                alias, token = line.strip().split()[0:2]
                almap[alias] = token
        return almap

    @staticmethod
    def read_generics(generics_file):
        """Read generic token set from given file"""
        gen_set = set()
        with open(generics_file) as gen_fd:
            for line in gen_fd:
                if line.startswith('#') or line == '\n':
                    continue
                gen_set.add(line.strip())
        return gen_set

    @staticmethod
    def read_avs(avs_file):
        """Read AV engine set from given file"""
        with open(avs_file) as fd:
            avs = set(map(str.strip, fd.readlines()))
        return avs

    @staticmethod
    def get_sample_info(data):
        """Parse and extract sample information from JSON data
           Returns a SampleInfo named tuple: md5, sha1, sha256, label_pairs

           This method has been improved to handle multiple JSON data formats
           Recognized formats include:
           - VT File Report
           - VT Notification
           - AVClass simplified JSON
        """
        def clean(value):
            return "".join([x for x in value if x in string.printable]).strip()

        if isinstance(data, Object):
            if data.type == 'file':
                label_pairs = list(
                    map(lambda r: (r['engine_name'], clean(r['result'])),
                        filter(lambda x: x['result'] is not None,
                               data.last_analysis_results.values()))
                )
                return SampleInfo(data.md5, data.sha1, data.sha256, label_pairs)
        elif isinstance(data, dict):
            try:
                if "response_code" in data:
                    # VT file report

                    if data["response_code"] == 0:
                        return None
                    label_pairs = [(av, clean(result["result"])) for av, result in data["scans"].items() if
                                   result["result"] is not None]
                elif "ruleset_name" in data:
                    # VT notification

                    label_pairs = [(av, clean(result)) for av, result in data["scans"].items() if result is not None]
                else:
                    label_pairs = data["av_labels"]
            except KeyError:
                return None

            return SampleInfo(str(data['md5']), str(data['sha1']), str(data['sha256']), label_pairs)
        return None

    @staticmethod
    def is_pup(av_label_pairs):
        """This function classifies the sample as PUP or not
           using the AV labels as explained in the paper:
           "Certified PUP: Abuse in Authenticode Code Signing"
           (ACM CCS 2015)
           It uses the AV labels of 11 specific AVs.
           The function checks for 13 keywords used to indicate PUP.
           Return:
              True/False/None
        """
        # If no AV labels, nothing to do, return
        if not av_label_pairs:
            return None
        # Initialize
        pup = False
        threshold = 0.5
        # AVs to use
        av_set = {'Malwarebytes', 'K7AntiVirus', 'Avast', 'AhnLab-V3', 'Kaspersky', 'K7GW', 'Ikarus', 'Fortinet',
                  'Antiy-AVL', 'Agnitum', 'ESET-NOD32'}
        # Tags that indicate PUP
        tags = {'PUA', 'Adware', 'PUP', 'Unwanted', 'Riskware', 'grayware', 'Unwnt', 'Adknowledge', 'toolbar', 'casino',
                'casonline', 'AdLoad', 'not-a-virus'}

        # Set with (AV name, Flagged/not flagged as PUP), for AVs in av_set
        bool_set = set([(pair[0], t.lower() in pair[1].lower()) for t in tags
                        for pair in av_label_pairs
                        if pair[0] in av_set])

        # Number of AVs that had a label for the sample
        av_detected = len([p[0] for p in av_label_pairs
                           if p[0] in av_set])

        # Number of AVs that flagged the sample as PUP
        av_pup = list(map(lambda x: x[1], bool_set)).count(True)  # python 2/3, inefficient on Py2

        # Flag as PUP according to a threshold
        if (float(av_pup) >= float(av_detected) * threshold) and av_pup != 0:
            pup = True
        return pup

    @staticmethod
    def __remove_suffixes(av_name, label):
        """Remove AV specific suffixes from given label

        :param av_name:
        :param label:
        :return: updated label
        """

        # Truncate after last '.'
        if av_name in {'Norman', 'Avast', 'Avira', 'Kaspersky', 'ESET-NOD32', 'Fortinet', 'Jiangmin', 'Comodo', 'GData',
                       'Avast', 'Sophos', 'TrendMicro-HouseCall', 'TrendMicro', 'NANO-Antivirus', 'Microsoft'}:
            label = label.rsplit('.', 1)[0]

        # Truncate after last '.'
        # if suffix only contains digits or uppercase (no lowercase) chars
        if av_name == 'AVG':
            tokens = label.rsplit('.', 1)
            if len(tokens) > 1 and re.match("^[A-Z0-9]+$", tokens[1]):
                label = tokens[0]

        # Truncate after last '!'
        if av_name in {'Agnitum', 'McAffee', 'McAffee-GW-Edition'}:
            label = label.rsplit('!', 1)[0]

        # Truncate after last '('
        if av_name in {'K7AntiVirus', 'K7GW'}:
            label = label.rsplit('(', 1)[0]

        # Truncate after last '@'
        # GData would belong here, but already trimmed earlier
        if av_name in {'Ad-Aware', 'BitDefender', 'Emsisoft', 'F-Secure', 'Microworld-eScan'}:
            label = label.rsplit('(', 1)[0]

        return label

    def update_aliases(self, alias_file):
        try:
            alias_map = AvLabels.read_aliases(alias_file)
            self.aliases_map.update(alias_map)
        except:
            logging.warning(traceback.format_exc())

    def __normalize(self, label, hashes):
        """Tokenize label, filter tokens, and replace aliases"""

        # If empty label, nothing to do
        if not label:
            return []

        # Initialize list of tokens to return
        ret = []

        # Split label into tokens and process each token
        for token in re.split("[^0-9a-zA-Z]", label):
            # Convert to lowercase
            token = token.lower()

            # Remove digits at the end
            end_len = len(re.findall("\d*$", token)[0])
            if end_len:
                token = token[:-end_len]

            # Ignore short token
            if len(token) < 4:
                continue

            # Remove generic tokens
            if token in self.gen_set:
                continue

            # Ignore token if prefix of a hash of the sample
            # Most AVs use MD5 prefixes in labels,
            # but we check SHA1 and SHA256 as well
            hash_token = False
            for hash_str in hashes:
                if hash_str[0:len(token)] == token:
                    hash_token = True
                    break
            if hash_token:
                continue

            # Replace alias
            token = self.aliases_map[token] if token in self.aliases_map \
                else token

            # Add token
            ret.append(token)
        return ret

    def get_family_ranking(self, sample_info):
        """Returns sorted dictionary of most likely family names for sample

        :param sample_info:
        :return:
        """
        # Extract info from named tuple
        av_label_pairs = sample_info[3]
        hashes = [sample_info[0], sample_info[1], sample_info[2]]

        # Whitelist the AVs to filter the ones with meaningful labels
        av_whitelist = self.avs

        # Initialize auxiliary data structures
        labels_seen = set()
        token_map = {}

        # Process each AV label
        for (av_name, label) in av_label_pairs:
            # If empty label, nothing to do
            if not label:
                continue

            ################
            # AV selection #
            ################
            if av_whitelist and av_name not in av_whitelist:
                continue

            #####################
            # Duplicate removal #
            #####################

            # If label ends in ' (B)', remove it
            if label.endswith(' (B)'):
                label = label[:-4]

            # If we have seen the label before, skip
            if label in labels_seen:
                continue
            # If not, we add it to the set of labels seen
            else:
                labels_seen.add(label)

            ##################
            # Suffix removal #
            ##################
            label = self.__remove_suffixes(av_name, label)

            ########################################################
            # Tokenization, token filtering, and alias replacement #
            ########################################################
            tokens = self.__normalize(label, hashes)

            # Increase token count in map
            for t in tokens:
                c = token_map[t] if t in token_map else 0
                token_map[t] = c + 1

        ##################################################################
        # Token ranking: sorts tokens by decreasing count and then token #
        ##################################################################
        sorted_tokens = sorted(token_map.items(),
                               key=lambda x: x[1] or 0,
                               reverse=True)

        # Delete the tokens appearing only in one AV, add rest to output
        sorted_dict = OrderedDict()
        for t, c in sorted_tokens:
            if c > 1:
                sorted_dict[t] = c
            else:
                break

        return sorted_dict


import os
import time
import json
import subprocess
import pdb
import consts

from infra.app import AppFactory

asic = os.environ.get('ASIC', 'capri')
HNTAP_CMD = 'build/x86_64/iris/' + asic + '/bin/nic_infra_hntap'
CONTAINER_HNTAP_CMD = "naples/nic/tools/restart-hntap.sh"

def preexec(): # Don't forward signals.
    os.setpgrp()

class Hntap(object):

    def __init__(self, conf_file):
        self._conf_file = conf_file

    def Run(self, nomodel=False):
        pass

    def Stop(self):
        pass

    def _wait_for_hntap_up(self, log_file=consts.hntap_log):
        log2 = open(log_file, "r")
        loop = 1
        time.sleep(2)

        # Wait until tap devices setup is complete
        while loop == 1:
            for line in log2.readlines():
                if "listening on" in line:
                    loop = 0
        log2.close()

class HntapLocal(Hntap):

    def __init__(self, conf_file):
        self._conf_file = conf_file
        self._process = None

    def Run(self, nomodel=False):
        log = open(consts.hntap_log, "w")
        cmd = [HNTAP_CMD, '-f', self._conf_file, '-n', '2']
        if nomodel:
            cmd.append("-s")
        os.chdir(consts.nic_dir)
        print (cmd)
        p = subprocess.Popen(cmd, stdout=log, stderr=log, preexec_fn = preexec)
        self._process = p
        print("* Starting Host-Network tapper, pid (%s)" % str(p.pid))
        print("    - Log file: " + consts.hntap_log + "\n")
        self._wait_for_hntap_up(consts.hntap_log)
        return

    def Stop(self):
        os.kill(int(self._process.pid), 9)

class HntapInContainer(Hntap):

    def __init__(self, conf_file, nic_container,
                container_conf_path=consts.hntap_container_conf_path):
        self._conf_file = conf_file
        self._nic_container_hntap_cfg_file = container_conf_path
        self._nic_container = nic_container
        self._process = None

    def Run(self, nomodel=False):
        log = open(consts.hntap_container_log, "w")
        log.close()
        #First copy the new hntap config file
        print ("Copying Hntap configuration file to nic container")
        copy_cmd = "docker cp " + self._conf_file + "  " + self._nic_container + ":" + self._nic_container_hntap_cfg_file
        p = subprocess.run(copy_cmd, shell=True)

        #Restart Hntap process
        print ("Restarting Hntap process in nic container")
        nic_container = AppFactory.Get("nic_container", id=self._nic_container)
        nic_container.RunCommand(CONTAINER_HNTAP_CMD, background=True, tty=False)
        self._wait_for_hntap_up(consts.hntap_container_log)

        #Move the interface to namespace 1
        print ("Moving interfaces to namespace 1")
        data = json.load(open(self._conf_file))
        for intf in data:
            print ("Moving interface " + intf["name"])
            nic_container.MoveInterface(intf["name"], 1)

    def Stop(self):
        pass

TypeHntapInContainer = HntapInContainer
TypeHntapLocal = HntapLocal

class HntapFactory():
    @staticmethod
    def Get(config_file, container=None):
        if container:
            return HntapInContainer(config_file, container)
        else:
            return HntapLocal(config_file)


class IMF_DataStructure(object):
    def __init__(self, database_id, json_data):
        self.database_id        = database_id
        self.structure          = json_data['Structure']

        # Parse out data from the Header section
        self.header             = self.structure['Header']
        self.contact_info       = {
            'telephone_number'  : self.header['Sender']['Contact']['Telephone']   ,
            'uri'               : self.header['Sender']['Contact']['URI']
        }
        self.available_frequencies  = {}
        self.geographical_areas     = {}
        self.indicators             = {}
        self.time_values            = {}

        #############################################################
        # Parse out data from the CodeLists section of the document #
        #############################################################
        raw_code_list           = self.structure['CodeLists']['CodeList'][0]['Code']
        self.codes              = {}
        for code in raw_code_list:
            key = code['@value']
            val = code['Description']['#text']
            self.codes[key] = val

        # Get Frequency Data
        for code in self.structure['CodeLists']['CodeList'][1]['Code']:
            key     = code['@value']
            value   = code['Description']['#text']
            self.available_frequencies[key] = value

        # Get Geographical Areas
        for code in self.structure['CodeLists']['CodeList'][2]['Code']:
            key     = code['@value']
            value   = code['Description']['#text']
            self.geographical_areas[key] = value

        # Get Indicators
        for code in self.structure['CodeLists']['CodeList'][3]['Code']:
            key     = code['@value']
            value   = code['Description']['#text']
            self.indicators[key] = value

        # Get Time Values
        for code in self.structure['CodeLists']['CodeList'][4]['Code']:
            key     = code['@value']
            value   = code['Description']['#text']
            self.time_values[key] = value

        #######################################################
        # Parse out any needed data from the Concepts section #
        #######################################################
        self.concept = self.structure['Concepts']['ConceptScheme']['Name']['#text']

        # Parse out any needed data from the KeyFamilies section
        #print(self.structure['KeyFamilies']['KeyFamily']['Components'])
        #print(self.structure['KeyFamilies']['KeyFamily']['Annotations'])
        raw_annotations = self.structure['KeyFamilies']['KeyFamily']['Annotations']['Annotation']
        self.annotations = {}
        for annotation in raw_annotations:
            key = annotation['AnnotationTitle']
            value = annotation['AnnotationText']['#text']
            self.annotations[key] = value
        #print(self.structure['KeyFamilies']['KeyFamily'].keys())

    def __str__(self):
        """
        Human readable representation of the data
        :return: Well Formatted Text Representation of this object
        """
        #print(self.structure.keys())
        #print(self.header.keys())
        desc  = "\n{0}\n".format(self.concept)
        desc += "{0:-<{width}}\n".format('-',width=len(self.concept))
        desc += "Database ID: {0}\n".format(self.database_id)
        desc += "Contact URI: {0}\n".format(self.contact_info['uri'])
        desc += "Contact Telephone: {0}\n".format(self.contact_info['telephone_number'])

        section_title = "Frequencies"
        desc += "\n{0}\n".format(section_title)
        for code, description in self.available_frequencies.items():
            desc += "{0}: {1}\n".format(code, description)

        section_title = "Annotations"
        desc += "\n{0}\n".format(section_title)
        for annotation, annotation_text in self.annotations.items():
            desc += "{0}: {1}\n".format(annotation, annotation_text)
        return desc


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from matplotlib import rcParams
rcParams["figure.figsize"] = 16, 6
plt.style.use("seaborn-white")

import os
from subprocess import check_output

import tensorflow as tf

PATH = "../input/"

print(check_output(["ls","../input/"]).decode("utf8"))

train1 = pd.read_csv(PATH + "stage1_train_labels.csv")
ss1 = pd.read_csv(PATH + "stage1_sample_submission.csv")

train1.head()

print("There are {} rows of data.".format(train1.shape[0]))

TARGET = "EncodedPixels"

ss1.head()

def dimg(idx):
    """
    Displays image corresponding to the id
    """
    img = mpimg.imread(PATH+"stage1_train/"+idx+"/"+"images/"+idx+".png")
    return img

def dmsk(idx):
    """
    Displays the masks corersponding to id
    """
    f = os.listdir(PATH + "stage1_train/" + idx + "/masks")[0]
    nim = mpimg.imread(PATH + "stage1_train/" + idx + "/masks/" + f)

    for m in os.listdir(PATH + "stage1_train/" + idx + "/masks")[1:]:
        nim += mpimg.imread(PATH + "stage1_train/" + idx + "/masks/" + m)

    return nim

def dbth(idx):
    """
    Display both the mask and the image
    """
    fig, ax = plt.subplots(ncols=2)
    ax[0].imshow(dimg(idx))
    ax[1].imshow(dmsk(idx), cmap="Purples")
    plt.show()

for ind in train1.sample(5)["ImageId"].index:
    print("Image ID:", ind)
    dbth(train1.iloc[ind,0])


"""Dataset functions.
"""
import os
import pathlib
import re
import tempfile
from datetime import datetime
from zipfile import ZipFile

import tensorflow as tf
import numpy as np
from object_detection.protos import string_int_label_map_pb2
from google.protobuf import text_format

from cvat.apps.dataset_manager.task import export_task as cvat_export_task
import collections
import json

tf.compat.v1.enable_eager_execution()

def _cvat_get_frame_path(base_dir, frame):
    """CVAT's image directory layout.

    Specified in cvat.engine.models.py Task class
    """
    d1 = str(int(frame) // 10000)
    d2 = str(int(frame) // 100)
    path = os.path.join(base_dir, d1, d2,
                        str(frame) + '.jpg')

    return path

def dump_cvat_task_annotations(db_task, db_user, scheme, host, format_name=None):
    """Use CVAT's utilities to dump annotations for a task."""
    timestamp = datetime.now().strftime('%Y_%m_%d_%H_%M_%S')

    if format_name is None:
        if db_task.mode == 'annotation':
            format_name = "CVAT for images 1.1"
        else:
            format_name = "CVAT for video 1.1"

    output_file_path = os.path.join(
        db_task.get_task_dirname(),
        '{}.{}.{}.zip'.format(db_task.id, db_user.username, timestamp)
    )

    cvat_export_task(
        task_id=db_task.id,
        dst_file=output_file_path,
        format_name=format_name,
        server_url=scheme + host,
        save_images=True,
    )
    return output_file_path

def fix_cvat_tfrecord(cvat_tf_record_zip, output_file_path):
    """Fix cvat tfrecord to comply with TF's object detection api.
    - change label id to invalid -1: so that TF is forced to use image/object/class/text
    """
    with tempfile.TemporaryDirectory() as temp_dir:
        with ZipFile(cvat_tf_record_zip) as cur_zip:
            cur_zip.extractall(temp_dir)

        tfrecord_files = list(pathlib.Path(temp_dir).glob('*.tfrecord'))
        tfrecord_files = [str(tfrecord_file) for tfrecord_file in tfrecord_files]
        dataset = tf.data.TFRecordDataset(tfrecord_files)
        with tf.io.TFRecordWriter(str(output_file_path)) as writer:
            for item in iter(dataset):
                example = tf.train.Example()
                example.ParseFromString(item.numpy())

                # change class label to -1. force TF to use class/text
                for i in range(len(example.features.feature['image/object/class/label'].int64_list.value)):
                    example.features.feature['image/object/class/label'].int64_list.value[i] = -1

                writer.write(example.SerializeToString())

def get_label_map_from_cvat_tfrecord_zip(cvat_tf_record_zip):
    """Extract label map from cvat tfrecord zip file.
    CVAT's tfrecord file contains:
    - label_map.pbtxt
    - *.tfrecord
    """
    labels = []
    with tempfile.TemporaryDirectory() as temp_dir:
        with ZipFile(cvat_tf_record_zip) as cur_zip:
            with cur_zip.open('label_map.pbtxt', 'r') as f:
                content = f.read().decode('utf-8')
                cur_label_map = string_int_label_map_pb2.StringIntLabelMap()
                text_format.Merge(content, cur_label_map)
                for item in cur_label_map.item:
                    if item.name not in labels:
                        labels.append(item.name)
    return labels

def dump_metadata(metadata, output_file_path):
    with open(output_file_path, 'w') as f:
        json.dump(metadata, f)

def _dump_labelmap_file(labels, output_file_path):
    """Write out labels as tensorflow object detection API's lable_map.txt.
    https://github.com/tensorflow/models/blob/master/research/object_detection/data/kitti_label_map.pbtxt

    Label id 0 is reserved for 'background', therefore this file starts with id 1.
    """
    label_ids = collections.OrderedDict((label, idx + 1)
                                        for idx, label in enumerate(labels))
    with open(output_file_path, 'w', encoding='utf-8') as f:
        for label, idx in label_ids.items():
            f.write(u'item {\n')
            f.write(u'\tid: {}\n'.format(idx))
            f.write(u"\tname: '{}'\n".format(label))
            f.write(u'}\n\n')

def dump_detector_annotations(db_detector, db_tasks, db_user, scheme, host):
    """Dump annotation data for detector training.

    Output is placed into the detector's ondisk dir.
    """
    output_dir = db_detector.get_training_data_dir()
    output_labelmap_file_path = output_dir / 'label_map.pbtxt'

    # see cvat.apps.dataset_manager.formats
    # dump_format = 'COCO 1.0'
    dump_format = 'TFRecord 1.0'

    labels = []
    # call cvat dump tool on each video in the trainset
    for db_task in db_tasks:
        task_annotations_file_path = dump_cvat_task_annotations(
                db_task, db_user, scheme, host, format_name=dump_format)

        # force label_id's to -1
        fix_cvat_tfrecord(task_annotations_file_path, output_dir / (
            os.path.splitext(
                os.path.basename(task_annotations_file_path))[0] + '.tfrecord')
        )
        task_labels = get_label_map_from_cvat_tfrecord_zip(
            task_annotations_file_path
        )
        for label in task_labels:
            if label not in labels:
                labels.append(label)
        os.remove(task_annotations_file_path)

    _dump_labelmap_file(labels,
                        output_labelmap_file_path)
    split_train_eval_tfrecord(output_dir)

def split_train_eval_tfrecord(data_dir):
    """Split tfrecord in the data_dir into train and eval sets."""
    tfrecord_files = data_dir.glob('*.tfrecord')
    tfrecord_files = [str(tfrecord_file) for tfrecord_file in tfrecord_files]
    dataset = tf.data.TFRecordDataset(tfrecord_files)
    output_train_tfrecord_file_path = str(data_dir / 'train.tfrecord')
    output_eval_tfrecord_file_path = str(data_dir / 'eval.tfrecord')

    # get train/eval item ids
    total_num = 0
    eval_percentage = 0.1
    meta_data = {
        'train_num': 0,
        'eval_num': 0
    }
    for item in iter(dataset):
        total_num += 1
    eval_ids = np.random.choice(total_num,
                                int(eval_percentage * total_num), replace=False)

    with tf.io.TFRecordWriter(output_train_tfrecord_file_path) as train_writer:
        with tf.io.TFRecordWriter(output_eval_tfrecord_file_path) as eval_writer:
            for idx, item in enumerate(iter(dataset)):
                if idx in eval_ids:
                    eval_writer.write(item.numpy())
                    meta_data['eval_num'] += 1
                else:
                    train_writer.write(item.numpy())
                    meta_data['train_num'] += 1
    dump_metadata(
        meta_data,
        data_dir / 'meta'
    )

# def prepare_coco_dataset(annotation_file_path, cvat_image_dir, output_dir):
#     """Create a on-disk coco dataset with both images and annotations.
#     """
#     from pycocotools import coco as coco_loader

#     annotation_file_path = pathlib.Path(annotation_file_path).resolve()
#     output_dir = pathlib.Path(output_dir).resolve()
#     output_dir.mkdir(parents=True, exist_ok=True)

#     # annotation file
#     annotation_file_name = 'annotation.json'
#     os.symlink(annotation_file_path, output_dir / annotation_file_name)

#     # image files
#     output_data_dir = output_dir / 'images'
#     shutil.rmtree(output_data_dir, ignore_errors=True)
#     output_data_dir.mkdir()
#     coco_dataset = coco_loader.COCO(str(annotation_file_path))
#     coco_images = coco_dataset.loadImgs(coco_dataset.getImgIds())
#     cvat_frame_id_regex = re.compile(r'\d+')
#     for coco_image in coco_images:
#         coco_file_name = coco_image['file_name']
#         # cvat uses "frame_{:06d}".format(frame) as default file name
#         # see cvat.annotations.annotation
#         cvat_frame_id = int(cvat_frame_id_regex.findall(coco_file_name)[0])
#         input_image_file_path = _cvat_get_frame_path(cvat_image_dir,
#                                                      cvat_frame_id)
#         output_image_file_path = output_data_dir / coco_file_name
#         os.symlink(input_image_file_path, output_image_file_path)