from seqlike import SeqLike
import numpy as np
import pytest

# ---- test sequence to index ---------
sequence_list = [
    ("ATGCGCGCG-", ([1.0, 4.0, 3.0, 2.0, 3.0, 2.0, 3.0, 2.0, 3.0, 0.0]), "dna"),
    ("AUGGUGCUGUUU-", ([1.0, 5.0, 3.0, 3.0, 5.0, 3.0, 2.0, 5.0, 3.0, 5.0, 5.0, 5.0, 0.0]), "rna"),
    ("AUAAUAUAU-", ([1.0, 5.0, 1.0, 1.0, 5.0, 1.0, 5.0, 1.0, 5.0, 0.0]), "rna"),
    ("LAIGHATY*-", ([11.0, 2.0, 9.0, 7.0, 8.0, 2.0, 18.0, 21.0, 0.0, 1.0]), "aa"),
    pytest.param(
        "#^@*)!&$&$)_#_",
        "#^@*)!&$&$)_#_",
        None,
        marks=pytest.mark.xfail(reason="Not a biological sequence."),
    ),
]


@pytest.mark.parametrize("sequence, indexed_array, seq_type", sequence_list)
def test_to_index(sequence, indexed_array, seq_type):
    idx_arr = SeqLike(sequence, seq_type=seq_type).to_index()
    assert np.array_equal(idx_arr, indexed_array)


# ----- test sequence to onehot ---------
DNA_ATGCGCGCGdash_onehot = np.array(
    [
        [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    ]
)

RNA_AUGGUGCUGdash_onehot = np.array(
    [
        [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
        [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
        [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
        [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    ]
)

AA_LAIGHATYstardash_onehot = np.array(
    [
        [
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            1.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
        ],
        [
            0.0,
            0.0,
            1.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
        ],
        [
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            1.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
        ],
        [
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            1.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
        ],
        [
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            1.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
        ],
        [
            0.0,
            0.0,
            1.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
        ],
        [
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            1.0,
            0.0,
            0.0,
            0.0,
            0.0,
        ],
        [
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            1.0,
            0.0,
        ],
        [
            1.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
        ],
        [
            0.0,
            1.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
        ],
    ]
)

sequence_list = [
    ("ATGCGCGCG-", DNA_ATGCGCGCGdash_onehot, "dna"),
    ("AUGGUGCUG-", RNA_AUGGUGCUGdash_onehot, "rna"),
    ("LAIGHATY*-", AA_LAIGHATYstardash_onehot, "aa"),
    pytest.param(
        "#^@*)!&$&$)_#_",
        "#^@*)!&$&$)_#_",
        None,
        marks=pytest.mark.xfail(reason="Not a biological sequence."),
    ),
]


@pytest.mark.parametrize("sequence, onehot_array, seq_type", sequence_list)
def test_to_onehot(sequence, onehot_array, seq_type):
    oh_arr = SeqLike(sequence, seq_type=seq_type).to_onehot()
    assert np.array_equal(oh_arr, onehot_array)
