import numpy as np
import cv2
import numexpr
import re
import torch
from PIL import Image

def parse_weight_string(string, max_frames):
    """
    Parses complex Deforum weight strings with math support (sin, cos, t).
    """
    string = re.sub(r'\s+', '', str(string))
    keyframes = {}
    parts = string.split(',')
    
    for part in parts:
        try:
            if ':' not in part: continue
            f_str, v_str = part.split(':', 1)
            keyframes[int(f_str)] = v_str.strip('()')
        except: continue
        
    if 0 not in keyframes: keyframes[0] = "0"
    
    series = np.zeros(int(max_frames))
    sorted_keys = sorted(keyframes.keys())
    
    for i in range(len(sorted_keys)):
        f_start = sorted_keys[i]
        f_end = sorted_keys[i+1] if i < len(sorted_keys)-1 else int(max_frames)
        formula = keyframes[f_start]
        
        for f in range(f_start, f_end):
            t = f
            try:
                val = numexpr.evaluate(formula, local_dict={'t': t, 'pi': np.pi, 'sin': np.sin, 'cos': np.cos, 'tan': np.tan, 'abs': np.abs})
                series[f] = float(val)
            except:
                try: series[f] = float(formula)
                except: series[f] = series[f-1] if f > 0 else 0.0
                
    return series

def get_border_mode(mode_str):
    return {
        'Reflect': cv2.BORDER_REFLECT_101,
        'Replicate': cv2.BORDER_REPLICATE,
        'Wrap': cv2.BORDER_WRAP,
        'Black': cv2.BORDER_CONSTANT
    }.get(mode_str, cv2.BORDER_REFLECT_101)

def maintain_colors(image, anchor, mode='LAB'):
    """
    Matches the color distribution of 'image' to 'anchor'.
    """
    if mode == 'None' or anchor is None: return image
    
    img_np = np.array(image).astype(np.uint8)
    anc_np = np.array(anchor).astype(np.uint8)
    
    if mode == 'LAB':
        img_cvt = cv2.cvtColor(img_np, cv2.COLOR_RGB2LAB)
        anc_cvt = cv2.cvtColor(anc_np, cv2.COLOR_RGB2LAB)
        for i in range(3):
            img_cvt[:,:,i] = np.clip(img_cvt[:,:,i] - img_cvt[:,:,i].mean() + anc_cvt[:,:,i].mean(), 0, 255)
        out = cv2.cvtColor(img_cvt, cv2.COLOR_LAB2RGB)
        
    elif mode == 'HSV':
        img_cvt = cv2.cvtColor(img_np, cv2.COLOR_RGB2HSV)
        anc_cvt = cv2.cvtColor(anc_np, cv2.COLOR_RGB2HSV)
        # Match S and V, keep Hue
        for i in [1, 2]:
            img_cvt[:,:,i] = np.clip(img_cvt[:,:,i] - img_cvt[:,:,i].mean() + anc_cvt[:,:,i].mean(), 0, 255)
        out = cv2.cvtColor(img_cvt, cv2.COLOR_HSV2RGB)
        
    elif mode == 'RGB':
        for i in range(3):
            img_np[:,:,i] = np.clip(img_np[:,:,i] - img_np[:,:,i].mean() + anc_np[:,:,i].mean(), 0, 255)
        out = img_np
        
    else:
        return image

    return Image.fromarray(out)

def anim_frame_warp_2d(prev_img, args, border_mode_str):
    """
    Applies 2D affine transformation (Zoom, Rotate, Pan).
    """
    if prev_img is None: return None
    cv_img = np.array(prev_img)
    h, w = cv_img.shape[:2]
    center = (w // 2, h // 2)
    
    angle = args.get('angle', 0)
    zoom = args.get('zoom', 1.0)
    tx = args.get('tx', 0)
    ty = args.get('ty', 0)
    
    # Create Matrix
    mat = cv2.getRotationMatrix2D(center, angle, zoom)
    mat[0, 2] += tx
    mat[1, 2] += ty
    
    border = get_border_mode(border_mode_str)
    warped = cv2.warpAffine(cv_img, mat, (w, h), borderMode=border)
    return Image.fromarray(warped)

def add_noise(img, noise_amt):
    if noise_amt <= 0: return img
    img_np = np.array(img).astype(np.float32)
    # np.random.normal will use the seed set in the engine loop
    noise = np.random.normal(0, noise_amt * 255, img_np.shape).astype(np.float32)
    noisy = np.clip(img_np + noise, 0, 255).astype(np.uint8)
    return Image.fromarray(noisy)