"""
Reformat the ecg report json data
- change metrics (variables) to descriptive names
- add "Finding"
"""

import sys
import csv
import json
import time
import g4f

llm="gpt-4o"

def generate(inputs, conclusion, model=llm):
    # find models at https://github.com/xtekky/gpt4free/blob/main/g4f/models.py

    template = """
        **Task: Drafting Clinical Reports**

        Your task is to draft comprehensive clinical reports, encompassing Clinical Findings and a subsequent Final Interpretation.

        **Example Clinical Findings:**

        "The patient's cardiac monitoring showed a minimum heart rate of 29 bpm, a maximum of 212 bpm, and an average of 72 bpm. The predominant underlying rhythm was sinus rhythm with first-degree AV block present. 192 episodes of second-degree AV block (Mobitz II) occurred, lasting a total of 1 day 4 hours. There were 4 ventricular tachycardia runs, with the fastest lasting 7 beats at a max rate of 182 bpm, and the longest lasting 16 beats with an average rate of 135 bpm. 5871 supraventricular tachycardia (SVT) runs were recorded, the fastest lasting 12 beats with a max rate of 187 bpm, and the longest lasting 35.3 seconds with an average rate of 133 bpm. Atrial fibrillation occurred with a 12% burden, ranging from 45-212 bpm (average 84 bpm). Three pauses occurred, the longest lasting 5.4 seconds (11 bpm). Supraventricular tachycardia, pauses, AV block, and atrial fibrillation were detected within 45 seconds of patient-triggered/diary events. Isolated SVEs were frequent (5.4%, 76752), SVE couplets were occasional (3.7%, 26323), and SVE triplets were occasional (1.7%, 7781). Isolated VEs were rare (0 to <1.0%, 5154), VE couplets were rare (0 to <1.0%, 19), and VE triplets were rare (0 to <1.0%, 1). Ventricular bigeminy and trigeminy were present. MD notification criteria for rapid atrial fibrillation and AV block were met, and the RN was notified on 03/15/2013 at 12:00 pm CST."
        
        **Example Final Interpretation:**

        "1- Sinus rhythm with first degree AV block. 
        2- Episodes of second degree 2:1 AV Block. 
        3- Sinus pauses with longest being 5.4 seconds. 
        4- Frequent isolated PACs with occasional atrial couplets and triplets. 
        5- Many runs of SVT many of which appear to be ectopic atrial tachycardia (5871) with longest being 35 sec. 
        6- Rare isolated PVCs, ventricular couplets and triplets. 
        7- 4 Runs of VT up to 16 beats rate range of 135-182/min. 
        8- Atrial fibrillation with a burden of 12% with rates of 45-212/ min. 
        9- Symptoms appear to correlate with SVT, sinus pauses, AV Block and atrial fibrillation."

        ---

        **Your Task:**

        As a knowledgeable clinician, you are provided with a diagnosis in the format "clinical metrics: diagnostic value."

        Based on the above examples of Clinical Findings and Final Interpretation, please:

        1. Generate the Clinical Findings in the same format as the above "Example Clinical Findings".
        2. Create the Final Interpretation in a similar format as the above "Example Final Interpretation", incorporating the provided information in "Must-include Interpretation".
        
        **Instructions:**
        1. The Clinical Findings should be detailed descriptive paragraphs. Do not simply list existing items.
        2. The Final Interpretation can be either a list or descriptive paragraphs.
        
        **Input:**
        - Diagnosis: [DIAGNOSIS]
        - Must-include Interpretation: [SHORT_INTERPRETATION]

        **Output:**

        **Clinical Findings:**
        [CLINICAL_FINDINGS]

        **Final Interpretation:**
        [FINAL_INTERPRETATION]

        
    """
               
    while True:
        try: 
            response = g4f.ChatCompletion.create(
                model=model,
                messages=[
                {
                    "role": "user",
                    "content": template.replace("[DIAGNOSIS]", inputs).replace("[SHORT_INTERPRETATION]", conclusion)
                }
                ],
                temperature=0,
                max_tokens=10000,
                top_p=0,
                frequency_penalty=0,
                presence_penalty=0
            )
            if "Findings:" in response and "Interpretation:" in response:
                return response
        
        except:
            pass


def load_csv_as_dict(csv_file_path="metric_dict.csv"):
    data_dict = {}
    with open(csv_file_path, mode='r', newline='') as csv_file:
        reader = csv.reader(csv_file)
        for row in reader:
            if len(row) == 2:  # Ensure the row has exactly two columns
                key, value = row
                data_dict[key] = value
    return data_dict

data_dict = load_csv_as_dict()


data_file_path = "data_result_0716.json"
save_file_name = f"data_result_{llm}.json"

data = []
with open(data_file_path, 'r') as file:
    json_list = json.load(file)

    for sample in json_list:
        metrics = sample["metrics"]
        
        new_metrics = []
        for pair in metrics.split("\n"):
            assert len(pair.split(": ")) == 2
            k, v = pair.split(": ")
            
            if k in data_dict:
                new_metrics.append(f"{data_dict[k]}: {v}")
            else:
                # print(f"{k} not exist in metric dict")
                new_metrics.append(f"{k}: {v}")
        
        new_metrics = "\n".join(new_metrics)
        response = generate(new_metrics, sample["conclusion"])
        assert "**Clinical Findings:**" in response and "**Final Interpretation:**" in response
        
        findings, interpretation = response.split("**Clinical Findings:**")[-1].split("**Final Interpretation:**")
        
        findings = findings.strip()
        interpretation = interpretation.strip()
        
        sample["findings"] = findings
        sample["interpretation"] = interpretation

        data.append(sample)
        print(sample)

with open(save_file_name, 'w') as f:
    json.dump(data, f, indent=4)
