from __future__ import annotations

import collections
from itertools import combinations
from numbers import Number
from unittest.mock import patch

import matplotlib.pyplot as plt
import numpy as np
import plotly.graph_objects as go
import pytest
from monty.json import MontyDecoder
from monty.serialization import dumpfn, loadfn
from numpy.testing import assert_allclose
from pytest import approx

from pymatgen.analysis.phase_diagram import (
    CompoundPhaseDiagram,
    GrandPotentialPhaseDiagram,
    GrandPotPDEntry,
    PatchedPhaseDiagram,
    PDEntry,
    PDPlotter,
    PhaseDiagram,
    ReactionDiagram,
    TransformedPDEntry,
    tet_coord,
    triangular_coord,
    uniquelines,
)
from pymatgen.core import Composition, DummySpecies, Element
from pymatgen.core.entries import ComputedEntry, EntrySet
from pymatgen.util.testing import TEST_FILES_DIR, MatSciTest

TEST_DIR = f"{TEST_FILES_DIR}/analysis"


class TestPDEntry:
    def setup_method(self):
        comp = Composition("LiFeO2")
        self.entry = PDEntry(comp, 53, name="mp-757614")
        self.gp_entry = GrandPotPDEntry(self.entry, {Element("O"): 1.5})

    def test_get_energy(self):
        assert self.entry.energy == 53, "Wrong energy!"
        assert self.gp_entry.energy == 50, "Wrong energy!"

    def test_get_chemical_energy(self):
        assert self.gp_entry.chemical_energy == 3, "Wrong energy!"

    def test_get_energy_per_atom(self):
        assert self.entry.energy_per_atom == approx(53.0 / 4), "Wrong energy per atom!"
        assert self.gp_entry.energy_per_atom == approx(50.0 / 2), "Wrong energy per atom!"

    def test_get_name(self):
        assert self.entry.name == "mp-757614"
        assert self.gp_entry.name == "mp-757614"

    def test_composition(self):
        comp = self.entry.composition
        expected_comp = Composition("LiFeO2")
        assert comp == expected_comp
        comp = self.gp_entry.composition
        expected_comp = Composition("LiFe")
        assert comp == expected_comp

    def test_formula(self):
        assert self.entry.formula == "Li1 Fe1 O2"
        assert self.gp_entry.formula == "Li1 Fe1 O2"

    def test_reduced_formula(self):
        assert self.entry.reduced_formula == "LiFeO2"
        assert self.gp_entry.reduced_formula == "LiFeO2"

    def test_elements(self):
        expected_elements = list(map(Element, ["Li", "Fe", "O"]))
        assert self.entry.elements == expected_elements

    def test_is_element(self):
        assert not self.entry.is_element
        assert not self.gp_entry.is_element

    def test_as_from_dict(self):
        dct = self.entry.as_dict()
        gpd = self.gp_entry.as_dict()
        entry = PDEntry.from_dict(dct)

        assert entry.name == "mp-757614"
        assert entry.energy_per_atom == approx(53.0 / 4)
        gp_entry = GrandPotPDEntry.from_dict(gpd)
        assert gp_entry.name == "mp-757614"
        assert gp_entry.energy_per_atom == approx(50.0 / 2)

        d_anon = dct.copy()
        del d_anon["name"]
        try:
            entry = PDEntry.from_dict(d_anon)
        except KeyError:
            self.fail("Should not need to supply name!")

    def test_str(self):
        assert str(self.entry) == "PDEntry : Li1 Fe1 O2 (mp-757614) with energy = 53.0000"
        pde = self.entry.as_dict()
        del pde["name"]
        pde = PDEntry.from_dict(pde)
        assert str(pde) == "PDEntry : Li1 Fe1 O2 with energy = 53.0000"

    def test_read_csv(self):
        entries = EntrySet.from_csv(f"{TEST_DIR}/pd_entries_test.csv")
        assert entries.chemsys == {"Li", "Fe", "O"}, "Wrong elements!"
        assert len(entries) == 490, "Wrong number of entries!"


class TestTransformedPDEntry:
    def setup_method(self):
        comp = Composition("LiFeO2")
        entry = PDEntry(comp, 53)

        terminal_compositions = [*map(Composition, ("Li2O", "FeO", "LiO8"))]

        sp_mapping = {}
        for idx, comp in enumerate(terminal_compositions):
            sp_mapping[comp] = DummySpecies(f"X{chr(102 + idx)}")

        self.transformed_entry = TransformedPDEntry(entry, sp_mapping)

    def test_get_energy(self):
        assert self.transformed_entry.energy == 53, "Wrong energy!"
        assert self.transformed_entry.original_entry.energy == approx(53.0)

    def test_get_energy_per_atom(self):
        assert self.transformed_entry.energy_per_atom == approx(53.0 / (23 / 15))

    def test_get_name(self):
        assert self.transformed_entry.name == "LiFeO2"

    def test_composition(self):
        comp = self.transformed_entry.composition
        expected_comp = Composition(
            {
                DummySpecies("Xf"): 14 / 30,
                DummySpecies("Xg"): 1.0,
                DummySpecies("Xh"): 2 / 30,
            }
        )
        assert comp == expected_comp

    def test_elements(self):
        expected_elements = list(map(Element, ["Li", "Fe", "O"]))
        assert self.transformed_entry.elements == expected_elements

    def test_is_element(self):
        assert self.transformed_entry.is_element is False
        assert self.transformed_entry.original_entry.is_element is False
        iron = Composition("Fe")
        assert TransformedPDEntry(PDEntry(iron, 0), {iron: iron}).is_element is True

    def test_as_from_dict(self):
        dct = self.transformed_entry.as_dict()
        entry = TransformedPDEntry.from_dict(dct)
        assert entry.name == "LiFeO2" == self.transformed_entry.name
        assert entry.energy_per_atom == approx(53.0 / (23 / 15)) == self.transformed_entry.energy_per_atom

    def test_str(self):
        assert str(self.transformed_entry).startswith("TransformedPDEntry Xf0+0.46666667 Xg")
        assert str(self.transformed_entry).endswith("with original composition Li1 Fe1 O2, energy = 53.0000")

    def test_normalize(self):
        norm_entry = self.transformed_entry.normalize(mode="atom")
        expected_comp = Composition(
            {
                DummySpecies("Xf"): 7 / 23,
                DummySpecies("Xg"): 15 / 23,
                DummySpecies("Xh"): 1 / 23,
            }
        )
        assert norm_entry.composition == expected_comp


