# Import the libraries
import os
import time
import json
import uuid
import joblib
import pandas as pd
import gradio as gr
import math

from huggingface_hub import CommitScheduler
from pathlib import Path


# Run the training script placed in the same directory as app.py
os.system("python train.py")

# The training script will train and persist a linear regression
# model with the filename 'model.joblib'




# Load the freshly trained model from disk
insurance_charge_predictor = joblib.load('model.joblib')

# Prepare the logging functionality
log_file = Path("logs/") / f"data_{uuid.uuid4()}.json"
log_folder = log_file.parent

scheduler = CommitScheduler(
    repo_id="anirudhabokil/insurance-charge-mlops-logs",  # provide a name "insurance-charge-mlops-logs" for the repo_id
    repo_type="dataset",
    folder_path=log_folder,
    path_in_repo="data",
    every=2
)

# Define the predict function which will take features, convert to dataframe and make predictions using the saved model
# the functions runs when 'Submit' is clicked or when a API request is made



# Set up UI components for input and output
age = gr.Number(label="Age")
bmi = gr.Number(label="BMI")
children = gr.Number(label="Children")
sex = gr.Dropdown(['male','female'], label="Sex")
smoker = gr.Dropdown(['yes','no'], label="Smoker")
region = gr.Dropdown(['southwest','southeast','northwest','northeast'], label="Region")

model_output = gr.Label(label="Insurance Charge")

def predict_insurance_charge(age, bmi, children, sex, smoker, region):
  sample = {
      'age': age,
      'bmi': bmi,
      'children': children,
      'sex': sex,
      'smoker': smoker,
      'region': region
  }
  df = pd.DataFrame([sample])
  print(sample)
  prediction = insurance_charge_predictor.predict(df).tolist()
  #print(prediction)


  # While the prediction is made, log both the inputs and outputs to a  log file
  # While writing to the log file, ensure that the commit scheduler is locked to avoid parallel
  # access

  with scheduler.lock:
      with log_file.open("a") as f:
          f.write(json.dumps(
              {
                  'age': age,
                  'bmi': bmi,
                  'children': children,
                  'sex': sex,
                  'smoker': smoker,
                  'region': region,
                  'prediction': prediction[0]
              }
          ))
          f.write("\n")
  return prediction[0]


# Create the gradio interface, make title "HealthyLife Insurance Charge Prediction"
demo = gr.Interface(fn=predict_insurance_charge,
                    inputs=[age, bmi, children, sex, smoker, region],
                    outputs=model_output,
                    title="HealthyLife Insurance Charge Prediction",
                    description="This API allows you to predict insurance charge",
                    flagging_mode="auto",
                    concurrency_limit=8)

# Launch with a load balancer
demo.queue()
demo.launch(share=True, debug=True)