from openai import OpenAI
from pathlib import Path
import os
import re
import json
import glob
import tqdm
from multiprocessing import Pool
from functools import partial
from collections import Counter

API_KEY = "sk-oYh3Xrhg8oDY2gW02c966f31C84449Ad86F9Cd9dF6E64a8d"
BASE_URL = "https://vip.apiyi.com/v1"
MODEL_GPT = "gpt-4o-mini"

# 确保输出为标准json格式字符串
def comfirm_json_string_gpt(json_string):
    client = OpenAI(api_key=API_KEY, base_url=BASE_URL)

    prompt = f"""
        You will read a <string>, please fix this string into a string that can be parsed by json.loads.

        Note:
        1. No descriptive text is required.
        2. Don't use markdown syntax.

        The <string>: {json_string}
    """

    response = client.chat.completions.create(
        model=MODEL_GPT,
        messages=[
            {"role": "system", "content": "You are an assistant who is proficient in material synthesis."},
            {"role": "user", "content": prompt}
        ]
    )

    return response.choices[0].message.content

# 确保输出为标准json格式字符串
def comfirm_json_string(json_string):
    json_string = re.sub(r'[“”]', '"', json_string)
    json_string = re.sub(r'\\', r'\\\\', json_string)
    json_string = re.sub(r'\\"', r'\"', json_string)
    json_string = json_string.replace("\n", "").replace("\r", "")
    # 去掉 Markdown 的语法包裹
    if json_string.startswith("```json"):
        json_string = json_string.strip("`json\n")
    json_string = json_string.strip('`\n')

    return json_string

# 文本分割
def split_by_heading(markdown_text, heading_level='#'):
    # `heading_level` could be '#', '##', '###', etc.
    pattern = r'(?=\n{})'.format(re.escape(heading_level))
    
    # 使用正则表达式进行切割，以包含标题的内容
    split_texts = re.split(pattern, markdown_text)
    
    # 去除空白的块
    return [block.strip() for block in split_texts if block.strip()]

# 文本段分类
def segment_classification(text_split):
    client = OpenAI(api_key=API_KEY, base_url=BASE_URL)

    prompt = f"""
        You will read a text segment about hydrophilic polymers. Please analyze which part of a paper this segment belongs to and give your classification result. The categories you can only choose are as follows:
        1. Abstract
        2. Introduction
        3. Materials and methods
        4. Results and discussion
        5. Conclusions
        6. References

        Please output the result using the following format:
        Category: Abstract/Introduction/Materials and methods/Results and discussion/Conclusions/References 

        Text segment as follows: {text_split}
    """

    response = client.chat.completions.create(
        model=MODEL_GPT,
        messages=[
            {"role": "system", "content": "You are an expert in interdisciplinary research involving materials chemistry, surface and interface science, and the functionalization of nanomaterials."},
            {"role": "user", "content": prompt}
        ]
    )

    return response.choices[0].message.content

# 处理单个md文件
def process_file(md_path, output_dir):
    chunks = []
    with open(md_path, 'r', encoding='utf-8') as file:
        md_content = file.read()

    # 将文本按heading分割
    content_splits = split_by_heading(md_content)

    id = 0
    for content_split in content_splits:
        id += 1
        chunk = {}
        result = segment_classification(content_split)
        chunk["id"] = id
        chunk["chunk"] = content_split
        chunk["category"] = result[9:]
        chunks.append(chunk)

    output_path = os.path.join(output_dir, os.path.basename(md_path).replace('.md', '.json'))
    with open(output_path, 'w', encoding='utf-8') as json_file:
        json.dump(chunks, json_file, ensure_ascii=False, indent=4)

# 获取已经处理过的md
def chunk_done(json_dir):
    jsons = os.listdir(json_dir)
    json_names = [json_name.replace('.json', '') for json_name in jsons]
    return json_names

