Uncertain Knowledge & Reasoning

"""
AI Track Unit 7: Probabilistic Reasoning & Decision Making
1. Bayesian Network Exact Inference by Enumeration
2. Markov Decision Process (MDP) Value Iteration Solver
Author: DevAscent Engineering Team
"""

from typing import Dict, List, Tuple
import itertools

# =====================================================================
# 1. BAYESIAN NETWORK INFERENCE BY ENUMERATION
# =====================================================================
# Classic Burglar Alarm Network CPTs
CPT = {
    'B': {(): 0.001},
    'E': {(): 0.002},
    'A': {
        (True, True): 0.95,
        (True, False): 0.94,
        (False, True): 0.29,
        (False, False): 0.001
    },
    'J': {
        (True,): 0.90,
        (False,): 0.05
    },
    'M': {
        (True,): 0.70,
        (False,): 0.01
    }
}

VARIABLES = ['B', 'E', 'A', 'J', 'M']

def get_probability(var: str, val: bool, assignment: Dict[str, bool]) -> float:
    """Lookup conditional probability P(var=val | parents)."""
    if var == 'B':
        prob_true = CPT['B'][()]
    elif var == 'E':
        prob_true = CPT['E'][()]
    elif var == 'A':
        parents = (assignment['B'], assignment['E'])
        prob_true = CPT['A'][parents]
    elif var == 'J':
        parents = (assignment['A'],)
        prob_true = CPT['J'][parents]
    elif var == 'M':
        parents = (assignment['A'],)
        prob_true = CPT['M'][parents]
    else:
        raise ValueError(f"Unknown variable {var}")

    return prob_true if val else (1.0 - prob_true)

def enumerate_all(vars_left: List[str], evidence: Dict[str, bool]) -> float:
    """Recursively sum out hidden variables over all possible truth assignments."""
    if not vars_left:
        return 1.0

    Y = vars_left[0]
    rest = vars_left[1:]

    if Y in evidence:
        prob = get_probability(Y, evidence[Y], evidence)
        return prob * enumerate_all(rest, evidence)
    else:
        # Sum out hidden variable Y across True and False
        total = 0.0
        for val in [True, False]:
            extended_evidence = dict(evidence)
            extended_evidence[Y] = val
            prob = get_probability(Y, val, extended_evidence)
            total += prob * enumerate_all(rest, extended_evidence)
        return total

def query_bayes_net(query_var: str, evidence: Dict[str, bool]) -> Dict[bool, float]:
    """Compute normalized posterior distribution P(Query | Evidence)."""
    distribution = {}
    for val in [True, False]:
        extended_evidence = dict(evidence)
        extended_evidence[query_var] = val
        distribution[val] = enumerate_all(VARIABLES, extended_evidence)

    # Normalize probabilities so they sum to 1.0
    total = sum(distribution.values())
    for val in distribution:
        distribution[val] /= total
    return distribution

