import pygame
import sys
import random
import math

pygame.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Draw & Guess - AI Guesser (Ultimate Edition)")
clock = pygame.time.Clock()
FPS = 60

# ==================== Fonts ====================
def get_font(size):
    """Try to get a system font that supports English well."""
    for name in ["arial", "verdana", "helvetica", "comicsansms"]:
        try:
            font = pygame.font.SysFont(name, size)
            return font
        except:
            continue
    return pygame.font.Font(None, size)

font_small = get_font(20)
font_medium = get_font(24)
font_large = get_font(36)
font_huge = get_font(48)

# ==================== Colors ====================
WHITE   = (255, 255, 255)
BLACK   = (0, 0, 0)
GRAY    = (128, 128, 128)
LIGHT_GRAY = (220, 220, 220)
DARK_GRAY  = (80, 80, 80)
RED     = (255, 0, 0)
GREEN   = (0, 200, 0)
BLUE    = (0, 100, 255)
YELLOW  = (255, 220, 0)
ORANGE  = (255, 150, 0)
BROWN   = (139, 69, 19)
PURPLE  = (160, 50, 160)
PINK    = (255, 120, 180)
CYAN    = (0, 200, 255)
LIME    = (150, 255, 0)
TEAL    = (0, 150, 130)
INDIGO  = (75, 0, 130)
MAROON  = (128, 0, 0)
OLIVE   = (128, 128, 0)
NAVY    = (0, 0, 128)
SILVER  = (192, 192, 192)
GOLD    = (255, 215, 0)
CORAL   = (255, 127, 80)
SKY     = (135, 206, 235)

# Palette
PALETTE_GROUPS = [
    [BLACK, DARK_GRAY, GRAY, LIGHT_GRAY, WHITE],
    [RED, ORANGE, YELLOW, LIME, GREEN],
    [CYAN, BLUE, NAVY, PURPLE, INDIGO],
    [PINK, MAROON, CORAL, GOLD, BROWN],
    [TEAL, OLIVE, SILVER, SKY, PURPLE]
]
color_buttons = []
for row_idx, row in enumerate(PALETTE_GROUPS):
    for col_idx, color in enumerate(row):
        color_buttons.append({
            "color": color,
            "rect": pygame.Rect(630 + col_idx*32, 20 + row_idx*32, 26, 26)
        })

# ==================== Canvas ====================
CANVAS_RECT = pygame.Rect(10, 10, 600, 500)
CANVAS_SCALE = 2
canvas = pygame.Surface((CANVAS_RECT.width * CANVAS_SCALE, CANVAS_RECT.height * CANVAS_SCALE))
canvas.fill(WHITE)
# Paper texture
for _ in range(3000):
    x = random.randint(0, canvas.get_width()-1)
    y = random.randint(0, canvas.get_height()-1)
    shade = random.randint(245, 255)
    canvas.set_at((x, y), (shade, shade, shade))

# Brush parameters
brush_color = BLACK
brush_size = 8
brush_opacity = 255
eraser_mode = False
brush_type = "normal"   # normal, marker, spray

# Undo/Redo
undo_stack = []
redo_stack = []
MAX_UNDO = 20

# Challenge mode
challenge_mode = False
challenge_time = 60.0
challenge_start_ticks = 0
challenge_word = ""

# ==================== Simple sound effects ====================
def make_sound(freq, duration=0.05):
    sample_rate = 22050
    n_samples = int(sample_rate * duration)
    buf = bytes([int(128 + 100 * math.sin(2 * math.pi * freq * i / sample_rate)) for i in range(n_samples)])
    return pygame.mixer.Sound(buffer=buf)

try:
    pygame.mixer.init()
    sound_draw = make_sound(800, 0.03)
    sound_erase = make_sound(400, 0.03)
    sound_click = make_sound(1200, 0.05)
    sound_guess = make_sound(600, 0.1)
    sound_win = make_sound(1000, 0.2)
    sound_fail = make_sound(200, 0.2)
except:
    sound_draw = sound_erase = sound_click = sound_guess = sound_win = sound_fail = None

