piano.py

#!/usr/bin/env python3
"""
DreamScaler Piano
Unified interactive tool: LED visualisation, scale/chord explorer, demos.
Replaces the old piano.py + scale_selector_gui.py + dreamscaler.py.

Usage:
    python piano.py <serial_port> [menu_choice]
    python piano.py COM5
    python piano.py COM5 12       # jump straight to GUI scale selector
"""

import sys
import time
import json
import math
import random
import atexit
import signal
import tkinter as tk
from tkinter import ttk
from pathlib import Path

from controller_api import LEDController, LEDControllerError
from config import COM_PORT, LED_INTENSITY, SCALE_DEGREE_COLORS
from arturia_keylab49_map import (
    PIANO_KEY_MAP,
    LED_COUNT,
    WHITE_KEY_COLOR,
    BLACK_KEY_COLOR,
    OCTAVE_COLORS,
    visualize_piano_layout,
    get_all_white_keys,
    get_all_black_keys,
    print_piano_map,
)


# ============================================================
# CONSTANTS
# ============================================================

NOTE_NAMES = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']

# GUI degree colours (hex)
DEGREE_COLORS_HEX = {
    1: '#FF4444', 2: '#FF8800', 3: '#44FF44',
    4: '#FFFF00', 5: '#4444FF', 6: '#00FFFF', 7: '#FF44FF',
}

# LED degree colours — defined in config.py

# True  = colour by key type (white/black), root in red
# False = colour by scale degree
USE_KEY_COLOR_MODE = False


# ============================================================
# CLEANUP
# ============================================================

_global_controller = None


def _cleanup():
    global _global_controller
    if _global_controller is not None:
        try:
            print('\n Cleanup: turning off LEDs and closing port...')
            _global_controller.clear_all()
            _global_controller.disconnect()
            _global_controller = None
            print('Done')
        except Exception:
            pass


def _signal_handler(signum, frame):
    print(f'\nSignal {signum} received - exiting...')
    _cleanup()
    sys.exit(0)


atexit.register(_cleanup)
signal.signal(signal.SIGINT, _signal_handler)
signal.signal(signal.SIGTERM, _signal_handler)


# ============================================================
# SCALE DATA
# ============================================================

def _flatten_scale(s: dict, lang: str = 'en') -> dict:
    """Flatten a unified-format scale entry to a flat dict for the given language."""
    flat = {k: v for k, v in s.items() if k not in ('en', 'cs')}
    flat.update(s.get(lang, s.get('en', {})))
    return flat


def load_scales_from_file(filepath='scales.json', lang: str = 'en'):
    """Load scales from unified JSON file, flattened to the specified language."""
    path = Path(filepath)
    if not path.exists():
        path = Path(__file__).parent / filepath
    try:
        with open(path, encoding='utf-8') as f:
            return [_flatten_scale(s, lang) for s in json.load(f)]
    except Exception as e:
        print(f'Error loading scales: {e}')
        return []


def get_scale_notes(root_note, intervals):
    """Return list of (note, degree) tuples for the scale."""
    notes = [(root_note, 1)]
    current = root_note
    degree = 2
    for interval in intervals[:-1]:
        current = (current + interval) % 12
        notes.append((current, degree))
        degree += 1
    return notes


# ============================================================
# SCALE SELECTOR GUI
# ============================================================

