"""
LoRA fine-tuning method for comparison experiments.

Trains a LoRA adapter on benchmark tasks using supervised fine-tuning (SFT)
on positive responses from contrastive pairs.

Optionally evaluates LoRA + steering by generating a steering vector on the
LoRA model and combining both methods.
"""

from __future__ import annotations

import gc
import json
import tempfile
from pathlib import Path
from typing import TYPE_CHECKING

import torch
from datasets import Dataset
from peft import LoraConfig, TaskType, get_peft_model
from trl import SFTTrainer, SFTConfig

from wisent.comparison.utils import (
    generate_contrastive_pairs,
    create_test_only_task,
    extract_accuracy,
    run_lm_eval_evaluation,
    run_ll_evaluation,
    load_model_and_tokenizer,
    apply_steering_to_model,
    remove_steering,
)
from wisent.core.utils.device import preferred_dtype

if TYPE_CHECKING:
    from wisent.core.models.wisent_model import WisentModel

__all__ = ["train_lora_adapter", "evaluate_lora", "evaluate_lora_with_steering", "apply_lora_to_model", "remove_lora"]


# Default LoRA configurations per model architecture
LORA_TARGET_MODULES = {
    "gemma": ["q_proj", "k_proj", "v_proj", "o_proj"],
    "llama": ["q_proj", "k_proj", "v_proj", "o_proj"],
    "mistral": ["q_proj", "k_proj", "v_proj", "o_proj"],
    "phi": ["q_proj", "k_proj", "v_proj", "dense"],
    "gpt_neo": ["q_proj", "v_proj"],
    "gpt2": ["c_attn"],
    "default": "all-linear",
}


def get_target_modules(model_name: str) -> str | list[str]:
    """Get LoRA target modules based on model architecture."""
    model_name_lower = model_name.lower()

    for arch, modules in LORA_TARGET_MODULES.items():
        if arch in model_name_lower:
            return modules

    return LORA_TARGET_MODULES["default"]


def prepare_sft_dataset(
    pairs: list[dict],
    tokenizer,
    max_length: int = 512,
) -> Dataset:
    """
    Prepare dataset for SFT from contrastive pairs.

    Uses only positive responses for training.

    Args:
        pairs: List of contrastive pairs
        tokenizer: Tokenizer for formatting
        max_length: Maximum sequence length

    Returns:
        HuggingFace Dataset ready for SFTTrainer
    """
    formatted_examples = []

    for pair in pairs:
        prompt = pair["prompt"]
        positive_response = pair["positive_response"]["model_response"]

        # Format as chat if tokenizer supports it, otherwise simple format
        if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template:
            messages = [
                {"role": "user", "content": prompt},
                {"role": "assistant", "content": positive_response},
            ]
            text = tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=False,
            )
        else:
            # Simple format for base models
            text = f"Q: {prompt}\nA: {positive_response}"

        formatted_examples.append({"text": text})

    return Dataset.from_list(formatted_examples)


