# Copyright 2018 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.
"""Tests for third_party.nucleus.util.io."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import sys
if 'google' in sys.modules and 'google.protobuf' not in sys.modules:
  del sys.modules['google']


import itertools

from absl.testing import absltest
from absl.testing import parameterized

import six

from tensorflow.python.platform import gfile
from third_party.nucleus.io import sam
from third_party.nucleus.io import tfrecord
from third_party.nucleus.protos import reads_pb2
from third_party.nucleus.protos import reference_pb2
from third_party.nucleus.testing import test_utils
from third_party.nucleus.util import ranges


class SamReaderTests(parameterized.TestCase):
  """Test the iteration functionality provided by io.SamReader."""

  def test_sam_iterate(self):
    reader = sam.SamReader(test_utils.genomics_core_testdata('test.sam'))
    with reader:
      self.assertEqual(test_utils.iterable_len(reader.iterate()), 6)

  def test_bam_iterate(self):
    reader = sam.SamReader(test_utils.genomics_core_testdata('test.bam'))
    with reader:
      self.assertEqual(test_utils.iterable_len(reader.iterate()), 106)

  def test_bam_iterate_partially(self):
    """Verify that iteration provides results incrementally, not all at once."""
    reader = sam.SamReader(test_utils.genomics_core_testdata('test.bam'))
    with reader:
      iterable = reader.iterate()
      # We expect 106 records in total.
      for _ in range(10):
        results = list(itertools.islice(iterable, 10))
        self.assertEqual(len(results), 10)
      results = list(itertools.islice(iterable, 10))
      self.assertEqual(len(results), 6)

  def test_sam_query(self):
    reader = sam.SamReader(test_utils.genomics_core_testdata('test.bam'))
    expected = [(ranges.parse_literal('chr20:10,000,000-10,000,100'), 106),
                (ranges.parse_literal('chr20:10,000,000-10,000,000'), 45)]
    with reader:
      for interval, n_expected in expected:
        with reader.query(interval) as iterable:
          self.assertEqual(test_utils.iterable_len(iterable), n_expected)

  def test_sam_query_alternate_index_name(self):
    reader = sam.SamReader(
        test_utils.genomics_core_testdata('test_alternate_index.bam'))
    expected = [(ranges.parse_literal('chr20:10,000,000-10,000,100'), 106),
                (ranges.parse_literal('chr20:10,000,000-10,000,000'), 45)]
    with reader:
      for interval, n_expected in expected:
        with reader.query(interval) as iterable:
          self.assertEqual(test_utils.iterable_len(iterable), n_expected)

  @parameterized.parameters(('\t'.join(x[0] for x in items), {
      k: v for t in items for k, v in t[1].items()
  }) for r in [1, 2] for items in itertools.permutations(
      [
          ('X1:i:0', {
              'X1': 0
          }),
          ('X2:i:127', {
              'X2': 127
          }),
          ('X3:i:255', {
              'X3': 255
          }),
          ('X4:A:a', {
              'X4': 'a'
          }),
          ('X5:f:1.234', {
              'X5': 1.234
          }),
          ('X6:Z:string', {
              'X6': 'string'
          }),
          ('X7:Z:with spaces', {
              'X7': 'with spaces'
          }),
          ('ZJ:Z:', {
              'ZJ': ''
          }),  # Empty string.
          # We skip H hex byte-array tags as they appear deprecated.
          ('X8:H:1AE301', {}),
          ('X9:B:i,1', {
              'X9': [1]
          }),
          ('XA:B:i,1,2', {
              'XA': [1, 2]
          }),
          ('XB:B:i,1,2,3', {
              'XB': [1, 2, 3]
          }),
          ('XC:B:i,1,2,3,4,5,6,7,8,9,10', {
              'XC': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
          }),
          ('XD:B:I,1,2,3', {
              'XD': [1, 2, 3]
          }),
          ('XE:B:c,1,2,3', {
              'XE': [1, 2, 3]
          }),
          ('XF:B:C,1,2,3', {
              'XF': [1, 2, 3]
          }),
          ('XG:B:f,0.12,0.34', {
              'XG': [0.12, 0.34]
          }),
          ('XH:B:s,1,2,3', {
              'XH': [1, 2, 3]
          }),
          ('XI:B:S,1,2,3', {
              'XI': [1, 2, 3]
          }),
      ],
      r=r))
  def test_parsing_aux_tags(self, tag_string, expected_info):
    # Minimal header line to create a valid SAM file.
    reads = self._parse_read_with_aux_tags(tag_string)
    self.assertLen(reads, 1)
    self.assertInfoMapEqual(reads[0].info, expected_info)

  @parameterized.parameters(
      '\t'.join(tags) for r in [1, 2, 3] for tags in itertools.permutations(
          [
              'X2:i:x',  # Integer with character value.
              'X3:f:string',  # A string instead of the expected float.
              'X4:A:',  # Supposed to be single char, but we none here.
              'X5:A:ab',  # Supposed to be single char, but we have two here.
              'X6:B:i',  # Empty byte array.
              'X7:B:i,1.23',  # Integer byte array with a float.
              'X8:B:i,string',  # Integer byte array with a string.
              'X8:B:f,string',  # Float byte array with a string.
              'X8:B:z,1,2,3',  # z is not a valid subtype.
          ],
          r=r))
  def test_survives_malformed_lines(self, tag_string):
    try:
      reads = self._parse_read_with_aux_tags(tag_string)
      # If we didn't detect the error, make sure we actually still parsed the
      # read itself.
      self.assertLen(reads, 1)
      self.assertEqual(reads[0].fragment_name, 'read_name')
      self.assertEqual(reads[0].aligned_sequence, 'CCC')
      self.assertEqual(reads[0].alignment.position.reference_name, 'chr1')
      self.assertEqual(reads[0].alignment.position.position, 0)
 
   except ValueError as e:
      if 'Failed to parse SAM record' not in str(e):
        self.fail('Parsing failed but unexpected exception was see
n: ' + str(e))

  def _parse_read_with_aux_tags(self, tag_string):
    # Minimal header line to create a valid SAM file.
    header_lines = '@HD\tVN:1.3\tSO:coordinate\n@SQ\tSN:chr1\tLN:248956422\n'
    # A single stock read we'll add our AUX fields to.
    read = 'read_name\t0\tchr1\t1\t0\t3M\t*\t0\t0\tCCC\tAAA\t' + tag_string
    path = test_utils.test_tmpfile('aux_tags.bam')
    with gfile.Open(path, 'w') as fout:
      fout.write(header_lines)
      fout.write(read + '\n')
    with sam.SamReader(path, parse_aux_fields=True) as reader:
      return list(reader.iterate())

  def assertInfoMapEqual(self, info_map, expected_info):
    self.assertCountEqual(
        info_map.keys(), expected_info.keys(),
        'info has {} keys but we expected {}'.format(info_map.keys(),
            