class ScaleSelectorGUI:
    """
    Tkinter window for interactive scale selection.
    Calls on_scale_selected(root_note, scale_name, intervals) on apply.
    root_note=None signals 'clear LEDs'.
    """

    def __init__(self, scales_data, on_scale_selected=None):
        self.scales_data = scales_data
        self.on_scale_selected = on_scale_selected
        self.current_scale = None
        self.current_root = 0

        self.categories = {}
        for scale in scales_data:
            cat = scale.get('category', 'Other')
            self.categories.setdefault(cat, []).append(scale)

        self.root = None
        self.is_running = False

    def run(self):
        self.root = tk.Tk()
        self.root.title('DreamScaler - Scale Selector')
        self.root.geometry('700x600')
        self.root.configure(bg='#2b2b2b')
        self.is_running = True
        self._create_widgets()

        self.root.update_idletasks()
        w, h = self.root.winfo_width(), self.root.winfo_height()
        x = (self.root.winfo_screenwidth() // 2) - (w // 2)
        y = (self.root.winfo_screenheight() // 2) - (h // 2)
        self.root.geometry(f'{w}x{h}+{x}+{y}')
        self.root.protocol('WM_DELETE_WINDOW', self._on_close)
        self.root.mainloop()

    def _on_close(self):
        self.is_running = False
        self.root.destroy()

    def _create_widgets(self):
        style = ttk.Style()
        style.theme_use('clam')
        style.configure('TFrame',            background='#2b2b2b')
        style.configure('TLabel',            background='#2b2b2b', foreground='white',   font=('Segoe UI', 10))
        style.configure('Title.TLabel',      font=('Segoe UI', 14, 'bold'))
        style.configure('Info.TLabel',       font=('Segoe UI', 9),  foreground='#aaaaaa')
        style.configure('TButton',           font=('Segoe UI', 10))
        style.configure('TCombobox',         font=('Segoe UI', 10))
        style.configure('TLabelframe',       background='#2b2b2b', foreground='white')
        style.configure('TLabelframe.Label', background='#2b2b2b', foreground='white', font=('Segoe UI', 11, 'bold'))

        main = ttk.Frame(self.root, padding='20')
        main.pack(fill=tk.BOTH, expand=True)

        ttk.Label(main, text='DreamScaler - Scale Selector', style='Title.TLabel').pack(pady=(0, 20))

        # Selection
        sel = ttk.LabelFrame(main, text='Scale Selection', padding='15')
        sel.pack(fill=tk.X, pady=(0, 15))

        cat_row = ttk.Frame(sel)
        cat_row.pack(fill=tk.X, pady=(0, 10))
        ttk.Label(cat_row, text='Category:').pack(side=tk.LEFT, padx=(0, 10))
        self.category_var = tk.StringVar()
        self.category_combo = ttk.Combobox(cat_row, textvariable=self.category_var,
                                           values=sorted(self.categories.keys()),
                                           state='readonly', width=30)
        self.category_combo.pack(side=tk.LEFT, fill=tk.X, expand=True)
        self.category_combo.bind('<<ComboboxSelected>>', self._on_category_change)

        sc_row = ttk.Frame(sel)
        sc_row.pack(fill=tk.X, pady=(0, 10))
        ttk.Label(sc_row, text='Scale:').pack(side=tk.LEFT, padx=(0, 10))
        self.scale_var = tk.StringVar()
        self.scale_combo = ttk.Combobox(sc_row, textvariable=self.scale_var,
                                        state='readonly', width=30)
        self.scale_combo.pack(side=tk.LEFT, fill=tk.X, expand=True)
        self.scale_combo.bind('<<ComboboxSelected>>', self._on_scale_change)

        root_row = ttk.Frame(sel)
        root_row.pack(fill=tk.X)
        ttk.Label(root_row, text='Root Note:').pack(side=tk.LEFT, padx=(0, 10))
        self.root_var = tk.StringVar(value='C')
        self.root_combo = ttk.Combobox(root_row, textvariable=self.root_var,
                                       values=NOTE_NAMES, state='readonly', width=10)
        self.root_combo.pack(side=tk.LEFT)
        self.root_combo.bind('<<ComboboxSelected>>', self._on_root_change)

        # Info
        info = ttk.LabelFrame(main, text='Scale Information', padding='15')
        info.pack(fill=tk.BOTH, expand=True, pady=(0, 15))

        self.info_feelings  = ttk.Label(info, text='', style='Info.TLabel', wraplength=600)
        self.info_feelings.pack(anchor=tk.W, pady=2)
        self.info_genre     = ttk.Label(info, text='', style='Info.TLabel', wraplength=600)
        self.info_genre.pack(anchor=tk.W, pady=2)
        self.info_usage     = ttk.Label(info, text='', style='Info.TLabel', wraplength=600)
        self.info_usage.pack(anchor=tk.W, pady=2)
        self.info_intervals = ttk.Label(info, text='', style='Info.TLabel')
        self.info_intervals.pack(anchor=tk.W, pady=2)

        notes_frame = ttk.Frame(info)
        notes_frame.pack(fill=tk.X, pady=(15, 0))
        ttk.Label(notes_frame, text='Notes in scale:').pack(anchor=tk.W)
        self.notes_display = tk.Frame(notes_frame, bg='#2b2b2b')
        self.notes_display.pack(fill=tk.X, pady=(5, 0))

        # Buttons
        btn_row = ttk.Frame(main)
        btn_row.pack(fill=tk.X)
        tk.Button(btn_row, text='Show on LED', command=self._apply_scale,
                  bg='#4CAF50', fg='white', font=('Segoe UI', 11, 'bold'),
                  padx=20, pady=10, cursor='hand2').pack(side=tk.LEFT, padx=(0, 10))
        tk.Button(btn_row, text='Clear LED', command=self._clear_leds,
                  bg='#f44336', fg='white', font=('Segoe UI', 11),
                  padx=20, pady=10, cursor='hand2').pack(side=tk.LEFT, padx=(0, 10))
        tk.Button(btn_row, text='Close', command=self._on_close,
                  bg='#555555', fg='white', font=('Segoe UI', 11),
                  padx=20, pady=10, cursor='hand2').pack(side=tk.RIGHT)

        self.status_var = tk.StringVar(value='Select a category and scale')
        ttk.Label(main, textvariable=self.status_var,
                  style='Info.TLabel', anchor=tk.W).pack(fill=tk.X, pady=(10, 0))

        if self.categories:
            first = sorted(self.categories.keys())[0]
            self.category_combo.set(first)
            self._on_category_change(None)

    def _on_category_change(self, _event):
        cat = self.category_var.get()
        if cat in self.categories:
            names = [s['name'] for s in self.categories[cat]]
            self.scale_combo['values'] = names
            if names:
                self.scale_combo.set(names[0])
                self._on_scale_change(None)

    def _on_scale_change(self, _event):
        name = self.scale_var.get()
        cat  = self.category_var.get()
        self.current_scale = None
        for s in self.categories.get(cat, []):
            if s['name'] == name:
                self.current_scale = s
                break
        self._update_info()
        self._update_notes()

    def _on_root_change(self, _event):
        self.current_root = NOTE_NAMES.index(self.root_var.get())
        self._update_notes()

    def _update_info(self):
        if self.current_scale:
            self.info_feelings .config(text=f"Feelings: {self.current_scale.get('feelings', 'N/A')}")
            self.info_genre    .config(text=f"Genre: {self.current_scale.get('genre', 'N/A')}")
            self.info_usage    .config(text=f"Usage: {self.current_scale.get('usage', 'N/A')}")
            self.info_intervals.config(text=f"Intervals: {self.current_scale.get('intervals', [])}")
            self.status_var.set(f"Selected: {self.current_scale['name']}")
        else:
            for lbl in (self.info_feelings, self.info_genre, self.info_usage, self.info_intervals):
                lbl.config(text='')

    def _update_notes(self):
        for w in self.notes_display.winfo_children():
            w.destroy()
        if not self.current_scale:
            return
        notes = get_scale_notes(self.current_root, self.current_scale.get('intervals', []))
        for note, deg in notes:
            color = DEGREE_COLORS_HEX.get(deg, '#888888')
            text_color = 'black' if deg in (3, 4, 6) else 'white'
            lbl = tk.Label(self.notes_display, text=f' {NOTE_NAMES[note]} ',
                           bg=color, fg=text_color,
                           font=('Segoe UI', 12, 'bold'), padx=8, pady=4)
            lbl.pack(side=tk.LEFT, padx=2)
            deg_labels = {1:'Root',2:'Second',3:'Third',4:'Fourth',5:'Fifth',6:'Sixth',7:'Seventh'}
            self._tooltip(lbl, f'Degree {deg} ({deg_labels.get(deg, "")})')

    def _tooltip(self, widget, text):
        def show(e):
            tip = tk.Toplevel(widget)
            tip.wm_overrideredirect(True)
            tip.wm_geometry(f'+{e.x_root+10}+{e.y_root+10}')
            tk.Label(tip, text=text, bg='#ffffe0', relief='solid', borderwidth=1,
                     font=('Segoe UI', 9), padx=5, pady=2).pack()
            widget._tip = tip
        def hide(_e):
            if hasattr(widget, '_tip'):
                widget._tip.destroy()
        widget.bind('<Enter>', show)
        widget.bind('<Leave>', hide)

    def _apply_scale(self):
        if not self.current_scale:
            self.status_var.set('Please select a scale first!')
            return
        self.status_var.set(f"Showing: {NOTE_NAMES[self.current_root]} {self.current_scale['name']}")
        if self.on_scale_selected:
            self.on_scale_selected(self.current_root, self.current_scale['name'],
                                   self.current_scale['intervals'])

    def _clear_leds(self):
        self.status_var.set('LED cleared')
        if self.on_scale_selected:
            self.on_scale_selected(None, None, None)


# ============================================================
# PIANO VISUALISATION
# ============================================================

def show_piano_keys(controller):
    print(f'\n=== Piano keys ({len(PIANO_KEY_MAP)} total) ===')
    visualize_piano_layout(controller)
    print(f'  {len(get_all_white_keys())} white + {len(get_all_black_keys())} black keys')


def show_white_keys_only(controller):
    print('\n=== White keys ===')
    controller.clear_all()
    for pos in get_all_white_keys():
        controller.set_pixel(pos, *WHITE_KEY_COLOR)
    print(f'  {len(get_all_white_keys())} keys shown')


def show_black_keys_only(controller):
    print('\n=== Black keys ===')
    controller.clear_all()
    for pos in get_all_black_keys():
        controller.set_pixel(pos, *BLACK_KEY_COLOR)
    print(f'  {len(get_all_black_keys())} keys shown')


def show_octaves(controller):
    print('\n=== Octave colours ===')
    controller.clear_all()
    for led_pos, _note, _is_white, octave in PIANO_KEY_MAP:
        if led_pos < LED_COUNT and octave in OCTAVE_COLORS:
            controller.set_pixel(led_pos, *OCTAVE_COLORS[octave])
    print('  Red=Oct2  Green=Oct3  Blue=Oct4  Yellow=Oct5  Magenta=Oct6')


def test_key_animation(controller):
    print('\n=== Key animation ===')
    controller.clear_all()
    for led_pos, _note, is_white, _octave in PIANO_KEY_MAP:
        if led_pos < LED_COUNT:
            controller.set_pixel(led_pos, *(WHITE_KEY_COLOR if is_white else BLACK_KEY_COLOR))
            time.sleep(0.05)
    time.sleep(2)
    controller.clear_all()
    print('  Done')


# ============================================================
# SCALE DISPLAY ON LED
# ============================================================

def show_scale(controller, root_note, scale_name, intervals):
    print(f'\n=== {NOTE_NAMES[root_note]} {scale_name} ===')
    controller.clear_all()
    scale_notes  = get_scale_notes(root_note, intervals)
    note_to_deg  = {n: d for n, d in scale_notes}
    ROOT_COLOR   = (LED_INTENSITY, 0, 0, 0)

    for led_pos, note, is_white, _octave in PIANO_KEY_MAP:
        if led_pos < LED_COUNT and note in note_to_deg:
            deg = note_to_deg[note]
            if USE_KEY_COLOR_MODE:
                color = ROOT_COLOR if deg == 1 else (WHITE_KEY_COLOR if is_white else BLACK_KEY_COLOR)
            else:
                color = SCALE_DEGREE_COLORS.get(deg, (LED_INTENSITY,) * 3 + (0,))
            controller.set_pixel(led_pos, *color)

    deg_names = {1:'Tonic',2:'Second',3:'Third',4:'Fourth',5:'Fifth',6:'Sixth',7:'Seventh'}
    for note, deg in scale_notes:
        print(f'  {deg}. {deg_names.get(deg,""):8} {NOTE_NAMES[note]}')


def show_all_scales_menu(controller):
    scales = load_scales_from_file()
    if not scales:
        return
    categories = {}
    for s in scales:
        categories.setdefault(s['category'], []).append(s)

    while True:
        print('\n' + '='*60 + '\nSCALES - choose category\n' + '='*60)
        cat_list = sorted(categories)
        for i, cat in enumerate(cat_list, 1):
            print(f'{i:2}. {cat:22} ({len(categories[cat])} scales)')
        print('\n 0. Back')
        choice = input('\nCategory: ').strip()
        if choice == '0':
            controller.clear_all()
            break
        try:
            idx = int(choice) - 1
            if 0 <= idx < len(cat_list):
                _show_category_menu(controller, cat_list[idx], categories[cat_list[idx]])
        except ValueError:
            pass


def _show_category_menu(controller, cat_name, scales_in_cat):
    while True:
        print(f'\n{"="*60}\n{cat_name}\n{"="*60}')
        for i, s in enumerate(scales_in_cat, 1):
            print(f'{i:2}. {s["name"]}  -  {s.get("feelings", "")}')
        print('\n 0. Back')
        choice = input('\nScale: ').strip()
        if choice == '0':
            break
        try:
            idx = int(choice) - 1
            if 0 <= idx < len(scales_in_cat):
                _show_root_menu(controller, scales_in_cat[idx])
        except ValueError:
            pass


def _show_root_menu(controller, scale):
    while True:
        print(f'\n{"="*60}\n{scale["name"]}\nFeelings: {scale.get("feelings","")}\nGenre: {scale.get("genre","")}\n{"="*60}')
        for i, n in enumerate(NOTE_NAMES, 1):
            print(f'{i:2}. {n} {scale["name"]}')
        print('\n99. Back')
        choice = input('\nRoot: ').strip()
        if choice == '99':
            break
        try:
            root = int(choice) - 1
            if 0 <= root <= 11:
                show_scale(controller, root, scale['name'], scale['intervals'])
                input('\nEnter to continue...')
        except ValueError:
            pass


def show_scale_selector_gui(controller):
    print('\nOpening GUI scale selector...')
    scales = load_scales_from_file()
    if not scales:
        print('Could not load scales')
        return

    def on_selected(root, name, intervals):
        if root is None:
            controller.clear_all()
        else:
            show_scale(controller, root, name, intervals)

    ScaleSelectorGUI(scales, on_selected).run()
    print('GUI closed')


# ============================================================
# CHORD PROGRESSIONS
# ============================================================

CHORD_PROGRESSIONS = {
    'I-IV-V-I':     {'name': 'Basic cadence',    'genre': 'Pop, Rock, Folk',      'chords': [1,4,5,1],                   'chord_types': ['maj','maj','maj','maj']},
    'I-V-vi-IV':    {'name': 'Pop progression',  'genre': 'Pop, Rock',            'chords': [1,5,6,4],                   'chord_types': ['maj','maj','min','maj']},
    'ii-V-I':       {'name': 'Jazz cadence',      'genre': 'Jazz, Bossa Nova',     'chords': [2,5,1],                     'chord_types': ['min7','dom7','maj7']},
    'I-vi-IV-V':    {'name': '50s progression',  'genre': 'Oldies, Doo-wop',      'chords': [1,6,4,5],                   'chord_types': ['maj','min','maj','maj']},
    'vi-IV-I-V':    {'name': 'Emotional',         'genre': 'Pop, Ballads',         'chords': [6,4,1,5],                   'chord_types': ['min','maj','maj','maj']},
    '12-bar-blues': {'name': '12-bar blues',      'genre': "Blues, Rock'n'roll",   'chords': [1,1,1,1,4,4,1,1,5,4,1,5],  'chord_types': ['dom7']*12},
}
CHORD_INTERVALS = {
    'maj': [0,4,7], 'min': [0,3,7], 'dim': [0,3,6], 'aug': [0,4,8],
    'maj7': [0,4,7,11], 'min7': [0,3,7,10], 'dom7': [0,4,7,10], 'dim7': [0,3,6,9],
}
SCALE_DEGREES_SEMITONES = {1:0, 2:2, 3:4, 4:5, 5:7, 6:9, 7:11, -7:10}
CHORD_COLORS = {
    1: (LED_INTENSITY*2, 0, 0, 0),                       # I   - Red
    2: (LED_INTENSITY,   LED_INTENSITY, 0, 0),            # ii  - Yellow
    3: (0, LED_INTENSITY*2, 0, 0),                        # iii - Green
    4: (0, LED_INTENSITY,   LED_INTENSITY, 0),            # IV  - Cyan
    5: (0, 0, LED_INTENSITY*2, 0),                        # V   - Blue
    6: (LED_INTENSITY, 0, LED_INTENSITY, 0),              # vi  - Magenta
    7: (LED_INTENSITY, LED_INTENSITY//2, 0, 0),           # vii - Orange
   -7: (LED_INTENSITY, 0, 0, LED_INTENSITY),              # bVII
}


def _chord_notes(root, degree, chord_type):
    chord_root = (root + SCALE_DEGREES_SEMITONES[degree]) % 12
    return [(chord_root + i) % 12 for i in CHORD_INTERVALS.get(chord_type, CHORD_INTERVALS['maj'])]


def _show_chord(controller, root, degree, chord_type, color):
    notes = set(_chord_notes(root, degree, chord_type))
    for led_pos, note, _w, _o in PIANO_KEY_MAP:
        if led_pos < LED_COUNT and note in notes:
            controller.set_pixel(led_pos, *color)


def show_chord_progressions_menu(controller):
    prog_list = list(CHORD_PROGRESSIONS.keys())
    while True:
        print('\n' + '='*60 + '\nCHORD PROGRESSIONS\n' + '='*60)
        for i, key in enumerate(prog_list, 1):
            p = CHORD_PROGRESSIONS[key]
            print(f'{i:2}. {key:22} {p["name"]}  -  {p["genre"]}')
        print('\n 0. Back')
        choice = input('\nProgression: ').strip()
        if choice == '0':
            controller.clear_all()
            break
        try:
            idx = int(choice) - 1
            if 0 <= idx < len(prog_list):
                _progression_root_menu(controller, prog_list[idx])
        except ValueError:
            pass


def _progression_root_menu(controller, prog_key):
    p = CHORD_PROGRESSIONS[prog_key]
    while True:
        print(f'\n{"="*60}\n{prog_key} - {p["name"]}\n{p["genre"]}\n{"="*60}')
        for i, n in enumerate(NOTE_NAMES, 1):
            print(f'{i:2}. {n}')
        print('\n99. Back')
        choice = input('\nKey: ').strip()
        if choice == '99':
            break
        try:
            root = int(choice) - 1
            if 0 <= root <= 11:
                _play_progression(controller, root, prog_key)
        except ValueError:
            pass


def _play_progression(controller, root_note, prog_key):
    p = CHORD_PROGRESSIONS[prog_key]
    chords, types = p['chords'], p['chord_types']
    print(f'\n{NOTE_NAMES[root_note]} - {p["name"]}  [Enter=next  a=auto  r=restart  q=quit]')
    idx = 0
    while True:
        controller.clear_all()
        deg, ct = chords[idx], types[idx]
        _show_chord(controller, root_note, deg, ct, CHORD_COLORS.get(deg, (LED_INTENSITY,)*3+(0,)))
        cn = _chord_notes(root_note, deg, ct)
        print(f'  {idx+1}/{len(chords)}  {NOTE_NAMES[cn[0]]} ({ct})')
        c = input('  [Enter/a/r/q]: ').strip().lower()
        if c == 'q':
            break
        elif c == 'r':
            idx = 0
        elif c == 'a':
            print('  Auto (1.5s/chord)...')
            for i in range(len(chords)):
                controller.clear_all()
                _show_chord(controller, root_note, chords[i], types[i],
                            CHORD_COLORS.get(chords[i], (LED_INTENSITY,)*3+(0,)))
                time.sleep(1.5)
            idx = 0
        else:
            idx = (idx + 1) % len(chords)
    controller.clear_all()


# ============================================================
# DEMOS / EFFECTS
# ============================================================

def demo_info(controller):
    info = controller.get_info()
    if info:
        print(f'  Protocol: {info["protocol_version"]}  LEDs: {info["led_count"]}  '
              f'Pin: {info["led_pin"]}  Brightness: {info["brightness"]}')
    else:
        print('  No response')


def demo_basic(controller):
    for r, g, b, w, label in [(255,0,0,0,'Red'),(0,255,0,0,'Green'),(0,0,255,0,'Blue'),(0,0,0,255,'White')]:
        print(f'  LED 0: {label}')
        controller.set_pixel(0, r, g, b, w)
        time.sleep(0.8)
    controller.clear_all()


def demo_gradient(controller):
    n = LED_COUNT - 1
    pairs = [
        ('Red to Blue',    255,0,0,0,  0,0,255,0),
        ('Green to Yellow',  0,255,0,0, 255,255,0,0),
    ]
    for name, *c in pairs:
        print(f'  {name}')
        controller.fill_gradient(0, n, *c)
        time.sleep(2)
    controller.clear_all()


def demo_rainbow(controller, duration=8):
    print(f'  Rainbow {duration}s...')
    start = time.time(); frames = 0
    n = LED_COUNT
    while time.time() - start < duration:
        offset = frames * 3
        pixels = []
        for i in range(n):
            h = (i * 360 / n + offset) % 360
            r, g, b = controller.hsv_to_rgb(h, 1.0, 0.3)
            pixels.append((r, g, b, 0))
        controller.stream_update(pixels)
        frames += 1
    controller.clear_all()
    print(f'  {frames} frames ({frames/duration:.1f} FPS)')


def demo_knight_rider(controller, duration=8):
    print(f'  Knight Rider {duration}s...')
    n = LED_COUNT
    start = time.time(); pos = 0; direction = 1; tail = 10
    while time.time() - start < duration:
        controller.clear_all()
        controller.set_pixel(pos, 255, 0, 0, 0)
        for i in range(1, tail):
            tp = pos - i * direction
            if 0 <= tp < n:
                bri = 255 - i * 255 // tail
                controller.set_pixel(tp, bri, 0, 0, 0)
        pos += direction
        if pos >= n - 1 or pos <= 0:
            direction *= -1
        time.sleep(0.01)
    controller.clear_all()


def demo_fire(controller, duration=8):
    print(f'  Fire {duration}s...')
    n = LED_COUNT
    start = time.time()
    while time.time() - start < duration:
        pixels = [(random.randint(100, 255), random.randint(30, 100), 0, 0) for _ in range(n)]
        controller.stream_update(pixels)
        time.sleep(0.05)
    controller.clear_all()


def demo_wave(controller, duration=8, wave_width=20):
    print(f'  Wave {duration}s  ({wave_width} active LEDs)...')
    n = LED_COUNT
    pos = 0
    direction = 1
    start = time.time()
    while time.time() - start < duration:
        pixels = [(0, 0, 0, 0)] * n
        for i in range(n):
            dist = abs(i - pos)
            if dist < wave_width:
                bri = int((1 - dist / wave_width) * 200)
                pixels[i] = (0, 0, bri, 0)
        controller.stream_update(pixels)
        pos += direction
        if pos >= n - 1 or pos <= 0:
            direction *= -1
        time.sleep(0.03)
    controller.clear_all()


def show_demos_menu(controller):
    items = [
        ('Info',            lambda: demo_info(controller)),
        ('Basic RGBW',      lambda: demo_basic(controller)),
        ('Gradient',        lambda: demo_gradient(controller)),
        ('Rainbow',         lambda: demo_rainbow(controller)),
        ('Knight Rider',    lambda: demo_knight_rider(controller)),
        ('Fire',            lambda: demo_fire(controller)),
        ('Wave',            lambda: demo_wave(controller)),
        ('Clear all',       lambda: controller.clear_all()),
    ]
    while True:
        print('\n' + '='*60 + '\nDEMOS\n' + '='*60)
        for i, (name, _) in enumerate(items, 1):
            print(f'{i:2}. {name}')
        print('\n 0. Back')
        choice = input('\nDemo: ').strip()
        if choice == '0':
            break
        try:
            idx = int(choice) - 1
            if 0 <= idx < len(items):
                items[idx][1]()
        except ValueError:
            pass


# ============================================================
# MAIN MENU
# ============================================================

def _print_menu():
    print('\n' + '='*60)
    print('DreamScaler Piano  -  Arturia KeyLab 49 MKII')
    print('='*60)
    print(f'\n  LED intensity: {LED_INTENSITY}')
    print('\nVISUALISATION:')
    print('  1  All keys')
    print('  2  White keys only')
    print('  3  Black keys only')
    print('  4  Octave colours')
    print('  5  Key animation')
    print('  6  Print key map')
    print('\nSCALES & CHORDS:')
    print(' 10  All scales (by category)')
    print(' 11  Chord progressions')
    print(' 12  GUI scale selector')
    print('\nDEMOS & EFFECTS:')
    print(' 20  Demos menu')
    print('\n  7  Clear all LEDs')
    print('  0  Quit')
    print('='*60)


def _execute(controller, choice):
    if choice == '1':
        show_piano_keys(controller)
        input('\nEnter to continue...')
        controller.clear_all()
    elif choice == '2':
        show_white_keys_only(controller)
        input('\nEnter to continue...')
        controller.clear_all()
    elif choice == '3':
        show_black_keys_only(controller)
        input('\nEnter to continue...')
        controller.clear_all()
    elif choice == '4':
        show_octaves(controller)
        input('\nEnter to continue...')
        controller.clear_all()
    elif choice == '5':
        test_key_animation(controller)
    elif choice == '6':
        print_piano_map()
        input('\nEnter to continue...')
    elif choice == '7':
        controller.clear_all()
        print('  Cleared')
    elif choice == '10':
        show_all_scales_menu(controller)
    elif choice == '11':
        show_chord_progressions_menu(controller)
    elif choice == '12':
        show_scale_selector_gui(controller)
    elif choice == '20':
        show_demos_menu(controller)
    elif choice == '0':
        controller.clear_all()
        return 'exit'
    else:
        print('  Invalid choice')
    return None


def main():
    global _global_controller

    if len(sys.argv) < 2:
        port = COM_PORT
        print(f'No port given — using default from config.py: {port}')
    else:
        port = sys.argv[1]
    auto = sys.argv[2] if len(sys.argv) > 2 else None

    print(f'Connecting to {port}...')
    controller = LEDController(port)
    _global_controller = controller

    if not controller.connect():
        print('Connection failed')
        _global_controller = None
        return

    print(f'Connected  |  {len(PIANO_KEY_MAP)} keys mapped')

    if auto:
        print(f'Auto-launching option {auto}')
        _execute(controller, auto)
        return

    while True:
        _print_menu()
        try:
            c = input('\nChoice: ').strip()
            if _execute(controller, c) == 'exit':
                break
        except KeyboardInterrupt:
            print('\nInterrupted')
            break
        except LEDControllerError as e:
            print(f'Controller error: {e}')
        except Exception as e:
            print(f'Error: {e}')


if __name__ == '__main__':
    main()