Blame for brain_lab/server.py

import io,json,re,sqlite3,tempfile,threading,time,uuid
from typing import Literal
from pathlib import Path
import numpy as np
import torch
from fastapi import FastAPI,HTTPException,Request,UploadFile,File,Form
from fastapi.responses import FileResponse,JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.exceptions import RequestValidationError
from pydantic import BaseModel,Field
from backend import ml
from backend import sketch_provider
from backend import tribe_service
from backend.credentials import set_global_key
from backend.networks import CATEGORIES
from backend.racing import Racer
from backend.mario import MarioLab
from backend.arcade_rl import ArcadeRL
from backend.neural_exhibits import Garden,Plasticity,PATTERNS,GARDEN_INPUTS
from backend import drugs,emotion_study,clip_library,mindeye,mindeye_live
from backend.anatomy import ANATOMY,SOURCES

ROOT=Path(__file__).resolve().parent
ROOT.joinpath('data').mkdir(exist_ok=True)
DB=ROOT/'data'/'lab.sqlite3'
with sqlite3.connect(DB) as c:
    c.executescript('CREATE TABLE IF NOT EXISTS orders(session TEXT, customer TEXT, drink TEXT, PRIMARY KEY(session,customer)); CREATE TABLE IF NOT EXISTS memories(id TEXT PRIMARY KEY, session TEXT, caption TEXT, image TEXT, feature BLOB);')
DEFAULT_MODULES={k:True for k in ['occipital','hippocampus','motor','frontal','association','temporal','wernicke','broca','basal','cerebellum','pineal','hypothalamus','amygdala','ras','thalamus','callosum']}
STATES={}; STATE_LOCK=threading.RLock()
app=FastAPI(title='Brain Lab')

@app.on_event('startup')
def resume_cinema_queue():
    clip_library.resume_pending()

@app.middleware('http')
async def fresh_local_assets(request,call_next):
    response=await call_next(request)
    response.headers['Cache-Control']='no-store'
    return response

def session(request):
    sid=request.headers.get('X-Lab-Session','classroom')
    if not re.fullmatch(r'[a-zA-Z0-9-]{1,80}',sid): raise HTTPException(400,'Invalid session.')
    with STATE_LOCK:
        if sid not in STATES:
            STATES[sid]={'modules':dict(DEFAULT_MODULES),'racer':None,'garden':None,'plasticity':None,'cinema':None,'sketch_config':{'provider':'openrouter'}}
    return sid,STATES[sid]

def require(state,*parts):
    off=[next((p['name'] for p in ANATOMY if p['id']==k),k) for k in parts if not state['modules'].get(k,True)]
    if off: raise HTTPException(409,'Switch on '+', '.join(off)+' to use this step.')

@app.exception_handler(RuntimeError)
async def model_error(request,exc): return JSONResponse(status_code=503,content={'detail':str(exc)})

@app.exception_handler(ValueError)
async def value_error(request,exc): return JSONResponse(status_code=400,content={'detail':str(exc)})

@app.exception_handler(RequestValidationError)
async def validation_error(request,exc):
    return JSONResponse(status_code=422,content={'detail':'Invalid request fields: '+', '.join('.'.join(str(x) for x in e['loc']) for e in exc.errors())})

class Modules(BaseModel):
    modules:dict[str,bool]
@app.post('/api/modules')
def set_modules(body:Modules,request:Request):
    _,s=session(request)
    with STATE_LOCK:
        for k,v in body.modules.items():
            if k not in DEFAULT_MODULES: raise HTTPException(400,'Unknown module: '+k)
            s['modules'][k]=v
    return s['modules']

@app.get('/api/state')
def state(request:Request):
    sid,s=session(request)
    return {'modules':s['modules'],'anatomy':ANATOMY,'sources':SOURCES,'categories':[],'sample_categories':CATEGORIES,'sketch_vocabulary':'open'}

