import unittest
import torch
import os
from parapred.model import Parapred, clean_output
from parapred.preprocessing import encode_parapred, encode_batch
from parapred.cnn import generate_mask

FPATH = os.path.dirname(os.path.abspath(__file__))
WEIGHTS_PATH = os.path.join(FPATH, "../weights/parapred_pytorch.h5")


class ParapredTest(unittest.TestCase):
    def setUp(self):
        self.model = Parapred()
        self.model.load_state_dict(torch.load(WEIGHTS_PATH))
        _ = self.model.eval()

        self.sequence = "YCQRYNRAPYTFG"
        self.max_length = 40
        self.num_features = 28

    def test_probailities(self):
        """
        Integration test for probability prediction
        """
        encoding, lengths = encode_batch([self.sequence], self.max_length)
        m = generate_mask(encoding, lengths)

        with torch.no_grad():
            pr = self.model(encoding, m, lengths)

        v = clean_output(pr[0], lengths[0].item())

        self.assertTrue(
            torch.allclose(
                torch.Tensor([0.03122, 0.00289, 0.01522, 0.03233, 0.91215, 0.82423, 0.87741,
                              0.77854, 0.24664, 0.76494, 0.00932, 0.00534, 0.00251]), v,
                rtol=1e-2
            )
        )

    def test_encoding(self):
        """
        Unit test the encoding function
        """
        # Deliberately not testing padding
        encoded_representation = encode_parapred(self.sequence, len(self.sequence))

        self.assertTrue(
            torch.allclose(
                torch.Tensor([[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 2.9400, 0.3000, 6.4700,
                               0.9600, 5.6600, 0.2500, 0.4100],
                              [1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.7700, 0.1300, 2.4300,
                               1.5400, 6.3500, 0.1700, 0.4100],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.5600, 0.1800, 3.9500,
                               -0.2200, 5.6500, 0.3600, 0.2500],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 2.3400, 0.2900, 6.1300,
                               -1.0100, 10.7400, 0.3600, 0.2500],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 2.9400, 0.3000, 6.4700,
                               0.9600, 5.6600, 0.2500, 0.4100],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.6000, 0.1300, 2.9500,
                               -0.6000, 6.5200, 0.2100, 0.2200],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 2.3400, 0.2900, 6.1300,
                               -1.0100, 10.7400, 0.3600, 0.2500],
                              [0.0000, 0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.2800, 0.0500, 1.0000,
                               0.3100, 6.1100, 0.4200, 0.2300],
                              [0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 2.6700, 0.0000, 2.7200,
                               0.7200, 6.8000, 0.1300, 0.3400],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 2.9400, 0.3000, 6.4700,
                               0.9600, 5.6600, 0.2500, 0.4100],
                              [0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 3.0300, 0.1100, 2.6000,
                               0.2600, 5.6000, 0.2100, 0.3600],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 2.9400, 0.2900, 5.8900,
                               1.7900, 5.6700, 0.3000, 0.3800],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 6.0700, 0.1300, 0.1500]]), encoded_representation
            )
        )

    def test_batch_prediction(self):
        """
        Integration testing for a batch of sequences
        """
        batch = ["SRWGGDGFYAMDYWG", "YCQRYNRAPYTFG"]
        encoding, lengths = encode_batch(batch, self.max_length)
        m = generate_mask(encoding, lengths)

        with torch.no_grad():
            pr = self.model(encoding, m, lengths)

        v1 = clean_output(pr[0], lengths[0].item())
        v2 = clean_output(pr[1], lengths[1].item())

        self.assertTrue(
            torch.allclose(
                torch.Tensor([0.04144, 0.34117, 0.97052, 0.67401, 0.9148, 0.93996, 0.81214,
                              0.78589, 0.94175, 0.04701, 0.06284, 0.18635, 0.08849, 0.00447,
                              0.00532]), v1,
                rtol=1e-2
            )
        )
        self.assertTrue(
            torch.allclose(
                torch.Tensor([0.03122, 0.00289, 0.01522, 0.03233, 0.91215, 0.82423, 0.87741,
                              0.77854, 0.24664, 0.76494, 0.00932, 0.00534, 0.00251]), v2,
                rtol=1e-2
            )
        )

