#!/usr/bin/env python3
"""
DBSCAN Clustering Filter

Filters clustering results based on specific criteria:
- For each cluster: select 50% of points
  - 25% from center region (closest to centroid)
  - 25% from border region (furthest from centroid)
- All noise points are selected
- Uses cosine distance metric
"""

import json
import numpy as np
from sklearn.preprocessing import normalize
from sklearn.metrics.pairwise import cosine_distances
import argparse
import os
from pathlib import Path
import random


class DBSCANFilter:
    def __init__(self, embeddings_path, clustering_results_path):
        """
        Initialize DBSCAN filter
        
        Args:
            embeddings_path: Path to embeddings JSON file
            clustering_results_path: Path to DBSCAN clustering results JSON
        """
        self.embeddings_path = embeddings_path
        self.clustering_results_path = clustering_results_path
        self.embeddings = None
        self.embeddings_normalized = None
        self.clustering_results = None
        self.filepath_to_embedding = {}
        
    def load_data(self):
        """Load embeddings and clustering results"""
        print("Loading embeddings...")
        with open(self.embeddings_path, 'r') as f:
            embeddings_data = json.load(f)
        
        # Create mapping from filepath to embedding
        embeddings_list = []
        filepaths = []
        for item in embeddings_data:
            self.filepath_to_embedding[item['filepath']] = item['embedding']
            embeddings_list.append(item['embedding'])
            filepaths.append(item['filepath'])
        
        self.embeddings = np.array(embeddings_list, dtype=np.float32)
        self.embeddings_normalized = normalize(self.embeddings, norm='l2')
        print(f"Loaded {len(embeddings_list)} embeddings")
        
        print("Loading clustering results...")
        with open(self.clustering_results_path, 'r') as f:
            self.clustering_results = json.load(f)
        
        print(f"Loaded clustering results: {self.clustering_results['n_clusters']} clusters, "
              f"{self.clustering_results['n_samples']} samples")
    
    def group_by_clusters(self):
        """Group data points by cluster labels"""
        clusters = {}
        noise_points = []
        
        for result in self.clustering_results['results']:
            cluster_id = result['cluster']
            filepath = result['filepath']
            if 'is_noise' in result:
                is_noise = result['is_noise']
            else:
                is_noise = False
            
            if is_noise or cluster_id == -1:
                noise_points.append({
                    'filepath': filepath,
                    'embedding': self.filepath_to_embedding[filepath]
                })
            else:
                if cluster_id not in clusters:
                    clusters[cluster_id] = []
                clusters[cluster_id].append({
                    'filepath': filepath,
                    'embedding': self.filepath_to_embedding[filepath]
                })
        
        return clusters, noise_points
    
    def calculate_cluster_centroid(self, cluster_points):
        """Calculate centroid of a cluster using normalized embeddings"""
        embeddings = np.array([point['embedding'] for point in cluster_points])
        embeddings_normalized = normalize(embeddings, norm='l2')
        
        # For cosine distance, centroid is the normalized mean
        centroid = np.mean(embeddings_normalized, axis=0)
        centroid_normalized = normalize(centroid.reshape(1, -1), norm='l2')[0]
        
        return centroid_normalized
    
    def calculate_cosine_distances_to_centroid(self, cluster_points, centroid):
        """Calculate cosine distances from each point to cluster centroid"""
        embeddings = np.array([point['embedding'] for point in cluster_points])
        embeddings_normalized = normalize(embeddings, norm='l2')
        
        # Calculate cosine distances to centroid
        distances = cosine_distances(embeddings_normalized, centroid.reshape(1, -1)).flatten()
        
        return distances
    
    # v3 
    def filter_cluster(self, cluster_points, selection_ratio=0.3):
    
        """
        Filter points from a cluster
        
        Args:
            cluster_points: List of points in the cluster
            selection_ratio: Ratio of points to select (default: 0.5 = 50%)
            
        Returns:
            List of selected points
        """
        if len(cluster_points) == 0:
            return []
        
        
        # Calculate how many points to select
        total_points = len(cluster_points)
        num_to_select = max(15, int(total_points * selection_ratio))
        
        # If we need to select all or almost all points, just return all
        if num_to_select >= total_points:
            return cluster_points
        
        # Calculate centroid
        centroid = self.calculate_cluster_centroid(cluster_points)
        
        # Calculate distances to centroid
        distances = self.calculate_cosine_distances_to_centroid(cluster_points, centroid)
        
        # Create list of (point, distance) pairs
        point_distance_pairs = list(zip(cluster_points, distances))
        
        # Sort by distance (closest to furthest from centroid)
        point_distance_pairs.sort(key=lambda x: x[1])
        
        dis = 0.1
        # count_center = sum(1 for pair in point_distance_pairs if pair[1] < dis)
        all_center_points = [pair[0] for pair in point_distance_pairs if pair[1] < dis]
        print(f"Number of center points (distance < {dis}): {len(all_center_points)}")

        # count_border = sum(1 for pair in point_distance_pairs if pair[1] >= dis)
        all_border_points = [pair[0] for pair in point_distance_pairs if pair[1] >= dis]
        print(f"Number of border points (distance >= {dis}): {len(all_border_points)}")

        # Calculate how many points to select from center and border
        n_center = len(all_center_points)
        center_count = max(1, int(n_center * 0.15))  
        
        n_border = len(all_border_points)
        if n_border < 70:
            border_count = n_border
        else:
            border_count = max(0, int(n_border * 0.3)) # remaining from border
            
        selected_points = []
        
        random.seed(42)
        # Select center points (closest to centroid)
        # center_points = [pair[0] for pair in point_distance_pairs[:center_count]]
        center_points = random.sample(all_center_points, center_count)
        selected_points.extend(center_points)
        
        # # Select border points (furthest from centroid)
        if border_count > 0:
            # border_points = [pair[0] for pair in point_distance_pairs[-border_count:]]
            border_points = random.sample(all_border_points, border_count)
            selected_points.extend(border_points)
        
        print(f"Cluster with {total_points} points -> selected {len(selected_points)} points "
              f"({center_count} center + {border_count} border)")
        
        return selected_points
    
    def filter_all_clusters(self):
        """Filter all clusters according to the specified criteria"""
        print("\n" + "="*60)
        print("FILTERING DBSCAN CLUSTERING RESULTS")
        print("="*60)
        
        clusters, noise_points = self.group_by_clusters()
        
        print(f"Found {len(clusters)} clusters and {len(noise_points)} noise points")
        
        filtered_results = []
        
        # Process each cluster
        for cluster_id, cluster_points in clusters.items():
            print(f"\nProcessing Cluster {cluster_id}:")
            filtered_points = self.filter_cluster(cluster_points)
            
            # Add cluster information
            for point in filtered_points:
                filtered_results.append({
                    'filepath': point['filepath'],
                    'cluster': cluster_id,
                    'is_noise': False,
                    'selection_type': 'cluster_filtered'
                })
        
        # Add all noise points
        print(f"\nAdding all {len(noise_points)} noise points...")
        for point in noise_points:
            filtered_results.append({
                'filepath': point['filepath'],
                'cluster': -1,
                'is_noise': True,
                'selection_type': 'noise'
            })
        
        return filtered_results

    def save_full_clusters_separately(self, output_dir=None):
        """Save ALL points of each cluster (and noise) into separate JSON files.

        The output format for each file matches the existing aggregated format
        produced by save_filtered_results (keys / structure unchanged). Only the
        content (subset of results) differs. Noise is treated as a cluster with id -1.

        Args:
            output_dir: Optional directory to place per-cluster files. If None,
                        files are written next to the clustering_results file.
        Returns:
            List of written file paths.
        """
        clusters, noise_points = self.group_by_clusters()

        base_name = Path(self.clustering_results_path).stem
        if output_dir is None:
            output_dir = Path('.')
        else:
            output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        written_files = []

        # Helper to build and write one cluster file
        def _write_cluster_file(cluster_id, points, is_noise=False):
            # Build results entries preserving field names
            results_entries = []
            for p in points:
                results_entries.append({
                    'filepath': p['filepath'],
                    'cluster': -1 if is_noise else cluster_id,
                    'is_noise': is_noise,
                    'selection_type': 'noise' if is_noise else 'cluster_full'
                })

            # Stats
            if is_noise:
                cluster_stats = {}
                noise_count = len(points)
            else:
                cluster_stats = {cluster_id: len(points)}
                noise_count = 0

            output_data = {
                "method": "DBSCAN_FILTERED",  # keep unchanged label
                "original_n_clusters": self.clustering_results['n_clusters'],
                "original_n_samples": self.clustering_results['n_samples'],
                "filtered_n_samples": len(results_entries),
                "filtering_criteria": {
                    # Keep same keys / structure (values kept as original literal values)
                    "cluster_selection_ratio": 0.5,
                    "center_points_ratio": 0.5,
                    "border_points_ratio": 0.5,
                    "noise_points": "all_selected"
                },
                "cluster_statistics": cluster_stats,
                "noise_points": noise_count,
                "results": results_entries
            }

            # Filename
            if is_noise:
                filename = f"{base_name}_cluster_noise.json"
            else:
                filename = f"{base_name}_cluster_{cluster_id}.json"
            out_path = output_dir / filename
            with open(out_path, 'w', encoding='utf-8') as f:
                json.dump(output_data, f, indent=4, ensure_ascii=False)
            written_files.append(str(out_path))
            print(f"Saved cluster {cluster_id if not is_noise else 'noise'} -> {out_path} ({len(points)} samples)")

        # Write each real cluster
        for cid, pts in clusters.items():
            _write_cluster_file(cid, pts, is_noise=False)

        # Noise as its own cluster
        if noise_points:
            _write_cluster_file(-1, noise_points, is_noise=True)

        return written_files
    
    def save_filtered_results(self, filtered_results, output_path=None):
        """Save filtered results to JSON file"""
        if output_path is None:
            # Generate output filename based on input
            base_name = Path(self.clustering_results_path).stem
            output_path = f"{base_name}_filtered.json"
        
        # Create summary statistics
        cluster_stats = {}
        noise_count = 0
        
        for result in filtered_results:
            if result['is_noise']:
                noise_count += 1
            else:
                cluster_id = result['cluster']
                if cluster_id not in cluster_stats:
                    cluster_stats[cluster_id] = 0
                cluster_stats[cluster_id] += 1
        
        # Prepare output data
        output_data = {
            "method": "DBSCAN_FILTERED",
            "original_n_clusters": self.clustering_results['n_clusters'],
            "original_n_samples": self.clustering_results['n_samples'],
            "filtered_n_samples": len(filtered_results),
            "filtering_criteria": {
                "cluster_selection_ratio": 0.5,
                "center_points_ratio": 0.5,  # 50% of selected points from center
                "border_points_ratio": 0.5,  # 50% of selected points from border
                "noise_points": "all_selected"
            },
            "cluster_statistics": cluster_stats,
            "noise_points": noise_count,
            "results": filtered_results
        }
        
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(output_data, f, indent=4, ensure_ascii=False)
        
        print("\n" + "="*60)
        print("FILTERING SUMMARY")
        print("="*60)
        print(f"Original samples: {self.clustering_results['n_samples']}")
        print(f"Filtered samples: {len(filtered_results)}")
        print(f"Reduction ratio: {len(filtered_results)/self.clustering_results['n_samples']:.2%}")
        print("\nCluster breakdown:")
        for cluster_id, count in sorted(cluster_stats.items()):
            print(f"  Cluster {cluster_id}: {count} points")
        print(f"  Noise points: {noise_count} points")
        print(f"\nFiltered results saved to: {output_path}")
        
        return output_path
    
    def create_filepath_list(self, filtered_results, output_txt_path=None):
        """Create a simple text file with filtered filepaths"""
        if output_txt_path is None:
            base_name = Path(self.clustering_results_path).stem
            output_txt_path = f"{base_name}_filtered_filepaths.txt"
        
        filepaths = [result['filepath'] for result in filtered_results]
        
        with open(output_txt_path, 'w', encoding='utf-8') as f:
            for filepath in filepaths:
                f.write(f"{filepath}\n")
        
        print(f"Filepath list saved to: {output_txt_path}")
        return output_txt_path


