import json
import pickle
import numpy as np
import pandas as pd
import tensorflow as tf
from pymongo import MongoClient
from math import sin, cos, sqrt, atan2, radians, asin

try:
    client = MongoClient("mongodb+srv://admin:admin@early-detection.jevtz.mongodb.net/test")
    db = client["Heal"]
    client.server_info()
    print("DB accessed !")

except Exception as e:
    print("*************************************************************")
    print(e)
    print("*************************************************************")
    
##################################### PARAMS #########################################

risk_level_mapping = {
                    0 : "situation doesn't exist",
                    1 : "Manageable",
                    2 : "Normal Healthy",
                    3 : "Risky"
                    }

excercise_plan_mapping = {
                        0: 'No exercises',
                        1: 'sleep,Rest, Yoga',
                        2: 'jogging, walking, Running',
                        3: 'swimming, Breathing exercises'
                        }

diet_plan_mapping = {
                    0: 'No food plan',
                    1: 'Low cholesterol food',
                    2: 'Balance Diet which contain all the nutrients',
                    3: 'Hot Beverages & Meals Food rich in potassium'
                    }

##################################### LOAD MODELS #########################################

risk_model = tf.keras.models.load_model('weights/risk_prediction.h5')
risk_model.compile(
                optimizer='adam', 
                loss='categorical_crossentropy', 
                metrics=[
                        tf.keras.metrics.CategoricalAccuracy(),
                        tf.keras.metrics.Precision(),
                        tf.keras.metrics.Recall(),
                        tf.keras.metrics.AUC()
                        ])

with open('weights/scalar - risk_prediction.pkl', 'rb') as f:
    risk_scalar = pickle.load(f)


diet_workout_model = tf.keras.models.load_model('weights/diet_workout_prediction.h5')
diet_workout_model.compile(
                        loss='sparse_categorical_crossentropy',
                        optimizer='adam',
                        metrics=['accuracy']
                        )

with open('weights/scalar - diet_workout_prediction.pkl', 'rb') as f:
    diet_workout_scalar = pickle.load(f)

##################################### PREDICTION #########################################

def risk_level_prediction():
    health_collection = db["health"]
    health_data = pd.DataFrame(list(health_collection.find()))
    sample_data = health_data[-50:][["HeartRate","BloodOxygenLevel","BodyTemperature"]]

    HeartRateArr = sample_data["HeartRate"].values.tolist()
    BloodOxygenLevelArr = sample_data["BloodOxygenLevel"].values.tolist()
    BodyTemperatureArr = sample_data["BodyTemperature"].values.tolist()

    Xinf = sample_data.values
    Xinf = risk_scalar.transform(Xinf)
    Xinf = Xinf.reshape(1, 50, 3)

    risk_level = risk_model.predict(Xinf).squeeze()
    risk_level = np.argmax(risk_level)
    risk_level = risk_level_mapping[risk_level]
    return risk_level, HeartRateArr, BloodOxygenLevelArr, BodyTemperatureArr

def diet_workout_prediction():
    health_collection = db["health"]
    health_data = pd.DataFrame(list(health_collection.find()))
    sample_data = health_data[-50:][["HeartRate","BloodOxygenLevel","BodyTemperature"]]

    HeartRateArr = sample_data["HeartRate"].values.tolist()
    BloodOxygenLevelArr = sample_data["BloodOxygenLevel"].values.tolist()
    BodyTemperatureArr = sample_data["BodyTemperature"].values.tolist()

    Xinf = sample_data.values
    Xinf = diet_workout_scalar.transform(Xinf)

    exercise_plan_idxs, diet_plan_idxs = diet_workout_model.predict(Xinf)
    exercise_plan_idxs = exercise_plan_idxs.squeeze()
    diet_plan_idxs = diet_plan_idxs.squeeze()

    exercise_plan_idxs = np.argmax(exercise_plan_idxs, axis=1)
    diet_plan_idxs = np.argmax(diet_plan_idxs, axis=1)

    exercise_plans = [excercise_plan_mapping[idx] for idx in exercise_plan_idxs]
    diet_plans = [diet_plan_mapping[idx] for idx in diet_plan_idxs]

    mode_exercise_plan = max(set(exercise_plans), key=exercise_plans.count)
    mode_diet_plan = max(set(diet_plans), key=diet_plans.count)
    return mode_exercise_plan, mode_diet_plan, HeartRateArr, BloodOxygenLevelArr, BodyTemperatureArr

def haversine(p1, p2):
    lat1, lon1 = p1
    lat2, lon2 = p2
    
    lat1 = float(lat1)
    lon1 = float(lon1)
    
    lat2 = float(lat2)
    lon2 = float(lon2)

    # convert decimal degrees to radians 
    lon1, lat1, lon2, lat2 = map(radians, [lon1, lat1, lon2, lat2])
 
    # haversine formula 
    dlon = lon2 - lon1 
    dlat = lat2 - lat1 
    
    a = sin(dlat/2)**2 + cos(lat1) * cos(lat2) * sin(dlon/2)**2
    
    c = 2 * asin(sqrt(a)) 
    r = 6371 # Radius of earth in kilometers. Use 3956 for miles. Determines return value units.
    distance_km = c * r
    return distance_km

def pharmacy_recommendation(user_location, medicine_list):
    user_location = user_location.split(',')
    user_location = (float(user_location[0]), float(user_location[1]))

    medicine_list = medicine_list.split(',')
    medicine_list = [medicine.strip() for medicine in medicine_list]
    
    pharmacy_collection = db["pharmacy"]
    df_pharmacy = pd.DataFrame(list(pharmacy_collection.find()))

    df_pharmacy[['latitude', 'longitude']] = df_pharmacy['location'].str.split(',', 1, expand=True)
    del df_pharmacy['location']

    df_pharmacy['latitude'] = df_pharmacy['latitude'].str.strip()
    df_pharmacy['longitude'] = df_pharmacy['longitude'].str.strip()

    df_pharmacy['latitude'] = df_pharmacy['latitude'].astype(float)
    df_pharmacy['longitude'] = df_pharmacy['longitude'].astype(float)

    df_pharmacy["medicine"] = df_pharmacy["medicine"].str.split(',')
    
    df_pharmacy = df_pharmacy.loc[df_pharmacy['medicine'].apply(lambda x: set(medicine_list).issubset([m.strip() for m in x]))]
    df_pharmacy['distance'] = df_pharmacy.apply(lambda row: haversine(user_location, (row['latitude'], row['longitude'])), axis=1)
    df_pharmacy.sort_values(by=['distance'], inplace=True)
    
    df_pharmacy = df_pharmacy.head(5)
    df_pharmacy = df_pharmacy[['name', 'distance', 'medicine']]
    df_pharmacy['distance'] = df_pharmacy['distance'].apply(lambda x: round(x, 2))
    rec_list = []
    for _, row in df_pharmacy.iterrows():
        rec_json = {}
        rec_json["name"] = f"{row['name']}"
        rec_json["distance"] = f"{row['distance']} km"
        rec_json["available medicines"] = f"{row['medicine']}"
        rec_list.append(rec_json)
    return rec_list