import pyxel
import math
import random

# ─── SETTINGS ────────────────────────────────────────────────────

WIDTH  = 256
HEIGHT = 256
FPS    = 60

LANES       = [48, 88, 128, 168, 208]
LANE_KEYS   = [pyxel.KEY_D, pyxel.KEY_F, pyxel.KEY_G, pyxel.KEY_J, pyxel.KEY_K]
LANE_LABELS = ["D", "F", "G", "J", "K"]

HIT_Y      = 215
NOTE_SPEED = 110
HIT_WIN    = {"perfect": 0.055, "great": 0.11, "good": 0.20}

COL_BG    = 0
COL_WHITE = 7
COL_GRAY  = 6
COL_YELLOW= 9

LANE_COLS  = [14, 12, 3, 11, 9]
JUDGE_COLS = {"PERFECT": 9, "GREAT": 12, "GOOD": 11, "MISS": 8}

# ─── STORY ───────────────────────────────────────────────────────

CHAR_COLS  = {"hana": 14, "kai": 12, "mio": 3, "ryu": 11}
CHAR_NAMES = {"hana": "Hana", "kai": "Kai", "mio": "Mio", "ryu": "Ryu"}

STORY = [
    {
        "chapter": 1, "title": "Beginnings", "level": "easy",
        "scenes": [
            ("hana", "Hey... you're the new kid who transferred here?"),
            ("kai",  "Heard you used to play rhythm games. That true?"),
            ("hana", "Our band PRISM needs a 5th member. We have a show soon."),
            ("mio",  "No pressure... but can you actually keep up?"),
        ],
        "after": [
            ("hana", "...You're actually good. Like, really good."),
            ("kai",  "Welcome to PRISM, I guess. Don't make us regret it."),
            ("ryu",  "First practice tomorrow. Don't. Be. Late."),
        ]
    },
    {
        "chapter": 2, "title": "The Rival", "level": "hard",
        "scenes": [
            ("mio",  "ECLIPSE won nationals last year. They'll be at our venue."),
            ("kai",  "So what? We've got heart. They've got... okay they're good."),
            ("hana", "Their rhythm section is flawless. We need to step it up."),
            ("ryu",  "One way to get ready. Hard setlist. All of it. Now."),
        ],
        "after": [
            ("kai",  "Okay. OKAY! That is what I'm talking about!"),
            ("mio",  "We're actually doing this. I can't believe it."),
            ("hana", "Three weeks left. We go harder. Together."),
        ]
    },
    {
        "chapter": 3, "title": "Showtime", "level": "master",
        "scenes": [
            ("hana", "This is it. Stage lights on. I can see the crowd..."),
            ("ryu",  "We played our hearts out to get here. That's everything."),
            ("kai",  "Deep breath. We're PRISM. Tonight, we shine."),
            ("mio",  "...One, two, three, four --"),
        ],
        "after": [
            ("hana", "They're cheering. They're actually cheering for us!"),
            ("kai",  "WE DID IT! WE DID IT!"),
            ("mio",  "I'm going to cry. I am totally going to cry right now."),
            ("ryu",  "Same time next year? Same time next year."),
        ]
    },
]

# ─── CHARTS ──────────────────────────────────────────────────────

CHARTS = {
    "easy": [
        (1.0,0),(1.8,1),(2.6,2),(3.4,3),(4.5,0,1.2),
        (6.5,2),(7.3,1),(8.1,3),(8.9,0),
        (10.0,1),(10.8,2),(11.6,3),(12.4,0),(13.5,2,1.0),
        (15.5,0),(16.3,1),(17.1,2),(17.9,3),
        (19.0,3),(19.8,2),(20.6,1),(21.4,0),(22.5,1,1.2),
        (24.5,0),(25.3,3),(26.1,2),(26.9,1),
        (28.0,0),(28.8,1),(29.6,2),(30.4,3),
    ],
    "hard": [
        (1.0,0),(1.6,1),(2.2,2),(3.0,3),(3.8,1,1.0),
        (5.5,3),(6.0,2),(6.5,1),(7.2,0,1.3),
        (9.0,0),(9.4,1),(9.8,2),(10.2,3),
        (11.0,2),(11.4,1),(11.8,0),(12.2,3),
        (13.0,0,1.0),(14.5,3,1.0),
        (16.5,1),(17.0,2),(17.5,1),(18.0,0),
        (19.0,3),(19.3,2),(19.6,1),(19.9,0),(21.0,0,1.5),
        (23.5,1),(24.0,2),(24.5,3),(25.0,0),
        (26.0,2),(26.3,1),(26.6,3),(26.9,0),
        (28.0,1,1.2),(30.0,2,1.2),
    ],
    "master": [
        (0.8,0),(1.3,1),(1.8,2),(2.3,3),(3.0,0,1.2),
        (4.5,2),(5.0,3),(5.5,1),(6.2,1,1.4),
        (8.0,3),(8.5,2),(9.0,0),(9.3,1),(9.6,2),(9.9,3),
        (10.5,3),(10.8,2),(11.1,1),(11.4,0),
        (12.0,0,1.0),(13.5,3,1.0),
        (15.0,1),(15.3,0),(15.6,2),(15.9,3),
        (16.5,0),(16.7,1),(16.9,2),(17.1,3),(18.0,2,1.5),
        (20.0,0),(20.2,1),(20.4,2),(20.6,3),
        (21.0,3),(21.2,2),(21.4,1),(21.6,0),
        (22.5,1),(22.7,3),(22.9,0),(23.1,2),
        (24.0,0,1.0),(25.5,1,1.0),(27.0,2,1.0),
        (29.0,0),(29.2,1),(29.4,2),(29.6,3),
        (30.0,3),(30.2,2),(30.4,1),(30.6,0),
    ],
}