class TestPhaseDiagram(MatSciTest):
    def setup_method(self):
        self.entries = EntrySet.from_csv(f"{TEST_DIR}/pd_entries_test.csv")
        self.pd = PhaseDiagram(self.entries)

    def test_init(self):
        # Ensure that a bad set of entries raises a PD error. Remove all Li
        # from self.entries.
        entries = filter(
            lambda e: (not e.composition.is_element) or e.elements[0] != Element("Li"),
            self.entries,
        )
        with pytest.raises(
            ValueError,
            match=r"Missing terminal entries for elements \['Fe', 'Li', 'O'\]",
        ):
            PhaseDiagram(entries)

    def test_repr(self):
        assert (
            str(self.pd) == "Li-Fe-O phase diagram\n11 stable phases: \nFe, FeO, "
            "Fe2O3, Fe3O4, LiFeO2, Li, Li2O, LiO, Li5FeO4, Li2FeO3, O"
        )

    def test_dim1(self):
        # Ensure that dim 1 PDs can be generated.
        for el in ("Li", "Fe", "O2"):
            entries = [entry for entry in self.entries if entry.reduced_formula == el]
            pd = PhaseDiagram(entries)
            assert len(pd.stable_entries) == 1

            for entry in entries:
                e_hull = pd.get_e_above_hull(entry)
                assert e_hull >= 0

            plotter = PDPlotter(pd)
            lines, *_ = plotter.pd_plot_data
            assert lines[0][1] == [0, 0]

    def test_ordering(self):
        # Test sorting of elements
        entries = [ComputedEntry(Composition(formula), 0) for formula in ["O", "N", "Fe"]]
        pd = PhaseDiagram(entries)
        sorted_elements = (Element("Fe"), Element("N"), Element("O"))
        assert tuple(pd.elements) == sorted_elements

        entries.reverse()
        pd = PhaseDiagram(entries)
        assert tuple(pd.elements) == sorted_elements

        # Test manual specification of order
        ordering = [Element(elt_string) for elt_string in ["O", "N", "Fe"]]
        pd = PhaseDiagram(entries, elements=ordering)
        assert tuple(pd.elements) == tuple(ordering)

    def test_stable_entries(self):
        stable_formulas = [ent.reduced_formula for ent in self.pd.stable_entries]
        expected_stable = "Fe2O3 Li5FeO4 LiFeO2 Fe3O4 Li Fe Li2O O2 FeO".split()
        for formula in expected_stable:
            assert formula in stable_formulas, f"{formula} not in stable entries!"

    def test_get_form_energy(self):
        expected_formation_energies = {
            "Li5FeO4": -164.8117344,
            "Li2O2": -14.119232793,
            "Fe2O3": -16.57416433,
            "FeO": -5.71415199,
            "Li": 0.0,
            "LiFeO2": -7.73275231,
            "Li2O": -6.229303868,
            "Fe": 0.0,
            "Fe3O4": -22.56571445,
            "Li2FeO3": -45.67166036,
            "O2": 0.0,
        }

        for entry in self.pd.stable_entries:
            formula = entry.reduced_formula
            expected = expected_formation_energies[formula]
            n_atoms = entry.composition.num_atoms
            # test get_form_energy
            assert self.pd.get_form_energy(entry) == approx(expected), formula
            # test get_form_energy_per_atom
            assert self.pd.get_form_energy_per_atom(entry) == approx(expected / n_atoms), formula

    def test_get_reference_energy(self):
        expected_ref_energies = {
            "Li5FeO4": -265.5546721,
            "Li2O2": -24.685172046,
            "Fe2O3": -51.93425725,
            "FeO": -21.70885048,
            "Li": -1.91301487,
            "LiFeO2": -17.02571825,
            "Li2O": -8.084307881,
            "Fe": -6.5961471,
            "Fe3O4": -73.6431077,
            "Li2FeO3": -92.78804506,
            "O2": -25.54966885,
        }
        for entry in self.pd.stable_entries:
            formula = entry.reduced_formula
            actual = self.pd.get_reference_energy(entry.composition)
            expected = expected_ref_energies[formula]
            assert actual == approx(expected), formula

    def test_all_entries_hulldata(self):
        assert len(self.pd.all_entries_hulldata) == 490

    def test_planar_inputs(self):
        elems = ["H", "He", "Li", "Be", "B", "Rb"]
        e1, e2, e3, e4, e5, e6 = (PDEntry(elem, 0) for elem in elems)

        pd = PhaseDiagram([e1, e2, e3, e4, e5, e6], map(Element, elems))

        assert len(pd.facets) == 1

    def test_str(self):
        assert (
            str(self.pd) == "Li-Fe-O phase diagram\n11 stable phases: \nFe, FeO, Fe2O3, Fe3O4,"
            " LiFeO2, Li, Li2O, LiO, Li5FeO4, Li2FeO3, O"
        )

    def test_get_e_above_hull(self):
        for entry in self.pd.all_entries:
            for entry in self.pd.stable_entries:
                decomp, e_hull = self.pd.get_decomp_and_e_above_hull(entry)
                assert e_hull < 1e-11, "Stable entries should have e_above_hull of zero!"
                assert decomp[entry] == 1, "Decomposition of stable entry should be itself."

            e_ah = self.pd.get_e_above_hull(entry)
            assert isinstance(e_ah, Number)
            assert e_ah >= 0

    def test_get_decomp_and_e_above_hull_on_error(self):
        for method, expected in (
            (self.pd.get_e_above_hull, None),
            (self.pd.get_decomp_and_e_above_hull, (None, None)),
        ):
            # test raises ValueError on entry with element not in the phase diagram
            U_entry = PDEntry("U", 0)
            with pytest.raises(
                ValueError,
                match="Unable to get decomposition for PDEntry : U1 with energy",
            ):
                method(U_entry)

            # test raises ValueError on entry with very negative energy
            too_neg_entry = PDEntry("Li", -1e6)
            match_msg = "No valid decomposition found for PDEntry : Li1 with energy"
            with pytest.raises(ValueError, match=match_msg):
                method(too_neg_entry)

            with pytest.warns(UserWarning, match=match_msg):
                out = method(too_neg_entry, on_error="warn")
            assert out == expected

            out = method(too_neg_entry, on_error="ignore")
            assert out == expected

    def test_downstream_methods_can_also_ignore_errors(self):
        # test that downstream methods get_e_above_hull() and get_phase_separation_energy()
        # can also ignore errors
        too_neg_entry = PDEntry("Li", -1e6)
        exotic_entry = PDEntry("U W", -1e6)

        # get_e_above_hull
        with pytest.raises(ValueError, match="No valid decomposition found for PDEntry "):
            self.pd.get_e_above_hull(too_neg_entry, on_error="raise")
        assert self.pd.get_e_above_hull(too_neg_entry, on_error="ignore") is None

        with pytest.raises(ValueError, match="Unable to get decomposition for PDEntry"):
            self.pd.get_e_above_hull(exotic_entry, on_error="raise")
        assert self.pd.get_e_above_hull(exotic_entry, on_error="ignore") is None

        # get_phase_separation_energy
        with pytest.raises(ValueError, match="Unable to get decomposition for PDEntry"):
            self.pd.get_phase_separation_energy(exotic_entry, on_error="raise")
        assert self.pd.get_phase_separation_energy(exotic_entry, on_error="ignore") is None

    def test_get_equilibrium_reaction_energy(self):
        for entry in self.pd.stable_entries:
            assert self.pd.get_equilibrium_reaction_energy(entry) <= 0, (
                "Stable entries should have negative equilibrium reaction energy!"
            )

    def test_get_phase_separation_energy(self):
        for entry in self.pd.unstable_entries:
            if entry.composition.fractional_composition not in [
                entry.composition.fractional_composition for entry in self.pd.stable_entries
            ]:
                assert self.pd.get_phase_separation_energy(entry) >= 0, (
                    "Unstable entries should have positive decomposition energy!"
                )
            elif entry.is_element:
                el_ref = self.pd.el_refs[entry.elements[0]]
                e_d = entry.energy_per_atom - el_ref.energy_per_atom
                assert self.pd.get_phase_separation_energy(entry) == approx(e_d)
            # NOTE the remaining materials would require explicit tests as they
            # could be either positive or negative

        for entry in self.pd.stable_entries:
            if entry.composition.is_element:
                assert self.pd.get_phase_separation_energy(entry) == 0, (
                    "Stable elemental entries should have decomposition energy of zero!"
                )
            else:
                assert self.pd.get_phase_separation_energy(entry) <= 0, (
                    "Stable entries should have negative decomposition energy!"
                )
                assert self.pd.get_phase_separation_energy(entry, stable_only=True) == approx(
                    self.pd.get_equilibrium_reaction_energy(entry)
                ), "Using `stable_only=True` should give decomposition energy equal to equilibrium reaction energy!"

        # Test that we get correct behavior with a polymorph
        toy_entries = {"Li": 0.0, "Li2O": -5, "LiO2": -4, "O2": 0.0}

        toy_pd = PhaseDiagram([PDEntry(c, e) for c, e in toy_entries.items()])

        # stable entry
        assert toy_pd.get_phase_separation_energy(PDEntry("Li2O", -5)) == approx(-1.0)
        # polymorph
        assert toy_pd.get_phase_separation_energy(PDEntry("Li2O", -4)) == approx(-2.0 / 3.0)

        # Test that the method works for novel entries
        novel_stable_entry = PDEntry("Li5FeO4", -999)
        assert self.pd.get_phase_separation_energy(novel_stable_entry) < 0, (
            "Novel stable entries should have negative decomposition energy!"
        )

        novel_unstable_entry = PDEntry("Li5FeO4", 999)
        assert self.pd.get_phase_separation_energy(novel_unstable_entry) > 0, (
            "Novel unstable entries should have positive decomposition energy!"
        )

        duplicate_entry = PDEntry("Li2O", -14.31361175)
        scaled_dup_entry = PDEntry("Li4O2", -14.31361175 * 2)
        stable_entry = next(entry for entry in self.pd.stable_entries if entry.name == "Li2O")

        assert self.pd.get_phase_separation_energy(duplicate_entry) == self.pd.get_phase_separation_energy(
            stable_entry
        ), "Novel duplicates of stable entries should have same decomposition energy!"

        assert self.pd.get_phase_separation_energy(scaled_dup_entry) == self.pd.get_phase_separation_energy(
            stable_entry
        ), "Novel scaled duplicates of stable entries should have same decomposition energy!"

    def test_inner_hull_reduction(self):
        """Regression test for two related bugs in get_decomp_and_phase_separation_energy:

        Bug 1: Line 997 passed _get_stable_entries_in_space(entry_elems) unfiltered into
        competing_entries, which re-introduced the same-composition stable entry that was
        explicitly excluded at line 980 via same_comp_mem_ids. The SLSQP would then
        trivially decompose the target into itself, producing a wrong (too large) energy.

        Bug 2: Line 999 overwrote the inner hull result with the original full set,
        making PhaseDiagram(reduced_space) dead code and defeating the space_limit
        optimization.

        The fix: filter same_comp_mem_ids when adding stable entries back at line 997,
        then remove the overwrite at line 999.
        """
        # Pick a stable non-elemental entry as the polymorph target.
        # The gate at line 976 requires a stable same-composition entry to exist,
        # which is satisfied since we're using a stable entry itself.
        entry = next(e for e in self.pd.stable_entries if not e.is_element)

        # space_limit=1 forces the inner hull branch for any non-trivial entry.
        result_reduced = self.pd.get_decomp_and_phase_separation_energy(entry, space_limit=1)
        result_full = self.pd.get_decomp_and_phase_separation_energy(entry, space_limit=10_000)

        # Both paths must produce the same phase separation energy.
        assert result_reduced[1] == approx(result_full[1], abs=1e-5), (
            "Inner hull reduction produced a different phase separation energy. "
            "Likely cause: same-composition stable entry leaked into competing_entries."
        )

        # The same-composition entry must NOT appear in the decomposition from either path.
        entry_frac = entry.composition.fractional_composition
        for decomp, _ in (result_reduced, result_full):
            for competing in decomp:
                assert competing.composition.fractional_composition != entry_frac, (
                    "Same-composition entry appeared in decomposition — should have been excluded."
                )

    def test_get_decomposition(self):
        for entry in self.pd.stable_entries:
            assert len(self.pd.get_decomposition(entry.composition)) == 1, (
                "Stable composition should have only 1 decomposition!"
            )
        dim = len(self.pd.elements)
        for entry in self.pd.all_entries:
            n_decomp = len(self.pd.get_decomposition(entry.composition))
            assert n_decomp > 0
            assert n_decomp <= dim, (
                "The number of decomposition phases can at most be equal to the number of components."
            )

        # Just to test decomposition for a fictitious composition
        actual = {entry.formula: amt for entry, amt in self.pd.get_decomposition(Composition("Li3Fe7O11")).items()}
        expected = {
            "Fe2 O2": 0.0952380952380949,
            "Li1 Fe1 O2": 0.5714285714285714,
            "Fe6 O8": 0.33333333333333393,
        }
        assert actual == approx(expected)

    def test_get_transition_chempots(self):
        for el in self.pd.elements:
            assert len(self.pd.get_transition_chempots(el)) <= len(self.pd.facets)

    def test_get_element_profile(self):
        for el in self.pd.elements:
            for entry in self.pd.stable_entries:
                if not entry.composition.is_element:
                    assert len(self.pd.get_element_profile(el, entry.composition)) <= len(self.pd.facets)

        expected = [
            {
                "evolution": 1.0,
                "chempot": -4.2582781416666666,
                "reaction": "Li2O + 0.5 O2 -> Li2O2",
            },
            {"evolution": 0, "chempot": -5.08859066, "reaction": "Li2O -> Li2O"},
            {
                "evolution": -1.0,
                "chempot": -10.48758201,
                "reaction": "Li2O -> 2 Li + 0.5 O2",
            },
        ]
        result = self.pd.get_element_profile(Element("O"), Composition("Li2O"))
        for d1, d2 in zip(expected, result, strict=True):
            assert d1["evolution"] == approx(d2["evolution"])
            assert d1["chempot"] == approx(d2["chempot"])
            assert d1["reaction"] == str(d2["reaction"])

    def test_get_get_chempot_range_map(self):
        elements = [el for el in self.pd.elements if el.symbol != "Fe"]
        assert len(self.pd.get_chempot_range_map(elements)) == 10

    def test_getmu_vertices_stability_phase(self):
        results = self.pd.getmu_vertices_stability_phase(Composition("LiFeO2"), Element("O"))
        assert len(results) == approx(6)
        test_equality = False
        for c in results:
            if (
                abs(c[Element("O")] + 7.115) < 1e-2
                and abs(c[Element("Fe")] + 6.596) < 1e-2
                and abs(c[Element("Li")] + 3.931) < 1e-2
            ):
                test_equality = True
        assert test_equality, "there is an expected vertex missing in the list"

    def test_getmu_range_stability_phase(self):
        results = self.pd.get_chempot_range_stability_phase(Composition("LiFeO2"), Element("O"))
        assert results[Element("O")][1] == approx(-4.450181224)
        assert results[Element("Fe")][0] == approx(-6.5961470)
        assert results[Element("Li")][0] == approx(-3.6250022625)

    def test_get_hull_energy(self):
        for entry in self.pd.stable_entries:
            h_e = self.pd.get_hull_energy(entry.composition)
            assert h_e == approx(entry.energy)
            n_h_e = self.pd.get_hull_energy(entry.composition.fractional_composition)
            assert n_h_e == approx(entry.energy_per_atom)

    def test_get_hull_energy_per_atom(self):
        for entry in self.pd.stable_entries:
            h_e = self.pd.get_hull_energy_per_atom(entry.composition)
            assert h_e == approx(entry.energy_per_atom)

    def test_1d_pd(self):
        entry = PDEntry("H", 0)
        pd = PhaseDiagram([entry])
        decomp, e_above_hull = pd.get_decomp_and_e_above_hull(PDEntry("H", 1))
        assert e_above_hull == 1
        assert decomp[entry] == approx(1.0)

    def test_get_critical_compositions_fractional(self):
        c1 = Composition("Fe2O3").fractional_composition
        c2 = Composition("Li3FeO4").fractional_composition
        c3 = Composition("Li2O").fractional_composition

        comps = self.pd.get_critical_compositions(c1, c2)
        expected = [
            Composition("Fe2O3").fractional_composition,
            Composition("Li0.3243244Fe0.1621621O0.51351349"),
            Composition("Li3FeO4").fractional_composition,
        ]
        for crit, exp in zip(comps, expected, strict=True):
            assert crit.almost_equals(exp, rtol=0, atol=1e-5)

        comps = self.pd.get_critical_compositions(c1, c3)
        expected = [
            Composition("Fe0.4O0.6"),
            Composition("LiFeO2").fractional_composition,
            Composition("Li5FeO4").fractional_composition,
            Composition("Li2O").fractional_composition,
        ]
        for crit, exp in zip(comps, expected, strict=True):
            assert crit.almost_equals(exp, rtol=0, atol=1e-5)

    def test_get_critical_compositions(self):
        c1 = Composition("Fe2O3")
        c2 = Composition("Li3FeO4")
        c3 = Composition("Li2O")

        comps = self.pd.get_critical_compositions(c1, c2)
        expected = [
            Composition("Fe2O3"),
            Composition("Li0.3243244Fe0.1621621O0.51351349") * 7.4,
            Composition("Li3FeO4"),
        ]
        for crit, exp in zip(comps, expected, strict=True):
            assert crit.almost_equals(exp, rtol=0, atol=1e-5)

        comps = self.pd.get_critical_compositions(c1, c3)
        expected = [
            Composition("Fe2O3"),
            Composition("LiFeO2"),
            Composition("Li5FeO4") / 3,
            Composition("Li2O"),
        ]
        for crit, exp in zip(comps, expected, strict=True):
            assert crit.almost_equals(exp, rtol=0, atol=1e-5)

        # Don't fail silently if input compositions aren't in phase diagram
        # Can be very confusing if you're working with a GrandPotentialPD
        with pytest.raises(ValueError, match="Xe1 has elements not in the phase diagram Li, Fe, O"):
            self.pd.get_critical_compositions(Composition("Xe"), Composition("Mn"))

        # For the moment, should also fail even if compositions are in the gppd
        # because it isn't handled properly
        gppd = GrandPotentialPhaseDiagram(self.pd.all_entries, {"Xe": 1}, [*self.pd.elements, Element("Xe")])
        with pytest.raises(
            ValueError,
            match="Li3 Fe1 O4 Xe1 has elements not in the phase diagram",
        ):
            gppd.get_critical_compositions(Composition("Fe2O3"), Composition("Li3FeO4Xe"))

        # check that the function still works though
        comps = gppd.get_critical_compositions(c1, c2)
        expected = [
            Composition("Fe2O3"),
            Composition("Li0.3243244Fe0.1621621O0.51351349") * 7.4,
            Composition("Li3FeO4"),
        ]
        for crit, exp in zip(comps, expected, strict=True):
            assert crit.almost_equals(exp, rtol=0, atol=1e-5)

        # case where the endpoints are identical
        assert self.pd.get_critical_compositions(c1, c1 * 2) == [c1, c1 * 2]

    def test_get_composition_chempots(self):
        c1 = Composition("Fe3.1O4")
        c2 = Composition("Fe3.2O4.1Li0.01")

        e1 = self.pd.get_hull_energy(c1)
        e2 = self.pd.get_hull_energy(c2)

        cp = self.pd.get_composition_chempots(c1)
        calc_e2 = e1 + sum(cp[k] * v for k, v in (c2 - c1).items())
        assert e2 == approx(calc_e2)

    def test_get_all_chempots(self):
        c1 = Composition("Fe3.1O4")
        c2 = Composition("FeO")

        cp1 = self.pd.get_all_chempots(c1)
        cp_result = {
            Element("Li"): -4.077061954,
            Element("Fe"): -6.741593864,
            Element("O"): -6.969907375,
        }

        for elem, energy in cp_result.items():
            assert cp1["Fe3O4-FeO-LiFeO2"][elem] == approx(energy)

        cp2 = self.pd.get_all_chempots(c2)
        cp_result = {
            Element("O"): -7.11535414,
            Element("Fe"): -6.5961471,
            Element("Li"): -3.93161518,
        }

        for elem, energy in cp_result.items():
            assert cp2["FeO-LiFeO2-Fe"][elem] == approx(energy)

    def test_get_plot(self):
        self.pd.get_plot()  # PDPlotter functionality is tested separately

    def test_as_from_dict(self):
        # test round-trip for other entry types such as ComputedEntry
        entry = ComputedEntry("H", 0.0, 0.0, entry_id="test")
        pd = PhaseDiagram([entry])
        pd_dict = pd.as_dict()
        pd_roundtrip = PhaseDiagram.from_dict(pd_dict)
        assert pd.all_entries[0].entry_id == pd_roundtrip.all_entries[0].entry_id

        pd_dict = self.pd.as_dict()
        reconstructed_pd = PhaseDiagram.from_dict(pd_dict)
        reconstructed_pd_dict = reconstructed_pd.as_dict()
        assert reconstructed_pd_dict == pd_dict
        assert isinstance(pd.to_json(), str)

        assert MontyDecoder().process_decoded(pd_dict).as_dict() == pd_dict

        for entry in self.pd.all_entries:
            # NOTE: allow_negative=True is necessary due to fp errors we see after serialization
            decomp_rpd, e_above_hull_rppd = reconstructed_pd.get_decomp_and_e_above_hull(entry, allow_negative=True)
            decomp_pd, e_above_hull_ppd = self.pd.get_decomp_and_e_above_hull(entry)

            # Check that decompositions sum to the original composition
            for decomp, name in [(decomp_pd, "pd"), (decomp_rpd, "rpd")]:
                decomp_comp = Composition({})
                for e, amount in decomp.items():
                    comp_scaled = e.composition.fractional_composition * amount
                    decomp_comp += comp_scaled
                assert decomp_comp.almost_equals(
                    entry.composition.fractional_composition, rtol=0, atol=Composition.amount_tolerance
                ), (
                    f"Decomposition for {entry} from {name} does not sum to original composition! "
                    f"Expected {entry.composition}, got {decomp_comp}"
                )

            # Compare decompositions by matching entries by composition
            # since serialization creates new entry objects with potentially slightly different energies
            decomp_pd_by_comp = {e.composition: amount for e, amount in decomp_pd.items()}
            decomp_rpd_by_comp = {e.composition: amount for e, amount in decomp_rpd.items()}
            assert decomp_pd_by_comp == approx(decomp_rpd_by_comp), f"Decomposition for {entry} is not correct!"
            assert np.isclose(e_above_hull_rppd, e_above_hull_ppd, rtol=1e-4, atol=1e-4)

    def test_read_json(self):
        dumpfn(self.pd, f"{self.tmp_path}/pd.json")
        pd = loadfn(f"{self.tmp_path}/pd.json")
        assert isinstance(pd, PhaseDiagram)
        assert pd.elements == self.pd.elements
        assert {*pd.as_dict()} == {*self.pd.as_dict()}

    def test_el_refs(self):
        # Create an imitation of pre_computed phase diagram with el_refs keys being
        # tuple[str, PDEntry] instead of tuple[Element, PDEntry].
        mock_el_refs = [(str(el), entry) for el, entry in self.pd.el_refs.items()]
        mock_computed_data = {**self.pd.computed_data, "el_refs": mock_el_refs}
        pd = PhaseDiagram(self.entries, computed_data=mock_computed_data)
        # Check the keys in el_refs dict have been updated to Element object via PhaseDiagram class.
        assert all(isinstance(el, Element) for el in pd.el_refs)

    def test_val_err_on_no_entries(self):
        # check that PhaseDiagram raises ValueError when building phase diagram with no entries
        for entries in [None, [], set(), ()]:
            with pytest.raises(ValueError, match="Unable to build phase diagram without entries."):
                PhaseDiagram(entries=entries)


