import pygame
import random
import sys
import math

pygame.init()
WIDTH, HEIGHT = 900, 650
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("植物大战僵尸 终极扩展中文版")
clock = pygame.time.Clock()
FPS = 60

# 基础颜色
WHITE = (255,255,255)
BLACK = (0,0,0)
RED = (220,0,0)
GREEN = (30,140,30)
YELLOW = (255,210,0)
SKY = (135,205,240)
BROWN = (110,65,15)
ICE_BLUE = (90,170,255)
CHERRY_RED = (255,35,35)
NUT_BROWN = (145,85,30)
POTATO_BROWN = (80,50,20)
CHOMP_PURPLE = (160,40,180)
DOUBLE_SUN = (255,180,0)
CABBAGE_GREEN = (60,190,80)
CONE_ORANGE = (255,130,0)
FLAG_RED = (190,0,0)
FOOTBALL_GRAY = (70,70,70)
LADDER_GRAY = (150,150,150)
BOSS_PURPLE = (80,20,120)

# 字体兼容中文
try:
    font = pygame.font.Font("simhei.ttf", 20)
    font_big = pygame.font.Font("simhei.ttf", 38)
except:
    font = pygame.font.SysFont("Microsoft YaHei", 20)
    font_big = pygame.font.SysFont("Microsoft YaHei", 38)

# 全局游戏变量
sun = 200
game_over = False
win = False
ROW = 5
COL = 10
CELL_W = WIDTH // COL
CELL_H = (HEIGHT - 130) // ROW
place_mode = 0
plant_group = pygame.sprite.Group()
bullet_group = pygame.sprite.Group()
zombie_group = pygame.sprite.Group()
explosion_group = pygame.sprite.Group()
cabbage_ball_group = pygame.sprite.Group()

# 波次设置 总12波+最终BOSS
max_wave = 12
wave_count = 0
zombie_spawn_timer = 0

# 植物父类
class Plant(pygame.sprite.Sprite):
    def __init__(self, col, row, cost, hp):
        super().__init__()
        self.col = col
        self.row = row
        self.rect = pygame.Rect(col*CELL_W+2, row*CELL_H+130+2, CELL_W-4, CELL_H-4)
        self.cost = cost
        self.hp = hp
        self.max_hp = hp
        self.color = GREEN

# 1.向日葵
class SunFlower(Plant):
    def __init__(self,c,r):
        super().__init__(c,r,50,100)
        self.color = YELLOW
        self.timer = 0
    def update(self):
        global sun
        self.timer +=1
        if self.timer >=330:
            sun +=25
            self.timer=0

# 2.双子向日葵
class TwinSunflower(Plant):
    def __init__(self,c,r):
        super().__init__(c,r,150,100)
        self.color = DOUBLE_SUN
        self.timer=0
    def update(self):
        global sun
        self.timer +=1
        if self.timer >=330:
            sun +=50
            self.timer=0

# 3.普通豌豆
class Peashooter(Plant):
    def __init__(self,c,r):
        super().__init__(c,r,100,100)
        self.color=GREEN
        self.cd=0
    def update(self,bullet_gp,zom_gp):
        self.cd+=1
        if self.cd>=90:
            if any(z.row==self.row for z in zom_gp):
                bullet_gp.add(Bullet(self.rect.centerx,self.rect.centery,self.row,25,False))
            self.cd=0

# 4.寒冰射手
class IcePea(Plant):
    def __init__(self,c,r):
        super().__init__(c,r,175,100)
        self.color=ICE_BLUE
        self.cd=0
    def update(self,bullet_gp,zom_gp):
        self.cd+=1
        if self.cd>=90:
            if any(z.row==self.row for z in zom_gp):
                bullet_gp.add(Bullet(self.rect.centerx,self.rect.centery,self.row,20,True))
            self.cd=0

# 5.坚果墙
class WallNut(Plant):
    def __init__(self,c,r):
        super().__init__(c,r,50,450)
        self.color=NUT_BROWN

# 6.土豆地雷
class PotatoMine(Plant):
    def __init__(self,c,r):
        super().__init__(c,r,25,80)
        self.color=POTATO_BROWN
        self.arm_cd=180
        self.triggered=False
    def update(self,zom_gp,exp_gp):
        if self.arm_cd>0:
            self.arm_cd-=1
            return
        if not self.triggered:
            for z in zom_gp:
                if self.rect.colliderect(z.rect):
                    self.triggered=True
                    exp_gp.add(Explosion(self.rect.centerx,self.rect.centery,CELL_W*2))
                    for zi in zom_gp:
                        dx=zi.rect.centerx-self.rect.centerx
                        dy=zi.rect.centery-self.rect.centery
                        if math.hypot(dx,dy) < CELL_W*2:
                            zi.hp -=180
                    self.kill()
                    break

