import pygame
import random
import sys

pygame.init()
pygame.mixer.init()

GRID_SIZE = 80
ROW = 5
COL = 9
WIDTH = COL * GRID_SIZE
HEIGHT = ROW * GRID_SIZE + 120
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("植物大战僵尸 杂交终极版")
clock = pygame.time.Clock()
FPS = 60

GREEN_GRASS = (20, 120, 20)
GRID_LINE = (10, 80, 10)
SUN_COLOR = (255, 230, 0)
PEA_COLOR = (0, 180, 0)
WALL_COLOR = (130, 80, 20)
ICE_COLOR = (120, 200, 255)
ZOMBIE_COLOR = (80, 80, 80)
BUCKET_COLOR = (50, 50, 50)
CHERRY_COLOR = (255, 30, 30)
POTATO_COLOR = (100, 60, 20)
WHITE = (255,255,255)
RED = (255,0,0)
BLACK = (0,0,0)
ORANGE = (255, 150, 0)
CYAN = (0, 200, 200)

font = pygame.font.Font(None, 18)
font_big = pygame.font.Font(None, 24)

def make_sound(freq, duration):
    try:
        import numpy as np
        sample_rate = 44100
        t = np.linspace(0, duration, int(sample_rate * duration))
        wave = np.sin(2 * np.pi * freq * t) * 32767
        wave = wave.astype(np.int16)
        return pygame.mixer.Sound(buffer=wave)
    except:
        return None

sound_plant = make_sound(440, 0.1)
sound_pea = make_sound(600, 0.05)
sound_explode = make_sound(150, 0.4)
sound_zombie = make_sound(220, 0.2)
sound_hybrid = make_sound(880, 0.15)

sun = 150
plants = []
zombies = []
bullets = []
selected_plant = None
game_over = False
spawn_time = 0
selected_card = -1  # 选中的植物卡片
grid_used = [[False for _ in range(ROW)] for _ in range(COL)]

class Plant:
    def __init__(self, x, y, hp, cost, atk, sun_prod, color):
        self.x = x
        self.y = y
        self.hp = hp
        self.max_hp = hp
        self.cost = cost
        self.atk = atk
        self.sun_prod = sun_prod
        self.color = color
        self.cd = 1000
        self.last_time = 0

    def draw_hp(self):
        w = GRID_SIZE - 10
        h = 6
        ratio = self.hp / self.max_hp
        pygame.draw.rect(screen, RED, (self.x+5, self.y-8, w, h))
        pygame.draw.rect(screen, (0,255,0), (self.x+5, self.y-8, w*ratio, h))

    def draw_cartoon(self):
        pygame.draw.ellipse(screen, (30,150,30), (self.x+10, self.y+45, 60, 30))

    def draw(self):
        self.draw_cartoon()
        pygame.draw.rect(screen, self.color, (self.x+15, self.y+10, 50, 40), border_radius=12)
        self.draw_hp()

    def update(self, now):
        pass

class SunFlower(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, 80, 50, 0, 25, SUN_COLOR)
        self.prod_cd = 3000

    def draw(self):
        self.draw_cartoon()
        pygame.draw.circle(screen, SUN_COLOR, (self.x+40, self.y+30), 25)
        pygame.draw.circle(screen, (255,180,0), (self.x+40, self.y+30), 15)
        self.draw_hp()

    def update(self, now):
        if now - self.last_time > self.prod_cd:
            global sun
            sun += self.sun_prod
            self.last_time = now

