from flask import Flask, jsonify, request, send_from_directory
from flask_cors import CORS
import pandas as pd
import os
import pickle
from transformers import DistilBertTokenizer
import torch
import numpy as np
# from your_model import load_model, predict_response  # assuming these functions are implemented

app = Flask(__name__)

# Load your machine learning model
# ai_only = load_model('ai_model_path.pkl')
# hitl = load_model('hitl_model_path.pkl')

# Load your questions and responses
questions_df = pd.read_csv('data/questions.csv')
# responses_df = pd.read_csv('data/responses.csv')

CORS(app, resources={r"/api/*": {"origins": "*"}})

react_folder= 'frontend'
directory= os.getcwd()+ f'/{react_folder}/build/static/'

def load_model(name):
    print("Loading model")

    with open(f'{name}_model.pkl', 'rb') as f:
        loaded_model = pickle.load(f)
        print("model loaded successfully")
        return loaded_model


def make_prediction(model, data):
    model= load_model(model)

    tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
    inputs = tokenizer(data, return_tensors="pt", truncation=True, padding=True, max_length=136)

    # Forward pass through the model to obtain logits
    with torch.no_grad():
        outputs = model(**inputs)

    # Get predicted probabilities using softmax function
    probs = torch.softmax(outputs.logits, dim=1)

    # Get predicted class (0 or 1)
    predicted_class = torch.argmax(probs, dim=1).item()

    # Map predicted class to corresponding label
    predicted_label = "AI" if predicted_class == 0 else "Human"
    print("Predicted Label:", predicted_label)
    print("Probability:", probs)

    return [predicted_label, max(probs[0][0], probs[0][1])]





@app.route('/DilemmaAI', endpoint='func1')
def index():
    path = os.getcwd()+f'/{react_folder}/build'
    print(path)
    return send_from_directory(directory=path, path='index.html')

@app.route('/static/<folder>/<file>')
def css(folder,file):
    path = folder+'/'+file
    return send_from_directory(directory= directory, path= path)

@app.route('/assets/<file>', endpoint='func2')
def css(folder,file):
    directory = os.getcwd()+f'/{react_folder}/build/assets'
    path = file
    return send_from_directory(directory=directory, path= path)

# @app.route('/api/questions', methods=['GET'])
# def get_questions():
#     # Sending a list of questions along with their IDs
#     questions = questions_df[['question_id', 'question']].to_dict(orient='records')
#     return jsonify(questions)


@app.route('/api/random-question', methods=['GET'])
def get_random_question():
    # Fetching a random question from the dataframe
    random_question = questions_df.sample().iloc[0][1]
    print(random_question)
    return jsonify(random_question)

@app.route('/api/predict', methods=['POST'])
def predict():
    # Get response from the request
    data = request.json
    response_text = data['response']
    model_type = data['mode']
    print('input = ', response_text)
    print('mode=',model_type)

    #make_prediction(model_type,response_text)
    # Perform prediction
    # prediction_label, confidence = predict_response(model, response_text)
    prediction_label, confidence = make_prediction(model_type, response_text)
    if not prediction_label:
        prediction_label = 'AI'
    if not confidence:
        confidence = 0

    return jsonify({
        'prediction': prediction_label,
        'confidence': f'{confidence*100:.2f}%'
    })


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5001)