import torch
from args import args, config
from tqdm import tqdm
from items_dataset import items_dataset

def test_predict(test_loader, device, model, min_label=1, max_label=3):
  model.eval()
  result = []
  
  for i, test_batch in enumerate(tqdm(test_loader)):
    batch_text = test_batch['batch_text']
    input_ids = test_batch['input_ids'].to(device)
    token_type_ids = test_batch['token_type_ids'].to(device)
    attention_mask = test_batch['attention_mask'].to(device)
    #labels = test_batch['labels'].to(device)
    crf_mask = test_batch["crf_mask"].to(device)
    sample_mapping = test_batch["overflow_to_sample_mapping"]
    output = model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask, labels=None, crf_mask=crf_mask)
    if args.use_crf:
      prediction = model.crf.decode(output[0], crf_mask)
    else:
      prediction = torch.max(output[0], -1).indices

    #make result of every sample
    sample_id = -1
    sample_result= {"text_a" : test_batch['batch_text'][0]}
    for batch_id in range(len(sample_mapping)):
        change_sample = False
        if sample_id != sample_mapping[batch_id]: change_sample = True
        #print(i, id)
        if change_sample:
          sample_id = sample_mapping[batch_id]
          sample_result= {"text_a" : test_batch['batch_text'][sample_id]}
          decode_span_table = torch.zeros(len(test_batch['batch_text'][sample_id]))

        spans = items_dataset.cal_agreement_span(None, agreement_table=prediction[batch_id], min_agree=min_label, max_agree=max_label)
        #decode spans
        for span in spans:
            #print(span)
            if span[0]==0: span[0]+=1
            if span[1]==1: span[1]+=1

            while(True):
              start = test_batch[batch_id].token_to_chars(span[0])
              if start != None or span[0]>=span[1]:
                break
              span[0]+=1
              
            while(True):
              end = test_batch[batch_id].token_to_chars(span[1])
              if end != None or span[0]>=span[1]:
                break
              span[1]-=1

            if span[0]<span[1]:
              de_start = test_batch[batch_id].token_to_chars(span[0])[0]
              de_end = test_batch[batch_id].token_to_chars(span[1]-1)[0]
              #print(de_start, de_end)
              #if(de_start>512): print(de_start, de_end)
              decode_span_table[de_start:de_end]=2 #insite
              decode_span_table[de_start]=1 #begin
        if change_sample:
          sample_result["predict_span_table"] = decode_span_table
          #sample_result["boundary"] = test_batch["boundary"][id]
          result.append(sample_result)
  model.train()
  return result

def add_sentence_table(result, pattern =":；。,，？!～！： ", threshold_num=5, threshold_rate=0.5):
  for sample in result:
    boundary_list = []
    for i, char in enumerate(sample['text_a']):
      if char in pattern:
        boundary_list.append(i)
    boundary_list.append(len(sample['text_a'])+1)
    start=0
    end =0
    fist_sentence = True
    sample["predict_sentence_table"] = torch.zeros(len(sample["predict_span_table"]))
    for boundary in boundary_list:
      end = boundary
      predict_num = sum(sample["predict_span_table"][start:end]>0)
      sentence_num = len(sample["predict_span_table"][start:end])
      if(predict_num > threshold_num) or (predict_num > sentence_num*threshold_rate):
        if fist_sentence:
          sample["predict_sentence_table"][start:end] = 2
          sample["predict_sentence_table"][start] = 1
          fist_sentence=False
        else: 
          sample["predict_sentence_table"][start-1:end] = 2
      else: fist_sentence =True
      start = end+1

def add_doc_id(result, test_data):
  #make dict {'text_a':"docid"}
  text_to_id = dict()
  for sample in test_data:
    text_to_id[sample["text_a"]] = sample["docid"]
  
  #add doc_id
  for sample in result:
    sample["docid"] = text_to_id[sample["text_a"]]