from CAI import CAI, RSCU, relative_adaptiveness
import pytest
import math


def test_bad_args():
    # make sure bad arguments raise errors
    with pytest.raises(TypeError):
        CAI("AAC")  # no reference data
    with pytest.raises(TypeError):
        CAI("AAC", reference=["AAC"], RSCUs=RSCU(["AAC"]))


def test_cai():
    # first, make sure all arguments get the same result
    assert (
        CAI("AAC", reference=["AAC"])
        == CAI("AAC", RSCUs=RSCU(["AAC"]))
        == CAI("AAC", weights=relative_adaptiveness(sequences=["AAC"]))
        == 1.0
    )

    # check other sequences
    assert CAI("AAT", reference=["AAC"]) == 0.5
    assert CAI("AATAAT", reference=["AAC"]) == 0.5
    assert CAI("AAT" * 100, reference=["AAC"]) == 0.5


def test_only_non_synonymous():
    assert math.isnan(CAI("ATG", reference=["AAC"]))


def test_stop_codon():
    # stop codons don't count for CAI
    assert (
        CAI("AAT", reference=["AAC"])
        == CAI("AATTAA", reference=["AAC"])
        == CAI("TAAAAT", reference=["AAC"])
    )


def test_alternate_genetic_code():
    # the differing stop codon
    assert CAI("AATTGA", reference=["AAC"]) != CAI(
        "AATTGA", reference=["AAC"], genetic_code=10
    )