"""
TAI 100, Week 1: Does remembering the weather help you decide about an umbrella?

The world has a hidden "season" (rainy or dry) that you never see directly.
Each morning you see the sky (cloudy or clear) and you remember whether it
rained on each of the last k days. You decide: take the umbrella or not.

A learner counts what happened in 20,000 training days, then is scored on
20,000 new days. We compare memory lengths k under two kinds of world:
  - seasons that last a long time (persistence 0.98, about 50 days)
  - seasons that flip at random every day (persistence 0.5)

Run:  python3 week01_umbrella.py
Only the Python standard library is needed.
"""
import random
from collections import defaultdict

RAIN_PROB = {"rainy": 0.7, "dry": 0.15}      # chance of rain in each season
CLOUDY_IF_RAIN, CLOUDY_IF_DRY = 0.6, 0.3      # how informative the morning sky is


def simulate(days, persistence, rng):
    """Return a list of (cloudy, rained) pairs, one per day."""
    season = "rainy"
    out = []
    for _ in range(days):
        if rng.random() > persistence:          # season sometimes switches
            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 what_agent_sees(days, t, k):
    """The sky today, plus how many of the last k days were rainy."""
    cloudy = days[t][0]
    recent_rain = sum(days[t - i][1] for i in range(1, k + 1))
    return (cloudy, recent_rain)


def train(days, k):
    """Count, for each thing the agent could see, how often it then rained."""
    counts = defaultdict(lambda: [0, 0])          # seen -> [rainy days, total days]
    for t in range(k, len(days)):
        seen = what_agent_sees(days, t, k)
        counts[seen][0] += days[t][1]
        counts[seen][1] += 1
    # Rule: take the umbrella if it rained more than half the time in this situation.
    return {seen: rain / total > 0.5 for seen, (rain, total) in counts.items()}


def score(rule, days, k):
    """Fraction of days on which the umbrella decision was right."""
    right = 0
    for t in range(k, len(days)):
        seen = what_agent_sees(days, t, k)
        take = rule.get(seen, False)              # never seen this before: leave it
        right += take == days[t][1]
    return right / (len(days) - k)


if __name__ == "__main__":
    for persistence in (0.98, 0.5):
        print(f"\nWorld with season persistence {persistence}")
        print("memory k   accuracy on new days")
        for k in (0, 1, 3, 7, 14):
            rng = random.Random(2026)
            train_days = simulate(20_000, persistence, rng)
            test_days = simulate(20_000, persistence, rng)
            rule = train(train_days, k)
            print(f"{k:>8}   {score(rule, test_days, k):.3f}")