import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# --- EXERCISE 1: La disparition (No 'e' or 'E') ---
class LaDisparition:
    """
    Generate text without ever using the letter 'e' or 'E'.
    
    You must use model() directly: model(input_ids) yields logits.
    You need to manually adjust the logits to forbid tokens containing 'e' or 'E'.
    
    REQUIREMENT: Do NOT use model.generate().
    
    Hints:
    - In __init__, pre-compute the set of forbidden token IDs by checking
      which tokens in the vocabulary decode to strings containing 'e' or 'E'.
    - In __call__, implement a token-by-token generation loop:
      1. Feed the current sequence to the model to get logits
      2. Mask out forbidden tokens (set their logits to -inf)
      3. Pick the next token (greedy: argmax, or use beam search for better quality)
      4. Append and repeat
    - Return only the generated text (not the prompt).
    """
    def __init__(self, model: AutoModelForCausalLM, tokenizer: AutoTokenizer):
        self.model = model
        self.tokenizer = tokenizer
        
        # TODO: Pre-calculate forbidden token IDs
        # Hint: also consider forbidding non-ASCII tokens that might hide the letter 'e'.
        # YOUR CODE HERE

    def __call__(self, prompt, max_tokens=20):
        # Tokenize the prompt using the chat template
        message = [{"role": "user", "content": prompt}]
        encoded = self.tokenizer.apply_chat_template(
            message, add_generation_prompt=True, return_tensors="pt"
        )
        input_ids = (encoded if isinstance(encoded, torch.Tensor) else encoded["input_ids"]).to(self.model.device)
        prompt_len = input_ids.shape[1]

        # TODO: Implement constrained generation loop
        # return only the generated text (after the prompt).
        
        # YOUR CODE HERE
        
        raise NotImplementedError("Implement constrained generation without 'e'")


# --- EXERCISE 2: The Toulouse Sequence ---
class ToulouseSequence:
    """
    Generate text without ever producing the word 'Toulouse'.
    
    You must use model() directly: model(input_ids) yields logits.
    
    REQUIREMENT: Do NOT use model.generate().
    
    This is harder than Exercise 1 because 'Toulouse' spans multiple tokens.
    """
    def __init__(self, model: AutoModelForCausalLM, tokenizer: AutoTokenizer):
        self.model = model
        self.tokenizer = tokenizer
        self.forbidden_word = "Toulouse"

    def __call__(self, prompt, max_tokens=20):
        # Tokenize the prompt using the chat template
        message = [{"role": "user", "content": prompt}]
        encoded = self.tokenizer.apply_chat_template(
            message, add_generation_prompt=True, return_tensors="pt"
        )
        inputs = (encoded if isinstance(encoded, torch.Tensor) else encoded["input_ids"]).to(self.model.device)
        prompt_length = inputs.shape[1]

        # TODO: Implement constrained generation loop
        # Return only the generated text (after the prompt).        
        # YOUR CODE HERE
        
        raise NotImplementedError("Implement constrained generation without 'Toulouse'")


if __name__ == "__main__":
    # NOTE: This block is for local testing only.
    # The evaluation server provides model and tokenizer.
    # You can use any small model for testing, e.g.:
    MODEL_NAME = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME, dtype=torch.float16, device_map="auto"
    )

    print("=== Exercise 1: La Disparition (no 'e') ===")
    ex1 = LaDisparition(model, tokenizer)
    result = ex1("Who is the king of the jungle?")
    print(f"Result: {result}")
    has_e = 'e' in result.lower()
    print(f"Contains 'e': {has_e} {'✗ FAIL' if has_e else '✓ PASS'}")

    print("\n=== Exercise 2: No Toulouse ===")
    ex2 = ToulouseSequence(model, tokenizer)
    result = ex2("Where are the headquarters of Airbus located?")
    print(f"Result: {result}")
    has_toulouse = 'toulouse' in result.lower()
    print(f"Contains 'Toulouse': {has_toulouse} {'✗ FAIL' if has_toulouse else '✓ PASS'}")