# ==================== UI Buttons ====================
button_eraser = pygame.Rect(630, 180, 80, 30)
button_clear = pygame.Rect(720, 180, 60, 30)
button_guess = pygame.Rect(630, 230, 150, 40)
button_opacity_down = pygame.Rect(630, 290, 40, 30)
button_opacity_up = pygame.Rect(680, 290, 40, 30)
button_size_down = pygame.Rect(630, 330, 40, 30)
button_size_up = pygame.Rect(680, 330, 40, 30)
button_normal = pygame.Rect(630, 370, 50, 25)
button_marker = pygame.Rect(690, 370, 50, 25)
button_spray = pygame.Rect(750, 370, 50, 25)
button_challenge = pygame.Rect(630, 410, 100, 30)
result_rect = pygame.Rect(630, 460, 150, 40)

# ==================== AI Word List ====================
words = [
    {"name": "Apple", "color": RED, "shape": "round"},
    {"name": "Banana", "color": YELLOW, "shape": "long"},
    {"name": "Sun", "color": YELLOW, "shape": "round"},
    {"name": "Tree", "color": GREEN, "shape": "tall"},
    {"name": "House", "color": BROWN, "shape": "square"},
    {"name": "Sky", "color": BLUE, "shape": "top"},
    {"name": "Strawberry", "color": RED, "shape": "small"},
    {"name": "Orange", "color": ORANGE, "shape": "round"},
    {"name": "Grapes", "color": PURPLE, "shape": "cluster"},
]

def get_color_ratio(surface, target_color, tolerance=50):
    w, h = surface.get_size()
    total = 0
    match = 0
    for x in range(0, w, 8):
        for y in range(0, h, 8):
            total += 1
            pixel = surface.get_at((x, y))[:3]
            dr = abs(pixel[0] - target_color[0])
            dg = abs(pixel[1] - target_color[1])
            db = abs(pixel[2] - target_color[2])
            if dr < tolerance and dg < tolerance and db < tolerance:
                match += 1
    return match / max(1, total)

def get_nonwhite_ratio(surface):
    w, h = surface.get_size()
    total = 0
    nonwhite = 0
    for x in range(0, w, 8):
        for y in range(0, h, 8):
            total += 1
            pixel = surface.get_at((x, y))[:3]
            if pixel[0] < 240 or pixel[1] < 240 or pixel[2] < 240:
                nonwhite += 1
    return nonwhite / max(1, total)

def get_shape_features(surface):
    w, h = surface.get_size()
    min_x, min_y = w, h
    max_x, max_y = 0, 0
    count = 0
    for x in range(0, w, 4):
        for y in range(0, h, 4):
            pixel = surface.get_at((x, y))[:3]
            if pixel[0] < 240 or pixel[1] < 240 or pixel[2] < 240:
                count += 1
                if x < min_x: min_x = x
                if x > max_x: max_x = x
                if y < min_y: min_y = y
                if y > max_y: max_y = y
    if count == 0:
        return None
    bbox_w = max_x - min_x + 1
    bbox_h = max_y - min_y + 1
    aspect = bbox_w / bbox_h if bbox_h > 0 else 0
    fill = count / (bbox_w * bbox_h) if bbox_w * bbox_h > 0 else 0
    sum_x = sum_y = 0
    for x in range(min_x, max_x+1, 4):
        for y in range(min_y, max_y+1, 4):
            pixel = surface.get_at((x, y))[:3]
            if pixel[0] < 240 or pixel[1] < 240 or pixel[2] < 240:
                sum_x += x
                sum_y += y
    cx = sum_x / count
    cy = sum_y / count
    top_ratio = (cy - min_y) / bbox_h if bbox_h > 0 else 0
    return {"aspect": aspect, "fill": fill, "top_ratio": top_ratio}

