import polars as pl
from scipy.sparse import dok_matrix
import numpy as np
from sklearn.cluster import AffinityPropagation

import scipy.sparse.csgraph
from scipy.sparse import csr_matrix
from sklearn.cluster import SpectralClustering
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE


def pairwise_sets(input_list_str):
  input_list = input_list_str.split(',')  
  return [','.join(sorted([input_list[i], input_list[j]])) for i in range(len(input_list)) for j in range(i+1, len(input_list))]


tags = pl.read_ipc("data/tags_filt.feather").filter(pl.col("sensitive") == False)#.filter(pl.col("language") == "en")

CONSENSUS_ACCOUNTS = 250

common_tags = tags.explode("tags").with_columns(
  pl.col("tags").str.to_lowercase()
).unique(
  ["host", "acct", "tags"]
).group_by("tags").count().filter(
  pl.col("count") >= CONSENSUS_ACCOUNTS
).sort("count", descending=True)

filt_tags = tags.explode("tags").with_columns(
  pl.col("tags").str.to_lowercase()
).filter(
  pl.col("tags").is_in(set(common_tags["tags"]))
).sort(["host", "acct", "id", "tags"])

filt_posts = filt_tags.filter(pl.col("sensitive").not_()).group_by(["host", "id", "acct", "created_at"]).agg([
  pl.col("tags")
]).filter(
  pl.col("tags").list.len() >= 2
).filter(
  pl.col("tags").list.len() <= 6
).with_columns(pl.col("tags").list.join(',').alias("tags")).unique(["host", "acct", "tags"])

pairs_posts = filt_posts.with_columns(
  pl.col("tags").map_elements(pairwise_sets).alias("pairs")
).explode("pairs").select(
  pl.col(["host", "acct", "pairs"])
).unique()
# TODO: ^

pairs = pairs_posts.group_by("pairs").count().sort("count", descending=True).with_columns(
  pl.col("pairs").map_elements(lambda x: x.split(",")).alias("pairs")
)

account_paired_tag_counts = pairs_posts.with_columns( # was filt_posts
  pl.col("pairs").alias("tags")
).with_columns(
  pl.col("tags").str.split(',')
).explode("tags").group_by("tags").agg(
  pl.col("host").len().alias("count")
).sort("count", descending=True)

total_posts = len(pairs_posts)

df = pairs.with_columns(
  pl.col("pairs").list.get(0).alias("first"),
  pl.col("pairs").list.get(1).alias("last")
).drop(["pairs"]).join(
  account_paired_tag_counts.rename({"tags":"first","count":"first_count"}),on="first",how="inner"
).join(
  account_paired_tag_counts.rename({"tags":"last","count":"last_count"}),on="last",how="inner"
)

p = df.with_columns(
  (pl.col("first_count") / total_posts).alias("p1"),
  (pl.col("last_count") / total_posts).alias("p2"),
  (pl.col("count") / total_posts).alias("p_joint"),
).with_columns(
  (pl.col("p_joint")/(pl.col("p1")*pl.col("p2"))).log(base=2).alias("pmi")
).sort("pmi")

common_pairs = p.filter(
  pl.col("count") >= 1
).sort("pmi", descending=True).filter(
  pl.col("pmi") > 0
).filter(
  pl.col("p_joint")/pl.col("p1") >= 0.01
).filter(
  pl.col("p_joint")/pl.col("p2") >= 0.01
)

most_common_tags = sorted(list(set(common_pairs["first"]).union(common_pairs["last"])))
string_to_index = {string: index for index, string in enumerate(most_common_tags)}
mapped_indices = [string_to_index[string] for string in most_common_tags]
S = csr_matrix((len(most_common_tags), len(most_common_tags)), dtype=np.float64)
#np.zeros((len(most_common_tags), len(most_common_tags)), dtype=np.float64) #* -100.0
for r in common_pairs.with_columns(
    pl.col("first").replace(string_to_index).cast(pl.UInt64),#.map_elements(lambda x: string_to_index[x]).alias("first"),
    pl.col("last").replace(string_to_index).cast(pl.UInt64)#.map_elements(lambda x: string_to_index[x]).alias("last"),
).iter_rows(named=True):
  S[r["first"], r["last"]] = r["pmi"]
  S[r["last"], r["first"]] = r["pmi"]

n_components, comp_labels = scipy.sparse.csgraph.connected_components(S, directed=False)
comp_S = S[:, comp_labels == 1][comp_labels == 1, :]

