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

"""迷子の地図屋・教材4/5。各ファイルを単独で実行できます。"""
import heapq

from dataclasses import dataclass, replace

DIRECTIONS = {'R': (1, 0), 'D': (0, 1), 'L': (-1, 0), 'U': (0, -1)}

@dataclass(frozen=True)
class Level:
    seed: str
    title: str
    nodes: tuple
    edges: tuple
    marks: tuple

    def points(self):
        return {n: (x, y) for n, x, y in self.nodes}

LEVELS = (Level('A-01-v1', '四角の向こう', (('H', 0, 0), ('A', 4, 0), ('B', 4, 3), ('C', 0, 3), ('D', 7, 0), ('E', 7, 3)), (('H', 'A'), ('A', 'B'), ('B', 'C'), ('C', 'H'), ('A', 'D'), ('D', 'E'), ('E', 'B')), ('B', 'E')), Level('B-01-v1', '近そうな長い横道', (('H', 0, 0), ('A', 6, 0), ('B', 6, 2), ('C', 0, 2), ('D', 6, 6), ('E', 0, 6)), (('H', 'A'), ('A', 'B'), ('B', 'C'), ('C', 'H'), ('B', 'D'), ('D', 'E'), ('E', 'C')), ('A', 'D')), Level('C-01-v1', '短い中央の接続', (('H', 0, 0), ('A', 3, 0), ('B', 8, 0), ('C', 0, 4), ('D', 3, 4), ('E', 8, 4)), (('H', 'A'), ('A', 'B'), ('H', 'C'), ('A', 'D'), ('B', 'E'), ('C', 'D'), ('D', 'E')), ('B', 'C')))

def length(l, i):
    p = l.points()
    u, v = l.edges[i]
    return abs(p[u][0] - p[v][0]) + abs(p[u][1] - p[v][1])

def edge_dir(l, i, n):
    p = l.points()
    u, v = l.edges[i]
    other = v if n == u else u
    x, y = p[n]
    xx, yy = p[other]
    return next((k for k, (dx, dy) in DIRECTIONS.items() if (xx - x > 0) - (xx - x < 0) == dx and (yy - y > 0) - (yy - y < 0) == dy))

@dataclass(frozen=True)
class State:
    node: str = 'H'
    edge: int = -1
    offset: int = 0
    origin: str = 'H'
    progress: int = 0
    seen: tuple = ('H',)
    measured: tuple = (0,) * 7
    lower: tuple = (0,) * 7
    marks: tuple = ()
    steps: int = 0
    tick: int = 0
    outcome: str = 'running'

def initial(l):
    return State(measured=(0,) * len(l.edges), lower=(0,) * len(l.edges))

def position(l, s):
    if s.node:
        return l.points()[s.node]
    u, v = l.edges[s.edge]
    x, y = l.points()[u]
    dx, dy = DIRECTIONS[edge_dir(l, s.edge, u)]
    return (x + dx * s.offset, y + dy * s.offset)

def move(s, l, a):
    if s.outcome != 'running':
        return (s, 'ended')
    if a == 'submit':
        if s.node != 'H':
            return (s, 'home_only')
        return (replace(s, tick=s.tick + 1, outcome='won' if len(s.marks) == 2 else 'partial'), 'submit')
    if a not in DIRECTIONS:
        return (s, 'invalid')
    if s.node:
        options = [i for i, e in enumerate(l.edges) if s.node in e and edge_dir(l, i, s.node) == a]
        if not options:
            return (s, 'wall')
        i = options[0]
        u, v = l.edges[i]
        origin = s.node
        offset = 1 if origin == u else length(l, i) - 1
        progress = 1
    else:
        i = s.edge
        u, v = l.edges[i]
        origin = s.origin
        forward = edge_dir(l, i, u)
        if a == forward:
            offset = s.offset + 1
        elif DIRECTIONS[a] == tuple((-d for d in DIRECTIONS[forward])):
            offset = s.offset - 1
        else:
            return (s, 'wall')
        progress = offset if origin == u else length(l, i) - offset
    low = list(s.lower)
    low[i] = max(low[i], progress)
    measured = list(s.measured)
    seen = s.seen
    marks = s.marks
    node = u if offset == 0 else v if offset == length(l, i) else ''
    if node and node != origin:
        measured[i] = progress
        if node not in seen:
            seen = seen + (node,)
        if node in l.marks and node not in marks:
            marks = marks + (node,)
    ns = State(node, -1 if node else i, 0 if node else offset, node or origin, 0 if node else progress, seen, tuple(measured), tuple(low), marks, s.steps + 1, s.tick + 1, s.outcome)
    return (ns, 'measured' if node and node != origin else 'walk')