class TestGrandPotentialPhaseDiagram:
    def setup_method(self):
        self.entries = EntrySet.from_csv(f"{TEST_DIR}/pd_entries_test.csv")
        self.pd = GrandPotentialPhaseDiagram(self.entries, {Element("O"): -5})
        self.pd6 = GrandPotentialPhaseDiagram(self.entries, {Element("O"): -6})

    def test_stable_entries(self):
        stable_formulas = [ent.original_entry.reduced_formula for ent in self.pd.stable_entries]
        expected_stable = ["Li5FeO4", "Li2FeO3", "LiFeO2", "Fe2O3", "Li2O2"]
        for formula in expected_stable:
            assert formula in stable_formulas, f"{formula} not in stable entries!"
        assert len(self.pd6.stable_entries) == 4

    def test_get_formation_energy(self):
        stable_formation_energies = {
            ent.original_entry.reduced_formula: self.pd.get_form_energy(ent) for ent in self.pd.stable_entries
        }
        expected_formation_energies = {
            "Fe2O3": 0.0,
            "Li5FeO4": -5.30551504,
            "Li2FeO3": -2.34247415,
            "LiFeO2": -0.4302639625,
            "Li2O2": 0.0,
        }
        for formula, energy in expected_formation_energies.items():
            assert energy == approx(stable_formation_energies[formula]), (
                f"Calculated formation for {formula} is not correct!"
            )

    def test_str(self):
        # using startswith since order of stable phases is random
        assert str(self.pd).startswith(
            "Fe-Li GrandPotentialPhaseDiagram with chempots = 'mu_O = -5.0000'5 stable phases: "
        )


