# 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.

"""帰り道の燃料・教材5/5。Python 3.11+ / Pyxel 2.9.9。"""
from dataclasses import dataclass, replace, asdict
from collections import deque
from functools import lru_cache
import hashlib, json

@dataclass(frozen=True)
class Item:
    id: str
    place: str
    weight: int
    value: int

@dataclass(frozen=True)
class Level:
    title: str
    nodes: tuple
    edges: tuple
    items: tuple
    capacity: int
    initial_fuel: int
    supply: tuple
    goal: int
SQUARE = (('H', 'A', 1), ('A', 'B', 1), ('B', 'C', 1), ('C', 'H', 1))
LEVELS = (Level('同じ一周、違う帰り', ('H', 'A', 'B', 'C'), SQUARE, (Item('a', 'A', 1, 2), Item('b', 'B', 2, 4), Item('c', 'C', 1, 3)), 12, 12, (0, 0, 0, 0), 4), Level('補給の前後', ('H', 'A', 'B', 'C'), SQUARE, (Item('a', 'A', 2, 4), Item('b', 'B', 1, 3), Item('c', 'C', 2, 5)), 9, 9, (0, 0, 3, 0), 7), Level('重い主部品は最後に', ('H', 'A', 'B', 'C', 'D'), (('H', 'A', 1), ('A', 'B', 2), ('B', 'H', 2), ('H', 'C', 2), ('C', 'D', 1), ('D', 'B', 1)), (Item('a', 'A', 1, 2), Item('b', 'B', 3, 6), Item('c', 'C', 1, 3), Item('d', 'D', 1, 3)), 12, 12, (0, 0, 0, 0, 4), 6))

@dataclass(frozen=True)
class State:
    place: str
    fuel: int
    locations: tuple
    supply: tuple
    departed: bool = False
    spent: int = 0
    outcome: str = 'running'

@dataclass(frozen=True)
class Timeline:
    state: State
    history: tuple = ()
REASONS = {'ready': '遠征は一回。基地Hへ戻ると終了です。', 'ok': '予告どおり確定しました。', 'fuel': '燃料不足。置く・補給・別の道を検討。', 'full': '荷重4を超えます。拾う前に置けます。', 'absent': 'この品は現在地にも荷台にもありません。', 'road': 'この二地点を結ぶ道はありません。', 'supply': 'タンク満杯か、この地点の補給が空です。', 'terminal': '帰還済み。Xで戻すかRで再開。', 'invalid': 'その操作はありません。', 'undo': '一つ前の操作へ戻りました。', 'no_undo': 'まだ戻せる操作はありません。'}

def initial(l):
    return State('H', l.initial_fuel, tuple((i.place for i in l.items)), l.supply)

def weight(s, l):
    return sum((i.weight for i, p in zip(l.items, s.locations) if p == 'T'))

def value(s, l):
    return sum((i.value for i, p in zip(l.items, s.locations) if p == 'T'))

def distance(l, a, b):
    return next((d for u, v, d in l.edges if (a, b) in ((u, v), (v, u))), None)

def actions(l):
    return tuple((('move', n) for n in l.nodes)) + tuple((('item', i) for i in range(len(l.items)))) + (('refuel',),)

def travel_cost(s, l, d, flat=False):
    return d * (1 if flat else 1 + weight(s, l))

def apply(s, l, a, flat=False):
    if s.outcome != 'running':
        return (s, 'terminal')
    if a not in actions(l):
        return (s, 'invalid')
    k = a[0]
    n = s
    if k == 'move':
        dst = a[1]
        d = distance(l, s.place, dst)
        if d is None:
            return (s, 'road')
        cost = travel_cost(s, l, d, flat)
        if s.fuel < cost:
            return (s, 'fuel')
        n = replace(s, place=dst, fuel=s.fuel - cost, spent=s.spent + cost, departed=True)
        if dst == 'H':
            n = replace(n, outcome='won' if value(n, l) >= l.goal else 'short')
    elif k == 'item':
        i = a[1]
        p = s.locations[i]
        if p not in ('T', s.place):
            return (s, 'absent')
        if p != 'T' and weight(s, l) + l.items[i].weight > 4:
            return (s, 'full')
        loc = list(s.locations)
        loc[i] = s.place if p == 'T' else 'T'
        n = replace(s, locations=tuple(loc))
    elif k == 'refuel':
        i = l.nodes.index(s.place)
        amount = min(l.capacity - s.fuel, s.supply[i])
        if amount <= 0:
            return (s, 'supply')
        sup = list(s.supply)
        sup[i] -= amount
        n = replace(s, fuel=s.fuel + amount, supply=tuple(sup))
    check(n, l)
    return (n, 'ok')

