import streamlit as st
from transformers import pipeline
import nltk
from nltk.tokenize import word_tokenize

# Buat objek terjemahan
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-id-en")
terjemah = pipeline("translation", model="Helsinki-NLP/opus-mt-en-id")
pipe = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")

# Antarmuka Streamlit
st.title("Diagnosa Berdasarkan Gejala Kehamilan")

# Masukkan gejala, usia, dan jenis kelamin
usia = st.number_input("Masukkan usia Anda:", min_value=0, max_value=120)
jenis_kelamin = st.selectbox("Masukkan jenis kelamin Anda:", ["Laki-laki", "Perempuan"])
gejala_id = st.text_area("Masukkan gejala Anda:")

if st.button("Diagnosa"):
    if gejala_id:
        # Terjemahkan gejala dari Bahasa Indonesia ke Bahasa Inggris
        gejala_en = translator(gejala_id, max_length=100)[0]["translation_text"]
        informasi_pasien = f"I am {usia} years old, {jenis_kelamin}. Current symptoms are: {gejala_en}"

        # Contoh kasus kehamilan untuk diagnosis
        pesan = [
            {"role": "system",
             "content": """
You are a doctor diagnosing pregnancy-related conditions based on symptoms.

Example 1:
Patient symptoms: Missed period, nausea, tender breasts
Diagnosis: Possible pregnancy

Example 2:
Patient symptoms: Severe abdominal pain, bleeding
Diagnosis: Possible miscarriage or ectopic pregnancy.
"""},
            {"role": "user", "content": f"Based on your assessment, {informasi_pasien}, what is your diagnosis?"}
        ]

        # Fungsi untuk mendapatkan konten dari role 'assistant'
        def get_assistant_content(response):
            return response[0]['generated_text']

        # Dapatkan respon dari pipe
        response = pipe(pesan, num_return_sequences=1, truncation=True)
        diagnosis = get_assistant_content(response)

        # Terjemahkan hasil ke Bahasa Indonesia
        diagnosa_terjemahan = terjemah(diagnosis, max_length=100)[0]["translation_text"]

        # Tampilkan hasil ke Streamlit
        st.subheader("Hasil Diagnosis:")
        st.write(f"Diagnosis: {diagnosa_terjemahan}")
    else:
        st.error("Harap isi semua kolom sebelum menekan tombol Diagnosa.")