###
#clustering = AffinityPropagation(affinity='precomputed', damping = 0.95, max_iter = 1000).fit_predict(S)
#cdata = pl.DataFrame({
#  "tag": most_common_tags,
#  "cluster": clustering
#}).sort("cluster")
#cdata.group_by("cluster").agg([pl.col("tag").len().alias("count")]).sort("count", descending=True)
#cdata.filter(pl.col("cluster") == 17)
###

#spec = SpectralClustering(n_clusters=500, affinity='precomputed', assign_labels='discretize')
#labels = spec.fit_predict(comp_S)

af_clust = AffinityPropagation(affinity='precomputed', damping = 0.95, max_iter = 1000).fit(comp_S.toarray())

clusters = pl.DataFrame({
  "tag": np.array(most_common_tags)[comp_labels == 1],
  "cluster": af_clust.labels_
}).sort("cluster")

clusters.write_ipc("data/scratch/tag_clusters.feather")

t = tags.explode("tags").rename({"tags":"tag"}).join(clusters, on="tag", how="inner").unique(["host", "id", "cluster"])
host_totals = t.unique(["host", "tag"]).group_by("host").count().rename({"count":"tag_count"}).join(
  t.group_by(["host"]).count().rename({"count":"term_count"}),
  on="host",
  how="inner"
)
tf = t.group_by(["host", "cluster"]).count().join(host_totals, on="host", how="inner").filter(
  pl.col("tag_count") >= 100
).with_columns(
  (pl.col("count")/pl.col("term_count")).alias("tf")
)
clust_host_occurs = t.filter(
  pl.col("host").is_in(set(tf["host"]))
).unique(["host", "cluster"]).group_by("cluster").count().sort("count")
idf = clust_host_occurs.with_columns((len(t.unique("host")) / pl.col("count")).log().alias("idf")).rename({"count":"doc_count"})
tf_idf = tf.join(idf, on=["cluster"], how="inner").with_columns(
  (pl.col("tf")*pl.col("idf")).alias("tfidf")
).sort("tfidf", descending=True)

host_names = list(tf_idf.select("host").unique().sort("host")["host"])
string_to_index_host = {string: index for index, string in enumerate(host_names)}
#mapped_indices_host = [string_to_index_host[string] for string in host_names]
host_cluster_matrix = np.zeros((len(host_names), clusters["cluster"].max() + 1), dtype=np.float64)
for row in tf_idf.iter_rows(named=True):
  host_cluster_matrix[string_to_index_host[row["host"]], row["cluster"]] = row["tfidf"]

#from sklearn.metrics.pairwise import linear_kernel
from sklearn.metrics.pairwise import cosine_similarity
host_similarity = cosine_similarity(host_cluster_matrix)

def most_similar(server):
  string_to_index_host = {string: index for index, string in enumerate(host_names)}
  return pl.DataFrame({
    "server": host_names,
    "sim": host_similarity[string_to_index_host[server]]
  }).sort("sim", descending=True)#[0:10]

# Viz
X_reduced = PCA(n_components=2).fit_transform(host_similarity)
X_embedded = TSNE(n_components=2, learning_rate='auto', init='random', perplexity=3).fit_transform(host_similarity)

af = AffinityPropagation(affinity='precomputed').fit(host_similarity)

pl.DataFrame({
  "server": host_names,
  "cluster": af.labels_,
  "x": X_reduced[:,0],
  "y": X_reduced[:,1]
}).write_ipc("data/scratch/host_pca.feather")

pl.DataFrame({
  "server": host_names,
  "cluster": af.labels_,
  "x": X_embedded[:,0],
  "y": X_embedded[:,1]
}).write_ipc("data/scratch/host_tsne.feather")

pl.DataFrame({
  "server": host_names,
  "x": X_embedded[:,0],
  "y": X_embedded[:,1]
}).write_csv("data/scratch/host_pca.csv")

# Using cosine similarity
"""
dist_matrix = np.around(1 - host_similarity, 10)

dist_matrix = np.where(np.abs(1 - host_similarity) < 0.000000001, 0, 1 - host_similarity)
X_sim = TSNE(n_components=2, learning_rate='auto', init='random', perplexity=3, metric='precomputed').fit_transform(dist_matrix)
pl.DataFrame({
  "server": host_names,
  "x": X_sim[:,0],
  "y": X_sim[:,1]
}).write_csv("data/scratch/host_tsne_sim.csv")

pl.DataFrame({
  "server": host_names,
  "x": X_sim[:,0],
  "y": X_sim[:,1]
}).filter(pl.col("server") == "scholar.social")
"""