# -*- coding: utf-8 -*-

"""This module contains tests for the SPIA exporter."""

import unittest

from pandas import DataFrame

from pybel.dsl import activity, composite_abundance, pmod, protein, rna
from pybel.examples.sialic_acid_example import (
    cd33, citation, evidence_1, shp1, shp2, sialic_acid_cd33_complex, sialic_acid_graph, trem2,
)
from pybel.io.spia import build_spia_matrices, get_matrix_index, to_spia_dfs, update_spia_matrices


class TestSpia(unittest.TestCase):
    """Test SPIA Exporter."""

    def setUp(self):
        self.sialic_acid_graph = sialic_acid_graph.copy()

    def test_build_matrix(self):
        """Test build empty matrix."""

        node_names = get_matrix_index(self.sialic_acid_graph)

        matrix_dict = build_spia_matrices(node_names)

        nodes = {'PTPN11', 'TREM2', 'PTPN6', 'TYROBP', 'CD33', 'SYK'}

        self.assertEqual(set(matrix_dict["activation"].columns), nodes)
        self.assertEqual(set(matrix_dict["repression"].index), nodes)

    def test_update_matrix_inhibition_ubiquination(self):
        """Test updating the matrix with an inhibition ubiquitination."""
        sub = protein(namespace='HGNC', name='A', identifier='1')
        obj = protein(namespace='HGNC', name='B', identifier='2', variants=[pmod('Ub')])

        index = {'A', 'B'}

        test_dict = {}

        test_matrix = DataFrame(0, index=index, columns=index)

        # Initialize matrix correctly
        self.assertEqual(test_matrix.values.all(), 0)

        test_dict["inhibition_ubiquination"] = test_matrix

        update_spia_matrices(test_dict, sub, obj, {'relation': 'decreases'})

        self.assertEqual(test_dict["inhibition_ubiquination"]['A']['B'], 1)
        self.assertEqual(test_dict["inhibition_ubiquination"]['A']['A'], 0)
        self.assertEqual(test_dict["inhibition_ubiquination"]['B']['A'], 0)
        self.assertEqual(test_dict["inhibition_ubiquination"]['B']['B'], 0)

    def test_update_matrix_activation_ubiquination(self):
        """Test updating the matrix with an activation ubiquitination."""
        sub = protein(namespace='HGNC', name='A', identifier='1')
        obj = protein(namespace='HGNC', name='B', identifier='2', variants=[pmod('Ub')])

        index = {'A', 'B'}

        test_dict = {}

        test_matrix = DataFrame(0, index=index, columns=index)

        test_dict["activation_ubiquination"] = test_matrix

        update_spia_matrices(test_dict, sub, obj, {'relation': 'increases'})

        self.assertEqual(test_dict["activation_ubiquination"]['A']['B'], 1)
        self.assertEqual(test_dict["activation_ubiquination"]['A']['A'], 0)
        self.assertEqual(test_dict["activation_ubiquination"]['B']['A'], 0)
        self.assertEqual(test_dict["activation_ubiquination"]['B']['B'], 0)

    def test_update_matrix_inhibition_phosphorylation(self):
        """Test updating the matrix with an inhibition phosphorylation."""
        sub = protein(namespace='HGNC', name='A', identifier='1')
        obj = protein(namespace='HGNC', name='B', identifier='2', variants=[pmod('Ph')])

        index = {'A', 'B'}

        test_dict = {}

        test_matrix = DataFrame(0, index=index, columns=index)

        test_dict["inhibition_phosphorylation"] = test_matrix

        update_spia_matrices(test_dict, sub, obj, {'relation': 'decreases'})

        self.assertEqual(test_dict["inhibition_phosphorylation"]['A']['B'], 1)
        self.assertEqual(test_dict["inhibition_phosphorylation"]['A']['A'], 0)
        self.assertEqual(test_dict["inhibition_phosphorylation"]['B']['A'], 0)
        self.assertEqual(test_dict["inhibition_phosphorylation"]['B']['B'], 0)

    def test_update_matrix_activation_phosphorylation(self):
        """Test updating the matrix with an activation phosphorylation."""

        sub = protein(namespace='HGNC', name='A', identifier='1')
        obj = protein(namespace='HGNC', name='B', identifier='2', variants=[pmod('Ph')])

        index = {'A', 'B'}

        test_dict = {}

        test_matrix = DataFrame(0, index=index, columns=index)

        test_dict["activation_phosphorylation"] = test_matrix

        update_spia_matrices(test_dict, sub, obj, {'relation': 'increases'})

        self.assertEqual(test_dict["activation_phosphorylation"]['A']['B'], 1)
        self.assertEqual(test_dict["activation_phosphorylation"]['A']['A'], 0)
        self.assertEqual(test_dict["activation_phosphorylation"]['B']['A'], 0)
        self.assertEqual(test_dict["activation_phosphorylation"]['B']['B'], 0)

    def test_update_matrix_expression(self):
        """Test updating the matrix with RNA expression."""
        sub = protein(namespace='HGNC', name='A', identifier='1')
        obj = rna(namespace='HGNC', name='B', identifier='2')

        index = {'A', 'B'}

        test_dict = {}

        test_matrix = DataFrame(0, index=index, columns=index)

        test_dict["expression"] = test_matrix

        update_spia_matrices(test_dict, sub, obj, {'relation': 'increases'})

        self.assertEqual(test_dict["expression"]['A']['B'], 1)
        self.assertEqual(test_dict["expression"]['A']['A'], 0)
        self.assertEqual(test_dict["expression"]['B']['A'], 0)
        self.assertEqual(test_dict["expression"]['B']['B'], 0)

    def test_update_matrix_repression(self):
        """Test updating the matrix with RNA repression."""
        sub = protein(namespace='HGNC', name='A', identifier='1')
        obj = rna(namespace='HGNC', name='B', identifier='2')

        index = {'A', 'B'}

        test_dict = {}

        test_matrix = DataFrame(0, index=index, columns=index)

        test_dict["repression"] = test_matrix

        update_spia_matrices(test_dict, sub, obj, {'relation': 'decreases'})

        self.assertEqual(test_dict["repression"]['A']['B'], 1)
        self.assertEqual(test_dict["repression"]['A']['A'], 0)
        self.assertEqual(test_dict["repression"]['B']['A'], 0)
        self.assertEqual(test_dict["repression"]['B']['B'], 0)

    def test_update_matrix_activation(self):
        """Test updating the matrix with activation."""
        sub = protein(namespace='HGNC', name='A', identifier='1')
        obj = protein(namespace='HGNC', name='B', identifier='2')

        index = {'A', 'B'}

        test_dict = {}

        test_matrix = DataFrame(0, index=index, columns=index)

        test_dict["activation"] = test_matrix

        update_spia_matrices(test_dict, sub, obj, {'relation': 'increases'})

        self.assertEqual(test_dict["activation"]['A']['B'], 1)
        self.assertEqual(test_dict["activation"]['A']['A'], 0)
        self.assertEqual(test_dict["activation"]['B']['A'], 0)
        self.assertEqual(test_dict["activation"]['B']['B'], 0)

    def test_update_matrix_inhibition(self):
        """Test updating the matrix with activation."""
        sub = protein(namespace='HGNC', name='A', identifier='1')
        obj = protein(namespace='HGNC', name='B', identifier='2')

        index = {'A', 'B'}

        test_dict = {}

        test_matrix = DataFrame(0, index=index, columns=index)

        test_dict["inhibition"] = test_matrix

        update_spia_matrices(test_dict, sub, obj, {'relation': 'decreases'})

        self.assertEqual(test_dict["inhibition"]['A']['B'], 1)
        self.assertEqual(test_dict["inhibition"]['A']['A'], 0)
        self.assertEqual(test_dict["inhibition"]['B']['A'], 0)
        self.assertEqual(test_dict["inhibition"]['B']['B'], 0)

    def test_update_matrix_association(self):
        """Test updating the matrix with association."""
        sub = protein(namespace='HGNC', name='A', identifier='1')
        obj = protein(namespace='HGNC', name='B', identifier='2')

        index = {'A', 'B'}

        test_dict = {}

        test_matrix = DataFrame(0, index=index, columns=index)

        test_dict["binding_association"] = test_matrix

        update_spia_matrices(test_dict, sub, obj, {'relation': 'association'})

        self.assertEqual(test_dict["binding_association"]['A']['B'], 1)
        self.assertEqual(test_dict["binding_association"]['A']['A'], 0)
        self.assertEqual(test_dict["binding_association"]['B']['A'], 0)
        self.assertEqual(test_dict["binding_association"]['B']['B'], 0)

    def test_update_matrix_pmods(self):
        """Test updating the matrix with multiple protein modifications."""
        sub = protein(namespace='HGNC', name='A', identifier='1')
        obj = protein(namespace='HGNC', name='B', identifier='2', variants=[pmod('Ub'), pmod('Ph')])

        index = {'A', 'B'}

        test_dict = {}

        test_matrix = DataFrame(0, index=index, columns=index)

        test_dict["activation_ubiquination"] = test_matrix
        test_dict["activation_phosphorylation"] = test_matrix

        update_spia_matrices(test_dict, sub, obj, {'relation': 'increases'})

        self.assertEqual(test_dict["activation_ubiquination"]['A']['B'], 1)
        self.assertEqual(test_dict["activation_ubiquination"]['A']['A'], 0)
        self.assertEqual(test_dict["activation_ubiquination"]['B']['A'], 0)
        self.assertEqual(test_dict["activation_ubiquination"]['B']['B'], 0)

        self.assertEqual(test_dict["activation_phosphorylation"]['A']['B'], 1)
        self.assertEqual(test_dict["activation_phosphorylation"]['A']['A'], 0)
        self.assertEqual(test_dict["activation_phosphorylation"]['B']['A'], 0)
        self.assertEqual(test_dict["activation_phosphorylation"]['B']['B'], 0)

    def test_spia_matrix_complexes(self):
        """Test handling of complexes."""
        self.sialic_acid_graph.add_increases(
            sialic_acid_cd33_complex,
            trem2,
            citation=citation,
            annotations={'Species': '9606', 'Confidence': 'High'},
            evidence=evidence_1,
            object_modifier=activity()
        )

        spia_dfs = to_spia_dfs(self.sialic_acid_graph)

        self.assertEqual(spia_dfs["activation"][cd33.name][trem2.name], 1)

    def test_spia_matrix_composites(self):
        """Test handling of composites."""
        shp = composite_abundance([shp1, shp2])

        self.sialic_acid_graph.add_increases(
            shp,
            trem2,
            citation=citation,
            annotations={'Species': '9606', 'Confidence': 'High'},
            evidence=evidence_1,
            object_modifier=activity()
        )

        spia_dfs = to_spia_dfs(self.sialic_acid_graph)

        self.assertEqual(spia_dfs["activation"][shp1.name][trem2.name], 1)
        self.assertEqual(spia_dfs["activation"][shp2.name][trem2.name], 1)
