# -*- coding: utf-8 -*-
import pandas as pd
import json

# 加载地质描述文本，提取prompt和label
with open('F:/GeoLLM/data/train_triples.json', 'r', encoding='utf-8') as f:
    data = json.load(f)
# 将data转换为DataFrame
df = pd.DataFrame(data)
# 提取prompt和label
text = df['text']
label = df['triple_list']
# print(prompt,label)

print(len(text))


from response_to_json import parse_llm_response, save_to_json, save_raw_response
from LLM import zero_shot
# from prompt_generate import generate_prompt_with_examples as generate_prompt
# from prompt_generate import generate_prompt_with_best_matches as generate_prompt_b

model_series = 'gpt'
# model_name = 'deepseek-ai/DeepSeek-V3'
model_name='deepseek-ai/DeepSeek-R1'

prompt = '''
你是一名专业经验丰富的工程地质领域专家，你的任务是从给定的输入文本中提取"实体-关系-实体"三元组。关系类型包括24种："出露于"、"位于"、"整合接触"、"不整合接触"、"假整合接触"、"断层接触"、"分布形态"、"大地构造位置"、"地层区划"、"出露地层"、"岩性"、"厚度"、"面积"、"坐标"、"长度"、"含有"、"所属年代"、"行政区划"、"发育"、"古生物"、"海拔"、"属于"、"吞噬"、"侵入"。提取过程请按照以下规范：
1.	输出格式：
严格遵循JSON数组，无额外文本，每个元素包含：
[
    {
        "entity1": "实体1",
        "relation": "关系",
        "entity2": "实体2"
    }
]
2.	复杂关系处理：
- 若同一实体参与多个关系，需分别列出不同三元组
3.	部分专业关系判断指引：
- 出露于：地质单元（如岩层、地层或构造体）在地表或接近地表的显露状况。
- 整合接触：同一地区的上下两套岩层，产状一致，并且在沉积上和在生物演化上都是连续的。
- 不整合接触：同一地区的上下两套岩层间有明显的沉积间断，且在古生物演化顺序上也不连续。
- 假整合接触：又称平行不整合接触，指具有相同产状的上下两套地层之间有一明显的沉积间断或缺失。
- 断层接触：一种构造接触，即侵入岩体与围岩间的界面就是断层面或断层带。
- 分布形态：地质体在空间中的排列和呈现方式。通常包括几何特征、沿某方向的延伸、断裂、弯曲等形态学特征。
- 大地构造位置：地质体或地质现象在全球或区域大地构造体系中的相对位置。
- 地层区划：根据地壳活动性、古地理条件、古气候条件和古生物群特征等因素,在各个地区的情况不一致所导致的地层发育上的差异划分出不同的地层区域。
- 出露地层：地表或地下露出的岩石或沉积物层系。
- 岩性：反映岩石特征的一些属性，如颜色、成分、结构、胶结物及胶结类型、特殊矿物等
- 所属年代：指的是岩石、岩石或其他地质现象所对应的地质历史时期。
- 发育：地质体在地质条件控制下经历形成、演化并达到特定形态或分布特征的过程。
- 古生物：泛指地质历史时期中曾经生活于地球上的生物。一般将更新世以前的生物称为古生物。
- 吞噬：岩浆同化捕虏体现象。
- 侵入：深部岩浆在向上运移过程中，侵入到周围岩石而未到达地表的过程。
'''
j=0
q=0

json_path = 'F:/GeoLLM/output/Knowledge-guided/one_shot/'+model_name+'.json'
j=len(json.load(open(json_path,'r',encoding='utf-8')))
q=len(json.load(open('F:/GeoLLM/output/Knowledge-guided/one_shot_raw/'+model_name+'.json','r',encoding='utf-8')))
# # 当q=j时才继续处理
if q==j:        
    print(j)
    for i in range(j,500):
        response = zero_shot(model_series, model_name, prompt+'\n'+text[i])
        # print(prompt+'\n'+text[i])
        # 解析响应
        formatted_triples = parse_llm_response(response)
        print(i)            
        # 保存结果  
        save_to_json(text[i], formatted_triples, model_series=model_name, output_dir='F:/GeoLLM/output/Knowledge-guided/one_shot/')
        # 保存原始响应为josn文件save_raw_response
        save_raw_response(response, text[i], model_series=model_name, output_dir='F:/GeoLLM/output/Knowledge-guided/one_shot_raw/')

else:
    print('q!=j')