import numpy as np
import numexpr
import re
import json

def parse_weight_schedule(string, max_frames):
    string = str(string).replace(' ', '')
    keyframes = {}
    pattern = r'((?P<frame>[0-9]+):[\s]*[\(](?P<param>[\S\s]*?)[\)])'
    for match in re.finditer(pattern, string):
        keyframes[int(match.group("frame"))] = match.group("param")
    if 0 not in keyframes: keyframes[0] = "0"
    
    series = np.zeros(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 max_frames
        formula = keyframes[f_start]
        for f in range(f_start, f_end):
            try:
                val = numexpr.evaluate(formula, local_dict={'t':f, 'pi':np.pi, 'sin':np.sin, 'cos':np.cos, 'tan':np.tan})
                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 parse_prompts(prompt_json):
    try:
        data = json.loads(prompt_json.replace("'", '"'))
        return {int(k): v for k, v in data.items()}
    except:
        return {0: "error parsing prompts"}