import streamlit as st
import joblib
import numpy as np
import pandas as pd
import plotly.express as px


# Load the models
scaler = joblib.load('scaler.joblib')
model = joblib.load('best_rf.joblib')

sensory_feature_column = ['sweet',
 'sour',
 'bitter',
 'aromatic_impact',
 'fruity_impact',
 'art_sweetener_chewing',
 'chews',
 'art_sweetener_after',
 'stickiness_with_fingers',
 'springiness_with_fingers',
 'abrasive',
 'hardness_with_molars',
 'uniformity_of_bite',
 'gritty',
 'springiness_during_chew',
 'cohesiveness_of_mass',
 'moistness_of_mass',
 'toothsticking',
 'toothpacking',
 'adhesiveness_to_molars',
 'oily_mouthcoating']

def scale_output(prediction_matrix):
    # Calculate the sum for each row
    row_sums = np.sum(prediction_matrix, axis=1)

    # Calculate the surplus or lack from 1 for each row
    delta = 1 - row_sums

    # Adjust each row proportionally
    adjusted_matrix = prediction_matrix + delta[:, np.newaxis] * (prediction_matrix / row_sums[:, np.newaxis])
    print(f'scaled probability: {adjusted_matrix.sum(axis=1)}')
    return adjusted_matrix

# Streamlit App
def main():
    st.title("Consumer liking score distribution prediction")
    st.markdown("Please adjust the values of 21 sensory attributes (assumed to have the range from 1 - 10), and see the liking score distribution below")
    
    with st.expander("Feature importance heatmap of sensory attributes on liking score"):
        st.image("mmr-corr.png")
    
    col1, col2, col3 = st.columns(3)
    
    with col1:
        sweet = st.slider('sweet', 0.00, 10.00, 0.01)
        aromatic_impact = st.slider('aromatic impact', 0.00, 10.00, 0.01)
        chew = st.slider('chew', 0.00, 10.00, 0.01)
        springiness_with_fingers = st.slider('springiness_with_fingers', 0.00, 10.00, 0.01)
        uniformity_of_bite = st.slider('uniformity_of_bite', 0.00, 10.00, 0.01)
        cohesiveness_of_mass = st.slider('cohesiveness_of_mass', 0.00, 10.00, 0.01)
        toothpacking = st.slider('toothpacking', 0.00, 10.00, 0.01)
        
    with col2:
        sour = st.slider('sour', 0.00, 10.00, 0.01)
        fruity_impact = st.slider('fruity impact',0.00, 10.00, 0.01)
        art_sweetener_after = st.slider('art_sweetener_after', 0.00, 10.00, 0.01)
        abrasiveness = st.slider('abrasiveness', 0.00, 10.00, 0.01)
        gritty = st.slider('gritty', 0.00, 10.00, 0.01)
        moisture_of_mass = st.slider('moisture_of_mass', 0.00, 10.00, 0.01)
        adhesiveness_to_molars = st.slider('adhesiveness_to_molars', 0.00, 10.00, 0.01)


    with col3:
        bitter = st.slider('bitter',  0.00, 10.00, 0.01)
        art_sweetener_chewing = st.slider('art_sweetener_chewing',0.00, 10.00, 0.01)
        stickiness_with_fingers = st.slider('stickiness_with_fingers', 0.00, 10.00, 0.01)
        hardness_with_molars = st.slider('hardness_with_molars', 0.00, 10.00, 0.01)
        springiness_during_chew = st.slider('springiness_during_chew', 0.00, 10.00, 0.01)
        toothsticking = st.slider('toothsticking', 0.00, 10.00, 0.01)
        oily_mouthcoating = st.slider('oily_mouthcoating', 0.00, 10.00, 0.01)


    data_input = [sweet, sour, bitter, aromatic_impact, fruity_impact, art_sweetener_chewing,chew,
                 art_sweetener_after, stickiness_with_fingers,  springiness_with_fingers , abrasiveness,  hardness_with_molars , uniformity_of_bite, gritty,
                 springiness_during_chew, cohesiveness_of_mass, moisture_of_mass , toothsticking , toothpacking, adhesiveness_to_molars, oily_mouthcoating]

    test_df = pd.DataFrame(np.array(data_input).reshape(1, 21))
    test_df.columns = sensory_feature_column
    
    scaled_data_input = scaler.transform(test_df)
    handout_prediction = scale_output(model.predict(scaled_data_input))
    
    handout_pred_df = pd.DataFrame(handout_prediction)
    handout_pred_df.columns = [str(i) for i in range(1,10)]
    handout_pred_df['prod'] = 'prediction'
    handout_pred_df = handout_pred_df.melt('prod')

    # st.write(f"prediction: {handout_prediction  }")
    st.write(f"sum of probability: {np.round(handout_prediction.sum(), 2)}")
    
    
    
    fig = px.histogram(data_frame=handout_pred_df, x= 'variable',y='value', nbins= 9,
                     text_auto=True, title="Probability prediction of liking score from sensory attributes",
                labels={"value": "probability",
                      "variable": "liking score"})
                
    st.plotly_chart(fig, theme='streamlit')
    

if __name__ == "__main__":
    main()