import random
import numpy as np
import networkx as nx
from joblib import Parallel, delayed
import itertools
from alias import create_alias_table, alias_sample
from tqdm import tqdm
import argparse
import multiprocessing
import logging
import os

def softmax(x, temperature=1.0):
    """
    计算带有温度参数的softmax，并加入防止溢出的技巧
    """
    x = np.array(x)
    x_max = np.max(x)
    exp_x = np.exp((x - x_max) / temperature)  # 加入temperature参数
    return exp_x / np.sum(exp_x)

class DeepWalk:
    def __init__(self, edge_file, node_tag_file, use_softmax=True, temperature=1.0, p_tag_walk=0.5):
        """
        初始化DeepWalk实例，构建图和标签索引，预处理alias采样表
        """
        logging.info(f"Initializing DeepWalk with edge file: {edge_file} and node-tag file: {node_tag_file}")
        self.graph = self.build_graph_from_edge_file(edge_file)
        if node_tag_file:
            self.node_to_tags, self.tag_to_nodes = self.build_tag_index(node_tag_file)
        else:
            self.node_to_tags = None
            self.tag_to_nodes = None
            
        self.alias_nodes = {}
        self.p_tag_walk = p_tag_walk
        logging.info(f"Graph built with {self.graph.number_of_nodes()} nodes and {self.graph.number_of_edges()} edges.")

        if use_softmax:
            logging.info(f"Using softmax with temperature: {temperature}")
            self.preprocess_transition_probs__softmax(temperature)
        else:
            logging.info("Using standard alias sampling.")
            self.preprocess_transition_probs()

    def build_graph_from_edge_file(self, edge_file):
        """
        从edge文件构建图
        edge文件格式: bid1 \t bid2:weight1,bid2:weight2,...
        """
        G = nx.Graph()

        # 打开edge文件并读取内容
        with open(edge_file, 'r') as f:
            for line in f:
                parts = line.strip().split('\t')
                if len(parts) != 2:
                    continue
                node, edges_str = parts
                edges = edges_str.split(',')
                
                for edge in edges:
                    nbr, weight = edge.split(':')
                    try:
                        node, nbr = int(node), int(nbr)
                    except ValueError:
                        continue
                    weight = float(weight)
                    
                    # 检查图中是否已存在这条边
                    if G.has_edge(node, nbr):
                        # 如果已经有这条边，更新权重，累加新权重
                        G[node][nbr]['weight'] += weight
                    else:
                        # 如果没有这条边，直接添加
                        G.add_edge(node, nbr, weight=weight)

        return G

    def build_tag_index(self, node_tag_file):
        """
        构建节点-标签的正排和倒排索引
        node_tag_file格式: book_id \t tag1,tag2,tag3
        """
        node_to_tags = {}
        tag_to_nodes = {}

        with open(node_tag_file, 'r') as f:
            for line in f:
                parts = line.strip().split('\t')
                if len(parts) != 2:
                    continue
                node, tags_str = parts
                try:
                    node = int(node)
                except ValueError:
                    continue
                # 只保留有过用户行为的node
                if not node in self.graph:
                    continue
                tags = tags_str.split(',')
                node_to_tags[node] = tags
                for tag in tags:
                    tag_to_nodes.setdefault(tag, []).append(node)

        return node_to_tags, tag_to_nodes

    def preprocess_transition_probs(self):
        """
        预处理节点的alias采样表，用于快速加权随机游走
        """
        G = self.graph

        for node in G.nodes():
            unnormalized_probs = [G[node][nbr].get('weight', 1.0) for nbr in G.neighbors(node)]
            norm_const = sum(unnormalized_probs)
            normalized_probs = [float(u_prob) / norm_const for u_prob in unnormalized_probs]
            self.alias_nodes[node] = create_alias_table(normalized_probs)

    def preprocess_transition_probs__softmax(self, temperature=1.0):
        """
        预处理节点的alias采样表，用于快速加权随机游走
        """
        G = self.graph

        for node in G.nodes():
            unnormalized_probs = [G[node][nbr].get('weight', 1.0) for nbr in G.neighbors(node)]
            normalized_probs = softmax(unnormalized_probs, temperature)
            self.alias_nodes[node] = create_alias_table(normalized_probs)

    def deepwalk_walk(self, walk_length, start_node):
        """
        执行一次DeepWalk随机游走，基于alias方法加速，支持通过标签游走
        """
        G = self.graph
        alias_nodes = self.alias_nodes
        walk = [start_node]

        while len(walk) < walk_length:
            cur = walk[-1]

            # 根据p_tag_walk的概率决定是通过邻居游走还是通过tag游走
            if self.node_to_tags and random.random() < self.p_tag_walk and cur in self.node_to_tags:
                walk = self.tag_based_walk(cur, walk)
            else:
                walk = self.neighbor_based_walk(cur, alias_nodes, walk)

            if not walk:
                break

        return walk

    def neighbor_based_walk(self, cur, alias_nodes, walk):
        """
        基于邻居的随机游走
        """
        G = self.graph
        cur_nbrs = list(G.neighbors(cur))
        if len(cur_nbrs) > 0:
            idx = alias_sample(alias_nodes[cur][0], alias_nodes[cur][1])
            walk.append(cur_nbrs[idx])
        else:
            return None
        return walk

    def tag_based_walk(self, cur, walk):
        """
        基于标签的随机游走
        """
        tags = self.node_to_tags[cur]
        if not tags:
            return None

        # 随机选择一个tag
        chosen_tag = random.choice(tags)

        # 获取该tag下的节点列表
        nodes_with_tag = self.tag_to_nodes.get(chosen_tag, [])
        if not nodes_with_tag:
            return None

        # 随机选择一个节点
        chosen_node = random.choice(nodes_with_tag)
        walk.append(chosen_node)
        return walk

    def simulate_walks(self, num_walks, walk_length, workers, output_file):
        """
        多进程模拟多次随机游走，并将游走结果保存到文件
        """
        G = self.graph
        nodes = list(G.nodes())
        num_walks_per_worker = max(1, num_walks // workers)
        logging.info(f"Starting simulation with {num_walks_per_worker} walks per node, walk length {walk_length}, using {workers} workers.")

        # 
        # results = Parallel(n_jobs=workers)(
        # results = Parallel(n_jobs=workers, backend='multiprocessing')(
        # results = Parallel(n_jobs=workers, backend='loky')(
        results = Parallel(n_jobs=workers)(
            delayed(self._simulate_walks)(nodes, num_walks_per_worker, walk_length)
            for _ in range(workers)
        )
        walks = list(itertools.chain(*results))

        # 保存游走结果到文件
        self.save_walks_to_file(walks, output_file)

    def _simulate_walks(self, nodes, num_walks, walk_length):
        
        """
        模拟多次随机游走
        """
        logging.info(f"_simulate_walks started, num_walks:{num_walks}, walk_length:{walk_length}")
        walks = []
        for i in range(num_walks):
            logging.info(f"_simulate_walks run num_walks of {i}.")
            random.shuffle(nodes)
            for node in nodes:
                walks.append(self.deepwalk_walk(walk_length=walk_length, start_node=node))
        return walks

    def save_walks_to_file(self, walks, output_file):
        """
        将游走结果保存到文件，按Word2Vec的输入格式
        """
        logging.info(f"Saving walks to file: {output_file}")
        with open(output_file, 'w') as f:
            for walk in walks:
                walk_str = ' '.join(map(str, walk))
                f.write(walk_str + '\n')
        logging.info(f"Successfully saved {len(walks)} walks to {output_file}.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run DeepWalk with tag-based random walks")
    parser.add_argument('--edge-file', type=str, required=True, help="Path to the edge file") # ../../fetch_data/data/edge.txt.20240923
    parser.add_argument('--node-tag-file', type=str, help="Path to the node-tag file")
    parser.add_argument('--num-walks', type=int, default=100, help="Number of walks per node (default: 10)")
    parser.add_argument('--walk-length', type=int, default=40, help="Length of each walk (default: 40)")
    parser.add_argument('--workers', type=int, default=multiprocessing.cpu_count() - 1, help="Number of workers (default: CPU cores - 1)")
    parser.add_argument('--use-softmax', action='store_true', help="Use softmax-based alias sampling (default: False)")
    parser.add_argument('--temperature', type=float, default=1.0, help="Temperature for softmax (default: 1.0)")
    parser.add_argument('--p-tag-walk', type=float, default=0.2, help="Probability to walk through tag-based neighbors (default: 0.5)")
    parser.add_argument('--output-file', type=str, required=True, help="Path to save the walks file")

    args = parser.parse_args()

    # 初始化日志记录
    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s', level=logging.INFO)

    # 初始化DeepWalk实例，传入边文件和节点标签文件
    deepwalk = DeepWalk(
        edge_file=args.edge_file,
        node_tag_file=args.node_tag_file,
        use_softmax=args.use_softmax,
        temperature=args.temperature,
        p_tag_walk=args.p_tag_walk
    )

    # 模拟随机游走并将结果保存到文件
    deepwalk.simulate_walks(
        num_walks=args.num_walks,
        walk_length=args.walk_length,
        workers=args.workers,
        output_file=args.output_file
    )