import kandinsky
import ion

# Dimensions ecran NumWorks
W = 320
H = 222

# Plan complexe : x et y entre -2 et 2
X_MIN, X_MAX = -2.0, 2.0
Y_MIN, Y_MAX = -1.5, 1.5

def complex_to_pixel(z):
    """Convertit un nombre complexe en coordonnees pixel."""
    px = int((z.real - X_MIN) / (X_MAX - X_MIN) * (W - 1))
    py = int((Y_MAX - z.imag) / (Y_MAX - Y_MIN) * (H - 1))
    return px, py

def find_roots(coeffs):
    """
    Trouve les racines d'un polynome par la methode de Durand-Kerner.
    coeffs[0] = coefficient de x^n, coeffs[-1] = terme constant.
    """
    n = len(coeffs) - 1
    if n == 0:
        return []
    if n == 1:
        if coeffs[0] != 0:
            return [complex(-coeffs[1] / coeffs[0])]
        return []

    # Normalisation
    a = [c / coeffs[0] for c in coeffs]

    # Initialisation des racines (cercle unite decale)
    import cmath
    roots = []
    for k in range(n):
        angle = 2 * cmath.pi * k / n + 0.1
        roots.append(0.4 * cmath.exp(complex(0, 1) * angle))

    # Iterations Durand-Kerner
    for _ in range(60):
        new_roots = []
        for i in range(n):
            # Evaluation du polynome en roots[i]
            p = complex(a[0])
            for j in range(1, n + 1):
                p = p * roots[i] + a[j]
            # Produit des differences
            denom = complex(1)
            for j in range(n):
                if j != i:
                    denom *= (roots[i] - roots[j])
            if abs(denom) < 1e-12:
                new_roots.append(roots[i])
            else:
                new_roots.append(roots[i] - p / denom)
        roots = new_roots

    return roots

def draw_axes():
    """Dessine les axes et le fond."""
    # Fond noir
    kandinsky.fill_rect(0, 0, W, H, (0, 0, 0))
    # Axe x (partie imaginaire = 0)
    py = int((Y_MAX) / (Y_MAX - Y_MIN) * (H - 1))
    kandinsky.fill_rect(0, py, W, 1, (60, 60, 60))
    # Axe y (partie reelle = 0)
    px = int((-X_MIN) / (X_MAX - X_MIN) * (W - 1))
    kandinsky.fill_rect(px, 0, 1, H, (60, 60, 60))
    # Cercle unite (approx avec points)
    import cmath
    for i in range(200):
        angle = 2 * cmath.pi * i / 200
        z = cmath.exp(complex(0, 1) * angle)
        px2, py2 = complex_to_pixel(z)
        if 0 <= px2 < W and 0 <= py2 < H:
            kandinsky.set_pixel(px2, py2, (40, 40, 80))

# Palette de couleurs par degre
COLORS = [
    (255, 100, 100),  # degre 1 : rouge
    (255, 180,  50),  # degre 2 : orange
    (255, 255,  50),  # degre 3 : jaune
    ( 80, 255,  80),  # degre 4 : vert
    ( 50, 200, 255),  # degre 5 : cyan
    (150, 100, 255),  # degre 6 : violet
    (255, 100, 200),  # degre 7 : rose
    (200, 255, 150),  # degre 8 : vert clair
]

def run():
    draw_axes()
    MAX_DEGREE = 10

    for deg in range(1, MAX_DEGREE + 1):
        color = COLORS[(deg - 1) % len(COLORS)]
        n_coeffs = deg + 1

        # Enumere tous les polynomes de Littlewood de ce degre
        total = 1 << n_coeffs  # 2^(deg+1)
        for mask in range(total):
            # Construit les coefficients : +1 ou -1
            coeffs = []
            for bit in range(n_coeffs - 1, -1, -1):
                if (mask >> bit) & 1:
                    coeffs.append(1)
                else:
                    coeffs.append(-1)

            # Calcule les racines
            roots = find_roots(coeffs)

            # Trace les racines
            for z in roots:
                px, py = complex_to_pixel(z)
                if 0 <= px < W and 0 <= py < H:
                    kandinsky.set_pixel(px, py, color)

run()