# Copyright 2017 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.
"""Utilities to help with testing DeepVariant code."""

import os



from third_party.nucleus.testing import test_utils as nucleus_test_utils

GENOMICS_DIR = 'learning/genomics'


def deeptrio_testdata(filename):
  """Gets the path to filename in genomics/deepvariant/testdata.

  These paths are only known at runtime, after flag parsing
  has occurred.

  Args:
    filename: The name of a testdata file in the core genomics testdata
      directory. For example, if you have a test file in
      "learning/genomics/deepvariant/testdata/foo.txt", filename should be
      "foo.txt" to get a path to it.

  Returns:
    The absolute path to a testdata file.
  """
  return nucleus_test_utils.genomics_testdata(
      os.path.join('deeptrio/testdata', filename), GENOMICS_DIR
  )


CHR20_FASTA = None
HG001_CHR20_BAM = None
NA12891_CHR20_BAM = None
NA12892_CHR20_BAM = None
GOLDEN_TRAINING_EXAMPLES = None
GOLDEN_CALLING_CANDIDATES = None
GOLDEN_CANDIDATE_POSITIONS = None
GOLDEN_CALLING_EXAMPLES = None
CONFIDENT_REGIONS_BED = None
TRUTH_VARIANTS_VCF = None
TRUTH_VARIANTS_VCF_WITH_TYPES = None
GOLDEN_POSTPROCESS_INPUT = None
GOLDEN_POSTPROCESS_OUTPUT = None
GOLDEN_POSTPROCESS_OUTPUT_COMPRESSED = None
GOLDEN_POSTPROCESS_GVCF_INPUT = None
GOLDEN_POSTPROCESS_GVCF_OUTPUT = None
GOLDEN_POSTPROCESS_GVCF_OUTPUT_COMPRESSED = None
GOLDEN_MAKE_EXAMPLES_RUN_INFO = None
WS_ALLELE_COUNT_LINEAR_MODEL = None
WS_ALLELE_COUNT_LINEAR_MODEL_PCKL = None
WS_VARIANT_READS_THRESHOLD_MODEL = None
# Test data for ONT
GRCH38_CHR0_FASTA = None
ONT_HG002_BAM = None
ONT_HG003_BAM = None
ONT_HG004_BAM = None
HG002_HIGH_CONFIDENCE_VCF = None
HG002_HIGH_CONFIDENCE_BED = None
HG002_DENOVO_BED = None
GOLDEN_ONT_MAKE_EXAMPLES_OUTPUT = None
GOLDEN_ONT_DENOVO_MAKE_EXAMPLES_OUTPUT = None

ONT_N_GOLDEN_TRAINING_EXAMPLES = 167
N_GOLDEN_TRAINING_EXAMPLES = 50
N_GOLDEN_CALLING_EXAMPLES = 103

CUSTOMIZED_CLASSES_GOLDEN_TRAINING_EXAMPLES = None
ALT_ALIGNED_PILEUP_GOLDEN_TRAINING_EXAMPLES = None
GOLDEN_VCF_CANDIDATE_IMPORTER_TRAINING_EXAMPLES = None
GOLDEN_VCF_CANDIDATE_IMPORTER_CALLING_EXAMPLES_CHILD = None