@app.get('/api/status')
def status():
    hf=ROOT/'models'/'hf'
    receipt=ROOT/'models'/'training_receipt.json'
    def available(name): return (hf/name/'model.safetensors').exists() or (hf/name/'pytorch_model.bin').exists() or (name=='image' and (hf/name/'unet'/'diffusion_pytorch_model.fp16.safetensors').exists())
    return {'models':{k:available(k) for k in ['language','vision','caption','speech','image']},'sketch':True,'sketch_model':sketch_provider.MODEL,'sketch_api':sketch_provider.config_status({'provider':'openrouter'}),'pong':(ROOT/'models'/'pong.pt').exists(),
        'pong_policy':ml.pong_config(),'device':ml.DEVICE,'training':json.loads(receipt.read_text()) if receipt.exists() else None,'tribe':tribe_service.status()}

class SketchRequest(BaseModel):
    image:str
    hint:str=Field(default='',max_length=500)
@app.post('/api/sketch')
def sketch(body:SketchRequest,request:Request):
    _,s=session(request); require(s,'occipital','thalamus')
    if body.hint: require(s,'wernicke','association','frontal')
    return sketch_provider.recognize(ml.decode_image(body.image),s['sketch_config'],body.hint)

class SketchConfig(BaseModel):
    provider:str=Field(pattern=r'^(openrouter|perplexity)$')
    api_key:str=Field(default='',max_length=500)
    save_globally:bool=False
@app.get('/api/sketch/config')
def sketch_config(request:Request): return sketch_provider.config_status(session(request)[1]['sketch_config'])
@app.post('/api/sketch/config')
def set_sketch_config(body:SketchConfig,request:Request):
    _,s=session(request);s['sketch_config']['provider']=body.provider
    if body.api_key:
        if body.save_globally: set_global_key(body.provider,body.api_key)
        else: s['sketch_config'].setdefault('api_keys',{})[body.provider]=body.api_key
    return sketch_provider.config_status(s['sketch_config'])

@app.get('/api/samples')
def samples():
    path=ROOT/'data'/'samples.json'
    return json.loads(path.read_text()) if path.exists() else {}

class Pong(BaseModel):
    state:list[float]=Field(min_length=5,max_length=5)
    threshold:float=Field(default=1.,ge=.2,le=4)
@app.get('/api/pong/config')
def pong_configuration():return ml.pong_config()

@app.post('/api/pong')
def pong(body:Pong,request:Request):
    _,s=session(request); require(s,'occipital','basal','ras')
    result=ml.pong(body.state,body.threshold)
    if not s['modules']['motor']: result['action']=0
    result['correction']=s['modules']['cerebellum']; return result

def racer_for(s):
    with STATE_LOCK:
        if s['racer'] is None: s['racer']=Racer()
        return s['racer']

class RaceTrack(BaseModel):
    points:list[list[float]]=Field(min_length=8,max_length=1500)
    width:float=Field(default=78,ge=36,le=110,allow_inf_nan=False)
    seed:int=Field(default=7,ge=0,le=2**31-1)
    population:int=Field(default=100,ge=1,le=150)

class RaceTraining(BaseModel):
    steps:int=Field(default=24,ge=1,le=64)

@app.get('/api/racer')
def race_state(request:Request):
    race=racer_for(session(request)[1])
    with race.lock: return race.status()

@app.post('/api/racer/track')
def race_track(body:RaceTrack,request:Request):
    _,s=session(request)
    race=Racer(body.points,body.width,body.seed,body.population)
    with STATE_LOCK: s['racer']=race
    return race.status()

@app.post('/api/racer/train')
def race_train(body:RaceTraining,request:Request):
    _,s=session(request); require(s,'occipital','motor','basal','ras','hippocampus')
    race=racer_for(s)
    with race.lock: return race.train(body.steps)

@app.post('/api/racer/evaluate')
def race_evaluate(request:Request):
    _,s=session(request); require(s,'occipital','motor','basal','ras')
    race=racer_for(s)
    with race.lock: return race.evaluate()

def exhibit_for(s,name,cls):
    with STATE_LOCK:
        if s[name] is None:s[name]=cls()
        return s[name]

class GardenRequest(BaseModel):
    pixels:list[float]=Field(default_factory=lambda:PATTERNS[0].tolist(),min_length=GARDEN_INPUTS,max_length=GARDEN_INPUTS)
    label:int=Field(default=0,ge=0,le=2)
    count:int=Field(default=12,ge=1,le=64)
    rate:float=Field(default=.025,ge=.001,le=.08,allow_inf_nan=False)

