import pandas as pd

product = pd.read_csv("./classification/product/hs92.csv")
product = product.drop(
    columns=["product_space_x", "product_space_y", "cluster_id"],
    errors="ignore",
)

ps_nodes = pd.read_csv("./classification/raw/umap_ps.csv", dtype={"product_code": str})[
    ["product_code", "x", "y", "cluster_name", "cluster_name_short"]
]
ps_edges = pd.read_csv(
    "./classification/raw/augmented_prox_nw_atlas_top_edges_backboned.csv",
    dtype={"product": str, "source": str, "target": str},
).drop(columns=["product", "rank"])

ps_nodes = (
    ps_nodes.merge(
        product[["code", "product_id"]],
        left_on="product_code",
        right_on="code",
        how="left",
    )
    .drop(columns=["product_code", "code"])
    .rename(columns={"x": "product_space_x", "y": "product_space_y"})
)

clusters = (
    ps_nodes[["cluster_name", "cluster_name_short"]]
    .drop_duplicates()
    .sort_values("cluster_name")
    .rename(columns={"cluster_name": "name_en", "cluster_name_short": "name_short_en"})
    .reset_index(drop=True)
    .reset_index(drop=False, names="cluster_id")
)

ps_nodes = ps_nodes.merge(
    clusters,
    left_on=["cluster_name", "cluster_name_short"],
    right_on=["name_en", "name_short_en"],
    how="left",
).drop(
    columns=["cluster_name", "cluster_name_short", "name_en", "name_short_en"],
    errors="ignore",
)

ps_edges = (
    ps_edges.merge(
        product[["code", "product_id"]], left_on="target", right_on="code", how="left"
    )
    .drop(columns=["target", "code"])
    .rename(columns={"product_id": "target_id"})
    .merge(
        product[["code", "product_id"]], left_on="source", right_on="code", how="left"
    )
    .drop(columns=["source", "code"])
)

product = product.merge(ps_nodes, on="product_id", how="left")
product.to_csv("./classification/product/hs92.csv", index=False)

ps_edges.to_csv("./classification/product/hs92_ps_edges.csv", index=False)
clusters.to_csv("./classification/product/hs92_ps_clusters.csv", index=False)