def init():
  """Initialize global variables from flag values."""
  global CHR20_FASTA
  global HG001_CHR20_BAM
  global NA12891_CHR20_BAM
  global NA12892_CHR20_BAM
  global GOLDEN_TRAINING_EXAMPLES
  global GOLDEN_CANDIDATE_POSITIONS
  global GOLDEN_CALLING_CANDIDATES
  global GOLDEN_CALLING_EXAMPLES
  global CONFIDENT_REGIONS_BED
  global TRUTH_VARIANTS_VCF
  global TRUTH_VARIANTS_VCF_WITH_TYPES
  global GOLDEN_POSTPROCESS_INPUT
  global GOLDEN_POSTPROCESS_OUTPUT
  global GOLDEN_POSTPROCESS_OUTPUT_COMPRESSED
  global GOLDEN_POSTPROCESS_GVCF_INPUT
  global GOLDEN_POSTPROCESS_GVCF_OUTPUT
  global GOLDEN_POSTPROCESS_GVCF_OUTPUT_COMPRESSED
  global GOLDEN_MAKE_EXAMPLES_RUN_INFO
  global WS_ALLELE_COUNT_LINEAR_MODEL
  global WS_ALLELE_COUNT_LINEAR_MODEL_PCKL
  global WS_VARIANT_READS_THRESHOLD_MODEL
  global GOLDEN_VCF_CANDIDATE_IMPORTER_TRAINING_EXAMPLES
  global GOLDEN_VCF_CANDIDATE_IMPORTER_CALLING_EXAMPLES_CHILD

  global GRCH38_CHR0_FASTA
  global ONT_HG002_BAM
  global ONT_HG003_BAM
  global ONT_HG004_BAM
  global HG002_HIGH_CONFIDENCE_VCF
  global HG002_HIGH_CONFIDENCE_BED
  global HG002_DENOVO_BED
  global GOLDEN_ONT_MAKE_EXAMPLES_OUTPUT
  global GOLDEN_ONT_DENOVO_MAKE_EXAMPLES_OUTPUT

  CHR20_FASTA = deeptrio_testdata('input/hs37d5.chr20.fa.gz')
  HG001_CHR20_BAM = deeptrio_testdata('input/HG001.chr20.10_10p1mb_sorted.bam')
  NA12891_CHR20_BAM = deeptrio_testdata(
      'input/NA12891.chr20.10_10p1mb_sorted.bam'
  )
  NA12892_CHR20_BAM = deeptrio_testdata(
      'input/NA12892.chr20.10_10p1mb_sorted.bam'
  )

  GOLDEN_TRAINING_EXAMPLES = deeptrio_testdata(
      'golden.training_examples.tfrecord.gz'
  )
  GOLDEN_CANDIDATE_POSITIONS = deeptrio_testdata(
      'golden_child.candidate_positions'
  )
  GOLDEN_CALLING_CANDIDATES = deeptrio_testdata(
      'golden_child.calling_examples.tfrecord.gz'
  )
  GOLDEN_CALLING_EXAMPLES = deeptrio_testdata(
      'golden_child.calling_examples.tfrecord.gz'
  )
  CONFIDENT_REGIONS_BED = deeptrio_testdata(
      'input/test_giab.b37_chr20_100kbp_at_10mb.bed'
  )
  TRUTH_VARIANTS_VCF = deeptrio_testdata(
      'input/HG001_chr20_GRCh37_GIAB_highconf_CG-IllFB-IllGATKHC-Ion-10X-SOLID_CHROM1-X_v.3.3.2_highconf_PGandRTGphasetransfer.vcf.gz'
  )
  TRUTH_VARIANTS_VCF_WITH_TYPES = deeptrio_testdata(
      'input/with_types.test_nist.b37_chr20_4kbp_at_10mb.vcf.gz'
  )
  GOLDEN_POSTPROCESS_INPUT = deeptrio_testdata(
      'golden.postprocess_single_site_input.tfrecord.gz'
  )
  GOLDEN_POSTPROCESS_OUTPUT = deeptrio_testdata(
      'golden.postprocess_single_site_output.vcf'
  )
  GOLDEN_POSTPROCESS_OUTPUT_COMPRESSED = deeptrio_testdata(
      'golden.postprocess_single_site_output.vcf.gz'
  )
  GOLDEN_POSTPROCESS_GVCF_INPUT = deeptrio_testdata(
      'golden_child.postprocess_gvcf_input.tfrecord.gz'
  )
  GOLDEN_POSTPROCESS_GVCF_OUTPUT = deeptrio_testdata(
      'golden.postprocess_gvcf_output.g.vcf'
  )
  GOLDEN_MAKE_EXAMPLES_RUN_INFO = deeptrio_testdata(
      'golden.training_examples.tfrecord.gz.run_info.pbtxt'
  )
  WS_ALLELE_COUNT_LINEAR_MODEL = deeptrio_testdata(
      'window_selector_allele_count_linear.pbtxt'
  )
  WS_ALLELE_COUNT_LINEAR_MODEL_PCKL = deeptrio_testdata(
      'window_selector_allele_count_linear.pckl'
  )
  WS_VARIANT_READS_THRESHOLD_MODEL = deeptrio_testdata(
      'window_selector_variant_read_threshold.pbtxt'
  )

  # For oxford nanopore
  GRCH38_CHR0_FASTA = deeptrio_testdata(
      'input/grch38.chr20_5050000_5075000.masked.fa.gz'
  )
  ONT_HG002_BAM = deeptrio_testdata('input/HG002_R10_chr20_5050000_5075000.bam')
  ONT_HG003_BAM = deeptrio_testdata('input/HG003_R10_chr20_5050000_5075000.bam')
  ONT_HG004_BAM = deeptrio_testdata('input/HG004_R10_chr20_5050000_5075000.bam')
  HG002_HIGH_CONFIDENCE_VCF = deeptrio_testdata(
      'input/HG002_GRCh38_1_22_v4.2.1_benchmark.chr20.vcf.gz'
  )
  HG002_HIGH_CONFIDENCE_BED = deeptrio_testdata(
      'input/HG002_GRCh38_1_22_v4.2.1_benchmark.chr20.bed'
  )
  HG002_DENOVO_BED = deeptrio_testdata(
      'input/HG002_GRCh38_1_22_v4.2.1_benchmark.chr20.denovo_regions.bed'
  )
  GOLDEN_ONT_MAKE_EXAMPLES_OUTPUT = deeptrio_testdata(
      'HG002_ONT_deeptrio.examples.tfrecord.gz'
  )
  GOLDEN_ONT_DENOVO_MAKE_EXAMPLES_OUTPUT = deeptrio_testdata(
      'HG002_ONT_deeptrio.denovo.examples.tfrecord.gz'
  )

  # For CustomizedClassesVariantLabeler.
  global CUSTOMIZED_CLASSES_GOLDEN_TRAINING_EXAMPLES
  CUSTOMIZED_CLASSES_GOLDEN_TRAINING_EXAMPLES = deeptrio_testdata(
      'customized_classes.golden.training_examples.tfrecord.gz'
  )

  # For alt-aligned pileups
  global ALT_ALIGNED_PILEUP_GOLDEN_TRAINING_EXAMPLES
  ALT_ALIGNED_PILEUP_GOLDEN_TRAINING_EXAMPLES = deeptrio_testdata(
      'alt_aligned_pileup.golden.training_examples.tfrecord.gz'
  )

  GOLDEN_VCF_CANDIDATE_IMPORTER_TRAINING_EXAMPLES = deeptrio_testdata(
      'golden.vcf_candidate_importer.training_examples.tfrecord.gz'
  )
  GOLDEN_VCF_CANDIDATE_IMPORTER_CALLING_EXAMPLES_CHILD = deeptrio_testdata(
      'golden_child.vcf_candidate_importer.calling_examples.tfrecord.gz'
  )