def train_lora_adapter(
    task: str,
    model_name: str,
    output_path: str | Path,
    trait_label: str = "correctness",
    num_pairs: int = 50,
    device: str = "cuda:0",
    keep_intermediate: bool = False,
    # LoRA-specific parameters
    lora_r: int = 16,
    lora_alpha: int = 32,
    lora_dropout: float = 0.05,
    learning_rate: float = 2e-4,
    num_epochs: int = 3,
    batch_size: int = 2,
    max_length: int = 512,
) -> Path:
    """
    Train a LoRA adapter using SFT on positive responses.

    Args:
        task: lm-eval task name (e.g., 'boolq', 'cb')
        model_name: HuggingFace model name
        output_path: Where to save the LoRA adapter
        trait_label: Label for the trait being trained
        num_pairs: Number of training examples to use
        device: Device to train on
        keep_intermediate: Whether to keep intermediate files
        lora_r: LoRA rank
        lora_alpha: LoRA alpha scaling factor
        lora_dropout: LoRA dropout
        learning_rate: Training learning rate
        num_epochs: Number of training epochs
        batch_size: Training batch size
        max_length: Maximum sequence length

    Returns:
        Path to the saved LoRA adapter directory
    """
    import gc

    output_path = Path(output_path)

    # Step 1: Generate contrastive pairs
    print(f"Step 1: Generating training data from task: {task}")
    pairs, pairs_file = generate_contrastive_pairs(task, num_pairs)
    print(f"   Loaded {len(pairs)} training examples")

    # Step 2: Load model and tokenizer
    print(f"\nStep 2: Loading model {model_name}...")
    model, tokenizer = load_model_and_tokenizer(model_name, device, eval_mode=False)

    # Step 3: Configure LoRA
    print(f"\nStep 3: Configuring LoRA (r={lora_r}, alpha={lora_alpha})...")

    target_modules = get_target_modules(model_name)
    print(f"   Target modules: {target_modules}")

    lora_config = LoraConfig(
        r=lora_r,
        lora_alpha=lora_alpha,
        target_modules=target_modules,
        lora_dropout=lora_dropout,
        bias="none",
        task_type=TaskType.CAUSAL_LM,
    )

    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # Step 4: Prepare dataset
    print(f"\nStep 4: Preparing SFT dataset...")
    train_dataset = prepare_sft_dataset(pairs, tokenizer, max_length=max_length)
    print(f"   Dataset size: {len(train_dataset)} examples")

    # Step 5: Training
    print(f"\nStep 5: Training LoRA adapter...")

    # Create temporary directory for training outputs
    training_output_dir = tempfile.mkdtemp(prefix="lora_training_")

    # Use device-optimized dtype (bfloat16 on CUDA, float16 on MPS, float32 on CPU)
    dtype = preferred_dtype(device)

    training_args = SFTConfig(
        output_dir=training_output_dir,
        num_train_epochs=num_epochs,
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=1,
        learning_rate=learning_rate,
        weight_decay=0.01,
        warmup_ratio=0.1,
        logging_steps=10,
        save_strategy="no",  # Don't save checkpoints
        bf16=(dtype == torch.bfloat16),
        fp16=(dtype == torch.float16),
        report_to="none",  # Disable wandb/tensorboard
        dataset_text_field="text",  # Field containing the text to train on
    )

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        processing_class=tokenizer,
    )

    trainer.train()

    # Step 6: Save LoRA adapter
    print(f"\nStep 6: Saving LoRA adapter to {output_path}...")
    output_path.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(output_path)
    tokenizer.save_pretrained(output_path)

    # Save metadata
    metadata = {
        "method": "lora",
        "model": model_name,
        "task": task,
        "trait_label": trait_label,
        "num_pairs": len(pairs),
        "lora_config": {
            "r": lora_r,
            "alpha": lora_alpha,
            "dropout": lora_dropout,
            "target_modules": target_modules if isinstance(target_modules, list) else [target_modules],
        },
        "training_config": {
            "learning_rate": learning_rate,
            "num_epochs": num_epochs,
            "batch_size": batch_size,
            "max_length": max_length,
        },
    }

    with open(output_path / "metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)

    # Cleanup
    del model, trainer
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

    if not keep_intermediate:
        import os
        os.unlink(pairs_file)
        import shutil
        shutil.rmtree(training_output_dir, ignore_errors=True)

    print(f"\nLoRA adapter saved to {output_path}")
    return output_path


def apply_lora_to_model(wisent_model: "WisentModel", lora_path: str | Path) -> None:
    """
    Apply a trained LoRA adapter to a WisentModel.

    Args:
        wisent_model: WisentModel instance
        lora_path: Path to the saved LoRA adapter
    """
    from peft import PeftModel

    lora_path = Path(lora_path)

    # Check if model already has adapters
    if hasattr(wisent_model.hf_model, 'peft_config'):
        # Model already has PEFT, just load new adapter
        wisent_model.hf_model.load_adapter(str(lora_path), adapter_name="steering")
        wisent_model.hf_model.set_adapter("steering")
    else:
        # Wrap model with PEFT
        wisent_model.hf_model = PeftModel.from_pretrained(
            wisent_model.hf_model,
            str(lora_path),
            adapter_name="steering",
        )

    print(f"LoRA adapter loaded from {lora_path}")


def remove_lora(wisent_model: "WisentModel") -> None:
    """
    Remove/disable LoRA adapter from a WisentModel.

    Args:
        wisent_model: WisentModel instance with LoRA applied
    """
    if hasattr(wisent_model.hf_model, 'disable_adapters'):
        try:
            wisent_model.hf_model.disable_adapters()
            print("LoRA adapter disabled")
        except ValueError:
            # No adapter was loaded
            pass
    elif hasattr(wisent_model.hf_model, 'base_model'):
        # Unwrap the model
        wisent_model.hf_model = wisent_model.hf_model.base_model.model
        print("LoRA adapter removed")


def evaluate_lora(
    model_name: str,
    lora_path: str | Path,
    task: str,
    train_ratio: float = 0.8,
    device: str = "cuda:0",
    batch_size: int = 1,
    max_batch_size: int = 8,
    limit: int | None = None,
    output_dir: str | Path = None,
) -> dict:
    """
    Evaluate a trained LoRA adapter comparing base vs LoRA performance.

    Args:
        model_name: HuggingFace model name
        lora_path: Path to trained LoRA adapter
        task: lm-eval task name
        train_ratio: Train/test split ratio
        device: Device to run on
        batch_size: Batch size for evaluation
        max_batch_size: Max batch size
        limit: Limit number of eval examples
        output_dir: Where to save results

    Returns:
        Dict with evaluation results
    """
    import gc

    from wisent.core.models.wisent_model import WisentModel

    lora_path = Path(lora_path)

    # Create test task
    print(f"\n{'='*60}")
    print(f"Creating test task for: {task}")
    print(f"{'='*60}")

    task_dict = create_test_only_task(task, train_ratio=train_ratio)

    # Load model
    print(f"\n{'='*60}")
    print(f"Loading model: {model_name}")
    print(f"{'='*60}")
    wisent_model = WisentModel(model_name=model_name, device=device)

    # BASE evaluation
    print(f"\n{'='*60}")
    print(f"Running BASE evaluation (no LoRA)")
    print(f"{'='*60}")

    base_results = run_lm_eval_evaluation(wisent_model, task_dict, task, batch_size, max_batch_size, limit)
    base_acc_lm_eval = extract_accuracy(base_results, task)
    print(f"Base accuracy (lm-eval): {base_acc_lm_eval:.4f}")

    base_acc_ll = run_ll_evaluation(wisent_model, task_dict, task, limit)
    print(f"Base accuracy (LL): {base_acc_ll:.4f}")

    # Apply LoRA
    print(f"\n{'='*60}")
    print(f"Applying LoRA adapter from: {lora_path}")
    print(f"{'='*60}")
    apply_lora_to_model(wisent_model, lora_path)

    # LORA evaluation
    print(f"\n{'='*60}")
    print(f"Running LORA evaluation")
    print(f"{'='*60}")

    lora_results = run_lm_eval_evaluation(wisent_model, task_dict, task, batch_size, max_batch_size, limit)
    lora_acc_lm_eval = extract_accuracy(lora_results, task)
    print(f"LoRA accuracy (lm-eval): {lora_acc_lm_eval:.4f}")

    lora_acc_ll = run_ll_evaluation(wisent_model, task_dict, task, limit)
    print(f"LoRA accuracy (LL): {lora_acc_ll:.4f}")

    # Cleanup
    remove_lora(wisent_model)
    del wisent_model
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # Results
    results = {
        "task": task,
        "model": model_name,
        "lora_path": str(lora_path),
        "train_ratio": train_ratio,
        "base_accuracy_lm_eval": base_acc_lm_eval,
        "base_accuracy_ll": base_acc_ll,
        "lora_accuracy_lm_eval": lora_acc_lm_eval,
        "lora_accuracy_ll": lora_acc_ll,
        "difference_lm_eval": lora_acc_lm_eval - base_acc_lm_eval,
        "difference_ll": lora_acc_ll - base_acc_ll,
    }

    # Print summary
    print(f"\n{'='*60}")
    print(f"RESULTS SUMMARY")
    print(f"{'='*60}")
    print(f"Task: {task}")
    print(f"Model: {model_name}")
    print(f"LoRA: {lora_path}")
    print(f"{'-'*60}")
    print(f"{'Metric':<20} {'Base':<12} {'LoRA':<12} {'Diff':<12}")
    print(f"{'-'*60}")
    print(f"{'lm-eval accuracy':<20} {base_acc_lm_eval:<12.4f} {lora_acc_lm_eval:<12.4f} {lora_acc_lm_eval - base_acc_lm_eval:+.4f}")
    print(f"{'LL accuracy':<20} {base_acc_ll:<12.4f} {lora_acc_ll:<12.4f} {lora_acc_ll - base_acc_ll:+.4f}")
    print(f"{'='*60}")

    # Save results
    if output_dir:
        output_dir = Path(output_dir)
        # Add model name to path (sanitize "/" -> "_")
        model_dir_name = model_name.replace("/", "_")
        output_dir = output_dir / model_dir_name
        output_dir.mkdir(parents=True, exist_ok=True)
        results_file = output_dir / f"{task}_lora_eval_results.json"
        with open(results_file, "w") as f:
            json.dump(results, f, indent=2)
        print(f"\nResults saved to: {results_file}")

    return results


def evaluate_lora_with_steering(
    model_name: str,
    lora_path: str | Path,
    task: str,
    method: str = "caa",
    layers: str = "12",
    num_pairs: int = 50,
    steering_scales: list[float] | None = None,
    train_ratio: float = 0.8,
    device: str = "cuda:0",
    batch_size: int = 1,
    max_batch_size: int = 8,
    limit: int | None = None,
    output_dir: str | Path = None,
    extraction_strategy: str = "mc_balanced",
) -> dict:
    """
    Evaluate LoRA + steering.

    Generates steering vector ON the LoRA model, then evaluates at multiple scales.

    Args:
        model_name: HuggingFace model name
        lora_path: Path to trained LoRA adapter
        task: lm-eval task name
        method: Steering method (caa or fgaa)
        layers: Layers for steering vector
        num_pairs: Number of pairs for steering generation
        steering_scales: List of steering scales to evaluate
        train_ratio: Train/test split ratio
        device: Device to run on
        batch_size: Batch size for evaluation
        max_batch_size: Max batch size
        limit: Limit number of eval examples
        output_dir: Where to save results
        extraction_strategy: Strategy for activation extraction

    Returns:
        Dict with evaluation results at all scales
    """
    from wisent.core.models.wisent_model import WisentModel
    from wisent.core.trainers.steering_trainer import WisentSteeringTrainer
    from wisent.core.steering_methods import get_steering_method
    from wisent.core.activations.extraction_strategy import ExtractionStrategy
    from wisent.core.contrastive_pairs.core.set import ContrastivePairSet
    from wisent.core.contrastive_pairs.core.pair import ContrastivePair
    from wisent.core.contrastive_pairs.core.response import PositiveResponse, NegativeResponse

    if steering_scales is None:
        steering_scales = [1.0, 2.0, 4.0]

    lora_path = Path(lora_path)

    # Setup output directory
    if output_dir:
        output_dir = Path(output_dir)
        model_dir_name = model_name.replace("/", "_")
        output_dir = output_dir / model_dir_name
        output_dir.mkdir(parents=True, exist_ok=True)
        vectors_dir = output_dir / "steering_vectors"
        vectors_dir.mkdir(parents=True, exist_ok=True)
    else:
        output_dir = Path(".")
        vectors_dir = Path(".")

    # Step 1: Load model
    print(f"\n{'='*60}")
    print(f"Loading model: {model_name}")
    print(f"{'='*60}")
    wisent_model = WisentModel(model_name=model_name, device=device)

    # Step 2: Apply LoRA
    print(f"\n{'='*60}")
    print(f"Applying LoRA adapter from: {lora_path}")
    print(f"{'='*60}")
    apply_lora_to_model(wisent_model, lora_path)

    # Step 3: Generate contrastive pairs
    print(f"\n{'='*60}")
    print(f"Generating {num_pairs} contrastive pairs from {task}")
    print(f"{'='*60}")
    pairs_data, pairs_file = generate_contrastive_pairs(task, num_pairs)

    # Convert to ContrastivePairSet
    pairs = []
    for p in pairs_data:
        pair = ContrastivePair(
            prompt=p["prompt"],
            positive_response=PositiveResponse(model_response=p["positive_response"]["model_response"]),
            negative_response=NegativeResponse(model_response=p["negative_response"]["model_response"]),
        )
        pairs.append(pair)
    pair_set = ContrastivePairSet(pairs=pairs)
    print(f"Created {len(pair_set)} contrastive pairs")

    # Step 4: Generate steering vector on LoRA model
    print(f"\n{'='*60}")
    print(f"Generating {method.upper()} steering vector on LoRA model")
    print(f"Layers: {layers}")
    print(f"{'='*60}")

    # Get steering method
    steering_method = get_steering_method(method, device=device)

    # Get extraction strategy
    strategy = ExtractionStrategy(extraction_strategy)

    # Create trainer and run
    trainer = WisentSteeringTrainer(
        model=wisent_model,
        pair_set=pair_set,
        steering_method=steering_method,
    )

    result = trainer.run(
        layers_spec=layers,
        strategy=strategy,
        accept_low_quality_vector=True,
    )

    # Convert to our format
    steering_vectors = {}
    for layer_name, tensor in result.steered_vectors.to_dict().items():
        if tensor is not None:
            steering_vectors[layer_name] = tensor.cpu().float().tolist()

    steering_data = {
        "steering_vectors": steering_vectors,
        "layers": list(steering_vectors.keys()),
        "model": model_name,
        "method": method,
        "task": task,
        "num_pairs": len(pairs),
        "generated_on": "lora_model",
    }

    # Save steering vector
    vector_path = vectors_dir / f"{task}_{method}_on_lora_steering_vector.json"
    with open(vector_path, "w") as f:
        json.dump(steering_data, f, indent=2)
    print(f"Saved steering vector to {vector_path}")

    # Cleanup temp file
    import os
    os.unlink(pairs_file)

    # Create test task
    print(f"\n{'='*60}")
    print(f"Creating test task for: {task}")
    print(f"{'='*60}")
    task_dict = create_test_only_task(task, train_ratio=train_ratio)

    # Step 5: Evaluate at each scale
    all_results = {
        "task": task,
        "model": model_name,
        "method": method,
        "lora_path": str(lora_path),
        "steering_vector_path": str(vector_path),
        "steering_layers": list(steering_vectors.keys()),
        "num_pairs": num_pairs,
        "train_ratio": train_ratio,
        "generated_on": "lora_model",
        "scales": {},
    }

    for scale in steering_scales:
        print(f"\n{'='*60}")
        print(f"Evaluating LoRA+{method.upper()} at scale={scale}")
        print(f"{'='*60}")

        # Apply steering at this scale
        apply_steering_to_model(wisent_model, steering_data, scale=scale)

        # Evaluate
        lora_steer_results = run_lm_eval_evaluation(wisent_model, task_dict, task, batch_size, max_batch_size, limit)
        lora_steer_acc_lm_eval = extract_accuracy(lora_steer_results, task)
        print(f"LoRA+{method.upper()} accuracy (lm-eval): {lora_steer_acc_lm_eval:.4f}")

        lora_steer_acc_ll = run_ll_evaluation(wisent_model, task_dict, task, limit)
        print(f"LoRA+{method.upper()} accuracy (LL): {lora_steer_acc_ll:.4f}")

        # Remove steering before next scale
        remove_steering(wisent_model)

        all_results["scales"][scale] = {
            "lora_steer_accuracy_lm_eval": lora_steer_acc_lm_eval,
            "lora_steer_accuracy_ll": lora_steer_acc_ll,
        }

    # Cleanup
    remove_lora(wisent_model)
    del wisent_model
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # Print summary
    print(f"\n{'='*70}")
    print(f"LORA+STEERING RESULTS SUMMARY")
    print(f"{'='*70}")
    print(f"Task: {task}")
    print(f"Model: {model_name}")
    print(f"LoRA: {lora_path}")
    print(f"Steering: {method.upper()} (layers={list(steering_vectors.keys())})")
    print(f"Generated on: LoRA model")
    print(f"{'-'*70}")
    print(f"{'Scale':<10} {'lm-eval acc':<15} {'LL acc':<15}")
    print(f"{'-'*70}")
    for scale, res in all_results["scales"].items():
        print(f"{scale:<10} {res['lora_steer_accuracy_lm_eval']:<15.4f} {res['lora_steer_accuracy_ll']:<15.4f}")
    print(f"{'='*70}")

    # Save results
    if output_dir:
        results_file = output_dir / f"{task}_lora_{method}_steering_results.json"
        with open(results_file, "w") as f:
            json.dump(all_results, f, indent=2)
        print(f"\nResults saved to: {results_file}")

    return all_results


def main():
    import argparse

    parser = argparse.ArgumentParser(description="Train and evaluate LoRA adapter on benchmark task")
    parser.add_argument("--model", required=True, help="HuggingFace model name")
    parser.add_argument("--task", default="boolq", help="lm-eval task name")
    parser.add_argument("--output-dir", default="/home/ubuntu/output", help="Output directory")
    parser.add_argument("--num-pairs", type=int, default=50, help="Number of training examples")
    parser.add_argument("--device", default="cuda:0", help="Device")
    parser.add_argument("--lora-r", type=int, default=16, help="LoRA rank")
    parser.add_argument("--lora-alpha", type=int, default=32, help="LoRA alpha")
    parser.add_argument("--lora-dropout", type=float, default=0.05, help="LoRA dropout")
    parser.add_argument("--learning-rate", type=float, default=2e-4, help="Learning rate")
    parser.add_argument("--num-epochs", type=int, default=3, help="Number of epochs")
    parser.add_argument("--batch-size", type=int, default=2, help="Training batch size")
    parser.add_argument("--max-length", type=int, default=512, help="Max sequence length")
    parser.add_argument("--keep-intermediate", action="store_true", help="Keep intermediate files")
    # Eval args
    parser.add_argument("--train-ratio", type=float, default=0.8, help="Train/test split ratio")
    parser.add_argument("--eval-batch-size", default="auto", help="Eval batch size (int or 'auto')")
    parser.add_argument("--eval-max-batch-size", type=int, default=64, help="Max eval batch size for auto")
    parser.add_argument("--eval-limit", type=int, default=None, help="Limit eval examples")
    parser.add_argument("--skip-eval", action="store_true", help="Skip evaluation after training")
    # LoRA + Steering args
    parser.add_argument("--with-steering", action="store_true", help="Also evaluate LoRA + steering")
    parser.add_argument("--steering-method", default="caa", choices=["caa", "fgaa"], help="Steering method")
    parser.add_argument("--steering-layers", default="12", help="Layers for steering vector")
    parser.add_argument("--steering-num-pairs", type=int, default=50, help="Number of pairs for steering")
    parser.add_argument("--steering-scales", default="1.0,2.0,4.0", help="Comma-separated steering scales")
    parser.add_argument("--extraction-strategy", default="mc_balanced", help="Extraction strategy for steering")

    args = parser.parse_args()

    output_path = Path(args.output_dir) / f"{args.task}_lora_adapter"

    # Train
    train_lora_adapter(
        task=args.task,
        model_name=args.model,
        output_path=output_path,
        num_pairs=args.num_pairs,
        device=args.device,
        keep_intermediate=args.keep_intermediate,
        lora_r=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        learning_rate=args.learning_rate,
        num_epochs=args.num_epochs,
        batch_size=args.batch_size,
        max_length=args.max_length,
    )

    # Evaluate base vs LoRA
    if not args.skip_eval:
        # Parse eval batch size (can be "auto" or int)
        eval_batch_size = args.eval_batch_size
        if eval_batch_size != "auto":
            eval_batch_size = int(eval_batch_size)

        evaluate_lora(
            model_name=args.model,
            lora_path=output_path,
            task=args.task,
            train_ratio=args.train_ratio,
            device=args.device,
            batch_size=eval_batch_size,
            max_batch_size=args.eval_max_batch_size,
            limit=args.eval_limit,
            output_dir=args.output_dir,
        )

    # Evaluate LoRA + steering
    if args.with_steering:
        # Parse steering scales
        steering_scales = [float(s.strip()) for s in args.steering_scales.split(",")]

        # Parse eval batch size
        eval_batch_size = args.eval_batch_size
        if eval_batch_size != "auto":
            eval_batch_size = int(eval_batch_size)

        evaluate_lora_with_steering(
            model_name=args.model,
            lora_path=output_path,
            task=args.task,
            method=args.steering_method,
            layers=args.steering_layers,
            num_pairs=args.steering_num_pairs,
            steering_scales=steering_scales,
            train_ratio=args.train_ratio,
            device=args.device,
            batch_size=eval_batch_size,
            max_batch_size=args.eval_max_batch_size,
            limit=args.eval_limit,
            output_dir=args.output_dir,
            extraction_strategy=args.extraction_strategy,
        )


if __name__ == "__main__":
    main()
