import psycopg2
import codecs
import xlwt
import re
import os
import xlrd
import pandas as pd

def getData(t="final_label_money"):
    conn = psycopg2.connect(dbname="BiddingKM_test_10000",user="postgres",password="postgres",host="192.168.2.101")
    cursor = conn.cursor()
    
    
    select_sql = " select A.doc_id,C.entity_id,C.label,case when C.label=0 then '招标人' when C.label=1 then '招标代理' when C.label=2 then '中标人/第一候选' when C.label=3 then '第二' when C.label=4 then '第三' else '无' end as 再标注,case when D.label=0 then '招标人' when D.label=1 then '招标代理' when D.label=2 then '中标人/第一候选' when D.label=3 then '第二' when D.label=4 then '第三' else '无' end as 原标注,B.entity_text,A.tokens[B.begin_index-10:B.begin_index],A.tokens[B.begin_index+1:B.end_index+1],A.tokens[B.end_index+2:B.end_index+12] "
    
    group_sql = " group by A.doc_id,C.entity_id,C.label,D.label,B.entity_text,B.begin_index,B.end_index,A.tokens,A.sentence_index "
    
    sql = select_sql+" from sentences A,entity_mention B,"+t+" C,label_guest_role D  where A.doc_id=B.doc_id and A.sentence_index=B.sentence_index and B.entity_id=C.entity_id and C.entity_id=D.entity_id and C.label!=D.label "+group_sql+"order by A.doc_id,A.sentence_index asc,D.label asc"
    cursor.execute(sql)
    
    result = []
    rows = cursor.fetchall()
    for row in rows:
        item = []
        for column in row:
            item.append(column)
        result.append(item)
    conn.close()
    return result

def labeling(datas=getData()):
    sum = 0
    row_index = 0
    begin_doc_id = str(input("开始文章是："))
    begin_index = 0
    end_index = len(datas)-1
    find_flag = False
    while(row_index<len(datas)):
        row = datas[row_index]
        if begin_doc_id!="" and begin_index==0:
            if row[0]==begin_doc_id:
                begin_index = row_index
            else:
                row_index += 1
                continue
        find_flag = True
        print(row[0])
        print(row[3],row[4],row[5])
        print("before",row[6])
        print("entity",row[7])
        print("after",row[8])
        while(True):
            l = str(input("标签为："))
            if l in ["0","1","2","3","4","5","","8","9"]:
                break
        if l=="0":
            row[2] = 0
        elif l=="1":
            row[2] = 1
        elif l=="2":
            row[2] = 2
        elif l=="3":
            row[2] = 3
        elif l=="4":
            row[2] = 4
        elif l=="5":
            row[2] = 5
        elif l=="":
            pass
        elif l=="8":
            row_index -= 1
            sum -= 1
            continue
        elif l=="9":
            end_index = row_index-1
            break
            
        sum += 1
        row_index += 1
        print("sum:",sum)
    if find_flag:
        with codecs.open("relabel.txt","a",encoding="utf8") as f:
            for row in datas[begin_index:end_index+1]:
                f.write(str(row[1]))
                f.write("\t")
                f.write(str(row[2]))
                f.write("\n")
            f.flush()
            f.close()
#设置表格样式
def set_style(name,height,bold=False):
    style = xlwt.XFStyle()
    font = xlwt.Font()
    font.name = name
    font.bold = bold
    font.color_index = 4
    font.height = height
    style.font = font
    return style