# 7.大嘴花
class Chomper(Plant):
    def __init__(self,c,r):
        super().__init__(c,r,150,120)
        self.color=CHOMP_PURPLE
        self.eat_cd=0
        self.eating_target=None
    def update(self,zom_gp):
        if self.eating_target:
            self.eat_cd+=1
            if self.eat_cd>=220:
                self.eating_target.kill()
                self.eating_target=None
                self.eat_cd=0
            return
        # 找正前方同行最近僵尸
        min_dist=9999
        target=None
        for z in zom_gp:
            if z.row==self.row and z.rect.x>self.rect.x:
                dist=z.rect.x-self.rect.x
                if dist<CELL_W*1.8 and dist<min_dist:
                    min_dist=dist
                    target=z
        if target:
            self.eating_target=target

# 8.卷心菜投手
class CabbagePitcher(Plant):
    def __init__(self,c,r):
        super().__init__(c,r,100,100)
        self.color=CABBAGE_GREEN
        self.cd=0
    def update(self,ball_gp,zom_gp):
        self.cd+=1
        if self.cd>=110:
            if any(z.row==self.row for z in zom_gp):
                ball_gp.add(CabbageBall(self.rect.centerx,self.rect.top))
            self.cd=0

# 9.樱桃炸弹
class CherryBomb(Plant):
    def __init__(self,c,r):
        super().__init__(c,r,150,80)
        self.color=CHERRY_RED
        self.count=160
    def update(self,zom_gp,exp_gp):
        self.count-=1
        if self.count<=0:
            exp_gp.add(Explosion(self.rect.centerx,self.rect.centery,CELL_W*3))
            for z in zom_gp:
                d=math.hypot(z.rect.centerx-self.rect.centerx,z.rect.centery-self.rect.centery)
                if d<CELL_W*3:
                    z.hp=0
            self.kill()

# 普通子弹
class Bullet(pygame.sprite.Sprite):
    def __init__(self,x,y,row,dmg,slow):
        super().__init__()
        self.image=pygame.Surface((10,10))
        self.image.fill(ICE_BLUE if slow else YELLOW)
        self.rect=self.image.get_rect(center=(x,y))
        self.speed=7
        self.row=row
        self.dmg=dmg
        self.slow=slow
    def update(self):
        self.rect.x+=self.speed
        if self.rect.left>WIDTH:
            self.kill()

# 卷心菜投掷球
class CabbageBall(pygame.sprite.Sprite):
    def __init__(self,x,y):
        super().__init__()
        self.r=8
        self.x=x
        self.y=y
        self.vx=4
        self.vy=-3.5
        self.dmg=30
    def update(self):
        self.x+=self.vx
        self.y+=self.vy
        self.vy+=0.18
        self.rect=pygame.Rect(self.x-self.r,self.y-self.r,self.r*2,self.r*2)
        if self.x>WIDTH or self.y>HEIGHT:
            self.kill()
    def draw(self,screen):
        pygame.draw.circle(screen,CABBAGE_GREEN,(int(self.x),int(self.y)),self.r)

# 爆炸特效
class Explosion(pygame.sprite.Sprite):
    def __init__(self,x,y,maxr):
        super().__init__()
        self.cx=x
        self.cy=y
        self.r=0
        self.maxr=maxr
        self.life=35
    def update(self):
        self.r+=6
        self.life-=1
        if self.life<=0:
            self.kill()
    def draw(self,screen):
        pygame.draw.circle(screen,(255,110,0),(int(self.cx),int(self.cy)),int(self.r),4)

# 僵尸父类
class Zombie(pygame.sprite.Sprite):
    def __init__(self,row,hp,speed,color):
        super().__init__()
        self.row=row
        self.rect=pygame.Rect(WIDTH-48,row*CELL_H+130,44,CELL_H-8)
        self.hp=hp
        self.max_hp=hp
        self.base_speed=speed
        self.speed=speed
        self.color=color
        self.eat_cd=0
        self.slow_timer=0
        self.ladder_plant=None #扶梯专用

    def update(self,plant_gp):
        # 减速判定
        if self.slow_timer>0:
            self.slow_timer-=1
            self.speed=self.base_speed*0.38
        else:
            self.speed=self.base_speed

        # 扶梯僵尸特殊逻辑：架梯子越过坚果
        if hasattr(self,"ladder_plant") and self.ladder_plant:
            self.rect.x -= self.speed*1.6
            if self.rect.right < self.ladder_plant.rect.left:
                self.ladder_plant=None
            return

        hit_p=None
        for p in plant_gp:
            if p.row==self.row and self.rect.colliderect(p.rect):
                hit_p=p
                break
        if hit_p:
            self.eat_cd+=1
            if self.eat_cd>=52:
                hit_p.hp-=12
                self.eat_cd=0
                if hit_p.hp<=0:
                    hit_p.kill()
        else:
            self.rect.x -= self.speed
        # 啃穿房子游戏结束
        if self.rect.left <= 5:
            global game_over
            game_over=True