def test_get_equilibrium_reaction_energy_open_element_in_original_composition():
    entries = [
        PDEntry(Composition("Li"), -2.0),
        PDEntry(Composition("Fe"), -8.0),
        PDEntry(Composition("O2"), -10.0),
        PDEntry(Composition("LiFe3O"), -40.0),
    ]
    gpd = GrandPotentialPhaseDiagram(entries, {"Li": -2.0})

    # Stable GP entry has projected composition Fe3O but original composition LiFe3O.
    gp_entry = next(
        entry
        for entry in gpd.stable_entries
        if entry.composition == Composition("Fe3O") and entry.original_comp == Composition("LiFe3O")
    )
    assert gp_entry in gpd.stable_entries

    ere = gpd.get_equilibrium_reaction_energy(gp_entry)
    assert isinstance(ere, float)
    assert ere <= 0


class TestCompoundPhaseDiagram:
    def setup_method(self):
        self.entries = EntrySet.from_csv(f"{TEST_DIR}/pd_entries_test.csv")
        self.pd = CompoundPhaseDiagram(self.entries, [Composition("Li2O"), Composition("Fe2O3")])

    def test_stable_entries(self):
        stable_formulas = [ent.name for ent in self.pd.stable_entries]
        expected_stable = ["Fe2O3", "Li5FeO4", "LiFeO2", "Li2O"]
        for formula in expected_stable:
            assert formula in stable_formulas

    def test_get_formation_energy(self):
        stable_formation_energies = {ent.name: self.pd.get_form_energy(ent) for ent in self.pd.stable_entries}
        expected_formation_energies = {
            "Li5FeO4": -7.0773284399999739,
            "Fe2O3": 0,
            "LiFeO2": -0.4745592975,
            "Li2O": 0,
        }
        for formula, energy in expected_formation_energies.items():
            assert energy == approx(stable_formation_energies[formula])

    def test_str(self):
        assert str(self.pd) == "Xf-Xg phase diagram\n4 stable phases: \nLiFeO2, Li2O, Li5FeO4, Fe2O3"

    def test_num2str(self):
        ret = set()

        num = 100
        for idx in range(num):
            val = CompoundPhaseDiagram.num2str(idx)
            assert val.isalpha()
            assert ord(val[0]) >= ord("f")
            ret.add(val)

        assert len(ret) == num


