"""
The ``predict`` subcommand allows you to make bulk JSON-to-JSON
or dataset to JSON predictions using a trained model and its
:class:`~allennlp.service.predictors.predictor.Predictor` wrapper.

.. code-block:: bash

    $ allennlp predict --help
    usage: allennlp predict [-h] [--output-file OUTPUT_FILE]
                            [--weights-file WEIGHTS_FILE]
                            [--batch-size BATCH_SIZE] [--silent]
                            [--cuda-device CUDA_DEVICE] [--use-dataset-reader]
                            [--dataset-reader-choice {train,validation}]
                            [-o OVERRIDES] [--predictor PREDICTOR]
                            [--include-package INCLUDE_PACKAGE]
                            archive_file input_file

    Run the specified model against a JSON-lines input file.

    positional arguments:
      archive_file          the archived model to make predictions with
      input_file            path to or url of the input file

    optional arguments:
      -h, --help            show this help message and exit
      --output-file OUTPUT_FILE
                            path to output file
      --weights-file WEIGHTS_FILE
                            a path that overrides which weights file to use
      --batch-size BATCH_SIZE
                            The batch size to use for processing
      --silent              do not print output to stdout
      --cuda-device CUDA_DEVICE
                            id of GPU to use (if any)
      --use-dataset-reader  Whether to use the dataset reader of the original
                            model to load Instances. The validation dataset reader
                            will be used if it exists, otherwise it will fall back
                            to the train dataset reader. This behavior can be
                            overridden with the --dataset-reader-choice flag.
      --dataset-reader-choice {train,validation}
                            Indicates which model dataset reader to use if the
                            --use-dataset-reader flag is set. (default =
                            validation)
      -o OVERRIDES, --overrides OVERRIDES
                            a JSON structure used to override the experiment
                            configuration
      --predictor PREDICTOR
                            optionally specify a specific predictor to use
      --include-package INCLUDE_PACKAGE
                            additional packages to include
"""
from typing import List, Iterator, Optional
import argparse
import sys
import torch
import json

from allennlp.commands.subcommand import Subcommand
from allennlp.common.checks import check_for_gpu, ConfigurationError
from allennlp.common.file_utils import cached_path
from allennlp.common.util import lazy_groups_of, sanitize
from allennlp.models.archival import load_archive
from allennlp.modules.scalar_mix import ScalarMix
from allennlp.predictors.predictor import Predictor, JsonDict
from allennlp.data import Instance
from tqdm import tqdm



def _get_predictor(args: argparse.Namespace) -> Predictor:
    check_for_gpu(args.cuda_device)
    print("loading from archive...")
    archive = load_archive(
        args.archive_file,
        weights_file=args.weights_file,
        cuda_device=args.cuda_device,
        overrides=args.overrides,
    )
    print("done!")
    return Predictor.from_archive(
        archive, args.predictor, dataset_reader_to_load=args.dataset_reader_choice
    )