def garden_pixels(body):
    if not np.isfinite(body.pixels).all() or any(x<0 or x>1 for x in body.pixels):raise ValueError('Pixels must be between zero and one.')

@app.get('/api/garden')
def garden_state(request:Request):
    g=exhibit_for(session(request)[1],'garden',Garden)
    with g.lock:return g.status()

@app.post('/api/garden/probe')
def garden_probe(body:GardenRequest,request:Request):
    _,s=session(request);require(s,'occipital');garden_pixels(body)
    g=exhibit_for(s,'garden',Garden)
    with g.lock:return g.status(body.pixels)

@app.post('/api/garden/teach')
def garden_teach(body:GardenRequest,request:Request):
    _,s=session(request);require(s,'occipital','hippocampus','association');garden_pixels(body)
    g=exhibit_for(s,'garden',Garden)
    with g.lock:return g.teach(body.pixels,body.label,body.count,body.rate)

@app.post('/api/garden/grow')
def garden_grow(body:GardenRequest,request:Request):
    _,s=session(request);require(s,'hippocampus','association')
    g=exhibit_for(s,'garden',Garden)
    with g.lock:
        g.grow(body.count,body.rate);return g.status(body.pixels)

@app.post('/api/garden/reset')
def garden_reset(request:Request):
    _,s=session(request)
    with STATE_LOCK:s['garden']=Garden()
    return s['garden'].status()

class PlasticityRequest(BaseModel):
    kind:str=Field(default='strongest',pattern=r'^(strongest|neuron|edge)$')
    neuron:int=Field(default=0,ge=0,le=7)
    layer:int=Field(default=0,ge=0,le=1)
    row:int=Field(default=0,ge=0,le=7)
    col:int=Field(default=0,ge=0,le=7)
    count:int=Field(default=16,ge=1,le=128)
    probe:list[float]=Field(default_factory=lambda:[.65,.65],min_length=2,max_length=2)

@app.get('/api/plasticity')
def plasticity_state(request:Request):
    p=exhibit_for(session(request)[1],'plasticity',Plasticity)
    with p.lock:return p.status()

@app.post('/api/plasticity/probe')
def plasticity_probe(body:PlasticityRequest,request:Request):
    _,s=session(request);require(s,'occipital')
    if not np.isfinite(body.probe).all() or any(abs(x)>1 for x in body.probe):raise ValueError('Probe outside the board.')
    p=exhibit_for(s,'plasticity',Plasticity)
    with p.lock:return p.status(body.probe)

@app.post('/api/plasticity/cut')
def plasticity_cut(body:PlasticityRequest,request:Request):
    p=exhibit_for(session(request)[1],'plasticity',Plasticity)
    with p.lock:return p.cut(body.kind,body.neuron,body.layer,body.row,body.col)

@app.post('/api/plasticity/retrain')
def plasticity_retrain(body:PlasticityRequest,request:Request):
    _,s=session(request);require(s,'hippocampus','association')
    p=exhibit_for(s,'plasticity',Plasticity)
    with p.lock:p.practice(body.count);return p.status()

@app.post('/api/plasticity/reset')
def plasticity_reset(request:Request):
    _,s=session(request)
    with STATE_LOCK:s['plasticity']=Plasticity()
    return s['plasticity'].status()

class DrugRequest(BaseModel):
    model_config={'extra':'forbid'}
    drug:str=Field(default='none',max_length=30)
    strength:float=Field(default=.7,ge=0,le=1,allow_inf_nan=False)

@app.get('/api/drugs')
def drug_catalog():
    return {'drugs':drugs.CATALOG,'kind':'Qualitative mechanism teaching model','clinical_prediction':False}

@app.post('/api/drugs/simulate')
def drug_simulate(body:DrugRequest):return drugs.simulate(body.drug,body.strength)

class Memory(BaseModel):
    image:str
    caption:str=Field(default='',max_length=300)
