"""Color lock for the image captions (Qwen-Image-2.1 only): the first color a caption gives a recurring subject is kept,
a later caption that gives it another color is rewritten to the first one, and a locked animal mentioned without a color
gets it filled in ('a large, wet octopus' -> 'a large, wet blue octopus'), because each picture is drawn from its own caption
alone. 'safe' mode locks only animals and a lone person's hair and worn clothing. Deterministic, stdlib only.
Assembled verbatim from scripts/eval/writer_metrics.py (word lists, stop_tok, lemma) and scripts/eval/caption_lock.py
(spans, fill, lock); measured there on 3 writers x 76 articles: recurring subjects that change color 43.8/39.3/39.7% ->
32.1/29.8/30.1% with the lock alone; the fill step keeps an animal's color in pictures whose captions omit it.
"""
import re

COLORS = {'red', 'orange', 'yellow', 'green', 'blue', 'purple', 'pink', 'brown', 'black', 'white', 'gray', 'grey', 'beige',
          'tan', 'teal', 'navy', 'maroon', 'gold', 'golden', 'silver', 'violet', 'turquoise', 'blonde', 'blond', 'auburn',
          'ginger', 'lavender', 'magenta', 'cyan', 'olive', 'cream', 'burgundy', 'crimson', 'lime', 'indigo', 'mint'}

NORM = {'grey': 'gray', 'blond': 'blonde', 'golden': 'gold', 'violet': 'purple', 'lavender': 'purple', 'crimson': 'red',
        'navy': 'blue', 'maroon': 'red', 'burgundy': 'red', 'ginger': 'orange', 'auburn': 'red', 'magenta': 'pink',
        'cyan': 'blue', 'turquoise': 'teal', 'lime': 'green', 'olive': 'green', 'mint': 'green', 'indigo': 'blue',
        'cream': 'white', 'beige': 'tan'}

STOP = {'a', 'an', 'the', 'with', 'in', 'on', 'of', 'to', 'and', 'or', 'under', 'over', 'near', 'from', 'by', 'at', 'into',
        'onto', 'while', 'as', 'that', 'which', 'is', 'are', 'its', 'their', 'his', 'her', 'for', 'next', 'beside', 'behind',
        'above', 'below', 'inside', 'around', 'against', 'between', 'toward', 'towards', 'through', 'across', 'up', 'down',
        'out', 'off', 'who', 'it', 'they', 'he', 'she', 'one', 'another', 'other', 'each', 'both', 'this', 'then', 'also',
        'holds', 'hold', 'sits', 'sit', 'stands', 'stand', 'wears', 'wear', 'places', 'place', 'points', 'point', 'looks',
        'look', 'lies', 'rests', 'hangs', 'appears', 'appear', 'shows', 'reads', 'says', 'has', 'have', 'uses', 'use',
        'puts', 'put', 'pours', 'pulls', 'pushes', 'gives', 'takes', 'moves', 'turns', 'covers', 'contains', 'fills',
        'floats', 'swims', 'rides', 'walks', 'runs', 'smiles', 'waves', 'reaches', 'lifts', 'drops', 'writes', 'draws',
        'sets', 'set', 'shaped', 'filled', 'covered', 'colored', 'coloured', 'labeled', 'labelled', 'marked', 'wrapped'}

VERBS = set("""displays overlays pets waves grabs scrubs rubs wipes washes feeds pats taps touches kneels leans lies sleeps
flies jumps climbs spreads stirs cuts chops slices folds opens closes presses clicks types writes draws paints sprays pours
empties drips sprinkles shakes squeezes twists lifts carries drags hands offers stares laughs cries frowns nods shrugs talks
speaks listens hugs kisses waits watches follows leads guides wraps ties lays adds mixes blends bakes boils fries cooks eats
drinks bites chews licks sniffs dives glows shines sparkles reflects emerges hovers peeks peers stretches bends crouches
dances flows drips spills leaks grows blooms sprouts perches soars glides tilts rolls spins slides zips unzips buttons
dangles drapes clings crawls hops""".split())

NOUN_OK = {'belly', 'jelly', 'family', 'butterfly', 'dragonfly', 'firefly', 'assembly', 'supply', 'bully', 'hundred',
           'seabed', 'flatbed'}

def stop_tok(w):
    """A token that ends a color noun phrase: stop words, common caption verbs, participles (-ed) and adverbs (-ly)."""
    if w in STOP or w in VERBS:
        return True
    if w in NOUN_OK:
        return False
    return len(w) > 4 and (w.endswith('ed') or w.endswith('ly'))