class TestPatchedPhaseDiagram:
    def setup_method(self):
        self.entries = EntrySet.from_csv(f"{TEST_DIR}/reaction_entries_test.csv")
        # NOTE add He to test for correct behavior despite no patches involving He
        self.no_patch_entry = he_entry = PDEntry("He", -1.23)
        self.entries.add(he_entry)

        self.pd = PhaseDiagram(entries=self.entries)
        self.ppd = PatchedPhaseDiagram(entries=self.entries)
        self.ppd_all = PatchedPhaseDiagram(entries=self.entries, keep_all_spaces=True)

        # novel entries not in any of the patches
        self.novel_comps = [Composition("H5C2OP"), Composition("V2PH4C")]
        for c in self.novel_comps:
            assert c.chemical_system not in self.ppd.spaces

        self.novel_entries = [PDEntry(c, -39.8) for c in self.novel_comps]

    def test_get_stable_entries(self):
        assert self.pd.stable_entries == self.ppd.stable_entries

    def test_get_qhull_entries(self):
        # NOTE qhull_entry is an specially sorted list due to it's construction, we
        # can't mimic this in ppd therefore just test if sorted versions are equal.
        assert sorted(self.pd.qhull_entries, key=lambda e: e.composition) == sorted(
            self.ppd.qhull_entries, key=lambda e: e.composition
        )

    def test_get_decomposition(self):
        for comp in self.novel_comps:
            decomp_pd = self.pd.get_decomposition(comp)
            decomp_ppd = self.ppd.get_decomposition(comp)
            assert decomp_pd == approx(decomp_ppd)

    def test_get_phase_separation_energy(self):
        for entry in self.novel_entries:
            e_phase_sep_pd = self.pd.get_phase_separation_energy(entry)
            e_phase_sep_ppd = self.ppd.get_phase_separation_energy(entry)
            assert np.isclose(e_phase_sep_pd, e_phase_sep_ppd)

    def test_get_equilibrium_reaction_energy(self):
        for entry in self.pd.stable_entries:
            e_equi_rxn_pd = self.pd.get_equilibrium_reaction_energy(entry)
            e_equi_rxn_pdd = self.ppd.get_equilibrium_reaction_energy(entry)
            assert np.isclose(e_equi_rxn_pd, e_equi_rxn_pdd)

    def test_get_form_energy(self):
        for entry in self.pd.stable_entries:
            e_form_pd = self.pd.get_form_energy(entry)
            e_form_ppd = self.ppd.get_form_energy(entry)
            assert np.isclose(e_form_pd, e_form_ppd)

    def test_dimensionality(self):
        assert self.pd.dim == self.ppd.dim

        # test dims of sub PDs
        dim_counts = collections.Counter(pd.dim for pd in self.ppd.pds.values())
        assert dim_counts == {4: 2, 3: 2}

        # test dims of sub PDs
        dim_counts = collections.Counter(pd.dim for pd in self.ppd_all.pds.values())
        assert dim_counts == {2: 8, 3: 7, 4: 2}

    def test_get_hull_energy(self):
        for comp in self.novel_comps:
            e_hull_pd = self.pd.get_hull_energy(comp)
            e_hull_ppd = self.ppd.get_hull_energy(comp)
            assert np.isclose(e_hull_pd, e_hull_ppd)

    def test_get_decomp_and_e_above_hull(self):
        for entry in self.pd.all_entries:
            decomp_pd, e_above_hull_pd = self.pd.get_decomp_and_e_above_hull(entry)
            decomp_ppd, e_above_hull_ppd = self.ppd.get_decomp_and_e_above_hull(entry, check_stable=True)

            # Check that decompositions sum to the original composition
            for decomp, name in [(decomp_pd, "pd"), (decomp_ppd, "ppd")]:
                decomp_comp = Composition({})
                for e, amount in decomp.items():
                    comp_scaled = e.composition.fractional_composition * amount
                    decomp_comp += comp_scaled
                assert decomp_comp.almost_equals(
                    entry.composition.fractional_composition, rtol=0, atol=Composition.amount_tolerance
                ), (
                    f"Decomposition for {entry} from {name} does not sum to original composition! "
                    f"Expected {entry.composition}, got {decomp_comp}"
                )

            assert decomp_pd == approx(decomp_ppd)
            assert np.isclose(e_above_hull_pd, e_above_hull_ppd)

    def test_repr(self):
        assert repr(self.ppd) == str(self.ppd) == "PatchedPhaseDiagram covering 4 sub-spaces"

    def test_as_from_dict(self):
        ppd_dict = self.ppd.as_dict()
        assert ppd_dict["@module"] == type(self.ppd).__module__
        assert ppd_dict["@class"] == type(self.ppd).__name__

        # Check format with computed_data
        assert "computed_data" in ppd_dict
        computed_data = ppd_dict["computed_data"]
        assert "all_entries" in computed_data
        assert isinstance(computed_data["all_entries"], list)
        assert all(isinstance(entry, dict) for entry in computed_data["all_entries"])

        # Verify entries can be reconstructed
        reconstructed_entries = [MontyDecoder().process_decoded(entry) for entry in computed_data["all_entries"]]
        assert len(reconstructed_entries) == len(self.ppd.all_entries)
        assert ppd_dict["elements"] == [elem.symbol for elem in self.ppd.elements]

        reconstructed_ppd = PatchedPhaseDiagram.from_dict(ppd_dict)
        reconstructed_dict = reconstructed_ppd.as_dict()
        assert ppd_dict == reconstructed_dict

        assert MontyDecoder().process_decoded(ppd_dict).as_dict() == ppd_dict

        for entry in self.ppd.all_entries:
            decomp_pd, e_above_hull_pd = self.pd.get_decomp_and_e_above_hull(entry)
            for pd, name in [
                (self.ppd, "ppd"),
                (reconstructed_ppd, "rppd"),
            ]:
                # NOTE: allow_negative=True is necessary due to fp errors we see after serialization
                decomp, e_above_hull = pd.get_decomp_and_e_above_hull(entry, check_stable=True, allow_negative=True)
                decomp_comp = Composition({})
                for e, amount in decomp.items():
                    comp_scaled = e.composition.fractional_composition * amount
                    decomp_comp += comp_scaled
                assert decomp_comp.almost_equals(
                    entry.composition.fractional_composition, rtol=0, atol=Composition.amount_tolerance
                ), (
                    f"Decomposition for {entry} from {name} does not sum to original composition! "
                    f"Expected {entry.composition}, got {decomp_comp}"
                )

                assert decomp == approx(decomp_pd), f"Decomposition for {entry} is not correct!"
                assert np.isclose(e_above_hull, e_above_hull_pd)

    def test_read_json(self, tmp_path):
        dumpfn(self.ppd, f"{tmp_path}/ppd.json")
        ppd = loadfn(f"{tmp_path}/ppd.json")
        assert isinstance(ppd, PatchedPhaseDiagram)
        assert ppd.elements == self.ppd.elements
        assert {*ppd.as_dict()} == {*self.ppd.as_dict()}

    def test_get_pd_for_entry(self):
        for entry in self.ppd.all_entries:
            if entry == self.no_patch_entry:
                continue
            pd = self.ppd.get_pd_for_entry(entry)
            # test that entry is in pd and pd can return valid decomp
            assert isinstance(pd.get_decomposition(entry.composition), dict)

        with pytest.raises(ValueError, match="No suitable PhaseDiagrams found for PDEntry"):
            self.ppd.get_pd_for_entry(self.no_patch_entry)

    def test_raises_on_missing_terminal_entries(self):
        entry = PDEntry("FeO", -1.23)
        with pytest.raises(ValueError, match=r"Missing terminal entries for elements \['Fe', 'O'\]"):
            PatchedPhaseDiagram(entries=[entry])

    def test_contains(self):
        for space in self.ppd.spaces:
            assert space in self.ppd
        unlikely_chem_space = frozenset(map(Element, "HBCNOFPS"))
        assert unlikely_chem_space not in self.ppd

    def test_getitem(self):
        chem_space = self.ppd.spaces[0]
        pd = self.ppd[chem_space]
        assert isinstance(pd, PhaseDiagram)
        assert chem_space in pd._qhull_spaces

        with pytest.raises(KeyError, match="frozenset"):
            self.ppd[frozenset(map(Element, "HBCNOFPS"))]

    def test_iter(self):
        for pd in self.ppd:
            assert isinstance(pd, PhaseDiagram)
        assert len(self.ppd) == len(self.ppd.pds)

    def test_len(self):
        assert len(self.ppd) == len(self.ppd.pds)

    def test_setitem_and_delitem(self):
        unlikely_chem_space = frozenset(map(Element, "HBCNOFPS"))
        self.ppd[unlikely_chem_space] = self.pd
        assert unlikely_chem_space in self.ppd
        assert self.ppd[unlikely_chem_space] == self.pd
        del self.ppd[unlikely_chem_space]  # test __delitem__() and restore original state

    def test_remove_redundant_spaces(self):
        spaces = tuple(frozenset(entry.elements) for entry in self.ppd.qhull_entries)
        # NOTE this is 5 not 4 as "He" is a non redundant space that gets dropped for other reasons
        assert len(self.ppd.remove_redundant_spaces(spaces)) == 5

        test = (
            list(combinations(range(1, 7), 4))
            + list(combinations(range(1, 10), 2))
            + list(combinations([1, 4, 7, 9, 2], 5))
        )
        test = [frozenset(t) for t in test]
        assert len(self.ppd.remove_redundant_spaces(test)) == 30

    def test_update_returns_new_instance(self):
        """update() must return a new PatchedPhaseDiagram, not mutate self."""
        original_stable = self.ppd.stable_entries.copy()
        ppd_new = self.ppd.update([PDEntry("V2O5", -500)])
        assert ppd_new is not self.ppd
        assert isinstance(ppd_new, PatchedPhaseDiagram)
        assert self.ppd.stable_entries == original_stable

    def test_update_with_empty_entries(self):
        ppd_new = self.ppd.update([])
        assert ppd_new is self.ppd

    def test_update_entry_in_existing_space(self):
        """Add a very low energy entry in an existing chemical space and verify
        it becomes stable and the patched PD matches a full rebuild."""
        new_entry = PDEntry("V2O5", -500)
        ppd_new, info = self.ppd.update([new_entry], return_info=True)
        assert new_entry in info["new_stable_entries"]
        assert len(info["new_spaces"]) == 0
        assert any(frozenset(new_entry.elements).issubset(s) for s in info["updated_spaces"])
        assert new_entry in ppd_new.stable_entries
        ppd_rebuilt = PatchedPhaseDiagram(entries=[*list(self.entries), new_entry])
        assert ppd_new.stable_entries == ppd_rebuilt.stable_entries

    def test_update_new_chemical_space(self):
        """Add entries that span a chemical space not present in any existing patch."""
        new_entries = [PDEntry("FeO", -10), PDEntry("Fe", -1), PDEntry("O", -2)]
        existing_elements = {el for entry in self.entries for el in entry.elements}
        fe_el = Element("Fe")
        o_el = Element("O")
        if fe_el in existing_elements and o_el in existing_elements:
            ppd = PatchedPhaseDiagram(entries=self.entries)
            fe_o_space = frozenset([fe_el, o_el])
            if not any(s.issuperset(fe_o_space) for s in ppd.spaces):
                ppd_new, info = ppd.update(new_entries, return_info=True)
                assert len(info["new_spaces"]) > 0
            else:
                ppd_full = PatchedPhaseDiagram(entries=list(self.entries) + new_entries)
                ppd_new = ppd.update(new_entries)
                assert ppd_new.stable_entries == ppd_full.stable_entries

    def test_update_matches_full_rebuild(self):
        """The updated PPD should give the same stable entries, decompositions,
        and hull energies as a PPD built from scratch with all entries."""
        new_entries = [PDEntry("V2O3", -200), PDEntry("PH3", -50)]
        ppd_updated = PatchedPhaseDiagram(entries=self.entries).update(new_entries)
        all_entries = list(self.entries) + new_entries
        ppd_full = PatchedPhaseDiagram(entries=all_entries)
        assert ppd_updated.stable_entries == ppd_full.stable_entries
        for comp in self.novel_comps:
            e_hull_updated = ppd_updated.get_hull_energy(comp)
            e_hull_full = ppd_full.get_hull_energy(comp)
            assert np.isclose(e_hull_updated, e_hull_full)

    def test_update_unstable_entry_no_change(self):
        """Adding a high-energy entry should not change stable entries."""
        stable_before = self.ppd.stable_entries.copy()
        ppd_new, info = self.ppd.update([PDEntry("V2O5", 1000)], return_info=True)
        assert info["new_stable_entries"] == set()
        assert info["removed_stable_entries"] == set()
        assert ppd_new.stable_entries == stable_before

    def test_update_multiple_batches(self):
        """Simulate a discovery campaign with multiple sequential updates."""
        ppd = PatchedPhaseDiagram(entries=self.entries)
        batch1 = [PDEntry("V2O5", -500)]
        ppd = ppd.update(batch1)
        batch2 = [PDEntry("PH3", -50)]
        ppd = ppd.update(batch2)
        all_entries = list(self.entries) + batch1 + batch2
        ppd_full = PatchedPhaseDiagram(entries=all_entries)
        assert ppd.stable_entries == ppd_full.stable_entries

    def test_update_rejects_new_elements(self):
        """Entries with elements not in the original PPD should raise."""
        with pytest.raises(ValueError, match="does not support adding new elements"):
            self.ppd.update([PDEntry("LiF", -10)])

    def test_update_returns_removed_stable(self):
        """When a new entry displaces an old stable entry, it should appear
        in removed_stable_entries."""
        stable_formulas = {e.reduced_formula for e in self.ppd.stable_entries}
        if "V2O5" in stable_formulas:
            target_comp = "V2O5"
        else:
            target_comp = next(e.reduced_formula for e in self.ppd.stable_entries if not e.composition.is_element)
        old_entry = next(e for e in self.ppd.stable_entries if e.reduced_formula == target_comp)
        crusher = PDEntry(target_comp, old_entry.energy * 10)
        _, info = self.ppd.update([crusher], return_info=True)
        assert crusher in info["new_stable_entries"]
        assert old_entry in info["removed_stable_entries"]

    def test_update_preserves_keep_all_spaces(self):
        """The keep_all_spaces setting should be stored and respected on update."""
        ppd_all = PatchedPhaseDiagram(entries=self.entries, keep_all_spaces=True)
        assert ppd_all.keep_all_spaces is True
        n_spaces_before = len(ppd_all.spaces)
        ppd_new = ppd_all.update([PDEntry("V2O5", -500)])
        assert ppd_new.keep_all_spaces is True
        assert len(ppd_new.spaces) >= n_spaces_before

    def test_update_shares_entry_objects(self):
        """Entry objects should be shared (same id) between old and new PPD."""
        original_entry_ids = {id(e) for e in self.ppd.all_entries}
        ppd_new = self.ppd.update([PDEntry("V2O5", -500)])
        carried_ids = {id(e) for e in ppd_new.all_entries} & original_entry_ids
        assert carried_ids == original_entry_ids

    def test_update_same_comp_lower_energy_replaces(self):
        """A new entry with the same composition but lower energy should
        replace the existing qhull entry and become stable."""
        target = next(
            e for e in self.ppd.stable_entries if not e.composition.is_element and len(e.composition.elements) >= 2
        )
        better = PDEntry(target.composition, target.energy * 10)
        ppd_new, info = self.ppd.update([better], return_info=True)
        assert better in ppd_new.stable_entries
        assert better in info["new_stable_entries"]
        assert target in info["removed_stable_entries"]
        rc = target.composition.reduced_composition
        qhull_for_comp = [e for e in ppd_new.qhull_entries if e.composition.reduced_composition == rc]
        assert len(qhull_for_comp) == 1
        assert qhull_for_comp[0] is better

    def test_update_same_comp_higher_energy_noop(self):
        """A new entry with the same composition but higher energy should
        not replace the existing qhull entry or change stability."""
        target = next(
            e for e in self.ppd.stable_entries if not e.composition.is_element and len(e.composition.elements) >= 2
        )
        worse = PDEntry(target.composition, target.energy / 10)
        ppd_new, info = self.ppd.update([worse], return_info=True)
        assert target in ppd_new.stable_entries
        assert worse not in ppd_new.stable_entries
        assert info["new_stable_entries"] == set()
        assert info["removed_stable_entries"] == set()
        rc = target.composition.reduced_composition
        qhull_for_comp = [e for e in ppd_new.qhull_entries if e.composition.reduced_composition == rc]
        assert len(qhull_for_comp) == 1
        assert qhull_for_comp[0] is target

    def test_update_lower_elemental_ref_raises(self):
        """By default, a new elemental entry with lower energy should raise."""
        el = next(iter(self.ppd.el_refs))
        old_ref = self.ppd.el_refs[el]
        new_ref = PDEntry(str(el), old_ref.energy_per_atom * 2)
        with pytest.raises(ValueError, match="lower than current el_ref"):
            self.ppd.update([new_ref])

    def test_update_lower_elemental_ref_ignore(self):
        """With on_new_el_ref='ignore', el_refs should stay unchanged."""
        el = next(iter(self.ppd.el_refs))
        old_ref = self.ppd.el_refs[el]
        new_ref = PDEntry(str(el), old_ref.energy_per_atom * 2)
        ppd_new = self.ppd.update([new_ref], on_new_el_ref="ignore")
        assert ppd_new.el_refs[el] is old_ref

    def test_update_lower_elemental_ref_recalculate(self):
        """With on_new_el_ref='recalculate', el_refs should update and
        affected spaces should be rebuilt."""
        el = next(iter(self.ppd.el_refs))
        old_ref = self.ppd.el_refs[el]
        new_ref = PDEntry(str(el), old_ref.energy_per_atom * 2)
        ppd_new, info = self.ppd.update([new_ref], return_info=True, on_new_el_ref="recalculate")
        assert ppd_new.el_refs[el] is new_ref
        affected_els = {el}
        for space in info["updated_spaces"]:
            assert space & affected_els
        ppd_full = PatchedPhaseDiagram(entries=[*list(self.entries), new_ref])
        assert ppd_new.stable_entries == ppd_full.stable_entries

    def test_update_qhull_dedup_across_batches(self):
        """After multiple updates adding entries for the same composition,
        qhull_entries should contain only one entry per reduced composition."""
        comp = "V2O5"
        ppd = PatchedPhaseDiagram(entries=self.entries)
        ppd = ppd.update([PDEntry(comp, -100)])
        ppd = ppd.update([PDEntry(comp, -500)])
        rc = Composition(comp).reduced_composition
        qhull_for_comp = [e for e in ppd.qhull_entries if e.composition.reduced_composition == rc]
        assert len(qhull_for_comp) == 1
        assert np.isclose(qhull_for_comp[0].energy, -500)

    def test_update_all_positive_formation_energy(self):
        """When all new entries have positive formation energy, the PPD
        should still return a new instance with all_entries updated but
        stable entries unchanged."""
        stable_before = self.ppd.stable_entries.copy()
        n_entries_before = len(self.ppd.all_entries)
        ppd_new = self.ppd.update([PDEntry("V2O5", 9999)])
        assert ppd_new.stable_entries == stable_before
        assert len(ppd_new.all_entries) == n_entries_before + 1


