# MIT License
# 
# Copyright (c) 2026 Pyxel 100 contributors
# 
# 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.
# 
# The bundled Noto-derived bitmap glyphs are separately licensed under SIL OFL.
# See assets/NOTO-LICENSE.txt.
# 教材画面はASCIIと図形のみ。字形データは使いません。

"""棚からぼたもちの独立教材。日本語の説明はLESSON.mdを参照。"""
from dataclasses import dataclass, replace, asdict
from fractions import Fraction as F
import hashlib
import json

FLOOR = 12

@dataclass(frozen=True)
class Body:
    id: str
    kind: str
    x: int
    y: int
    w: int = 1
    mass: F = F(0)
    capacity: F | None = None
    @property
    def right(self): return self.x + self.w
    @property
    def fixed(self): return self.kind in ('anchor','plate')
    @property
    def removable(self): return self.kind == 'board'

@dataclass(frozen=True)
class State:
    level: int
    bodies: tuple[Body, ...]
    removed: tuple[str, ...] = ()
    broken: tuple[str, ...] = ()
    delivered: tuple[str, ...] = ()
    status: str = 'ready'
    reason: str = 'ready'
    danger: tuple[str, ...] = ()
    moves: int = 0
    def __post_init__(self):
        object.__setattr__(self,'bodies',tuple(sorted(self.bodies,key=lambda b:b.id)))
        for name in ('removed','broken','delivered','danger'):
            object.__setattr__(self,name,tuple(sorted(getattr(self,name))))

@dataclass(frozen=True)
class Frame:
    state: State
    event: str
    ids: tuple[str, ...] = ()

@dataclass(frozen=True)
class Result:
    state: State
    frames: tuple[Frame, ...]

def board(id,x,y,w,capacity=4): return Body(id,'board',x,y,w,F(1),None if capacity is None else F(capacity))

def anchor(id,x,y): return Body(id,'anchor',x,y)

def mochi(id,x,y,mass=1): return Body(id,'mochi',x,y,1,F(mass))

def plate(id,x,y): return Body(id,'plate',x,y)

# 原作Cの固定配置そのもの。level=2は元作品の0始まりの番号。
def initial():
    return State(2, (
        mochi('M1',5,2), mochi('M2',10,2),
        board('A1',4,3,3), board('A2',9,3,3),
        anchor('a1L',4,4), anchor('a1R',6,4),
        anchor('a2L',9,4), anchor('a2R',11,4),
        board('D',3,7,10,None), board('E',2,8,2,F(5,4)),
        board('F',12,8,2,F(5,4)), anchor('eL',2,9),
        anchor('eR',3,9), anchor('fL',12,9), anchor('fR',13,9),
        plate('P1',5,9), plate('P2',10,9)))

def overlap(a,b): return min(a.right,b.right)-max(a.x,b.x) >= 1

def contacts(state):
    """Actual adjacent cell contacts, never authored graph shortcuts."""
    return {b.id:tuple(c.id for c in state.bodies if c.id != b.id and c.y == b.y+1 and overlap(b,c))
            for b in state.bodies if not b.fixed}

def validate_geometry(state):
    for b in state.bodies:
        assert 0<=b.x<b.right<=16 and 0<=b.y<FLOOR,(b,'bounds')
    for i,b in enumerate(state.bodies):
        for c in state.bodies[i+1:]:assert not (b.y==c.y and overlap(b,c)),('overlap',b,c)
    return True

def state_dict(state):
    data=asdict(state)
    for b in data['bodies']:
        for k in ('mass','capacity'):
            if b[k] is not None:b[k]=str(b[k])
    return data

def state_hash(state):
    return hashlib.sha256(json.dumps(state_dict(state),sort_keys=True,ensure_ascii=True,separators=(',',':')).encode('utf-8')).hexdigest()


def support_edges(state):
    by={b.id:b for b in state.bodies}; contact=contacts(state); out={}
    for b in state.bodies:
        if b.fixed: continue
        below=[by[i] for i in contact[b.id]]
        if b.kind=='board':
            left=[c for c in below if c.x<=b.x<c.right]
            right=[c for c in below if c.x<=b.right-1<c.right]
            shares={}
            sides=[side for side in (left,right) if side]
            if sides:
                for side in sides:
                    for c in side: shares[c.id]=shares.get(c.id,F(0))+F(1,len(sides)*len(side))
            elif below:
                shares={c.id:F(1,len(below)) for c in below}
            out[b.id]=shares
        else:
            under=[c for c in below if c.x<=b.x<c.right]
            out[b.id]={c.id:F(1,len(under)) for c in under}
    return out

