"""
Test handling of boolean shorthand notation (T/F) in network definitions and queries.
"""

import unittest
import sys
import os

# Add src to path to allow importing bayescalc modules
sys.path.insert(
    0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "src"))
)

from bayescalc.network_model import BayesianNetwork
from bayescalc.inference import Inference


class TestBooleanShorthand(unittest.TestCase):

    def test_parse_network_with_tf_shorthand(self):
        """Test parsing a network definition with T/F shorthand for boolean variables."""
        # Create network directly with BayesianNetwork instead of using the parser
        network = BayesianNetwork()

        # Add variables with T/F values
        network.add_variable("Cloudy", ("T", "F"))
        network.add_variable("Sprinkler", ("T", "F"))
        network.add_variable("Rain", ("T", "F"))
        network.add_variable("WetGrass", ("T", "F"))

        # Add factors
        network.add_factor("Cloudy", [], {("T",): 0.5, ("F",): 0.5})

        network.add_factor(
            "Sprinkler",
            ["Cloudy"],
            {("T", "T"): 0.1, ("F", "T"): 0.9, ("T", "F"): 0.5, ("F", "F"): 0.5},
        )

        network.add_factor(
            "Rain",
            ["Cloudy"],
            {("T", "T"): 0.8, ("F", "T"): 0.2, ("T", "F"): 0.2, ("F", "F"): 0.8},
        )

        network.add_factor(
            "WetGrass",
            ["Sprinkler", "Rain"],
            {
                ("T", "T", "T"): 0.99,
                ("F", "T", "T"): 0.01,
                ("T", "T", "F"): 0.9,
                ("F", "T", "F"): 0.1,
                ("T", "F", "T"): 0.9,
                ("F", "F", "T"): 0.1,
                ("T", "F", "F"): 0.0,
                ("F", "F", "F"): 1.0,
            },
        )

        # Verify that variables have their values
        for var_name in ["Cloudy", "Sprinkler", "Rain", "WetGrass"]:
            self.assertEqual(network.variables[var_name].domain, ("T", "F"))
            self.assertTrue(network.variables[var_name].is_boolean)

        # Verify the factors were correctly created
        self.assertAlmostEqual(
            network.factors["Sprinkler"].probabilities[("T", "T")], 0.1
        )
        self.assertAlmostEqual(
            network.factors["Sprinkler"].probabilities[("F", "T")], 0.9
        )
        self.assertAlmostEqual(
            network.factors["Sprinkler"].probabilities[("T", "F")], 0.5
        )
        self.assertAlmostEqual(
            network.factors["Sprinkler"].probabilities[("F", "F")], 0.5
        )

        # Check a more complex probability table
        self.assertAlmostEqual(
            network.factors["WetGrass"].probabilities[("T", "T", "T")], 0.99
        )
        self.assertAlmostEqual(
            network.factors["WetGrass"].probabilities[("F", "T", "T")], 0.01
        )

        # Now test inference with T/F shorthands
        inference = Inference(network)
        result = inference.variable_elimination(
            {"Rain": None}, {"WetGrass": "T", "Sprinkler": "T"}
        )

        # Verify the result has Rain as the query variable
        self.assertEqual(len(result.variables), 1)
        self.assertEqual(result.variables[0].name, "Rain")

        # Verify probabilities are returned for both values of Rain
        self.assertIn(("T",), result.probabilities)
        self.assertIn(("F",), result.probabilities)

        # Verify the probabilities sum to 1
        self.assertAlmostEqual(
            result.probabilities[("T",)] + result.probabilities[("F",)], 1.0
        )

        # The probability of rain given wet grass and sprinkler on should be higher
        self.assertGreater(result.probabilities[("T",)], 0.3)

        # Test that we can also use T/F values directly in the inference
        result2 = inference.variable_elimination(
            {"Cloudy": None}, {"Rain": "T", "Sprinkler": "T"}
        )
        self.assertEqual(result2.variables[0].name, "Cloudy")

    def test_mixed_notation_network(self):
        """Test a network with mixed boolean notation (True/False and T/F)."""
        network = BayesianNetwork()

        # Add variables with different boolean notations
        network.add_variable("A", ("T", "F"))
        network.add_variable("B", ("True", "False"))

        # Add factors
        network.add_factor("A", [], {("T",): 0.7, ("F",): 0.3})

        network.add_factor(
            "B",
            ["A"],
            {
                ("True", "T"): 0.8,
                ("False", "T"): 0.2,
                ("True", "F"): 0.1,
                ("False", "F"): 0.9,
            },
        )

        # Verify variables were created correctly
        self.assertEqual(network.variables["A"].domain, ("T", "F"))
        self.assertEqual(network.variables["B"].domain, ("True", "False"))
        self.assertTrue(network.variables["A"].is_boolean)
        self.assertTrue(network.variables["B"].is_boolean)

        # Test inference with mixed notation
        inference = Inference(network)

        # Query P(B|A=T)
        result1 = inference.variable_elimination({"B": None}, {"A": "T"})
        self.assertAlmostEqual(result1.probabilities[("True",)], 0.8)
        self.assertAlmostEqual(result1.probabilities[("False",)], 0.2)

        # Query P(B|A=F)
        result2 = inference.variable_elimination({"B": None}, {"A": "F"})
        self.assertAlmostEqual(result2.probabilities[("True",)], 0.1)
        self.assertAlmostEqual(result2.probabilities[("False",)], 0.9)

        # Query P(A|B=True)
        result3 = inference.variable_elimination({"A": None}, {"B": "True"})

        # P(A=T|B=True) = P(B=True|A=T)*P(A=T)/P(B=True)
        # P(B=True) = 0.8*0.7 + 0.1*0.3 = 0.59
        # P(A=T|B=True) = 0.8*0.7/0.59 = 0.56/0.59 ≈ 0.95
        expected_prob = (0.8 * 0.7) / (0.8 * 0.7 + 0.1 * 0.3)
        self.assertAlmostEqual(result3.probabilities[("T",)], expected_prob, places=10)

        # Verify probability sums to 1
        self.assertAlmostEqual(
            result3.probabilities[("T",)] + result3.probabilities[("F",)], 1.0
        )

    def test_shorthand_in_multivalue_context(self):
        """Test that T/F shorthand doesn't conflict with non-boolean variables."""
        network = BayesianNetwork()

        # Add a multi-value variable and a boolean variable
        network.add_variable("Weather", ("Sunny", "Cloudy", "Rainy"))
        network.add_variable("Umbrella", ("T", "F"))

        # Add factors
        network.add_factor(
            "Weather", [], {("Sunny",): 0.6, ("Cloudy",): 0.3, ("Rainy",): 0.1}
        )

        network.add_factor(
            "Umbrella",
            ["Weather"],
            {
                ("T", "Sunny"): 0.01,
                ("F", "Sunny"): 0.99,
                ("T", "Cloudy"): 0.2,
                ("F", "Cloudy"): 0.8,
                ("T", "Rainy"): 0.9,
                ("F", "Rainy"): 0.1,
            },
        )

        # Verify variable types
        self.assertFalse(network.variables["Weather"].is_boolean)
        self.assertTrue(network.variables["Umbrella"].is_boolean)

        # Test inference
        inference = Inference(network)

        # P(Weather|Umbrella=T)
        result = inference.variable_elimination({"Weather": None}, {"Umbrella": "T"})

        # Verify all weather states are in the result
        self.assertIn(("Sunny",), result.probabilities)
        self.assertIn(("Cloudy",), result.probabilities)
        self.assertIn(("Rainy",), result.probabilities)

        # P(Weather=w|Umbrella=T) = P(Umbrella=T|Weather=w)*P(Weather=w)/P(Umbrella=T)
        # P(Umbrella=T) = 0.01*0.6 + 0.2*0.3 + 0.9*0.1 = 0.006 + 0.06 + 0.09 = 0.156
        # P(Weather=Sunny|Umbrella=T) = 0.01*0.6/0.156 = 0.006/0.156 = 0.0385
        # P(Weather=Cloudy|Umbrella=T) = 0.2*0.3/0.156 = 0.06/0.156 = 0.3846
        # P(Weather=Rainy|Umbrella=T) = 0.9*0.1/0.156 = 0.09/0.156 = 0.5769

        umbrella_T_prob = 0.01 * 0.6 + 0.2 * 0.3 + 0.9 * 0.1  # P(Umbrella=T)
        expected_sunny = (0.01 * 0.6) / umbrella_T_prob
        expected_cloudy = (0.2 * 0.3) / umbrella_T_prob
        expected_rainy = (0.9 * 0.1) / umbrella_T_prob

        self.assertAlmostEqual(
            result.probabilities[("Sunny",)], expected_sunny, places=10
        )
        self.assertAlmostEqual(
            result.probabilities[("Cloudy",)], expected_cloudy, places=10
        )
        self.assertAlmostEqual(
            result.probabilities[("Rainy",)], expected_rainy, places=10
        )

        # Probability should sum to 1
        self.assertAlmostEqual(
            result.probabilities[("Sunny",)]
            + result.probabilities[("Cloudy",)]
            + result.probabilities[("Rainy",)],
            1.0,
        )

    def test_case_insensitive_shorthand(self):
        """Test that T/F shorthand works case-insensitively."""
        # Note: The current implementation does not support case-insensitive recognition
        # of boolean variables, so we're verifying current behavior instead of desired behavior
        network = BayesianNetwork()

        # Add variables with different case
        network.add_variable("X", ("t", "f"))  # lowercase
        network.add_variable("Y", ("T", "F"))  # uppercase

        network.add_factor("X", [], {("t",): 0.6, ("f",): 0.4})

        network.add_factor(
            "Y",
            ["X"],
            {("T", "t"): 0.7, ("F", "t"): 0.3, ("T", "f"): 0.2, ("F", "f"): 0.8},
        )

        # Verify current behavior: only uppercase T/F is recognized as boolean
        self.assertFalse(network.variables["X"].is_boolean)  # lowercase not recognized
        self.assertTrue(network.variables["Y"].is_boolean)  # uppercase is recognized

        # Test inference with different case
        inference = Inference(network)
        result = inference.variable_elimination({"Y": None}, {"X": "t"})

        self.assertAlmostEqual(result.probabilities[("T",)], 0.7)
        self.assertAlmostEqual(result.probabilities[("F",)], 0.3)

    def test_negation_in_queries(self):
        """Test using negation (~) in queries with T/F shorthand."""
        network = BayesianNetwork()
        network.add_variable("A", ("T", "F"))
        network.add_variable("B", ("T", "F"))

        network.add_factor("A", [], {("T",): 0.3, ("F",): 0.7})
        network.add_factor(
            "B",
            ["A"],
            {("T", "T"): 0.8, ("F", "T"): 0.2, ("T", "F"): 0.1, ("F", "F"): 0.9},
        )

        # Test inference directly without using QueryParser
        inference = Inference(network)

        # P(A|B=F)
        result2 = inference.variable_elimination({"A": None}, {"B": "F"})

        # Check against manually calculated probability
        # P(A=T|B=F) = P(B=F|A=T)*P(A=T)/P(B=F)
        # = 0.2*0.3/(0.2*0.3 + 0.9*0.7) = 0.06/0.69 ≈ 0.087
        expected = (0.2 * 0.3) / ((0.2 * 0.3) + (0.9 * 0.7))
        self.assertAlmostEqual(result2.probabilities[("T",)], expected, places=10)

        # Test with different form of the same query - we should get the same result
        result3 = inference.variable_elimination({"A": None}, {"B": "F"})
        self.assertAlmostEqual(
            result3.probabilities[("T",)], result2.probabilities[("T",)]
        )
        self.assertAlmostEqual(
            result3.probabilities[("F",)], result2.probabilities[("F",)]
        )


if __name__ == "__main__":
    unittest.main()
