import re
from collections import Counter


def extract(text):
    text = text.split('</think>')[-1].strip()
    index = text.rfind('\\boxed')
    if index < 0:
        index = text.rfind('\\fbox')
    if index < 0:
        return None
    key = '\\boxed' if text.startswith('\\boxed', index) else '\\fbox'
    rest = text[index + len(key):]
    stripped = rest.lstrip()
    if stripped.startswith('{'):
        depth = 0
        for i, char in enumerate(stripped):
            depth += (char == '{') - (char == '}')
            if depth == 0:
                return stripped[1:i].strip() or None
    elif key == '\\boxed' and rest.startswith(' '):
        return rest.strip().split('$')[0].split('\n')[0].strip() or None
    return None


def correct(prediction, gold):
    if prediction is None:
        return False
    from math_verify import parse, verify
    reference = parse('\\boxed{' + gold + '}', parsing_timeout=5)
    candidate = parse('\\boxed{' + prediction + '}', parsing_timeout=5)
    return bool(reference and candidate and verify(reference, candidate, timeout_seconds=5))


def repetition(text):
    words = text.split()[-1000:]
    grams = [tuple(words[i:i + 8]) for i in range(max(0, len(words) - 7))]
    duplicate = 1 - len(set(grams)) / max(1, len(grams))
    tail = words[-200:]
    four = [tuple(tail[i:i + 4]) for i in range(max(0, len(tail) - 3))]
    counts = Counter(four)
    repetitive = bool(four) and sum(counts[g] >= 3 for g in four) / len(four) >= .5
    return {'duplicate_8gram_tail': duplicate, 'repetitive_end': repetitive,
            'whitespace_tail': bool(text) and not text[-1000:].strip(),
            'thinking_closed': '</think>' in text}


def grade_ifeval(row, text):
    import nltk
    from common import ROOT
    from langdetect import DetectorFactory
    from instruction_following_eval import instructions_registry as registry
    nltk.data.path[:] = [str(ROOT / 'data/nltk')]
    DetectorFactory.seed = 0
    response = text.split('</think>')[-1].strip() if '</think>' in text else ''
    flags = []
    for instruction_id, kwargs in zip(row['instruction_id_list'], row['kwargs']):
        checker = registry.INSTRUCTION_DICT[instruction_id](instruction_id)
        checker.build_description(**{k: v for k, v in (kwargs or {}).items() if v is not None})
        args = checker.get_instruction_args()
        if args and 'prompt' in args:
            checker.build_description(prompt=row['prompt'])
        flags.append(bool(response) and bool(checker.check_following(response)))
    return all(flags)