def main():
    parser = argparse.ArgumentParser(description="Filter DBSCAN clustering results")
    parser.add_argument("--embeddings_path", required=True,
                        help="Path to embeddings JSON file")
    parser.add_argument("--clustering_results_path", required=True,
                        help="Path to DBSCAN clustering results JSON file")
    parser.add_argument("--output_path", 
                        help="Output path for filtered results (optional)")
    parser.add_argument("--create_filepath_list", action="store_true",
                        help="Also create a simple text file with filtered filepaths")
    parser.add_argument("--save_full_clusters", action="store_true",
                        help="Save ALL points of each cluster (and noise) into separate JSON files without altering format")
    parser.add_argument("--clusters_output_dir",
                        help="Directory to store per-cluster JSON files (default: current directory)")
    
    args = parser.parse_args()
    
    # Validate input files exist
    if not os.path.exists(args.embeddings_path):
        print(f"Error: Embeddings file not found: {args.embeddings_path}")
        return
    
    if not os.path.exists(args.clustering_results_path):
        print(f"Error: Clustering results file not found: {args.clustering_results_path}")
        return
    
    # Initialize filter
    filter_obj = DBSCANFilter(args.embeddings_path, args.clustering_results_path)
    
    # Load data
    filter_obj.load_data()
    
    # Filter clusters
    filtered_results = filter_obj.filter_all_clusters()
    
    # Save results
    filter_obj.save_filtered_results(filtered_results, args.output_path)

    # Save full clusters separately if requested
    if args.save_full_clusters:
        print("\nSaving full clusters (no filtering) separately...")
        filter_obj.save_full_clusters_separately(args.clusters_output_dir)
    
    # Create filepath list if requested
    if args.create_filepath_list:
        filter_obj.create_filepath_list(filtered_results)
    
    print("\nFiltering completed successfully!")


if __name__ == "__main__":
    main()
