import polars as pl
import numpy as np
from scipy.sparse import csr_matrix
from scipy.sparse import lil_matrix
from gensim.models import HdpModel
from gensim.corpora import Dictionary
from gensim.models import TfidfModel
import scipy
from sklearn.cluster import AffinityPropagation
from sklearn.metrics.pairwise import cosine_similarity
from gensim.models import LsiModel

CONSENSUS_ACCOUNTS = 100
CONSENSUS_SERVERS = 25

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))]


def common_tags_by_account(tags: pl.DataFrame) -> pl.DataFrame:
  return tags.explode("tags").with_columns(
    pl.col("tags").str.to_lowercase()
  ).unique(
    ["host", "acct", "tags"]
  ).group_by("tags").agg([
    pl.col("acct").len().alias("account_count")
  ]).sort("account_count", descending=True)

def common_tags_by_server(tags: pl.DataFrame) -> pl.DataFrame:
  return tags.explode("tags").with_columns(
    pl.col("tags").str.to_lowercase()
  ).unique(
    ["host", "tags"]
  ).group_by("tags").agg([
    pl.col("host").len().alias("server_count")
  ]).sort("server_count", descending=True)


def filter_tag_posts_by_tag_count(tags: pl.DataFrame, common_tags: pl.DataFrame, min_tags: int, max_tags: int) -> pl.DataFrame:
  return tags.filter(pl.col("tags").list.len() <= max_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"]).group_by(["host", "id", "acct"]).agg([
    pl.col("tags")
  ]).filter(pl.col("tags").list.len() >= min_tags)#.with_columns(
  #  pl.col("tags").list.join(',').alias("tags")
  #)#.unique(["host", "acct", "tags"])


def document_matrix(term_list: list[str], doc_list: list[list[str]]):
  term_index = {tag: i for i, tag in enumerate(sorted(term_list))}
  m = lil_matrix((len(doc_list), len(term_list)), dtype=np.int64)
  for i, post in enumerate(doc_list):
    for tag in post:
      m[i, term_index[tag]] += 1
  return m, term_index


def calculate_pmi(d: pl.DataFrame, min_count=10):
  doc_freq = d.explode("tags").unique(["host", "id", "tags"]).group_by("tags").count().filter(
    pl.col("count") >= min_count
  )
  return d.with_columns(
    pl.col("tags").map_elements(pairwise_sets).alias("pairs")
  ).explode("pairs").group_by("pairs").count().sort("count", descending=True).with_columns(
    pl.col("pairs").str.split(",").list.get(0).alias("first"),
    pl.col("pairs").str.split(",").list.get(1).alias("last")
  ).drop(["pairs"]).join(
    doc_freq.rename({"tags":"first","count":"first_count"}),on="first",how="inner"
  ).join(
    doc_freq.rename({"tags":"last","count":"last_count"}),on="last",how="inner"
  ).with_columns( #pmi
    ((pl.col("count") / len(d)) / ((pl.col("first_count") / len(d)) * (pl.col("last_count") / len(d)))).log(base=2).alias("pmi")
  ).sort("pmi", descending=True)


def construct_pmi_matrix(terms_list: list[str], d: pl.DataFrame):
  term_index = {tag: i for i, tag in enumerate(terms_list)}
  m = lil_matrix((len(terms_list), len(terms_list)), dtype=np.float64)
  for row in d.iter_rows(named=True):
    m[term_index[row["first"]], term_index[row["last"]]] = row["pmi"]
    m[term_index[row["last"]], term_index[row["first"]]] = row["pmi"]
  return m, term_index



def host_clusters(df, group_col="host"):
  t = df.explode("tags").rename({"tags":"tag"}).join(
    clusters, on="tag", how="inner"
  ).unique([group_col, "id", "cluster"])
  host_totals = t.unique([group_col, "cluster"]).group_by(group_col).count().rename({"count":"unique_cluster_count"}).join(
    t.unique([group_col, "id"]).group_by([group_col]).count().rename({"count":"host_docs"}),
    on=group_col,
    how="inner"
  ).join(
    t.group_by(group_col).count().rename({"count":"total_cluster_count"}),
    on=group_col,
    how="inner"
  )
  tf = t.select(
    pl.col([group_col, "id", "cluster"])
  ).group_by([group_col, "cluster"]).count().join(host_totals, on=group_col, how="inner").filter(
    #pl.col("host_docs") >= 100
    pl.col("total_cluster_count") >= 0
  ).with_columns(
    (pl.col("count")/pl.col("total_cluster_count")).alias("tf")
  )
  clust_host_occurs = t.filter(
    pl.col(group_col).is_in(set(tf[group_col]))
  ).unique([group_col, "cluster"]).group_by("cluster").count().sort("count")
  idf = clust_host_occurs.with_columns(
    (len(t.unique(group_col)) / pl.col("count")).log().alias("idf")).rename({"count":"cluster_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 = sorted(tf_idf.select(group_col).unique().sort(group_col)[group_col])
  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[group_col]], row["cluster"]] = row["tfidf"]
  return host_cluster_matrix, tf_idf, host_names, idf


