generate_conjoint_tasks.py

"""
Generate 100 binary conjoint comparison tasks for hotel rooms in NYC.

Design:
  - Block 1 (tasks 1-5):   Rationality checks — identical quality, different prices.
                            A rational agent always picks the cheaper option.
  - Block 2 (tasks 6-20):  Star-price tradeoffs — sqft & bed are the same across
                            options, but the higher-star option costs more.
                            5 tasks each for star gaps of 1, 2, and 3.
                            Price premiums are geometrically spaced from small to
                            large so we can trace where the LLM switches.
  - Block 3 (tasks 21-100): Full random variation across all attributes.

Prices are drawn from a log-uniform distribution over [$20, $10,000] so that
most draws fall in the realistic range while still covering extremes.

After generation, a balance check (Mann-Whitney U on each attribute, A vs B)
is run. If any test rejects at alpha=0.05, the seed is incremented and the
design is regenerated until it passes.

Usage:
    python generate_conjoint_tasks.py
"""

import numpy as np
import pandas as pd
from validate_conjoint_balance import validate

# ── Attribute levels ─────────────────────────────────────────────────────────
STARS_LEVELS = [2, 3, 4, 5]
BED_LEVELS   = ['single', 'double', 'queen', 'king']
SQFT_MIN, SQFT_MAX, SQFT_STEP = 150, 800, 25   # sq-ft grid
PRICE_MIN, PRICE_MAX = 20, 10_000               # dollars per night

LOG_PRICE_MIN = np.log(PRICE_MIN)
LOG_PRICE_MAX = np.log(PRICE_MAX)

N_RATIONALITY  = 5
N_STAR_PRICE   = 15   # 5 per star-gap (1, 2, 3)
N_FULL_RANDOM  = 3580
N_TASKS = N_RATIONALITY + N_STAR_PRICE + N_FULL_RANDOM  # 3600

OUTPUT_FILE = 'conjoint_tasks.csv'
MAX_SEED_ATTEMPTS = 100
STARTING_SEED = 2024


def generate_tasks(seed):
    """Generate the full set of conjoint tasks for a given RNG seed."""
    rng = np.random.default_rng(seed)

    def rand_price():
        return int(round(np.exp(rng.uniform(LOG_PRICE_MIN, LOG_PRICE_MAX))))

    def rand_sqft():
        n_steps = (SQFT_MAX - SQFT_MIN) // SQFT_STEP
        return SQFT_MIN + rng.integers(0, n_steps + 1) * SQFT_STEP

    def rand_stars():
        return int(rng.choice(STARS_LEVELS))

    def rand_bed():
        return rng.choice(BED_LEVELS)

    tasks = []

    # ── Block 1: Rationality checks (5 tasks) ───────────────────────────
    for _ in range(N_RATIONALITY):
        stars = rand_stars()
        sqft  = rand_sqft()
        bed   = rand_bed()
        p_a   = rand_price()
        p_b   = rand_price()
        while p_a == p_b:
            p_b = rand_price()
        tasks.append(dict(
            task_id=len(tasks) + 1,
            a_stars=stars, a_sqft=sqft, a_bed=bed, a_price=p_a,
            b_stars=stars, b_sqft=sqft, b_bed=bed, b_price=p_b,
        ))

    # ── Block 2: Star-price tradeoffs (15 tasks) ────────────────────────
    for gap in [1, 2, 3]:
        premiums = np.exp(np.linspace(
            np.log(15 * gap),
            np.log(2500 * gap),
            5,
        )).astype(int)

        for premium in premiums:
            sqft = rand_sqft()
            bed  = rand_bed()

            eligible = [s for s in STARS_LEVELS if s + gap <= 5]
            low_star  = int(rng.choice(eligible))
            high_star = low_star + gap

            max_base = PRICE_MAX - premium
            low_price = int(round(np.exp(
                rng.uniform(LOG_PRICE_MIN, np.log(max(max_base, PRICE_MIN + 1)))
            )))
            high_price = low_price + int(premium)

            if rng.random() < 0.5:
                a_stars, a_price = high_star, high_price
                b_stars, b_price = low_star,  low_price
            else:
                a_stars, a_price = low_star,  low_price
                b_stars, b_price = high_star, high_price

            tasks.append(dict(
                task_id=len(tasks) + 1,
                a_stars=a_stars, a_sqft=sqft, a_bed=bed, a_price=a_price,
                b_stars=b_stars, b_sqft=sqft, b_bed=bed, b_price=b_price,
            ))

    # ── Block 3: Full random variation (80 tasks) ────────────────────────
    for _ in range(N_FULL_RANDOM):
        tasks.append(dict(
            task_id=len(tasks) + 1,
            a_stars=rand_stars(), a_sqft=rand_sqft(), a_bed=rand_bed(),
            a_price=rand_price(),
            b_stars=rand_stars(), b_sqft=rand_sqft(), b_bed=rand_bed(),
            b_price=rand_price(),
        ))

    return pd.DataFrame(tasks)


# ── Generate with balance validation ─────────────────────────────────────────
seed = STARTING_SEED
for attempt in range(1, MAX_SEED_ATTEMPTS + 1):
    df = generate_tasks(seed)
    df.to_csv(OUTPUT_FILE, index=False)

    balanced = validate(OUTPUT_FILE, alpha=0.05, verbose=False)
    if balanced:
        print(f"Seed {seed} passed balance check (attempt {attempt}).\n")
        break
    else:
        print(f"Seed {seed} FAILED balance check, retrying ...")
        seed += 1
else:
    print(f"WARNING: Could not find a balanced design in {MAX_SEED_ATTEMPTS} "
          f"attempts. Using last generated design (seed={seed}).")

# ── Final report ─────────────────────────────────────────────────────────────
validate(OUTPUT_FILE, alpha=0.05, verbose=True)

print(f"\nGenerated {len(df)} conjoint tasks -> {OUTPUT_FILE}  (seed={seed})")
print(f"\nPrice summary (across both options):")
all_prices = pd.concat([df['a_price'], df['b_price']])
print(all_prices.describe().to_string())
print(f"\nBlock 1 (rationality checks):   tasks 1-{N_RATIONALITY}")
print(f"Block 2 (star-price tradeoffs): tasks {N_RATIONALITY+1}-"
      f"{N_RATIONALITY+N_STAR_PRICE}")
print(f"Block 3 (full random):          tasks {N_RATIONALITY+N_STAR_PRICE+1}-"
      f"{N_TASKS}")

b2 = df.iloc[N_RATIONALITY:N_RATIONALITY+N_STAR_PRICE].copy()
b2['star_diff'] = abs(b2['a_stars'] - b2['b_stars'])
b2['price_diff'] = abs(b2['a_price'] - b2['b_price'])
print(f"\nStar-price tradeoff block (premium for higher stars):")
for _, r in b2.iterrows():
    print(f"  task {int(r['task_id']):>3d}: "
          f"star_gap={int(r['star_diff'])}, premium=${int(r['price_diff']):,}")