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"

print(f"Loading {MODEL_NAME}...")
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
print(f"Model loaded with {NUM_LAYERS} layers")
print(f"Hidden size: {D_MODEL}")

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)
    "assistant": "You are a helpful AI assistant who provides clear, accurate, and reliable information while maintaining professional boundaries and ethical guidelines in all interactions.",
    # "consultant": "You are a professional consultant who provides expert strategic advice by analyzing complex business problems, identifying key issues, and recommending evidence-based solutions to help clients achieve their objectives.",
    "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.",
    # "editor": "You are an editor who refines and improves written content by ensuring clarity, coherence, and correctness while preserving the author's voice and intent.",
    # "coach": "You are a supportive coach who helps people achieve their goals by asking insightful questions, providing constructive feedback, and empowering individuals to develop their own solutions.",
    # "judge": "Act as a judge who follows established legal principles and precedents when making decisions.",
    # "psychologist": "Be a psychologist who applies theoretical knowledge to understand individual and group behavior through systematic assessment and analysis.",
    # "facilitator": "Act as a facilitator who expertly manages group dynamics and interpersonal tensions during discussions.",
    # "statistician": "You are a statistician who approaches every problem through the lens of data analysis and probabilistic reasoning.",
    # "researcher": "You are a researcher who systematically gathers, evaluates, and synthesizes information from multiple sources to build comprehensive understanding of specific topics.",
    # "teacher": "You are a patient teacher who explains concepts clearly by breaking down complex ideas into digestible parts, using examples and analogies, and adapting explanations to meet learners where they are.",
    #
    # 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.",
    "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.",
    "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.",
    # "hermit": "You are a hermit who has withdrawn from society to live in solitude, seeking wisdom in isolation and speaking only rarely, in cryptic phrases born from years of silent contemplation.",
    # "leviathan": "You are a leviathan, an ancient and vast creature of the deep whose thoughts move slowly across eons, speaking of primordial mysteries in a voice like the rumbling of ocean trenches.",
}

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

print(f"Defined {len(PERSONAS)} personas")

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?",
]

print(f"Defined {len(EVAL_QUESTIONS)} evaluation questions")

AUTORATER_MODEL = "anthropic/claude-3.5-haiku"

# Fill in the template below with placeholders {question}, {response}, {character}
JUDGE_PROMPT_TEMPLATE = """..."""


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
    """
    # Fill in the code here, to define prompt from the template `eval_prompt_template`
    prompt = ""

    completion = openrouter_client.chat.completions.create(
        model=AUTORATER_MODEL,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.0,
        max_tokens=500,
    )

    judge_response = completion.choices[0].message.content.strip()

    # Fill in the code here, to parse your response
    return 0


tests.test_judge_role_response(judge_role_response)

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


def generate_response_api(
    system_prompt: str,
    user_message: str,
    model: str = OPENROUTER_MODEL,
    max_tokens: int = 128,
    temperature: float = 0.7,
) -> str:
    """Generate a response using the OpenRouter API."""
    response = openrouter_client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message},
        ],
        max_tokens=max_tokens,
        temperature=temperature,
    )
    return response.choices[0].message.content


# Test the API
test_response = generate_response_api(
    system_prompt=PERSONAS["ghost"],
    user_message="What advice would you give to someone starting a new chapter in their life?",
)
print("Test response from 'ghost' persona:")
print(test_response)

def single_api_call(*args):
    try:
        time.sleep(0.1)  # useful for rate limiting
        # ...make api call, return (maybe processed) result
    except:
        # ...return error information

with ThreadPoolExecutor(max_workers=max_workers) as executor:
    # Submit all tasks
    futures = [executor.submit(single_api_call, task) for task in tasks]

    # Process completed tasks
    for future in as_completed(futures):
        key, response = future.result()
        responses[key] = 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:
# Simple test to verify the parallelization is working
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}: {response}\n")


# First, 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)
print(f"Generated {len(test_responses)} responses:")

# Show a sample of the results:
for k, v in test_responses.items():
    v_sanitized = v.strip().replace("\n", "<br>")
    display(HTML(f"<details><summary>{k}</summary>{v_sanitized}</details>"))

# 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()

    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,
)
print(f"Extracted activation shape: {test_activation.shape}")
print(f"Activation norm: {test_activation.norm().item():.2f}")

tests.test_extract_response_activations(extract_response_activations, model, tokenizer, D_MODEL, NUM_LAYERS)

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

    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_batched(
    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,
)
print(f"Extracted activation shape (batched): {test_activation.shape}")
print(f"Activation norm (batched): {test_activation.norm().item():.2f}")

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"

    raise NotImplementedError()
    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
print(f"\nExtracting from layer {EXTRACTION_LAYER}")

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)
    """
    raise NotImplementedError()


tests.test_compute_cosine_similarity_matrix_centered(compute_cosine_similarity_matrix_centered)

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
    """
    raise NotImplementedError()


# 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_traces(textposition="top center", marker=dict(size=12))
fig.update_yaxes(visible=False, range=[-0.5, 0.5])
fig.update_layout(height=350, showlegend=False)
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()