def read_tags_file(file: str, accounts: set[str]) -> pl.DataFrame:
  return pl.read_ipc(file).with_columns(
    pl.concat_str([pl.col("host"), pl.lit("_"), pl.col("acct")]).alias("account_id")
  ).filter(
    pl.col("account_id").is_in(accounts)
  ).filter(#pl.col("sensitive") == False).filter(
    pl.col("language") == "en").filter(
    pl.col("tags").list.len() <= 8
  ).unique(["host", "id"]).filter(
    pl.col("host").is_in(set(metadata["server"]))
  ).select(pl.col(["host", "acct", "id", "sensitive", "created_at", "tags"])).explode("tags").with_columns(
    pl.col("tags").str.to_lowercase()
  ).with_columns(
    pl.when(pl.col("sensitive")).then(pl.col("tags") + "_sensitive").otherwise(pl.col("tags")).alias("tags")
  ).group_by(["host", "acct", "id", "created_at"]).agg([
    pl.col("tags")
  ])



from code.load_accounts import read_metadata_file
metadata = read_metadata_file("data/metadata-2024-01-31.feather").select(pl.col(["server", "user_count"])).filter(pl.col("user_count") >= 100)
accounts = pl.scan_ipc("data/scratch/all_accounts.feather").select(
  pl.col(["server", "acct", "bot", "noindex", "followers_count", "suspended"])
).filter(pl.col("bot") == False).filter(
  pl.col("noindex") == False).filter(
  pl.col("followers_count") > 1).filter(
  pl.col("suspended").fill_null(False) == False
  ).collect().rename({"server":"host"}).select(pl.col(["host", "acct"])).unique(["host", "acct"]).with_columns(
  pl.concat_str([pl.col("host"), pl.lit("_"), pl.col("acct")]).alias("id")
)
#5_789_169
all_tag_posts = pl.concat([
  read_tags_file("data/tags_filt.feather", set(accounts["id"].to_list())),
  read_tags_file("data/tags-202302-202308.feather", set(accounts["id"].to_list()))
]).unique(["host", "id"])#.filter(pl.col("created_at") >= pl.date(2023, 1, 1)).filter(pl.col("created_at") < pl.date(2023, 8, 1))
all_tag_posts.write_ipc("data/scratch/all_tag_posts.feather")

all_tag_posts = pl.read_ipc("data/scratch/all_tag_posts.feather").filter(pl.col("created_at") >= pl.date(2023, 2, 1)).filter(pl.col("created_at") < pl.date(2023, 8, 1))

common_tags_account = common_tags_by_account(all_tag_posts)
common_tags_server = common_tags_by_server(all_tag_posts)
common_tags = common_tags_account.join(
  common_tags_server, on="tags", how="inner"
).filter(pl.col("server_count") >= CONSENSUS_SERVERS).filter(pl.col("account_count") >= CONSENSUS_ACCOUNTS)

common_tags.write_ipc("data/scratch/common_tags.feather")
all_tag_posts_filtered = filter_tag_posts_by_tag_count(all_tag_posts, common_tags, 2, 5)
all_tag_posts = None
pmi = calculate_pmi(all_tag_posts_filtered, min_count=1).filter(pl.col("pmi") > 0)#.filter(pl.col("count") >= 0).filter(pl.col("first_count") >= 5).filter(pl.col("last_count") >= 5)