# ─── NOTE ────────────────────────────────────────────────────────

class Note:
    def __init__(self, lane, time, length=0):
        self.lane      = lane
        self.time      = time
        self.length    = length
        self.hit       = False
        self.holding   = False
        self.completed = False

# ─── GAME ────────────────────────────────────────────────────────

class Game:

    def __init__(self):
        pyxel.init(WIDTH, HEIGHT, title="PRISM: Rhythm Story", fps=FPS)
        self._setup_sounds()

        self.state        = "title"
        self.menu_cursor  = 0
        self.chapter_idx  = 0
        self.scene_idx    = 0
        self.scene_phase  = "pre"
        self.freeplay_lvl = 0

        self.time         = 0.0
        self.notes        = []
        self.active_notes = []
        self.combo        = 0
        self.max_combo    = 0
        self.score        = 0
        self.hp           = 100
        self.judge_text   = ""
        self.judge_timer  = 0
        self.total_hits   = 0
        self.total_acc    = 0.0
        self.lane_flash   = [0] * 5

        self.tw_text   = ""
        self.tw_target = ""
        self.tw_timer  = 0

        self.particles   = []
        self.burst       = []
        self.grid_offset = 0
        self.anim_t      = 0
        self.beat_flash  = 0

        self.stars_far  = [(random.randint(0,WIDTH), random.randint(0,HEIGHT)) for _ in range(40)]
        self.stars_near = [(random.randint(0,WIDTH), random.randint(0,HEIGHT)) for _ in range(20)]

        pyxel.run(self.update, self.draw)

    # ─── SOUNDS ──────────────────────────────────────────────────
    # NOTE: pyxel 2.x max octave is 4. No octave-5 notes allowed.
    # sound.set(notes, tones, volumes, effects, speed)
    # tones: t=triangle s=square p=pulse n=noise
    # effects: n=none v=vibrato s=slide f=fadeout

    def _setup_sounds(self):
        # Melody easy
        pyxel.sounds[0].set(
            "e3 e3 g3 a3 b3 a3 g3 e3 d3 d3 f3 g3 a3 g3 f3 d3",
            "tttttttttttttttt", "5566554455665544", "nnnnnnnnnnnnnnnn", 20)
        pyxel.sounds[13].set(
            "c3 c3 e3 g3 a3 g3 e3 c3 b2 b2 d3 f3 g3 f3 d3 b2",
            "tttttttttttttttt", "5566554455665544", "nnnnnnnnnnnnnnnn", 20)
        # Melody hard
        pyxel.sounds[1].set(
            "a3 c4 e4 c4 a3 g3 f3 g3 a3 b3 c4 b3 a3 g3 f3 e3",
            "ssssssssssssssss", "6677665566776655", "nnvnnnvnnnvnnnvn", 18)
        pyxel.sounds[14].set(
            "d4 e4 f4 e4 d4 c4 b3 a3 g3 a3 b3 c4 d4 e4 f4 e4",
            "ssssssssssssssss", "6677665566776655", "nnvnnnvnnnvnnnvn", 18)
        # Melody master (octave capped at 4)
        pyxel.sounds[2].set(
            "e4 g4 b4 g4 e4 d4 c4 b3 a3 c4 e4 c4 a3 g3 f3 e3",
            "pppppppppppppppp", "7777666677776666", "nnnsnnnsnnnsnnns", 14)
        pyxel.sounds[15].set(
            "f4 a4 c4 a4 f4 e4 d4 c4 b3 d4 f4 d4 b3 a3 g3 f3",
            "pppppppppppppppp", "7777666677776666", "nnnsnnnsnnnsnnns", 14)
        # Bass easy
        pyxel.sounds[3].set(
            "e2 r e2 r g2 r e2 r d2 r d2 r f2 r d2 r",
            "tttttttttttttttt", "6060606060606060", "nnnnnnnnnnnnnnnn", 20)
        pyxel.sounds[16].set(
            "c2 r c2 r e2 r c2 r b1 r b1 r d2 r b1 r",
            "tttttttttttttttt", "6060606060606060", "nnnnnnnnnnnnnnnn", 20)
        # Bass hard
        pyxel.sounds[4].set(
            "a2 a2 r a2 f2 f2 r f2 g2 g2 r g2 e2 e2 r e2",
            "ssssssssssssssss", "5500550055005500", "nnnnnnnnnnnnnnnn", 18)
        pyxel.sounds[17].set(
            "d2 d2 r d2 b1 b1 r b1 c2 c2 r c2 a1 a1 r a1",
            "ssssssssssssssss", "5500550055005500", "nnnnnnnnnnnnnnnn", 18)
        # Bass master
        pyxel.sounds[5].set(
            "e2 e2 e2 r b1 b1 b1 r a1 a1 a1 r c2 c2 c2 r",
            "pppppppppppppppp", "6660666066606660", "nnnnnnnnnnnnnnnn", 14)
        pyxel.sounds[18].set(
            "f2 f2 f2 r d2 d2 d2 r c2 c2 c2 r e2 e2 e2 r",
            "pppppppppppppppp", "6660666066606660", "nnnnnnnnnnnnnnnn", 14)
        # Arpeggios
        pyxel.sounds[19].set(
            "e3 g3 b3 e4 g3 b3 e4 g4 d3 f3 a3 d4 f3 a3 d4 f4",
            "tttttttttttttttt", "4444444444444444", "nnnnnnnnnnnnnnnn", 20)
        pyxel.sounds[20].set(
            "a3 c4 e4 a4 c4 e4 a4 c4 f3 a3 c4 f4 a3 c4 f4 a4",
            "pppppppppppppppp", "4444444444444444", "nnnnnnnnnnnnnnnn", 18)
        pyxel.sounds[21].set(
            "e4 g4 b4 e4 b4 g4 e4 b3 a3 c4 e4 a4 e4 c4 a3 e3",
            "ssssssssssssssss", "5555555555555555", "nnnnnnnnnnnnnnnn", 14)
        # Percussion
        pyxel.sounds[6].set(
            "c1 r c1 r c1 r c1 r c1 r c1 r c1 r c1 r",
            "nnnnnnnnnnnnnnnn", "7070707070707070", "nnnnnnnnnnnnnnnn", 20)
        pyxel.sounds[7].set(
            "c1 r c1 c1 r c1 r c1 c1 r c1 c1 r c1 r c1",
            "nnnnnnnnnnnnnnnn", "7070707070707070", "nnnnnnnnnnnnnnnn", 16)
        # Menu music
        pyxel.sounds[22].set(
            "c3 e3 g3 c4 e3 g3 c4 e4 a2 c3 e3 a3 c3 e3 a3 c4",
            "tttttttttttttttt", "5555555555555555", "nnnnnnnnnnnnnnnn", 25)
        pyxel.sounds[23].set(
            "c2 r r r a1 r r r f1 r r r g1 r r r",
            "tttttttttttttttt", "6000600060006000", "nnnnnnnnnnnnnnnn", 25)
        # SFX (perfect: rising two-note chime)
        pyxel.sounds[8].set("a4 e4",  "tt", "75", "nf", 8)
        pyxel.sounds[9].set("g4",     "t",  "6",  "f",  8)
        pyxel.sounds[10].set("e4",    "s",  "5",  "f",  8)
        pyxel.sounds[11].set("c2 b1", "nn", "53", "nf", 8)
        pyxel.sounds[12].set("c4",    "t",  "5",  "f",  6)

        # Music tracks
        pyxel.musics[0].set([0,13,0,13], [3,16,3,16], [19,19,19,19], [6,6,6,6])
        pyxel.musics[1].set([1,14,1,14], [4,17,4,17], [20,20,20,20], [7,7,7,7])
        pyxel.musics[2].set([2,15,2,15], [5,18,5,18], [21,21,21,21], [7,7,7,7])
        pyxel.musics[3].set([22,22,22,22], [23,23,23,23], [], [])

        pyxel.playm(3, loop=True)

    # ─── LEVEL MANAGEMENT ────────────────────────────────────────

    def _music_idx(self, level):
        return {"easy": 0, "hard": 1, "master": 2}[level]

    def start_game(self, level):
        self.notes        = [Note(n[1], n[0], n[2] if len(n)>2 else 0) for n in CHARTS[level]]
        self.active_notes = []
        self.time         = 0.0
        self.combo        = 0
        self.max_combo    = 0
        self.score        = 0
        self.hp           = 100
        self.total_hits   = 0
        self.total_acc    = 0.0
        self.judge_text   = ""
        self.judge_timer  = 0
        self.lane_flash   = [0] * 5
        self.particles    = []
        self.burst        = []
        self.grid_offset  = 0
        self.state        = "game"
        pyxel.stop()
        pyxel.playm(self._music_idx(level), loop=True)

    # ─── HELPERS ─────────────────────────────────────────────────

    def _note_y(self, note):
        return int(HIT_Y - (note.time - self.time) * NOTE_SPEED)

    def _acc(self):
        return 0.0 if self.total_hits == 0 else self.total_acc / self.total_hits * 100

    def _rank(self, a):
        if a > 97: return "S+"
        if a > 92: return "S"
        if a > 85: return "A"
        if a > 70: return "B"
        if a > 50: return "C"
        return "D"

    def _spawn_particles(self, x, y, col, count=8):
        for _ in range(count):
            ang = random.uniform(0, math.pi * 2)
            spd = random.uniform(20, 80)
            self.particles.append({
                "x": x, "y": y,
                "vx": math.cos(ang)*spd, "vy": math.sin(ang)*spd - 10,
                "life": 1.0, "col": col, "size": random.choice([1,1,2])
            })

    def _play_sfx(self, result):
        sfx = {"PERFECT": 8, "GREAT": 9, "GOOD": 10, "MISS": 11}
        pyxel.play(3, sfx.get(result, 11))

    def _start_tw(self, txt):
        self.tw_target = txt
        self.tw_text   = ""
        self.tw_timer  = 0

    # ─── JUDGE ───────────────────────────────────────────────────

    def _judge_lane(self, lane):
        best, best_diff = None, 999.0
        for n in self.active_notes:
            if n.lane != lane or n.hit:
                continue
            d = abs(n.time - self.time)
            if d < best_diff:
                best_diff, best = d, n
        if best and best_diff < HIT_WIN["good"] + 0.05:
            self._apply_judge(best_diff, best)
        else:
            self.combo = 0
            self.hp = max(0, self.hp - 6)
            self.judge_text  = "MISS"
            self.judge_timer = 40
            self._play_sfx("MISS")

    def _apply_judge(self, diff, note, is_long=False):
        if diff < HIT_WIN["perfect"]:
            result, pts, acc = "PERFECT", 400 if is_long else 300, 1.0
        elif diff < HIT_WIN["great"]:
            result, pts, acc = "GREAT",   250 if is_long else 200, 0.7
        elif diff < HIT_WIN["good"]:
            result, pts, acc = "GOOD",    120 if is_long else 100, 0.4
        else:
            result, pts, acc = "MISS", 0, 0.0

        self.judge_text  = result
        self.judge_timer = 40
        self.total_hits += 1
        self.total_acc  += acc

        if result != "MISS":
            self.combo += 1
            bonus       = 1 + (self.combo // 20) * 0.1
            self.score += int(pts * bonus)
            note.hit    = True
            if note.length > 0:
                note.holding = True
            self._spawn_particles(LANES[note.lane], HIT_Y, LANE_COLS[note.lane],
                                  12 if result=="PERFECT" else 7)
            self.lane_flash[note.lane] = 12
            if result == "PERFECT":
                self.beat_flash = 6
        else:
            self.combo = 0
            self.hp    = max(0, self.hp - 6)

        self.max_combo = max(self.max_combo, self.combo)
        self._play_sfx(result)

    def _update_long_notes(self):
        for n in self.active_notes:
            if not n.holding:
                continue
            if pyxel.btn(LANE_KEYS[n.lane]):
                if self.time >= n.time + n.length:
                    self._apply_judge(abs(self.time - n.time - n.length), n, is_long=True)
                    n.holding = False; n.completed = True
            else:
                self.combo = 0; self.hp = max(0, self.hp-10)
                self.judge_text  = "MISS"; self.judge_timer = 40
                n.holding = False; n.completed = True
                self._play_sfx("MISS")

    def _init_burst(self):
        self.burst = []
        for i in range(60):
            ang = random.uniform(0, math.pi*2)
            spd = random.uniform(30, 120)
            self.burst.append({
                "x": WIDTH//2, "y": HEIGHT//2,
                "vx": math.cos(ang)*spd, "vy": math.sin(ang)*spd,
                "life": 1.0, "col": LANE_COLS[i%5], "size": random.choice([1,1,2,2,3])
            })

    # ─── UPDATE ──────────────────────────────────────────────────

    def update(self):
        if len(self.tw_text) < len(self.tw_target):
            self.tw_timer += 1
            if self.tw_timer % 2 == 0:
                self.tw_text = self.tw_target[:len(self.tw_text)+1]
        self.anim_t += 1

        if   self.state == "title":         self._update_title()
        elif self.state == "story_menu":    self._update_story_menu()
        elif self.state == "freeplay_menu": self._update_freeplay_menu()
        elif self.state == "scene":         self._update_scene()
        elif self.state == "game":          self._update_game()
        else:                               self._update_result()

    def _update_title(self):
        if pyxel.btnp(pyxel.KEY_UP):
            self.menu_cursor = (self.menu_cursor-1)%2; pyxel.play(3,12)
        if pyxel.btnp(pyxel.KEY_DOWN):
            self.menu_cursor = (self.menu_cursor+1)%2; pyxel.play(3,12)
        if pyxel.btnp(pyxel.KEY_SPACE) or pyxel.btnp(pyxel.KEY_RETURN):
            pyxel.play(3,12); pyxel.stop()
            if self.menu_cursor == 0:
                self.state = "story_menu"; self.chapter_idx = 0
            else:
                self.state = "freeplay_menu"; self.freeplay_lvl = 0

    def _update_story_menu(self):
        if pyxel.btnp(pyxel.KEY_UP):
            self.chapter_idx = (self.chapter_idx-1)%len(STORY); pyxel.play(3,12)
        if pyxel.btnp(pyxel.KEY_DOWN):
            self.chapter_idx = (self.chapter_idx+1)%len(STORY); pyxel.play(3,12)
        if pyxel.btnp(pyxel.KEY_SPACE) or pyxel.btnp(pyxel.KEY_RETURN):
            pyxel.play(3,12)
            self.scene_idx=0; self.scene_phase="pre"
            self._start_tw(STORY[self.chapter_idx]["scenes"][0][1])
            self.state = "scene"
        if pyxel.btnp(pyxel.KEY_ESCAPE):
            self.state = "title"; pyxel.playm(3, loop=True)

    def _update_freeplay_menu(self):
        if pyxel.btnp(pyxel.KEY_UP):
            self.freeplay_lvl = (self.freeplay_lvl-1)%3; pyxel.play(3,12)
        if pyxel.btnp(pyxel.KEY_DOWN):
            self.freeplay_lvl = (self.freeplay_lvl+1)%3; pyxel.play(3,12)
        if pyxel.btnp(pyxel.KEY_SPACE) or pyxel.btnp(pyxel.KEY_RETURN):
            pyxel.play(3,12)
            self.scene_phase = "free"
            self.start_game(["easy","hard","master"][self.freeplay_lvl])
        if pyxel.btnp(pyxel.KEY_ESCAPE):
            self.state = "title"; pyxel.playm(3, loop=True)

    def _update_scene(self):
        if pyxel.btnp(pyxel.KEY_SPACE) or pyxel.btnp(pyxel.KEY_RETURN):
            if self.tw_text != self.tw_target:
                self.tw_text = self.tw_target; return
            pyxel.play(3,12)
            ch     = STORY[self.chapter_idx]
            scenes = ch["scenes"] if self.scene_phase=="pre" else ch["after"]
            self.scene_idx += 1
            if self.scene_idx >= len(scenes):
                if self.scene_phase == "pre":
                    self.start_game(ch["level"])
                else:
                    self.state = "story_menu"; pyxel.playm(3, loop=True)
            else:
                self._start_tw(scenes[self.scene_idx][1])
        if pyxel.btnp(pyxel.KEY_ESCAPE):
            self.state = "story_menu"

    def _update_game(self):
        self.time += 1/FPS
        self.grid_offset = (self.grid_offset + NOTE_SPEED/FPS) % 24
        if self.beat_flash > 0: self.beat_flash -= 1

        spawn_off = HIT_Y / NOTE_SPEED
        for n in self.notes:
            if n not in self.active_notes and (n.time - spawn_off) <= self.time:
                self.active_notes.append(n)

        for n in self.active_notes:
            if not n.hit and not n.completed:
                if self._note_y(n) > HEIGHT + 8:
                    n.hit=True; self.combo=0; self.hp=max(0,self.hp-8)
                    self.judge_text="MISS"; self.judge_timer=40
                    self._play_sfx("MISS")

        for i, key in enumerate(LANE_KEYS):
            if pyxel.btnp(key):
                self._judge_lane(i)

        self._update_long_notes()

        for i in range(5):
            if self.lane_flash[i] > 0: self.lane_flash[i] -= 1
        if self.judge_timer > 0: self.judge_timer -= 1

        for p in self.particles[:]:
            p["x"] += p["vx"]/FPS; p["y"] += p["vy"]/FPS
            p["vy"] += 90/FPS;     p["life"] -= 1.8/FPS
            if p["life"] <= 0: self.particles.remove(p)

        if self.hp <= 0:
            pyxel.stop(); self._init_burst(); self.state="fail"; return

        all_done = all(n.hit or n.completed for n in self.notes)
        last_t   = max(n.time + n.length for n in self.notes)
        if (all_done and self.time > last_t) or self.time > last_t + 2:
            pyxel.stop(); self._init_burst(); self.state="clear"

    def _update_result(self):
        for p in self.burst[:]:
            p["x"] += p["vx"]/FPS; p["y"] += p["vy"]/FPS
            p["vy"] += 50/FPS;     p["life"] -= 0.6/FPS
            if p["life"] <= 0: self.burst.remove(p)

        if pyxel.btnp(pyxel.KEY_SPACE) or pyxel.btnp(pyxel.KEY_RETURN):
            pyxel.play(3,12)
            if self.scene_phase == "free":
                self.state = "freeplay_menu"; pyxel.playm(3, loop=True)
            elif self.state == "clear":
                ch=STORY[self.chapter_idx]; self.scene_phase="post"
                self.scene_idx=0; self._start_tw(ch["after"][0][1]); self.state="scene"
            else:
                self.state="story_menu"; pyxel.playm(3, loop=True)

    # ─── DRAW HELPERS ────────────────────────────────────────────

    def _draw_stars(self, speed=1.0):
        for i,(x,y) in enumerate(self.stars_far):
            ny = int(y + self.anim_t * 0.15 * speed) % HEIGHT
            self.stars_far[i] = (x, ny)
            pyxel.pset(x, ny, 1)
        for i,(x,y) in enumerate(self.stars_near):
            ny = int(y + self.anim_t * 0.4 * speed) % HEIGHT
            self.stars_near[i] = (x, ny)
            pyxel.pset(x, ny, 6 if i%3!=0 else 7)

    # ─── DRAW ────────────────────────────────────────────────────

    def draw(self):
        pyxel.cls(0)
        if   self.state == "title":         self._draw_title()
        elif self.state == "story_menu":    self._draw_story_menu()
        elif self.state == "freeplay_menu": self._draw_freeplay_menu()
        elif self.state == "scene":         self._draw_scene()
        elif self.state == "game":          self._draw_game()
        else:                               self._draw_result()

    def _draw_title(self):
        self._draw_stars()

        # Rainbow bar
        for i in range(WIDTH):
            c = LANE_COLS[int((i/WIDTH*5 + self.anim_t*0.02))%5]
            pyxel.line(i, 46, i, 49, c)

        # Title text with pulse
        col = 14 if (self.anim_t//20)%2 else 9
        pyxel.text(WIDTH//2-19, 21, "PRISM", 1)   # shadow
        pyxel.text(WIDTH//2-20, 20, "PRISM", col)
        pyxel.text(WIDTH//2-31, 33, "RHYTHM  STORY", 12)

        # Lane dots
        for i,lc in enumerate(LANE_COLS):
            dx = WIDTH//2 - 28 + i*14
            pyxel.circ(dx, 58, 3, lc)
            pyxel.circb(dx, 58, 4, lc)

        # Menu options
        opts = ["  STORY  MODE  ", "  FREE  PLAY  "]
        for i, label in enumerate(opts):
            y   = 90 + i*28
            sel = self.menu_cursor == i
            bc  = LANE_COLS[i*2]
            lw  = len(label)*4 + 4
            cx  = WIDTH//2 - lw//2
            if sel:
                pyxel.rect(cx-2, y-4, lw+4, 16, 1)
                pyxel.rectb(cx-2, y-4, lw+4, 16, bc)
                pyxel.line(cx-1, y-3, cx+lw+2, y-3, bc)
                pyxel.text(cx-8, y, ">", bc)
            pyxel.text(WIDTH//2 - len(label)*2, y, label, bc if sel else 6)

        pyxel.text(4, HEIGHT-18, "D F G J K  to  hit  notes", 5)
        pyxel.text(4, HEIGHT-10, "SPACE/ENTER select   ESC back", 5)

        # Floating note icons
        for i in range(4):
            nx = 10 + i*60
            ny = int(HEIGHT-35 + math.sin(self.anim_t*0.06 + i*1.5)*4)
            pyxel.text(nx, ny, "^", LANE_COLS[i])

    def _draw_story_menu(self):
        self._draw_stars()
        pyxel.rect(0,0,WIDTH,18,1)
        pyxel.text(WIDTH//2-20, 6, "STORY  MODE", 12)
        pyxel.line(0,17,WIDTH,17,5)

        for i,ch in enumerate(STORY):
            y   = 26 + i*68
            sel = self.chapter_idx == i
            bc  = LANE_COLS[i]
            if sel:
                pyxel.rect(4,y,WIDTH-8,60,1)
                pyxel.rectb(4,y,WIDTH-8,60,bc)
                pyxel.rect(4,y,3,60,bc)
            pyxel.text(14, y+6,  f"CHAPTER {ch['chapter']}", bc if sel else 5)
            pyxel.text(14, y+16, ch["title"].upper(), 7 if sel else 6)
            star_col = bc if sel else 5
            for s in range(i+1):
                sx = 14+s*10
                pyxel.circb(sx, y+30, 3, star_col)
                if sel: pyxel.circ(sx, y+30, 2, star_col)
            bx = WIDTH-60
            pyxel.rect(bx, y+22, 52, 12, bc if sel else 1)
            pyxel.text(bx+4, y+25, ch["level"].upper(), 0 if sel else bc)
            if sel: pyxel.text(WIDTH-32, y+44, "PLAY >", bc)

        pyxel.text(4, HEIGHT-10, "UP/DOWN  SPACE play  ESC back", 5)

    def _draw_freeplay_menu(self):
        self._draw_stars()
        pyxel.rect(0,0,WIDTH,18,1)
        pyxel.text(WIDTH//2-20, 6, "FREE  PLAY", 9)
        pyxel.line(0,17,WIDTH,17,5)

        lvls  = ["EASY",        "HARD",            "MASTER"]
        descs = ["Chill vibes", "Getting serious",  "No mercy"]
        bpms  = ["~120 BPM",   "~150 BPM",         "~180 BPM"]
        cols  = [11, 9, 14]

        for i in range(3):
            y   = 30 + i*66
            sel = self.freeplay_lvl == i
            bc  = cols[i]
            if sel:
                pyxel.rect(8,y,WIDTH-16,56,1)
                pyxel.rectb(8,y,WIDTH-16,56,bc)
                pyxel.rect(8,y,3,56,bc)
            for d in range(i+1):
                dx = WIDTH-30-d*14
                if sel: pyxel.circ(dx, y+28, 5, bc)
                else:   pyxel.circb(dx, y+28, 4, bc)
            pyxel.text(18, y+10, lvls[i],  bc if sel else 6)
            pyxel.text(18, y+22, descs[i], 7 if sel else 5)
            pyxel.text(18, y+36, bpms[i],  6 if sel else 1)

        pyxel.text(4, HEIGHT-10, "UP/DOWN  SPACE play  ESC back", 5)

    def _draw_scene(self):
        self._draw_stars(speed=0.3)
        ch      = STORY[self.chapter_idx]
        scenes  = ch["scenes"] if self.scene_phase=="pre" else ch["after"]
        if self.scene_idx >= len(scenes): return
        char_key, _ = scenes[self.scene_idx]
        char_name   = CHAR_NAMES[char_key]
        char_col    = CHAR_COLS[char_key]

        # Header
        pyxel.rect(0,0,WIDTH,14,1)
        pyxel.line(0,14,WIDTH,14,char_col)
        pyxel.text(4, 4, f"CH.{ch['chapter']}  {ch['title'].upper()}", char_col)
        pyxel.text(WIDTH-28, 4, f"{self.scene_idx+1}/{len(scenes)}", 6)

        # Character portrait
        cx, cy = WIDTH//2, 96
        pyxel.circb(cx, cy-10, 30, 1)
        pyxel.circb(cx, cy-20, 18, char_col)
        pyxel.circ( cx, cy-20, 14, char_col)
        pyxel.circ( cx, cy-20, 11, 1)
        # Eyes
        pyxel.pset(cx-5, cy-22, 7); pyxel.pset(cx+5, cy-22, 7)
        pyxel.pset(cx-5, cy-21, 7); pyxel.pset(cx+5, cy-21, 7)
        # Mouth per character
        mouths = {
            "hana": [(cx-2,cy-14),(cx,cy-14),(cx+2,cy-14)],
            "kai":  [(cx-3,cy-13),(cx-1,cy-14),(cx+1,cy-14),(cx+3,cy-13)],
            "mio":  [(cx-2,cy-14),(cx,cy-13),(cx+2,cy-14)],
            "ryu":  [(cx-3,cy-14),(cx+3,cy-14)],
        }
        for px,py2 in mouths.get(char_key,[]):
            pyxel.pset(px, py2, 7)
        # Body
        pyxel.rect(cx-14, cy-6, 28, 32, char_col)
        pyxel.line(cx-14, cy-6, cx-14, cy+6, 7)
        pyxel.line(cx+13, cy-6, cx+13, cy+6, 7)
        pyxel.rect(cx-6, cy-34, 12, 4, char_col)
        # Name tag
        nw = len(char_name)*4+8
        pyxel.rect(cx-nw//2, cy+28, nw, 12, char_col)
        pyxel.text(cx-len(char_name)*2, cy+31, char_name, 0)

        # Dialogue box
        bx,by,bw,bh = 2, HEIGHT-78, WIDTH-4, 68
        pyxel.rect(bx+2, by+2, bw, bh, 0)
        pyxel.rect(bx, by, bw, bh, 1)
        pyxel.rectb(bx, by, bw, bh, char_col)
        pyxel.line(bx+1, by+1, bx+bw-2, by+1, char_col)

        # Word-wrapped typewriter
        words  = self.tw_text.split(" ")
        line   = ""
        line_y = by+8
        for w in words:
            test = (line+" "+w).strip()
            if len(test)*4 > bw-14 and line:
                pyxel.text(bx+7, line_y, line, 7)
                line_y += 9; line = w
            else:
                line = test
        if line: pyxel.text(bx+7, line_y, line, 7)

        # Advance prompt
        if self.tw_text == self.tw_target and (self.anim_t//20)%2:
            pyxel.text(bw-14, by+bh-10, ">>", char_col)

    def _draw_game(self):
        # Side panels
        pyxel.rect(0, 0, LANES[0]-11, HEIGHT, 0)
        pyxel.rect(LANES[4]+11, 0, WIDTH-LANES[4]-11, HEIGHT, 0)

        # Lanes
        for i,x in enumerate(LANES):
            fl = self.lane_flash[i]
            pyxel.rect(x-10, 0, 20, HEIGHT, 1)
            # Scrolling grid lines on lane
            off = int(self.grid_offset)
            for gy in range(off%24, HIT_Y, 24):
                pyxel.line(x-9, gy, x+9, gy, 0)
            # Flash effect
            if fl > 0:
                for fy in range(0, HIT_Y, 3):
                    if random.random() < fl/12.0 * 0.4:
                        pyxel.line(x-9, fy, x+9, fy, LANE_COLS[i])

        # Lane separators
        for i in range(len(LANES)-1):
            mx = (LANES[i]+LANES[i+1])//2
            pyxel.line(mx, 0, mx, HIT_Y, 0)

        # Hit line glow
        pyxel.line(LANES[0]-10, HIT_Y-1, LANES[4]+10, HIT_Y-1, 1)
        pyxel.line(LANES[0]-10, HIT_Y,   LANES[4]+10, HIT_Y,   7)
        pyxel.line(LANES[0]-10, HIT_Y+1, LANES[4]+10, HIT_Y+1, 1)

        # Hit circles
        for i,x in enumerate(LANES):
            held = pyxel.btn(LANE_KEYS[i])
            hc   = LANE_COLS[i] if held else 5
            pyxel.circb(x, HIT_Y, 9, hc)
            if held:
                pyxel.circ(x, HIT_Y, 7, hc)
                pyxel.circ(x, HIT_Y, 3, 7)
            else:
                pyxel.circb(x, HIT_Y, 6, 1)
            pyxel.text(x-2, HIT_Y+12, LANE_LABELS[i], hc)

        # Notes
        for n in self.active_notes:
            if n.completed: continue
            x   = LANES[n.lane]
            col = LANE_COLS[n.lane]
            y   = self._note_y(n)
            if n.length > 0:
                tail_h = int(n.length * NOTE_SPEED)
                pyxel.rect(x-5, y-tail_h, 10, tail_h, 5)
                pyxel.rectb(x-5, y-tail_h, 10, tail_h, col)
                if n.holding:
                    rem = max(0, min(int((n.time+n.length-self.time)*NOTE_SPEED), HIT_Y))
                    pyxel.rect(x-5, HIT_Y-rem, 10, rem, col)
                    if (self.anim_t%8)<4:
                        pyxel.line(x-5, HIT_Y-rem, x+5, HIT_Y-rem, 7)
            # Note head
            pyxel.rect(x-9, y-5, 18, 10, col)
            pyxel.rect(x-7, y-6, 14, 12, col)
            pyxel.rect(x-6, y-4, 10, 3,  7)   # shine
            pyxel.line(x-9, y-5, x-9, y+5, 7)  # left highlight

        # Particles
        for p in self.particles:
            if p["life"] > 0:
                s = p.get("size",1)
                if s > 1: pyxel.rect(int(p["x"]), int(p["y"]), s, s, p["col"])
                else:      pyxel.pset(int(p["x"]), int(p["y"]), p["col"])

        self._draw_hud()

    def _draw_hud(self):
        # Top bar
        pyxel.rect(0,0,WIDTH,12,0)
        pyxel.line(0,12,WIDTH,12,1)
        # Score
        sc = f"{self.score:07d}"
        pyxel.text(WIDTH-len(sc)*4-4, 3, sc, 9)
        # Accuracy
        pyxel.text(4, 3, f"{self._acc():.1f}%", 6)
        # Combo
        if self.combo > 1:
            cs  = f"{self.combo}x"
            col = 9 if self.combo>=50 else (12 if self.combo>=20 else 7)
            cx  = WIDTH//2 - len(cs)*2
            pyxel.text(cx-1, 16, cs, 1)
            pyxel.text(cx+1, 16, cs, 1)
            pyxel.text(cx,   16, cs, col)
        # HP bar
        bw, bx, by = WIDTH-40, 20, HEIGHT-12
        pyxel.rect(bx-1, by-1, bw+2, 8, 0)
        pyxel.rect(bx,   by,   bw,   6, 1)
        hw = int(bw * max(0,self.hp) / 100)
        hc = 11 if self.hp>60 else (9 if self.hp>30 else (14 if (self.anim_t//6)%2 else 8))
        pyxel.rect(bx, by, hw, 6, hc)
        pyxel.rectb(bx-1, by-1, bw+2, 8, 5)
        pyxel.text(4, by, "HP", 5)
        # Judge text
        if self.judge_timer > 0:
            jt = self.judge_text
            jc = JUDGE_COLS.get(jt, 7)
            pyxel.text(WIDTH//2-len(jt)*2+1, HIT_Y-34+1, jt, 0)
            pyxel.text(WIDTH//2-len(jt)*2,   HIT_Y-34,   jt, jc)

    def _draw_result(self):
        self._draw_stars(speed=0.5 if self.state=="clear" else 0.1)

        for p in self.burst:
            if p["life"] > 0:
                s = p.get("size",1)
                if s > 1: pyxel.rect(int(p["x"]), int(p["y"]), s, s, p["col"])
                else:      pyxel.pset(int(p["x"]), int(p["y"]), p["col"])

        acc  = self._acc()
        rank = self._rank(acc)
        is_clear = self.state == "clear"

        header = "STAGE  CLEAR!" if is_clear else "GAME  OVER"
        hc     = 9 if is_clear else 8
        pyxel.text(WIDTH//2-len(header)*2+1, 27, header, 0)
        pyxel.text(WIDTH//2-len(header)*2,   26, header, hc)
        pyxel.line(20, 38, WIDTH-20, 38, hc)

        rank_cols = {"S+":9,"S":9,"A":12,"B":3,"C":11,"D":8}
        rc    = rank_cols.get(rank, 7)
        blink = (self.anim_t//15)%2
        pyxel.circ( WIDTH//2, 62, 16, 1)
        pyxel.circb(WIDTH//2, 62, 16, rc)
        rx = WIDTH//2 - len(rank)*4
        pyxel.text(rx+1, 57, rank, 0)
        pyxel.text(rx,   56, rank, rc if blink else 7)

        rows = [
            ("SCORE",     f"{self.score:07d}", 9),
            ("MAX COMBO", f"{self.max_combo}x", 12),
            ("ACCURACY",  f"{acc:.2f}%",        11),
        ]
        for i,(label,val,vc) in enumerate(rows):
            y = 90 + i*24
            pyxel.rect(16, y, WIDTH-32, 18, 1)
            pyxel.rectb(16, y, WIDTH-32, 18, 5)
            pyxel.rect(16, y, 3, 18, vc)
            pyxel.text(28, y+5, label, 6)
            pyxel.text(16+WIDTH-32-len(val)*4-6, y+5, val, vc)

        msgs = {"S+":"FLAWLESS!","S":"INCREDIBLE!","A":"GREAT JOB!",
                "B":"NICE!","C":"KEEP TRYING","D":"PRACTICE MORE"}
        msg = msgs.get(rank,"")
        pyxel.text(WIDTH//2-len(msg)*2, 168, msg, rc)

        if (self.anim_t//25)%2:
            hint = "SPACE -> continue"
            pyxel.text(WIDTH//2-len(hint)*2, HEIGHT-14, hint, 6)


Game()