from CAI import RSCU, relative_adaptiveness
import pytest


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


def test_arg_equivalence():
    # should be able to take either reference sequences or an RSCU dict
    assert relative_adaptiveness(sequences=["AAC"]) == relative_adaptiveness(
        RSCUs=RSCU(["AAC"])
    )


def test_weights():
    assert relative_adaptiveness(sequences=["AAC"]) == {
        "AAA": 1.0,
        "AAC": 1.0,
        "AAG": 1.0,
        "AAT": (0.5 / (0.5 * (1 + 0.5))) / (1 / (0.5 * (1 + 0.5))),
        "ACA": 1.0,
        "ACC": 1.0,
        "ACG": 1.0,
        "ACT": 1.0,
        "AGA": 1.0,
        "AGC": 1.0,
        "AGG": 1.0,
        "AGT": 1.0,
        "ATA": 1.0,
        "ATC": 1.0,
        "ATG": 1.0,
        "ATT": 1.0,
        "CAA": 1.0,
        "CAC": 1.0,
        "CAG": 1.0,
        "CAT": 1.0,
        "CCA": 1.0,
        "CCC": 1.0,
        "CCG": 1.0,
        "CCT": 1.0,
        "CGA": 1.0,
        "CGC": 1.0,
        "CGG": 1.0,
        "CGT": 1.0,
        "CTA": 1.0,
        "CTC": 1.0,
        "CTG": 1.0,
        "CTT": 1.0,
        "GAA": 1.0,
        "GAC": 1.0,
        "GAG": 1.0,
        "GAT": 1.0,
        "GCA": 1.0,
        "GCC": 1.0,
        "GCG": 1.0,
        "GCT": 1.0,
        "GGA": 1.0,
        "GGC": 1.0,
        "GGG": 1.0,
        "GGT": 1.0,
        "GTA": 1.0,
        "GTC": 1.0,
        "GTG": 1.0,
        "GTT": 1.0,
        "TAC": 1.0,
        "TAT": 1.0,
        "TCA": 1.0,
        "TCC": 1.0,
        "TCG": 1.0,
        "TCT": 1.0,
        "TGC": 1.0,
        "TGG": 1.0,
        "TGT": 1.0,
        "TTA": 1.0,
        "TTC": 1.0,
        "TTG": 1.0,
        "TTT": 1.0,
    }


def test_alternate_genetic_code():
    assert relative_adaptiveness(sequences=["AAC"], genetic_code=10) == {
        "AAA": 1.0,
        "AAC": 1.0,
        "AAG": 1.0,
        "AAT": (0.5 / (0.5 * (1 + 0.5))) / (1 / (0.5 * (1 + 0.5))),
        "ACA": 1.0,
        "ACC": 1.0,
        "ACG": 1.0,
        "ACT": 1.0,
        "AGA": 1.0,
        "AGC": 1.0,
        "AGG": 1.0,
        "AGT": 1.0,
        "ATA": 1.0,
        "ATC": 1.0,
        "ATG": 1.0,
        "ATT": 1.0,
        "CAA": 1.0,
        "CAC": 1.0,
        "CAG": 1.0,
        "CAT": 1.0,
        "CCA": 1.0,
        "CCC": 1.0,
        "CCG": 1.0,
        "CCT": 1.0,
        "CGA": 1.0,
        "CGC": 1.0,
        "CGG": 1.0,
        "CGT": 1.0,
        "CTA": 1.0,
        "CTC": 1.0,
        "CTG": 1.0,
        "CTT": 1.0,
        "GAA": 1.0,
        "GAC": 1.0,
        "GAG": 1.0,
        "GAT": 1.0,
        "GCA": 1.0,
        "GCC": 1.0,
        "GCG": 1.0,
        "GCT": 1.0,
        "GGA": 1.0,
        "GGC": 1.0,
        "GGG": 1.0,
        "GGT": 1.0,
        "GTA": 1.0,
        "GTC": 1.0,
        "GTG": 1.0,
        "GTT": 1.0,
        "TAC": 1.0,
        "TAT": 1.0,
        "TCA": 1.0,
        "TCC": 1.0,
        "TCG": 1.0,
        "TCT": 1.0,
        "TGA": 1.0,  # this is not a stop codon in genetic code 10
        "TGC": 1.0,
        "TGG": 1.0,
        "TGT": 1.0,
        "TTA": 1.0,
        "TTC": 1.0,
        "TTG": 1.0,
        "TTT": 1.0,
    }