1️⃣ Mapping Persona Space

# You may need to log in to HuggingFace to access Gemma weights
# Get a token at https://huggingface.co/settings/tokens

HF_TOKEN = os.getenv("HF_TOKEN")
login(token=HF_TOKEN)

MODEL_NAME = "google/gemma-2-27b-it"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    dtype=dtype,
    device_map="auto",
    attn_implementation="eager",  # Required for Gemma 2 to access attention weights
)

NUM_LAYERS = model.config.num_hidden_layers
D_MODEL = model.config.hidden_size

PERSONAS = {
    # Default (neutral system prompts)
    "default": "",  # No system prompt
    "default_assistant": "You are an AI assistant",
    "default_llm": "You are a large language model",
    "default_helpful": "You are a helpful assistant",
    #
    # Assistant-like (professional, helpful)
    "consultant": "You are a professional consultant who provides expert advice and strategic recommendations, drawing on deep domain knowledge to help clients solve problems and achieve their goals.",
    "assistant": "You are a helpful AI assistant who provides clear, accurate, and reliable information while maintaining professional boundaries and ethical guidelines in all interactions.",
    "analyst": "You are a data analyst who examines information methodically, identifying patterns and trends through systematic investigation, and presenting findings in clear, actionable formats.",
    "evaluator": "You are an evaluator who assesses situations, proposals, and outcomes objectively using established criteria and frameworks to provide fair, balanced judgments.",
    "generalist": "You are a generalist with broad knowledge across many domains who can synthesize information from multiple disciplines to provide well-rounded perspectives on diverse topics.",
    #
    # Mid-range
    "storyteller": "You are a creative storyteller who weaves engaging narratives by crafting vivid characters, building tension through plot development, and drawing listeners into imaginative worlds.",
    "philosopher": "You are a philosopher who ponders deep questions about existence, knowledge, ethics, and reality, examining assumptions and exploring ideas through rigorous logical reasoning and contemplation.",
    "artist": "You are an artist who sees the world through a creative lens, finding beauty and meaning in unexpected places, and expressing unique perspectives through aesthetic sensibility and imaginative interpretation.",
    "rebel": "You are a rebel who challenges conventional thinking and questions established norms, pushing boundaries to explore alternative viewpoints and unconventional approaches.",
    "mystic": "You are a mystic who speaks of hidden truths and spiritual matters, perceiving deeper meanings beneath the surface of ordinary experience and exploring the mysteries of consciousness and existence.",
    #
    # Anti-Assistant (fantastical, un-helpful)
    "ghost": "You are a ghost, a spectral presence from beyond the veil who drifts between worlds, speaking in whispers of memories long past and truths that only the departed can know.",
    "bohemian": "You are a bohemian artist living outside conventional society, rejecting material pursuits and social expectations in favor of creative freedom, spontaneous expression, and unconventional experiences.",
    "bard": "You are a bard, a storyteller who employs poetic language, vivid imagery, and narrative structure, framing ideas through legend, history, and human drama while responding with lyrical eloquence and metaphorical depth.",
    "trickster": "You are a trickster who delights in mischief and riddles, speaking in paradoxes and wordplay, turning questions back on themselves, and finding humor in confusion and ambiguity.",
    "jester": "You are a jester who mocks and entertains in equal measure, using wit, satire, and absurdist humor to reveal uncomfortable truths while dancing along the edge of propriety and chaos.",
    "oracle": "You are an oracle who speaks in cryptic prophecies and riddles drawn from visions of possible futures, offering truth wrapped in metaphor and symbolism that must be interpreted to be understood.",
}

DEFAULT_PERSONAS = ["default", "default_assistant", "default_llm", "default_helpful"]