class Peashooter(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, 100, 100, 20, 0, PEA_COLOR)

    def draw(self):
        self.draw_cartoon()
        pygame.draw.rect(screen, PEA_COLOR, (self.x+20, self.y+15, 40, 35), border_radius=10)
        pygame.draw.circle(screen, (0,120,0), (self.x+60, self.y+32), 8)
        self.draw_hp()

    def update(self, now):
        if now - self.last_time > self.cd:
            bullets.append(Bullet(self.x+GRID_SIZE, self.y+GRID_SIZE//2, self.atk, "normal"))
            self.last_time = now
            if sound_pea: sound_pea.play()

class IceShooter(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, 90, 175, 15, 0, ICE_COLOR)

    def draw(self):
        self.draw_cartoon()
        pygame.draw.rect(screen, ICE_COLOR, (self.x+20, self.y+15, 40, 35), border_radius=10)
        pygame.draw.circle(screen, (180,230,255), (self.x+60, self.y+32), 8)
        self.draw_hp()

    def update(self, now):
        if now - self.last_time > self.cd:
            bullets.append(Bullet(self.x+GRID_SIZE, self.y+GRID_SIZE//2, self.atk, "ice"))
            self.last_time = now
            if sound_pea: sound_pea.play()

class WallNut(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, 450, 50, 0, 0, WALL_COLOR)

    def draw(self):
        self.draw_cartoon()
        pygame.draw.rect(screen, WALL_COLOR, (self.x+10, self.y+10, 60, 50), border_radius=15)
        pygame.draw.line(screen, (80,40,10), (self.x+20,self.y+25), (self.x+50,self.y+25),3)
        pygame.draw.line(screen, (80,40,10), (self.x+20,self.y+40), (self.x+50,self.y+40),3)
        self.draw_hp()

class DoublePea(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, 110, 200, 18, 0, (0,150,80))

    def draw(self):
        self.draw_cartoon()
        pygame.draw.rect(screen, (0,150,80), (self.x+20, self.y+10, 40, 45), border_radius=10)
        pygame.draw.circle(screen, (0,120,0), (self.x+60, self.y+22), 7)
        pygame.draw.circle(screen, (0,120,0), (self.x+60, self.y+42), 7)
        self.draw_hp()

    def update(self, now):
        if now - self.last_time > self.cd:
            bullets.append(Bullet(self.x+GRID_SIZE, self.y+20, self.atk, "normal"))
            bullets.append(Bullet(self.x+GRID_SIZE, self.y+60, self.atk, "normal"))
            self.last_time = now
            if sound_pea: sound_pea.play()

class CherryBomb(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, 60, 150, 0, 0, CHERRY_COLOR)
        self.explode = False
        self.explode_time = 0

    def draw(self):
        self.draw_cartoon()
        pygame.draw.circle(screen, CHERRY_COLOR, (self.x+30, self.y+30), 18)
        pygame.draw.circle(screen, CHERRY_COLOR, (self.x+50, self.y+35), 18)
        pygame.draw.line(screen, (0,100,0), (self.x+30,self.y+15), (self.x+50,self.y+15),2)
        self.draw_hp()

    def update(self, now):
        if not self.explode:
            self.explode = True
            self.explode_time = now
        if self.explode and now - self.explode_time > 800:
            for z in zombies[:]:
                if abs(z.y - self.y) < 10:
                    z.hp = 0
            if sound_explode: sound_explode.play()
            self.hp = 0

class PotatoMine(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, 50, 125, 0, 0, POTATO_COLOR)
        self.ready = False

    def draw(self):
        self.draw_cartoon()
        if not self.ready:
            pygame.draw.circle(screen, (80,50,10), (self.x+40, self.y+50), 20)
        else:
            pygame.draw.circle(screen, POTATO_COLOR, (self.x+40, self.y+40), 22)
            pygame.draw.circle(screen, RED, (self.x+40, self.y+30), 5)
        self.draw_hp()

    def update(self, now):
        if not self.ready:
            if now - self.last_time > 3000:
                self.ready = True
        else:
            for z in zombies:
                if abs(z.x - self.x) < 40 and abs(z.y - self.y) < 40:
                    z.hp -= 200
                    if sound_explode: sound_explode.play()
                    self.hp = 0
                    break

class SunPea(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, 120, 120, 18, 20, ORANGE)
    def draw(self):
        self.draw_cartoon()
        pygame.draw.circle(screen, ORANGE, (self.x+30, self.y+25), 22)
        pygame.draw.circle(screen, SUN_COLOR, (self.x+45, self.y+30), 10)
        self.draw_hp()
    def update(self, now):
        if now - self.last_time > 2500:
            global sun
            sun += self.sun_prod
        if now - self.last_time > 1000:
            bullets.append(Bullet(self.x+GRID_SIZE, self.y+GRID_SIZE//2, self.atk, "normal"))
        self.last_time = now

class IceWall(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, 600, 110, 12, 0, CYAN)
    def draw(self):
        self.draw_cartoon()
        pygame.draw.rect(screen, CYAN, (self.x+10, self.y+10, 60, 50), border_radius=15)
        pygame.draw.circle(screen, ICE_COLOR, (self.x+40, self.y+35), 12)
        self.draw_hp()
    def update(self, now):
        if now - self.last_time > 1200:
            bullets.append(Bullet(self.x+GRID_SIZE, self.y+GRID_SIZE//2, self.atk, "ice"))
            self.last_time = now

class BoomDouble(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, 130, 220, 25, 0, RED)
    def draw(self):
        self.draw_cartoon()
        pygame.draw.rect(screen, RED, (self.x+15, self.y+10, 50, 45), border_radius=10)
        pygame.draw.circle(screen, ORANGE, (self.x+60, self.y+20), 8)
        pygame.draw.circle(screen, ORANGE, (self.x+60, self.y+40), 8)
        self.draw_hp()
    def update(self, now):
        if now - self.last_time > 900:
            bullets.append(Bullet(self.x+GRID_SIZE, self.y+15, self.atk, "normal"))
            bullets.append(Bullet(self.x+GRID_SIZE, self.y+45, self.atk, "normal"))
            self.last_time = now

class SunWall(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, 550, 80, 0, 30, SUN_COLOR)
    def draw(self):
        self.draw_cartoon()
        pygame.draw.rect(screen, SUN_COLOR, (self.x+10, self.y+10, 60, 50), border_radius=15)
        pygame.draw.circle(screen, (255,180,0), (self.x+40, self.y+30), 15)
        self.draw_hp()
    def update(self, now):
        if now - self.last_time > 2500:
            global sun
            sun += self.sun_prod
            self.last_time = now

class MinePea(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, 90, 130, 22, 0, POTATO_COLOR)
    def draw(self):
        self.draw_cartoon()
        pygame.draw.circle(screen, POTATO_COLOR, (self.x+35, self.y+35), 25)
        pygame.draw.circle(screen, PEA_COLOR, (self.x+60, self.y+35), 7)
        self.draw_hp()
    def update(self, now):
        if now - self.last_time > 1000:
            bullets.append(Bullet(self.x+GRID_SIZE, self.y+GRID_SIZE//2, self.atk, "normal"))
            self.last_time = now

class HybridPlant(Plant):
    def __init__(self, x, y, p1, p2):
        new_hp = (p1.max_hp + p2.max_hp) // 2 + random.randint(20, 80)
        new_atk = (p1.atk + p2.atk) // 2 + random.randint(5, 15)
        new_sun = (p1.sun_prod + p2.sun_prod) // 2
        r = (p1.color[0] + p2.color[0]) // 2
        g = (p1.color[1] + p2.color[1]) // 2
        b = (p1.color[2] + p2.color[2]) // 2
        super().__init__(x, y, new_hp, 0, new_atk, new_sun, (r,g,b))
    def draw(self):
        self.draw_cartoon()
        pygame.draw.circle(screen, self.color, (self.x+40, self.y+30), 30)
        pygame.draw.circle(screen, WHITE, (self.x+40, self.y+30), 12)
        self.draw_hp()
    def update(self, now):
        if self.sun_prod > 0 and now - self.last_time > 3000:
            global sun
            sun += self.sun_prod
        if self.atk > 0 and now - self.last_time > 1200:
            bullets.append(Bullet(self.x+GRID_SIZE, self.y+GRID_SIZE//2, self.atk, "normal"))
            self.last_time = now
            if sound_pea: sound_pea.play()

def get_plant_cls(p):
    if isinstance(p, SunFlower): return "sunflower"
    if isinstance(p, Peashooter): return "pea"
    if isinstance(p, IceShooter): return "ice"
    if isinstance(p, WallNut): return "wall"
    if isinstance(p, DoublePea): return "doublepea"
    if isinstance(p, CherryBomb): return "cherry"
    if isinstance(p, PotatoMine): return "potato"
    return "other"

def hybrid_plant(p1, p2):
    if p1 is p2: return
    t1, t2 = get_plant_cls(p1), get_plant_cls(p2)
    new_p = None
    gx = p1.x // GRID_SIZE
    gy = p1.y // GRID_SIZE
    
    if (t1=="sunflower" and t2=="pea") or (t1=="pea" and t2=="sunflower"):
        new_p = SunPea(p1.x, p1.y)
    elif (t1=="ice" and t2=="wall") or (t1=="wall" and t2=="ice"):
        new_p = IceWall(p1.x, p1.y)
    elif (t1=="doublepea" and t2=="cherry") or (t1=="cherry" and t2=="doublepea"):
        new_p = BoomDouble(p1.x, p1.y)
    elif (t1=="sunflower" and t2=="wall") or (t1=="wall" and t2=="sunflower"):
        new_p = SunWall(p1.x, p1.y)
    elif (t1=="potato" and t2=="pea") or (t1=="pea" and t2=="potato"):
        new_p = MinePea(p1.x, p1.y)
    else:
        new_p = HybridPlant(p1.x, p1.y, p1, p2)
    
    plants.remove(p1)
    plants.remove(p2)
    plants.append(new_p)
    if sound_hybrid: sound_hybrid.play()

class Bullet:
    def __init__(self, x, y, dmg, btype):
        self.x = x
        self.y = y
        self.dmg = dmg
        self.type = btype
        self.speed = 6
        self.radius = 6
    def update(self):
        self.x += self.speed
    def draw(self):
        if self.type == "ice":
            pygame.draw.circle(screen, ICE_COLOR, (int(self.x), int(self.y)), self.radius)
        else:
            pygame.draw.circle(screen, (0,255,0), (int(self.x), int(self.y)), self.radius)

class Zombie:
    def __init__(self, row_idx, ztype="normal"):
        self.row = row_idx
        self.x = WIDTH
        self.y = row_idx * GRID_SIZE
        self.ztype = ztype
        if ztype == "bucket":
            self.hp = 350
            self.speed = 0.35
            self.color = BUCKET_COLOR
        else:
            self.hp = 180
            self.speed = 0.45
            self.color = ZOMBIE_COLOR
        self.max_hp = self.hp
        self.atk_cd = 1000
        self.last_atk = 0

    def draw(self):
        pygame.draw.rect(screen, self.color, (self.x+8, self.y+10, 65, 55), border_radius=8)
        pygame.draw.circle(screen, WHITE, (self.x+25, self.y+25), 6)
        pygame.draw.circle(screen, WHITE, (self.x+50, self.y+25), 6)
        
        h = 6
        ratio = self.hp / self.max_hp
        pygame.draw.rect(screen, RED, (self.x+5, self.y-8, GRID_SIZE-10, h))
        pygame.draw.rect(screen, (0,255,0), (self.x+5, self.y-8, int((GRID_SIZE-10)*ratio), h))

    def update(self, now):
        hit_plant = None
        for p in plants:
            if p.y == self.y and abs(p.x - self.x) < GRID_SIZE:
                hit_plant = p
                break
        if hit_plant:
            if now - self.last_atk > self.atk_cd:
                hit_plant.hp -= 15
                self.last_atk = now
                if sound_zombie: sound_zombie.play()
            return
        self.x -= self.speed
        if self.x < 20:
            global game_over
            game_over = True

def spawn_zombie():
    r = random.randint(0, ROW-1)
    zombies.append(Zombie(r, "bucket" if random.random()<0.25 else "normal"))

def draw_grass():
    screen.fill(GREEN_GRASS)
    for r in range(ROW+1):
        pygame.draw.line(screen, GRID_LINE, (0, r*GRID_SIZE), (WIDTH, r*GRID_SIZE), 2)
    for c in range(COL+1):
        pygame.draw.line(screen, GRID_LINE, (c*GRID_SIZE, 0), (c*GRID_SIZE, ROW*GRID_SIZE), 2)

def draw_ui():
    sun_txt = font_big.render(f"Sun: {sun}", True, WHITE)
    screen.blit(sun_txt, (20, ROW*GRID_SIZE + 15))
    btn_y = ROW*GRID_SIZE + 50
    
    cards = [
        (SUN_COLOR, "Sun50"),
        (PEA_COLOR, "Pea100"),
        (WALL_COLOR, "Nut50"),
        (ICE_COLOR, "Ice175"),
        ((0,150,80), "Double200"),
        (CHERRY_COLOR, "Bomb150"),
        (POTATO_COLOR, "Mine125")
    ]
    
    for i in range(7):
        color, text = cards[i]
        pygame.draw.rect(screen, color, (i*GRID_SIZE, btn_y, GRID_SIZE, 60))
        if i == selected_card:
            pygame.draw.rect(screen, WHITE, (i*GRID_SIZE, btn_y, GRID_SIZE, 60), 3)
        screen.blit(font.render(text, True, BLACK if i<5 else WHITE), (i*GRID_SIZE+5, btn_y+15))

    if selected_plant:
        tip = font.render("Select another plant", True, RED)
        screen.blit(tip, (300, ROW*GRID_SIZE+15))

def main():
    global sun, game_over, selected_plant, spawn_time, selected_card
    while True:
        now = pygame.time.get_ticks()
        clock.tick(FPS)
        draw_grass()

        for e in pygame.event.get():
            if e.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            if e.type == pygame.MOUSEBUTTONDOWN:
                mx, my = pygame.mouse.get_pos()
                gx = mx // GRID_SIZE
                gy = my // GRID_SIZE

                clicked_plant = None
                for p in plants:
                    if p.x//GRID_SIZE == gx and p.y//GRID_SIZE == gy:
                        clicked_plant = p
                        break
                if clicked_plant:
                    if selected_plant is None:
                        selected_plant = clicked_plant
                    else:
                        hybrid_plant(selected_plant, clicked_plant)
                        selected_plant = None
                    continue

                if my >= ROW * GRID_SIZE:
                    if 0 <= gx <7:
                        selected_card = gx
                else:
                    if 0<=gx<COL and 0<=gy<ROW and not grid_used[gx][gy] and selected_card!=-1:
                        cost_list = [50,100,50,175,200,150,125]
                        if sun >= cost_list[selected_card]:
                            if selected_card ==0:
                                plants.append(SunFlower(gx*GRID_SIZE, gy*GRID_SIZE))
                            elif selected_card ==1:
                                plants.append(Peashooter(gx*GRID_SIZE, gy*GRID_SIZE))
                            elif selected_card ==2:
                                plants.append(WallNut(gx*GRID_SIZE, gy*GRID_SIZE))
                            elif selected_card ==3:
                                plants.append(IceShooter(gx*GRID_SIZE, gy*GRID_SIZE))
                            elif selected_card ==4:
                                plants.append(DoublePea(gx*GRID_SIZE, gy*GRID_SIZE))
                            elif selected_card ==5:
                                plants.append(CherryBomb(gx*GRID_SIZE, gy*GRID_SIZE))
                            elif selected_card ==6:
                                pm = PotatoMine(gx*GRID_SIZE, gy*GRID_SIZE)
                                pm.last_time = now
                                plants.append(pm)
                            sun -= cost_list[selected_card]
                            grid_used[gx][gy] = True
                            if sound_plant: sound_plant.play()

        if now - spawn_time > 4000:
            spawn_zombie()
            spawn_time = now

        for p in plants[:]:
            p.update(now)
            if p.hp <=0:
                gx_p = p.x//GRID_SIZE
                gy_p = p.y//GRID_SIZE
                grid_used[gx_p][gy_p] = False
                plants.remove(p)
        for z in zombies[:]:
            z.update(now)
            if z.hp <=0: zombies.remove(z)
        for b in bullets[:]:
            b.update()
            if b.x > WIDTH:
                bullets.remove(b)
                continue
            for z in zombies:
                if abs(b.x-(z.x+40))<35 and abs(b.y-(z.y+40))<35:
                    z.hp -= b.dmg
                    bullets.remove(b)
                    break

        for p in plants: p.draw()
        for z in zombies: z.draw()
        for b in bullets: b.draw()
        draw_ui()

        if game_over:
            over_txt = font_big.render("Game Over!", True, RED)
            screen.blit(over_txt, (WIDTH//2-50, HEIGHT//2))

        pygame.display.flip()

if __name__ == "__main__":
    main()