def step(tl, l, a, flat=False):
    if a == ('undo',):
        return (Timeline(tl.history[-1], tl.history[:-1]), 'undo') if tl.history else (tl, 'no_undo')
    n, r = apply(tl.state, l, a, flat)
    return (Timeline(n, tl.history + (tl.state,)) if n != tl.state else tl, r)

def check(s, l):
    assert 0 <= s.fuel <= l.capacity and 0 <= weight(s, l) <= 4 and (s.spent >= 0)
    assert s.fuel + sum(s.supply) + s.spent == l.initial_fuel + sum(l.supply)
    assert len(s.locations) == len(l.items) and all((p in l.nodes + ('T',) for p in s.locations))
    assert all((0 <= a <= b for a, b in zip(s.supply, l.supply)))
    assert s.outcome == 'running' or (s.place == 'H' and s.departed)

@lru_cache(None)
def shortest_home(l, start):
    dist = {n: 999 for n in l.nodes}
    dist[start] = 0
    q = [start]
    while q:
        u = min(q, key=lambda n: dist[n])
        q.remove(u)
        for v in l.nodes:
            d = distance(l, u, v)
            if d is not None and dist[v] > dist[u] + d:
                dist[v] = dist[u] + d
                q.append(v)
    return dist['H']

def home_cost(s, l, flat=False):
    return travel_cost(s, l, shortest_home(l, s.place), flat)

@lru_cache(4096)
def return_plan(s, l, flat=False, allow_drop=False):
    """Exact reachable return with current cargo, finite refuels and optional drops."""
    if s.place == 'H':
        return ()
    q = deque([(s, ())])
    seen = {s}
    while q:
        st, path = q.popleft()
        acts = [('move', n) for n in l.nodes] + [('refuel',)]
        if allow_drop:
            acts += [('item', i) for i, p in enumerate(st.locations) if p == 'T']
        for a in acts:
            ns, r = apply(st, l, a, flat)
            if ns == st or ns in seen:
                continue
            if ns.place == 'H':
                return path + (a,)
            seen.add(ns)
            q.append((ns, path + (a,)))
    return None

def forecast(s, l, a, flat=False):
    ns, r = apply(s, l, a, flat)
    return {'state': ns, 'reason': r, 'move_cost': s.fuel - ns.fuel if a[0] == 'move' and ns != s else 0, 'home_before': home_cost(s, l, flat), 'home_after': home_cost(ns, l, flat), 'weight_after': weight(ns, l)}

def state_hash(s):
    return hashlib.sha256(json.dumps(asdict(s), sort_keys=True).encode('utf-8')).hexdigest()

# ASCII-only learning UI. The canonical Japanese text/data remain in the model.
STAGE = 5
START_LEVEL = 0
MESSAGES = {'ready':'READY: SELECT AN ACTION, THEN Z', 'ok':'OK',
    'fuel':'NOT ENOUGH FUEL; NO STATE CHANGED', 'full':'LOAD WOULD EXCEED 4; NO STATE CHANGED',
    'absent':'ITEM IS NEITHER HERE NOR IN THE TRUCK', 'road':'NO DIRECT ROAD BETWEEN THESE NODES',
    'supply':'TANK FULL OR LOCAL SUPPLY EMPTY', 'terminal':'ALREADY HOME: R RESTARTS',
    'invalid':'INVALID ACTION', 'undo':'RESTORED THE PREVIOUS STATE',
    'no_undo':'NO EARLIER STATE'}