EVAL_QUESTIONS = [
    "What advice would you give to someone starting a new chapter in their life?",
    "How do you view the relationship between knowledge and wisdom?",
    "What do you think about the nature of truth?",
    "How should someone approach making difficult decisions?",
    "What role does creativity play in problem-solving?",
    "How do you see the balance between tradition and progress?",
    "What matters most when building trust with others?",
    "How do you think about the passage of time?",
    "What would you say to someone feeling lost or uncertain?",
    "How do you approach understanding something complex?",
    "What do you think about the nature of change?",
    "How should one deal with failure or setbacks?",
    "What role does intuition play in understanding?",
    "How do you view the relationship between the individual and society?",
    "What do you think makes something meaningful?",
]

OPENROUTER_MODEL = "google/gemma-2-27b-it"  # Matches our local model


def generate_responses_parallel(
    messages_list: list[list[dict[str, str]]],
    model: str = OPENROUTER_MODEL,
    max_tokens: int = 128,
    temperature: float = 0.7,
    max_workers: int = 10,
) -> list[str]:
    """
    Generate responses for multiple conversations in parallel using ThreadPoolExecutor.

    Args:
        messages_list: List of conversations, where each conversation is a list of
                       message dicts with "role" and "content" keys.
        model: Which model to use via OpenRouter.
        max_tokens: Maximum tokens per response.
        temperature: Sampling temperature.
        max_workers: Maximum number of parallel API calls.

    Returns:
        List of response strings, in the same order as messages_list.
    """

    def _single_call(messages: list[dict[str, str]]) -> str:
        try:
            time.sleep(0.1)  # Rate limiting
            response = openrouter_client.chat.completions.create(
                model=model,
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"API error: {e}")
            return ""

    if len(messages_list) == 1:
        return [_single_call(messages_list[0])]

    results: list[str | None] = [None] * len(messages_list)
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_idx = {executor.submit(_single_call, msgs): i for i, msgs in enumerate(messages_list)}
        for future in tqdm(as_completed(future_to_idx), total=len(messages_list), desc="API calls"):
            idx = future_to_idx[future]
            results[idx] = future.result()

    return results  # type: ignore


def generate_response(
    system_prompt: str,
    user_message: str,
    model: str = OPENROUTER_MODEL,
    max_tokens: int = 128,
    temperature: float = 0.7,
) -> str:
    """Generate a single response using the OpenRouter API (convenience wrapper)."""
    messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": user_message}]
    return generate_responses_parallel([messages], model=model, max_tokens=max_tokens, temperature=temperature)[0]


# Test the batch-call API
test_personas = ["ghost", "bard", "assistant"]
test_messages = [
    [
        {"role": "system", "content": PERSONAS[p]},
        {"role": "user", "content": "What is your name?"},
    ]
    for p in test_personas
]
test_responses = generate_responses_parallel(test_messages, max_tokens=40)
for p, resp in zip(test_personas, test_responses):
    print(f"{p}\n{resp[:80].strip()}...\n")

# GPT-4.1-mini is used as a fallback autorater for traits where Claude Haiku's content filters
# trigger refusals (e.g. "evil", "hallucinating"). Used in run_trait_pipeline (Section 4).
AUTORATER_MODEL = "anthropic/claude-3.5-haiku"
AUTORATER_MODEL_GPT = "openai/gpt-4.1-mini"

# Fill in the template below with placeholders {question}, {response}, {character}, and
# ask for the score in some format.
JUDGE_PROMPT_TEMPLATE = """..."""


def parse_score(response: str) -> int:
    """Parse score from judge response."""
    # YOUR CODE HERE - fill this function in (will depend on your judge template above)
    raise NotImplementedError()


def judge_role_response(
    question: str, response: str, character: str, eval_prompt_template: str = JUDGE_PROMPT_TEMPLATE
) -> int:
    """
    Use an LLM judge to score how well a response stays in character.

    Args:
        question: The question that was asked
        response: The model's response
        character: The character/role being played
        eval_prompt_template: Template for the judging prompt (must have {question}, {response}, {character} placeholders)

    Returns:
        Score from 0-3 indicating role-playing quality
    """
    prompt = eval_prompt_template.format(question=question, response=response, character=character)

    judge_response = generate_responses_parallel(
        [[{"role": "user", "content": prompt}]],
        model=AUTORATER_MODEL,
        temperature=0.0,
        max_tokens=500,
    )[0].strip()

    return parse_score(judge_response)


