import numpy as np
import json
import re
from typing import List, Dict, Any, Optional
from numba import njit
import time
import numpy as np
import cProfile
import pstats

import random


class Word:

    def __init__(self, word: str, hints: list[str], difficulty: float):
        self.word = word
        self.hints = hints
        self.difficulty = difficulty
    
    def copy_with_random_hint(self) -> "Word":
        if not self.hints:
            return Word(self.word, [], self.difficulty)
        hint = random.choice(self.hints)
        return Word(self.word, [hint], self.difficulty)


    def to_dict(self) -> dict[str, any]:
        return {
            "word": self.word,
            "hints": self.hints,
            "difficulty": self.difficulty
        }

    @staticmethod
    def from_dict(data: dict[str, any]) -> "Word":
        return Word(
            word=data["word"],
            hints=data["hints"],
            difficulty=data["difficulty"]
        )

class Dictionary:
    def __init__(self):
        self.words: List[Word] = []
        self._length_index: Dict[int, List[Word]] = {}
        self._max_word_length = 0
        self._pos_index_list = None  # Will be built on demand
        self._index_built = False

    def add_word(self, word: Word):
        self.words.append(word)
        l = len(word.word)
        if l not in self._length_index:
            self._length_index[l] = []
        self._length_index[l].append(word)
        if l > self._max_word_length:
            self._max_word_length = l
            self._index_built = False  # Invalidate index
        elif self._index_built:
            # just add the word to the existing index
            if self._pos_index_list is not None:
                for i, c in enumerate(word.word):
                    idx = ord(c) - ord('a')
                    self._pos_index_list[l][i][idx].add(word)

    def _build_pos_index_list(self):
        # pos_index_list[length][pos][char_index] = set(words)
        self._pos_index_list = [None] * (self._max_word_length + 1)
        for length in range(1, self._max_word_length + 1):
            pos_list = []
            for pos in range(length):
                char_list = [set() for _ in range(26)]  # a-z
                pos_list.append(char_list)
            self._pos_index_list[length] = pos_list
        for word in self.words:
            l = len(word.word)
            for i, c in enumerate(word.word):
                idx = ord(c) - ord('a')
                self._pos_index_list[l][i][idx].add(word)
        self._index_built = True

    def to_dict(self) -> dict[str, any]:
        return {
            "words": [w.to_dict() for w in self.words]
        }

    @staticmethod
    def from_dict(data: dict[str, any]) -> "Dictionary":
        d = Dictionary()
        for word_data in data.get("words", []):
            d.add_word(Word.from_dict(word_data))
        return d

    def to_json(self) -> str:
        import json
        return json.dumps(self.to_dict(), indent=2)

    @staticmethod
    def from_json(json_str: str) -> "Dictionary":
        import json
        data = json.loads(json_str)
        return Dictionary.from_dict(data)

    def find_by_pattern(self, pattern: str) -> List[Word]:
        length = len(pattern)
        if not self._index_built:
            self._build_pos_index_list()
        if length > self._max_word_length or self._pos_index_list[length] is None:
            return []
        fixed_positions = [(i, c) for i, c in enumerate(pattern) if c != '*']
        if not fixed_positions:
            return [w for w in self._length_index.get(length, [])]
        sets = []
        for i, c in fixed_positions:
            idx = ord(c) - ord('a')
            sets.append(self._pos_index_list[length][i][idx])
        if not sets:
            return []
        sets.sort(key=len)
        candidates = set.intersection(*sets)
        return candidates



def main():

    def generate_dummy_words(n, min_length=4, max_length=8):
        # Generates n random lowercase words of random length between min_length and max_length
        rng = np.random.default_rng(42)
        alphabet = np.array([ord(c) for c in "abcdefghijklmnopqrstuvwxyz"], dtype=np.uint8)
        words = []
        lengths = rng.integers(min_length, max_length + 1, size=n)
        for l in lengths:
            chars = rng.choice(alphabet, size=l)
            word = ''.join([chr(c) for c in chars])
            words.append(word)
        return words

    N = 100_000
    MIN_WORD_LENGTH = 4
    MAX_WORD_LENGTH = 8
    dummy_words = generate_dummy_words(N, MIN_WORD_LENGTH, MAX_WORD_LENGTH)

    # Fill Dictionary (position-indexed)
    dict_obj = Dictionary()
    for w in dummy_words:
        dict_obj.add_word(Word(w, hints=[], difficulty=1.0))

    # Prepare 1000 patterns (repeat the same 10 patterns 100 times)
    base_patterns = ["**a*e", "b****", "c***d*", "*e***", "a****z", "****", "*****", "******", "*******", "********"]
    patterns = base_patterns * 100  # 10 * 100 = 1000

    dict_obj._build_pos_index_list()  # Ensure the index is built before testing

    print(f"Testing {len(patterns)} patterns on {N} words with lengths {MIN_WORD_LENGTH}-{MAX_WORD_LENGTH}.")

    # --- Profiling block ---
    profiler = cProfile.Profile()
    profiler.enable()

    t0 = time.time()
    total_matches_dict = 0
    for pattern in patterns:
        result1 = dict_obj.find_by_pattern(pattern)
        total_matches_dict += len(result1)
    t1 = time.time()

    profiler.disable()
    print(f"Dictionary (position-indexed) total matches: {total_matches_dict} in {t1-t0:.6f} seconds for 1000 queries.")

    # Print top 20 time-consuming functions
    stats = pstats.Stats(profiler).sort_stats('cumtime')
    stats.print_stats(20)

if __name__ == "__main__":
    main()