def loads(state):
    """Return received payload; add own mass only when forwarding downward."""
    edges=support_edges(state); payload={b.id:F(0) for b in state.bodies}
    for b in sorted(state.bodies,key=lambda b:(b.y,b.id)):
        for target,share in edges.get(b.id,{}).items(): payload[target]+=(payload[b.id]+b.mass)*share
    return payload

def example_state(count):
    """原作の落下フレームから取り出した、D上のもち0/1/2個の実座標。"""
    if count not in (0,1,2):
        raise ValueError('count must be 0, 1 or 2')
    removed = ('A1','A2')[:count]
    moved = ('M1','M2')[:count]
    return replace(initial(),
        bodies=tuple(replace(b,y=6) if b.id in moved else b
                     for b in initial().bodies if b.id not in removed),
        removed=removed, moves=count)


def overloaded(state):
    # 必ず同じ旧状態の荷重で全候補を調べる。等しい場合は割れない。
    payload = loads(state)
    return tuple(sorted(b.id for b in state.bodies
        if b.kind == 'board' and b.capacity is not None
        and payload[b.id] > b.capacity))


def break_once(state):
    ids = overloaded(state)
    if not ids:
        return Frame(state,'no_break')
    new = replace(state,
        bodies=tuple(b for b in state.bodies if b.id not in ids),
        broken=tuple(sorted(state.broken+ids)))
    return Frame(new,'break',ids)


def settle_once(state, goals=False):
    """1回の再計算。接触危険→同時破断→同時落下→静止/成功。"""
    by = {b.id:b for b in state.bodies}
    contact = contacts(state)
    # 原作と同じ下側優先。危険と到着が同時なら危険を先に返す。
    for b in sorted(state.bodies,key=lambda b:(-b.y,b.id)):
        if b.fixed:
            continue
        for cid in contact[b.id]:
            c = by[cid]
            reason = ('dish_hit' if c.kind == 'plate' else
                      'crushed' if c.kind == 'mochi' else None)
            if b.kind in ('board','stone') and reason:
                ids = tuple(sorted((b.id,c.id)))
                return Frame(replace(state,status='failure',reason=reason,danger=ids),reason,ids)
        if b.kind == 'mochi' and b.y+1 >= FLOOR:
            return Frame(replace(state,status='failure',reason='floor',danger=(b.id,)),'floor',(b.id,))
    broken = break_once(state)
    if broken.ids:
        return broken
    edges = support_edges(state)
    stable = {b.id for b in state.bodies if b.fixed}
    falling = []
    # 旧状態を下から調べる。リスト中のBodyをこのループで動かさない。
    for b in sorted(state.bodies,key=lambda b:(-b.y,b.id)):
        if b.fixed:
            continue
        if b.y+1 >= FLOOR and b.kind != 'mochi':
            stable.add(b.id)
        elif edges.get(b.id) and any(i in stable for i in edges[b.id]):
            stable.add(b.id)
        else:
            falling.append(b.id)
    if falling:
        new = replace(state,bodies=tuple(replace(b,y=b.y+1) if b.id in falling else b
                                         for b in state.bodies))
        validate_geometry(new)
        return Frame(new,'fall',tuple(sorted(falling)))
    delivered = tuple(sorted(b.id for b in state.bodies if b.kind == 'mochi'
        and any(by[i].kind == 'plate' for i in contact[b.id]))) if goals else ()
    won = goals and len(delivered) == sum(b.kind == 'mochi' for b in state.bodies)
    new = replace(state,delivered=delivered,status='success' if won else 'ready',
                  reason='success' if won else 'ready')
    return Frame(new,'plate' if delivered else 'ready',delivered)


def remove_board(state, action):
    if state.status != 'ready':
        return None
    if not any(b.id == action and b.removable for b in state.bodies):
        return None
    new = replace(state,bodies=tuple(b for b in state.bodies if b.id != action),
                  removed=tuple(sorted(state.removed+(action,))),moves=state.moves+1)
    return Frame(new,'removed',(action,))


def resolve(state):
    frames = []
    for _ in range(128):
        frame = settle_once(state,goals=True)
        frames.append(frame)
        state = frame.state
        if frame.event not in ('break','fall'):
            return Result(state,tuple(frames))
    state = replace(state,status='failure',reason='cycle',danger=())
    return Result(state,tuple(frames)+(Frame(state,'cycle'),))


def transition(state, action):
    removed = remove_board(state,action)
    if removed is None:
        return Result(state,())
    result = resolve(removed.state)
    return Result(result.state,(removed,)+result.frames)


def actions(state):
    return tuple(b.id for b in state.bodies if b.removable) if state.status == 'ready' else ()

STAGE = 5
STAGE_TITLE = 'COMPLETE PUZZLE'
GUIDE = 'Arrows: board / SPACE: remove + resolve / Z: undo'


