import os
import math
from PIL import Image

C_PUPIL = (202, 0, 0, 255)
C_WHITE = (255, 255, 255, 255)
C_BODY = (47, 96, 130, 255)

def animate_rat_v5(img, direction):
    w, h = img.size
    frames = []
    
    pixels = img.load()
    bg_color = (125, 125, 125, 255)
    
    for frame_idx in range(4):
        frame = Image.new('RGBA', (w, h), bg_color)
        fp = frame.load()
        
        offset_mag = 0
        bounce_y = 0
        is_blink = False
        
        if frame_idx == 1:
            offset_mag = -8  
            bounce_y = -1
        elif frame_idx == 3:
            offset_mag = 8   
            bounce_y = -1
            is_blink = True  
            
        for y in range(h):
            for x in range(w):
                px = pixels[x, y]
                if px != bg_color and px[3] > 0:
                    is_tail = False
                    if direction == 'UP' and y > 40: is_tail = True
                    if direction == 'DOWN' and y < 24: is_tail = True
                    if direction == 'LEFT' and x > 40: is_tail = True
                    if direction == 'RIGHT' and x < 24: is_tail = True
                    
                    nx, ny = x, y
                    
                    if is_tail:
                        if direction == 'UP':
                            factor = (y - 40) / 20.0
                            nx = int(x + offset_mag * (factor**1.8))
                            ny = y + bounce_y # FIX: tail bounces too!
                        elif direction == 'DOWN':
                            factor = (24 - y) / 20.0
                            nx = int(x + offset_mag * (factor**1.8))
                            ny = y + bounce_y # FIX: tail bounces too!
                        elif direction == 'LEFT':
                            factor = (x - 40) / 20.0
                            ny = int(y + offset_mag * (factor**1.8)) + bounce_y
                        elif direction == 'RIGHT':
                            factor = (24 - x) / 20.0
                            ny = int(y + offset_mag * (factor**1.8)) + bounce_y
                    else:
                        ny = y + bounce_y
                        
                    if 0 <= nx < w and 0 <= ny < h:
                        if is_blink and px in [C_WHITE, C_PUPIL]:
                            fp[nx, ny] = C_BODY
                        else:
                            fp[nx, ny] = px
                            
        frames.append(frame)
    return frames

os.makedirs('/home/enne2/dev/mice/assets/Rat_Animated', exist_ok=True)

for d in ['UP', 'DOWN', 'LEFT', 'RIGHT']:
    img = Image.open(f'/home/enne2/dev/mice/assets/Rat/BMP_MALE_{d}.png').convert('RGBA')
    frames = animate_rat_v5(img, d)
    
    w, h = img.size
    sheet = Image.new('RGBA', (w * 4, h))
    for i, f in enumerate(frames):
        sheet.paste(f, (i * w, 0))
        
    sheet.save(f'/home/enne2/dev/mice/assets/Rat_Animated/MALE_{d}_anim.png')

print("Animations fixed: no more gap between tail and body!")