editdist_sp = [ (sp1,sp2,editdistance.eval(sp1,sp2)) for sp1,sp2 in itertools.combinations(read_annot["species_fillna"].unique(),2) ]
editdist_df = pd.DataFrame.from_records(editdist_sp,columns=["sp1","sp2","edit_distance"])
editdist_df["similarity"] = editdist_df.apply(lambda r: (max(len(r["sp1"]),len(r["sp2"]))-r["edit_distance"])/ max(len(r["sp1"]),len(r["sp2"])),axis=1)
editdist_df.head()

cluster_membership = {}
clusters = []

#Single linkage clustering
for _,row in editdist_df[editdist_df.similarity > 0.65 ].sort_values("similarity",ascending=False).iterrows():
    if row["sp1"] not in cluster_membership and row["sp2"] not in cluster_membership:
        #Create new cluster and add them both
        clusters.append([row["sp1"],row["sp2"]])
        cluster_membership[row["sp1"]] = len(clusters)-1
        cluster_membership[row["sp2"]] = len(clusters)-1
    elif row["sp1"] in cluster_membership:
        if row["sp2"] not in cluster_membership:
            #Add sp2 to sp1 cluster
            clusters[cluster_membership[row["sp1"]]].append(row["sp2"])
            cluster_membership[row["sp2"]] = cluster_membership[row["sp1"]]
        else:
            if cluster_membership[row["sp1"]] != cluster_membership[row["sp2"]]:
                #Combine clusters!
                c1_pos = cluster_membership[row["sp1"]]
                c2_pos = cluster_membership[row["sp2"]]
                clusters[c1_pos] += clusters[c2_pos]
                for c in clusters[c2_pos]:
                    cluster_membership[c] = c1_pos
                clusters[c2_pos] = None
    else:
        #Row 2 is in cluster and row1 is not
        clusters[cluster_membership[row["sp2"]]].append(row["sp1"])
        cluster_membership[row["sp1"]] = cluster_membership[row["sp2"]]

cluster_names = [";".join(c) if c else None for c in clusters]

"""
Explanation: Combine similar annotations into a single 'clade'
E.g Enterobacteria phage phi92 and phi91 should be grouped
End of explanation
"""


read_annot["sp_group"] = read_annot["species_fillna"].apply(lambda x: cluster_names[cluster_membership[x]] if x in cluster_membership and "unclassified" not in x else x)
df_counts = read_annot.groupby("sp_group").size().reset_index()
df_counts.columns = ["species","read_count"]
print(df_counts[df_counts.read_count >= 2].shape)
df_counts[df_counts.read_count >= 2].sort_values("read_count",ascending=False)

virus_count_histogram = df_counts[df_counts.read_count >= 3].groupby("read_count").size().reset_index()
sns.barplot(x=virus_count_histogram["read_count"],y=virus_count_histogram[0])

df_counts[df_counts.read_count >= 2].sort_values("read_count",ascending=False).to_csv("454_seqs_kaiju_greedy.filt_species.tsv",sep="\t",index=False)

!ls

"""
Explanation: Use new clade groups
End of explanation
"""


