#!/usr/bin/python
"""
Tsetlin Automata team experiment
--------------------------------
We have 5 independent learning automata. Each one can only choose "No" or "Yes".
Together they form a team. The environment looks at how many said "Yes" (called M)
and then rewards or penalizes them. Over many rounds they learn to keep M near 3,
because that is where the reward probability is highest (0.6).
"""

import random

NUM_AUTOMATA = 5          # team size
STATES_PER_ACTION = 3     # states per action (No / Yes)
ITERATIONS = 10000        # learning rounds


class Environment:
    """Gives reward/penalty based only on M = how many automata said Yes."""

    def __init__(self, c_1, c_2):
        """c_1 = 0.2 (step), c_2 = 0.6 (peak reward at M=3)."""
        self.c_1 = c_1
        self.c_2 = c_2

    def give_penalty(self, M):
        """
        Sample reward/penalty for one automaton from M.
          M = 0..3: P(reward) = M * 0.2
          M = 4..5: P(reward) = 0.6 - (M-3)*0.2
        Returns True = penalty, False = reward.
        Best case is M=3 (P=0.6); too few or too many Yes votes get worse odds.
        """
        if M == 0 or M == 1 or M == 2 or M == 3:
            if random.random() <= M * self.c_1:
                return False
            return True
        elif M == 4 or M == 5:
            if random.random() <= self.c_2 - (M - 3) * self.c_1:
                return False
            return True

    def reward_probability(self, M):
        """Same formulas without randomness — used to print the payoff curve."""
        if M == 0 or M == 1 or M == 2 or M == 3:
            return M * self.c_1
        elif M == 4 or M == 5:
            return self.c_2 - (M - 3) * self.c_1


class Tsetlin:
    """
    One automaton with actions No/Yes.
    With n=3: states 1-3 = No, 4-6 = Yes.
    Lower No / higher Yes = more confident. Does not see M, only reward/penalty.
    """

    def __init__(self, n):
        """Start at the weak boundary (state n or n+1)."""
        self.n = n
        self.state = random.choice([self.n, self.n + 1])

    def reward(self):
        """Current action was good — move deeper into it (more confident)."""
        if self.state <= self.n and self.state > 1:
            self.state -= 1
        elif self.state > self.n and self.state < 2 * self.n:
            self.state += 1

    def penalize(self):
        """Current action was bad — move toward the other action (may flip)."""
        if self.state <= self.n:
            self.state += 1
        else:
            self.state -= 1

    def make_decision(self):
        """Return 'No' or 'Yes' from the current state."""
        if self.state <= self.n:
            return "No"
        return "Yes"


# Build environment and create 5 independent automata
env = Environment(0.2, 0.6)
automata = [Tsetlin(STATES_PER_ACTION) for _ in range(NUM_AUTOMATA)]
m_count = [0] * (NUM_AUTOMATA + 1)  # how often each M occurs

# Learning loop: decide -> count Yes (M) -> reward/penalize everyone from M
for _ in range(ITERATIONS):
    actions = [la.make_decision() for la in automata]
    M = actions.count("Yes")
    m_count[M] += 1

    print("M:", M, end=" ")

    # Shared team feedback: probability depends on M, not on this automaton alone
    for j, la in enumerate(automata):
        print(f"TA{j + 1} State: {la.state} Action: {actions[j]}", end=" ")

        if env.give_penalty(M):
            print("Penalty", end=" ")
            la.penalize()
        else:
            print("Reward", end=" ")
            la.reward()

        print("New State:", la.state, end=" ")

    print()

# Results
print("Reward probability P(reward | M):")  # theory: peak at M=3
for m in range(NUM_AUTOMATA + 1):
    print(f"  M={m}: {env.reward_probability(m):.1f}")

print()
print("Final automaton states:")  # usually about 3 Yes and 2 No
for i, ta in enumerate(automata, start=1):
    print(f"  Automaton {i}: state={ta.state} action={ta.make_decision()}")

print()
print(f"Distribution of M over all {ITERATIONS} iterations:")  # M=3 should dominate
for m in range(NUM_AUTOMATA + 1):
    share = 100.0 * m_count[m] / ITERATIONS
    print(f"  M={m}: {m_count[m]:6d}  ({share:6.2f}%)")

print()
average_m = sum(m * m_count[m] for m in range(NUM_AUTOMATA + 1)) / ITERATIONS
print("Average M over all iterations:")  # should be close to 3
print(f"  {average_m:.3f}")

print()
print("Optimum of the payoff function is M=3 (P(reward)=0.6).")