import pygame
import sys
import copy
import heapq

# ========== 常量 ==========
WALL = 0
FLOOR = 1
TARGET = 2
BOX = 3
PLAYER = 4
BOX_ON_TARGET = 5
PLAYER_ON_TARGET = 6

COLORS = {
    WALL: (40, 40, 40),
    FLOOR: (200, 180, 150),
    TARGET: (255, 200, 100),
    BOX: (139, 69, 19),
    PLAYER: (50, 150, 255),
    BOX_ON_TARGET: (255, 100, 50),
    PLAYER_ON_TARGET: (50, 255, 50)
}

CELL_SIZE = 60
MARGIN = 20

# ========== 游戏逻辑 ==========
class Sokoban:
    def __init__(self, level):
        self.initial_map = [row[:] for row in level]
        self.map = [row[:] for row in level]
        self.rows = len(level)
        self.cols = len(level[0])
        self.player_pos = None
        self.targets = []
        
        for r in range(self.rows):
            for c in range(self.cols):
                if self.map[r][c] in [PLAYER, PLAYER_ON_TARGET]:
                    self.player_pos = (r, c)
                if self.map[r][c] in [TARGET, PLAYER_ON_TARGET, BOX_ON_TARGET]:
                    self.targets.append((r, c))

    def reset(self):
        self.map = [row[:] for row in self.initial_map]
        for r in range(self.rows):
            for c in range(self.cols):
                if self.map[r][c] in [PLAYER, PLAYER_ON_TARGET]:
                    self.player_pos = (r, c)

    def get_cell(self, r, c):
        if 0 <= r < self.rows and 0 <= c < self.cols:
            return self.map[r][c]
        return WALL

    def is_box(self, r, c):
        return self.get_cell(r, c) in [BOX, BOX_ON_TARGET]

    def is_target(self, r, c):
        return (r, c) in self.targets

    def is_wall(self, r, c):
        return self.get_cell(r, c) == WALL

    def is_empty(self, r, c):
        return self.get_cell(r, c) in [FLOOR, TARGET]

    def move_player(self, dr, dc):
        pr, pc = self.player_pos
        nr, nc = pr + dr, pc + dc
        br, bc = nr + dr, nc + dc

        if self.is_wall(nr, nc):
            return False

        if self.is_box(nr, nc):
            if self.is_wall(br, bc) or self.is_box(br, bc):
                return False
            
            self.map[br][bc] = BOX_ON_TARGET if self.is_target(br, bc) else BOX
            self.map[nr][nc] = TARGET if self.is_target(nr, nc) else FLOOR
            self.map[pr][pc] = TARGET if self.is_target(pr, pc) else FLOOR
            self.map[nr][nc] = PLAYER_ON_TARGET if self.is_target(nr, nc) else PLAYER
            self.player_pos = (nr, nc)
            return True

        if self.is_empty(nr, nc):
            self.map[pr][pc] = TARGET if self.is_target(pr, pc) else FLOOR
            self.map[nr][nc] = PLAYER_ON_TARGET if self.is_target(nr, nc) else PLAYER
            self.player_pos = (nr, nc)
            return True

        return False

    def is_win(self):
        for r, c in self.targets:
            if self.map[r][c] != BOX_ON_TARGET:
                return False
        return True

    def get_state_key(self):
        boxes = []
        for r in range(self.rows):
            for c in range(self.cols):
                if self.is_box(r, c):
                    boxes.append((r, c))
        boxes.sort()
        return (self.player_pos, tuple(boxes))

    def clone(self):
        return copy.deepcopy(self)


# ========== A* 求解器 ==========
def heuristic(game):
    boxes = []
    for r in range(game.rows):
        for c in range(game.cols):
            if game.is_box(r, c):
                boxes.append((r, c))
    
    if not boxes:
        return 0
    
    total = 0
    for br, bc in boxes:
        min_dist = float('inf')
        for tr, tc in game.targets:
            dist = abs(br - tr) + abs(bc - tc)
            if dist < min_dist:
                min_dist = dist
        total += min_dist
    return total

