from math import sin, cos, pi
import kandinsky as kd
import ion
import time

# --- Paramètres écran ---
WIDTH = 320
HEIGHT = 240
CENTER_X = WIDTH // 2
CENTER_Y = HEIGHT // 2
SCALE_X = 20
SCALE_Y = 50

# --- Fonctions mathématiques ---
def sinc(t):
    if t == 0:
        return 1
    return sin(pi * t) / (pi * t)

def rect(t, largeur=1):
    return 1 if abs(t) < largeur/2 else 0

def dirac_approx(t):
    return 10 if abs(t) < 0.05 else 0

# --- Affichage axes ---
def draw_axes():
    for x in range(WIDTH):
        kd.set_pixel(x, CENTER_Y, (200, 200, 200))
    for y in range(HEIGHT):
        kd.set_pixel(CENTER_X, y, (200, 200, 200))

# --- Tracé de fonction ---
def draw_function(func, color=(0,0,255), t_min=-5, t_max=5):
    prev = None
    for px in range(WIDTH):
        t = t_min + (t_max - t_min) * px / WIDTH
        y = func(t)
        py = CENTER_Y - int(y * SCALE_Y)
        if 0 <= py < HEIGHT:
            kd.set_pixel(px, py, color)
        if prev:
            y0 = prev[1]
            for py2 in range(min(py, y0), max(py, y0)+1):
                if 0 <= py2 < HEIGHT:
                    kd.set_pixel(px, py2, color)
        prev = (px, py)

# --- Nettoyage écran ---
def clear_screen():
    kd.fill_rect(0, 0, WIDTH, HEIGHT, (255,255,255))

# --- Menu simple ---
def menu():
    clear_screen()
    kd.draw_string("Choisir la fonction :", 20, 20)
    kd.draw_string("1. sin(t)", 20, 50)
    kd.draw_string("2. cos(t)", 20, 70)
    kd.draw_string("3. sinc(t)", 20, 90)
    kd.draw_string("4. rect(t)", 20, 110)
    kd.draw_string("5. dirac approx.", 20, 130)
    kd.draw_string("OK pour quitter", 20, 170)

# --- Boucle principale ---
def main():
    while True:
        menu()
        while True:
            if ion.keypad.is_pressed(ion.keypad.KEY_ONE):
                plot_func(sin, (0,0,255), "sin(t)")
                break
            elif ion.keypad.is_pressed(ion.keypad.KEY_TWO):
                plot_func(cos, (0,128,0), "cos(t)")
                break
            elif ion.keypad.is_pressed(ion.keypad.KEY_THREE):
                plot_func(sinc, (255,0,0), "sinc(t)")
                break
            elif ion.keypad.is_pressed(ion.keypad.KEY_FOUR):
                plot_func(lambda t: rect(t, 2), (255,165,0), "rect(t)")
                break
            elif ion.keypad.is_pressed(ion.keypad.KEY_FIVE):
                plot_func(dirac_approx, (128,0,128), "dirac approx")
                break
            elif ion.keypad.is_pressed(ion.keypad.KEY_OK):
                return
            time.sleep(0.1)

# --- Affichage d'une fonction avec titre ---
def plot_func(func, color, label):
    clear_screen()
    draw_axes()
    draw_function(func, color)
    kd.draw_string(label, 10, 10)
    kd.draw_string("OK pour retour", 10, 200)
    while not ion.keypad.is_pressed(ion.keypad.KEY_OK):
        time.sleep(0.1)

main()