"""
Generates metadata CSVs for the DeepFashion retrieval project.

Before running:
1. Download the DeepFashion (In-shop Clothes Retrieval) high-res images
   and place them locally, preserving the original folder structure:
       <DATASET_PATH>/<gender>/<clothing_category>/<item_id>/<image>.jpg
2. Update DATASET_PATH below to point to your local copy.

Outputs (written to the current directory):
- original_metadata.csv          all images found
- full_metadata.csv               with segmentation images removed
- original_metadata_filtered.csv  with known bad/corrupted item_ids removed
"""

import os
import pandas as pd

# ---- CONFIG: update this to your local dataset path ----
DATASET_PATH = "/mnt/c/Users/User/Downloads/img_highres"

# item_ids excluded due to corrupted/mislabeled images found during data cleaning
EXCLUDED_ITEM_IDS = [
    "id_00003615", "id_00006951", "id_00004776", "id_00004850",
    "id_00007573", "id_00000773", "id_00006128", "id_00006574",
    "id_00001453", "id_00001995", "id_00002205", "id_00003020",
]


def build_metadata(dataset_path: str) -> pd.DataFrame:
    data = []
    for gender in os.listdir(dataset_path):
        gender_path = os.path.join(dataset_path, gender)
        if not os.path.isdir(gender_path):
            continue

        for clothing_category in os.listdir(gender_path):
            clothing_cat_path = os.path.join(gender_path, clothing_category)
            if not os.path.isdir(clothing_cat_path):
                continue

            for item_id in os.listdir(clothing_cat_path):
                item_id_path = os.path.join(clothing_cat_path, item_id)
                if not os.path.isdir(item_id_path):
                    continue

                for image in os.listdir(item_id_path):
                    image_path = os.path.join(item_id_path, image)
                    data.append({
                        "gender": gender,
                        "clothing_category": clothing_category,
                        "item_id": item_id,
                        "image_path": image_path,
                    })

    return pd.DataFrame(data)


def main():
    print(f"Scanning dataset at: {DATASET_PATH}")
    original = build_metadata(DATASET_PATH)
    original.to_csv("original_metadata.csv", index=False)
    print(f"original_metadata.csv written ({len(original)} rows)")

    full = original[~original.image_path.str.contains("segment")]
    full.to_csv("full_metadata.csv", index=False)
    print(f"full_metadata.csv written ({len(full)} rows)")

    filtered = full[~full.item_id.isin(EXCLUDED_ITEM_IDS)]
    filtered.to_csv("original_metadata_filtered.csv", index=False)
    print(f"original_metadata_filtered.csv written ({len(filtered)} rows)")


if __name__ == "__main__":
    main()