#将数据保存到excel中进行标注
def getDatasToExcel():
    moneys = ["0_预算金额","1_中投标","2_无(非预算及中投标金额)"]
    conn = psycopg2.connect(dbname="BiddingKM_test_10000",user="postgres",password="postgres",host="192.168.2.101")
    cursor = conn.cursor()
    
    nums = 3
    
    for money in moneys:
        
        
        select_sql = " select C.entity_id,C.label,A.tokens[B.begin_index-10:B.begin_index],A.tokens[B.begin_index+1:B.end_index+1],A.tokens[B.end_index+2:B.end_index+12],case when C.label=0 then '预算金额' when C.label=1 then '中/投标' else '无(非预算及中投标金额)' end as 再标注 "
        
        sql = select_sql+" from sentences A,entity_mention_copy B,hand_label_money C  where B.entity_type='money' and C.label="+money.split("_")[0]+" and A.doc_id=B.doc_id and A.sentence_index=B.sentence_index and B.entity_id=C.entity_id order by C.label asc,A.doc_id,A.sentence_index asc limit 7000"
        
        pattern_tenderee = re.compile("报价上限|限价|造价|控制(总?价|金额)|预算|概算|(?:造?价|投资|规模)预?估?算|预?估算?(?:造?价|投资|规模|金额)|(?:总|项目|计划)(?:[估预概]算|投资)|(?:投资|采购)(?:单价|总)?(?:额|金额)|投资约")
        pattern_wintenderer = re.compile("[\(（]?(?:中标|成交|评标|评审|投标|报价|合同|入围)[\),）]?(候选人|单位|人|候选单位|供应商|候选)?后?[单总]?(?:价|金额|价格|报价|标价)|报价|第[一二三](中标|投标|候选|名|成交)|(公司\s*[:：]?$)")
    
        
        print(sql)
        cursor.execute(sql)
        
        rows = cursor.fetchall()
        
        parts = len(rows)//3
        
        for nums_i in range(nums):
            file = xlwt.Workbook()
            sheet = file.add_sheet("标注"+money,cell_overwrite_ok=True)
            row_head = ["entity_id","标注id","实体前","实体","实体后","角色","找茬-错误标0(正确留空)"]
            row_index = 0
            style = set_style('Times New Roman',220,True)
            for i in range(len(row_head)):
                sheet.write(row_index,i,row_head[i],style)
                
            row_index += 1
            if nums_i<nums-1:
        
                for row in rows[nums_i*parts:(nums_i+1)*parts]:
                    for i in range(len(row)):
                        sheet.write(row_index,i,row[i],style)
                    row_index += 1
        
            else:
                for row in rows[nums_i*parts:]:
                    for i in range(len(row)):
                        sheet.write(row_index,i,row[i],style)
                    row_index += 1
        
            file.save("标注"+money.split("_")[1]+str(nums_i)+".xls")
    conn.close()
    
def getDatasToExcel_centainity():
    moneys = ["0_预算金额","1_中投标","2_无(非预算及中投标金额)"]
    conn = psycopg2.connect(dbname="BiddingKM_test_10000",user="postgres",password="postgres",host="192.168.2.101")
    cursor = conn.cursor()
    
    nums = 3
    
    for money in moneys:
        
        
        select_sql = " select C.entity_id,C.label,A.tokens[B.begin_index-10:B.begin_index],A.tokens[B.begin_index+1:B.end_index+1],A.tokens[B.end_index+2:B.end_index+12],case when C.label=0 then '预算金额' when C.label=1 then '中/投标' else '无(非预算及中投标金额)' end as 再标注 "
        
        sql = select_sql+" from sentences A,entity_mention B,final_label_money C  where B.entity_type='money' and C.label="+money.split("_")[0]+" and A.doc_id=B.doc_id and A.sentence_index=B.sentence_index and B.entity_id=C.entity_id order by C.label asc,A.doc_id,A.sentence_index asc limit 7000"
            
        
        print(sql)
        cursor.execute(sql)
        
        rows = cursor.fetchall()
        
        parts = len(rows)//3
        
        for nums_i in range(nums):
            file = xlwt.Workbook()
            sheet = file.add_sheet("标注"+money,cell_overwrite_ok=True)
            row_head = ["entity_id","标注id","实体前","实体","实体后","角色","找茬-错误标0(正确留空)"]
            row_index = 0
            style = set_style('Times New Roman',220,True)
            for i in range(len(row_head)):
                sheet.write(row_index,i,row_head[i],style)
                
            row_index += 1
            if nums_i<nums-1:
        
                for row in rows[nums_i*parts:(nums_i+1)*parts]:
                    for i in range(len(row)):
                        sheet.write(row_index,i,row[i],style)
                    row_index += 1
        
            else:
                for row in rows[nums_i*parts:]:
                    for i in range(len(row)):
                        sheet.write(row_index,i,row[i],style)
                    row_index += 1
        
            file.save("标注"+money.split("_")[1]+str(nums_i)+".xls")
    conn.close()
    
