from schemas_hf.wizard import (
    PlanKind,
    suggest_plan_for_warp_dataset,
)


def test_plan_dmt_brains_conditions_and_signals():
    plan = suggest_plan_for_warp_dataset("neuro/dmt-brains")
    assert plan.pattern == PlanKind.TABLE_AND_SIGNALS
    assert "uwunion/dmt-brains-conditions" in plan.hf_ids
    assert "uwunion/dmt-brains-signals" in plan.hf_ids


def test_plan_vctk_audio_and_speakers():
    plan = suggest_plan_for_warp_dataset("audio/vctk")
    assert plan.pattern == PlanKind.TABLE_AND_SIGNALS
    assert "uwunion/vctk-utterances" in plan.hf_ids
    assert "uwunion/vctk-speakers" in plan.hf_ids


def test_plan_cmu_mocap_table_and_raw():
    plan = suggest_plan_for_warp_dataset("mocap/cmu")
    assert plan.pattern == PlanKind.TABLE_AND_RAW
    assert "uwunion/cmu-mocap" in plan.hf_ids
    assert "uwunion/cmu-mocap-raw" in plan.hf_ids


def test_plan_celeba_attrs_image_dataset():
    plan = suggest_plan_for_warp_dataset("vision/celeba-attrs")
    assert plan.pattern == PlanKind.IMAGE_WITH_IMAGES
    assert plan.hf_ids == ["uwunion/celeba-attrs"]


def test_plan_coco_train_image_dataset():
    plan = suggest_plan_for_warp_dataset("vision/coco-train")
    assert plan.pattern == PlanKind.IMAGE_WITH_IMAGES
    assert plan.hf_ids == ["uwunion/coco-train"]


def test_plan_imagenet_images_image_dataset():
    plan = suggest_plan_for_warp_dataset("vision/imagenet-1k-images")
    assert plan.pattern == PlanKind.IMAGE_WITH_IMAGES
    assert plan.hf_ids == ["uwunion/imagenet-1k-images"]


def test_plan_sam_3d_body_metadata_only():
    plan = suggest_plan_for_warp_dataset("vision/sam-3d-body")
    assert plan.pattern == PlanKind.METADATA_ONLY
    assert plan.hf_ids == ["uwunion/sam-3d-body-metadata"]


def test_plan_laion_url_metadata_only():
    plan = suggest_plan_for_warp_dataset("vision/laion-10m")
    assert plan.pattern == PlanKind.PLAIN_TABLE
    assert plan.hf_ids == ["uwunion/laion-10m"]


def test_plan_finance_stocks_daily_plain():
    plan = suggest_plan_for_warp_dataset("stocks/stocks-daily")
    assert plan.pattern == PlanKind.PLAIN_TABLE
    assert plan.hf_ids == ["uwunion/stocks-daily"]


def test_plan_finance_hourly_ohlc_plain():
    plan = suggest_plan_for_warp_dataset("stocks/hourly-ohlc")
    assert plan.pattern == PlanKind.PLAIN_TABLE
    assert plan.hf_ids == ["uwunion/hourly-ohlc"]


def test_plan_finance_sec_filings_plain():
    plan = suggest_plan_for_warp_dataset("filings/sec-filings-text")
    assert plan.pattern == PlanKind.PLAIN_TABLE
    assert plan.hf_ids == ["uwunion/sec-filings-text"]


def test_plan_binance_plain():
    plan = suggest_plan_for_warp_dataset("crypto/binance/klines-1h")
    assert plan.pattern == PlanKind.PLAIN_TABLE
    # Uses the last segment as the HF id suffix.
    assert plan.hf_ids == ["uwunion/klines-1h"]


def test_plan_coingecko_plain():
    plan = suggest_plan_for_warp_dataset("crypto/coingecko/hourly")
    assert plan.pattern == PlanKind.PLAIN_TABLE
    assert plan.hf_ids == ["uwunion/hourly"]


def test_plan_mathx_plain():
    plan = suggest_plan_for_warp_dataset("math/mathx-100k")
    assert plan.pattern == PlanKind.PLAIN_TABLE
    # Replace underscore with dash in recommended id.
    assert plan.hf_ids == ["uwunion/mathx-100k"]


def test_plan_truthfulqa_plain():
    plan = suggest_plan_for_warp_dataset("nlp/truthfulqa")
    assert plan.pattern == PlanKind.PLAIN_TABLE
    assert plan.hf_ids == ["uwunion/truthfulqa"]
