from emmet.core.io.pymatgen import get_el_sp

from emmet.core.featurization.robocrys.util import (
    get_el,
    get_formatted_el,
    htmlify_spacegroup,
    superscript_number,
    unicodeify_spacegroup,
)
import pytest


def test_common_formulas():
    from emmet.core.featurization.robocrys.util import common_formulas

    assert common_formulas["H6CN"] == "CH3NH3"


def test_connected_geometries():
    from emmet.core.featurization.robocrys.util import connected_geometries

    assert "octahedral" in connected_geometries


def test_geometry_to_polyhedra():
    from emmet.core.featurization.robocrys.util import geometry_to_polyhedra

    assert geometry_to_polyhedra["octahedral"] == "octahedra"


def test_polyhedra_plurals():
    from emmet.core.featurization.robocrys.util import polyhedra_plurals

    assert polyhedra_plurals["pentagonal pyramid"] == "pentagonal pyramids"


def test_dimensionality_to_shape():
    from emmet.core.featurization.robocrys.util import dimensionality_to_shape

    assert dimensionality_to_shape[3] == "framework"


def test_superscript_number():
    assert superscript_number("+23") == "⁺²³"


def test_get_al():
    """Test getting element names."""
    species = get_el_sp("Sn2+")
    assert get_el(species) == "Sn"

    element = get_el_sp("Sn")
    assert get_el(element) == "Sn"

    assert get_el(50) == "Sn"


def test_get_formatted_el():
    """Test getting formatted element strings."""
    species = get_el_sp("Sn2+")
    form_el = get_formatted_el(species, "")
    assert form_el == "Sn2+"

    element = get_el_sp("Sn")
    form_el = get_formatted_el(element, "")
    assert form_el == "Sn"

    form_el = get_formatted_el(get_el(50), "")
    assert form_el == "Sn"

    for species, symm, use_oxi, use_sym, fmt, expected_label in [
        ("Sn2+", "(1,2)", True, True, "raw", "Sn(1,2)2+"),
        ("Sn2+", "(1,2)", False, True, "raw", "Sn(1,2)"),
        ("Sn2+", "(1,2)", True, False, "raw", "Sn2+"),
        ("Sn2+", "(1,2)", False, False, "raw", "Sn"),
        ("Sn2+", "(1,2)", True, True, "latex", "Sn(1,2)^{2+}"),
        ("Sn2+", "(1,2)", True, True, "html", "Sn(1,2)<sup>2+</sup>"),
        ("Sn2+", "(1,2)", True, True, "unicode", "Sn(1,2)²⁺"),
        ("Sn1.33-", "(3)", True, True, "unicode", "Sn(3)¹\u1427³³⁻"),
    ]:

        assert (
            get_formatted_el(
                species, symm, use_oxi_state=use_oxi, use_sym_label=use_sym, fmt=fmt
            )
            == expected_label
        )


@pytest.mark.xfail
def test_unicodeify_spacegroup():
    spg_symbol = unicodeify_spacegroup("P-42_1m")
    assert spg_symbol == "P4̅2₁m"

    spg_symbol2 = unicodeify_spacegroup("P2_12_12_1")
    assert spg_symbol2 == "P2₁2₁2₁"


@pytest.mark.xfail
def test_htmlify_spacegroup():
    spg_symbol = htmlify_spacegroup("P-42_1m")
    assert spg_symbol == "P4̅2<sub>1</sub>m"

    spg_symbol2 = htmlify_spacegroup("P2_12_12_1")
    assert spg_symbol2 == "P2<sub>1</sub>2<sub>1</sub>2<sub>1</sub>"


def test_load_condense_structure_json(test_condensed_structures):

    condensed_structure = test_condensed_structures["SnO2"].copy()

    # check that the site keys are correctly converted from str to int
    site_keys = list(condensed_structure["sites"].keys())
    assert isinstance(site_keys[0], int)

    component_keys = list(condensed_structure["components"].keys())
    assert isinstance(component_keys[0], int)
