import json
import numpy as np
from typing import List, Tuple
from sklearn.mixture import GaussianMixture
from sklearn.metrics import silhouette_samples
from tqdm import tqdm
import argparse
import os

def load_embeddings_from_json(json_path: str) -> Tuple[List[str], np.ndarray]:
    """Loads file paths and embeddings from a JSON file."""
    print(f"Loading embeddings from {json_path}...")
    with open(json_path, "r") as f:
        embedding_data = json.load(f)
    
    file_paths = [record["filepath"] for record in embedding_data]
    embeddings = np.array([record["embedding"] for record in embedding_data], dtype=np.float32)
    
    print(f"Loaded {len(file_paths)} samples with embedding dimension {embeddings.shape[1]}.")
    return file_paths, embeddings

def create_gmm_and_bic_score(embeddings: np.ndarray, n_clusters: int) -> Tuple[GaussianMixture, float]:
    """Fits a GMM and adds cluster labels to the dataset."""
    field_name = f"gmm_cluster_{n_clusters}"
    print(f"Running GMM with {n_clusters} components. Results will be in field '{field_name}'...")
    
    gmm = GaussianMixture(n_components=n_clusters, random_state=42)
    try:
        gmm.fit(embeddings)
        print("GMM clustering complete.")
    except:
        print("Passing this, error")
        return None, np.inf
    print("BIC score:", np.abs(gmm.bic(embeddings)))
    return gmm, np.abs(gmm.bic(embeddings))

def analyze_and_select_gmm(
    embeddings: np.ndarray,
    min_num_clusters: int,
    max_num_clusters: int,
    step_size: int = 10
):
    n_clusters_list = list(range(min_num_clusters, max_num_clusters + 1, step_size))
    bic_list = []
    gmm_list = []
    for n_clusters in tqdm(n_clusters_list, desc = "GMM analysis based on number of clusters"):
        gmm, bic_score = create_gmm_and_bic_score(embeddings=embeddings, n_clusters=n_clusters)
        gmm_list.append(gmm)
        bic_list.append(bic_score)
    
    bic_min_index = np.argmin(bic_list)
    selected_gmm = gmm_list[bic_min_index]
    return selected_gmm

def select_samples(
    embeddings: np.ndarray,
    gaussian_mixtures: GaussianMixture,
    filepaths: List[str]
):
    percentage_to_label:float = 0.5
    percentage_representative: float = 0.33
    
    total_number_images = len(filepaths)
    total_images_to_label = int(total_number_images * percentage_to_label)
    
    num_high_confidence = int(total_images_to_label * percentage_representative)
    num_low_confidence = total_images_to_label - num_high_confidence
    
    print(f"Total images: {total_number_images}")
    print(f"Total to label ({percentage_to_label*100}%): {total_images_to_label}")
    print(f"  - Highly representative samples: {num_high_confidence}")
    print(f"  - Least representative samples: {num_low_confidence}")
    
    labels = gaussian_mixtures.predict(embeddings)
    individual_scores = silhouette_samples(embeddings, labels)
    
    sorted_indices = np.argsort(individual_scores)
    # 3. Identify the indices for low and high confidence samples
    low_confidence_indices = sorted_indices[:num_low_confidence]
    high_confidence_indices = sorted_indices[-num_high_confidence:]

    # 4. Combine them into a final list of indices to be labeled
    final_indices_to_label = np.concatenate([low_confidence_indices, high_confidence_indices])

    # 5. Get the corresponding image file paths
    labeling_queue = [filepaths[i] for i in final_indices_to_label]
    
    return labeling_queue

def main():
    """Main function to run the analysis pipeline."""
    parser = argparse.ArgumentParser(description="Cluster and visualize document embeddings with FiftyOne.")
    
    parser.add_argument("json_path", type=str, help="Path to the JSON file containing embeddings.")
    parser.add_argument("json_label_path", type=str, help="Path to the JSON label file.")
    # parser.add_argument("--bill-type", type = str, help="Health invoice document type, eg. Ostéopathie, lentille, etc.")
    parser.add_argument("--min_num_clusters", type=int, default=10, help="Minimum numbers of clusters for GMM.")
    parser.add_argument("--max_num_clusters", type=int, default=100, help="Maximum numbers of clusters for GMM.")
    parser.add_argument("--step_size_clusters", type=int, default=10, help="Maximum numbers of clusters for GMM.")
    args = parser.parse_args()
    
    # 1. Load data
    filepaths, embeddings = load_embeddings_from_json(args.json_path)
    with open(args.json_label_path) as file:
        label_data = json.load(file) 
    
    gaussian_mixture = analyze_and_select_gmm(
        embeddings=embeddings,
        min_num_clusters=args.min_num_clusters,
        max_num_clusters=args.max_num_clusters,
        step_size=args.step_size_clusters
    )
    
    selected_filepaths = select_samples(
        embeddings=embeddings, gaussian_mixtures=gaussian_mixture, filepaths=filepaths
    )
    
    selected_label = []
    selected_prefixes = list(set([os.path.splitext(os.path.split(datum)[1])[0].split("_scale")[0][:-2] for datum in selected_filepaths]))
    for prefix in tqdm(selected_prefixes):
        try:
            label = list(filter(lambda x: bool(prefix.count(x["image"])), label_data))[0]
            selected_label.append(label)
        except:
            continue
        
    with open('./output_files.json', "w") as file:
        file.write(json.dumps(selected_label, indent=4))

if __name__ == "__main__":
    main()