% Medical Diagnosis System with Enhanced Knowledge Base and Improved Questioning Logic

% Declare dynamic predicates for flexibility
:- dynamic symptom/2, duration/2, severity/2, related_symptoms/2, diagnosis/2.

% Start of the diagnosis process
start_diagnosis :-
    format('Welcome to the Medical Diagnosis Assistant.~n'),
    ask_symptoms(user).

% Ask for symptoms and gather additional details
ask_symptoms(PatientName) :-
    write('Tell me what symptoms you have. Type a symptom and press enter. Type \'done\' when finished: '), nl,
    read_line_to_string(user, SymptomStr), 
    string_lower(SymptomStr, LowerSymptomStr),
    strip_quotes(LowerSymptomStr, Symptom),
    (   Symptom == "done"
    ->  collect_additional_info  % Call collect_additional_info directly
    ;   Symptom \= "",
        assertz(symptom(PatientName, Symptom)),
        ask_symptoms(PatientName)
    ).

% Helper predicate to strip quotes from the input string
strip_quotes(Str, Stripped) :-
    sub_string(Str, 0, 1, _, "'"),  % Check if starts with a quote
    sub_string(Str, _, 1, 0, "'"),  % Check if ends with a quote
    sub_string(Str, 1, _, 1, Stripped),  % Remove quotes
    !.
strip_quotes(Str, Str) :- !.  % If no quotes, return the string as is


collect_additional_info :-
    findall(Sym, symptom(user, Sym), Symptoms),
    maplist(ask_symptom_details, Symptoms),
    diagnose_patient.

ask_symptom_details(Symptom) :-
    ask_duration(Symptom),
    ask_severity(Symptom),
    inquire_related_symptoms(Symptom).

% Asking for symptom duration and severity
ask_duration(Symptom) :-
    format('How long have you had ~w? (e.g., 2 days, 3 hours) ', [Symptom]),
    read(Duration),
    assertz(duration(user, Symptom, Duration)).

ask_severity(Symptom) :-
    format('On a scale of 1-10, how severe is your ~w? ', [Symptom]),
    read(Severity),
    assertz(severity(user, Symptom, Severity)).

% Only inquire about related symptoms that are relevant
inquire_related_symptoms(Symptom) :-
    findall(S, symptom(user, S), ReportedSymptoms),
    related_symptoms(Symptom, Related),
    intersection(ReportedSymptoms, Related, RelevantRelated),
    ask_if_has_related_symptoms(RelevantRelated).

% Asking about related symptoms only if they haven't been reported
ask_if_has_related_symptoms([]).
ask_if_has_related_symptoms([R|Rs]) :-
    format('Are you experiencing ~w? (yes/no) ', [R]),
    read(Answer),
    (   Answer == yes -> assertz(symptom(user, R))
    ;   true
    ),
    ask_if_has_related_symptoms(Rs).

% Define related symptoms for better questioning
related_symptoms(fever, [cough, sore_throat, runny_nose, body_ache, headache, chills, fatigue]).
related_symptoms(cough, [fever, shortness_of_breath, wheezing, chest_pain, fatigue, night_sweats]).
related_symptoms(shortness_of_breath, [cough, chest_pain, wheezing, fatigue, dizziness]).
related_symptoms(chest_pain, [shortness_of_breath, cough, dizziness, sweating, nausea]).
related_symptoms(headache, [fever, nausea, vision_problems, sensitivity_to_light, stiff_neck]).
related_symptoms(nausea, [vomiting, diarrhea, fever, abdominal_pain, headache]).
related_symptoms(abdominal_pain, [nausea, vomiting, diarrhea, fever, bloating]).
related_symptoms(fatigue, [weakness, dizziness, headache, blurred_vision, shortness_of_breath]).
related_symptoms(dizziness, [nausea, headache, fainting, blurred_vision, loss_of_balance]).
related_symptoms(wheezing, [cough, shortness_of_breath, chest_tightness, fatigue]).
related_symptoms(diarrhea, [nausea, vomiting, abdominal_pain, fever, dehydration]).
related_symptoms(vomiting, [nausea, diarrhea, abdominal_pain, fever, dizziness]).
related_symptoms(joint_pain, [swelling, redness, warmth, stiffness, fever]).
related_symptoms(swelling_feet, [shortness_of_breath, fatigue, weight_gain, increased_urination]).
related_symptoms(high_blood_pressure, [headache, dizziness, blurred_vision, chest_pain]).