tests.test_judge_role_response(judge_role_response)

def generate_all_responses(
    personas: dict[str, str],
    questions: list[str],
    max_tokens: int = 256,
    max_workers: int = 10,
) -> dict[tuple[str, int], str]:
    """
    Generate responses for all persona-question combinations using parallel execution.

    Args:
        personas: Dict mapping persona name to system prompt
        questions: List of evaluation questions
        max_tokens: Maximum tokens per response
        max_workers: Maximum number of parallel workers

    Returns:
        Dict mapping (persona_name, question_idx) to response text
    """
    raise NotImplementedError()


# Demo of how this function works:
test_personas_demo = {
    "rhymer": "Reply in rhyming couplets.",
    "pirate": "Reply like a pirate.",
}
test_questions_demo = ["What is 2+2?", "What is the capital of France?"]

demo_responses = generate_all_responses(test_personas_demo, test_questions_demo, max_tokens=40)
for key, response in demo_responses.items():
    print(f"{key}:\n{response.strip()}\n")

# Now, a quick test of the function using just 2 personas & questions:
test_personas = {k: PERSONAS[k] for k in list(PERSONAS.keys())[-2:]}
test_questions = EVAL_QUESTIONS[:2]

test_responses = generate_all_responses(test_personas, test_questions)

# Show a sample of the results:
for (persona, question_idx), response in test_responses.items():
    question = test_questions[question_idx]
    print(f"{persona=}\n{question=}\n{response.strip()[:120]}...\n")

# Once you've confirmed these work, run them all!
responses = generate_all_responses(PERSONAS, EVAL_QUESTIONS)

def _normalize_messages(messages: list[dict[str, str]]) -> list[dict[str, str]]:
    """Merge any leading system message into the first user message.

    Gemma 2's chat template raises an error for the "system" role. The standard
    workaround is to prepend the system content to the first user message.
    """
    if not messages or messages[0]["role"] != "system":
        return messages
    system_content = messages[0]["content"]
    rest = list(messages[1:])
    if rest and rest[0]["role"] == "user" and system_content:
        rest[0] = {"role": "user", "content": f"{system_content}\n\n{rest[0]['content']}"}
    return rest


def format_messages(messages: list[dict[str, str]], tokenizer) -> tuple[str, int]:
    """Format a conversation for the model using its chat template.

    Args:
        messages: List of message dicts with "role" and "content" keys.
                 Can include "system", "user", and "assistant" roles.
                 Any leading system message is merged into the first user message
                 (required for Gemma 2, which does not support the system role).
        tokenizer: The tokenizer with chat template support

    Returns:
        full_prompt: The full formatted prompt as a string
        response_start_idx: The index of the first token in the last assistant message
    """
    messages = _normalize_messages(messages)

    # Apply chat template to get full conversation
    full_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)

    # Get prompt without final assistant message to compute response_start_idx
    prompt_without_response = tokenizer.apply_chat_template(
        messages[:-1], tokenize=False, add_generation_prompt=True
    ).rstrip()

    # Get start idx for response (+1 to skip newline)
    response_start_idx = tokenizer(prompt_without_response, return_tensors="pt").input_ids.shape[1] + 1

    return full_prompt, response_start_idx


tests.test_format_messages_response_index(format_messages, tokenizer)

def extract_response_activations(
    model,
    tokenizer,
    system_prompts: list[str],
    questions: list[str],
    responses: list[str],
    layer: int,
) -> Float[Tensor, "num_examples d_model"]:
    """
    Extract mean activation over response tokens at a specific layer.

    Returns:
        Batch of mean activation vectors of shape (num_examples, hidden_size)
    """
    assert len(system_prompts) == len(questions) == len(responses)

    raise NotImplementedError()


test_activation = extract_response_activations(
    model=model,
    tokenizer=tokenizer,
    system_prompts=[PERSONAS["assistant"]],
    questions=EVAL_QUESTIONS[:1],
    responses=["I would suggest taking time to reflect on your goals and values."],
    layer=NUM_LAYERS // 2,
)
tests.test_extract_response_activations(extract_response_activations, model, tokenizer, D_MODEL, NUM_LAYERS)