# 将文本段分割分类并保存为json
def md_segment():
    md_paths = glob.glob("/home/ubuntu/50T/fsy/wl/articles/mds/**/*.md", recursive=True)
    print("md文件数量：", len(md_paths))
    # 过滤已经过处理的文件
    output_dir = "/home/ubuntu/50T/fsy/wl/task2-chunks"
    json_names = chunk_done(output_dir)
    md_paths = [md_path for md_path in md_paths if os.path.basename(md_path).replace(".md", "") not in json_names]
    print("过滤后md文件数量：", len(md_paths))

    for path in tqdm.tqdm(md_paths):
        try:
            process_file(path, output_dir)
        except Exception as e:
            print(f"处理 {path} 时出错: {e}")

    # 设置多进程池
    # pool = Pool(processes=32)

    # process_func = partial(process_file, output_dir=output_dir)

    # # imap_unordered 将逐步从 md_paths 传给 process_func 进行并行处理
    # for _ in tqdm.tqdm(pool.imap_unordered(process_func, md_paths), total=len(md_paths)):
    #     pass

    # pool.close()
    # pool.join()

# 提取合成方案
def get_protocol(text):
    client = OpenAI(api_key=API_KEY, base_url=BASE_URL)

    prompt = f"""
        You will read a text excerpt from an article on hydrophilic polymer synthesis. Please extract the specific experimental protocol for synthesis of hydrophilic polymers with functional groups.

        Note:
        1. The information you extract must come from the text excerpt(example not included), and fabrication of information is strictly prohibited.
        2. Don't use markdown syntax.

        Please output the result using the following format:
        {{
            "protocol": ""
        }}

        The text except: {text}
    """

    response = client.chat.completions.create(
        model=MODEL_GPT,
        messages=[
             {"role": "system", "content": "You are an expert in developing hydrophilic polymers for applications such as biomedical hydrogels or water filtration membranes."},
            # {"role": "system", "content": "You are an expert in researching surface modification of black phosphorus."},
            {"role": "user", "content": prompt}
        ]
    )

    return  response.choices[0].message.content

# 提取实验方案
def extract_info(chunks_path):
    with open(chunks_path, 'r', encoding='utf-8') as file:
        chunks = json.load(file)

    protocol_dict = {"protocol" : ""}  # 存放最终输出
    for chunk in chunks:
        chunk_text = chunk['chunk']
        category = chunk['category']
        try:
            # 提取分子做亲水性聚合物合成的具体实验步骤
            if category == ' Introduction' or category == ' Materials and methods':
                intermediate_result = get_protocol(chunk_text)
                intermediate_result = comfirm_json_string(intermediate_result)
                try:
                    result_protocol = json.loads(intermediate_result)
                except json.JSONDecodeError as e:
                    # 修复json字符串(gpt)
                    escaped_protocol = comfirm_json_string_gpt(intermediate_result)
                    try:
                        result_protocol = json.loads(escaped_protocol)
                    except Exception as e:
                        print(e)
                        print(escaped_protocol)
                        return
                if result_protocol == "":
                    continue
                if isinstance(result_protocol, dict):
                    protocol_dict["protocol"] += result_protocol["protocol"]
                else:
                    print("result_protocol不是一个字典")
                    print(result_protocol)
                    return
        except Exception as e:
            print(e)
            return

    output_path = os.path.join(output_dir, os.path.basename(chunks_path))
    with open(output_path, 'w') as json_file:
        json.dump(protocol_dict, json_file, ensure_ascii=False, indent=4)

if __name__ == '__main__':
    chunks_dir = "/home/ubuntu/50T/fsy/wl/task2-chunks"
    paths = [os.path.join(chunks_dir, path) for path in os.listdir(chunks_dir)]
    print("chunks文件数量：", len(paths))

    # 过滤已处理的文件
    output_dir = "/home/ubuntu/50T/fsy/wl/task2-paper-info"
    proccessed_files = [path for path in os.listdir(output_dir)]
    paths = [path for path in paths if os.path.basename(path) not in proccessed_files]
    print("过滤后chunks文件数量：", len(paths))

    # md_segment()

    for path in tqdm.tqdm(paths):
        try:
            extract_info(path)
        except Exception as e:
            print(f"处理 {path} 时出错: {e}")

    # # 设置多进程池
    # pool = Pool(processes=32)

    # process_func = partial(extract_info)

    # # imap_unordered 将逐步从 md_paths 传给 process_func 进行并行处理
    # for _ in tqdm.tqdm(pool.imap_unordered(process_func, paths), total=len(paths)):
    #     pass

    # pool.close()
    # pool.join()