@app.post('/api/memories')
def memory(body:Memory,request:Request):
    sid,s=session(request); require(s,'hippocampus','occipital')
    im=ml.decode_image(body.image); im.thumbnail((512,512))
    desc=body.caption or ml.caption(im); feature=ml.image_features(im); image=ml.encode_image(im)
    ident=uuid.uuid4().hex
    with sqlite3.connect(DB) as c: c.execute('INSERT INTO memories VALUES(?,?,?,?,?)',(ident,sid,desc,image,feature.astype('float32').tobytes()))
    return {'id':ident,'caption':desc,'image':image}

@app.get('/api/memories')
def memories(request:Request):
    sid,_=session(request)
    with sqlite3.connect(DB) as c: rows=c.execute('SELECT id,caption,image FROM memories WHERE session=?',(sid,)).fetchall()
    return [dict(zip(['id','caption','image'],r)) for r in rows]

class Dream(BaseModel):
    theme:str=Field(default='a surreal floating world',max_length=300)
    ids:list[str]=Field(default_factory=list,max_length=10)
    seed:int=Field(default=21,ge=0,le=2**31-1)
@app.post('/api/dream')
def dream(body:Dream,request:Request):
    sid,s=session(request); require(s,'association','frontal','pineal')
    with sqlite3.connect(DB) as c: rows=c.execute('SELECT id,caption,image,feature FROM memories WHERE session=?',(sid,)).fetchall()
    if not rows: raise HTTPException(400,'Save some image memories first.')
    if body.ids: rows=[r for r in rows if r[0] in body.ids]
    if not rows: raise HTTPException(400,'Selected memories are unavailable.')
    query=ml.text_features([body.theme])[0]
    ranked=sorted(rows,key=lambda r:float(np.frombuffer(r[3],dtype='float32')@query),reverse=True)[:3]
    captions='; '.join(r[1] for r in ranked)
    prompt=ml.language(f'Write one vivid image prompt combining ALL these memories: {captions}. Theme: {body.theme}. Keep it under 60 words. Describe the scene directly.',100)
    image=ml.generate_image(prompt,body.seed)
    return {'prompt':prompt,'image':image,'memories':[{'id':r[0],'caption':r[1],'image':r[2]} for r in ranked],'seed':body.seed,'model':'SD-Turbo, 2 inference steps'}

class Telephone(BaseModel):
    text:str=Field(min_length=1,max_length=500)
    rounds:int=Field(default=3,ge=1,le=3)
    seed:int=Field(default=21,ge=0,le=2**31-1)
@app.post('/api/telephone')
def telephone(body:Telephone,request:Request):
    _,s=session(request); require(s,'frontal','association','motor','occipital','wernicke')
    if body.rounds>1: require(s,'callosum')
    text=body.text; rounds=[]
    for i in range(body.rounds):
        instructions=ml.language(f'Turn this description into a short, clear image-generation prompt. Include its objects, actions, and accessories. Description: {text}. Output only the drawing prompt.',90)
        image=ml.generate_image('a clear colorful illustration of '+instructions,body.seed+i)
        # The recognizer receives ONLY pixels, never text or prior rounds.
        recognized=ml.caption(ml.decode_image(image))
        rounds.append({'input':text,'instructions':instructions,'image':image,'recognized':recognized}); text=recognized
    return {'original':body.text,'rounds':rounds,'final':text,'recognizer_input':'image pixels only','memory_saved':s['modules']['hippocampus']}

@app.post('/api/speech')
async def speech(request:Request,audio:UploadFile=File(...)):
    _,s=session(request); require(s,'temporal','wernicke')
    raw=await audio.read(5_000_001)
    if len(raw)>5_000_000: raise HTTPException(400,'Audio file is too large.')
    with tempfile.NamedTemporaryFile(suffix='.webm') as f:
        f.write(raw); f.flush(); text=await __import__('starlette.concurrency',fromlist=['run_in_threadpool']).run_in_threadpool(ml.transcribe,f.name)
    return {'text':text,'model':'Whisper tiny.en'}

class CinemaFrame(BaseModel):
    image:str
    audio_on:bool=False
    text:str=Field(default='',max_length=300)