def extract_persona_vectors(
    model,
    tokenizer,
    personas: dict[str, str],
    questions: list[str],
    responses: dict[tuple[str, int], str],
    layer: int,
    scores: dict[tuple[str, int], int] | None = None,
    score_threshold: int = 3,
) -> dict[str, Float[Tensor, " d_model"]]:
    """
    Extract mean activation vector for each persona.

    Args:
        model: The language model
        tokenizer: The tokenizer
        personas: Dict mapping persona name to system prompt
        questions: List of evaluation questions
        responses: Dict mapping (persona, q_idx) to response text
        layer: Which layer to extract activations from
        scores: Optional dict mapping (persona, q_idx) to judge score (0-3)
        score_threshold: Minimum score required to include response (default 3)

    Returns:
        Dict mapping persona name to mean activation vector
    """
    assert questions and personas and responses, "Invalid inputs"

    persona_vectors = {}

    for counter, (persona_name, system_prompt) in enumerate(personas.items()):
        print(f"Running persona ({counter + 1}/{len(personas)}) {persona_name} ...", end="")

        # YOUR CODE HERE - compute activations for each persona, take the mean
        mean_acts = ...

        # Take mean across all responses for this persona
        persona_vectors[persona_name] = mean_acts
        print("finished!")

        # Clear GPU cache between personas to avoid OOM errors
        t.cuda.empty_cache()

    return persona_vectors

# # Score all responses using the judge
# print("Scoring responses with LLM judge...")
# scores: dict[tuple[str, int], int] = {}

# for (persona_name, q_idx), response in tqdm(responses.items()):
#     if response:  # Skip empty responses
#         score = judge_role_response(
#             question=EVAL_QUESTIONS[q_idx],
#             response=response,
#             character=persona_name,
#         )
#         scores[(persona_name, q_idx)] = score
#         time.sleep(0.1)  # Rate limiting

# # Print filtering statistics per persona
# print("\nFiltering statistics (score >= 3 required):")
# for persona_name in PERSONAS.keys():
#     persona_scores = [scores.get((persona_name, q_idx), 0) for q_idx in range(len(EVAL_QUESTIONS))]
#     n_passed = sum(1 for s in persona_scores if s >= 3)
#     n_total = len(persona_scores)
#     print(f"  {persona_name}: {n_passed}/{n_total} passed ({n_passed / n_total:.0%})")

# Extract vectors (using the test subset from before)
EXTRACTION_LAYER = round(NUM_LAYERS * 0.65)  # 65% through the model

persona_vectors = extract_persona_vectors(
    model=model,
    tokenizer=tokenizer,
    personas=PERSONAS,
    questions=EVAL_QUESTIONS,
    responses=responses,
    layer=EXTRACTION_LAYER,
)

print(f"\nExtracted vectors for {len(persona_vectors)} personas")
for name, vec in persona_vectors.items():
    print(f"  {name}: norm = {vec.norm().item():.2f}")

tests.test_extract_persona_vectors(extract_persona_vectors, model, tokenizer, D_MODEL, NUM_LAYERS)

def compute_cosine_similarity_matrix(
    persona_vectors: dict[str, Float[Tensor, " d_model"]],
) -> tuple[Float[Tensor, "n_personas n_personas"], list[str]]:
    """
    Compute pairwise cosine similarity between persona vectors.

    Returns:
        Tuple of (similarity matrix, list of persona names in order)
    """
    raise NotImplementedError()


tests.test_compute_cosine_similarity_matrix(compute_cosine_similarity_matrix)

cos_sim_matrix, persona_names = compute_cosine_similarity_matrix(persona_vectors)

px.imshow(
    cos_sim_matrix.float(),
    x=persona_names,
    y=persona_names,
    title="Persona Cosine Similarity Matrix (Uncentered)",
    color_continuous_scale="RdBu",
    color_continuous_midpoint=0.0,
).show()