pmi_terms = sorted(set(pmi["first"]).union(set(pmi["last"])))
pmi_matrix, term_index = construct_pmi_matrix(pmi_terms, pmi)
af_clust = AffinityPropagation(affinity='precomputed', damping = 0.85, max_iter = 1000).fit(pmi_matrix.toarray())
#n_components, comp_labels = scipy.sparse.csgraph.connected_components(pmi_matrix, directed=False)
clusters = pl.DataFrame({
  "tag": pmi_terms,
  "cluster": af_clust.labels_
}).sort("cluster")
#clusters = clusters.with_row_index().with_columns(pl.col("index").alias("cluster")).drop("index")
host_cluster_matrix, tf_idf, host_names, idf = host_clusters(all_tag_posts_filtered, group_col="host")
host_similarity = cosine_similarity(host_cluster_matrix)


def most_similar_servers(server: str):
  string_to_index_host = {string: index for index, string in enumerate(host_names)}
  return pl.DataFrame({
    "server": host_names,
    "similarity": host_similarity[string_to_index_host[server]]
  }).sort("similarity", descending=True)

def build_similarity_df():
  keys = []
  values = []
  for i in range(len(host_names)):
    for j in range(i+1, len(host_names)):
      if host_similarity[i, j] > 0.0:
        keys.append(','.join(sorted([host_names[i], host_names[j]])))
        values.append(host_similarity[i, j])
  df = pl.DataFrame({
    "servers": keys,
    "similarity": values
  }).sort("similarity", descending=True)
  return df.with_row_count().with_columns(
    pl.col("servers").str.split(',').alias("server"),
    # id column from row number
  ).explode("server").join(
    metadata,
    on="server",
    how="inner"
  ).group_by(["row_nr", "similarity"]).agg([
    pl.col("server")
  ]).filter(
    pl.col("server").list.len() == 2
  )

server_similarities = build_similarity_df()


server_similarities.filter(
  pl.col("similarity") > 0.1
).with_columns(
  pl.col("server").list.get(0).alias("Source"),
  pl.col("server").list.get(1).alias("Target"),
).rename({"similarity":"Weight"}).select(pl.col(["Source", "Target", "Weight"])).write_csv("data/scratch/similar_servers.csv")

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

from sklearn.decomposition import PCA
pca = PCA(n_components=3)
p = pca.fit_transform(host_cluster_matrix)

from sklearn.decomposition import NMF
model = NMF(n_components=32, init='random', random_state=0)
H = model.fit_transform(host_cluster_matrix)
W = model.components_

###
#from sklearn.cluster import AgglomerativeClustering
#server_clust_ap = AffinityPropagation(affinity='precomputed', damping = 0.95, max_iter = 1000).fit(host_similarity)
#server_clusters = pl.DataFrame({
#  "server": host_names,
#  "cluster": server_clust_ap.labels_
#}).sort("cluster")
#server_clusters.filter(pl.col("cluster") == server_clusters.filter(pl.col("server") == "econtwitter.net")["cluster"][0])
#
###
"""
dist_matrix = np.around(1 - host_similarity, 10)
ac = AgglomerativeClustering(n_clusters=1500, metric='precomputed', linkage='average').fit(dist_matrix)
np.array(host_names)[np.where(ac.labels_ == ac.labels_[host_names.index("hci.social")])]

np.array(host_names)[np.where(ac.labels_ == ac.labels_[host_names.index("mastodon.social")])]

###
# Can we see what servers are most similar to a users' history?

x = pl.scan_ipc("data/tags-2020-2022.feather").filter(
  pl.col("host") == "mastodon.gamedev.place"
).filter(pl.col("acct") == "raccoonformality").collect().explode("tags").with_columns(pl.col("tags").str.to_lowercase()).rename({"tags":"tag"}).join(
  clusters, on="tag", how="inner"
)
z = x.group_by(["cluster"]).count().sort("count", descending=True).with_columns(
  (pl.col("count") / len(x)).alias("tf")
).join(idf, on="cluster", how="inner").with_columns(
  (pl.col("tf") * pl.col("idf")).alias("tfidf")
).sort("tfidf", descending=True)

account_matrix = np.zeros((1, clusters["cluster"].max() + 1), dtype=np.float64)
for row in z.iter_rows(named=True):
  account_matrix[0, row["cluster"]] = row["tfidf"]

acc_similarity = pl.DataFrame({
  "host": host_names,
  "similarity": cosine_similarity(account_matrix, host_cluster_matrix)[0]
}).sort("similarity", descending=True)

###
#cluster_names = clusters.join(common_tags.rename({"tags":"tag"}), on="tag", how="inner").sort("count", descending=True).unique("cluster").select(pl.col(["tag", "cluster"])).sort("cluster")
"""