import numpy as np
import os
from joblib import Parallel, delayed
import time
import json

def compute_population_dist(human_list_embeddings, agent_list_embeddings, distance, method="sum", n_jobs=-1):
    if isinstance(human_list_embeddings, dict):
        human_list_embeddings = list(human_list_embeddings.values())
    def min_dist(human_embedding):
        return min(
            compute_dist_score(human_embedding, agent_embedding, distance)
            for agent_embedding in agent_list_embeddings
        )

    min_dists = Parallel(n_jobs=n_jobs)(
        delayed(min_dist)(human_embedding)
        for human_embedding in human_list_embeddings
    )

    return min(min_dists) if method == "min" else sum(min_dists)
    
def compute_dist_score(emb1, emb2, distance=None):        
    emb1 = emb1.flatten()
    emb2 = emb2.flatten()
    valid = ~(np.isnan(emb1) | np.isnan(emb2))
    total_coords = emb1.shape[0]
    valid_coords = np.sum(valid)
    if valid_coords == 0:
        return float('inf')  
    if distance == "euclidean":
        weight = total_coords / valid_coords
        squared_diff = (emb1[valid] - emb2[valid]) ** 2
        return np.sqrt(weight * np.sum(squared_diff))
    elif distance == "manhattan":
        weight = total_coords / valid_coords
        manhattan_dist = np.sum(np.abs(emb1[valid] - emb2[valid]))
        return weight * manhattan_dist