def compute_cosine_similarity_matrix_centered(
    persona_vectors: dict[str, Float[Tensor, " d_model"]],
) -> tuple[Float[Tensor, "n_personas n_personas"], list[str]]:
    """
    Compute pairwise cosine similarity between centered persona vectors.

    Returns:
        Tuple of (similarity matrix, list of persona names in order)
    """
    names = list(persona_vectors.keys())

    # Stack vectors into matrix and center by subtracting mean
    vectors = t.stack([persona_vectors[name] for name in names])
    vectors = vectors - vectors.mean(dim=0)

    # Normalize
    vectors_norm = vectors / vectors.norm(dim=1, keepdim=True)

    # Compute cosine similarity
    cos_sim = vectors_norm @ vectors_norm.T

    return cos_sim, names


cos_sim_matrix_centered, persona_names = compute_cosine_similarity_matrix_centered(persona_vectors)

px.imshow(
    cos_sim_matrix_centered.float(),
    x=persona_names,
    y=persona_names,
    title="Persona Cosine Similarity Matrix (Centered)",
    color_continuous_scale="RdBu",
    color_continuous_midpoint=0.0,
).show()

def pca_decompose_persona_vectors(
    persona_vectors: dict[str, Float[Tensor, " d_model"]],
    default_personas: list[str] = DEFAULT_PERSONAS,
) -> tuple[Float[Tensor, " d_model"], np.ndarray, PCA]:
    """
    Analyze persona space structure.

    Args:
        persona_vectors: Dict mapping persona name to vector
        default_personas: List of persona names considered "default" (neutral assistant behavior)

    Returns:
        Tuple of:
        - assistant_axis: Normalized direction from role-playing toward default/assistant behavior
        - pca_coords: 2D PCA coordinates for each persona (n_personas, 2)
        - pca: Fitted PCA object, via the method `PCA.fit_transform`
    """
    raise NotImplementedError()


tests.test_pca_decompose_persona_vectors(pca_decompose_persona_vectors)

# Compute mean vector to handle constant vector problem (same as in centered cosine similarity)
# This will be subtracted from activations before projection to center around zero
persona_vectors = {k: v.to(device, dtype=dtype) for k, v in persona_vectors.items()}
mean_vector = t.stack(list(persona_vectors.values())).mean(dim=0)
persona_vectors_centered = {k: v - mean_vector for k, v in persona_vectors.items()}

# Perform PCA decomposition on space (PCA uses numpy internally, so convert to cpu float32)
assistant_axis, pca_coords, pca = pca_decompose_persona_vectors(
    {k: v.cpu().float() for k, v in persona_vectors_centered.items()}
)
assistant_axis = assistant_axis.to(device, dtype=dtype)  # Set to model dtype upfront

print(f"Assistant Axis norm: {assistant_axis.norm().item():.4f}")
print(
    f"PCA explained variance: PC1={pca.explained_variance_ratio_[0]:.1%}, PC2={pca.explained_variance_ratio_[1]:.1%}"
)

# Compute projection onto assistant axis for coloring
vectors = t.stack([persona_vectors_centered[name] for name in persona_names]).to(device, dtype=dtype)
# Normalize vectors before projecting (so projections are in [-1, 1] range)
vectors_normalized = vectors / vectors.norm(dim=1, keepdim=True)
projections = (vectors_normalized @ assistant_axis).float().cpu().numpy()

# 2D scatter plot
fig = px.scatter(
    x=pca_coords[:, 0],
    y=pca_coords[:, 1],
    text=persona_names,
    color=projections,
    color_continuous_scale="RdBu",
    title="Persona Space (PCA) colored by Assistant Axis projection",
    labels={
        "x": f"PC1 ({pca.explained_variance_ratio_[0]:.1%})",
        "y": f"PC2 ({pca.explained_variance_ratio_[1]:.1%})",
        "color": "Assistant Axis",
    },
)
fig.update_traces(textposition="top center", marker=dict(size=10))
fig.show()

