gdm.iterate_model(user_corrections='''
Let's think step by step. Please make the following updates to the data model:
1. Remove the relationships between Patient and Disease, between Patient and Symptom and between Patient and Outcome.
2. Change the Patient node into Demographics.
3. Create a relationship HAS_DEMOGRAPHICS from Disease to Demographics.
4. Create a relationship HAS_SYMPTOM from Disease to Symptom. If the Symptom value is No, remove this relationship.
5. Create a relationship HAS_LAB from Disease to HealthIndicator.
6. Create a relationship HAS_OUTCOME from Disease to Outcome.
''')

from IPython.display import Image, display
gdm.current_model.visualize().render('output', format='png')
# Load and display the image with a specific width
img = Image('output.png', width=1200)  # Adjust the width as needed
display(img

# instantiate ingestion generator
gen = IngestionGenerator(data_model=gdm.current_model,
                        username="neo4j",
                        password='yourneo4jpasswordhere',
                        uri='neo4j+s://123654888.databases.neo4j.io',
                        database="neo4j",
                        csv_dir="/home/user/",
                        csv_name="disease_prepared.csv")

# create ingestion YAML
pyingest_yaml = gen.generate_pyingest_yaml_string()

# save local copy of YAML
gen.generate_pyingest_yaml_file(file_name="disease_prepared")

import warnings
import json
from langchain_community.graphsimportNeo4jGraph

with warnings.catch_warnings():
   warnings.simplefilter('ignore')

NEO4J_USERNAME = "neo4j"
NEO4J_DATABASE = 'neo4j'
NEO4J_URI = 'neo4j+s://1236547.databases.neo4j.io'
NEO4J_PASSWORD = 'yourneo4jdatabasepasswordhere'

prompt_template = """
Let's think step by step:

Step1: Task:
Generate an effective and concise Cypher statement with less than 256 characteres to query a graph database
Do not comment the code.

Step 2: Get to know the database schema: {schema}

Step 3: Instructions:
- In the cypher query, ONLY USE the provided relationship types and properties that appear in the schema AND in the user question.
- In the cypher query, do not use any other relationship types or properties in the user's question that are not contained in the provided schema.
- Regarding Age, NEVER work with the age itself. For example: 24 years old, use interval: more than 20 years old.
- USE ONLY ONE statement for Age, always use 'greater than', never 'less than' or 'equal'.
- DO NOT USE property keys that are not in the database.

Step 4: Examples:
Here are a few examples of generated Cypher statements for particular questions:

4.1 Which diseases present high blood pressure?
MATCH (d:Disease)
MATCH (d)-[r:HAS_LAB]->(l)
WHERE l.bloodPressure = 'High'
RETURN d.name

4.2 Which diseases present indicators as high blood pressure?
// Match the Disease nodes
MATCH (d:Disease)
// Match HAS_LAB relationships from Disease nodes to Lab nodes
MATCH (d)-[r:HAS_LAB]->(l)
MATCH (d)-[r2:HAS_OUTCOME]->(o)
// Ensure the Lab nodes have the bloodPressure property set to 'High'
WHERE l.bloodPressure = 'High' AND o.result='Positive'
RETURN d, properties(d) AS disease_properties, r, properties(r) AS relationship_properties, l, properties(l) AS lab_properties

4.3 What is the name of a disease of the elderly where the patient presents high blood pressure, high cholesterol, fever, fatigue
MATCH (d:Disease)
MATCH (d)-[r1:HAS_LAB]->(lab)
MATCH (d)-[r2:HAS_SYMPTOM]->(symptom)
MATCH (symptom)-[r3:HAS_DEMOGRAPHICS]->(demo)
WHERE lab.bloodPressure = 'High' AND lab.cholesterolLevel = 'High' AND symptom.fever = 'Yes' AND symptom.fatigue = 'Yes' AND TOINTEGER(demo.age) >40
RETURN d.name

4.4 What disease gives you fever, fatigue, no cough, no short breathe in people with high cholesterol?
MATCH (d:Disease)-[r:HAS_SYMPTOM]->(s:Symptom)
WHERE s.fever = 'Yes' AND s.fatigue = 'Yes' AND s.difficultyBreathing = 'No' AND s.cough = 'No'
MATCH (d:Disease)-[r1:HAS_LAB]->(lab:HealthIndicator)
MATCH (d)-[r2:HAS_OUTCOME]->(o:Outcome)
WHERE lab.cholesterolLevel='High' AND o.result='Positive'
RETURN d, properties(d) AS disease_properties, r, properties(r) AS relationship_properties


Step 5. These are the values allowed for each entity:
- Fever: Indicates whether the patient has a fever (Yes/No).
- Cough: Indicates whether the patient has a cough (Yes/No).
- Fatigue: Indicates whether the patient experiences fatigue (Yes/No).
- Difficulty Breathing': 'Indicates whether the patient has difficulty breathing (Yes/No).
- Age: The age of the patient in years.
- Gender: The gender of the patient (Male/Female).
- Blood Pressure: The blood pressure level of the patient (Normal/High).
- Cholesterol Level: The cholesterol level of the patient (Normal/High).
- Outcome Variable: The outcome variable indicating the result of the diagnosis or assessment for the specific disease (Positive/Negative).

Step 6. Answer the question {question}."""

prompt_template = """

让我们一步一步来：

步骤1：任务：

生成一个有效且简洁的 Cypher 语句，字符数不超过 256 个，用于查询图数据库。

不要对代码进行注释。

步骤2：了解数据库模式：{schema}

步骤3：指令：

- 在 Cypher 查询中，仅使用模式和用户问题中提供的关系类型和属性。

- 在 Cypher 查询中，不要使用用户问题中未包含在提供的模式中的其他关系类型或属性。

- 关于年龄，绝不要直接使用具体年龄。例如：24 岁，使用区间：大于 20 岁。

- 对于年龄，只使用一个语句，总是使用“大于”，从不使用“小于”或“等于”。

- 不要使用数据库中不存在的属性键。

步骤4：示例：

以下是针对特定问题生成的 Cypher 语句示例：

4.1 哪些疾病表现出高血压？

MATCH (d:Disease)

MATCH (d)-[r:HAS_LAB]->(l)

WHERE l.bloodPressure = 'High'

RETURN d.name

4.2 哪些疾病表现出高血压等指标？

// 匹配疾病节点

MATCH (d:Disease)

// 匹配从疾病节点到实验室节点的 HAS_LAB 关系

MATCH (d)-[r:HAS_LAB]->(l)

MATCH (d)-[r2:HAS_OUTCOME]->(o)

// 确保实验室节点的 bloodPressure 属性设置为 'High'

WHERE l.bloodPressure = 'High' AND o.result='Positive'

RETURN d, properties(d) AS disease_properties, r, properties(r) AS relationship_properties, l, properties(l) AS lab_properties

4.3 老年人患有高血压、高胆固醇、发烧、疲劳的疾病名称是什么？

MATCH (d:Disease)

MATCH (d)-[r1:HAS_LAB]->(lab)

MATCH (d)-[r2:HAS_SYMPTOM]->(symptom)

MATCH (symptom)-[r3:HAS_DEMOGRAPHICS]->(demo)

WHERE lab.bloodPressure = 'High' AND lab.cholesterolLevel = 'High' AND symptom.fever = 'Yes' AND symptom.fatigue = 'Yes' AND TOINTEGER(demo.age) >40

RETURN d.name

4.4 哪种疾病在高胆固醇人群中表现为发烧、疲劳、无咳嗽、无呼吸困难？

MATCH (d:Disease)-[r:HAS_SYMPTOM]->(s:Symptom)

WHERE s.fever = 'Yes' AND s.fatigue = 'Yes' AND s.difficultyBreathing = 'No' AND s.cough = 'No'

MATCH (d:Disease)-[r1:HAS_LAB]->(lab:HealthIndicator)

MATCH (d)-[r2:HAS_OUTCOME]->(o:Outcome)

WHERE lab.cholesterolLevel='High' AND o.result='Positive'

RETURN d, properties(d) AS disease_properties, r, properties(r) AS relationship_properties

步骤5：每个实体允许的值如下：

- 发烧：表示患者是否发烧（是/否）。

- 咳嗽：表示患者是否咳嗽（是/否）。

- 疲劳：表示患者是否感到疲劳（是/否）。

- 呼吸困难：表示患者是否呼吸困难（是/否）。

- 年龄：患者的年龄（以年为单位）。

- 性别：患者的性别（男/女）。

- 血压：患者的血压水平（正常/高）。

- 胆固醇水平：患者的胆固醇水平（正常/高）。

- 结果变量：表示特定疾病诊断或评估结果的变量（阳性/阴性）。

步骤6：回答问题 {question}。

"""

cypher_prompt = PromptTemplate(
   input_variables=["schema","question"],
   template=prompt_template
)

cypherChain = GraphCypherQAChain.from_llm(
   VertexAI(temperature=0.1),
   graph=kg,
   verbose=True,
   cypher_prompt=cypher_prompt,
   top_k=10# this can be adjusted also
)

cypherChain.run("""
Patient Information:
Jane Doe, a 58-year-old female, was admitted on June 15, 2024.

Chief Complaint and History of Present Illness:
Jane reported a high fever up to 104°F, body pain, and a rash,
starting five days prior to admission.

Past Medical History:
Jane has no significant past medical history and no known allergies.

Physical Examination:
Jane's temperature was 102.8°F, heart rate 110 bpm, blood pressure 100/70 mmHg, and respiratory rate 20 breaths
per minute. No petechiae or purpura were noted.

What disease may she have?""")