from subgraph import get_subgraph_edge,get_subgraph_node
from util import load_json_file, matrix_visualize
import time
import numpy as np
from tqdm import tqdm
from gen_static_distance_weight import load_static

import argparse
import json
import os

static_probability_matrix=[]

def get_probability_matrix(w_weight):
    n=len(w_weight)
    probability_matrix=np.zeros((n,n))
    for i in range(n):
        for j in range(n):
            probability_matrix[i][j]=w_weight[i]*w_weight[j]
    return probability_matrix

def build_graph(li):

    data=eval(li[0])
    w_tfidf_weight=eval(li[1])
    w_bert_weight=eval(li[2])
    w_dynamic_weight=eval(li[3])

    n=len(data['psg'])
    probability_matrix=np.zeros((n,n))

    strategy_cnt=0
    if args.strategy[0]=='1':probability_matrix1=get_probability_matrix(w_bert_weight[0:n]) ;strategy_cnt+=1
    if args.strategy[1]=='1':probability_matrix2=get_probability_matrix(w_dynamic_weight);strategy_cnt+=1
    if args.strategy[2]=='1':probability_matrix3=get_probability_matrix(w_tfidf_weight);strategy_cnt+=1
    if args.strategy[3]=='1':probability_matrix4=static_probability_matrix[n];strategy_cnt+=1

    #if strategy_cnt!=0:probability_matrix/=strategy_cnt
    probability_matrix = 0.2 * probability_matrix1 + 0.3 * probability_matrix2 + 0.2 * probability_matrix3 + 0.3 * probability_matrix4

    # make sum(probability_matrix)/n/n == args.sparsity
    probability_matrix*=args.sparsity*n*n/np.sum(probability_matrix)

    # build graph
    graph=np.zeros((n,n))
    for i in range(n):
        for j in range(0,i+1):
            coin=np.random.rand(1)
            if 0<=coin<=probability_matrix[i][j]:
                graph[i][j]=graph[j][i]=1
    
    matrix_visualize("graph",graph)

    # print("Actual sparsity: ",np.sum(graph)/n/n)
    
    subgraph1=get_subgraph_edge(data['psg'],graph.copy(),args.subgraph_number,max_subgraph_node=128)
    subgraph2=get_subgraph_node(data['psg'],graph.copy(),args.subgraph_number,max_subgraph_node=128)

    # concat the original passages slice
    ori_slice=[]
    for j in range((n+127)//128):
        if j*128+128 < n: en=j*128+128
        else: en=n
        ori_slice.append(data['psg'][j*128:en])
    while(len(ori_slice)<16):ori_slice.append([])

    ans={'qry':data['qry'],'psg1':ori_slice+subgraph1,'psg2':ori_slice+subgraph2,'label':data['label']}

    return ans


if __name__=="__main__":
    parser=argparse.ArgumentParser()
    parser.add_argument("--sparsity",
                        default=1 - 0.93,
                        type=float,
                        help="Sparsity of the social network.")
    parser.add_argument("--strategy",
                        default="1111",
                        type=str,
                        help="Whether to use some strategy or not, the order is bert,dynamic,tfidf,static")
    parser.add_argument("--subgraph_number",
                        default=8,
                        type=int,
                        help="The number of subgraph to choose")
    parser.add_argument("--write_path",
                        default="results/all.txt",
                        type=str)
    parser.add_argument("--read_path",
                        default="weights/total.txt",
                        type=str,
                        help="The path of file generated by paste.sh")
    args = parser.parse_args()
    
    # Because calculating static weight is time-cosuming, we calculate static weight in advance
    print("----------------loading static weight-------------------")
    static_probability_matrix.append(0)
    for i in tqdm(range(1,2049),leave=False):
        static_probability_matrix.append(load_static(i,root_path="weights/static/"))

    start=time.time()
    with open(args.read_path,"r") as fr:
        with open(args.write_path,"w") as fw:   
            start=time.time()  
            for li in tqdm(fr): 
                ans=build_graph(li.split("$"))
                json_str = json.dumps(ans)
                fw.write(json_str)           
                print(file=fw)