% Diagnosis process with enhanced reasoning
diagnose_patient :-
    findall(Disease, diagnosis(user, Disease), PossibleDiseases),
    diagnose_based_on(PossibleDiseases).

diagnose_based_on([]) :-
    write('A specific diagnosis could not be determined based on the symptoms. Please consult a healthcare professional.').

diagnose_based_on([Disease|_]) :-
    advice(Disease, Advice),
    treatment(Disease, Treatment),
    format('Based on the symptoms, the possible diagnosis is: ~w~n', [Disease]),
    format('Advice: ~w~n', [Advice]),
    format('Treatment: ~w~n', [Treatment]).

% Utility to convert various duration formats to days
convert_duration_to_days(User, Symptom, Days) :-
    duration(User, Symptom, Duration),
    parse_duration(Duration, Days).

% Parse duration input and convert to days
parse_duration(DurationStr, Days) :-
    split_string(DurationStr, " ", "", [NumStr, Unit]),
    number_string(Num, NumStr),
    convert_unit_to_days(Unit, Num, Days).

% Convert different time units to days
convert_unit_to_days("minute", Num, Days) :- Days is Num / 1440. % 1 day = 1440 minutes
convert_unit_to_days("minutes", Num, Days) :- Days is Num / 1440.
convert_unit_to_days("hour", Num, Days) :- Days is Num / 24. % 1 day = 24 hours
convert_unit_to_days("hours", Num, Days) :- Days is Num / 24.
convert_unit_to_days("day", Num, Days) :- Days is Num.
convert_unit_to_days("days", Num, Days) :- Days is Num.
convert_unit_to_days("week", Num, Days) :- Days is Num * 7.
convert_unit_to_days("weeks", Num, Days) :- Days is Num * 7.
convert_unit_to_days("month", Num, Days) :- Days is Num * 30. % Approximation
convert_unit_to_days("months", Num, Days) :- Days is Num * 30.
convert_unit_to_days("year", Num, Days) :- Days is Num * 365. % Approximation
convert_unit_to_days("years", Num, Days) :- Days is Num * 365.

% Fallback for unrecognized units
convert_unit_to_days(Unit, _, _) :-
    format('Unrecognized time unit: ~w. Please use minutes, hours, days, weeks, months, or years.', [Unit]),
    fail.

% Define symptom_duration
symptom_duration(User, Symptom, Duration, Severity) :-
    convert_duration_to_days(User, Symptom, Duration),
    severity(User, Symptom, Severity).

% Define diseases and their associated symptoms
% Acute Heart Failure
diagnosis(user, acute_heart_failure) :-
    symptom_duration(user, shortness_of_breath, Duration1, severe),
    symptom_duration(user, fatigue, Duration2, moderate_to_severe),
    symptom_duration(user, swelling_feet, Duration3, moderate),
    Duration1 < 2_weeks, Duration2 < 2_weeks, Duration3 < 2_weeks.

advice(acute_heart_failure, 'Limit fluid intake, avoid alcohol, and monitor weight daily').
treatment(acute_heart_failure, 'Diuretics, ACE inhibitors, and beta-blockers').

% Define the prior probability of Acute Heart Failure
prior_probability(acute_heart_failure, 0.01).

% Define the likelihood of symptoms for Acute Heart Failure
symptom_likelihood(acute_heart_failure, shortness_of_breath, severe, 0.8).
symptom_likelihood(acute_heart_failure, fatigue, moderate_to_severe, 0.7).
symptom_likelihood(acute_heart_failure, swelling_feet, moderate, 0.6).

% Update disease probability based on a symptom
update_probability(Disease, Symptom, Severity, UpdatedProbability) :-
    prior_probability(Disease, Prior),
    symptom_likelihood(Disease, Symptom, Severity, Likelihood),
    UpdatedProbability is Prior * Likelihood.

