# Copyright (C) 2024 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

"""Compositional prompt ensemble for WinCLIP.

This module provides prompt templates and utilities for generating prompt ensembles
used by the WinCLIP model. The prompts are used to query CLIP about normal and
anomalous states of objects.

The module contains:
    - Lists of normal and anomalous state descriptors
    - Templates for constructing image description prompts
    - Functions to generate prompt ensembles by combining states and templates

Example:
    >>> from anomalib.models.image.winclip.prompting import create_prompt_ensemble
    >>> normal, anomalous = create_prompt_ensemble("transistor")  # doctest: +SKIP
    >>> print(normal[0])  # doctest: +SKIP
    'a photo of a transistor.'

See Also:
    - :class:`WinClip`: Main model class using these prompts
    - :class:`WinClipModel`: PyTorch model implementation
"""

NORMAL_STATES = [
    "{}",
    "flawless {}",
    "perfect {}",
    "unblemished {}",
    "{} without flaw",
    "{} without defect",
    "{} without damage",
]

ANOMALOUS_STATES = [
    "damaged {}",
    "{} with flaw",
    "{} with defect",
    "{} with damage",
]

TEMPLATES = [
    "a cropped photo of the {}.",
    "a close-up photo of a {}.",
    "a close-up photo of the {}.",
    "a bright photo of a {}.",
    "a bright photo of the {}.",
    "a dark photo of the {}.",
    "a dark photo of a {}.",
    "a jpeg corrupted photo of the {}.",
    "a jpeg corrupted photo of the {}.",
    "a blurry photo of the {}.",
    "a blurry photo of a {}.",
    "a photo of a {}.",
    "a photo of the {}.",
    "a photo of a small {}.",
    "a photo of the small {}.",
    "a photo of a large {}.",
    "a photo of the large {}.",
    "a photo of the {} for visual inspection.",
    "a photo of a {} for visual inspection.",
    "a photo of the {} for anomaly detection.",
    "a photo of a {} for anomaly detection.",
]


def create_prompt_ensemble(class_name: str = "object") -> tuple[list[str], list[str]]:
    """Create prompt ensemble for WinCLIP model.

    This function generates a comprehensive set of text prompts used by the WinCLIP model for
    zero-shot anomaly detection. It creates two sets of prompts:

    1. Normal prompts describing non-anomalous objects
    2. Anomalous prompts describing objects with defects

    The prompts are generated by combining predefined states (e.g., "flawless", "damaged")
    with templates (e.g., "a photo of a {}") for the given object class.

    Args:
        class_name (str, optional): Name of the object class to use in the prompts.
            Defaults to ``"object"``.

    Returns:
        tuple[list[str], list[str]]: A tuple containing:
            - List of normal prompts describing non-anomalous objects
            - List of anomalous prompts describing defective objects

    Example:
        Generate prompts for the "bottle" class:

        >>> normal_prompts, anomalous_prompts = create_prompt_ensemble("bottle")
        >>> print(normal_prompts[0])
        'a cropped photo of the bottle.'
        >>> print(anomalous_prompts[0])
        'a cropped photo of the damaged bottle.'

    See Also:
        - :data:`NORMAL_STATES`: Predefined states for normal objects
        - :data:`ANOMALOUS_STATES`: Predefined states for anomalous objects
        - :data:`TEMPLATES`: Predefined templates for prompt generation
    """
    normal_states = [state.format(class_name) for state in NORMAL_STATES]
    normal_ensemble = [template.format(state) for state in normal_states for template in TEMPLATES]

    anomalous_states = [state.format(class_name) for state in ANOMALOUS_STATES]
    anomalous_ensemble = [template.format(state) for state in anomalous_states for template in TEMPLATES]
    return normal_ensemble, anomalous_ensemble