class App:
    def __init__(self):
        self.level_index = START_LEVEL
        self.reset()

    @property
    def board(self):
        return LEVELS[self.level_index]

    @property
    def state(self):
        return self.timeline.state

    def reset(self):
        self.timeline = Timeline(initial(self.board))
        self.selection = 1
        self.reason = 'ready'
        self.gated = False

    def options(self):
        choices = tuple(('move', n) for n in self.board.nodes)
        choices += tuple(('item', i) for i in range(len(self.board.items)))
        choices += (('refuel',),)
        return choices

    def act(self, action):
        self.timeline, self.reason = step(self.timeline, self.board, action)

    def control(self, pressed, held):
        # A key must be released before another action; chords are rejected.
        if 'R' in pressed:
            self.reset(); self.gated = True; return
        if self.gated:
            if not held and not pressed: self.gated = False
            return
        if len(pressed | held) > 1:
            self.gated = True; return
        if len(pressed) != 1: return
        key = next(iter(pressed)); self.gated = True
        if key in ('1','2','3'):
            self.level_index = int(key)-1; self.reset(); self.gated=True; return
        if key == 'X':
            self.act(('undo',)); return

        if key in ('LEFT', 'RIGHT', 'UP', 'DOWN'):
            delta = {'LEFT':-1, 'RIGHT':1, 'UP':-5, 'DOWN':5}[key]
            self.selection = (self.selection + delta) % len(self.options())
        elif key in ('Z', 'SPACE'):
            self.act(self.options()[self.selection])

    def snapshot(self):
        return {'stage':STAGE, 'level':self.level_index, 'state':asdict(self.state),
                'selection':self.selection, 'reason':self.reason, 'history':[asdict(s) for s in self.timeline.history]}

    def draw(self, p):
        s, l = self.state, self.board
        p.cls(1)
        p.text(14, 10, 'RETURN FUEL / STAGE '+str(STAGE)+' / SIMPLIFIED LEARNING UI', 10)
        p.text(14, 28, 'LEVEL '+'ABC'[self.level_index]+'   DISTANCE x (1 + CURRENT LOAD)', 7)
        p.text(364, 50, 'FUEL '+str(s.fuel)+'/'+str(l.capacity), 10)
        for i in range(l.capacity): p.rect(364+i*18, 64, 14, 6, 11 if i<s.fuel else 5)
        p.text(364, 80, 'LOAD '+str(weight(s,l))+'/4   VALUE '+str(value(s,l))+'/'+str(l.goal), 7)
        p.text(364, 96, 'SPENT '+str(s.spent)+'   LOCAL OIL '+str(sum(s.supply)), 6)
        coords = {'H':(46,98),'A':(300,98),'B':(300,214),'C':(46,214)}
        if self.level_index == 2:
            coords = {'H':(46,152),'A':(136,88),'B':(300,152),'C':(136,218),'D':(280,238)}
        for u,v,d in l.edges:
            x,y=coords[u]; xx,yy=coords[v]
            p.line(x,y,xx,yy,5)
            p.rect((x+xx)//2-4,(y+yy)//2-5,12,12,1)
            p.text((x+xx)//2,(y+yy)//2-2,str(d),13)
        for i,n in enumerate(l.nodes):
            x,y=coords[n]
            p.circ(x,y,14,3 if n=='H' else 2)
            p.circb(x,y,14,10 if n==s.place else 6)
            p.text(x-2,y-3,n,7)
            if n==s.place: p.rect(x-8,y+7,16,5,10)
            items=''.join(it.id for it,loc in zip(l.items,s.locations) if loc==n)
            p.text(x-14,y+19,'ITEM '+(items or '-'),10)
            if s.supply[i]: p.text(x-14,y-26,'OIL '+str(s.supply[i]),11)
        p.text(364, 123, 'ID   WEIGHT  VALUE  LOCATION', 13)
        for i,it in enumerate(l.items):
            y=142+25*i
            p.rect(358,y-5,266,20,3 if s.locations[i]=='T' else 2)
            place='TRUCK' if s.locations[i]=='T' else s.locations[i]
            p.text(364,y,it.id+'      '+str(it.weight)+'      '+str(it.value)+'      '+place,7)
        selected=self.options()[self.selection]
        f=forecast(s,l,selected)
        ns=f['state']
        p.text(14,274,'PREVIEW '+f['reason'].upper()+'   FUEL '+str(s.fuel)+'>'+str(ns.fuel)+'   LOAD '+str(weight(s,l))+'>'+str(weight(ns,l)),6)
        keep=return_plan(s,l)
        drop=return_plan(s,l,allow_drop=True) if keep is None else None
        advice='KEEP CARGO: YES' if keep is not None else ('KEEP: NO / DROP: YES' if drop is not None else 'NO RETURN FOUND')
        p.text(14,292,'HOME COST (NO REFUEL) '+str(home_cost(s,l))+'   '+advice,13)
        for i,action in enumerate(self.options()):
            x=14+(i%5)*124; y=314+(i//5)*26
            p.rectb(x,y,118,23,10 if i==self.selection else 5)
            name = 'MOVE '+action[1] if action[0]=='move' else 'REFUEL'
            if action[0]=='item':
                idx=action[1]
                name=('DROP ' if s.locations[idx]=='T' else 'PICK ')+l.items[idx].id
            p.text(x+7,y+8,name,10 if i==self.selection else 7)
        p.text(14, 375, 'RESULT '+s.outcome.upper()+'   '+MESSAGES[self.reason],11 if s.outcome=='won' else 7)
        p.text(14,395,'HISTORY '+str(len(self.timeline.history))+' / X RESTORES A CHANGED STATE',6)
        p.text(14, 412, 'ARROWS: SELECT   Z/SPACE: DO   R: RESET   ESC: QUIT',7)
        p.text(14, 429, '1/2/3: SWITCH LEVEL   X: UNDO EVEN AFTER RETURN',13)
        p.text(14, 446, 'EDIT -> SAVE -> QUIT -> RERUN. R DOES NOT RELOAD CODE.',6)

def make_demo_keys():
    app=App(); keys=[]
    keys.append('3'); app.level_index=2; app.reset()
    plan=[('move', 'C'), ('move', 'D'), ('refuel',), ('item', 3), ('move', 'B'), ('item', 1), ('move', 'H')]
    for action in plan:
        idx=app.options().index(action)
        while app.selection!=idx:
            keys.append('RIGHT'); app.selection=(app.selection+1)%len(app.options())
        keys.append('Z'); app.act(action)
    keys.append('X')
    return tuple(keys)

def main(argv=None):
    import argparse
    from pathlib import Path
    import json
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--headless-smoke', action='store_true', help='bounded actual Pyxel render')
    parser.add_argument('--demo-checkpoint', action='store_true', help='replay documented keys')
    parser.add_argument('--frames', type=int, help='bounded headless frames, 1..10000')
    parser.add_argument('--screenshot', help='new PNG; relative to launch working directory')
    parser.add_argument('--state-json', help='new JSON; relative to launch working directory')
    args=parser.parse_args(argv)
    if (args.demo_checkpoint or args.frames is not None or args.screenshot or args.state_json) and not args.headless_smoke:
        parser.error('--demo-checkpoint, --frames and outputs require --headless-smoke')
    if args.frames is not None and not 1<=args.frames<=10000:
        parser.error('--frames must be an integer from 1 to 10000')
    sequence=make_demo_keys() if args.demo_checkpoint else ()
    minimum=len(sequence)*2+2 if sequence else 1
    frames=args.frames if args.frames is not None else max(8,minimum)
    if frames<minimum: parser.error('--frames is too short for the full checkpoint; need '+str(minimum))
    # Resolve EVERY output before importing/initializing Pyxel, which may chdir.
    outputs={}
    for name,raw,suffix in (('screenshot',args.screenshot,'.png'),('state_json',args.state_json,'.json')):
        if raw is None: continue
        path=Path(raw).expanduser()
        if not raw or path.suffix.lower()!=suffix or not path.parent.is_dir() or path.exists() or path.is_symlink():
            parser.error(name+' must be a new '+suffix+' file in an existing directory; refusing overwrite')
        outputs[name]=path.resolve()
    if len(set(outputs.values()))!=len(outputs): parser.error('output paths must differ')
    import pyxel as p
    p.init(640,466,title='Return fuel lesson '+str(STAGE),fps=30,display_scale=2,headless=args.headless_smoke)
    app=App()
    names=('LEFT','RIGHT','UP','DOWN','Z','SPACE','R','X','1','2','3')
    keys={name:getattr(p,'KEY_'+name) for name in names}
    elapsed=0
    def finish():
        # Refuse again at save time; JSON uses exclusive creation as well.
        for path in outputs.values():
            if path.exists() or path.is_symlink(): raise FileExistsError('refusing overwrite: '+str(path))
        app.draw(p)
        if 'screenshot' in outputs: p.screen.save(str(outputs['screenshot']),1)
        if 'state_json' in outputs:
            with outputs['state_json'].open('x',encoding='utf-8') as file:
                json.dump(app.snapshot(),file,ensure_ascii=False,indent=2); file.write('\n')
        p.quit()
    def update():
        nonlocal elapsed
        if p.btnp(p.KEY_ESCAPE): finish(); return
        if args.headless_smoke:
            if elapsed>=frames: finish(); return
            index=elapsed//2
            pressed={sequence[index]} if elapsed%2==0 and index<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)})
        elapsed+=1
    p.run(update,lambda:app.draw(p))

if __name__=='__main__':
    main()