# Copyright 2021 Google LLC.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from this
#    software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
"""A prototype to create tumor-normal images (tf.Example protos)."""

import logging
import os

from absl import app
from absl import flags

from deepvariant import dv_constants
from deepvariant import logging_level
from deepvariant import make_examples_core
from deepvariant import make_examples_options
from deepvariant.protos import deepvariant_pb2
from third_party.nucleus.io.python import hts_verbose
from third_party.nucleus.util import errors
from third_party.nucleus.util import proto_utils

# Sentinel command line flag value indicating no downsampling should occur.
NO_DOWNSAMPLING = 0.0

# 1 is the tumor, 0 is the normal match.
# Tumor sample is the "main" sample because the goal here is somatic calling.
NORMAL_SAMPLE_INDEX = (
    0  # If the normal sample is present, it will appear first.
)

FLAGS = flags.FLAGS

# Adopt more general flags from make_examples_options.
flags.adopt_module_key_flags(make_examples_options)

# Flags related to samples in DeepSomatic:
_READS_TUMOR = flags.DEFINE_string(
    'reads_tumor',
    None,
    (
        'Required. Reads from the tumor sample. '
        'Aligned, sorted, indexed BAM file. '
        'Should be aligned to a reference genome compatible with --ref. '
        'Can provide multiple BAMs (comma-separated).'
    ),
)
_READS_NORMAL = flags.DEFINE_string(
    'reads_normal',
    None,
    (
        'Required. Reads from the normal matched sample. '
        'Aligned, sorted, indexed BAM file. '
        'Should be aligned to a reference genome compatible with --ref. '
        'Can provide multiple BAMs (comma-separated).'
    ),
)
_SAMPLE_NAME_TUMOR = flags.DEFINE_string(
    'sample_name_tumor',
    '',
    (
        'Sample name for tumor to use for our sample_name in the output'
        ' Variant/DeepVariantCall protos. If not specified, will be inferred'
        ' from the header information from --reads_tumor.'
    ),
)
_SAMPLE_NAME_NORMAL = flags.DEFINE_string(
    'sample_name_normal',
    '',
    (
        'Sample name for normal match to use for our sample_name in the output'
        ' Variant/DeepVariantCall protos. If not specified, will be inferred'
        ' from the header information from --reads_normal.'
    ),
)
_DOWNSAMPLE_FRACTION_TUMOR = flags.DEFINE_float(
    'downsample_fraction_tumor',
    NO_DOWNSAMPLING,
    'If not '
    + str(NO_DOWNSAMPLING)
    + ' must be a value between 0.0 and 1.0. '
    'Reads will be kept (randomly) with a probability of downsample_fraction '
    'from the input tumor sample BAM. This argument makes it easy to create '
    'examples as though the input BAM had less coverage.',
)
_DOWNSAMPLE_FRACTION_NORMAL = flags.DEFINE_float(
    'downsample_fraction_normal',
    NO_DOWNSAMPLING,
    'If not '
    + str(NO_DOWNSAMPLING)
    + ' must be a value between 0.0 and 1.0. '
    'Reads will be kept (randomly) with a probability of downsample_fraction '
    'from the input normal matched BAMs. This argument makes it easy to create '
    'examples as though the input BAMs had less coverage.',
)
_PILEUP_IMAGE_HEIGHT_TUMOR = flags.DEFINE_integer(
    'pileup_image_height_tumor',
    dv_constants.PILEUP_DEFAULT_HEIGHT,
    (
        'Height for the part of the pileup image showing reads from the tumor. '
        f'Uses {dv_constants.PILEUP_DEFAULT_HEIGHT} by default.'
    ),
)
_PILEUP_IMAGE_HEIGHT_NORMAL = flags.DEFINE_integer(
    'pileup_image_height_normal',
    dv_constants.PILEUP_DEFAULT_HEIGHT,
    (
        'Height for the part of the pileup image showing reads from the matched'
        f' normal. Uses {dv_constants.PILEUP_DEFAULT_HEIGHT} by default.'
    ),
)
_CANDIDATE_POSITIONS = flags.DEFINE_string(
    'candidate_positions',
    '',
    'Path to the binary file containing candidate positions.',
)

# Change any flag defaults that differ for DeepSomatic.
# I'm setting this to float('inf') because we don't want to include any
# candidates from the non-target (i.e., normal) sample.
FLAGS.set_default('vsc_min_fraction_multiplier', float('inf'))