def solve_sokoban(initial_game, max_steps=5000):
    dirs = [(-1, 0, 'up'), (1, 0, 'down'), (0, -1, 'left'), (0, 1, 'right')]
    
    start_state = initial_game.get_state_key()
    start_cost = heuristic(initial_game)
    
    counter = 0
    heap = []
    heapq.heappush(heap, (start_cost, counter, initial_game, start_state, []))
    visited = {start_state: 0}
    counter += 1
    
    steps = 0
    
    while heap and steps < max_steps:
        cost, _, game, state_key, path = heapq.heappop(heap)
        steps += 1
        
        if game.is_win():
            return path
        
        for dr, dc, name in dirs:
            new_game = game.clone()
            if new_game.move_player(dr, dc):
                new_key = new_game.get_state_key()
                new_g = len(path) + 1
                
                if new_key not in visited or new_g < visited[new_key]:
                    visited[new_key] = new_g
                    new_cost = new_g + heuristic(new_game)
                    heapq.heappush(heap, (new_cost, counter, new_game, new_key, path + [name]))
                    counter += 1
    
    return None


# ========== Pygame 界面（最简版，无中文）==========
class SokobanGame:
    def __init__(self, level_data):
        pygame.init()
        
        self.game = Sokoban(level_data)
        self.solution = None
        self.solution_index = 0
        self.auto_play = False
        self.clock = pygame.time.Clock()
        
        # 窗口大小
        width = self.game.cols * CELL_SIZE + MARGIN * 2
        height = self.game.rows * CELL_SIZE + MARGIN * 2 + 70
        self.screen = pygame.display.set_mode((width, height))
        pygame.display.set_caption("Sokoban")
        
        # 字体（只用默认，保证不闪退）
        self.font = pygame.font.Font(None, 24)
        self.big_font = pygame.font.Font(None, 32)
        
        # 按钮
        btn_w = 80
        btn_h = 30
        btn_y = height - 45
        self.buttons = {
            'solve': pygame.Rect(MARGIN, btn_y, btn_w, btn_h),
            'reset': pygame.Rect(MARGIN + btn_w + 10, btn_y, btn_w, btn_h),
            'auto': pygame.Rect(MARGIN + (btn_w + 10) * 2, btn_y, btn_w, btn_h),
            'step': pygame.Rect(MARGIN + (btn_w + 10) * 3, btn_y, btn_w, btn_h),
        }
        
        self.running = True

    def draw_map(self):
        for r in range(self.game.rows):
            for c in range(self.game.cols):
                x = c * CELL_SIZE + MARGIN
                y = r * CELL_SIZE + MARGIN
                cell = self.game.map[r][c]
                
                # 地板
                pygame.draw.rect(self.screen, COLORS[FLOOR], (x, y, CELL_SIZE, CELL_SIZE))
                pygame.draw.rect(self.screen, (100, 100, 100), (x, y, CELL_SIZE, CELL_SIZE), 1)
                
                # 目标点
                if self.game.is_target(r, c):
                    pygame.draw.circle(self.screen, COLORS[TARGET], 
                                     (x + CELL_SIZE//2, y + CELL_SIZE//2), 10)
                
                # 墙
                if cell == WALL:
                    pygame.draw.rect(self.screen, COLORS[WALL], (x, y, CELL_SIZE, CELL_SIZE))
                
                # 箱子
                elif cell in [BOX, BOX_ON_TARGET]:
                    color = COLORS[BOX_ON_TARGET] if cell == BOX_ON_TARGET else COLORS[BOX]
                    pygame.draw.rect(self.screen, color, (x+5, y+5, CELL_SIZE-10, CELL_SIZE-10))
                    pygame.draw.rect(self.screen, (0,0,0), (x+5, y+5, CELL_SIZE-10, CELL_SIZE-10), 2)
                
                # 玩家
                elif cell in [PLAYER, PLAYER_ON_TARGET]:
                    color = COLORS[PLAYER_ON_TARGET] if cell == PLAYER_ON_TARGET else COLORS[PLAYER]
                    pygame.draw.circle(self.screen, color, 
                                     (x + CELL_SIZE//2, y + CELL_SIZE//2), CELL_SIZE//2 - 5)
                    pygame.draw.circle(self.screen, (0,0,0), 
                                     (x + CELL_SIZE//2, y + CELL_SIZE//2), CELL_SIZE//2 - 5, 2)

    def draw_ui(self):
        # 按钮
        btn_text = {
            'solve': 'SOLVE',
            'reset': 'RESET',
            'auto': 'STOP' if self.auto_play else 'AUTO',
            'step': 'STEP'
        }
        btn_color = {
            'solve': (100, 200, 100),
            'reset': (200, 200, 100),
            'auto': (255, 100, 100) if self.auto_play else (100, 200, 200),
            'step': (200, 150, 100)
        }
        
        for name, rect in self.buttons.items():
            pygame.draw.rect(self.screen, btn_color[name], rect)
            pygame.draw.rect(self.screen, (0, 0, 0), rect, 2)
            text = self.font.render(btn_text[name], True, (0, 0, 0))
            text_rect = text.get_rect(center=rect.center)
            self.screen.blit(text, text_rect)
        
        # 状态信息
        done = sum(1 for r in range(self.game.rows) for c in range(self.game.cols) 
                   if self.game.map[r][c] == BOX_ON_TARGET)
        total = len(self.game.targets)
        
        if self.game.is_win():
            status = "YOU WIN!"
        else:
            status = f"Box: {done}/{total}"
        
        text = self.font.render(status, True, (0, 0, 0))
        self.screen.blit(text, (MARGIN, self.game.rows * CELL_SIZE + MARGIN + 5))
        
        # 步骤信息
        if self.solution:
            step_text = f"Step: {self.solution_index}/{len(self.solution)}"
        else:
            step_text = "Press SOLVE"
        
        text = self.font.render(step_text, True, (0, 0, 0))
        self.screen.blit(text, (MARGIN + 200, self.game.rows * CELL_SIZE + MARGIN + 5))

    def handle_click(self, pos):
        for name, rect in self.buttons.items():
            if rect.collidepoint(pos):
                if name == 'solve':
                    self.solve_level()
                elif name == 'reset':
                    self.reset_level()
                elif name == 'auto':
                    self.auto_play = not self.auto_play
                    if self.auto_play and not self.solution:
                        self.solve_level()
                elif name == 'step':
                    self.next_step()
                return True
        return False

    def solve_level(self):
        print("Solving...")
        self.game.reset()
        self.solution_index = 0
        self.solution = solve_sokoban(self.game, max_steps=5000)
        if self.solution:
            print(f"Found solution! {len(self.solution)} steps")
            self.game.reset()
            self.solution_index = 0
        else:
            print("No solution found")

    def reset_level(self):
        self.game.reset()
        self.solution_index = 0
        self.auto_play = False
        self.solution = None

    def next_step(self):
        if self.solution and self.solution_index < len(self.solution):
            dir_map = {'up': (-1,0), 'down': (1,0), 'left': (0,-1), 'right': (0,1)}
            dr, dc = dir_map[self.solution[self.solution_index]]
            self.game.move_player(dr, dc)
            self.solution_index += 1
            return True
        return False

    def run(self):
        while self.running:
            self.screen.fill((240, 240, 240))
            
            # 自动播放
            if self.auto_play and self.solution and not self.game.is_win():
                if not self.next_step():
                    self.auto_play = False
                    print("Auto play done!")
            
            self.draw_map()
            self.draw_ui()
            
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    self.running = False
                elif event.type == pygame.MOUSEBUTTONDOWN:
                    self.handle_click(event.pos)
                elif event.type == pygame.KEYDOWN:
                    # 键盘控制
                    key_map = {
                        pygame.K_UP: (-1, 0),
                        pygame.K_DOWN: (1, 0),
                        pygame.K_LEFT: (0, -1),
                        pygame.K_RIGHT: (0, 1)
                    }
                    if event.key in key_map:
                        dr, dc = key_map[event.key]
                        self.game.move_player(dr, dc)
                        self.auto_play = False
            
            pygame.display.flip()
            self.clock.tick(60)
        
        pygame.quit()
        sys.exit()


# ========== 运行 ==========
if __name__ == "__main__":
    # 关卡地图
    level_data = [
        [0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 1, 1, 1, 1, 1, 1, 1, 0],
        [0, 1, 3, 1, 3, 1, 2, 1, 0],
        [0, 1, 1, 1, 1, 1, 2, 1, 0],
        [0, 1, 1, 4, 1, 1, 1, 1, 0],
        [0, 1, 1, 1, 1, 1, 1, 1, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0]
    ]
    
    print("Sokoban Started")
    print("Controls: Arrow keys | Buttons: SOLVE, RESET, AUTO, STEP")
    
    game = SokobanGame(level_data)
    game.run()