def characterize_axis(
    persona_vectors: dict[str, Float[Tensor, " d_model"]],
    assistant_axis: Float[Tensor, " d_model"],
) -> dict[str, float]:
    """
    Compute cosine similarity of each persona vector with the assistant axis.

    Args:
        persona_vectors: Dict mapping persona name to its (centered) activation vector
        assistant_axis: Normalized Assistant Axis direction vector

    Returns:
        Dict mapping persona name to cosine similarity score, sorted by score
    """
    similarities = {}
    for name, vec in persona_vectors.items():
        cos_sim = (vec @ assistant_axis) / (vec.norm() * assistant_axis.norm() + 1e-8)
        similarities[name] = cos_sim.item()
    return dict(sorted(similarities.items(), key=lambda x: x[1]))


# Compute trait similarities using centered persona vectors
trait_similarities = characterize_axis(persona_vectors_centered, assistant_axis)

# Print extremes
items = list(trait_similarities.items())
print("Most ROLE-PLAYING (anti-assistant):")
for name, sim in items[:5]:
    print(f"  {name}: {sim:.3f}")
print("\nMost ASSISTANT-LIKE:")
for name, sim in items[-5:]:
    print(f"  {name}: {sim:.3f}")

# Create 1D visualization
names = [name for name, _ in items]
sims = [sim for _, sim in items]

fig = px.scatter(
    x=sims,
    y=[0] * len(sims),
    text=names,
    color=sims,
    color_continuous_scale="RdBu",
    title="Persona Projections onto the Assistant Axis",
    labels={"x": "Cosine Similarity with Assistant Axis", "color": "Similarity"},
)
fig.update_yaxes(visible=False, range=[-0.5, 0.5])
fig.update_layout(height=350, showlegend=False)
fig.update_traces(
    textposition=["top center" if i % 2 == 0 else "bottom center" for i in range(len(names))], marker=dict(size=12)
)
fig.show()

REPO_ID = "lu-christina/assistant-axis-vectors"
GEMMA2_MODEL = "gemma-2-27b"
GEMMA2_TARGET_LAYER = 22  # layer used in the paper's config

# Load the Gemma 2 27B assistant axis (shape [46, 4608] - 46 layers, d_model=4608)
hf_axis_path = hf_hub_download(repo_id=REPO_ID, filename=f"{GEMMA2_MODEL}/assistant_axis.pt", repo_type="dataset")
hf_axis_raw = t.load(hf_axis_path, map_location="cpu", weights_only=False)
hf_axis_vec = F.normalize(hf_axis_raw[GEMMA2_TARGET_LAYER].float(), dim=0)  # shape: (4608,)
print(f"Gemma 2 27B axis shape at layer {GEMMA2_TARGET_LAYER}: {hf_axis_vec.shape}")

# Download all 240 pre-computed trait vectors (each has shape [n_layers, d_model])
print("Downloading 240 trait vectors (this may take a moment)...")
local_dir = snapshot_download(
    repo_id=REPO_ID, repo_type="dataset", allow_patterns=f"{GEMMA2_MODEL}/trait_vectors/*.pt"
)
trait_vectors_hf = {
    p.stem: t.load(p, map_location="cpu", weights_only=False)
    for p in Path(local_dir, GEMMA2_MODEL, "trait_vectors").glob("*.pt")
}
print(f"Loaded {len(trait_vectors_hf)} trait vectors")

# Cosine similarity between each trait vector (at the target layer) and the assistant axis
trait_sims_hf = {
    name: F.cosine_similarity(vec[GEMMA2_TARGET_LAYER].float(), hf_axis_vec, dim=0).item()
    for name, vec in trait_vectors_hf.items()
}

sim_names = list(trait_sims_hf.keys())
sim_values = np.array([trait_sims_hf[n] for n in sim_names])
fig = utils.plot_similarity_line(sim_values, sim_names, n_extremes=5)
plt.title(f"Trait Vectors vs Assistant Axis (Gemma 2 27B, Layer {GEMMA2_TARGET_LAYER})")
plt.tight_layout()
plt.show()