# tk_Colorful_Tetris.py ZZZ

import tkinter as tk
import random
import json
import os
import time

BOARD_WIDTH = 10
BOARD_HEIGHT = 20
BLOCK_SIZE = 30
SHAPES = [
    [[1, 1, 1, 1]],
    [[1, 1], [1, 1]],
    [[1, 1, 1], [1, 0, 0]],
    [[1, 1, 1], [0, 0, 1]],
    [[1, 1, 0], [0, 1, 1]],
    [[0, 1, 1], [1, 1, 0]]
]
COLORS = ['cyan', 'yellow', 'purple', 'orange', 'blue', 'green', 'red']

board = [[0 for _ in range(BOARD_WIDTH)] for _ in range(BOARD_HEIGHT)]
current_piece = None
next_piece = None
highscore = 0
score = 0
lines_cleared = 0
ai_countdown = 0
mode = "player"
ai_weights = {
    'lines': 1.0,
    'height': -0.5,
    'holes': -0.7,
    'bumpiness': -0.3
}
ai_countdown_timer = None
game_timer = None

AI_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tetris.json')

def load_data():
    global highscore, ai_weights
    if os.path.exists(AI_PATH):
        try:
            with open(AI_PATH, 'r') as f:
                data = json.load(f)
                highscore = data.get("highscore", 0)
                ai_weights = data.get("ai_weights", ai_weights)
        except Exception:
            pass

def save_data():
    with open(AI_PATH, 'w') as f:
        json.dump({"highscore": highscore, "ai_weights": ai_weights}, f)

def save_highscore():
    save_data()

