# MIT License
# 
# Copyright (c) 2026
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# 
# Third-party exception: bitmap glyph data derived from Noto Sans CJK is licensed
# under the SIL Open Font License 1.1. See assets/NOTO-LICENSE.txt.

"""076 紙の水族館: 教材ステージ 1。架空の整数生態系。"""
STAGE = 1

from dataclasses import dataclass, replace

TICKS = 8
D_CAP = 12
FLOW = (1, 2, 1)

def _int(value):
    return isinstance(value, int) and not isinstance(value, bool)

def _vector(value, length, lo, hi):
    return (isinstance(value, tuple) and len(value) == length
            and all(_int(v) and lo <= v <= hi for v in value))

def _shelters(value):
    if not _vector(value, 3, 0, 2) or sum(value) != 2:
        raise ValueError("隠れ家は整数の3区画へ合計2チップ置いてください")

def drift(d, shelters):
    """Return (new counts, permitted (source,destination) moves, blocked moves).

    Receiving slots are frozen before any departure and reserved in L/M/R
    source order. This intentionally cannot reuse outgoing capacity this tick.
    """
    if not _vector(d, 3, 0, D_CAP):
        raise ValueError("daphnia must be a bounded integer triple")
    _shelters(shelters)
    slots = [D_CAP - value for value in d]
    moves, blocked = [], []
    for source, dest in enumerate(FLOW):
        if d[source] > 2 * shelters[source]:
            if slots[dest] > 0:
                moves.append((source, dest))
                slots[dest] -= 1
            else:
                blocked.append((source, dest))
    result = list(d)
    for source, dest in moves:
        result[source] -= 1
        result[dest] += 1
    return tuple(result), tuple(moves), tuple(blocked)

@dataclass(frozen=True)
class Lab:
    d: tuple = (5, 0, 3)
    shelters: tuple = (1, 0, 1)
    t: int = 0
    moves: tuple = ()
    blocked: tuple = ()

    @property
    def ended(self):
        return self.t == TICKS


def initial():
    return Lab()


def advance(state):
    if state.ended:
        return state
    d, moves, blocked = drift(state.d, state.shelters)
    return replace(state, d=d, t=state.t + 1, moves=moves, blocked=blocked)


# Teaching interface: independent of the original Controller.
from dataclasses import asdict