def tumor_normal_samples_from_flags(flags_obj):
  """Collects sample-related options into a list of samples."""
  samples_in_order = []

  def setup_sample(role):
    pileup_height = (
        flags_obj.pileup_image_height_tumor
        if role == 'tumor'
        else flags_obj.pileup_image_height_normal
    )
    read_filenames = (
        flags_obj.reads_tumor if role == 'tumor' else flags_obj.reads_normal
    )
    sample_name_flag = (
        flags_obj.sample_name_tumor
        if role == 'tumor'
        else flags_obj.sample_name_normal
    )
    reads_filenames_split = read_filenames.split(',')
    sample_name = make_examples_core.assign_sample_name(
        sample_name_flag=sample_name_flag,
        reads_filenames=read_filenames,
    )
    skip_output_generation = True if role == 'normal' else False

    sample_options = deepvariant_pb2.SampleOptions(
        role=role,
        name=sample_name,
        variant_caller_options=make_examples_core.make_vc_options(
            sample_name=sample_name, flags_obj=flags_obj
        ),
        skip_output_generation=skip_output_generation,
        pileup_height=pileup_height,
        reads_filenames=reads_filenames_split,
    )
    if role == 'tumor':
      if flags_obj.reads_normal and flags_obj.reads_tumor:
        sample_options.order.extend([0, 1])
      else:
        sample_options.order.extend([0])

    downsample_fraction = (
        flags_obj.downsample_fraction_tumor
        if role == 'tumor'
        else flags_obj.downsample_fraction_normal
    )
    if downsample_fraction != NO_DOWNSAMPLING:
      sample_options.downsample_fraction = downsample_fraction
    samples_in_order.append(sample_options)

  if flags_obj.reads_normal:
    setup_sample('normal')

  setup_sample('tumor')

  return samples_in_order, 'tumor'


def default_options(main_sample_index, add_flags=True, flags_obj=None):
  """Creates a MakeExamplesOptions proto populated with reasonable defaults.

  Args:
    main_sample_index: int. Indicates the position of the tumor sample.
    add_flags: bool. defaults to True. If True, we will push the value of
      certain FLAGS into our options. If False, those option fields are left
      uninitialized.
    flags_obj: object.  If not None, use as the source of flags, else use global
      FLAGS.

  Returns:
    deepvariant_pb2.MakeExamplesOptions protobuf.

  Raises:
    ValueError: If we observe invalid flag values.
  """
  if not flags_obj:
    flags_obj = FLAGS

  samples_in_order, sample_role_to_train = tumor_normal_samples_from_flags(
      flags_obj=flags_obj
  )
  samples_in_order[main_sample_index].candidate_positions = (
      flags_obj.candidate_positions
  )

  options = make_examples_options.shared_flags_to_options(
      add_flags=add_flags,
      flags_obj=flags_obj,
      samples_in_order=samples_in_order,
      sample_role_to_train=sample_role_to_train,
      main_sample_index=main_sample_index,
  )

  if _READS_NORMAL.value:
    options.bam_fname = f'{os.path.basename(flags_obj.reads_tumor)}|{os.path.basename(flags_obj.reads_normal)}'
  else:
    options.bam_fname = os.path.basename(flags_obj.reads_tumor)

  return options


def check_options_are_valid(options, main_sample_index):
  """Checks that all the options chosen make sense together."""

  # Check for general flags (shared for DeepVariant and DeepTrio).
  make_examples_options.check_options_are_valid(
      options, main_sample_index=main_sample_index
  )

  tumor = options.sample_options[main_sample_index]
  if _READS_NORMAL.value:
    normal = options.sample_options[NORMAL_SAMPLE_INDEX]

    if (
        tumor.variant_caller_options.sample_name
        == normal.variant_caller_options.sample_name
    ):
      errors.log_and_raise(
          (
              'Sample names of tumor and normal samples cannot be the same. Use'
              ' --sample_name_tumor and --sample_name_normal with different'
              ' names '
          ),
          errors.CommandLineError,
      )

  if options.sample_options[main_sample_index].candidate_positions:
    if options.max_reads_per_partition:
      logging.warning(
          'Since candidate_positions is set, we use '
          'max_reads_for_dynamic_bases_per_region instead of '
          'max_reads_per_partition. This is due to the dynamic nature of the '
          'partition size when candidate_positions is enabled, making a fixed '
          'max_reads_per_partition unsuitable.'
      )
      options.max_reads_for_dynamic_bases_per_region = (
          options.max_reads_per_partition
      )
      options.max_reads_per_partition = 0


def main(argv=()):
  with errors.clean_commandline_error_exit():
    if len(argv) > 1:
      errors.log_and_raise(
          'Command line parsing failure: make_examples does not accept '
          'positional arguments but some are present on the command line: '
          '"{}".'.format(str(argv)),
          errors.CommandLineError,
      )
    del argv  # Unused.

    proto_utils.uses_fast_cpp_protos_or_die()

    logging_level.set_from_flag()
    hts_verbose.set(hts_verbose.htsLogLevel.HTS_LOG_WARNING)

    # Set up options; may do I/O.
    is_tumor_only = _READS_TUMOR.value and not _READS_NORMAL.value
    main_sample_index = 0 if is_tumor_only else 1
    options = default_options(main_sample_index, flags_obj=FLAGS)
    check_options_are_valid(options, main_sample_index)

    # Run!
    make_examples_core.make_examples_runner(options)


if __name__ == '__main__':
  flags.mark_flags_as_required([
      'examples',
      'mode',
      'reads_tumor',
      'ref',
  ])
  app.run(main)