# to install bayespy go to anaconda prompt 

# python -m pip install --upgrade pip --user

# pip install bayespy


import bayespy as bp

import numpy as np

import csv 

from colorama import init

from colorama import Fore, Back, Style

init()


# Define Parameter Enum values

#Age

ageEnum = {'SuperSeniorCitizen':0, 'SeniorCitizen':1, 'MiddleAged':2, 'Youth':3, 'Teen':4}

# Gender

genderEnum = {'Male':0, 'Female':1}

# FamilyHistory

familyHistoryEnum = {'Yes':0, 'No':1}

# Diet(Calorie Intake)

dietEnum = {'High':0, 'Medium':1, 'Low':2}

# LifeStyle

lifeStyleEnum = {'Athlete':0, 'Active':1, 'Moderate':2, 'Sedetary':3}

# Cholesterol

cholesterolEnum = {'High':0, 'BorderLine':1, 'Normal':2}

# HeartDisease

heartDiseaseEnum = {'Yes':0, 'No':1}

#heart_disease_data.csv

with open('heart_disease_data.csv') as csvfile:

    lines = csv.reader(csvfile)

    dataset = list(lines)


    data = []

for x in dataset:

    data.append([ageEnum[x[0]],genderEnum[x[1]],familyHistoryEnum[x[2]],dietEnum[x[3]],lifeStyleEnum[x[4]],cholesterolEnum[x[5]],heartDiseaseEnum[x[6]]])

# Training data for machine learning todo: should import from csv

data = np.array(data)

N = len(data)

print("the values of N is",N)

# Input data column assignment

p_age = bp.nodes.Dirichlet(1.0*np.ones(5))


print("the page",p_age)

age = bp.nodes.Categorical(p_age, plates=(N,))

print("age",age)

age.observe(data[:,0])


p_gender = bp.nodes.Dirichlet(1.0*np.ones(2))

gender = bp.nodes.Categorical(p_gender, plates=(N,))

print("gender",gender)

gender.observe(data[:,1])


p_familyhistory = bp.nodes.Dirichlet(1.0*np.ones(2))

familyhistory = bp.nodes.Categorical(p_familyhistory, plates=(N,))

print("family his",familyhistory)

familyhistory.observe(data[:,2])


p_diet = bp.nodes.Dirichlet(1.0*np.ones(3))

diet = bp.nodes.Categorical(p_diet, plates=(N,))


diet.observe(data[:,3])


p_lifestyle = bp.nodes.Dirichlet(1.0*np.ones(4))

lifestyle = bp.nodes.Categorical(p_lifestyle, plates=(N,))


lifestyle.observe(data[:,4])


p_cholesterol = bp.nodes.Dirichlet(1.0*np.ones(3))


cholesterol = bp.nodes.Categorical(p_cholesterol, plates=(N,))

cholesterol.observe(data[:,5])


# Prepare nodes and establish edges

# np.ones(2) ->HeartDisease has 2 options Yes/No

# plates(5, 2, 2, 3, 4, 3)  ->  corresponds to options present for domain values 

p_heartdisease = bp.nodes.Dirichlet(np.ones(2), plates=(5, 2, 2, 3, 4, 3))

heartdisease = bp.nodes.MultiMixture([age, gender, familyhistory, diet, lifestyle, cholesterol], bp.nodes.Categorical, p_heartdisease)

heartdisease.observe(data[:,6])

p_heartdisease.update()


# Sample Test with hardcoded values

#print("Sample Probability")

#print("Probability(HeartDisease|Age=SuperSeniorCitizen, Gender=Female, FamilyHistory=Yes, DietIntake=Medium, LifeStyle=Sedetary, Cholesterol=High)")

#print(bp.nodes.MultiMixture([ageEnum['SuperSeniorCitizen'], genderEnum['Female'], familyHistoryEnum['Yes'], dietEnum['Medium'], lifeStyleEnum['Sedetary'], cholesterolEnum['High']], bp.nodes.Categorical, p_heartdisease).get_moments()[0][heartDiseaseEnum['Yes']])


# Interactive Test 

m = 0

while m == 0:

    print("\n")

    res = bp.nodes.MultiMixture([int(input('Enter Age: ' + str(ageEnum))), int(input('Enter Gender: ' + str(genderEnum))), int(input('Enter FamilyHistory: ' + str(familyHistoryEnum))), int(input('Enter dietEnum: ' + str(dietEnum))), int(input('Enter LifeStyle: ' + str(lifeStyleEnum))), int(input('Enter Cholesterol: ' + str(cholesterolEnum)))], bp.nodes.Categorical, p_heartdisease).get_moments()[0][heartDiseaseEnum['Yes']]

    print("Probability(HeartDisease) = " +  str(res))

    #print(Style.RESET_ALL)

    m = int(input("Enter for Continue:0, Exit :1  "))