def known_return(s, l):
    adj = {n: [] for n in s.seen}
    for i, d in enumerate(s.measured):
        if d:
            u, v = l.edges[i]
            adj[u].append((v, d))
            adj[v].append((u, d))
    starts = [(0, s.node)] if s.node else [(s.progress, s.origin)]
    if not s.node and s.measured[s.edge]:
        u, v = l.edges[s.edge]
        starts = [(s.offset, u), (s.measured[s.edge] - s.offset, v)]
    heap = starts[:]
    heapq.heapify(heap)
    best = {}
    while heap:
        d, n = heapq.heappop(heap)
        if n in best:
            continue
        best[n] = d
        if n == 'H':
            return d
        for nxt, cost in adj[n]:
            heapq.heappush(heap, (d + cost, nxt))
    return None

def view(s, l, public_lengths=False):
    """The only map UI interface. Unknown endpoints and world coordinates omitted."""
    rows = []
    for n in s.seen:
        exits = []
        for i, e in enumerate(l.edges):
            if n not in e:
                continue
            d = s.measured[i]
            other = e[1] if n == e[0] else e[0]
            exits.append({'direction': edge_dir(l, i, n), 'target': other if d else '?', 'distance': length(l, i) if public_lengths else d, 'lower': 0 if d or public_lengths else s.lower[i]})
        exits.sort(key=lambda e: 'URDL'.index(e['direction']))
        rows.append({'node': n, 'mark': n in s.marks, 'exits': exits})
    tiles = {}
    p = l.points()
    for i, (u, v) in enumerate(l.edges):
        x, y = p[u]
        dx, dy = DIRECTIONS[edge_dir(l, i, u)]
        for k in range(length(l, i) + 1):
            tiles[x + k * dx, y + k * dy] = True
    at = position(l, s)
    front = [(at, 0)]
    visited = {at}
    local = []
    while front:
        (x, y), d = front.pop(0)
        node = next((n for n, px, py in l.nodes if (px, py) == (x, y)), None)
        local.append({'x': x - at[0], 'y': y - at[1], 'node': node if node in s.seen else None, 'mark': bool(node in l.marks), 'room': node is not None})
        if d < LOCAL_RADIUS:
            for dx, dy in DIRECTIONS.values():
                pos = (x + dx, y + dy)
                if pos in tiles and pos not in visited:
                    visited.add(pos)
                    front.append((pos, d + 1))
    return {'rows': rows, 'local': local, 'steps': s.steps, 'count': sum((bool(d) for d in s.measured)), 'marks': s.marks, 'progress': s.progress, 'origin': s.origin, 'return': known_return(s, l)}