class _PredictManager:
    def __init__(
        self,
        predictor: Predictor,
        input_file: str,
        output_file: Optional[str],
        batch_size: int,
        print_to_console: bool,
        has_dataset_reader: bool,
    ) -> None:

        self._predictor = predictor
        self._input_file = input_file
        if output_file is not None:
            self._output_file = output_file
        else:
            self._output_file = None
        self._batch_size = batch_size
        self._print_to_console = print_to_console
        if has_dataset_reader:
            self._dataset_reader = predictor._dataset_reader
        else:
            self._dataset_reader = None
        num_layers = len(self._predictor._model.vae.encoder._linear_layers) + 1  # pylint: disable=protected-access
        initial_params = [1] + [-20] * (num_layers - 2) + [1]
        self.scalar_mix = ScalarMix(
                num_layers,
                do_layer_norm=False,
                initial_scalar_parameters=initial_params,
                trainable=False)

    def _predict_json(self, batch_data: List[JsonDict]) -> Iterator[str]:
        if len(batch_data) == 1:
            results = [self._predictor.predict_json(batch_data[0])]
        else:
            results = self._predictor.predict_batch_json(batch_data)
        for output in results:
            yield output

    def _predict_instances(self, batch_data: List[Instance]) -> Iterator[str]:
        if len(batch_data) == 1:
            results = [self._predictor.predict_instance(batch_data[0])]
        else:
            results = self._predictor.predict_batch_instance(batch_data)
        for output in results:
            yield output

    def _maybe_print_to_console_and_file(
        self, index: int, prediction: str, model_input: str = None
    ) -> None:
        if self._print_to_console:
            if model_input is not None:
                print(f"input {index}: ", model_input)
            print("prediction: ", prediction)
        if self._output_file is not None:
            self._output_file.write(prediction)

    def _get_json_data(self) -> Iterator[JsonDict]:
        if self._input_file == "-":
            for line in sys.stdin:
                if not line.isspace():
                    yield self._predictor.load_line(line)
        else:
            input_file = cached_path(self._input_file)
            with open(input_file, "r") as file_input:
                for line in file_input:
                    if not line.isspace():
                        yield self._predictor.load_line(line)

    def _get_instance_data(self) -> Iterator[Instance]:
        if self._input_file == "-":
            raise ConfigurationError("stdin is not an option when using a DatasetReader.")
        elif self._dataset_reader is None:
            raise ConfigurationError("To generate instances directly, pass a DatasetReader.")
        else:
            yield from self._dataset_reader.read(self._input_file)

    def run(self) -> None:
        has_reader = self._dataset_reader is not None
        index = 0
        if has_reader:
            for batch in lazy_groups_of(self._get_instance_data(), self._batch_size):
                for model_input_instance, result in zip(batch, self._predict_instances(batch)):
                    self._maybe_print_to_console_and_file(index, result, str(model_input_instance))
                    index = index + 1
        else:
            ids_ = []
            vecs = []
            for batch_json in tqdm(lazy_groups_of(self._get_json_data(), self._batch_size)):
                for model_input_json, result in zip(batch_json, self._predict_json(batch_json)):
                    scalar_mix = (torch.Tensor(result['activation_encoder_layer_0']).unsqueeze(0)
                                  + -20 * torch.Tensor(result['activation_encoder_layer_1']).unsqueeze(0)
                                  + torch.Tensor(result['activation_theta']).unsqueeze(0))
                    vecs.append(scalar_mix)
                    ids_.append(torch.IntTensor([model_input_json['index']]).unsqueeze(0))
                    index = index + 1
        torch.save((torch.cat(ids_,0), torch.cat(vecs, 0)), self._output_file)

def _predict(args: argparse.Namespace) -> None:
    from allennlp.common.util import import_submodules
    import_submodules(args.include_package)
    predictor = _get_predictor(args)

    if args.silent and not args.output_file:
        print("--silent specified without --output-file.")
        print("Exiting early because no output will be created.")
        sys.exit(0)

    manager = _PredictManager(
        predictor,
        args.input_file,
        args.output_file,
        args.batch_size,
        not args.silent,
        args.use_dataset_reader,
    )
    manager.run()

if __name__ == '__main__':
    description = """Run the specified model against a JSON-lines input file."""
    subparser = argparse.ArgumentParser(
            'predictor', description=description 
    )

    subparser.add_argument(
        "archive_file", type=str, help="the archived model to make predictions with"
    )
    subparser.add_argument("input_file", type=str, help="path to or url of the input file")

    subparser.add_argument("--output-file", type=str, help="path to output file")
    subparser.add_argument(
        "--weights-file", type=str, help="a path that overrides which weights file to use"
    )

    batch_size = subparser.add_mutually_exclusive_group(required=False)
    batch_size.add_argument(
        "--batch-size", type=int, default=1, help="The batch size to use for processing"
    )

    subparser.add_argument(
        "--silent", action="store_true", help="do not print output to stdout"
    )

    cuda_device = subparser.add_mutually_exclusive_group(required=False)
    cuda_device.add_argument(
        "--cuda-device", type=int, default=-1, help="id of GPU to use (if any)"
    )

    subparser.add_argument(
        "--use-dataset-reader",
        action="store_true",
        help="Whether to use the dataset reader of the original model to load Instances. "
        "The validation dataset reader will be used if it exists, otherwise it will "
        "fall back to the train dataset reader. This behavior can be overridden "
        "with the --dataset-reader-choice flag.",
    )

    subparser.add_argument(
        "--dataset-reader-choice",
        type=str,
        choices=["train", "validation"],
        default="validation",
        help="Indicates which model dataset reader to use if the --use-dataset-reader "
        "flag is set.",
    )

    subparser.add_argument(
        "-o",
        "--overrides",
        type=str,
        default="",
        help="a JSON structure used to override the experiment configuration",
    )

    subparser.add_argument(
        "--predictor", type=str, help="optionally specify a specific predictor to use"
    )
    subparser.add_argument(
                    "--include-package", type=str, help="optionally specify a specific predictor to use"
                        )

    subparser.set_defaults(func=_predict)
    _predict(subparser.parse_args())