class TestReactionDiagram:
    def setup_method(self):
        self.entries = list(EntrySet.from_csv(f"{TEST_DIR}/reaction_entries_test.csv").entries)
        for entry in self.entries:
            if entry.reduced_formula == "VPO5":
                entry1 = entry
            elif entry.reduced_formula == "H4(CO)3":
                entry2 = entry
        self.rd = ReactionDiagram(entry1=entry1, entry2=entry2, all_entries=self.entries[2:])

    def test_get_compound_pd(self):
        self.rd.get_compound_pd()

    def test_formula(self):
        for entry in self.rd.rxn_entries:
            assert Element.V in entry.composition
            assert Element.O in entry.composition
            assert Element.C in entry.composition
            assert Element.P in entry.composition
            assert Element.H in entry.composition
        # formed_formula = [entry.reduced_formula for entry in self.rd.rxn_entries]
        # expected_formula = [
        #     "V0.12707182P0.12707182H0.0441989C0.03314917O0.66850829",
        #     "V0.125P0.125H0.05C0.0375O0.6625",
        #     "V0.12230216P0.12230216H0.05755396C0.04316547O0.65467626",
        #     "V0.11340206P0.11340206H0.08247423C0.06185567O0.62886598",
        #     "V0.11267606P0.11267606H0.08450704C0.06338028O0.62676056",
        #     "V0.11229947P0.11229947H0.0855615C0.06417112O0.62566845",
        #     "V0.09677419P0.09677419H0.12903226C0.09677419O0.58064516",
        #     "V0.05882353P0.05882353H0.23529412C0.17647059O0.47058824",
        #     "V0.04225352P0.04225352H0.28169014C0.21126761O0.42253521",
        # ]

        # # Please do not uncomment this test. This test is far too fragile because of numerical precision errors.
        # # Unless someone wants to make an effort to write a PROPER test which do not fail with changes in
        # # OS or numpy versions, DO NOT UNCOMMENT!
        # for formula in expected_formula:
        #     assert formula in formed_formula, f"{formed_formula=} not in {expected_formula=}"