#从excel中获取标注数据
def getDatasFromExcel():
    home = "./label_done/"
    col_entity_id = 0
    col_label = 1
    col_flag = 6
    table = "hand_label_money"
    conn = psycopg2.connect(dbname="BiddingKM_test_10000",user="postgres",password="postgres",host="192.168.2.101")
    cursor = conn.cursor()
    
    cursor.execute(" select to_regclass('"+table+"') is null ")
    notExists = cursor.fetchall()[0][0]
    '''
    if notExists:
        cursor.execute(" create table "+table+" (entity_id text,label int)")
    else:
        cursor.execute(" delete from "+table)
    conn.commit()
    '''
    
    for file in os.listdir(home):
        if os.path.isfile(home+file):
            book = xlrd.open_workbook(home+file)
            sheet = book.sheet_by_index(0)
            for row_index in range(1,sheet.nrows):
                print(row_index,file)
                if len(re.sub('\s*','',str(sheet.cell_value(row_index,col_flag))))==0 or (int(sheet.cell_value(row_index,col_flag))>0):
                    #sql = " insert into "+table+"(entity_id,label) values('"+str(sheet.cell_value(row_index,col_entity_id))+"',"+str(int(sheet.cell_value(row_index,col_label)))+")"
                    #cursor.execute(sql)
                    print(str(sheet.cell_value(row_index,col_entity_id))+"',"+str(int(sheet.cell_value(row_index,col_label))))
    conn.commit()
    conn.close()
    
def getHandLabelData():
        
    def spanWindow(tokens,begin_index,end_index,size):
        '''
        @summary:取得某个实体的上下文词汇
        @param:
            tokens:句子分词list
            begin_index:实体的开始index
            end_index:实体的结束index
            size:左右两边各取多少个词
        @return: list,实体的上下文词汇
        '''  
        length_tokens = len(tokens)
        if begin_index>size:
            begin = begin_index-size
        else:
            begin = 0
        if end_index+size<length_tokens:
            end = end_index+size+1
        else:
            end = length_tokens
        result = []
        result.append(tokens[begin:begin_index])
        result.append(tokens[begin_index:end_index+1])
        result.append(tokens[end_index+1:end])
        #print(result)
        return result
    
    conn = psycopg2.connect(dbname="BiddingKM_test_10000",user="postgres",password="postgres",host="192.168.2.101")
    cursor = conn.cursor()
    
    sql = '''
     select A.entity_id,A.entity_text,A.begin_index,A.end_index,C.label,B.tokens 
     from entity_mention_copy A,sentences B,hand_label_money C 
     where A.doc_id=B.doc_id and A.sentence_index=B.sentence_index
     and A.entity_type in ('money')
     and A.entity_id=C.entity_id
      order by C.label 
    '''
    
    cursor.execute(sql)
    rows = cursor.fetchall()
    
    list_entity_id = []
    list_before = []
    list_after = []
    list_text = []
    list_label = []
    repeat = set()
    for row in rows:
        entity_id = row[0]
        #entity_text = row[1]
        begin_index = row[2]
        end_index = row[3]
        label = int(row[4])
        tokens = row[5]
        
        beforeafter = spanWindow(tokens, begin_index, end_index, 10)
        
        if ("".join(beforeafter[0]),"".join(beforeafter[1]),"".join(beforeafter[2])) in repeat:
            continue
        
        repeat.add(("".join(beforeafter[0]),"".join(beforeafter[1]),"".join(beforeafter[2])))
        
        list_entity_id.append(entity_id)
        list_before.append("".join(beforeafter[0]))
        list_after.append("".join(beforeafter[2]))
        list_text.append("".join(beforeafter[1]))
        list_label.append(label)
    print("len",len(list_entity_id))
    parts = 1
    parts_num = len(list_entity_id)//parts
    for i in range(parts-1):
        
        data = {"entity_id":list_entity_id[i*parts_num:(i+1)*parts_num],"list_before":list_before[i*parts_num:(i+1)*parts_num],"list_after":list_after[i*parts_num:(i+1)*parts_num],"list_text":list_text[i*parts_num:(i+1)*parts_num],"list_label":list_label[i*parts_num:(i+1)*parts_num]}
        df = pd.DataFrame(data)
        df.to_excel("未标注错误_"+str(i)+".xls",columns=["entity_id","list_before","list_text","list_after","list_label","list_prob"])
    i = parts - 1
    data = {"entity_id":list_entity_id[i*parts_num:],"list_before":list_before[i*parts_num:],"list_after":list_after[i*parts_num:],"list_text":list_text[i*parts_num:],"list_label":list_label[i*parts_num:]}
    df = pd.DataFrame(data)
    df.to_excel("wrong_money_"+str(i)+".xls",columns=["entity_id","list_before","list_text","list_after","list_label","list_prob"])

    
if __name__=="__main__":
    #labeling()
    #getDatasToExcel()
    #getDatasFromExcel()
    getHandLabelData()