% Diagnose Acute Heart Failure based on symptoms
diagnose_acute_heart_failure(user, Probability) :-
    findall(Prob, (symptom(user, Symptom, Severity), update_probability(acute_heart_failure, Symptom, Severity, Prob)), Probabilities),
    sum_list(Probabilities, TotalProbability),
    length(Probabilities, Count),
    Count > 0,
    Probability is TotalProbability / Count.

% Diabetes Mellitus Diagnosis
diagnosis(user, diabetes_mellitus) :-
    symptom_duration(user, frequent_urination, Duration1, persistent),
    symptom_duration(user, increased_thirst, Duration2, persistent),
    symptom_duration(user, unexplained_weight_loss, Duration3, significant),
    ((fasting_blood_glucose(user, Level), Level > 126);
     (hemoglobin_a1c(user, Level), Level > 6.5);
     (plasma_glucose(user, Level), Level > 200)).

% Define advice and treatment for Diabetes Mellitus
advice(diabetes_mellitus, 'Monitor blood sugar levels, maintain a balanced diet, and exercise regularly').
treatment(diabetes_mellitus, 'Metformin, insulin therapy, and lifestyle changes').

% Define the prior probability of Diabetes Mellitus
prior_probability(diabetes_mellitus, 0.01).

% Define the likelihood of symptoms for Diabetes Mellitus
symptom_likelihood(diabetes_mellitus, frequent_urination, persistent, 0.7).
symptom_likelihood(diabetes_mellitus, increased_thirst, persistent, 0.6).
symptom_likelihood(diabetes_mellitus, unexplained_weight_loss, significant, 0.5).

% Define likelihoods for blood glucose levels
symptom_likelihood(diabetes_mellitus, fasting_blood_glucose, Level, 0.8) :- Level > 126.
symptom_likelihood(diabetes_mellitus, fasting_blood_glucose, Level, 0.2) :- Level =< 126.
symptom_likelihood(diabetes_mellitus, hemoglobin_a1c, Level, 0.85) :- Level > 6.5.
symptom_likelihood(diabetes_mellitus, hemoglobin_a1c, Level, 0.15) :- Level =< 6.5.
symptom_likelihood(diabetes_mellitus, plasma_glucose, Level, 0.9) :- Level > 200.
symptom_likelihood(diabetes_mellitus, plasma_glucose, Level, 0.1) :- Level =< 200.

% Update disease probability based on a symptom or test result
update_probability(Disease, Symptom, Value, UpdatedProbability) :-
    prior_probability(Disease, Prior),
    symptom_likelihood(Disease, Symptom, Value, Likelihood),
    UpdatedProbability is Prior * Likelihood.

% Diagnose Diabetes Mellitus based on symptoms and test results
diagnose_diabetes_mellitus(Patient, Probability) :-
    findall(Prob, (
        (symptom(Patient, Symptom, Severity), update_probability(diabetes_mellitus, Symptom, Severity, Prob));
        (blood_test(Patient, Test, Value), update_probability(diabetes_mellitus, Test, Value, Prob))
    ), Probabilities),
    sum_list(Probabilities, TotalProbability),
    length(Probabilities, Count),
    Count > 0,
    Probability is TotalProbability / Count.

% Hypertension Diagnosis
diagnosis(user, hypertension) :-
    symptom_duration(user, high_blood_pressure, Duration, persistent),
    Duration > 180, % Assuming 6 months as 180 days
    ((symptom_duration(user, headache, _, moderate) ; symptom_duration(user, shortness_of_breath, _, moderate)),
    (symptom(user, blurred_vision) ; symptom(user, chest_pain)),
    (symptom(user, dizziness) ; symptom(user, nausea))).

advice(hypertension, 'Reduce salt intake, exercise regularly, and avoid tobacco use').
treatment(hypertension, 'ACE inhibitors, calcium channel blockers, and lifestyle modifications').

% Define the prior probability of Hypertension
prior_probability(hypertension, 0.05).