class TestPDPlotter:
    def setup_method(self):
        entries = list(EntrySet.from_csv(f"{TEST_DIR}/pd_entries_test.csv"))

        elemental_entries = [entry for entry in entries if entry.elements == [Element("Li")]]
        self.pd_unary = PhaseDiagram(elemental_entries)
        self.plotter_unary_plotly = PDPlotter(self.pd_unary, backend="plotly")

        entries_LiO = [entry for entry in entries if "Fe" not in entry.composition]
        self.pd_binary = PhaseDiagram(entries_LiO)
        self.plotter_binary_mpl = PDPlotter(self.pd_binary, backend="matplotlib")
        self.plotter_binary_plotly = PDPlotter(self.pd_binary, backend="plotly")

        self.pd_ternary = PhaseDiagram(entries)
        self.plotter_ternary_mpl = PDPlotter(self.pd_ternary, backend="matplotlib")
        self.plotter_ternary_plotly_2d = PDPlotter(self.pd_ternary, backend="plotly", ternary_style="2d")
        self.plotter_ternary_plotly_3d = PDPlotter(self.pd_ternary, backend="plotly", ternary_style="3d")

        entries.append(PDEntry("C", 0))
        self.pd_quaternary = PhaseDiagram(entries)
        self.plotter_quaternary_mpl = PDPlotter(self.pd_quaternary, backend="matplotlib")
        self.plotter_quaternary_plotly = PDPlotter(self.pd_quaternary, backend="plotly")

    def test_plot_pd_with_no_unstable(self):
        # https://github.com/materialsproject/pymatgen/issues/2885
        pd_entries = [PDEntry(comp, 0) for comp in ["Li", "Co", "O"]]
        pd = PhaseDiagram(pd_entries)
        plotter = PDPlotter(pd, backend="plotly", show_unstable=False)
        ax = plotter.get_plot()
        assert isinstance(ax, go.Figure)
        assert len(ax.data) == 4

    def test_pd_plot_data(self):
        lines, labels, unstable_entries = self.plotter_ternary_mpl.pd_plot_data
        assert len(lines) == 22
        assert len(labels) == len(self.pd_ternary.stable_entries), "Incorrect number of lines generated!"
        assert len(unstable_entries) == len(self.pd_ternary.all_entries) - len(self.pd_ternary.stable_entries), (
            "Incorrect number of lines generated!"
        )
        lines, labels, unstable_entries = self.plotter_quaternary_mpl.pd_plot_data
        assert len(lines) == 33
        assert len(labels) == len(self.pd_quaternary.stable_entries)
        assert len(unstable_entries) == len(self.pd_quaternary.all_entries) - len(self.pd_quaternary.stable_entries)
        lines, labels, unstable_entries = self.plotter_binary_mpl.pd_plot_data
        assert len(lines) == 3
        assert len(labels) == len(self.pd_binary.stable_entries)

    def test_matplotlib_plots(self):
        # Some very basic ("non")-tests. Just to make sure the methods are callable.
        for plotter in (
            self.plotter_binary_mpl,
            self.plotter_ternary_mpl,
            self.plotter_quaternary_mpl,
        ):
            ax = plotter.get_plot()
            assert isinstance(ax, plt.Axes)

        self.plotter_ternary_mpl.get_contour_pd_plot()
        self.plotter_ternary_mpl.get_chempot_range_map_plot([Element("Li"), Element("O")])
        self.plotter_ternary_mpl.plot_element_profile(Element("O"), Composition("Li2O"))

        # test show()
        assert self.plotter_ternary_mpl.show() is None

    def test_plotly_plots(self):
        # Also very basic tests. Ensures callability and 2D vs 3D properties.
        for plotter in (
            self.plotter_unary_plotly,
            self.plotter_binary_plotly,
            self.plotter_ternary_plotly_2d,
            self.plotter_ternary_plotly_3d,
            self.plotter_quaternary_plotly,
        ):
            fig = plotter.get_plot()
            assert isinstance(fig, go.Figure)

        # test show()
        # suppress default plotly behavior of opening figure in browser by patching plotly.io.show to noop
        with patch("plotly.io.show") as mock_show:
            assert self.plotter_ternary_plotly_2d.show() is None
            mock_show.assert_called_once()


