"""Does the AI model follow pond.py's rules, whatever order the moves come in?

    python wording.py              the wording pond.py has now
    python wording.py first        the first wording we wrote, which failed

pond.py tells a hunter one of four things about your fish. The Python rules in
pick() give the right move for each. This asks the model the same four things, as
each of the three kinds of hunter, with the three moves in all six orders: 72
questions. It prints the share the model gave the right move, and every wrong answer.

"first" is the wording of our first try. There a bigger fish was also "near" or
"far away": five things, 90 questions. The model followed those two words and did not
flee. It is kept here so you can see it fail.

It needs what the model game needs: the folder "model" beside pond.py, and
    pip install "decider-ai[gguf]"
The wording of RULES and MOVES was chosen with this test, so a clean result is a
fit and not a score.
"""
import itertools
import pathlib
import sys
import time

ASSETS = pathlib.Path(__file__).resolve().parent.parent      # pond.py: one folder up from checks/

# hidden, smaller, near: every mix that pond.py's facts() can give.
FACTS = [(True, True, True), (True, True, False), (True, False, True), (True, False, False),
         (False, True, True), (False, True, False), (False, False, True), (False, False, False)]

FIRST_MOVES = {
    "hunt": "Swim to the gold fish and eat it. Do this when it is smaller and near.",
    "flee": "Swim away from the gold fish. Do this when it is bigger than you.",
    "wander": "Swim around slowly. Do this when it is far away, or you cannot see it.",
}


def first_words(hidden, smaller, near):
    """What a hunter saw in our first try: a bigger fish was also near or far away."""
    if hidden:
        return "You cannot see the gold fish."
    sees = "The gold fish is bigger than you."
    if smaller:
        sees = "The gold fish is smaller than you."
    if near:
        return sees + " It is near."
    return sees + " It is far away."


def main():
    sys.path.insert(0, str(ASSETS))
    import pond
    started = time.perf_counter()
    from decider.infer import Decider
    import_seconds = time.perf_counter() - started
    first = sys.argv[1:] == ["first"]
    in_words = pond.in_words
    moves = dict(pond.MOVES)
    if first:
        in_words = first_words
        moves = dict(FIRST_MOVES)
    # Two mixes of facts can give the same sentence. Ask about each sentence once.
    right_for = {}
    for hidden, smaller, near in FACTS:
        right_for[in_words(hidden, smaller, near)] = pond.pick(hidden, smaller, near)
    started = time.perf_counter()
    model = Decider(pond.model_file(), gguf_options={"n_ctx": 2048})
    print(f"The library took {import_seconds:.1f} s to import, and the model {time.perf_counter() - started:.1f} s to load.")
    asked = wrong = calls = 0
    tokens = []
    lowest = 1.0
    started = time.perf_counter()
    for kind in ("fish", "snake", "lizard"):
        for sees, right in right_for.items():
            shares = []
            for order in itertools.permutations(moves):
                pond.MOVES = {name: moves[name] for name in order}      # ask() reads this
                who = "You are a hungry " + kind + " in a pond. A gold fish swims in the pond too."
                question = {"type": "choice", "instructions": "What should you do right now?",
                            "criteria": pond.MOVES}
                state = who + " Rules: " + pond.RULES + "\n\nSituation: " + sees
                result = model.system_one(state, {"move": question})
                answer = result["answers"]["move"]
                tokens.append(result["usage"]["input_tokens"])
                asked += 1
                calls += 1
                shares.append(answer["probabilities"][right])
                if answer["choice"] != right:
                    wrong += 1
                    print(f"   WRONG  {kind:6} moves in the order {', '.join(order)}: "
                          f"the model said {answer['choice']} {answer['probabilities']}")
                    continue
                calls += 1
                if pond.ask(model, kind, sees) != right:
                    wrong += 1
                    print(f"   WRONG  {kind:6} pond.ask() gave another move than the same question here")
            lowest = min(lowest, min(shares))
            print(f"{kind:6} {right:6} {min(shares):.3f} to {max(shares):.3f}   {sees}")
    took = time.perf_counter() - started
    pond.MOVES = moves
    print(f"{asked} questions, {wrong} wrong. The lowest share of a right move: {lowest:.3f}.")
    print(f"A question is {min(tokens)} to {max(tokens)} tokens long.")
    print(f"About {took / calls:.1f} s for one question on this computer.")
    return 1 if wrong else 0


if __name__ == "__main__":
    sys.exit(main())