# =====================================================================
# 2. MARKOV DECISION PROCESS (MDP) — 3x4 GRIDWORLD VALUE ITERATION
# =====================================================================
class GridWorldMDP:
    def __init__(self, gamma: float = 0.9, noise: float = 0.2):
        self.gamma = gamma
        self.noise = noise
        self.rows = 3
        self.cols = 4

        # States are (r, c)
        self.wall = (1, 1)
        self.terminals = {(0, 3): 1.0, (1, 3): -1.0}
        self.living_reward = -0.04

        self.actions = ['UP', 'DOWN', 'LEFT', 'RIGHT']
        self.action_deltas = {
            'UP': (-1, 0),
            'DOWN': (1, 0),
            'LEFT': (0, -1),
            'RIGHT': (0, 1)
        }

    def get_states(self) -> List[Tuple[int, int]]:
        states = []
        for r in range(self.rows):
            for c in range(self.cols):
                if (r, c) != self.wall:
                    states.append((r, c))
        return states

    def move(self, state: Tuple[int, int], action: str) -> Tuple[int, int]:
        dr, dc = self.action_deltas[action]
        nr, nc = state[0] + dr, state[1] + dc
        # Check boundary collision or wall
        if nr < 0 or nr >= self.rows or nc < 0 or nc >= self.cols or (nr, nc) == self.wall:
            return state  # Bounce back into place
        return (nr, nc)

    def get_transitions(self, state: Tuple[int, int], action: str) -> List[Tuple[float, Tuple[int, int]]]:
        """Stochastic transition: 80% intended action, 10% veer left, 10% veer right."""
        if state in self.terminals:
            return []

        perpendiculars = {
            'UP': ['LEFT', 'RIGHT'],
            'DOWN': ['LEFT', 'RIGHT'],
            'LEFT': ['UP', 'DOWN'],
            'RIGHT': ['UP', 'DOWN']
        }

        transitions = [
            (1.0 - self.noise, self.move(state, action)),
            (self.noise / 2, self.move(state, perpendiculars[action][0])),
            (self.noise / 2, self.move(state, perpendiculars[action][1]))
        ]
        return transitions

    def value_iteration(self, iterations: int = 100) -> Tuple[Dict, Dict]:
        states = self.get_states()
        V = {s: 0.0 for s in states}
        for s in self.terminals:
            V[s] = self.terminals[s]

        for _ in range(iterations):
            new_V = dict(V)
            for s in states:
                if s in self.terminals:
                    continue

                action_values = []
                for a in self.actions:
                    expected_future_val = 0.0
                    for prob, next_s in self.get_transitions(s, a):
                        reward = self.terminals.get(next_s, self.living_reward)
                        expected_future_val += prob * (reward + self.gamma * V[next_s])
                    action_values.append(expected_future_val)

                new_V[s] = max(action_values)
            V = new_V

        # Extract optimal policy pi*(s)
        policy = {}
        for s in states:
            if s in self.terminals:
                policy[s] = 'GOAL' if self.terminals[s] > 0 else 'EXIT'
                continue

            best_a = None
            best_q = -float('inf')
            for a in self.actions:
                q_val = 0.0
                for prob, next_s in self.get_transitions(s, a):
                    reward = self.terminals.get(next_s, self.living_reward)
                    q_val += prob * (reward + self.gamma * V[next_s])
                if q_val > best_q:
                    best_q = q_val
                    best_a = a
            policy[s] = best_a

        return V, policy

# =====================================================================
# DEMONSTRATION HARNESS
# =====================================================================
if __name__ == "__main__":
    print("=" * 68)
    print("DEMO 1: BAYESIAN NETWORK EXACT INFERENCE BY ENUMERATION")
    print("Scenario: JohnCalls = True, MaryCalls = True. Query: P(Burglary)?")
    print("=" * 68)

    evidence = {'J': True, 'M': True}
    posterior_burglary = query_bayes_net('B', evidence)
    print(f"P(Burglary = True  | JohnCalls, MaryCalls) = {posterior_burglary[True]:.4f} (28.4% chance)")
    print(f"P(Burglary = False | JohnCalls, MaryCalls) = {posterior_burglary[False]:.4f} (71.6% chance)")

    print("\n" + "=" * 68)
    print("DEMO 2: MARKOV DECISION PROCESS (MDP) VALUE ITERATION")
    print("3x4 Stochastic GridWorld with Obstacle at (1,1) & Noise = 20%")
    print("=" * 68)

    mdp = GridWorldMDP(gamma=0.9, noise=0.2)
    V_opt, policy_opt = mdp.value_iteration(iterations=50)

    print("Optimal State Values V*(s):")
    for r in range(3):
        row_str = " | ".join(f"{V_opt.get((r, c), 0.0):6.2f}" if (r, c) != (1, 1) else "  WALL" for c in range(4))
        print(f"  [{row_str}]")

    print("\nOptimal Policy pi*(s) Directional Decisions:")
    for r in range(3):
        row_str = " | ".join(f"{policy_opt.get((r, c), ' '):5}" if (r, c) != (1, 1) else " WALL" for c in range(4))
        print(f"  [{row_str}]")