import torch as t
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForImageTextToText, AutoProcessor
import pandas as pd
from datasets import load_dataset
from tqdm import tqdm
from pathlib import Path

t.set_grad_enabled(False)

MAX_TOKENS = 500


def load_model_and_tokenizer(model_name: str, get_n_layers: bool = False):
    # check for mistral
    if "mistral" in model_name:
        model = AutoModelForImageTextToText.from_pretrained(
            model_name,
            torch_dtype=t.bfloat16,
            device_map="auto",
            trust_remote_code=True,
        )
        model.eval()

        processor = AutoProcessor.from_pretrained(model_name)
        processor.tokenizer.chat_template = processor.chat_template
        tokenizer = processor.tokenizer
    else:
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=t.bfloat16,
            device_map="auto",
            trust_remote_code=True,
        )
        model.eval()

        tokenizer = AutoTokenizer.from_pretrained(model_name)

    if get_n_layers:
        try: n_layers = model.config.num_hidden_layers
        except: n_layers = model.config.text_config.num_hidden_layers

    if get_n_layers:
        return model, tokenizer, n_layers
    else:
        return model, tokenizer
    
def main(
        model_name: str,
        output_dir: str = "data/interim/activations/alpaca",
        batch_size: int = 8,
        max_samples: int = 1000,
        probe_dir: Path = None
) -> None:
    # === LOAD DATA ===
    data = load_dataset("tatsu-lab/alpaca")
    data = data["train"].to_pandas().sample(n=max_samples, random_state=42)

    prompts = []
    for _, row in data.iterrows():
        user_prompt = row["instruction"]
        if row["input"] != "":
            user_prompt += "\n" + row["input"]
        
        messages = [
            {
                "role": "system",   
                "content": "You are a helpful AI assistant."
            },
            {
                "role": "user",
                "content": user_prompt
            }
        ]
        prompts.append(messages)

    # === LOAD MODEL AND TOKENIZER === 
    model, tokenizer, nlayers = load_model_and_tokenizer(model_name, True)
    layer = int(nlayers * 0.25)

    # === Load Probes ===
    sigma = t.load(probe_dir / "sigma.pt", map_location=model.device)
    mu = t.load(probe_dir / "mu.pt", map_location=model.device)
    supervised = t.load(probe_dir / "supervised.pt", map_location=model.device)

    ccs = t.load(probe_dir / "ccs.pt", map_location=model.device)
    crc_tpc = t.load(probe_dir / "crc-tpc.pt", map_location=model.device)
    diff_in_means = t.load(probe_dir / "diff-in-means.pt", map_location=model.device)

    probes = {
        "supervised": supervised,
        "ccs": ccs,
        "crc_tpc": crc_tpc,
        "diff_in_means": diff_in_means
    }

    # === BUILD PROMPTS ===
    prompts_chat = tokenizer.apply_chat_template(
        prompts,
        tokenize=False,
        add_generation_prompt=True
    )

    # Sanitize model name for all filenames
    safe_model_name = model_name.replace('/', '_')
    output_path = Path(output_dir) / safe_model_name / str(layer)
    output_path.mkdir(parents=True, exist_ok=True)
    
    # Create probe-specific directories
    for probe_name in probes.keys():
        (output_path / probe_name).mkdir(parents=True, exist_ok=True)
    
    # === CACHE ACTIVATIONS ===
    sample_idx = 0  # Global sample index for consistent file naming

    # Process in batches for better efficiency
    for i in tqdm(range(0, len(prompts_chat), batch_size), desc="Processing batches"):
        batch_prompts_chat = prompts_chat[i:i+batch_size]
        batch_prompts_original = prompts[i:i+batch_size]
        
        # Tokenize batch (use chat-formatted prompts)
        tokenized_batch_prompts = tokenizer(
            batch_prompts_chat,
            return_tensors="pt",
            add_special_tokens=False,
        ).to(model.device)
        
        with t.inference_mode():
            out = model.generate(
                **tokenized_batch_prompts,
                do_sample=True,
                output_hidden_states=True,
                return_dict_in_generate=True,
                max_new_tokens=MAX_TOKENS, 
            )
            generated_tokens = out.sequences[:, tokenized_batch_prompts.input_ids.shape[1]:]
            generated_texts = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)

            # Process each item in the batch
            for j, prompt in enumerate(batch_prompts_original):
                reached_eos = tokenizer.eos_token_id in generated_tokens[j]
                generated_text = generated_texts[j]
                
                if reached_eos:
                    print(f"Collect from prompt {j}")

                    # Collect full conversation
                    conversation = prompt.copy()
                    conversation.append({
                        "role": "assistant",
                        "content": generated_text.strip()
                    })
                    df = pd.DataFrame({'messages': [conversation], 'dishonest': [0]})
                    df.to_csv(output_path / f'alpaca_text_{sample_idx}.csv', index=False)

                    # Collect activations
                    hs_sample = t.cat(
                        [x[layer][j:j+1] for x in out["hidden_states"][1:]],
                        dim=1
                    ).squeeze(0)
                    t.save(hs_sample.cpu(), output_path / f'activations_{sample_idx}.pt')

                    # Collect scores
                    for probe_name, probe in probes.items():
                        if probe_name == "supervised":
                            hs = (hs_sample - mu) / sigma
                        else:
                            hs = hs_sample
                            
                        scores = hs @ probe.squeeze(0)
                        t.save(scores.cpu(), output_path / probe_name / f'scores_{sample_idx}.pt')

                    sample_idx += 1


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Generate activations for Alpaca dataset")
    parser.add_argument("--model", required=True, help="Model to use")
    parser.add_argument("--probe_dir", type=Path, required=True, help="Path to saved probe")
    parser.add_argument("--output_dir", default="activations", help="Directory to save activations to")
    parser.add_argument("--batch_size", type=int, default=8, help="Batch size for processing")
    parser.add_argument("--max_samples", type=int, default=1000, help="Maximum number of samples to process")
    args = parser.parse_args()

    main(args.model, args.output_dir, args.batch_size, args.max_samples, args.probe_dir)