import numpy as np

def calculate_distribution_similarity(dist1, dist2):            
    all_types = set(dist1.keys()) | set(dist2.keys())

                   
    smooth = 1e-10
    p1 = np.array([dist1.get(t, 0) + smooth for t in all_types])
    p2 = np.array([dist2.get(t, 0) + smooth for t in all_types])

         
    p1 = p1 / p1.sum()
    p2 = p2 / p2.sum()

              
    kl_div = np.sum(p1 * np.log(p1 / p2)) + np.sum(p2 * np.log(p2 / p1))
    return kl_div


def calculate_weighted_distribution_similarity(dist1, dist2, weight_power=2.0):           
    all_types = set(dist1.keys()) | set(dist2.keys())

    if not all_types:
        return 0.0

                   
    smooth = 1e-10
    total1 = sum(dist1.values()) + len(all_types) * smooth
    total2 = sum(dist2.values()) + len(all_types) * smooth

    p1 = {t: (dist1.get(t, 0) + smooth) / total1 for t in all_types}
    p2 = {t: (dist2.get(t, 0) + smooth) / total2 for t in all_types}

                       
    weighted_distance = 0.0
    total_weight = 0.0

    for t in all_types:
                           
        avg_freq = (p1[t] + p2[t]) / 2
        weight = avg_freq ** weight_power

                         
        diff = (p1[t] - p2[t]) ** 2

        weighted_distance += weight * diff
        total_weight += weight

         
    if total_weight > 0:
        weighted_distance = weighted_distance / total_weight

    return weighted_distance


def calculate_jensen_shannon_weighted(dist1, dist2, weight_power=1.5):     
    all_types = set(dist1.keys()) | set(dist2.keys())

    if not all_types:
        return 0.0

             
    smooth = 1e-10
    total1 = sum(dist1.values()) + len(all_types) * smooth
    total2 = sum(dist2.values()) + len(all_types) * smooth

    p1 = np.array([dist1.get(t, 0) + smooth for t in all_types]) / total1
    p2 = np.array([dist2.get(t, 0) + smooth for t in all_types]) / total2

            
    m = (p1 + p2) / 2

          
    weights = m ** weight_power
    weights = weights / weights.sum()         

              
    kl1 = np.sum(weights * p1 * np.log(p1 / m))
    kl2 = np.sum(weights * p2 * np.log(p2 / m))

                      
    js_div = 0.5 * kl1 + 0.5 * kl2
    return js_div


def calculate_similarity_by_method(dist1, dist2, method='kl', weight_power=1.5):     
    if method == 'kl':
        divergence = calculate_distribution_similarity(dist1, dist2)
        method_name = "KL Divergence"
    elif method == 'weighted':
        divergence = calculate_weighted_distribution_similarity(dist1, dist2, weight_power)
        method_name = "Weighted Squared Difference"
    elif method == 'js_weighted':
        divergence = calculate_jensen_shannon_weighted(dist1, dist2, weight_power)
        method_name = "Weighted JS Divergence"
    else:
        raise ValueError(f"Unsupported similarity calculation method: {method}")

    return divergence, method_name


              
SUPPORTED_SIMILARITY_METHODS = ['kl', 'weighted', 'js_weighted']

      
METHOD_DESCRIPTIONS = {
    'kl': 'Standard symmetric KL divergence, treats all categories equally',
    'weighted': 'Weighted squared difference, gives higher weight to high-frequency categories',
    'js_weighted': 'Weighted Jensen-Shannon divergence, well-balanced, gives higher weight to high-frequency categories (recommended)'
}