# 合成入力も実キーボード入力も、この同じ入口へ渡す。
CONTROLS = ('LEFT','RIGHT','SPACE','X','Z','R','1','2','3')
REPLAY = {'L':'LEFT','T':'RIGHT','S':'SPACE','X':'X','Z':'Z','R':'R','1':'1','2':'2','3':'3'}

class Controller:
    def __init__(self):
        self.gate = False
        self.reset()

    def reset(self):
        self.state = initial()
        self.selected = 'A1'
        self.reason = 'ready'
        self.event_ids = ()
        self.pending = False
        self.history = []
        self.last_frames = ()
        # Rでもgateは消さない。全操作キーを離してから再入力。

    def select(self,direction):
        options = self.selectable()
        if options:
            index = options.index(self.selected) if self.selected in options else 0
            self.selected = options[(index+direction)%len(options)]

    def selectable(self):
        return [b.id for b in self.state.bodies if b.removable]

    def tick(self,pressed=(),held=()):
        pressed = set(pressed) & set(CONTROLS)
        held = set(held) & set(CONTROLS)
        if self.gate:
            if not held:
                self.gate = False
            return
        if not pressed and len(held) <= 1:
            return
        if len(pressed) != 1 or pressed != held:
            self.gate = True
            self.reason = 'chord_rejected'
            return
        self.gate = True
        key = next(iter(pressed))
        if key == 'R':
            self.reset()
            return
        if key == 'Z':
            if self.history:
                self.state,self.selected = self.history.pop()
                self.reason = 'undo'
                self.event_ids = ()
                self.last_frames = ()
            else:
                self.reason = 'no_history'
            return
        if self.state.status != 'ready':
            self.reason = self.state.reason
            return
        if key in ('LEFT','RIGHT'):
            self.select(-1 if key == 'LEFT' else 1)
            self.reason = 'select'
        elif key == 'SPACE':
            result = transition(self.state,self.selected)
            if result.frames:
                self.history.append((self.state,self.selected))
                self.state = result.state
                self.last_frames = result.frames
                self.reason = self.state.reason
                self.event_ids = ()
                options = self.selectable()
                if options and self.selected not in options:
                    self.selected = options[0]
            else:
                self.reason = 'unselected'
        else:
            self.reason = 'unused' 

    def snapshot(self):
        return {'state':state_dict(self.state),'hash':state_hash(self.state),
                'selected':self.selected,'event':self.reason,'event_ids':self.event_ids,
                'pending':self.pending,'history_depth':len(self.history),
                'contacts':contacts(self.state),
                'loads':{k:str(v) for k,v in loads(self.state).items()},
                'shares':{k:{a:str(b) for a,b in v.items()} for k,v in support_edges(self.state).items()},
                'last_frames':[{'event':f.event,'ids':f.ids,'hash':state_hash(f.state)}
                               for f in self.last_frames]}


