"""
TAI 100, Week 3: Information theory.

Part A  Surprise and entropy for coins, dice and the weather.
Part B  Mutual information: how many bits does each clue carry about rain?
        Explains the Week 1 results for both umbrella worlds.
Part C  Scoring forecasters in bits (cross-entropy).
Part D  Compression: real compressors against the limit set by the best
        predictor.

Run:  python3 week03_information.py
Only the Python standard library is needed.
"""
import lzma
import math
import random
import zlib
from collections import Counter

# Same umbrella world as Weeks 1 and 2
RAIN_PROB = {"rainy": 0.7, "dry": 0.15}
CLOUDY_IF_RAIN, CLOUDY_IF_DRY = 0.6, 0.3


def surprise(p):
    """Bits of surprise when something of probability p happens."""
    return math.log2(1 / p)


def entropy(probs):
    """Average surprise, in bits, of a distribution given as a list."""
    return sum(p * surprise(p) for p in probs if p > 0)


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 mutual_information(pairs):
    """I(X; Y) in bits, estimated by counting (x, y) pairs."""
    n = len(pairs)
    joint = Counter(pairs)
    px = Counter(x for x, _ in pairs)
    py = Counter(y for _, y in pairs)
    return sum(c / n * math.log2(c * n / (px[x] * py[y])) for (x, y), c in joint.items())


def clue(days, t, k):
    """What the Week 1 agent saw: today's sky and rainy days in the last k."""
    return (days[t][0], sum(days[t - i][1] for i in range(1, k + 1)))


def info_in_clues(days, k):
    """Bits the Week 1 clue (sky + last k days) carries about today's rain."""
    return mutual_information([(clue(days, t, k), days[t][1]) for t in range(14, len(days))])


# ---- forecasters: each returns P(rain today) every morning -----------------
def always_average(days):
    return [0.425 for _ in days]


def sky_only(days):
    return [0.6 if cloudy else 0.3 for cloudy, _ in days]


def belief_forecasts(days, persistence, use_sky):
    """The Week 2 belief agent, now giving a probability instead of a yes/no."""
    belief, out = 0.5, []
    for cloudy, rained in days:
        belief = persistence * belief + (1 - persistence) * (1 - belief)
        p = belief * RAIN_PROB["rainy"] + (1 - belief) * RAIN_PROB["dry"]
        if use_sky:
            a = (CLOUDY_IF_RAIN if cloudy else 1 - CLOUDY_IF_RAIN) * p
            b = (CLOUDY_IF_DRY if cloudy else 1 - CLOUDY_IF_DRY) * (1 - p)
            p = a / (a + b)
        out.append(p)
        like_r = RAIN_PROB["rainy"] if rained else 1 - RAIN_PROB["rainy"]
        like_d = RAIN_PROB["dry"] if rained else 1 - RAIN_PROB["dry"]
        belief = like_r * belief / (like_r * belief + like_d * (1 - belief))
    return out


def bits_per_day(forecasts, days):
    """Cross-entropy: average surprise at what actually happened."""
    return sum(surprise(p if rained else 1 - p) for p, (_, rained) in zip(forecasts, days)) / len(days)


def compressed_bits_per_day(days):
    """Write the rain record as a string of 0s and 1s and compress it."""
    text = "".join("1" if rained else "0" for _, rained in days).encode()
    return {"zlib": 8 * len(zlib.compress(text, 9)) / len(days),
            "lzma": 8 * len(lzma.compress(text, preset=9)) / len(days)}


if __name__ == "__main__":
    print("Part A: surprise and entropy (bits)")
    print(f"  surprise of a fair coin landing heads: {surprise(0.5):.3f}")
    print(f"  surprise of rolling a six:             {surprise(1/6):.3f}")
    print(f"  entropy of a fair coin:                {entropy([0.5, 0.5]):.3f}")
    print(f"  entropy of a fair die:                 {entropy([1/6] * 6):.3f}")
    print(f"  entropy of a 90/10 coin:               {entropy([0.9, 0.1]):.3f}")
    print(f"  entropy of rain (p = 0.425):           {entropy([0.425, 0.575]):.3f}")
    print("  entropy of a coin with P(heads) = p:")
    for p in (0.5, 0.6, 0.7, 0.8, 0.9, 0.95, 0.99, 1.0):
        print(f"    p = {p:<4}  {entropy([p, 1 - p]):.3f}")
    # exact mutual information between sky and rain, from Week 2's Table 1
    joint = {(1, 1): 0.255, (1, 0): 0.1725, (0, 1): 0.17, (0, 0): 0.4025}
    h_rain = entropy([0.425, 0.575])
    h_rain_given_sky = (0.4275 * entropy([0.255 / 0.4275, 0.1725 / 0.4275])
                        + 0.5725 * entropy([0.17 / 0.5725, 0.4025 / 0.5725]))
    print(f"  H(rain | sky) = {h_rain_given_sky:.3f},  I(rain; sky) = {h_rain - h_rain_given_sky:.3f}")

    print("\nPart B: bits each clue carries about today's rain (200,000 days)")
    rng = random.Random(2026)
    worlds = {0.98: simulate(200_000, 0.98, rng), 0.5: simulate(200_000, 0.5, rng)}
    print("  k   long seasons (0.98)   random seasons (0.5)")
    for k in (0, 1, 3, 7, 14):
        print(f"  {k:>2}  {info_in_clues(worlds[0.98], k):>10.3f}           {info_in_clues(worlds[0.5], k):>10.3f}")

    print("\nPart C: forecasters scored in bits per day (long seasons, lower is better)")
    days = worlds[0.98]
    print(f"  always say 0.425:          {bits_per_day(always_average(days), days):.3f}")
    print(f"  sky only:                  {bits_per_day(sky_only(days), days):.3f}")
    print(f"  belief agent, no sky:      {bits_per_day(belief_forecasts(days, 0.98, False), days):.3f}")
    print(f"  belief agent with sky:     {bits_per_day(belief_forecasts(days, 0.98, True), days):.3f}")
    print(f"  overconfident (0.99/0.01): {bits_per_day([0.99 if c else 0.01 for c, _ in days], days):.3f}")

    print("\nPart D: compressing the rain record (bits per day)")
    for pers, d in worlds.items():
        c = compressed_bits_per_day(d)
        best = bits_per_day(belief_forecasts(d, pers, False), d)
        print(f"  persistence {pers}: uncompressed 8.000, zlib {c['zlib']:.3f}, lzma {c['lzma']:.3f}, "
              f"best predictor {best:.3f}")