@app.post('/api/cinema/frame')
def cinema_frame(body:CinemaFrame,request:Request):
    _,s=session(request); require(s,'occipital')
    labels=['a colorful moving geometric animation','a person talking','a natural landscape','text on a screen','a dark empty screen']
    candidates,_=ml.classify_image(ml.decode_image(body.image),labels)
    return {'candidates':candidates,'activity':{'occipital':1,'temporal':int(body.audio_on and s['modules']['temporal']),'wernicke':int(bool(body.text) and s['modules']['wernicke']),'association':int(s['modules']['association'])},'kind':'Measured software execution indicators, not predicted human brain responses.'}

@app.post('/api/cinema/import')
async def cinema_import(request:Request,archive:UploadFile=File(...)):
    _,s=session(request); data=await archive.read(80_000_001)
    if len(data)>80_000_000: raise HTTPException(400,'Prediction archive exceeds 80 MB.')
    try:result=tribe_service.payload(io.BytesIO(data))
    except (KeyError,TypeError,ValueError) as e:raise HTTPException(400,'Expected a valid TRIBE archive with predictions, cortical coordinates, and metadata.') from e
    s['cinema']=result
    return result

@app.get('/api/tribe/emotions')
def tribe_emotions():
    return {'emotions':emotion_study.EMOTIONS,'source':emotion_study.SOURCE,
            'note':'These are intended clip categories chosen by you, not emotions decoded by TRIBE.'}

@app.get('/api/tribe/study')
def tribe_study(request:Request):return tribe_service.study_jobs(session(request)[0])

@app.get('/api/tribe/collection')
def tribe_collection():return clip_library.collection()

@app.get('/api/tribe/collection/{ident}/media')
def tribe_collection_media(ident:str):return FileResponse(clip_library.media_path(ident),media_type='video/mp4')

@app.get('/api/tribe/collection/{ident}/poster')
def tribe_collection_poster(ident:str):return FileResponse(clip_library.media_path(ident,poster=True),media_type='image/jpeg')

@app.post('/api/tribe/collection/prepare')
def tribe_collection_prepare_all(request:Request):
    _,s=session(request);require(s,'occipital','temporal')
    return {'jobs':clip_library.prepare_all()}

@app.post('/api/tribe/collection/{ident}/prepare')
def tribe_collection_prepare(ident:str,request:Request):
    _,s=session(request);require(s,'occipital','temporal')
    return clip_library.prepare(ident)

@app.get('/api/tribe/media/{ident}')
def tribe_media(ident:str):return FileResponse(tribe_service.media_path(ident))

@app.get('/api/tribe/status')
def tribe_status():return tribe_service.status()

@app.post('/api/tribe/jobs')
async def tribe_start(request:Request,mode:str='visual',title:str='',emotion:str='mixed',media:UploadFile=File(...)):
    sid,s=session(request);require(s,'occipital' if mode!='audio' else 'temporal')
    study=emotion_study.label(title or Path(media.filename or 'My clip').stem,emotion)
    raw=await media.read(100_000_001)
    if len(raw)>100_000_000:raise HTTPException(400,'Use a clip under 100 MB.')
    source=tribe_service.save_upload(raw,Path(media.filename or '').suffix.lower())
    return tribe_service.start(source,mode,study=study,owner_session=sid)

@app.get('/api/tribe/jobs/{ident}')
def tribe_job(ident:str):return tribe_service.get(ident)

@app.get('/api/tribe/results/{ident}')
def tribe_result(ident:str):
    result=tribe_service.payload(tribe_service.archive(ident))
    if ident not in ['visual','audiovisual']:
        study=tribe_service.get(ident).get('study')
        if study:result['study']=study
    return result

@app.post('/api/tribe/demo/{mode}')
def tribe_demo(mode:str):
    if mode not in ['visual','audiovisual']:raise HTTPException(400,'Unknown demo mode.')
    path=tribe_service.archive(mode) if (tribe_service.ROOT/'data'/'tribe_demo'/(mode+'.npz')).exists() else None
    if path:return {'state':'complete','result_id':mode}
    status_path=tribe_service.ROOT/'data'/'tribe_demo'/(mode+'.status.json')
    if status_path.exists():
        existing=tribe_service.get('demo_'+mode)
        if existing['state']=='running':return existing
    return tribe_service.start(tribe_service.ROOT/'data'/'tribe_demo'/'motion_with_tone.mp4',mode)