STAGE = 4
LOCAL_RADIUS = 2
WIDTH, HEIGHT = 640, 440

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

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

    def reset(self):
        self.state = initial(self.level)
        self.message = 'ready'
        self.gated = False

    def commit(self, action):
        self.state, self.message = move(self.state, self.level, action)

    def control(self, pressed, held):
        # A new reset wins even during a release gate. Held R is not a new reset.
        if 'R' in pressed:
            self.reset()
            self.gated = True
            return
        if self.gated:
            if not pressed and not held:
                self.gated = False
            return
        if len(pressed | held) > 1:
            self.message = 'release all keys'
            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
        arrows = {'UP': 'U', 'RIGHT': 'R', 'DOWN': 'D', 'LEFT': 'L'}
        if key in arrows: self.commit(arrows[key])
        elif key == 'Z': self.commit('submit')

    def draw(self, p):
        s = self.state
        p.cls(1)
        p.text(14, 12, 'MEASURED MAP / STAGE '+str(STAGE)+' / '+self.level.seed, 10)
        p.text(14, 31, 'ARROW = ONE WALKING STEP / RELEASE BEFORE NEXT KEY', 7)
        p.text(14, 50, 'STEPS '+bounded(s.steps), 11)
        projected = view(s, self.level)
        rows = projected['rows']
        p.text(14, 74, 'LOCAL: '+str(LOCAL_RADIUS)+' WALKING STEPS', 6)
        p.rectb(12, 91, 224, 218, 5)
        cx, cy = 124, 198
        for cell in projected['local']:
            x, y = cx + cell['x']*36, cy + cell['y']*36
            p.rect(x-15, y-15, 31, 31, 3 if cell['room'] else 5)
            p.rectb(x-15, y-15, 31, 31, 6)
            if cell['room']: p.circb(x, y, 10, 10 if cell['mark'] else 13)
            if cell['node']: p.text(x-2, y-3, cell['node'], 7)
            if cell['mark']: p.text(x+8, y-14, '*', 10)
        p.circ(cx, cy+5, 4, 10)
        p.tri(cx-4, cy+2, cx+4, cy+2, cx, cy-5, 7)
        p.text(20, 291, 'ONLY RELATIVE LOCAL OFFSETS', 13)
        p.text(253, 74, 'LEDGER: LAYOUT IS NOT DISTANCE', 6)
        for row_index, row in enumerate(rows):
            y = 94 + row_index*34
            p.rect(252, y, 374, 30, 3 if row['node'] == s.node else 2)
            p.text(258, y+5, row['node']+('*' if row['mark'] else ''), 10)
            for exit_index, item in enumerate(row['exits']):
                number = str(item['distance']) if item['distance'] else (
                    '>='+str(item['lower']) if item['lower'] else '?')
                label = item['direction']+' '+number+' -> '+item['target']
                p.text(283+exit_index*108, y+5, label, 11 if item['distance'] else 13)
            p.text(283, y+18, 'UNMEASURED TARGETS STAY ?', 13)
        p.text(14, 325, 'ROOM '+s.node if s.node else 'FROM '+s.origin+' / PROGRESS '+str(s.progress), 10)
        p.text(253, 325, 'MEASURED '+str(sum(bool(x) for x in s.measured))+'/'+str(len(self.level.edges)), 11)
        p.text(14, 344, 'KNOWN RETURN TO H: '+str(projected['return'])+' STEPS', 6)
        p.text(340, 344, 'MARKS '+str(len(s.marks))+'/2 / '+s.outcome.upper(), 10)
        p.text(14, 367, self.message.upper(), 10)
        p.text(14, 390, 'ARROWS: WALK / R: RESET SAME BANK / ESC: QUIT', 7)
        p.text(14, 408, 'Z: SUBMIT AT H / 1-3: FIXED BANKS', 7)

def bounded(n):
    return str(n) if n <= 9999 else '9999+'

def make_demo_keys():
    return ('RIGHT',)*7+('DOWN',)*3+('LEFT',)*7+('UP',)*3+('Z',)

def main(headless=False, screenshot=None, demo=False, bank=1):
    import pyxel as p
    p.init(WIDTH, HEIGHT, title='Measured map lesson '+str(STAGE), fps=30,
           display_scale=1, headless=headless)
    app = App(bank-1)
    sequence = make_demo_keys() if demo else ()
    keys = {k: getattr(p, 'KEY_'+k) for k in
            ('UP','RIGHT','DOWN','LEFT','Z','X','R','1','2','3')}
    limit = max(20, 2*len(sequence)+8)
    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)
            p.quit()
            return
        if headless:
            if demo:
                index = p.frame_count//2
                down = {sequence[index]} if p.frame_count%2 == 0 and index < len(sequence) else set()
                app.control(down, down)
        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=__doc__)
    parser.add_argument('--bank', type=int, choices=range(1, len(LEVELS)+1), default=1)
    parser.add_argument('--headless-smoke', action='store_true')
    parser.add_argument('--demo-checkpoint', action='store_true')
    parser.add_argument('--screenshot')
    args = parser.parse_args()
    if args.demo_checkpoint and not args.headless_smoke:
        parser.error('demo requires headless mode')
    if args.screenshot and not args.headless_smoke:
        parser.error('screenshot requires headless mode')
    if args.screenshot and (Path(args.screenshot).suffix.lower() != '.png' or
                            not Path(args.screenshot).parent.is_dir() or Path(args.screenshot).exists()):
        parser.error('use a new .png in an existing directory')
    # Pyxel may change cwd during initialization. Anchor the validated path first.
    screenshot = str(Path(args.screenshot).resolve()) if args.screenshot else None
    main(args.headless_smoke, screenshot, args.demo_checkpoint, args.bank)