class App:
    def __init__(self):
        self.level = 'A'
        self.algae_enabled = True
        self.gated = False
        self.reset()

    def reset(self):
        self.state = initial(self.level, self.algae_enabled) if STAGE >= 3 else initial()
        self.cursor = 0
        self.amount = 0
        self.message = 'READY / SELECTING DOES NOT ADVANCE TIME'

    def press(self, key):
        if key == 'R':
            self.reset()
            return
        if STAGE >= 3 and key in ('1', '2', '3'):
            self.level = 'ABC'[int(key) - 1]
            self.reset()
            return
        if STAGE >= 4 and key == 'B':
            self.algae_enabled = not self.algae_enabled
            self.reset()
            return
        if STAGE == 5 and key == 'X':
            self.state = undo(self.state)
            self.message = 'UNDO / REPLAYED EARLIER INPUTS'
            return
        if STAGE >= 3 and key == 'TAB':
            if self.state.t:
                self.message = 'SHELTERS LOCKED AFTER TICK 0'
            else:
                index = SHELTER_LAYOUTS.index(self.state.shelters)
                self.state = edit_shelters(self.state, SHELTER_LAYOUTS[(index + 1) % len(SHELTER_LAYOUTS)])
                self.message = 'TWO SHELTER CHIPS / TAB CYCLES SIX LAYOUTS'
            return
        if key in ('LEFT', 'RIGHT'):
            self.cursor = (self.cursor + (1 if key == 'RIGHT' else -1)) % 3
        elif STAGE >= 3 and key in ('UP', 'DOWN'):
            self.amount = max(0, min(2, self.amount + (1 if key == 'UP' else -1)))
        elif key in ('SPACE', 'RETURN'):
            if self.state.ended:
                self.message = 'EIGHT TICKS FINISHED / R RESTARTS'
                return
            if STAGE < 3:
                self.state = advance(self.state)
            else:
                amount = self.amount if key == 'RETURN' else 0
                try:
                    self.state = step(self.state, self.cursor, amount)
                except ValueError:
                    self.message = 'NOT ENOUGH FOOD / TICK DID NOT ADVANCE'
                    return
            self.message = 'ONE TICK COMMITTED / RELEASE BEFORE NEXT PRESS'

    def control(self, pressed, held):
        if self.gated:
            if not pressed and not held:
                self.gated = False
            return
        if len(pressed | held) > 1:
            self.gated = True
            self.message = 'CHORD REJECTED / RELEASE ALL KEYS'
            return
        if len(pressed) == 1:
            self.press(next(iter(pressed)))
            self.gated = True

    def prediction(self):
        if STAGE != 5 or self.state.ended:
            return None
        try:
            return forecast(self.state, self.cursor, self.amount)
        except ValueError:
            return None

    def draw(self, p):
        s = self.state
        p.cls(1)
        p.text(12, 12, 'PAPER AQUARIUM / STAGE ' + str(STAGE) + ' / ' + ('LAB' if STAGE < 3 else 'LEVEL ' + self.level), 10)
        p.text(12, 27, 'FICTIONAL INTEGER ECOSYSTEM / NOT AQUARIUM CARE ADVICE', 6)
        p.text(12, 45, 'TICK ' + str(s.t) + '/8  FLOW: L>M, M>R, R>M', 7)
        if STAGE >= 3:
            p.text(320, 45, 'FOOD LEFT ' + str(s.remaining_food) + '/10  MEALS ' + str(s.meals), 10)
        for z in range(3):
            x = 12 + z * 196
            p.rect(x, 68, 188, 216, 5)
            p.rect(x + 2, 70, 184, 203, 1)
            p.rectb(x, 68, 188, 216, 10 if z == self.cursor else 6)
            p.text(x + 8, 79, 'LMR'[z] + ' -> ' + 'MRM'[z], 10)
            p.text(x + 66, 79, 'D ' + str(s.d[z]) + ' / CAP 12', 11)
            for chip in range(s.shelters[z]):
                sx = x + 49 + chip * 72
                p.tri(sx - 21, 267, sx, 237, sx + 21, 267, 13)
                p.tri(sx - 10, 267, sx, 254, sx + 10, 267, 1)
            for i in range(s.d[z]):
                dx, dy = x + 21 + (i % 6) * 28, 194 + (i // 6) * 22
                p.circ(dx, dy, 3, 11)
                p.line(dx - 2, dy - 2, dx - 4, dy - 6, 7)
            if STAGE >= 2:
                for f in s.fish:
                    if f.zone != z:
                        continue
                    fx, fy = x + 46 + (f.id - 1) * 70, 142 + (f.id - 1) * 17
                    color = 10 if f.id == 1 else 14
                    p.tri(fx + 10, fy, fx - 8, fy - 8, fx - 8, fy + 8, color)
                    p.tri(fx - 7, fy, fx - 19, fy - 7, fx - 19, fy + 7, color)
                    p.pset(fx + 4, fy - 2, 0)
                    p.text(fx - 4, fy - 19, 'F' + str(f.id), color)
            if STAGE >= 3:
                particles = [f for f in s.food if f.zone == z]
                for i, f in enumerate(particles):
                    p.rect(x + 15 + (i % 8) * 20, 101 + (i // 8) * 9, 4, 4, 9 if f.age else 10)
                p.text(x + 8, 290, 'SHELTER ' + str(s.shelters[z]) + '  ALGAE ' + str(s.algae[z]), 6)
                p.text(x + 8, 302, 'FOOD ' + str(len(particles)) + '  PROTECTED ' + str(2 * s.shelters[z]), 7)
            else:
                p.text(x + 8, 290, 'SHELTER ' + str(s.shelters[z]) + '  PROTECTED ' + str(2 * s.shelters[z]), 6)
        if STAGE < 3:
            p.text(12, 325, 'DRIFT MOVES ' + str(s.moves) + ' / BLOCKED ' + str(s.blocked), 7)
            p.text(12, 341, 'CONSERVES TOTAL D = ' + str(sum(s.d)) + ' / NO MEALS OR BIRTHS', 11)
            if STAGE == 2:
                p.text(12, 357, 'POST-DRIFT SCORES ' + str(s.scores) + ' / BOTH USE OLD FISH POSITIONS' if s.t else 'NO FISH PLAN YET / SPACE TO COMPUTE THE FIRST TICK', 7)
        else:
            p.text(12, 325, 'NEXT INPUT: ' + 'LMR'[self.cursor] + ' / ' + str(self.amount) + ' PARTICLES', 10)
            if s.history:
                last = s.history[-1]
                p.text(12, 341, 'SCORES ' + str(last.scores) + ' / MEAL ORDER ' + str(last.meal_order), 7)
                kinds = ','.join(e.kind for e in s.events if e.kind in ('algae','birth','sink','eat_d','algae_full'))
                p.text(12, 357, 'EVENTS: ' + (kinds or 'feed / drift / meals')[:84], 11)
            if STAGE == 3:
                p.text(12, 377, 'STAGE 03: FOOD DOES NOT AGE / NO ALGAE CONVERSION OR BIRTHS', 6)
            else:
                p.text(12, 377, 'FOOD -> ALGAE: ' + ('ON' if s.algae_enabled else 'OFF') + ' / B RESTARTS COMPARISON', 6)
            if STAGE == 5:
                pred = self.prediction()
                if pred:
                    p.text(12, 396, 'NEXT TICK ONLY: FISH ' + str(pred.fish_positions) + ' MEALS ' + str(pred.meals) + ' D ' + str(pred.d), 10)
                else:
                    p.text(12, 396, 'NO NEXT FORECAST: FINISHED OR INSUFFICIENT FOOD', 8)
            if s.ended:
                p.text(12, 415, 'SUCCESS: BOTH MEALS >=3 AND TOTAL D >=2' if success(s) else 'FINISHED: CHECK MEALS >=3 EACH AND TOTAL D >=2', 11 if success(s) else 8)
        p.text(12, 440, self.message, 10)
        p.text(12, 461, 'ENTER COMMIT / SPACE WAIT / R RESET / ESC QUIT / RELEASE EACH KEY', 7)
        if STAGE >= 3:
            p.text(12, 477, 'LEFT/RIGHT ZONE / UP/DOWN 0-2 FOOD / TAB SHELTERS AT T0 / 1-3 LEVEL', 6)
        if STAGE == 5:
            p.text(12, 493, 'X UNDO BY REPLAY / R DOES NOT RELOAD CODE', 6)

def demo_keys():
    if STAGE < 3:
        return ('SPACE',)
    return ('RIGHT', 'UP', 'UP', 'RETURN', 'RIGHT', 'RETURN', 'LEFT', 'LEFT', 'RETURN', 'SPACE')

def main(headless=False, screenshot=None, state_json=None, demo=False):
    import pyxel as p
    import json
    from pathlib import Path
    p.init(612, 516, title='Paper aquarium lesson ' + str(STAGE), fps=30, display_scale=1, headless=headless)
    app = App()
    sequence = demo_keys() if demo else ()
    keynames = ('LEFT','RIGHT','UP','DOWN','RETURN','SPACE','TAB','X','R','1','2','3','B')
    keys = {k: getattr(p, 'KEY_' + k) for k in keynames}
    limit = max(8, len(sequence) * 2 + 6)
    def update():
        if p.btnp(p.KEY_ESCAPE) or headless and p.frame_count >= limit:
            if screenshot:
                app.draw(p)
                p.screen.save(screenshot, 1)
            if state_json:
                Path(state_json).write_text(json.dumps({'stage': STAGE, 'state': asdict(app.state), 'scope': 'Real Pyxel; synthetic controls, not physical keyboard or human playtest'}, ensure_ascii=False, indent=2) + '\n', encoding='utf-8')
            p.quit()
        if headless:
            i = p.frame_count // 2
            pressed = {sequence[i]} if demo and p.frame_count % 2 == 0 and i < len(sequence) else set()
            app.control(pressed, pressed)
        else:
            app.control({k for k,v in keys.items() if p.btnp(v)}, {k for k,v in keys.items() if p.btn(v)})
    p.run(update, lambda: app.draw(p))

if __name__ == '__main__':
    import argparse
    from pathlib import Path
    parser = argparse.ArgumentParser(description='Standalone paper aquarium teaching stage')
    parser.add_argument('--headless-smoke', action='store_true')
    parser.add_argument('--demo-checkpoint', action='store_true')
    parser.add_argument('--screenshot')
    parser.add_argument('--state-json')
    args = parser.parse_args()
    if (args.demo_checkpoint or args.screenshot or args.state_json) and not args.headless_smoke:
        parser.error('demo and outputs require --headless-smoke')
    # Freeze the caller cwd BEFORE importing or initializing Pyxel.
    for field, suffix in (('screenshot', '.png'), ('state_json', '.json')):
        value = getattr(args, field)
        if value:
            target = Path(value).resolve()
            if target.suffix.lower() != suffix or not target.parent.is_dir() or target.exists():
                parser.error('output must be a new ' + suffix + ' in an existing directory')
            setattr(args, field, str(target))
    main(args.headless_smoke, args.screenshot, args.state_json, args.demo_checkpoint)