import unittest
import torch
import os
from parapred.model import Parapred, clean_output
from parapred.preprocessing import encode_parapred, encode_batch
from parapred.cnn import generate_mask

FPATH = os.path.dirname(os.path.abspath(__file__))
WEIGHTS_PATH = os.path.join(FPATH, "../weights/parapred_pytorch.h5")


class ParapredTest(unittest.TestCase):
    def setUp(self):
        self.model = Parapred()
        self.model.load_state_dict(torch.load(WEIGHTS_PATH))
        _ = self.model.eval()

        self.sequence = "YCQRYNRAPYTFG"
        self.max_length = 40
        self.num_features = 28

    def test_probailities(self):
        """
        Integration test for probability prediction
        """
        encoding, lengths = encode_batch([self.sequence], self.max_length)
        m = generate_mask(encoding, lengths)

        with torch.no_grad():
            pr = self.model(encoding, m, lengths)

        v = clean_output(pr[0], lengths[0].item())

        self.assertTrue(
            torch.allclose(
                torch.Tensor([0.03122, 0.00289, 0.01522, 0.03233, 0.91215, 0.82423, 0.87741,
                              0.77854, 0.24664, 0.76494, 0.00932, 0.00534, 0.00251]), v,
                rtol=1e-2
            )
        )

    def test_encoding(self):
        """
        Unit test the encoding function
        """
        # Deliberately not testing padding
        encoded_representation = encode_parapred(self.sequence, len(self.sequence))

        self.assertTrue(
            torch.allclose(
                torch.Tensor([[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 2.9400, 0.3000, 6.4700,
                               0.9600, 5.6600, 0.2500, 0.4100],
                              [1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.7700, 0.1300, 2.4300,
                               1.5400, 6.3500, 0.1700, 0.4100],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.5600, 0.1800, 3.9500,
                               -0.2200, 5.6500, 0.3600, 0.2500],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 2.3400, 0.2900, 6.1300,
                               -1.0100, 10.7400, 0.3600, 0.2500],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 2.9400, 0.3000, 6.4700,
                               0.9600, 5.6600, 0.2500, 0.4100],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.6000, 0.1300, 2.9500,
                               -0.6000, 6.5200, 0.2100, 0.2200],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 2.3400, 0.2900, 6.1300,
                               -1.0100, 10.7400, 0.3600, 0.2500],
                              [0.0000, 0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.2800, 0.0500, 1.0000,
                               0.3100, 6.1100, 0.4200, 0.2300],
                              [0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 2.6700, 0.0000, 2.7200,
                               0.7200, 6.8000, 0.1300, 0.3400],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 2.9400, 0.3000, 6.4700,
                               0.9600, 5.6600, 0.2500, 0.4100],
                              [0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 3.0300, 0.1100, 2.6000,
                               0.2600, 5.6000, 0.2100, 0.3600],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 2.9400, 0.2900, 5.8900,
                               1.7900, 5.6700, 0.3000, 0.3800],
                              [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                               0.0000, 6.0700, 0.1300, 0.1500]]), encoded_representation
            )
        )

    def test_batch_prediction(self):
        """
        Integration testing for a batch of sequences
        """
        batch = ["SRWGGDGFYAMDYWG", "YCQRYNRAPYTFG"]
        encoding, lengths = encode_batch(batch, self.max_length)
        m = generate_mask(encoding, lengths)

        with torch.no_grad():
            pr = self.model(encoding, m, lengths)

        v1 = clean_output(pr[0], lengths[0].item())
        v2 = clean_output(pr[1], lengths[1].item())

        self.assertTrue(
            torch.allclose(
                torch.Tensor([0.04144, 0.34117, 0.97052, 0.67401, 0.9148, 0.93996, 0.81214,
                              0.78589, 0.94175, 0.04701, 0.06284, 0.18635, 0.08849, 0.00447,
                              0.00532]), v1,
                rtol=1e-2
            )
        )
        self.assertTrue(
            torch.allclose(
                torch.Tensor([0.03122, 0.00289, 0.01522, 0.03233, 0.91215, 0.82423, 0.87741,
                              0.77854, 0.24664, 0.76494, 0.00932, 0.00534, 0.00251]), v2,
                rtol=1e-2
            )
        )

Integration test for probability prediction Unit test the encoding function # Deliberately not testing padding Integration testing for a batch of sequences