"""Calibration data loading for sensitivity analysis."""

import torch
import numpy as np


def load_llm_calibration(
    tokenizer,
    n_samples: int = 128,
    seq_len: int = 512,
    seed: int = 42,
) -> list[torch.Tensor]:
    """Load calibration data for LLM sensitivity analysis.

    Uses WikiText-2 validation set — the standard calibration dataset
    used across quantization literature (GPTQ, AWQ, mlx-lm).
    Returns list of token tensors, each shape (seq_len,).
    """
    from datasets import load_dataset

    rng = np.random.RandomState(seed)

    ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="validation")
    texts = [t for t in ds["text"] if len(t.strip()) > 100]

    if not texts:
        raise RuntimeError(
            "WikiText-2 validation set returned no usable texts. "
            "Check your internet connection and datasets library installation."
        )

    # Tokenize and chunk into fixed-length sequences
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    all_tokens = []
    for text in texts:
        tokens = tokenizer.encode(text, add_special_tokens=False)
        all_tokens.extend(tokens)
        if len(all_tokens) >= n_samples * seq_len * 2:
            break

    if len(all_tokens) < seq_len:
        raise RuntimeError(
            f"Not enough tokens from WikiText-2 ({len(all_tokens)} < {seq_len}). "
            "Dataset may be corrupted."
        )

    # Chunk into sequences
    samples = []
    indices = list(range(0, len(all_tokens) - seq_len, seq_len))
    rng.shuffle(indices)
    for start in indices[:n_samples]:
        chunk = all_tokens[start:start + seq_len]
        samples.append(torch.tensor(chunk, dtype=torch.long))

    if len(samples) == 0:
        raise RuntimeError(
            f"Could not create any calibration sequences of length {seq_len} "
            f"from {len(all_tokens)} tokens."
        )

    print(f"  Loaded {len(samples)} calibration sequences (seq_len={seq_len}) from WikiText-2")
    return samples


def load_vlm_calibration(
    processor,
    n_samples: int = 4,
    seed: int = 42,
) -> list[dict]:
    """Load multimodal calibration data for VLM sensitivity analysis.

    Creates image+text pairs from COCO val2017 images with simple prompts.
    Returns list of processed input dicts ready for model forward pass.
    """
    from datasets import load_dataset
    from PIL import Image

    rng = np.random.RandomState(seed)

    print("  Loading COCO images for VLM calibration...")
    ds = load_dataset("detection-datasets/coco", split="val")

    total = len(ds)
    indices = rng.choice(total, size=min(n_samples * 3, total), replace=False)
    indices.sort()

    prompts = [
        "Describe what you see in this image.",
        "What objects are visible in this image?",
        "Provide a brief description of the image.",
    ]

    samples = []
    for i, idx in enumerate(indices):
        if len(samples) >= n_samples:
            break
        item = ds[int(idx)]
        img = item["image"]
        if not isinstance(img, Image.Image):
            img = Image.open(img)
        img = img.convert("RGB")

        prompt = prompts[i % len(prompts)]

        try:
            # Try chat template format first (works for most VLMs)
            if hasattr(processor, "apply_chat_template"):
                messages = [
                    {
                        "role": "user",
                        "content": [
                            {"type": "image"},
                            {"type": "text", "text": prompt},
                        ],
                    }
                ]
                text = processor.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True
                )
                inputs = processor(text=[text], images=[img], return_tensors="pt")
            else:
                inputs = processor(text=prompt, images=img, return_tensors="pt")

            samples.append(dict(inputs))
        except Exception as e:
            print(f"    Skipping image {idx}: {e}")
            continue

    if not samples:
        raise RuntimeError("Could not create any VLM calibration samples")

    print(f"  Created {len(samples)} multimodal calibration samples")
    return samples


def load_audio_calibration(
    model,
    n_samples: int = 4,
    seed: int = 42,
) -> list:
    """Load audio calibration data for Qwen3-ASR sensitivity analysis.

    Uses LibriSpeech clean validation set. Returns list of (args, kwargs)
    tuples ready for MLX model forward passes.

    Args:
        model: Qwen3-ASR MLX model (with _preprocess_audio and _build_prompt)
        n_samples: Number of calibration samples
        seed: Random seed for sample selection

    Returns:
        list of (args_tuple, kwargs_dict) for model forward passes
    """
    import mlx.core as mx
    from datasets import load_dataset

    rng = np.random.RandomState(seed)

    print("  Loading LibriSpeech validation for audio calibration...")
    ds = load_dataset("librispeech_asr", "clean", split="validation")

    total = len(ds)
    indices = rng.choice(total, size=min(n_samples * 3, total), replace=False)
    indices.sort()

    samples = []
    for idx in indices:
        if len(samples) >= n_samples:
            break
        item = ds[int(idx)]
        audio = item["audio"]["array"].astype(np.float32)

        try:
            # Preprocess audio through model's pipeline
            input_features, feature_attention_mask, num_audio_tokens = (
                model._preprocess_audio(audio)
            )
            # Build prompt with audio tokens
            input_ids = model._build_prompt(num_audio_tokens, language="English")

            # Create (args, kwargs) tuple for model forward pass
            samples.append((
                (input_ids,),  # args
                {
                    "input_features": input_features,
                    "feature_attention_mask": feature_attention_mask,
                },  # kwargs
            ))
        except Exception as e:
            print(f"    Skipping sample {idx}: {e}")
            continue

    if not samples:
        raise RuntimeError("Could not create any audio calibration samples")

    print(f"  Created {len(samples)} audio calibration samples from LibriSpeech")
    return samples
