# coding=utf-8
# Copyright 2022 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for align_transforms."""

import gin
import tensorflow as tf

from dedal import vocabulary
from dedal.data import align_transforms


class AlignTransformsTest(tf.test.TestCase):

  def setUp(self):
    super().setUp()
    gin.clear_config()
    tf.random.set_seed(0)
    self.vocab = vocabulary.alternative
    self.sampler = vocabulary.Sampler(vocab=self.vocab)

  def test_project_msa_rows(self):
    token = '-'
    seq1 = tf.convert_to_tensor(self.vocab.encode('XX--X-XX-X'), tf.int32)
    seq2 = tf.convert_to_tensor(self.vocab.encode('Y-Y-Y-YYYY'), tf.int32)
    mc = tf.convert_to_tensor([False] + 8 * [True] + [False], tf.bool)

    proj_msa_fn = align_transforms.ProjectMSARows(token=token, vocab=self.vocab)

    out1, out2 = proj_msa_fn.call(seq1, seq2)
    self.assertAllEqual(out1, self.vocab.encode('XX-XXX-X'))
    self.assertAllEqual(out2, self.vocab.encode('Y-YYYYYY'))

    out1, out2 = proj_msa_fn.call(seq1, seq2, mc)
    self.assertAllEqual(out1, self.vocab.encode('X-XXX-'))
    self.assertAllEqual(out2, self.vocab.encode('-YYYYY'))

  def test_pid1(self):
    token = '-'
    pid_fn = align_transforms.PID(definition=1, token=token, vocab=self.vocab)

    seq1 = tf.convert_to_tensor(self.vocab.encode('XYXXYXXY'), tf.int32)
    seq2 = tf.convert_to_tensor(self.vocab.encode('XYXXYXXY'), tf.int32)
    pid = pid_fn.call(seq1, seq2)
    self.assertEqual(pid, 1.0)

    seq1 = tf.convert_to_tensor(self.vocab.encode('XYXXYXXY'), tf.int32)
    seq2 = tf.convert_to_tensor(self.vocab.encode('YXXXYXYX'), tf.int32)
    pid = pid_fn.call(seq1, seq2)
    self.assertEqual(pid, 0.5)

    seq1 = tf.convert_to_tensor(self.vocab.encode('XYXYYXXY'), tf.int32)
    seq2 = tf.convert_to_tensor(self.vocab.encode('YXYXXYYX'), tf.int32)
    pid = pid_fn.call(seq1, seq2)
    self.assertEqual(pid, 0.0)

    seq1 = tf.convert_to_tensor(self.vocab.encode('XYXXYXXY'), tf.int32)
    seq2 = tf.convert_to_tensor(self.vocab.encode('---X-XXY'), tf.int32)
    pid = pid_fn.call(seq1, seq2)
    self.assertEqual(pid, 0.8)

    seq1 = tf.convert_to_tensor(self.vocab.encode('XYXXYXXY'), tf.int32)
    seq2 = tf.convert_to_tensor(self.vocab.encode('---X-YYX'), tf.int32)
    pid = pid_fn.call(seq1, seq2)
    self.assertEqual(pid, 0.2)

    seq1 = tf.convert_to_tensor(self.vocab.encode('XYXXYXXY----'), tf.int32)
    seq2 = tf.convert_to_tensor(self.vocab.encode('---X-YYXYXXX'), tf.int32)
    pid = pid_fn.call(seq1, seq2)
    self.assertEqual(pid, 0.2)

  def test_pid3(self):
    token = '-'
    pid_fn = align_transforms.PID(definition=3, token=token, vocab=self.vocab)

    seq1 = tf.convert_to_tensor(self.vocab.encode('XYXXYXXY'), tf.int32)
    seq2 = tf.convert_to_tensor(self.vocab.encode('XYXXYXXY'), tf.int32)
    pid = pid_fn.call(seq1, seq2)
    self.assertEqual(pid, 1.0)

    seq1 = tf.convert_to_tensor(self.vocab.encode('XYXXYXXY'), tf.int32)
    seq2 = tf.convert_to_tensor(self.vocab.encode('YXXXYXYX'), tf.int32)
    pid = pid_fn.call(seq1, seq2)
    self.assertEqual(pid, 0.5)

    seq1 = tf.convert_to_tensor(self.vocab.encode('XYXYYXXY'), tf.int32)
    seq2 = tf.convert_to_tensor(self.vocab.encode('YXYXXYYX'), tf.int32)