# 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。路線の共有重量と毎刻の出発予約を分ける。"""
from dataclasses import dataclass,replace,asdict
import itertools
STATIONS = ('A', 'B', 'X', 'Y', 'D', 'E')
EDGES = (('A', 'X', 1), ('B', 'X', 1), ('X', 'Y', 1), ('Y', 'D', 1), ('Y', 'E', 1), ('A', 'D', 3), ('B', 'E', 3), ('D', 'E', 1))

@dataclass(frozen=True)
class Level:
    id: str
    title: str
    edges: tuple
    demand: tuple
    capacity: int = 7
    limit: int = 32

LEVELS = (Level('A', '二つの直通', EDGES, ((0, 'A', 'D'), (0, 'B', 'E'), (1, 'A', 'D'), (1, 'B', 'E'), (2, 'A', 'D'), (2, 'B', 'E'))), Level('B', '行先が交じる夜', EDGES, ((0, 'A', 'D'), (0, 'B', 'E'), (1, 'A', 'E'), (1, 'B', 'D'), (2, 'A', 'E'), (2, 'B', 'D'))), Level('C', '一方に寄った便', EDGES + (('B', 'D', 2),), ((0, 'A', 'E'), (0, 'B', 'D'), (1, 'A', 'E'), (1, 'B', 'D'), (2, 'B', 'D'), (3, 'B', 'D'))))

@dataclass(frozen=True)
class Star:
    id: int
    born: int
    source: str
    destination: str
    path: tuple = ()
    lines: tuple = ()
    index: int = 0
    eta: int = 0
    departed: int = 0
    arrived: int = 0
    ready: int = 1
    delivered: int = -1
    wait: int = 0
    reason: str = '出発待ち'

@dataclass(frozen=True)
class State:
    routes: tuple = ((), ())
    t: int = 0
    stars: tuple = ()
    status: str = 'edit'
    reason: str = '二本までの路線で、六つの星を届けよう。'
    station_wait: tuple = (0, 0, 0, 0, 0, 0)
    departures: tuple = ()

def initial(level):
    return State()

def edge_lengths(level):
    return {tuple(sorted((a, b))): d for a, b, d in level.edges}

def route_edges(route):
    return tuple((tuple(sorted((a, b))) for a, b in zip(route, route[1:])))

def weight(level, routes):
    return sum((edge_lengths(level)[e] for e in set(itertools.chain.from_iterable((route_edges(r) for r in routes)))))

def intersections(routes):
    return set(routes[0]) & set(routes[1])

def valid_routes(level, routes):
    return len(routes) == 2 and all((len(r) <= 5 and len(set(r)) == len(r) and all((n in STATIONS for n in r)) and all((e in edge_lengths(level) for e in route_edges(r))) for r in routes))

def edit_reason(level, s, line, node):
    r = s.routes[line]
    if len(r) >= 5:
        return '一本は最大4辺です。'
    if node in r:
        return '同じ駅を二度通れません。'
    if r and tuple(sorted((r[-1], node))) not in edge_lengths(level):
        return 'その駅へ使える辺がありません。'
    return ''

def path_for(level, routes, source, dest, penalty=1):
    edges = edge_lengths(level)
    cross = intersections(routes)
    adj = {n: [] for n in STATIONS}
    uses = {}
    for line, r in enumerate(routes):
        for a, b in zip(r, r[1:]):
            e = tuple(sorted((a, b)))
            uses.setdefault(e, []).append(line)
            if b not in adj[a]:
                adj[a].append(b)
                adj[b].append(a)
    candidates = []

    def walk(path):
        a = path[-1]
        if a == dest:
            options = [uses[tuple(sorted(e))] for e in zip(path, path[1:])]
            for lines in itertools.product(*options):
                cost = sum((edges[tuple(sorted(e))] for e in zip(path, path[1:]))) + max(0, len(path) - 2) + penalty * sum((n in cross for n in path[1:-1]))
                transitions = len(lines) + sum((a != b for a, b in zip(lines, lines[1:])))
                candidates.append((cost, transitions, tuple(path), tuple(lines)))
            return
        for b in sorted(adj[a]):
            if b not in path:
                walk(path + [b])
    walk([source])
    return min(candidates) if candidates else None

def planned_arrival(level, routes, z, penalty=1):
    if not z.path:
        return None
    edges = edge_lengths(level)
    cross = intersections(routes)
    return z.born + 1 + sum((edges[tuple(sorted(e))] for e in zip(z.path, z.path[1:]))) + max(0, len(z.path) - 2) + penalty * sum((n in cross for n in z.path[1:-1]))

def spawn(level, s, t, penalty):
    out = list(s.stars)
    for i, (born, source, dest) in enumerate(level.demand):
        if born != t:
            continue
        plan = path_for(level, s.routes, source, dest, penalty)
        out.append(Star(i, born, source, dest, plan[2] if plan else (), plan[3] if plan else (), arrived=born, ready=born + 1, reason='出発待ち' if plan else '接続なし'))
    return tuple(sorted(out, key=lambda z: z.id))

