visualization.py

import argparse
from hashring import HashRing
import random
import statistics
import matplotlib.pyplot as plt

def visualize_distribution(counts, title="Key Distribution"):
    """Create bar chart of key distribution."""
    nodes = list(counts.keys())
    values = list(counts.values())
    
    plt.figure(figsize=(10, 6))
    plt.bar(nodes, values, color='skyblue', edgecolor='black')
    plt.axhline(y=statistics.mean(values), color='red', linestyle='--', 
                label=f'Mean: {statistics.mean(values):.1f}')
    plt.xlabel('Nodes')
    plt.ylabel('Number of Keys')
    plt.title(title)
    plt.legend()
    plt.xticks(rotation=45)
    plt.tight_layout()
    filename = f'distribution_{title.lower().replace(" ", "_")}.png'
    plt.savefig(filename)
    print(f"Saved visualization to: {filename}")
    plt.close()

def calculate_stats(counts):
    """Calculate mean, stdev, and coefficient of variation."""
    values = list(counts.values())
    mean = statistics.mean(values)
    stdev = statistics.stdev(values) if len(values) > 1 else 0
    cv = (stdev / mean * 100) if mean > 0 else 0
    return mean, stdev, cv

def run_simulation(nodes, replicas, total_keys):
    # === Initial Ring ===
    node_list = [f"node_{i}" for i in range(nodes)]
    hr = HashRing(nodes=nodes, default_replicas=replicas)
    
    counts = {node: 0 for node in node_list}
    initial_mapping = {}
    
    for i in range(total_keys):
        key = f"key_{i}"
        node = hr.get_node(key)
        initial_mapping[key] = node
        counts[node] += 1
    
    # Print initial stats
    mean, stdev, cv = calculate_stats(counts)
    print(f"\n{'='*60}")
    print(f"INITIAL STATE ({nodes} nodes, {total_keys} keys)")
    print(f"{'='*60}")
    print(f"Mean keys per node: {mean:.1f}")
    print(f"Standard deviation: {stdev:.2f}")
    print(f"Coefficient of variation: {cv:.2f}%")
    print(f"Distribution: {counts}")

    visualize_distribution(counts, "Initial Distribution")

    # ========== ADD A NODE ==========
    print(f"\n{'='*60}")
    print(f"ADDING NODE: node_{nodes}")
    print(f"{'='*60}")

    node_to_add = f"node_{nodes}"
    hr.add_node(node_to_add)
    node_list.append(node_to_add)
    
    counts = {node: 0 for node in node_list}
    mapping_after_add = {}
    moved_keys = 0
    
    for key, old_node in initial_mapping.items():
        new_node = hr.get_node(key)
        counts[new_node] += 1
        mapping_after_add[key] = new_node
        if old_node != new_node:
            moved_keys += 1

    mean, stdev, cv = calculate_stats(counts)
    
    print(f"Keys moved: {moved_keys} ({moved_keys/total_keys*100:.2f}%)")
    print(f"Expected movement: ~{total_keys/(nodes+1):.0f} keys ({100/(nodes+1):.2f}%)")
    print(f"New distribution: {counts}")
    print(f"Mean keys per node: {mean:.1f}")
    print(f"Coefficient of variation: {cv:.2f}%")

    visualize_distribution(counts, "After Adding Node")
    
    # ========== REMOVE A NODE ==========
    node_to_remove = random.choice(node_list)
    
    print(f"\n{'='*60}")
    print(f"REMOVING NODE: {node_to_remove}")
    print(f"{'='*60}")

    hr.remove_node(node_to_remove)
    node_list.remove(node_to_remove)

    counts = {node: 0 for node in node_list}
    moved_keys = 0

    for key, old_node in mapping_after_add.items():
        new_node = hr.get_node(key)
        counts[new_node] += 1
        if old_node != new_node:
            moved_keys += 1
    
    mean, stdev, cv = calculate_stats(counts)
    
    print(f"Keys moved: {moved_keys} ({moved_keys/total_keys*100:.2f}%)")
    print(f"Expected movement: ~{total_keys/len(node_list):.0f} keys (~{100/len(node_list):.2f}%)")
    print(f"New distribution: {counts}")
    print(f"Mean keys per node: {mean:.1f}")
    print(f"Coefficient of variation: {cv:.2f}%")

    visualize_distribution(counts, "After Removing Node")

def main():
    parser = argparse.ArgumentParser(description="Test to run consistent hashing simulation based on input number of nodes, replicas and keys")

    parser.add_argument('-n', '--nodes', type=int, default=5, help='Number of nodes')
    parser.add_argument('-r', '--replicas', type=int, default=100, help='Number of replicas per node')
    parser.add_argument('-k', '--keys', type=int, default=10000, help='Number of keys')

    args = parser.parse_args()
    run_simulation(args.nodes, args.replicas, args.keys)

if __name__ == "__main__":
    main()