% Define the likelihood of symptoms for Hypertension
symptom_likelihood(hypertension, high_blood_pressure, persistent, 0.8).
symptom_likelihood(hypertension, headache, moderate, 0.5).
symptom_likelihood(hypertension, shortness_of_breath, moderate, 0.5).

% Update disease probability based on a symptom
update_probability_hypertension(Disease, Symptom, Severity, UpdatedProbability) :-
    prior_probability(Disease, Prior),
    symptom_likelihood(Disease, Symptom, Severity, Likelihood),
    UpdatedProbability is Prior * Likelihood.

% Diagnose Hypertension based on symptoms
diagnose_hypertension(Patient, Probability) :-
    findall(Prob, (symptom(Patient, Symptom, Severity), update_probability_hypertension(hypertension, Symptom, Severity, Prob)), Probabilities),
    sum_list(Probabilities, TotalProbability),
    length(Probabilities, Count),
    Count > 0,
    Probability is TotalProbability / Count.

% Upper Respiratory Tract Infection (URTI)
diagnosis(user, urti) :-
    symptom_duration(user, sore_throat, Duration1, moderate),
    symptom_duration(user, runny_nose, Duration2, moderate),
    symptom_duration(user, cough, Duration3, moderate),
    symptom_duration(user, sneezing, Duration4, mild),
    symptom_duration(user, fever, Duration5, mild_to_moderate),
    Duration1 < 2_weeks, Duration2 < 2_weeks, Duration3 < 2_weeks, Duration4 < 2_weeks, Duration5 < 2_weeks.

advice(urti, 'Rest, stay hydrated, and avoid spreading infection').
treatment(urti, 'Symptomatic relief with over-the-counter medications, throat lozenges, and increased fluid intake').

% Define the prior probability of URTI
prior_probability(urti, 0.05).

% Define the likelihood of symptoms for URTI
symptom_likelihood(urti, sore_throat, moderate, 0.6).
symptom_likelihood(urti, runny_nose, moderate, 0.7).
symptom_likelihood(urti, cough, moderate, 0.65).
symptom_likelihood(urti, sneezing, mild, 0.5).
symptom_likelihood(urti, fever, mild_to_moderate, 0.55).

% Update disease probability based on a symptom
update_probability(Disease, Symptom, Severity, UpdatedProbability) :-
    prior_probability(Disease, Prior),
    symptom_likelihood(Disease, Symptom, Severity, Likelihood),
    UpdatedProbability is Prior * Likelihood.

% Diagnose URTI based on symptoms
diagnose_urti(Patient, Probability) :-
    findall(Prob, (symptom(Patient, Symptom, Severity), update_probability(urti, Symptom, Severity, Prob)), Probabilities),
    sum_list(Probabilities, TotalProbability),
    length(Probabilities, Count),
    Count > 0,
    Probability is TotalProbability / Count.


% Flu Virus
diagnosis(user, flu_virus) :-
    symptom_duration(user, high_fever, Duration1, severe),
    symptom_duration(user, body_aches, Duration2, moderate),
    symptom_duration(user, cough, Duration3, moderate),
    symptom_duration(user, fatigue, Duration4, moderate),
    symptom_duration(user, sore_throat, Duration5, moderate),
    symptom_duration(user, headache, Duration6, moderate),
    Duration1 < 2_weeks, Duration2 < 2_weeks, Duration3 < 2_weeks, Duration4 < 2_weeks, Duration5 < 2_weeks, Duration6 < 2_weeks.

advice(flu_virus, 'Rest, stay hydrated, and avoid close contact with others').
treatment(flu_virus, 'Antiviral medications, symptomatic relief, and plenty of rest').

% Define the prior probability of Flu Virus
prior_probability(flu_virus, 0.05).

% Define the likelihood of symptoms for Flu Virus
symptom_likelihood(flu_virus, high_fever, severe, 0.7).
symptom_likelihood(flu_virus, body_aches, moderate, 0.6).
symptom_likelihood(flu_virus, cough, moderate, 0.65).
symptom_likelihood(flu_virus, fatigue, moderate, 0.6).
symptom_likelihood(flu_virus, sore_throat, moderate, 0.55).
symptom_likelihood(flu_virus, headache, moderate, 0.6).