MODS = {'small', 'large', 'big', 'tiny', 'little', 'long', 'short', 'curly', 'straight', 'wavy', 'dark', 'light', 'bright',
        'pale', 'striped', 'plaid', 'spotted', 'clear', 'round', 'square', 'thick', 'thin', 'shiny', 'fluffy', 'cute',
        'old', 'young', 'new', 'tall', 'wide', 'messy', 'neat', 'soft', 'hard', 'deep', 'full', 'empty', 'half', 'whole',
        'medium', 'giant', 'huge', 'fat', 'slim', 'cartoon', 'smiling', 'happy', 'sad', 'angry', 'friendly'}

ANNOT = {'arrow', 'check', 'checkmark', 'mark', 'x', 'circle', 'line', 'dot', 'outline', 'text', 'letter', 'lettering',
         'ink', 'background', 'surface', 'backdrop', 'highlight', 'glow', 'sparkle', 'star', 'cross', 'underline', 'word',
         'label', 'symbol', 'icon', 'box', 'border', 'bubble', 'number', 'font', 'stroke', 'swirl', 'motion', 'light',
         'shade', 'shadow', 'tint', 'color', 'colour', 'hue', 'area', 'side', 'part', 'section', 'stripe', 'marker',
         'sticker', 'tag', 'emoji', 'heart', 'lightning', 'bolt', 'plus', 'minus', 'sign', 'squiggle', 'splash',
         # settings and backgrounds: style, not a recurring subject
         'gradient', 'environment', 'sky', 'void', 'setting', 'scene', 'space', 'room', 'wall', 'floor', 'ground', 'tone',
         'backdrop', 'field', 'glare'}

PERSON_ATTR = {'hair', 'shirt', 't-shirt', 'tshirt', 'sweater', 'jacket', 'hoodie', 'dress', 'apron', 'cap', 'hat', 'glove',
               'pant', 'pants', 'jean', 'jeans', 'coat', 'short', 'shorts', 'skirt', 'top', 'blouse', 'uniform', 'vest',
               'scarf', 'suit', 'tie', 'sleeve', 'shoe', 'boot', 'sock', 'beard', 'mustache', 'sweatshirt', 'cardigan',
               'overall', 'overalls', 'robe', 'polo', 'tank', 'glasses', 'bow'}

MALE = {'man', 'boy', 'guy', 'gentleman', 'father', 'dad', 'grandfather', 'husband', 'king', 'businessman'}

FEMALE = {'woman', 'girl', 'lady', 'mother', 'mom', 'grandmother', 'wife', 'queen', 'businesswoman'}

def lemma(w):
    w = w.strip("-'")
    if w.endswith("'s"):
        w = w[:-2]
    if len(w) > 3 and w.endswith('es') and w[:-2].endswith(('sh', 'ch', 'x')):
        return w[:-2]
    if len(w) > 3 and w.endswith('s') and not w.endswith(('ss', 'us', 'is', 'ys')):
        return w[:-1]
    return w


ANIMALS = set('''octopus cat kitten dog puppy llama alpaca dragon raccoon goldfish fish crow raven pigeon bird parrot chicken
hen rooster duck goose horse pony cow pig goat sheep lamb rabbit bunny hamster guinea mouse rat squirrel fox wolf bear
panda tiger lion monkey frog toad turtle tortoise snake lizard gecko spider bee butterfly owl eagle hawk seagull penguin
dolphin whale shark crab lobster squid jellyfish snail worm deer moose elk camel elephant giraffe zebra unicorn ghost
zombie robot alien monster creature goblin troll fairy wizard dinosaur dino'''.split())


PEOPLE = MALE | FEMALE | {'person', 'child', 'kid', 'toddler', 'baby', 'people', 'friend', 'friends', 'two', 'couple', 'group',
                          'another', 'mayor', 'waiter', 'teacher', 'student', 'roommate', 'customer'}


WORN_CUE = {'in', 'wearing', 'wears', 'wore', 'with', 'and', 'dressed'}


TOK = re.compile(r"[A-Za-z][A-Za-z'\-]*|[.,;:!?()]")


SHADES = {'dark', 'light', 'bright', 'pale'}


_QUOTES = re.compile(r'"[^"\n]{1,200}"|“[^”\n]{1,200}”|‘[^’\n]{1,200}’|(?<![A-Za-z0-9])\'(?:[^\'\n]|(?<=[A-Za-z])\'(?=[A-Za-z])){1,200}\'(?![A-Za-z0-9])')


def quoted_spans(cap):
    """Character ranges of quoted text ('Pigeon Ring', \"OPEN\"): lettering the picture should show, never a subject's look."""
    return [(m.start(), m.end()) for m in _QUOTES.finditer(cap)]


