#!/usr/bin/env python3
"""
Script to filter label_data.json based on filepaths in dbscan_filtered_results.json

This script will keep only the samples from label_data.json where the 'image' field
matches with the filepath names found in the 'results' section of dbscan_filtered_results.json
"""

import json
import os
from pathlib import Path
from typing import List, Dict, Set


def extract_image_names_from_filepath(filepath: str) -> str:
    """
    Extract the base image name from a filepath for comparison.
    
    Args:
        filepath: Full filepath from dbscan results
        
    Returns:
        Base filename without extension and path
    """
    return Path(filepath).stem


def load_dbscan_filtered_results(filepath: str) -> Set[str]:
    """
    Load dbscan filtered results and extract image names.
    
    Args:
        filepath: Path to dbscan_filtered_results.json
        
    Returns:
        Set of image names extracted from filepaths
    """
    with open(filepath, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    image_names = set()
    
    # Extract image names from all results
    for result in data.get('results', []):
        filepath_str = result.get('filepath', '')
        if filepath_str:
            # Extract filename without extension
            image_name = extract_image_names_from_filepath(filepath_str)
            image_names.add(image_name)
            
            # Also try to match different variations
            # Some images might have different naming patterns
            filename = Path(filepath_str).name
            # Remove extension and add to set
            filename_without_ext = Path(filename).stem
            image_names.add(filename_without_ext)
    
    return image_names


def load_label_data(filepath: str) -> List[Dict]:
    """
    Load label data from JSON file.
    
    Args:
        filepath: Path to label_data.json
        
    Returns:
        List of label data samples
    """
    with open(filepath, 'r', encoding='utf-8') as f:
        return json.load(f)


def filter_label_data(label_data: List[Dict], valid_image_names: Set[str]) -> List[Dict]:
    """
    Filter label data to keep only samples with matching image names.
    
    Args:
        label_data: List of label data samples
        valid_image_names: Set of valid image names from dbscan results
        
    Returns:
        Filtered list of label data samples
    """
    filtered_data = []
    
    for sample in label_data:
        image_name = sample.get('image', '')
        
        # Check if the image name matches any of the valid names
        # We'll check various matching strategies
        if (image_name in valid_image_names or
            # Check if any valid name is contained in the image name
            any(valid_name in image_name for valid_name in valid_image_names) or
            # Check if image name is contained in any valid name
            any(image_name in valid_name for valid_name in valid_image_names)):
            filtered_data.append(sample)
    
    return filtered_data


def save_filtered_data(filtered_data: List[Dict], output_filepath: str):
    """
    Save filtered data to JSON file.
    
    Args:
        filtered_data: Filtered label data
        output_filepath: Path to save the filtered data
    """
    with open(output_filepath, 'w', encoding='utf-8') as f:
        json.dump(filtered_data, f, ensure_ascii=False, indent=2)


def main():
    """Main function to execute the filtering process."""
    
    # Input file paths
    # dbscan_results_path = '014_dbscan_filtered_results_temp.json'
    # dbscan_results_path = '/home/nguyendc/sonnh/embedding-clustering/filter/027_dbscan_v3_filtered_results_temp.json'
    dbscan_results_path = '/home/nguyendc/sonnh/embedding-clustering/filter/029_dbscan_v3_filtered_results_temp_30.json'
    
    print(dbscan_results_path)
    label_data_path = '/home/nguyendc/sonnh/embedding-clustering/filter/label_data.json'
    
    # Output file path
    # output_path = '014_filtered_label_data_dbscan_v2_30_percent_amount_data.json'
    
    output_path = '030_dbscan_v3_retrain_014.json'
    
    # Check if input files exist
    if not os.path.exists(dbscan_results_path):
        print(f"Error: {dbscan_results_path} not found!")
        return
    
    if not os.path.exists(label_data_path):
        print(f"Error: {label_data_path} not found!")
        return
    
    print("Loading dbscan filtered results...")
    valid_image_names = load_dbscan_filtered_results(dbscan_results_path)
    print(f"Found {len(valid_image_names)} unique image names in dbscan results")
    
    print("Loading label data...")
    label_data = load_label_data(label_data_path)
    print(f"Loaded {len(label_data)} samples from label data")
    
    print("Filtering label data...")
    filtered_data = filter_label_data(label_data, valid_image_names)
    print(f"Filtered to {len(filtered_data)} samples")
    
    print(f"Saving filtered data to {output_path}...")
    save_filtered_data(filtered_data, output_path)
    
    print("Filtering completed successfully!")
    print(f"Original samples: {len(label_data)}")
    print(f"Filtered samples: {len(filtered_data)}")
    print(f"Reduction: {len(label_data) - len(filtered_data)} samples removed")
    print(f"Retention rate: {len(filtered_data)/len(label_data)*100:.2f}%")


if __name__ == "__main__":
    main()
