import polars as pl
from codebase.load_accounts import read_metadata_file

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


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

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