# %% API call1
import time
import json
import os
from datetime import datetime
import pandas as pd
from openai import OpenAI
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# === CONFIGURATION ===
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL")
MODEL_NAME = "GPT-OSS-120B"
HEALTH_URL = f"{OPENAI_BASE_URL}/health"  # Placeholder - actual health check would need to be implemented
CHAT_URL = f"{OPENAI_BASE_URL}/chat/completions"

# File paths
INPUT_CSV = "/home/shahin/Lab/Doktorarbeit/Barcelona/Data/MS_Briefe_400_with_unique_id_SHA3_explore_cleaned_unique.csv"
EDSS_INSTRUCTIONS_PATH = "/home/shahin/Lab/Doktorarbeit/Barcelona/attach/Komplett.txt"
#GRAMMAR_FILE = "/home/shahin/Lab/Doktorarbeit/Barcelona/attach/just_edss_schema.gbnf"

# Initialize OpenAI client
client = OpenAI(
    api_key=OPENAI_API_KEY,
    base_url=OPENAI_BASE_URL
)

# Read EDSS instructions from file
with open(EDSS_INSTRUCTIONS_PATH, 'r') as f:
    EDSS_INSTRUCTIONS = f.read().strip()

# === RUN INFERENCE ===
#def run_inference(patient_text):
#    prompt = f''' Du bist ein medizinischer Assistent, der spezialisiert darauf ist, EDSS-Scores (Expanded Disability Status Scale) aus klinischen Berichten zu extrahieren.
#
#### Regeln für die Ausgabe:
#1. **Reason**: Erstelle eine prägnante Zusammenfassung (max. 400 Zeichen) der Befunde, die zur Einstufung führen. Die Zusammenfassung muss auf **DEUTSCH** verfasst sein.
#2. **nicht_klassifizierbar**:
#   - Setze dies auf **true**, NUR wenn die klinischen Daten unzureichend, widersprüchlich oder fehlend sind, um einen EDSS-Wert zu bestimmen.
#   - Setze dies auf **false**, wenn ein spezifischer EDSS-Wert identifiziert oder berechnet werden kann.
#3. **EDSS**:
#   - Dieses Feld ist **VERPFLICHTEND**, wenn "nicht_klassifizierbar" auf false steht.
#   - Es muss eine Zahl zwischen 0.0 und 10.0 sein.
#   - Dieses Feld **DARF NICHT ERSCHEINEN**, wenn "nicht_klassifizierbar" auf true steht.
#
#### Einschränkungen:
#- Erfinde keine Daten.
#- Wenn der Bericht zweideutig ist, markiere "nicht_klassifizierbar": true.
#- Halte dich strikt an die JSON-Struktur.
#
#EDSS-Bewertungsrichtlinien:
#{EDSS_INSTRUCTIONS}
#
#Patientenbericht:
#{patient_text}
#'''
#
#    start_time = time.time()
#
#    try:
#        # Make API call using OpenAI client
#        response = client.chat.completions.create(
#            messages=[
#                {
#                    "role": "system",
#                    "content": "You extract EDSS scores."
#                },
#                {
#                    "role": "user",
#                    "content": prompt
#                }
#            ],
#            model=MODEL_NAME,
#            max_tokens=2048,
#            temperature=0.0,
#            response_format={"type": "json_object"}
#        )
#
#        # Extract content from response
#        content = response.choices[0].message.content
#
#        # Parse the JSON response
#        parsed = json.loads(content)
#
#        inference_time = time.time() - start_time
#
#        return {
#            "success": True,
#            "result": parsed,
#            "inference_time_sec": inference_time
#        }
#
#    except Exception as e:
#        print(f"Inference error: {e}")
#        return {
#            "success": False,
#            "error": str(e),
#            "inference_time_sec": -1
#        }
#

