from pathlib import Path
import json

# ===================== EDIT THESE PATHS =====================

# Path to nerfstudio transforms.json
TRANSFORMS_PATH = Path('workspace/transforms.json')

# Path to COLMAP images.txt (text model, under workspace/sparse/0)
COLMAP_IMAGES_TXT = Path('workspace/sparse/0/images.txt')

# Text file listing the images you want as TRAIN
# One path per line. Can be:
#   - exactly as in images.txt (e.g. "dense/0/images/center/0.png")
#   - or any suffix of that (e.g. "center/0.png")
TRAIN_LIST_PATH = Path('train_list.txt')

# Text file listing the images you want as TEST/EVAL
# Same rules as above (exact or suffix)
TEST_LIST_PATH = Path('test_list.txt')

# ============================================================


def load_colmap_name_to_id(images_txt_path: Path) -> dict:
    """
    Read COLMAP's images.txt and build a mapping:
        image_name (str) -> image_id (int)

    Lines look like:
        IMAGE_ID QW QX QY QZ TX TY TZ CAMERA_ID NAME
    Comments start with '#'.
    """
    name_to_id = {}

    with images_txt_path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue

            parts = line.split()
            # First token should be an integer ID, last token is image name
            try:
                image_id = int(parts[0])
            except ValueError:
                # Not an image line
                continue

            if len(parts) < 9:
                # Unexpected format, skip
                continue

            image_name = parts[-1]
            name_to_id[image_name] = image_id

    return name_to_id


def load_specs(list_path: Path) -> list:
    """
    Load user-specified image identifiers from a text file.
    Each line is used as a string pattern.

    Examples of valid lines:
      dense/0/images/center/0.png
      center/0.png
      left/3.png
    """
    specs = []
    with list_path.open("r", encoding="utf-8") as f:
        for line in f:
            s = line.strip()
            if not s or s.startswith("#"):
                continue
            specs.append(s)
    return specs


def map_specs_to_ids(specs: list, name_to_id: dict, list_label: str) -> set:
    """
    For each spec string, try:
      1) exact match to a COLMAP image name
      2) if not found, match by suffix: image_name.endswith(spec)

    Returns a set of COLMAP image IDs.
    """
    ids = set()
    missing = []

    all_names = list(name_to_id.keys())

    for spec in specs:
        matched_ids = []

        # Exact match first
        if spec in name_to_id:
            matched_ids.append(name_to_id[spec])
        else:
            # Suffix match
            for name in all_names:
                if name.endswith(spec):
                    matched_ids.append(name_to_id[name])

        if len(matched_ids) == 0:
            missing.append(spec)
        else:
            for mid in matched_ids:
                ids.add(mid)

    if missing:
        print(f"Warning: the following entries in {list_label} "
              f"did not match any image name in images.txt:")
        for m in missing:
            print("  ", m)

    return ids


def main():
    # ----------- Basic checks -----------
    if not TRANSFORMS_PATH.is_file():
        print(f"Error: transforms.json not found at {TRANSFORMS_PATH}")
        return

    if not COLMAP_IMAGES_TXT.is_file():
        print(f"Error: COLMAP images.txt not found at {COLMAP_IMAGES_TXT}")
        return

    if not TRAIN_LIST_PATH.is_file():
        print(f"Error: train_list.txt not found at {TRAIN_LIST_PATH}")
        return

    if not TEST_LIST_PATH.is_file():
        print(f"Error: test_list.txt not found at {TEST_LIST_PATH}")
        return

    # ----------- Load COLMAP mapping -----------
    name_to_id = load_colmap_name_to_id(COLMAP_IMAGES_TXT)
    if not name_to_id:
        print(f"Error: no image entries found in {COLMAP_IMAGES_TXT}")
        return

    print(f"Loaded {len(name_to_id)} image entries from images.txt")

    # ----------- Load train and test specs -----------
    train_specs = load_specs(TRAIN_LIST_PATH)
    if not train_specs:
        print("Error: train_list.txt is empty or only comments.")
        return
    print(f"Loaded {len(train_specs)} train specs from train_list.txt")

    test_specs = load_specs(TEST_LIST_PATH)
    if not test_specs:
        print("Error: test_list.txt is empty or only comments.")
        return
    print(f"Loaded {len(test_specs)} test specs from test_list.txt")

    # ----------- Map specs to COLMAP IDs -----------
    train_ids = map_specs_to_ids(train_specs, name_to_id, "train_list.txt")
    if not train_ids:
        print("Error: none of the train specs could be matched to COLMAP image IDs.")
        return

    test_ids = map_specs_to_ids(test_specs, name_to_id, "test_list.txt")
    if not test_ids:
        print("Error: none of the test specs could be matched to COLMAP image IDs.")
        return

    # Handle overlap: if an image is in both train and test lists,
    # keep it in TRAIN and remove from TEST (and warn).
    overlap = train_ids & test_ids
    if overlap:
        print("Warning: some images are listed in BOTH train and test.")
        print("They will be kept ONLY in TRAIN and removed from TEST.")
        print("Overlapping COLMAP IDs:", sorted(overlap))
        test_ids -= train_ids

    print(f"Final COLMAP IDs for training: {len(train_ids)}")
    print(f"Final COLMAP IDs for testing : {len(test_ids)}")

    # ----------- Load transforms.json -----------
    transforms_path = TRANSFORMS_PATH.resolve()
    original_text = transforms_path.read_text(encoding="utf-8")
    meta = json.loads(original_text)

    if "frames" not in meta:
        print("Error: 'frames' key not found in transforms.json.")
        return

    frames = meta["frames"]

    # ----------- Split frames based on colmap_im_id -----------
    train_files = []
    test_files = []
    ignored_count = 0

    for frame in frames:
        colmap_id = frame.get("colmap_im_id")
        file_path = frame.get("file_path")

        if colmap_id is None:
            print("Error: a frame is missing 'colmap_im_id'. "
                  "This script assumes ns-process-data came from COLMAP.")
            return
        if file_path is None:
            print("Error: a frame is missing 'file_path'.")
            return

        if colmap_id in train_ids:
            train_files.append(file_path)
        elif colmap_id in test_ids:
            test_files.append(file_path)
        else:
            ignored_count += 1

    # Deduplicate and sort
    train_files = sorted(set(train_files))
    test_files = sorted(set(test_files))

    print(f"Frames in train split : {len(train_files)}")
    print(f"Frames in test split  : {len(test_files)}")
    print(f"Frames ignored (in neither list): {ignored_count}")

    if not train_files:
        print("Error: train split ended up empty. Aborting.")
        return
    if not test_files:
        print("Warning: test split ended up empty.")

    # ----------- Backup and write updated transforms.json -----------
    backup_path = transforms_path.with_suffix(".backup.json")
    if not backup_path.exists():
        backup_path.write_text(original_text, encoding="utf-8")
        print(f"Backup of original transforms.json written to {backup_path}")

    # Overwrite split fields
    meta["train_filenames"] = train_files
    meta["val_filenames"] = test_files   # use same set for val and test
    meta["test_filenames"] = test_files

    with transforms_path.open("w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2)

    print("Updated transforms.json with new train/test split.")
    print("Sample train filenames:")
    for fn in train_files[:5]:
        print("  ", fn)
    if len(train_files) > 5:
        print("  ...")

    print("Sample test filenames:")
    for fn in test_files[:5]:
        print("  ", fn)
    if len(test_files) > 5:
        print("  ...")


if __name__ == "__main__":
    main()
