from fastapi import FastAPI, Request, Response
import json
import uuid
import algo as rezon_algo
import tensorflow as tf
from tensorflow.keras.models import model_from_json
import tensorflow.keras as keras
import numpy as np
import os

# This is a demo API for what we would actually generate for users
# Example URL: https://javier-is-cool.cloud.rezon.ai

# TODO: in the MVP we want to allow the user to pass param names to replace the defailt aliases (Variable 1, etc) --
# this could be done via query params or the user could send it in the request payload


app = FastAPI()

# -------- MODEL CONSTRUCTION --------
# Note: we want to do this before any requests can get to the root handler so that we dont re-construct the model on every request

# Load json and create model
with open('./models/testModel.json', 'r') as json_file:
    loaded_user_model_in_json = json_file.read()

# Load the initial model
loaded_user_model = model_from_json(loaded_user_model_in_json)

# Load weights into the new model
loaded_user_model.load_weights("./models/testModel.h5")

@app.post("/")
async def root(request: Request, response: Response):
    # Parse our inputs...
    payload = await request.json()
    model_input = payload.get('input') if payload.get('input') is not None else False
    should_explain = True if payload.get('include_explanations') is None else payload.get('include_explanations')


    # If we did not get a model input we want to throw a 400 error
    if not model_input:
        # We should add a better function for setting the response status code in fastAPI...currently feels too verbose
        response.status_code = 400
        return {"error": f"INVALID_MODEL_INPUT", "status": response.status_code, "message": "No model input was provided"}

    # Running the model and getting outputs could fail...(ex: an invalid input is sent)
    try:
        # Attempt to get outputs...
        # Get outputs (we would then send this to the client)
        data = loaded_user_model.predict(model_input)
    except Exception as e:
        # Logging should be much better but this is fine for now
        # TODO: We should have a logger that logs to a database and is shown in the Rezon UI (we could also have a /logs endpoint)
        print(f"[ERROR]: {e}")
        response.status_code = 400
        # Output an error response
        # TODO: At some point we should have standardized error responses for each `ERROR_CODE` that way
        # we can have documentation on what each error code means (like FAILED_MODEL_PERDICTION likely means ur inputs were invalid)
        return {"error": f"FAILED_MODEL_PERDICTION", "status": response.status_code, "message": f"Error occurred while running model perdiction: {e}"}

    # -------  SAVE INPUTS TO THE DATABASE FILE ------- -- Note: in the real mvp this will all be cached or "dumped" into a database for now this will be done by simply caching inside of a .json or .txt file
    # We do this because explanations are potentially very expensive and we want to make sure that we don't have to recompute them

    # Sanitize our data and convert to json
    sanitized_data = json.dumps(data.tolist())
    # Check to make sure that the user wants explanations
    # If they do, we want to do the computation and send it back to them
    if should_explain:
        # ------- RUN REZON-CORE ALGO -------
        # (This is where we want to check if the user wants an explanation or not)
        # Remember, this computation couild likely be expensive and we want to avoid running it if we don't need to

        # (1) - Study the model and compute the base info we need to obtain explainations...
        specs, arr_types_afs, arr_weights, arr_biases = rezon_algo.study_model(loaded_user_model)

        # (2) - Run the explainations algo and get the output
        output_explanations = rezon_algo.run_explanations(model_input, loaded_user_model, specs, arr_types_afs, arr_weights, arr_biases)

        # TODO: this logging is decent but should be improved in MVP
        print("[NEW_EXPLANATION]")

        for i in range(len(model_input)):
            print("Input Values: ", output_explanations.input_values[i])
            print("Output Pred. by NN: ", output_explanations.output_pred_NN[i])
            print("Output Pred. by LF: ", output_explanations.output_pred_LF[i])
            print()
            for j in range(specs.n_outs):
                print("  For Output Dimension ", j+1)
                print("     Linear Func. Dep. Term: ", output_explanations.LF_dep_terms[i][j])
                print("     Linear Func. Indep. Term: ", output_explanations.LF_indep_term[i][j])
                print("     Input Contributions: ", output_explanations.input_contribs[i][j])
                print("     Text Explanations: ")
                print("         ", output_explanations.text_explanation[i][j][0])
                print("         ", output_explanations.text_explanation[i][j][1])
                print("         ", output_explanations.text_explanation[i][j][2])
                print("Reference ID: ", output_explanations.reference_id[i])
                print()

        print("[END_EXPLANATION]")

        # Generate a unique explanation uuid (this will be useful for us to identify the explanation in the future)
        explanation_uuid = str(uuid.uuid4())

        # (3) - Sanitize the output and send it back to the client
        sanitized_explanations = {
            "input_values": output_explanations.input_values,
            "input_contributions": output_explanations.input_contribs,
            "input_dependencies": output_explanations.LF_dep_terms,
            "independent_term": output_explanations.LF_indep_term,
            "neural_network_output": output_explanations.output_pred_NN,
            "linear_function_output": output_explanations.output_pred_LF,
            "text_explanations": output_explanations.text_explanation, # TODO: we should change the data format for text_explanations (they are hard to parse when all grouped in one giant blob)
            "reference_id": explanation_uuid
        }

        return {"output": sanitized_data, "explanation": sanitized_explanations}

    return {"output": sanitized_data}
