import json
import numpy as np
import os
import argparse
from pathlib import Path
from sklearn.metrics.pairwise import cosine_similarity
from tqdm import tqdm

def parse_args():
    parser = argparse.ArgumentParser(description="Cluster speakers based on feature similarity (emo_vec or condition).")
    parser.add_argument("--input", type=Path, required=True, help="Input manifest (singles).")
    parser.add_argument("--output", type=Path, required=True, help="Output manifest (clustered).")
    parser.add_argument("--threshold", type=float, default=0.92, help="Cosine similarity threshold.")
    parser.add_argument("--feature", type=str, default="emo_vec", choices=["emo_vec", "condition", "combined"], 
                        help="Which feature to use for clustering finding. 'condition' is usually better for speaker identity.")
    return parser.parse_args()

def main():
    args = parse_args()
    
    print(f"Loading manifest from {args.input}...")
    entries = []
    with args.input.open('r', encoding='utf-8') as f:
        for line in f:
            if line.strip():
                entries.append(json.loads(line))

    print(f"Found {len(entries)} files. Using feature: '{args.feature}' with threshold {args.threshold}...")

    valid_entries = []
    current_speaker_id = 0
    last_vec = None
    
    dataset_root = args.input.parent

    for i, entry in enumerate(tqdm(entries)):
        # Determine which paths to load based on --feature
        paths_to_load = []
        if args.feature == "emo_vec":
            paths_to_load.append(entry.get('emo_vec_path'))
        elif args.feature == "condition":
            paths_to_load.append(entry.get('condition_path'))
        elif args.feature == "combined":
            paths_to_load.append(entry.get('condition_path'))
            paths_to_load.append(entry.get('emo_vec_path'))

        current_vec_list = []
        load_failed = False
        
        for rel_path in paths_to_load:
            if not rel_path:
                load_failed = True
                break
            vec_path = dataset_root / rel_path
            if not vec_path.exists():
                load_failed = True
                break
            
            try:
                # Load feature
                vec = np.load(vec_path)
                # Flatten/Average logic
                # If [Time, Dim] (e.g. 32, 256) -> Average to [256]
                if vec.ndim > 1:
                    vec = vec.mean(axis=0)
                # If still multidimensional or scalar, flatten
                vec = vec.reshape(-1) 
                current_vec_list.append(vec)
            except Exception as e:
                # print(f"Error: {e}")
                load_failed = True
                break
        
        if load_failed or not current_vec_list:
            continue

        # Concatenate if combined, else just take the first
        current_vec = np.concatenate(current_vec_list, axis=0)
        current_vec = current_vec.reshape(1, -1) # Shape for sklearn (1, TotalDim)

        if last_vec is None:
            # First file
            entry['speaker'] = f"auto_speaker_{current_speaker_id}"
            last_vec = current_vec
        else:
            # Compare
            score = cosine_similarity(last_vec, current_vec)[0][0]
            
            if score > args.threshold:
                # Same speaker
                entry['speaker'] = f"auto_speaker_{current_speaker_id}"
                # Keep last_vec (Cluster Head Strategy) - prevents drift
            else:
                # New speaker
                current_speaker_id += 1
                entry['speaker'] = f"auto_speaker_{current_speaker_id}"
                last_vec = current_vec

        valid_entries.append(entry)

    print(f"Clustering complete. Found {current_speaker_id + 1} distinct speaker groups.")
    print(f"Saving to {args.output}...")

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open('w', encoding='utf-8') as f:
        for entry in valid_entries:
            f.write(json.dumps(entry, ensure_ascii=False) + "\n")

    print("Done!")

if __name__ == "__main__":
    main()