% Update disease probability based on a symptom
update_probability(Disease, Symptom, Severity, UpdatedProbability) :-
    prior_probability(Disease, Prior),
    symptom_likelihood(Disease, Symptom, Severity, Likelihood),
    UpdatedProbability is Prior * Likelihood.

% Diagnose Flu Virus based on symptoms
diagnose_flu_virus(Patient, Probability) :-
    findall(Prob, (symptom(Patient, Symptom, Severity), update_probability(flu_virus, Symptom, Severity, Prob)), Probabilities),
    sum_list(Probabilities, TotalProbability),
    length(Probabilities, Count),
    Count > 0,
    Probability is TotalProbability / Count.

% Pneumonia
diagnosis(user, pneumonia) :-
    symptom_duration(user, cough_with_phlegm, Duration1, severe),
    symptom_duration(user, fever, Duration2, moderate_to_severe),
    symptom_duration(user, shortness_of_breath, Duration3, severe),
    symptom_duration(user, chest_pain, Duration4, moderate_to_severe),
    Duration1 < 1_month, Duration2 < 1_month, Duration3 < 1_month, Duration4 < 1_month.

advice(pneumonia, 'Rest, stay hydrated, and follow-up with healthcare provider').
treatment(pneumonia, 'Antibiotics, fever reducers, and cough medicine').

% Define the prior probability of Pneumonia
prior_probability(pneumonia, 0.03).

% Define the likelihood of symptoms for Pneumonia
symptom_likelihood(pneumonia, cough_with_phlegm, severe, 0.8).
symptom_likelihood(pneumonia, fever, moderate_to_severe, 0.7).
symptom_likelihood(pneumonia, shortness_of_breath, severe, 0.75).
symptom_likelihood(pneumonia, chest_pain, moderate_to_severe, 0.65).

% Update disease probability based on a symptom
update_probability(Disease, Symptom, Severity, UpdatedProbability) :-
    prior_probability(Disease, Prior),
    symptom_likelihood(Disease, Symptom, Severity, Likelihood),
    UpdatedProbability is Prior * Likelihood.

% Diagnose Pneumonia based on symptoms
diagnose_pneumonia(Patient, Probability) :-
    findall(Prob, (symptom(Patient, Symptom, Severity), update_probability(pneumonia, Symptom, Severity, Prob)), Probabilities),
    sum_list(Probabilities, TotalProbability),
    length(Probabilities, Count),
    Count > 0,
    Probability is TotalProbability / Count.

% COVID-19 Infection
diagnosis(user, covid_19) :-
    symptom_duration(user, fever, Duration1, moderate_to_severe),
    symptom_duration(user, cough, Duration2, moderate),
    symptom_duration(user, shortness_of_breath, Duration3, moderate_to_severe),
    symptom_duration(user, loss_of_taste_or_smell, Duration4, mild_to_moderate),
    symptom_duration(user, fatigue, Duration5, moderate),
    Duration1 < 2_weeks, Duration2 < 2_weeks, Duration3 < 2_weeks, Duration4 < 2_weeks, Duration5 < 2_weeks.

advice(covid_19, 'Isolate, monitor symptoms, and stay hydrated').
treatment(covid_19, 'Symptomatic relief and potential antiviral medication').

% Define the prior probability of COVID-19
prior_probability(covid_19, 0.05).

% Define the likelihood of symptoms for COVID-19
symptom_likelihood(covid_19, fever, moderate_to_severe, 0.7).
symptom_likelihood(covid_19, cough, moderate, 0.6).
symptom_likelihood(covid_19, shortness_of_breath, moderate_to_severe, 0.65).
symptom_likelihood(covid_19, loss_of_taste_or_smell, mild_to_moderate, 0.75).
symptom_likelihood(covid_19, fatigue, moderate, 0.5).

% Update disease probability based on a symptom
update_probability(Disease, Symptom, Severity, UpdatedProbability) :-
    prior_probability(Disease, Prior),
    symptom_likelihood(Disease, Symptom, Severity, Likelihood),
    UpdatedProbability is Prior * Likelihood.

