import pytest
from modrna_params import _data


def test_residue_set_and_mapping():
    assert _data.RESIDUES == (
        "5MC", "PSU", "5MU", "2MG", "1MG", "2MA", "1MA", "4AC", "6IA")
    assert set(_data.MODXNA_BASE) == set(_data.RESIDUES)
    assert _data.MODXNA_BASE["PSU"] == "PUU"   # not modXNA's "PSU" (thiouracil)
    assert _data.MODXNA_BASE["2MG"] == "M2G"   # not modXNA's "2MG" (dimethyl)


def test_termini_and_targets():
    assert _data.TERMINI == ("internal", "5prime", "3prime")
    assert _data.NET_CHARGE["internal"] == -1.0
    assert _data.NET_CHARGE["5prime"] == pytest.approx(-0.320348)
    assert _data.NET_CHARGE["3prime"] == pytest.approx(-0.679652)


def test_unit_name():
    assert _data.unit_name("5MC") == "5MC"
    assert _data.unit_name("5MC", "5prime") == "5MC5"
    assert _data.unit_name("5MC", "3prime") == "5MC3"
    with pytest.raises(KeyError):
        _data.unit_name("XXX")
    with pytest.raises(KeyError):
        _data.unit_name("5MC", "middle")


def test_expected_atom_count():
    # internal 5MC = 34; 5' cap removes phosphate + adds HO5' (-2); 3' cap adds HO3' (+1)
    assert _data.expected_atom_count("5MC") == 34
    assert _data.expected_atom_count("5MC", "5prime") == 32
    assert _data.expected_atom_count("5MC", "3prime") == 35
    assert _data.expected_atom_count("6IA") == 46