def _in(pos, ranges):
    return any(a <= pos < b for a, b in ranges)


def spans(cap, mode='safe'):
    """[(start, end, head, colors)] for each color phrase in a caption, with character offsets of the color words."""
    q = quoted_spans(cap)
    toks = [(m.group(0).lower(), m.start(), m.end()) for m in TOK.finditer(cap)]
    n_people = sum(1 for t, _, _ in toks if t in PEOPLE)
    out, i = [], 0
    while i < len(toks):
        w = toks[i][0]
        if w in COLORS or (w in SHADES and i + 1 < len(toks) and toks[i + 1][0] in COLORS):
            s, j, cols = toks[i][1], i, set()
            while j < len(toks) and (toks[j][0] in COLORS or toks[j][0] in SHADES or toks[j][0] in ('and', 'or')):
                if toks[j][0] in COLORS:
                    cols.add(NORM.get(toks[j][0], toks[j][0]))
                j += 1
            while j > i and toks[j - 1][0] in ('and', 'or'):
                j -= 1
            e = toks[j - 1][2]
            run, k = [], j
            while k < len(toks) and len(run) < 3 and re.match(r'[a-z]', toks[k][0]) and not stop_tok(toks[k][0]) \
                    and not (toks[k][0].endswith('ing') and len(toks[k][0]) > 5) and toks[k][0] not in COLORS:
                if toks[k][0] in MODS and not run:
                    k += 1
                    continue
                run.append(toks[k][0]); k += 1
            if run and cols:
                head = lemma(run[-1])
                ok = head == 'hair' or (head not in ANNOT and lemma(run[0]) not in ANNOT)
                if ok and mode == 'safe':
                    before = {t for t, _, _ in toks[max(0, i - 4):i]}
                    ok = head in ANIMALS or (head in PERSON_ATTR and n_people <= 1 and (head == 'hair' or before & WORN_CUE))
                if ok and not _in(s, q):  # a color word inside quoted text is lettering: never a look, never edited
                    out.append((s, e, head, cols))
            i = max(j, i + 1)
        else:
            i += 1
    return out


def fill(cap, first):
    """Give a locked animal its color when a caption mentions it with none ('a large, wet octopus' -> 'a large, wet blue
    octopus'): the illustrator sees only this caption and would otherwise fall back to a default color."""
    have = {h for _, _, h, _ in spans(cap, 'all')}
    added = 0
    for head, (_, phrase) in first.items():
        if head not in ANIMALS or head in have:
            continue
        q = quoted_spans(cap)
        m = next((x for x in re.finditer(r'\b' + re.escape(head) + r'(e?s)?\b', cap, re.I) if not _in(x.start(), q)), None)
        if not m:  # only mentioned inside quoted text ('Pigeon Ring'): leave the lettering alone
            continue
        before = cap[:m.start()]
        prev = [w.lower() for w in re.findall(r"[A-Za-z'\-]+", before)[-3:]]
        if any(w in COLORS or w.endswith('-colored') for w in prev):  # already has a color nearby ('brown tabby cat')
            continue
        a = re.search(r'\b([Aa])(n?)\s+$', before)
        if a:
            before = before[:a.start()] + a.group(1) + ('n' if phrase[:1].lower() in 'aeiou' else '') + ' '
        cap = before + phrase.lower() + ' ' + cap[m.start():]
        added += 1
    return cap, added


def lock(caps, mode='safe', do_fill=False):
    """Returns (new captions, number of color phrases changed or added)."""
    first, new, changed = {}, [], 0
    for cap in caps:
        edits = []
        sp = spans(cap, mode)
        multi = {h for h in {x[2] for x in sp} if len({frozenset(x[3]) for x in sp if x[2] == h}) > 1}  # two of them here
        for s, e, head, cols in sp:
            if head in multi:
                continue
            if head not in first:
                first[head] = (cols, cap[s:e])
            elif not (cols & first[head][0]):
                edits.append((s, e, first[head][1]))
        for s, e, text in sorted(edits, reverse=True):
            before = cap[:s]
            m = re.search(r'\b([Aa])(n?)\s+$', before)  # fix the article: 'a orange' -> 'an orange'
            if m:
                art = m.group(1) + ('n' if text[:1].lower() in 'aeiou' else '')
                before = before[:m.start()] + art + ' '
            cap = before + text.lower() + cap[e:]
            changed += 1
        if do_fill:
            cap, k = fill(cap, first)
            changed += k
        new.append(cap)
    return new, changed