@app.get('/api/tribe/demo-media')
def tribe_demo_media():return FileResponse(tribe_service.ROOT/'data'/'tribe_demo'/'motion_with_tone.mp4',media_type='video/mp4')

def mindeye_for(request):
    state=session(request)[1]
    with STATE_LOCK:
        if state.get('mindeye') is None:state['mindeye']=mindeye.MindsEye()
        return state['mindeye']

@app.get('/api/mindeye')
def mindeye_information():return mindeye.information()

@app.get('/api/mindeye/state')
def mindeye_state(request:Request):return mindeye_for(request).snapshot()

class MindEyeRound(BaseModel):
    training:Literal['one-hour','forty-hours']='one-hour'
    choices:Literal[4,8]=4

@app.post('/api/mindeye/round')
def mindeye_round(body:MindEyeRound,request:Request):
    mindeye.information()
    return mindeye_for(request).begin(body.training,body.choices)

class MindEyeAnswer(BaseModel):
    round_id:str=Field(max_length=32,min_length=32)
    choice_id:str=Field(max_length=32,min_length=32)

@app.post('/api/mindeye/answer')
def mindeye_answer(body:MindEyeAnswer,request:Request):return mindeye_for(request).answer(body.round_id,body.choice_id)

class MindEyeTraining(BaseModel):
    round_id:str=Field(max_length=32,min_length=32)
    training:Literal['one-hour','forty-hours']

@app.post('/api/mindeye/training')
def mindeye_training(body:MindEyeTraining,request:Request):return mindeye_for(request).training(body.round_id,body.training)

@app.post('/api/mindeye/reset')
def mindeye_reset(request:Request):return mindeye_for(request).reset()

@app.get('/api/mindeye/examples/{index}')
def mindeye_comparison(index:int):return mindeye.comparison(index)

@app.get('/api/mindeye/live/status')
def mindeye_live_status():return mindeye_live.status()

@app.get('/api/mindeye/live/jobs')
def mindeye_live_jobs(request:Request):return mindeye_live.jobs(session(request)[0])

@app.post('/api/mindeye/live/jobs')
async def mindeye_live_start(request:Request,fmri:UploadFile=File(...),access_confirmed:bool=Form(False)):
    raw=await fmri.read(2_000_001)
    return mindeye_live.start(raw,Path(fmri.filename or '').suffix.lower(),session(request)[0],access_confirmed)

@app.get('/api/mindeye/live/jobs/{ident}')
def mindeye_live_job(ident:str,request:Request):return mindeye_live.public(mindeye_live.get(ident,session(request)[0]))

@app.post('/api/mindeye/live/examples/{ident}')
def mindeye_live_example(ident:str,request:Request):return mindeye_live.start_recorded(ident,session(request)[0])

@app.get('/api/mindeye/live/media/{ident}')
def mindeye_live_image(ident:str,token:str=''):return FileResponse(mindeye_live.media(ident,token),media_type='image/png')

@app.post('/api/reset/modules')
def reset_modules(request:Request):
    _,s=session(request); s['modules']=dict(DEFAULT_MODULES); return s['modules']

def mario_for(request):
    sid,s=session(request)
    with STATE_LOCK:
        if s.get('mario') is None:s['mario']=MarioLab(sid)
        return s['mario']

class MarioReset(BaseModel):
    population:int=Field(32,ge=1,le=128)
    course:Literal['campaign','practice','gaps','original','bowser','level-1','level-2','level-3','level-4','level-5','level-6','level-7','level-8','level-9','level-10','level-11','level-12','level-13','level-14','level-15']='campaign'
    seed:int=Field(7,ge=0,le=1000000)
class MarioTrain(BaseModel):
    ticks:int=Field(30,ge=1,le=1200)
class MarioConfig(BaseModel):
    exploration:float=Field(.35,ge=0,le=1)
    death_penalty:float=Field(12,ge=0,le=50)
class MarioProgression(BaseModel):
    auto_advance:bool=True
class MarioHuman(BaseModel):
    action:int=Field(4,ge=0,le=5)