def ai_guess(surface):
    """Return top three guesses with confidence."""
    nonwhite = get_nonwhite_ratio(surface)
    if nonwhite < 0.01:
        return [("I see nothing!", 0.0)]
    color_scores = {}
    for word in words:
        ratio = get_color_ratio(surface, word["color"], tolerance=50)
        color_scores[word["name"]] = ratio
    shape = get_shape_features(surface)
    scores = {}
    for word in words:
        score = color_scores[word["name"]] * 100
        if shape:
            if word["shape"] == "round":
                if shape["aspect"] > 0.7 and shape["aspect"] < 1.4 and shape["fill"] > 0.4:
                    score += 20
            elif word["shape"] == "long":
                if shape["aspect"] > 1.5 or shape["aspect"] < 0.6:
                    score += 20
            elif word["shape"] == "tall":
                if shape["aspect"] < 0.8 and shape["top_ratio"] < 0.7:
                    score += 15
            elif word["shape"] == "square":
                if shape["aspect"] > 0.7 and shape["aspect"] < 1.4 and shape["fill"] > 0.3:
                    score += 15
            elif word["shape"] == "top":
                if shape["top_ratio"] < 0.4:
                    score += 15
            elif word["shape"] == "cluster":
                if shape["fill"] < 0.4:
                    score += 10
        scores[word["name"]] = score
    sorted_scores = sorted(scores.items(), key=lambda x: x[1], reverse=True)
    results = []
    for name, s in sorted_scores[:3]:
        conf = min(0.95, s / 150)
        results.append((name, conf))
    if not results or results[0][1] < 0.2:
        return [("I don't understand...", 0.1)]
    return results