# 1.普通僵尸
class NormalZombie(Zombie):
    def __init__(self,row):
        super().__init__(row,100,0.3,BROWN)

# 2.路障僵尸
class ConeZombie(Zombie):
    def __init__(self,row):
        super().__init__(row,230,0.27,CONE_ORANGE)

# 3.旗帜僵尸
class FlagZombie(Zombie):
    def __init__(self,row):
        super().__init__(row,140,0.48,FLAG_RED)

# 4.橄榄球僵尸（高移速高血量）
class FootballZombie(Zombie):
    def __init__(self,row):
        super().__init__(row,380,0.62,FOOTBALL_GRAY)

# 5.扶梯僵尸（可以搭梯子越过前排坚果）
class LadderZombie(Zombie):
    def __init__(self,row):
        super().__init__(row,180,0.33,LADDER_GRAY)
    def update(self,plant_gp):
        # 首次碰到坚果就架梯
        if not self.ladder_plant:
            for p in plant_gp:
                if isinstance(p,WallNut) and p.row==self.row and self.rect.colliderect(p.rect):
                    self.ladder_plant=p
                    break
        super().update(plant_gp)

# 6.BOSS巨型僵尸（超大血量慢速）
class BossZombie(Zombie):
    def __init__(self,row):
        super().__init__(row,1200,0.15,BOSS_PURPLE)
        self.rect.width=70
        self.rect.height=CELL_H-4

# 按钮配置
btn_info = [
    {"name":"向日葵50","mode":1,"col":YELLOW,"rect":pygame.Rect(10,30,95,34)},
    {"name":"双子150","mode":2,"col":DOUBLE_SUN,"rect":pygame.Rect(110,30,95,34)},
    {"name":"豌豆100","mode":3,"col":GREEN,"rect":pygame.Rect(210,30,95,34)},
    {"name":"寒冰175","mode":4,"col":ICE_BLUE,"rect":pygame.Rect(310,30,95,34)},
    {"name":"坚果50","mode":5,"col":NUT_BROWN,"rect":pygame.Rect(410,30,95,34)},
    {"name":"土豆25","mode":6,"col":POTATO_BROWN,"rect":pygame.Rect(510,30,95,34)},
    {"name":"大嘴150","mode":7,"col":CHOMP_PURPLE,"rect":pygame.Rect(610,30,95,34)},
    {"name":"卷心菜100","mode":8,"col":CABBAGE_GREEN,"rect":pygame.Rect(710,30,95,34)},
    {"name":"樱桃150","mode":9,"col":CHERRY_RED,"rect":pygame.Rect(810,30,95,34)},
]