class TestUtilityFunction:
    def test_unique_lines(self):
        testdata = [
            [5, 53, 353],
            [399, 20, 52],
            [399, 400, 20],
            [13, 399, 52],
            [21, 400, 353],
            [393, 5, 353],
            [400, 393, 353],
            [393, 400, 399],
            [393, 13, 5],
            [13, 393, 399],
            [400, 17, 20],
            [21, 17, 400],
        ]
        expected = {
            (5, 393),
            (21, 353),
            (353, 400),
            (5, 13),
            (17, 20),
            (21, 400),
            (17, 400),
            (52, 399),
            (393, 399),
            (20, 52),
            (353, 393),
            (5, 353),
            (5, 53),
            (13, 399),
            (393, 400),
            (13, 52),
            (53, 353),
            (17, 21),
            (13, 393),
            (20, 399),
            (399, 400),
            (20, 400),
        }
        assert uniquelines(testdata) == expected

    def test_triangular_coord(self):
        coord = [0.5, 0.5]
        coord = triangular_coord(coord)
        assert_allclose(coord, [0.75, 0.4330127])

    def test_tet_coord(self):
        coord = [0.5, 0.5, 0.5]
        coord = tet_coord(coord)
        assert_allclose(coord, [1.0, 0.57735027, 0.40824829])