def create_piece():
    i = random.randint(0, len(SHAPES)-1)
    shape = [row[:] for row in SHAPES[i]]
    color = COLORS[i]
    return {'shape': shape, 'color': color, 'x': BOARD_WIDTH // 2 - len(shape[0]) // 2, 'y': 0}

def rotate_piece(piece):
    shape = piece['shape']
    rotated = [list(row) for row in zip(*shape[::-1])]
    return rotated

def collision(piece, dx=0, dy=0, test_shape=None):
    shape = test_shape if test_shape else piece['shape']
    for y, row in enumerate(shape):
        for x, cell in enumerate(row):
            if cell:
                nx = piece['x'] + x + dx
                ny = piece['y'] + y + dy
                if nx < 0 or nx >= BOARD_WIDTH or ny >= BOARD_HEIGHT:
                    return True
                if ny >= 0 and board[ny][nx]:
                    return True
    return False

def merge(piece):
    for y, row in enumerate(piece['shape']):
        for x, cell in enumerate(row):
            if cell and piece['y'] + y >= 0:
                board[piece['y'] + y][piece['x'] + x] = piece['color']

def clear_lines():
    global board, score, lines_cleared
    new_board = [row for row in board if any(cell == 0 for cell in row)]
    cleared = BOARD_HEIGHT - len(new_board)
    for _ in range(cleared):
        new_board.insert(0, [0]*BOARD_WIDTH)
    if cleared:
        score += cleared * 100
        lines_cleared += cleared
    board = new_board

def get_board_stats(test_board):
    heights = [0]*BOARD_WIDTH
    holes = 0
    bumpiness = 0
    for x in range(BOARD_WIDTH):
        col_height = 0
        block_found = False
        for y in range(BOARD_HEIGHT):
            if test_board[y][x]:
                if not block_found:
                    col_height = BOARD_HEIGHT - y
                    block_found = True
            elif block_found:
                holes += 1
        heights[x] = col_height
    for i in range(BOARD_WIDTH-1):
        bumpiness += abs(heights[i] - heights[i+1])
    total_height = sum(heights)
    return total_height, holes, bumpiness

def simulate_drop(piece, x, rotation):
    test_piece = {'shape': [row[:] for row in piece['shape']],
                  'color': piece['color'],
                  'x': x, 'y': 0}
    for _ in range(rotation):
        test_piece['shape'] = rotate_piece(test_piece['shape'])
    
    if collision(test_piece, dx=0, dy=0):
        return -1, float('inf'), float('inf'), float('inf'), test_piece
    
    while not collision(test_piece, dy=1):
        test_piece['y'] += 1
    
    test_board = [row[:] for row in board]
    for y, row in enumerate(test_piece['shape']):
        for xx, cell in enumerate(row):
            if cell:
                by = test_piece['y'] + y
                bx = test_piece['x'] + xx
                if 0 <= by < BOARD_HEIGHT and 0 <= bx < BOARD_WIDTH:
                    test_board[by][bx] = 1
    
    lines = 0
    temp_board = [r[:] for r in test_board]
    for y in range(BOARD_HEIGHT-1, -1, -1):
        if all(temp_board[y]):
            lines += 1
            del temp_board[y]
            temp_board.insert(0, [0]*BOARD_WIDTH)
    
    height, holes, bumpiness = get_board_stats(temp_board)
    return lines, height, holes, bumpiness, test_piece

def ai_best_move():
    best_score = -float('inf')
    best_move = (current_piece['x'], 0)
    piece = {'shape': [row[:] for row in current_piece['shape']],
             'color': current_piece['color'],
             'x': 0, 'y': 0}
    
    for rotation in range(4):
        shape = piece['shape']
        width = len(shape[0])
        for x in range(-min([i for i, v in enumerate(shape[0]) if v]), BOARD_WIDTH - width + 1):
            lines, height, holes, bumpiness, _ = simulate_drop(piece, x, rotation)
            
            if lines == -1:
                continue
                
            score_eval = (lines * ai_weights['lines'] +
                          height * ai_weights['height'] +
                          holes * ai_weights['holes'] +
                          bumpiness * ai_weights['bumpiness'])
            
            if score_eval > best_score:
                best_score = score_eval
                best_move = (x, rotation)
        
        piece['shape'] = rotate_piece(piece)
    
    return best_move

def draw():
    game_canvas.delete('all')
    for y in range(BOARD_HEIGHT):
        for x in range(BOARD_WIDTH):
            color = board[y][x]
            if color:
                game_canvas.create_rectangle(
                    x*BLOCK_SIZE, y*BLOCK_SIZE,
                    (x+1)*BLOCK_SIZE, (y+1)*BLOCK_SIZE,
                    fill=color, outline='black'
                )
    if current_piece:
        for y, row in enumerate(current_piece['shape']):
            for x, cell in enumerate(row):
                if cell:
                    game_canvas.create_rectangle(
                        (current_piece['x']+x)*BLOCK_SIZE,
                        (current_piece['y']+y)*BLOCK_SIZE,
                        (current_piece['x']+x+1)*BLOCK_SIZE,
                        (current_piece['y']+y+1)*BLOCK_SIZE,
                        fill=current_piece['color'], outline='black'
                    )
    game_canvas.create_text(5, 5, anchor='nw', text=f"Score: {score}\nLines: {lines_cleared}", fill='white', font=('Arial', 14))
    next_canvas.delete('all')
    next_canvas.create_text(5, 5, anchor='nw', text=f"Highscore:\n {highscore}", fill='yellow', font=('Arial', 14))
    if next_piece:
        shape = next_piece['shape']
        color = next_piece['color']
        offset_x = (6 - len(shape[0])) // 2
        offset_y = (6 - len(shape)) // 2
        for y, row in enumerate(shape):
            for x, cell in enumerate(row):
                if cell:
                    next_canvas.create_rectangle(
                        (x+offset_x)*BLOCK_SIZE,
                        (y+offset_y)*BLOCK_SIZE + 30,
                        (x+offset_x+1)*BLOCK_SIZE,
                        (y+offset_y+1)*BLOCK_SIZE + 30,
                        fill=color, outline='black'
                    )
    if mode == "ai" and ai_countdown > 0:
        next_canvas.create_text(90, 70, text=f"AI Restart: {ai_countdown}", fill='red', font=('Arial', 18), anchor='n')

def cancel_all_timers():
    global game_timer, ai_countdown_timer
    if game_timer:
        root.after_cancel(game_timer)
        game_timer = None
    if ai_countdown_timer:
        root.after_cancel(ai_countdown_timer)
        ai_countdown_timer = None

def on_key(event):
    global current_piece
    if mode != "player" or not current_piece:
        return
    if event.keysym == 'Left':
        if not collision(current_piece, dx=-1):
            current_piece['x'] -= 1
    elif event.keysym == 'Right':
        if not collision(current_piece, dx=1):
            current_piece['x'] += 1
    elif event.keysym == 'Down':
        if not collision(current_piece, dy=1):
            current_piece['y'] += 1
    elif event.keysym == 'Up':
        rotated = rotate_piece(current_piece)
        if not collision(current_piece, test_shape=rotated):
            current_piece['shape'] = rotated
    elif event.keysym == 'space':
        while not collision(current_piece, dy=1):
            current_piece['y'] += 1
    draw()

def end_game():
    global highscore
    if score > highscore:
        highscore = score
        save_highscore()
    draw()
    game_canvas.create_text(BOARD_WIDTH*BLOCK_SIZE//2, BOARD_HEIGHT*BLOCK_SIZE//2, text="GAME OVER", fill='red', font=('Arial', 24, 'bold'))

def player_game_step():
    global current_piece, next_piece, score, lines_cleared, highscore, game_timer
    if not current_piece:
        return
    if not collision(current_piece, dy=1):
        current_piece['y'] += 1
    else:
        merge(current_piece)
        clear_lines()
        if score > highscore:
            highscore = score
            save_highscore()
        current_piece = next_piece
        next_piece = create_piece()
        if collision(current_piece):
            end_game()
            return
    draw()
    if current_piece:
        game_timer = root.after(500, player_game_step)

def start_player_game():
    global board, score, lines_cleared, current_piece, next_piece, mode
    cancel_all_timers()
    mode = "player"
    board = [[0 for _ in range(BOARD_WIDTH)] for _ in range(BOARD_HEIGHT)]
    score = 0
    lines_cleared = 0
    current_piece = create_piece()
    next_piece = create_piece()
    draw()
    root.after(1, player_game_step)

ai_move_plan = None
ai_move_step = 0
ai_drop_step = 0

def ai_prepare_move():
    global ai_move_plan, ai_move_step, ai_drop_step
    x_target, rot_target = ai_best_move()
    ai_move_plan = {'x': x_target, 'rot': rot_target}
    ai_move_step = 0
    ai_drop_step = 0

def ai_animate_move():
    global ai_move_plan, ai_move_step, ai_drop_step, current_piece, next_piece, score, lines_cleared, game_timer
    
    if not current_piece or ai_move_plan is None:
        return
    
    if ai_move_step < ai_move_plan['rot']:
        rotated = rotate_piece(current_piece)
        if not collision(current_piece, test_shape=rotated):
            current_piece['shape'] = rotated
            ai_move_step += 1
        else:
            ai_move_step = ai_move_plan['rot']
        draw()
        game_timer = root.after(100, ai_animate_move)
        return
    
    if current_piece['x'] < ai_move_plan['x']:
        if not collision(current_piece, dx=1):
            current_piece['x'] += 1
        elif current_piece['x'] < ai_move_plan['x'] - 1:
            current_piece['x'] += 1
        draw()
        game_timer = root.after(50, ai_animate_move)
        return
    elif current_piece['x'] > ai_move_plan['x']:
        if not collision(current_piece, dx=-1):
            current_piece['x'] -= 1
        elif current_piece['x'] > ai_move_plan['x'] + 1:
            current_piece['x'] -= 1
        draw()
        game_timer = root.after(50, ai_animate_move)
        return
    
    if not collision(current_piece, dy=1):
        current_piece['y'] += 1
        draw()
        game_timer = root.after(30, ai_animate_move)
        return
    
    merge(current_piece)
    clear_lines()
    if lines_cleared and random.random() < 0.1:
        for k in ai_weights:
            ai_weights[k] += random.uniform(-0.05, 0.05)
        save_data()
    
    current_piece = next_piece
    next_piece = create_piece()
    if collision(current_piece):
        end_game()
        start_ai_countdown()
        return
    
    ai_prepare_move()
    draw()
    game_timer = root.after(100, ai_animate_move)

def ai_game_step():
    ai_prepare_move()
    ai_animate_move()

def start_ai_game():
    global board, score, lines_cleared, current_piece, next_piece, mode
    cancel_all_timers()
    mode = "ai"
    board = [[0 for _ in range(BOARD_WIDTH)] for _ in range(BOARD_HEIGHT)]
    score = 0
    lines_cleared = 0
    current_piece = create_piece()
    next_piece = create_piece()
    draw()
    root.after(1, ai_game_step)

def start_ai_countdown():
    global ai_countdown, ai_countdown_timer
    ai_countdown = 6
    if ai_countdown_timer:
        root.after_cancel(ai_countdown_timer)
    ai_countdown_tick()

def ai_countdown_tick():
    global ai_countdown, ai_countdown_timer
    if ai_countdown > 0:
        draw()
        ai_countdown -= 1
        ai_countdown_timer = root.after(1000, ai_countdown_tick)
    else:
        ai_countdown_timer = None
        start_ai_game()

if __name__ == "__main__":
    load_data()
    root = tk.Tk()
    root.title("Colorful Tetris")
    root.geometry("+0+0")
    frame = tk.Frame(root)
    frame.pack()
    game_canvas = tk.Canvas(frame, width=BOARD_WIDTH*BLOCK_SIZE, height=BOARD_HEIGHT*BLOCK_SIZE, bg='black')
    game_canvas.pack(side='left')
    right_frame = tk.Frame(frame)
    right_frame.pack(side='top', padx=20)
    next_canvas = tk.Canvas(right_frame, width=6*BLOCK_SIZE, height=6*BLOCK_SIZE, bg='gray20')
    next_canvas.pack(side='top', pady=10)
    btn_frame = tk.Frame(right_frame)
    btn_frame.pack(side='top', pady=10)
    btn_player = tk.Button(btn_frame, text="Player Restart", command=start_player_game, width=15)
    btn_player.pack(side='top', pady=5)
    btn_ai = tk.Button(btn_frame, text="AI", command=start_ai_game, width=15)
    btn_ai.pack(side='top', pady=5)
    root.bind("<Key>", on_key)
    root.focus_set()
    start_player_game()
    root.mainloop()