def advance(level, s, penalty=1):
    if s.status != 'run':
        return s
    t = s.t + 1
    cross = intersections(s.routes)
    stars = []
    for z in s.stars:
        if z.eta and z.eta == t:
            node = z.path[z.index]
            if node == z.destination:
                z = replace(z, eta=0, arrived=t, delivered=t, reason='配達済み')
            else:
                z = replace(z, eta=0, arrived=t, ready=t + 1 + (penalty if node in cross else 0), reason='一拍待ち' if node in cross and penalty else '出発待ち')
        stars.append(z)
    tmp = replace(s, stars=tuple(stars))
    stars = list(spawn(level, tmp, t, penalty))
    edge_used = set()
    station_used = set()
    chosen = set()
    rejections = {}
    departures = []
    waiting = list(s.station_wait)
    candidates = sorted((z for z in stars if not z.eta and z.delivered < 0 and z.path and (t >= z.ready)), key=lambda z: (z.arrived, z.id))
    edges = edge_lengths(level)
    for z in candidates:
        a, b = z.path[z.index:z.index + 2]
        e = (a, b)
        if a in cross and a in station_used:
            rejections[z.id] = '交差の列待ち'
            continue
        if e in edge_used:
            rejections[z.id] = '辺の列待ち'
            continue
        chosen.add(z.id)
        edge_used.add(e)
        if a in cross:
            station_used.add(a)
        departures.append((z.id, a, b))
    out = []
    for z in stars:
        if z.id in chosen:
            a, b = z.path[z.index:z.index + 2]
            z = replace(z, index=z.index + 1, eta=t + edges[tuple(sorted((a, b)))], departed=t, reason='移動中')
        elif not z.eta and z.delivered < 0:
            node = z.path[z.index] if z.path else z.source
            reason = '接続なし' if not z.path else rejections.get(z.id, '一拍待ち' if t < z.ready and z.ready - z.arrived > 1 else '出発待ち')
            if t > z.arrived:
                waiting[STATIONS.index(node)] += 1
                z = replace(z, wait=z.wait + 1)
            z = replace(z, reason=reason)
        out.append(z)
    done = len(out) == len(level.demand) and all((z.delivered >= 0 for z in out))
    status = 'done' if done else 'time' if t >= level.limit else 'run'
    return replace(s, t=t, stars=tuple(out), station_wait=tuple(waiting), departures=tuple(departures), status=status, reason='六つの星が、夜の向こうへ届きました。' if done else '32刻で試運転終了。接続と列を見直そう。' if status == 'time' else '星の札は、実際の待機理由を示します。')

def apply_action(level,state,action,penalty=1):
    if state.status=='edit':
        bits=action.split(':')
        if len(bits)==3 and bits[0]=='add' and bits[1] in ('0','1') and bits[2] in STATIONS:
            line=int(bits[1]);node=bits[2]
            if not edit_reason(level,state,line,node):
                routes=list(state.routes);routes[line]=routes[line]+(node,)
                return replace(state,routes=tuple(routes))
        if len(bits)==2 and bits[0]=='pop' and bits[1] in ('0','1'):
            line=int(bits[1]);routes=list(state.routes);routes[line]=routes[line][:-1]
            return replace(state,routes=tuple(routes))
        if action=='start' and weight(level,state.routes)<=level.capacity:
            started=replace(state,status='run',reason='Spaceで一刻。Pで連続運転・停止。')
            return replace(started,stars=spawn(level,started,0,penalty))
    if action=='tick':return advance(level,state,penalty)
    return state

def forecast(level,state,penalty=1):
    if state.status=='edit':state=apply_action(level,state,'start',penalty)
    return advance(level,state,penalty) if state.status=='run' else None

STAGE = 4
PENALTY = 1
LEVEL = LEVELS[0]
START_ROUTES = (('A', 'X', 'Y', 'D'), ('B', 'X', 'Y', 'E'))
POS = {'A':(60,100),'B':(60,240),'X':(175,170),'Y':(290,170),'D':(405,100),'E':(405,240)}

