"""
TAI 100, Week 2: Probability, Bayes' rule and beliefs that update.

Part A  The clinic test. Checks Bayes' rule against plain counting.
Part B  The umbrella world from Week 1, now with an agent that keeps a single
        belief ("how likely is it that we are in the rainy season?") and
        updates it every day with Bayes' rule. We compare it with the
        Week 1 learner that remembers the last k days.

Run:  python3 week02_bayes.py
Only the Python standard library is needed.
"""
import random

# --------------------------------------------------------------------------
# Part A: the clinic test (made-up but realistic numbers)
# --------------------------------------------------------------------------
HAS_MALARIA = 0.10        # 10% of feverish patients at this clinic have malaria
SENSITIVITY = 0.95        # the test is positive for 95% of people who have it
FALSE_POSITIVE = 0.10     # and positive for 10% of people who do not


def bayes_positive():
    """P(malaria | positive test), from Bayes' rule."""
    yes = SENSITIVITY * HAS_MALARIA
    no = FALSE_POSITIVE * (1 - HAS_MALARIA)
    return yes / (yes + no)


def count_positive(patients, rng):
    """The same quantity, estimated by simulating patients and counting."""
    positives = sick_and_positive = 0
    for _ in range(patients):
        sick = rng.random() < HAS_MALARIA
        positive = rng.random() < (SENSITIVITY if sick else FALSE_POSITIVE)
        if positive:
            positives += 1
            sick_and_positive += sick
    return sick_and_positive / positives


# --------------------------------------------------------------------------
# Part B: the umbrella world (same numbers as Week 1)
# --------------------------------------------------------------------------
RAIN_PROB = {"rainy": 0.7, "dry": 0.15}
CLOUDY_IF_RAIN, CLOUDY_IF_DRY = 0.6, 0.3


def simulate(days, persistence, rng):
    season = "rainy"
    out = []
    for _ in range(days):
        if rng.random() > persistence:
            season = "dry" if season == "rainy" else "rainy"
        rained = rng.random() < RAIN_PROB[season]
        cloudy = rng.random() < (CLOUDY_IF_RAIN if rained else CLOUDY_IF_DRY)
        out.append((cloudy, rained))
    return out


def predict(belief, persistence):
    """Overnight the season may switch. belief = P(rainy season)."""
    return persistence * belief + (1 - persistence) * (1 - belief)


def update(belief, rained):
    """After seeing whether it rained, apply Bayes' rule to the season."""
    like_rainy = RAIN_PROB["rainy"] if rained else 1 - RAIN_PROB["rainy"]
    like_dry = RAIN_PROB["dry"] if rained else 1 - RAIN_PROB["dry"]
    top = like_rainy * belief
    return top / (top + like_dry * (1 - belief))


def chance_of_rain(belief, cloudy):
    """P(rain today | belief about season, this morning's sky)."""
    prior = belief * RAIN_PROB["rainy"] + (1 - belief) * RAIN_PROB["dry"]
    if cloudy:
        yes, no = CLOUDY_IF_RAIN * prior, CLOUDY_IF_DRY * (1 - prior)
    else:
        yes, no = (1 - CLOUDY_IF_RAIN) * prior, (1 - CLOUDY_IF_DRY) * (1 - prior)
    return yes / (yes + no)


def belief_agent(days, persistence):
    """Keeps one number. Returns accuracy and the belief on each day."""
    belief, right, trace = 0.5, 0, []
    for cloudy, rained in days:
        belief = predict(belief, persistence)        # morning: season may have changed
        take = chance_of_rain(belief, cloudy) > 0.5  # decide
        right += take == rained
        trace.append(belief)
        belief = update(belief, rained)              # evening: learn from today
    return right / len(days), trace


def memory_agent(train_days, test_days, k):
    """The Week 1 learner: counts outcomes for (sky, rainy days in last k)."""
    counts = {}
    for t in range(k, len(train_days)):
        seen = (train_days[t][0], sum(train_days[t - i][1] for i in range(1, k + 1)))
        rain, total = counts.get(seen, (0, 0))
        counts[seen] = (rain + train_days[t][1], total + 1)
    rule = {s: r / n > 0.5 for s, (r, n) in counts.items()}
    right = 0
    for t in range(k, len(test_days)):
        seen = (test_days[t][0], sum(test_days[t - i][1] for i in range(1, k + 1)))
        right += rule.get(seen, False) == test_days[t][1]
    return right / (len(test_days) - k)


if __name__ == "__main__":
    print("Part A: the clinic test")
    print(f"  Bayes' rule:            P(malaria | positive) = {bayes_positive():.3f}")
    print(f"  Counting 1,000,000 sims: P(malaria | positive) = {count_positive(1_000_000, random.Random(1)):.3f}")

    print("\nPart B: umbrella world, seasons last ~50 days (persistence 0.98)")
    rng = random.Random(2026)
    train_days = simulate(20_000, 0.98, rng)
    test_days = simulate(20_000, 0.98, rng)
    acc, trace = belief_agent(test_days, 0.98)
    print(f"  Week 1 learner, k = 7 days of memory: {memory_agent(train_days, test_days, 7):.3f}")
    print(f"  Belief agent, one number of memory:   {acc:.3f}")

    print("\n  First 12 days of the test period:")
    print("  day  rained  belief P(rainy season) that morning")
    for t in range(12):
        print(f"  {t + 1:>3}  {'yes' if test_days[t][1] else 'no ':>6}  {trace[t]:.2f}")