def draw_board(p,c):
    """描画だけ。state・履歴・入力ゲートへ書き込まない。"""
    p.cls(1)
    p.text(12,10,'SHELF MOCHI / STEP '+str(STAGE)+' / '+STAGE_TITLE,10)
    p.text(12,25,'Original case C / exact grid / no rotation or sliding',6)
    ox,oy,cell = 24,52,24
    for y in range(13):
        p.text(9,oy+y*cell+8,str(y),13)
        p.line(ox,oy+y*cell,ox+16*cell,oy+y*cell,5)
    for x in range(17):
        p.line(ox+x*cell,oy,ox+x*cell,oy+12*cell,5)
        if x < 16:
            p.text(ox+x*cell+8,oy-12,str(x),13)
    state = c.state
    by = {b.id:b for b in state.bodies}
    for b in sorted(state.bodies,key=lambda b:(-b.y,b.id)):
        x,y,w = ox+b.x*cell,oy+b.y*cell,b.w*cell
        # 点線枠が占有する格子。図形の余白や皿の薄さは接触判定を変えない。
        p.rectb(x,y,w,cell,13)
        if b.kind == 'anchor':
            p.rect(x+2,y+2,w-4,cell-4,5)
            p.line(x+3,y+3,x+w-4,y+cell-4,6)
            p.text(x+3,y+8,b.id,7)
        elif b.kind == 'board':
            p.rect(x+1,y+2,w-2,cell-4,9 if b.id not in ('E','F') else 8)
            p.line(x+3,y+5,x+w-4,y+5,10)
            p.text(x+w//2-len(b.id)*2,y+10,b.id,1)
        elif b.kind == 'plate':
            p.rect(x+2,y+1,w-4,4,7)
            p.text(x+7,y+10,b.id,6)
        elif b.kind == 'mochi':
            p.rect(x+4,y+3,w-8,cell-6,15)
            p.pset(x+8,y+8,1);p.pset(x+15,y+8,1)
            p.text(x+7,y+13,b.id,1)
        else:
            p.rect(x+2,y+2,w-4,cell-4,13)
            p.text(x+8,y+8,b.id,1)
    selected = by.get(c.selected)
    if selected:
        p.rectb(ox+selected.x*cell-1,oy+selected.y*cell-1,selected.w*cell+2,cell+2,10)
    p.rect(417,40,211,319,0)
    p.text(429,51,'INSPECT '+c.selected,10)
    if selected:
        p.text(429,68,f'x={selected.x} y={selected.y} w={selected.w}',7)
        p.text(429,85,'kind: '+selected.kind,7)
        p.text(429,102,'CONTACTS:',6)
        contact = contacts(state).get(selected.id,())
        p.text(429,115,','.join(contact) or '(none / fixed)',7)
        payload = loads(state)
        p.text(429,134,'PAYLOAD: '+str(payload[selected.id]),10)
        p.text(429,150,'SELF MASS: '+str(selected.mass),7)
        p.text(429,166,'CAPACITY: '+('UNLIMITED' if selected.capacity is None else str(selected.capacity)),7)
        p.text(429,184,'OUTGOING SHARE:',6)
        for row,(target,share) in enumerate(support_edges(state).get(selected.id,{}).items()):
            p.text(429,199+row*12,target+': '+str(share),7)
    p.text(429,241,'REMOVED: '+','.join(state.removed),6)
    p.text(429,257,'BROKEN: '+','.join(state.broken),8)
    p.text(429,273,'DELIVERED: '+','.join(state.delivered),11)
    p.text(429,291,'MOVES: '+str(state.moves),7)
    p.text(429,307,'STATUS: '+state.status.upper(),10)
    p.text(429,324,'HISTORY: '+str(len(c.history)),7)
    p.text(429,340,'PENDING: '+('YES - USE X' if c.pending else 'NO'),10 if c.pending else 13)
    p.text(12,372,'EVENT: '+c.reason.upper()+' '+','.join(c.event_ids),10)
    p.text(12,389,GUIDE,7)
    p.text(12,406,'R: reset  ESC: quit / release ALL keys before next input',6)
    p.text(12,424,'Fictional non-rotating model. NOT a real structural safety calculation.',13)
    p.text(12,440,'Cell outlines = collision cells; artwork is only decoration. Silent.',13)


def main(headless=False,sequence='',screenshot=None,report=None):
    from pathlib import Path
    import time
    import pyxel  # import時には不要。実画面を開く時だけ読み込む。
    controller = Controller()
    keycodes = {name:getattr(pyxel,'KEY_'+name) for name in CONTROLS}
    frame = 0
    events = []
    last_frame = max(12,len(sequence)*2+2)
    start = time.perf_counter()
    def update():
        nonlocal frame
        frame += 1
        if headless:
            index = (frame-1)//2
            if frame % 2 and index < len(sequence):
                key = REPLAY[sequence[index]]
                controller.tick((key,),(key,))
                events.append({'input':sequence[index],'snapshot':controller.snapshot()})
            else:
                controller.tick()
        else:
            controller.tick([n for n,k in keycodes.items() if pyxel.btnp(k)],
                            [n for n,k in keycodes.items() if pyxel.btn(k)])
    def draw():
        draw_board(pyxel,controller)
        if headless and frame >= last_frame:
            if screenshot:
                pyxel.screenshot(str(Path(screenshot).resolve()),scale=1)
            if report:
                Path(report).write_text(json.dumps({
                    'label':'actual Pyxel headless / synthetic controller input / silent',
                    'stage':STAGE,'frames':frame,'logical_seconds':frame/30,
                    'program_sha256':hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
                    'wall_seconds_before_exit':round(time.perf_counter()-start,6),
                    'sequence':sequence,'events':events,'snapshot':controller.snapshot()
                },ensure_ascii=False,indent=2)+'\n',encoding='utf-8')
            pyxel.quit()
    pyxel.init(640,460,title='Shelf Mochi - Step '+str(STAGE),fps=30,headless=headless)
    pyxel.run(update,draw)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--headless-smoke',action='store_true')
    parser.add_argument('--headless-replay',help='Synthetic keys: '+str(REPLAY))
    parser.add_argument('--screenshot')
    parser.add_argument('--report')
    args = parser.parse_args()
    if args.headless_smoke and args.headless_replay is not None:
        parser.error('Choose smoke or replay')
    headless = args.headless_smoke or args.headless_replay is not None
    if not headless and (args.screenshot or args.report):
        parser.error('Outputs require headless mode')
    sequence = args.headless_replay or ''
    if any(char not in REPLAY for char in sequence):
        parser.error('Unsupported replay character')
    main(headless,sequence,args.screenshot,args.report)