@app.get('/api/mario')
def mario_state(request:Request):return mario_for(request).state()
@app.post('/api/mario/reset')
def mario_reset(body:MarioReset,request:Request):return mario_for(request).reset(body.population,body.course,body.seed)
@app.post('/api/mario/train')
def mario_train(body:MarioTrain,request:Request):return mario_for(request).train(body.ticks)
@app.post('/api/mario/config')
def mario_config(body:MarioConfig,request:Request):return mario_for(request).config(body.exploration,body.death_penalty)
@app.post('/api/mario/progression')
def mario_progression(body:MarioProgression,request:Request):return mario_for(request).progression(body.auto_advance)
@app.post('/api/mario/next')
def mario_next(request:Request):return mario_for(request).next_level()
@app.post('/api/mario/human/next')
def mario_human_next(request:Request):return mario_for(request).human_next()
@app.post('/api/mario/evaluate')
def mario_evaluate(request:Request):return mario_for(request).evaluate()
@app.post('/api/mario/human/reset')
def mario_human_reset(request:Request):return mario_for(request).human_reset()
@app.post('/api/mario/human')
def mario_human(body:MarioHuman,request:Request):return mario_for(request).human(body.action)
@app.post('/api/mario/save')
def mario_save(request:Request):return mario_for(request).save()
@app.post('/api/mario/load')
def mario_load(request:Request):return mario_for(request).load()

ArcadeGame=Literal['flappy','crossy']
def arcade_rl_for(request,game):
    sid,s=session(request)
    with STATE_LOCK:
        labs=s.setdefault('arcade_rl',{})
        if game not in labs:labs[game]=ArcadeRL(game,sid,course_length=None)
        return labs[game]

class ArcadeRLReset(BaseModel):
    population:int=Field(32,ge=1,le=128)
    difficulty:Literal['easy','standard']='easy'
    seed:int=Field(7,ge=0,le=1000000)
    course_length:int|None=Field(None,ge=1,le=1000000)
class ArcadeRLTrain(BaseModel):
    ticks:int=Field(30,ge=1,le=1200)
class ArcadeRLConfig(BaseModel):
    exploration:float=Field(.2,ge=0,le=1)
    auto_exploration:bool|None=None
class ArcadeRLHuman(BaseModel):
    action:int=Field(0,ge=0,le=4)

@app.get('/api/arcade-rl/{game}')
def arcade_rl_state(game:ArcadeGame,request:Request):return arcade_rl_for(request,game).state()
@app.post('/api/arcade-rl/{game}/reset')
def arcade_rl_reset(game:ArcadeGame,body:ArcadeRLReset,request:Request):return arcade_rl_for(request,game).reset(body.population,body.difficulty,body.seed,body.course_length)
@app.post('/api/arcade-rl/{game}/train')
def arcade_rl_train(game:ArcadeGame,body:ArcadeRLTrain,request:Request):return arcade_rl_for(request,game).train(body.ticks)
@app.post('/api/arcade-rl/{game}/config')
def arcade_rl_config(game:ArcadeGame,body:ArcadeRLConfig,request:Request):return arcade_rl_for(request,game).config(body.exploration,body.auto_exploration)
@app.post('/api/arcade-rl/{game}/evaluate')
def arcade_rl_evaluate(game:ArcadeGame,request:Request):return arcade_rl_for(request,game).evaluate()
@app.post('/api/arcade-rl/{game}/human/reset')
def arcade_rl_human_reset(game:ArcadeGame,request:Request):return arcade_rl_for(request,game).human_reset()
@app.post('/api/arcade-rl/{game}/human')
def arcade_rl_human(game:ArcadeGame,body:ArcadeRLHuman,request:Request):
    if game=='flappy' and body.action>1:raise ValueError('Flappy Bird has only Coast and Flap actions.')
    return arcade_rl_for(request,game).human(body.action)
@app.post('/api/arcade-rl/{game}/save')
def arcade_rl_save(game:ArcadeGame,request:Request):return arcade_rl_for(request,game).save()
@app.post('/api/arcade-rl/{game}/load')
def arcade_rl_load(game:ArcadeGame,request:Request):return arcade_rl_for(request,game).load()

@app.get('/')
def index(): return FileResponse(ROOT/'static'/'index.html')

app.mount('/static',StaticFiles(directory=ROOT/'static'),name='static')

if __name__=='__main__':
    import uvicorn
    uvicorn.run(app,host='127.0.0.1',port=8769)