def draw_scene():
    screen.fill(SKY)
    # 绘制草地网格
    for r in range(ROW):
        for c in range(COL):
            rc=pygame.Rect(c*CELL_W,r*CELL_H+130,CELL_W,CELL_H)
            pygame.draw.rect(screen,(105,172,70),rc,2)
    # 顶部文字
    top_text=font.render(f"阳光:{sun}  波次:{wave_count}/{max_wave}  共12波+最终BOSS",True,RED)
    screen.blit(top_text,(10,5))
    # 绘制种植按钮
    global place_mode
    for b in btn_info:
        pygame.draw.rect(screen,b["col"],b["rect"])
        t=font.render(b["name"],True,BLACK)
        screen.blit(t,(b["rect"].x+4,b["rect"].y+4))
    # 绘制植物+血条
    for p in plant_group:
        pygame.draw.rect(screen,p.color,p.rect)
        hp_w=p.rect.width*(p.hp/p.max_hp)
        pygame.draw.rect(screen,RED,(p.rect.x,p.rect.y-5,hp_w,4))
    # 子弹
    bullet_group.draw(screen)
    # 卷心菜球手动绘制
    for cb in cabbage_ball_group:
        cb.draw(screen)
    # 僵尸+血条
    for z in zombie_group:
        pygame.draw.rect(screen,z.color,z.rect)
        hp_w=z.rect.width*(z.hp/z.max_hp)
        pygame.draw.rect(screen,RED,(z.rect.x,z.rect.y-5,hp_w,4))
    # 爆炸特效
    for exp in explosion_group:
        exp.draw(screen)
    # 结局提示
    if game_over:
        over=font_big.render("游戏失败！僵尸吃掉脑子",True,RED)
        screen.blit(over,(WIDTH//2-270,HEIGHT//2))
    if win:
        wint=font_big.render("全部僵尸清除，通关胜利！",True,GREEN)
        screen.blit(wint,(WIDTH//2-290,HEIGHT//2))

# 主循环
running=True
while running:
    clock.tick(FPS)
    draw_scene()
    # 事件捕获
    for e in pygame.event.get():
        if e.type==pygame.QUIT:
            running=False
        if e.type==pygame.MOUSEBUTTONDOWN and not game_over and not win:
            mx,my=pygame.mouse.get_pos()
            # 点击选择植物
            for b in btn_info:
                if b["rect"].collidepoint(mx,my):
                    place_mode=b["mode"]
            # 点击地块种植
            if my>130:
                c_idx=mx//CELL_W
                r_idx=(my-130)//CELL_H
                if 0<=c_idx<COL and 0<=r_idx<ROW:
                    occupied=False
                    for p in plant_group:
                        if p.col==c_idx and p.row==r_idx:
                            occupied=True
                            break
                    if occupied:
                        continue
                    # 按模式生成植物
                    if place_mode==1 and sun>=50:
                        plant_group.add(SunFlower(c_idx,r_idx));sun-=50
                    elif place_mode==2 and sun>=150:
                        plant_group.add(TwinSunflower(c_idx,r_idx));sun-=150
                    elif place_mode==3 and sun>=100:
                        plant_group.add(Peashooter(c_idx,r_idx));sun-=100
                    elif place_mode==4 and sun>=175:
                        plant_group.add(IcePea(c_idx,r_idx));sun-=175
                    elif place_mode==5 and sun>=50:
                        plant_group.add(WallNut(c_idx,r_idx));sun-=50
                    elif place_mode==6 and sun>=25:
                        plant_group.add(PotatoMine(c_idx,r_idx));sun-=25
                    elif place_mode==7 and sun>=150:
                        plant_group.add(Chomper(c_idx,r_idx));sun-=150
                    elif place_mode==8 and sun>=100:
                        plant_group.add(CabbagePitcher(c_idx,r_idx));sun-=100
                    elif place_mode==9 and sun>=150:
                        plant_group.add(CherryBomb(c_idx,r_idx));sun-=150

    if not game_over and not win:
        # 批量更新所有植物
        for pl in plant_group:
            if isinstance(pl,(SunFlower,TwinSunflower)):
                pl.update()
            elif isinstance(pl,(Peashooter,IcePea)):
                pl.update(bullet_group,zombie_group)
            elif isinstance(pl,PotatoMine):
                pl.update(zombie_group,explosion_group)
            elif isinstance(pl,Chomper):
                pl.update(zombie_group)
            elif isinstance(pl,CabbagePitcher):
                pl.update(cabbage_ball_group,zombie_group)
            elif isinstance(pl,CherryBomb):
                pl.update(zombie_group,explosion_group)

        # 子弹碰撞僵尸
        bullet_group.update()
        hit_bullet=pygame.sprite.groupcollide(zombie_group,bullet_group,False,True)
        for z,bullets in hit_bullet.items():
            for b in bullets:
                z.hp -= b.dmg
                if b.slow:
                    z.slow_timer=190
                if z.hp<=0:
                    z.kill()

        # 卷心菜球碰撞
        for cb in cabbage_ball_group:
            cb.update()
            for z in zombie_group:
                if cb.rect and cb.rect.colliderect(z.rect):
                    z.hp -= cb.dmg
                    cb.kill()
                    if z.hp<=0:
                        z.kill()

        # 僵尸刷新逻辑，波次随进程加快刷新速度
        zombie_spawn_timer +=1
        base_interval=850-wave_count*45
        base_interval=max(220,base_interval)
        if zombie_spawn_timer>=base_interval and wave_count<max_wave:
            zombie_spawn_timer=0
            wave_count+=1
            r=random.randint(0,ROW-1)
            rd=random.random()
            if rd<0.32:
                zombie_group.add(NormalZombie(r))
            elif rd<0.55:
                zombie_group.add(ConeZombie(r))
            elif rd<0.72:
                zombie_group.add(FlagZombie(r))
            elif rd<0.88:
                zombie_group.add(FootballZombie(r))
            else:
                zombie_group.add(LadderZombie(r))
            # 第12波结束刷最终BOSS
            if wave_count==max_wave:
                boss_row=random.randint(0,ROW-1)
                zombie_group.add(BossZombie(boss_row))

        # 僵尸整体AI
        zombie_group.update(plant_group)
        explosion_group.update()

        # 胜利判定：12波打完+场上无任何僵尸
        if wave_count>=max_wave and len(zombie_group)==0:
            win=True

    pygame.display.flip()

pygame.quit()
sys.exit()