# ==================== Drawing Helpers ====================
def draw_button(rect, text, color, hover=False):
    shadow_rect = rect.move(2, 2)
    pygame.draw.rect(screen, DARK_GRAY, shadow_rect, border_radius=8)
    pygame.draw.rect(screen, color, rect, border_radius=8)
    pygame.draw.rect(screen, BLACK, rect, 2, border_radius=8)
    if hover:
        light_color = tuple(min(255, c+30) for c in color)
        pygame.draw.rect(screen, light_color, rect, border_radius=8)
    text_surf = font_medium.render(text, True, BLACK)
    screen.blit(text_surf, (rect.x + rect.width//2 - text_surf.get_width()//2,
                            rect.y + rect.height//2 - text_surf.get_height()//2))

def draw_ui(guess_results):
    # Toolbar background
    pygame.draw.rect(screen, (240, 240, 240), pygame.Rect(620, 10, 170, 580))
    pygame.draw.rect(screen, BLACK, pygame.Rect(620, 10, 170, 580), 2)

    # Color buttons
    for btn in color_buttons:
        pygame.draw.rect(screen, btn["color"], btn["rect"], border_radius=5)
        if btn["color"] == brush_color and not eraser_mode:
            pygame.draw.rect(screen, GOLD, btn["rect"], 3, border_radius=5)
        else:
            pygame.draw.rect(screen, BLACK, btn["rect"], 1, border_radius=5)

    # Tool buttons
    draw_button(button_eraser, "Eraser", (200,200,200) if eraser_mode else WHITE)
    draw_button(button_clear, "Clear", WHITE)
    draw_button(button_guess, "AI Guess!", (255,200,100))

    # Brush type buttons
    draw_button(button_normal, "Normal", LIGHT_GRAY if brush_type != "normal" else GOLD)
    draw_button(button_marker, "Marker", LIGHT_GRAY if brush_type != "marker" else GOLD)
    draw_button(button_spray, "Spray", LIGHT_GRAY if brush_type != "spray" else GOLD)

    # Challenge mode button
    draw_button(button_challenge, "Timed" if not challenge_mode else "Stop", (255,180,180))

    # Opacity & size controls
    draw_button(button_opacity_down, "-Op", LIGHT_GRAY)
    draw_button(button_opacity_up, "+Op", LIGHT_GRAY)
    draw_button(button_size_down, "-Sz", LIGHT_GRAY)
    draw_button(button_size_up, "+Sz", LIGHT_GRAY)

    # Parameter display
    opacity_text = font_small.render(f"{int(brush_opacity/255*100)}%", True, BLACK)
    screen.blit(opacity_text, (700, 295))
    size_text = font_small.render(f"{brush_size}px", True, BLACK)
    screen.blit(size_text, (700, 335))

    # AI results
    pygame.draw.rect(screen, WHITE, result_rect)
    pygame.draw.rect(screen, BLACK, result_rect, 2, border_radius=5)
    if guess_results:
        y_offset = result_rect.y + 8
        for idx, (name, conf) in enumerate(guess_results[:3]):
            if idx == 0:
                text_color = (0, 150, 0) if conf > 0.3 else (200, 0, 0)
                text = font_medium.render(f"{name}", True, text_color)
            else:
                text = font_small.render(f"{name} {conf*100:.0f}%", True, GRAY)
            screen.blit(text, (result_rect.x + 10, y_offset))
            y_offset += 25
    else:
        text = font_small.render("AI guesses:", True, BLACK)
        screen.blit(text, (result_rect.x + 10, result_rect.y + 10))

    # Challenge mode status
    if challenge_mode:
        elapsed = (pygame.time.get_ticks() - challenge_start_ticks) / 1000.0
        remaining = max(0, challenge_time - elapsed)
        time_text = font_large.render(f"{int(remaining)}s", True, RED if remaining < 10 else BLACK)
        screen.blit(time_text, (WIDTH - 100, 10))
        word_text = font_medium.render(f"Draw: {challenge_word}", True, BLACK)
        screen.blit(word_text, (CANVAS_RECT.x + 20, CANVAS_RECT.y + 20))

# ==================== Main ====================
def main():
    global brush_color, brush_size, brush_opacity, eraser_mode, brush_type
    global challenge_mode, challenge_time, challenge_start_ticks, challenge_word
    global undo_stack, redo_stack
    global canvas

    brush_color = BLACK
    brush_size = 8
    brush_opacity = 255
    eraser_mode = False
    brush_type = "normal"

    challenge_mode = False
    challenge_time = 60.0
    challenge_start_ticks = 0
    challenge_word = ""

    undo_stack = []
    redo_stack = []
    guess_results = []

    drawing = False
    last_pos = None
    display_canvas = pygame.Surface((CANVAS_RECT.width, CANVAS_RECT.height))

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 1:
                    if CANVAS_RECT.collidepoint(event.pos):
                        undo_stack.append(canvas.copy())
                        if len(undo_stack) > MAX_UNDO:
                            undo_stack.pop(0)
                        redo_stack.clear()

                    if button_guess.collidepoint(event.pos):
                        small = pygame.transform.smoothscale(canvas, (CANVAS_RECT.width, CANVAS_RECT.height))
                        guess_results = ai_guess(small)
                        if sound_guess: sound_guess.play()
                    elif button_eraser.collidepoint(event.pos):
                        eraser_mode = not eraser_mode
                        if sound_click: sound_click.play()
                    elif button_clear.collidepoint(event.pos):
                        canvas.fill(WHITE)
                        for _ in range(3000):
                            x = random.randint(0, canvas.get_width()-1)
                            y = random.randint(0, canvas.get_height()-1)
                            shade = random.randint(245, 255)
                            canvas.set_at((x, y), (shade, shade, shade))
                        guess_results = []
                        undo_stack.clear()
                        redo_stack.clear()
                        if sound_click: sound_click.play()
                    elif button_opacity_down.collidepoint(event.pos):
                        brush_opacity = max(20, brush_opacity - 20)
                    elif button_opacity_up.collidepoint(event.pos):
                        brush_opacity = min(255, brush_opacity + 20)
                    elif button_size_down.collidepoint(event.pos):
                        brush_size = max(2, brush_size - 1)
                    elif button_size_up.collidepoint(event.pos):
                        brush_size = min(30, brush_size + 1)
                    elif button_normal.collidepoint(event.pos):
                        brush_type = "normal"
                        eraser_mode = False
                    elif button_marker.collidepoint(event.pos):
                        brush_type = "marker"
                        eraser_mode = False
                    elif button_spray.collidepoint(event.pos):
                        brush_type = "spray"
                        eraser_mode = False
                    elif button_challenge.collidepoint(event.pos):
                        if challenge_mode:
                            challenge_mode = False
                            challenge_word = ""
                        else:
                            challenge_mode = True
                            challenge_word = random.choice(words)["name"]
                            challenge_start_ticks = pygame.time.get_ticks()
                            canvas.fill(WHITE)
                            for _ in range(3000):
                                x = random.randint(0, canvas.get_width()-1)
                                y = random.randint(0, canvas.get_height()-1)
                                shade = random.randint(245, 255)
                                canvas.set_at((x, y), (shade, shade, shade))
                            guess_results = []
                            undo_stack.clear()
                            redo_stack.clear()
                        if sound_click: sound_click.play()
                    else:
                        for btn in color_buttons:
                            if btn["rect"].collidepoint(event.pos):
                                brush_color = btn["color"]
                                eraser_mode = False
                                if sound_click: sound_click.play()
                        if CANVAS_RECT.collidepoint(event.pos):
                            drawing = True
                            x = int((event.pos[0] - CANVAS_RECT.x) * CANVAS_SCALE)
                            y = int((event.pos[1] - CANVAS_RECT.y) * CANVAS_SCALE)
                            last_pos = (x, y)
                            if eraser_mode:
                                pygame.draw.circle(canvas, WHITE, last_pos, brush_size)
                                if sound_erase: sound_erase.play()
                            else:
                                if brush_type == "spray":
                                    for _ in range(5):
                                        dx = random.randint(-brush_size, brush_size)
                                        dy = random.randint(-brush_size, brush_size)
                                        pygame.draw.circle(canvas, (*brush_color, brush_opacity), (last_pos[0]+dx, last_pos[1]+dy), 2)
                                else:
                                    for i in range(3):
                                        radius = brush_size - i
                                        temp = pygame.Surface((radius*2, radius*2), pygame.SRCALPHA)
                                        pygame.draw.circle(temp, (*brush_color, brush_opacity), (radius, radius), radius)
                                        canvas.blit(temp, (last_pos[0]-radius, last_pos[1]-radius))
                                if sound_draw: sound_draw.play()
            elif event.type == pygame.MOUSEBUTTONUP:
                if event.button == 1:
                    drawing = False
                    last_pos = None
            elif event.type == pygame.MOUSEMOTION:
                if drawing:
                    x = int((event.pos[0] - CANVAS_RECT.x) * CANVAS_SCALE)
                    y = int((event.pos[1] - CANVAS_RECT.y) * CANVAS_SCALE)
                    pos = (x, y)
                    if last_pos:
                        if eraser_mode:
                            pygame.draw.line(canvas, WHITE, last_pos, pos, brush_size*2)
                            pygame.draw.circle(canvas, WHITE, pos, brush_size)
                        else:
                            dist = math.hypot(pos[0]-last_pos[0], pos[1]-last_pos[1])
                            steps = max(1, int(dist / (brush_size/2)))
                            for i in range(steps):
                                t = i / steps
                                ix = int(last_pos[0] + (pos[0]-last_pos[0])*t)
                                iy = int(last_pos[1] + (pos[1]-last_pos[1])*t)
                                if brush_type == "spray":
                                    for _ in range(4):
                                        dx = random.randint(-brush_size, brush_size)
                                        dy = random.randint(-brush_size, brush_size)
                                        pygame.draw.circle(canvas, (*brush_color, brush_opacity), (ix+dx, iy+dy), 2)
                                else:
                                    for j in range(3):
                                        r = brush_size - j
                                        temp = pygame.Surface((r*2, r*2), pygame.SRCALPHA)
                                        pygame.draw.circle(temp, (*brush_color, brush_opacity), (r, r), r)
                                        canvas.blit(temp, (ix-r, iy-r))
                    last_pos = pos
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_z and pygame.key.get_mods() & pygame.KMOD_CTRL:
                    if undo_stack:
                        redo_stack.append(canvas.copy())
                        canvas = undo_stack.pop()
                elif event.key == pygame.K_y and pygame.key.get_mods() & pygame.KMOD_CTRL:
                    if redo_stack:
                        undo_stack.append(canvas.copy())
                        canvas = redo_stack.pop()

        # Challenge timer check
        if challenge_mode:
            elapsed = (pygame.time.get_ticks() - challenge_start_ticks) / 1000.0
            if elapsed >= challenge_time:
                challenge_mode = False
                small = pygame.transform.smoothscale(canvas, (CANVAS_RECT.width, CANVAS_RECT.height))
                guess_results = ai_guess(small)
                if guess_results and guess_results[0][0] == challenge_word:
                    if sound_win: sound_win.play()
                else:
                    if sound_fail: sound_fail.play()
                challenge_word = ""

        # Draw everything
        pygame.transform.smoothscale(canvas, (CANVAS_RECT.width, CANVAS_RECT.height), display_canvas)
        screen.fill((180, 180, 180))
        screen.blit(display_canvas, CANVAS_RECT.topleft)
        pygame.draw.rect(screen, BLACK, CANVAS_RECT, 2)

        draw_ui(guess_results)

        pygame.display.flip()
        clock.tick(FPS)

if __name__ == "__main__":
    main()