class App:
    def __init__(self):self.reset()
    def reset(self):
        self.state=replace(initial(LEVEL),routes=START_ROUTES)
        self.station=0
        self.line=0
        self.gated=False
        self.message='READY'

    def commit(self,action):
        before=self.state
        self.state=apply_action(LEVEL,self.state,action,PENALTY)
        self.message=action.upper() if self.state!=before else 'NO CHANGE / REFUSED'

    def control(self,pressed,held):
        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.gated=True;self.message='RELEASE ALL KEYS';return
        if len(pressed)!=1:return
        key=next(iter(pressed));self.gated=True
        if key=='RETURN':self.commit('start');return
        if key=='SPACE':self.commit('tick');return
        if self.state.status!='edit':return
        if key in ('LEFT','RIGHT'):self.station=(self.station+(-1 if key=='LEFT' else 1))%len(STATIONS)
        elif key=='TAB':self.line=1-self.line
        elif key=='Z':self.commit('add:'+str(self.line)+':'+STATIONS[self.station])
        elif key=='BACKSPACE':self.commit('pop:'+str(self.line))

    def draw(self,p):
        state=self.state;p.cls(1)
        p.text(12,8,'STAR RAILWAYS / STAGE '+str(STAGE)+' / A',10)
        p.text(12,25,'WEIGHT '+str(weight(LEVEL,state.routes))+' / CAP '+str(LEVEL.capacity)+' / '+state.status.upper(),7)
        for a,b,length in LEVEL.edges:
            x,y=POS[a];xx,yy=POS[b];p.line(x,y,xx,yy,5);p.text((x+xx)//2+4,(y+yy)//2+4,str(length),6)
        for line,route in enumerate(state.routes):
            for a,b in zip(route,route[1:]):
                x,y=POS[a];xx,yy=POS[b];p.line(x,y+line*3,xx,yy+line*3,(10,12)[line])
        cross=intersections(state.routes)
        for i,n in enumerate(STATIONS):
            x,y=POS[n];p.circ(x,y,11,2 if n in cross else 1);p.circb(x,y,11,10 if i==self.station else 7);p.text(x-2,y-2,n,7)
            if n in cross:p.text(x-10,y+17,'CROSS',14)
        p.text(12,274,'LINE '+str(self.line+1)+' / STATION '+STATIONS[self.station],10)
        for i,route in enumerate(state.routes):p.text(12,291+i*15,'L'+str(i+1)+': '+('-'.join(route) or '(empty)'),(10,12)[i])
        delivered=sum(z.delivered>=0 for z in state.stars)
        p.text(470,50,'TICK '+str(state.t),10)
        p.text(470,67,'DELIVERED '+str(delivered)+'/'+str(len(LEVEL.demand)),11)
        p.text(470,84,'TOTAL WAIT '+str(sum(state.station_wait)),14)
        for z in state.stars:
            y=110+z.id*25
            p.text(470,y,'#'+str(z.id)+' '+z.source+'>'+z.destination+' '+('DONE '+str(z.delivered) if z.delivered>=0 else 'ETA '+str(z.eta) if z.eta else 'WAIT '+str(z.wait)),7)
            labels={'移動中':'MOVING','配達済み':'DELIVERED','出発待ち':'READY WAIT','接続なし':'NO PATH','辺の列待ち':'EDGE QUEUE','交差の列待ち':'STATION QUEUE','一拍待ち':'EXTRA REST'}
            p.text(478,y+10,labels.get(z.reason,'WAIT'),6)
            if z.path:
                if z.eta:
                    a,b=z.path[z.index-1:z.index+1];x,y0=POS[a];xx,yy=POS[b]
                    duration=z.eta-z.departed;u=(state.t-z.departed)/duration
                    x,y0=x+(xx-x)*u,y0+(yy-y0)*u
                else:x,y0=POS[z.path[z.index]]
                p.circ(x,y0-15-z.id*4,2,(8,10,12,11,9,14)[z.id])
        q=forecast(LEVEL,state,PENALTY)
        p.text(12,330,'NEXT DEPARTURES '+str(q.departures) if q else 'NO NEXT TICK',6)
        p.text(12,347,'NEXT DELIVERED '+str(sum(z.delivered>=0 for z in q.stars)) if q else 'END / OR TOO HEAVY',6)
        p.text(12,367,'CROSS: ONE DEPARTURE/TICK + EXTRA REST1',14)
        p.text(12,406,self.message,10)
        p.text(12,430,'LEFT/RIGHT: STATION LIST / TAB: LINE / Z: ADD / BACKSPACE: POP',7)
        p.text(12,447,'R: RESET TO THIS STAGE EXAMPLE / ESC: QUIT',7)
        p.text(12,464,'ENTER: START / SPACE: TICK'+(' / X: UNDO; DECISIONS STAY' if STAGE==5 else ''),7)

def make_demo_keys():
    return ('RETURN',)+('SPACE',)*13

def main(headless=False,screenshot=None,demo=False):
    import pyxel as p
    p.init(672,492,title='Star railway lesson '+str(STAGE),fps=30,display_scale=2,headless=headless)
    app=App();sequence=make_demo_keys() if demo else ()
    keys={'LEFT':p.KEY_LEFT,'RIGHT':p.KEY_RIGHT,'TAB':p.KEY_TAB,'Z':p.KEY_Z,'BACKSPACE':p.KEY_BACKSPACE,'RETURN':p.KEY_RETURN,'SPACE':p.KEY_SPACE,'X':p.KEY_X,'R':p.KEY_R}
    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()
        if headless:
            if demo:
                index=p.frame_count//2;pressed={sequence[index]} if p.frame_count%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)})
    p.run(update,lambda:app.draw(p))

if __name__=='__main__':
    import argparse
    from pathlib import Path
    parser=argparse.ArgumentParser(description=__doc__);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')
    if args.screenshot:args.screenshot=str(Path(args.screenshot).resolve())
    main(args.headless_smoke,args.screenshot,args.demo_checkpoint)