from hw2skeleton import cluster
from hw2skeleton import io
import os
import numpy as np
def test_similarity():
    filename_a = os.path.join("data", "276.pdb")
    filename_b = os.path.join("data", "4629.pdb")

    activesite_a = io.read_active_site(filename_a)
    activesite_b = io.read_active_site(filename_b)
    assert cluster.compute_similarity(activesite_a.counts, activesite_b.counts) == np.linalg.norm(activesite_a.counts - activesite_b.counts)
    assert cluster.compute_similarity(activesite_a.counts, activesite_a.counts) == 0
    assert cluster.compute_similarity(activesite_a.counts, activesite_b.counts) == cluster.compute_similarity(activesite_b.counts, activesite_a.counts)

def test_partition_clustering():
    # tractable subset
    pdb_ids = [276, 4629, 10701]

    active_sites = []
    for id in pdb_ids:
        filepath = os.path.join("data", "%i.pdb"%id)
        active_sites.append(io.read_active_site(filepath))


    assert len(cluster.cluster_by_partitioning(active_sites).keys()) >= 2

def test_hierarchical_clustering():
    # tractable subset
    pdb_ids = [276, 4629, 10701]

    active_sites = []
    for id in pdb_ids:
        filepath = os.path.join("data", "%i.pdb"%id)
        active_sites.append(io.read_active_site(filepath))

    print()
    assert len(cluster.cluster_hierarchically(active_sites).keys()) >=2