% Diagnose COVID-19 based on symptoms
diagnose_covid_19(Patient, Probability) :-
    findall(Prob, (symptom(Patient, Symptom, Severity), update_probability(covid_19, Symptom, Severity, Prob)), Probabilities),
    sum_list(Probabilities, TotalProbability),
    length(Probabilities, Count),
    Count > 0,
    Probability is TotalProbability / Count.

% Urinary Tract Infection (UTI)
diagnosis(Patient, uti) :-
    symptom_duration(Patient, burning_feeling_when_peeing, Duration1, Severity1),
    symptom_duration(Patient, peeing_more_often_than_usual, Duration2, Severity2),
    symptom_duration(Patient, pee_that_looks_cloudy, Duration3, Severity3),
    symptom_duration(Patient, pee_that_smells_strong, Duration4, Severity4),
    Duration1 < 2_weeks, Duration2 < 2_weeks, Duration3 < 2_weeks, Duration4 < 2_weeks,
    Severity1 >= moderate, Severity2 >= moderate, Severity3 >= mild, Severity4 >= mild.

advice(uti, 'Increase fluid intake, practice good hygiene, and avoid irritants like spicy foods').
treatment(uti, 'Antibiotics such as trimethoprim/sulfamethoxazole, nitrofurantoin, or fosfomycin').

% Define the prior probability of UTI
prior_probability(uti, 0.03).

% Define the likelihood of symptoms for UTI
symptom_likelihood(uti, burning_feeling_when_peeing, moderate, 0.8).
symptom_likelihood(uti, peeing_more_often_than_usual, moderate, 0.7).
symptom_likelihood(uti, pee_that_looks_cloudy, mild, 0.6).
symptom_likelihood(uti, pee_that_smells_strong, mild, 0.5).

% Update disease probability based on a symptom
update_probability(Disease, Symptom, Severity, UpdatedProbability) :-
    prior_probability(Disease, Prior),
    symptom_likelihood(Disease, Symptom, Severity, Likelihood),
    UpdatedProbability is Prior * Likelihood.

% Diagnose UTI based on symptoms
diagnose_uti(Patient, Probability) :-
    findall(Prob, (symptom(Patient, Symptom, Severity), update_probability(uti, Symptom, Severity, Prob)), Probabilities),
    sum_list(Probabilities, TotalProbability),
    length(Probabilities, Count),
    Count > 0,
    Probability is TotalProbability / Count.

% Migraine Headaches
diagnosis(Patient, migraine_headaches) :-
    symptom_duration(Patient, severe_pulsating_headache, Duration1, Severity1),
    symptom_duration(Patient, sensitivity_to_light_or_sound, Duration2, Severity2),
    symptom_duration(Patient, nausea_or_vomiting, Duration3, Severity3),
    symptom_duration(Patient, headache_worsens_with_activity, Duration4, Severity4),
    Duration1 < 72_hours, Duration2 < 72_hours, Duration3 < 72_hours, Duration4 < 72_hours,
    Severity1 >= moderate, Severity2 >= mild, Severity3 >= mild, Severity4 >= mild.

advice(migraine_headaches, 'Identify and avoid triggers, stress management, regular sleep schedule').
treatment(migraine_headaches, 'Pain relief with NSAIDs, triptans like sumatriptan, preventive medications like propranolol').

% Define the prior probability of Migraine Headaches
prior_probability(migraine_headaches, 0.05).

% Define the likelihood of symptoms for Migraine Headaches
symptom_likelihood(migraine_headaches, severe_pulsating_headache, moderate, 0.8).
symptom_likelihood(migraine_headaches, sensitivity_to_light_or_sound, mild, 0.6).
symptom_likelihood(migraine_headaches, nausea_or_vomiting, mild, 0.5).
symptom_likelihood(migraine_headaches, headache_worsens_with_activity, mild, 0.4).

% Update disease probability based on a symptom
update_probability(Disease, Symptom, Severity, UpdatedProbability) :-
    prior_probability(Disease, Prior),
    symptom_likelihood(Disease, Symptom, Severity, Likelihood),
    UpdatedProbability is Prior * Likelihood.