# === RUN INFERENCE 2 ===
def run_inference(patient_text):
    prompt = f''' Du bist ein medizinischer Assistent, der spezialisiert darauf ist, EDSS-Scores (Expanded Disability Status Scale) aus klinischen Berichten zu extrahieren.

### Regeln für die Ausgabe:
1. **Reason**: Erstelle eine prägnante Zusammenfassung (max. 400 Zeichen) der Befunde auf **DEUTSCH**, die zur Einstufung führen.
2. **klassifizierbar**:
   - Setze dies auf **true**, wenn ein EDSS-Wert identifiziert, berechnet oder basierend auf den klinischen Hinweisen plausibel geschätzt werden kann.
   - Setze dies auf **false**, NUR wenn die Daten absolut unzureichend oder so widersprüchlich sind, dass keinerlei Einstufung möglich ist.
3. **EDSS**:
   - Dieses Feld ist **VERPFLICHTEND**, wenn "klassifizierbar" auf true steht.
   - Es muss eine Zahl zwischen 0.0 und 10.0 sein.
   - Versuche stets, den EDSS-Wert so präzise wie möglich zu bestimmen, auch wenn die Datenlage dünn ist (nutze verfügbare Informationen zu Gehstrecke und Funktionssystemen).
   - Dieses Feld **DARF NICHT ERSCHEINEN**, wenn "klassifizierbar" auf false steht.

### Einschränkungen:
- Erfinde keine Fakten, aber nutze klinische Herleitungen aus dem Bericht, um den EDSS zu bestimmen.
- Priorisiere die Vergabe eines EDSS-Wertes gegenüber der Markierung als nicht klassifizierbar.
- Halte dich strikt an die JSON-Struktur.

EDSS-Bewertungsrichtlinien:
{EDSS_INSTRUCTIONS}

Patientenbericht:
{patient_text}
'''

    start_time = time.time()

    try:
        # Make API call using OpenAI client
        response = client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": "You extract EDSS scores. You prioritize providing a score even if data is partial, by using clinical inference."
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            model=MODEL_NAME,
            max_tokens=2048,
            temperature=0.0,
            response_format={"type": "json_object"}
        )

        # Extract content from response
        content = response.choices[0].message.content

        # Parse the JSON response
        parsed = json.loads(content)

        inference_time = time.time() - start_time

        return {
            "success": True,
            "result": parsed,
            "inference_time_sec": inference_time
        }

    except Exception as e:
        print(f"Inference error: {e}")
        return {
            "success": False,
            "error": str(e),
            "inference_time_sec": -1
        }
# === BUILD PATIENT TEXT ===
def build_patient_text(row):
    return (
        str(row["T_Zusammenfassung"]) + "\n" +
        str(row["Diagnosen"]) + "\n" +
        str(row["T_KlinBef"]) + "\n" +
        str(row["T_Befunde"]) + "\n"
    )

if __name__ == "__main__":
    # Read CSV file ONLY inside main block
    df = pd.read_csv(INPUT_CSV, sep=';')
    results = []

    # Process each row
    for idx, row in df.iterrows():
        print(f"Processing row {idx + 1}/{len(df)}")
        try:
            patient_text = build_patient_text(row)
            result = run_inference(patient_text)

            # Add unique_id and MedDatum to result for tracking
            result["unique_id"] = row.get("unique_id", f"row_{idx}")
            result["MedDatum"] = row.get("MedDatum", None)

            results.append(result)
            print(json.dumps(result, indent=2))
        except Exception as e:
            print(f"Error processing row {idx}: {e}")
            results.append({
                "success": False,
                "error": str(e),
                "unique_id": row.get("unique_id", f"row_{idx}"),
                "MedDatum": row.get("MedDatum", None)
            })

    # Save results to a JSON file
    output_json = INPUT_CSV.replace(".csv", "_results2.json")
    with open(output_json, 'w') as f:
        json.dump(results, f, indent=2)
    print(f"Results saved to {output_json}")
##