import pandas as pd
import random
import sys
import math
from gensim.models import Word2Vec

def _to_edges_df(item):
    if isinstance(item, pd.DataFrame):
        df = item.copy()
    else:
        try:
            df = pd.DataFrame(item)
        except Exception:
            df = pd.DataFrame.from_dict(item)

    if df.shape[1] < 2:
        raise ValueError("Edge list must have at least 2 columns")

    cols = list(df.columns)
    if "source" in cols and "target" in cols:
        src_col, tgt_col = "source", "target"
    else:
        src_col, tgt_col = cols[0], cols[1]

    if "weight" in cols:
        w_col = "weight"
    elif len(cols) >= 3:
        w_col = cols[2]
    else:
        w_col = None

    keep = [src_col, tgt_col] + ([w_col] if w_col is not None else [])
    df = df[keep].copy()
    df.columns = ["source", "target"] + (["weight"] if w_col is not None else [])

    df["source"] = df["source"].astype(str)
    df["target"] = df["target"].astype(str)

    if "weight" in df.columns:
        df["weight"] = pd.to_numeric(df["weight"], errors="coerce").fillna(1.0).astype(float)
    else:
        df["weight"] = 1.0

    df.loc[df["weight"].isna(), "weight"] = 0.0
    df.loc[~df["weight"].map(lambda x: math.isfinite(float(x)) if x is not None else False), "weight"] = 0.0
    df.loc[df["weight"] < 0, "weight"] = 0.0

    return df

def _build_graph(df):
    adj = {}
    und = {}

    for s, t, w in zip(df["source"].tolist(), df["target"].tolist(), df["weight"].tolist()):
        if s not in adj:
            adj[s] = {}
        adj[s][t] = adj[s].get(t, 0.0) + float(w)

        und.setdefault(s, set()).add(t)
        und.setdefault(t, set()).add(s)

        adj.setdefault(t, {})

    nodes = list(und.keys())
    return adj, und, nodes

def _safe_positive(x, default=1.0, eps=1e-12):
    try:
        x = float(x)
    except Exception:
        x = float(default)
    if not math.isfinite(x) or x <= 0:
        x = float(default)
    if x < eps:
        x = eps
    return x

def _pick_neighbor(nbrs, weights):
    w = []
    for val in weights:
        try:
            val = float(val)
        except Exception:
            val = 0.0
        if not math.isfinite(val) or val < 0:
            val = 0.0
        w.append(val)
    if sum(w) <= 0:
        return random.choice(nbrs)
    return random.choices(nbrs, weights=w, k=1)[0]

def _one_walk(start, adj, und, p, q, length):
    walk = [start]
    if length <= 1:
        return walk

    cur = start
    prev = None
    p = _safe_positive(p, default=1.0)
    q = _safe_positive(q, default=1.0)

    for _ in range(length - 1):
        nbrs_dict = adj.get(cur, {})
        nbrs = list(nbrs_dict.keys())
        if not nbrs:
            break

        if prev is None:
            weights = [nbrs_dict.get(x, 1.0) for x in nbrs]
        else:
            prev_und = und.get(prev, set())
            weights = []
            for x in nbrs:
                base_w = nbrs_dict.get(x, 1.0)
                if x == prev:
                    alpha = 1.0 / p
                elif x in prev_und:
                    alpha = 1.0
                else:
                    alpha = 1.0 / q
                weights.append(base_w * alpha)

        nxt = _pick_neighbor(nbrs, weights)
        prev, cur = cur, nxt
        walk.append(cur)

    return walk

def _generate_walks(nodes, adj, und, p, q, n, length):
    walks = []
    for node in nodes:
        for _ in range(n):
            walks.append(_one_walk(node, adj, und, p, q, length))
    return walks

def net_embedding(net_list, p, q, num, l, vector_size, window, min_count, sg_model, workers, epochs):
    num = int(num)
    l = int(l)
    vector_size = int(vector_size)
    window = int(window)
    min_count = int(min_count)
    sg_model = int(sg_model)
    workers = int(workers)
    epochs = int(epochs)
    p = float(p)
    q = float(q)

    if isinstance(net_list, dict):
        net_items = list(net_list.items())
    else:
        net_items = [(str(i), x) for i, x in enumerate(net_list)]

    out = {}
    total = len(net_items)

    for idx, (name, item) in enumerate(net_items, start=1):
        print(f"Analyzing network {idx} of {total}: {name}")
        sys.stdout.flush()

        df = _to_edges_df(item)

        print("Building graph")
        sys.stdout.flush()
        adj, und, nodes = _build_graph(df)

        print("Starting biased random walk")
        sys.stdout.flush()
        walks = _generate_walks(nodes, adj, und, p, q, num, l)
        str_walks = [[str(n) for n in walk] for walk in walks]

        print("Starting Word2vec")
        sys.stdout.flush()
        model = Word2Vec(
            str_walks,
            vector_size=vector_size,
            window=window,
            min_count=min_count,
            sg=sg_model,
            workers=workers,
            epochs=epochs
        )

        node_ids = model.wv.index_to_key
        node_embeddings = model.wv.vectors
        out[name] = pd.DataFrame(node_embeddings, index=node_ids)

    return out