% Diagnose Migraine Headaches based on symptoms
diagnose_migraine_headaches(Patient, Probability) :-
    findall(Prob, (symptom(Patient, Symptom, Severity), update_probability(migraine_headaches, Symptom, Severity, Prob)), Probabilities),
    sum_list(Probabilities, TotalProbability),
    length(Probabilities, Count),
    Count > 0,
    Probability is TotalProbability / Count.

% Diarrheal Disease
diagnosis(Patient, diarrheal_disease) :-
    symptom_duration(Patient, frequent_loose_or_watery_stools, Duration1, Severity1),
    symptom_duration(Patient, abdominal_pain_or_cramping, Duration2, Severity2),
    symptom_duration(Patient, nausea_or_vomiting, Duration3, Severity3),
    symptom_duration(Patient, dehydration_signs, Duration4, Severity4),
    Duration1 < 2_weeks, Duration2 < 2_weeks, Duration3 < 2_weeks, Duration4 < 2_weeks,
    Severity1 >= moderate, Severity2 >= mild, Severity3 >= mild, Severity4 >= mild.

advice(diarrheal_disease, 'Staying hydrated, eating bland foods, practicing good hygiene').
treatment(diarrheal_disease, 'Oral rehydration solutions, anti-diarrheal medications, antibiotics if bacterial').

% Define the prior probability of Diarrheal Disease
prior_probability(diarrheal_disease, 0.05).

% Define the likelihood of symptoms for Diarrheal Disease
symptom_likelihood(diarrheal_disease, frequent_loose_or_watery_stools, moderate, 0.8).
symptom_likelihood(diarrheal_disease, abdominal_pain_or_cramping, mild, 0.6).
symptom_likelihood(diarrheal_disease, nausea_or_vomiting, mild, 0.5).
symptom_likelihood(diarrheal_disease, dehydration_signs, mild, 0.4).

% Update disease probability based on a symptom
update_probability(Disease, Symptom, Severity, UpdatedProbability) :-
    prior_probability(Disease, Prior),
    symptom_likelihood(Disease, Symptom, Severity, Likelihood),
    UpdatedProbability is Prior * Likelihood.

% Diagnose Diarrheal Disease based on symptoms
diagnose_diarrheal_disease(Patient, Probability) :-
    findall(Prob, (symptom(Patient, Symptom, Severity), update_probability(diarrheal_disease, Symptom, Severity, Prob)), Probabilities),
    sum_list(Probabilities, TotalProbability),
    length(Probabilities, Count),
    Count > 0,
    Probability is TotalProbability / Count.

% Integrated Diagnosis Process
integrated_diagnosis(Patient, Disease, Advice, Treatment, Probability) :-
    (   diagnosis(Patient, Disease)
    ->  advice(Disease, Advice),
        treatment(Disease, Treatment)
    ;   diagnose_disease_bayesian(Patient, Disease, Probability),
        Probability > 0.5, % Threshold for considering the diagnosis
        advice(Disease, Advice),
        treatment(Disease, Treatment)
    ).

% Bayesian Diagnosis for Different Diseases
diagnose_disease_bayesian(Patient, Disease, Probability) :-
    (   Disease = diabetes_mellitus,
        diagnose_diabetes_mellitus(Patient, Probability)
    ;   Disease = acute_heart_failure,
        diagnose_acute_heart_failure(Patient, Probability)
    ;   Disease = urti,
        diagnose_urti(Patient, Probability)
    ;   Disease = hypertension,
        diagnose_hypertension(Patient, Probability)
    ;   Disease = flu_virus,
        diagnose_flu_virus(Patient, Probability)
    ;   Disease = pneumonia,
        diagnose_pneumonia(Patient, Probability)
    ;   Disease = covid_19,
        diagnose_covid_19(Patient, Probability)
    ;   Disease = uti,
        diagnose_uti(Patient, Probability)
    ;   Disease = migraine_headaches,
        diagnose_migraine(Patient, Probability)
    ;   Disease = diarrheal_disease,
        diagnose_diarrheal_disease(Patient, Probability)
    % Add other diseases here if needed
    ).

% Main entry point
:- initialization(start_diagnosis).

