import os
import io
import uvicorn
import pandas as pd
import tensorflow as tf
from fastapi.responses import RedirectResponse
import numpy as np
from fastapi import FastAPI, File, UploadFile
from starlette.middleware.cors import CORSMiddleware
from tensorflow.keras.models import load_model
from PIL import Image
from tensorflow.keras.applications.vgg16 import preprocess_input


df = pd.read_json("plant_details.json")
# Initialize FastAPI app
app = FastAPI()


@app.get("/")
async def root():
    return RedirectResponse(url="/docs")


# Add CORS middleware to allow cross-origin requests
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Adjust this to limit which domains can make requests
    allow_methods=["*"],
    allow_headers=["*"],
)
# Check if GPU is available and use it for TensorFlow operations
gpus = tf.config.experimental.list_physical_devices("GPU")
if gpus:
    try:
        # Currently, memory growth needs to be the same across GPUs
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        logical_gpus = tf.config.experimental.list_logical_devices("GPU")
        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
    except RuntimeError as e:
        # Memory growth must be set before GPUs have been initialized
        print(e)

# Load the pre-trained model during app startup
model = load_model(
    "efficientnetb3-Plant Village Disease-99.71.h5"
)  # Replace with your actual model path


model_27 = load_model("efficientnetb3-Botanical-Garden-v1-98.88.h5")

model_plant_leaves = load_model("efficientnetb3-Plant-leaves-v1-0.97.h5")

model_merge_plant_leaves = load_model("efficientnetb3-Plant-leaves-merge-v2-0.96.h5")


model_demo = load_model("my_model.h5")
# Define class labels
class_labels = {
    0: "Apple___Apple_scab",
    1: "Apple___Black_rot",
    2: "Apple___Cedar_apple_rust",
    3: "Apple___healthy",
    4: "Blueberry___healthy",
    5: "Cherry_(including_sour)___Powdery_mildew",
    6: "Cherry_(including_sour)___healthy",
    7: "Corn_(maize)___Cercospora_leaf_spot Gray_leaf_spot",
    8: "Corn_(maize)___Common_rust_",
    9: "Corn_(maize)___Northern_Leaf_Blight",
    10: "Corn_(maize)___healthy",
    11: "Grape___Black_rot",
    12: "Grape___Esca_(Black_Measles)",
    13: "Grape___Leaf_blight_(Isariopsis_Leaf_Spot)",
    14: "Grape___healthy",
    15: "Orange___Haunglongbing_(Citrus_greening)",
    16: "Peach___Bacterial_spot",
    17: "Peach___healthy",
    18: "Pepper,_bell___Bacterial_spot",
    19: "Pepper,_bell___healthy",
    20: "Potato___Early_blight",
    21: "Potato___Late_blight",
    22: "Potato___healthy",
    23: "Raspberry___healthy",
    24: "Soybean___healthy",
    25: "Squash___Powdery_mildew",
    26: "Strawberry___Leaf_scorch",
    27: "Strawberry___healthy",
    28: "Tomato___Bacterial_spot",
    29: "Tomato___Early_blight",
    30: "Tomato___Late_blight",
    31: "Tomato___Leaf_Mold",
    32: "Tomato___Septoria_leaf_spot",
    33: "Tomato___Spider_mites Two-spotted_spider_mite",
    34: "Tomato___Target_Spot",
    35: "Tomato___Tomato_Yellow_Leaf_Curl_Virus",
    36: "Tomato___Tomato_mosaic_virus",
    37: "Tomato___healthy",
}

class_labels_27 = {
    0: "Bamboo",
    1: "Bambusa Ventricosa",
    2: "Bougainvillea",
    3: "Butterfly_ginger_lily",
    4: "Chinese_palm",
    5: "Clove",
    6: "Dalchini",
    7: "Grape",
    8: "Hing",
    9: "Kajubadam",
    10: "Kathal",
    11: "Kola_halodhi",
    12: "Pimenta_racemosa_spice",
    13: "Saururaceae",
    14: "Tamarindus",
    15: "Vitaceae",
    16: "Zingiberaceae",
    17: "Amla",
    18: "Elaichi",
    19: "Isobgul",
    20: "Kendu",
    21: "Malvaceae",
    22: "Mousambi",
    23: "Owtenga",
    24: "Plantago",
    25: "Pologonoceae",
    26: "Sida_acuta",
}

class_plant_leaves = {
    0: "Alstonia Scholaris diseased (P2a)",
    1: "Alstonia Scholaris healthy (P2b)",
    2: "Arjun diseased (P1a)",
    3: "Arjun healthy (P1b)",
    4: "Bael diseased (P4b)",
    5: "Basil healthy (P8)",
    6: "Chinar diseased (P11b)",
    7: "Chinar healthy (P11a)",
    8: "Gauva diseased (P3b)",
    9: "Gauva healthy (P3a)",
    10: "Jamun diseased (P5b)",
    11: "Jamun healthy (P5a)",
    12: "Jatropha diseased (P6b)",
    13: "Jatropha healthy (P6a)",
    14: "Lemon diseased (P10b)",
    15: "Lemon healthy (P10a)",
    16: "Mango diseased (P0b)",
    17: "Mango healthy (P0a)",
    18: "Pomegranate diseased (P9b)",
    19: "Pomegranate healthy (P9a)",
    20: "Pongamia Pinnata diseased (P7b)",
    21: "Pongamia Pinnata healthy (P7a)",
}

class_plant_leaves_merged = {
    0: "Alstonia Scholaris diseased (P2a)",
    1: "Alstonia Scholaris healthy (P2b)",
    2: "Arjun diseased (P1a)",
    3: "Arjun healthy (P1b)",
    4: "Bael diseased (P4b)",
    5: "Basil healthy (P8)",
    6: "Chinar diseased (P11b)",
    7: "Chinar healthy (P11a)",
    8: "Curry Patta",
    9: "Gauva diseased (P3b)",
    10: " Gauva healthy (P3a)",
    11: " Jamun diseased (P5b)",
    12: " Jamun healthy (P5a)",
    13: " Jatropha diseased (P6b)",
    14: " Jatropha healthy (P6a)",
    15: " Kola Jamun",
    16: " Lemon",
    17: " Lemon diseased (P10b)",
    18: " Lemon healthy (P10a)",
    19: " Mango diseased (P0b)",
    20: " Mango healthy (P0a)",
    21: " Pomegranate diseased (P9b)",
    22: " Pomegranate healthy (P9a)",
    23: " Pongamia Pinnata diseased (P7b)",
    24: " Pongamia Pinnata healthy (P7a)",
    25: " Shewali",
}
class_demo_mapping = {0: "Curry_leaf", 1: "Kola_Jamun", 2: "Lemon", 3: "shewali"}


# Define a route to accept single image uploads and make predictions
@app.post("/predict/")
async def predict_single_image(file: UploadFile):
    try:
        # Read the uploaded image
        image = await file.read()
        img = Image.open(io.BytesIO(image))

        # Model Params
        img_size = (224, 224)
        channels = 3  # either BGR or Grayscale
        color = "rgb"
        img_shape = (img_size[0], img_size[1], channels)

        # Preprocess the image to match the input requirements of your model
        img = img.resize((224, 224))  # Adjust the size as needed
        img = np.asarray(img)
        img = preprocess_input(img)

        # Make a prediction using your pre-trained model
        predictions = model.predict(np.expand_dims(img, axis=0))

        # Convert prediction indices to class labels
        predicted_class_index = np.argmax(predictions)
        predicted_class_label = class_labels.get(predicted_class_index, "Unknown")

        # You can format the predictions as needed and return them
        result = {
            "predicted_class": predicted_class_label,
            "class_probabilities": predictions.tolist(),
        }

        return result

    except Exception as e:
        return {"error": str(e)}


@app.post("/predict-self-dataset-27/")
async def predict_single_image_self(file: UploadFile):
    try:
        # Read the uploaded image
        image = await file.read()
        img = Image.open(io.BytesIO(image))

        # Model Params
        img_size = (224, 224)
        channels = 3  # either BGR or Grayscale
        color = "rgb"
        img_shape = (img_size[0], img_size[1], channels)

        # Preprocess the image to match the input requirements of your model
        img = img.resize((224, 224))  # Adjust the size as needed
        img = np.asarray(img)
        img = preprocess_input(img)

        # Make a prediction using your pre-trained model
        predictions = model_27.predict(np.expand_dims(img, axis=0))

        # Convert prediction indices to class labels
        predicted_class_index = np.argmax(predictions)
        predicted_class_label = class_labels_27.get(predicted_class_index, "Unknown")

        matched_rows = df[df["id"] == predicted_class_index]
        matched_rows_json = matched_rows.to_json(orient="records")

        # You can format the predictions as needed and return them
        result = {
            "predicted_class": predicted_class_label,
            "class_probabilities": predictions.tolist(),
            "details": matched_rows_json,
        }

        return result

    except Exception as e:
        return {"error": str(e)}


def h():
    return "Hello World"


@app.post("/predict-plant-dataset-merge-25/")
async def predict_single_image_self(file: UploadFile):
    try:
        # Read the uploaded image
        image = await file.read()
        img = Image.open(io.BytesIO(image))

        # Model Params
        img_size = (224, 224)
        channels = 3  # either BGR or Grayscale
        color = "rgb"
        img_shape = (img_size[0], img_size[1], channels)

        # Preprocess the image to match the input requirements of your model
        img = img.resize((224, 224))  # Adjust the size as needed
        img = np.asarray(img)
        img = preprocess_input(img)

        # Make a prediction using your pre-trained model
        predictions = model_plant_leaves.predict(np.expand_dims(img, axis=0))

        # Convert prediction indices to class labels
        predicted_class_index = np.argmax(predictions)
        predicted_class_label = class_plant_leaves.get(predicted_class_index, "Unknown")

        matched_rows = df[df["id"] == predicted_class_index]
        matched_rows_json = matched_rows.to_json(orient="records")

        if not matched_rows.empty:
            matched_rows_json = matched_rows.to_json(orient="records")

            result = {
                "predicted_class": predicted_class_label,
                "class_probabilities": predictions.tolist(),
                "details": matched_rows_json,
            }
        else:
            result = {
                "predicted_class": predicted_class_label,
                "class_probabilities": predictions.tolist(),
            }

        return result

    except Exception as e:
        return {"error": str(e)}


@app.post("/predict-plant-dataset-22/")
async def predict_single_image_self(file: UploadFile):
    try:
        # Read the uploaded image
        image = await file.read()
        img = Image.open(io.BytesIO(image))

        # Model Params
        img_size = (224, 224)
        channels = 3  # either BGR or Grayscale
        color = "rgb"
        img_shape = (img_size[0], img_size[1], channels)

        # Preprocess the image to match the input requirements of your model
        img = img.resize((224, 224))  # Adjust the size as needed
        img = np.asarray(img)
        img = preprocess_input(img)

        # Make a prediction using your pre-trained model
        predictions = model_merge_plant_leaves.predict(np.expand_dims(img, axis=0))

        # Convert prediction indices to class labels
        predicted_class_index = np.argmax(predictions)
        predicted_class_label = class_plant_leaves_merged.get(
            predicted_class_index, "Unknown"
        )

        matched_rows = df[df["Plant Name"] == predicted_class_label.lower()]
        matched_rows_json = matched_rows.to_json(orient="records")

        if not matched_rows.empty:
            matched_rows_json = matched_rows.to_json(orient="records")

            result = {
                "predicted_class": predicted_class_label,
                "class_probabilities": predictions.tolist(),
                "details": matched_rows_json,
            }
        else:
            result = {
                "predicted_class": predicted_class_label,
                "class_probabilities": predictions.tolist(),
            }

        return result

    except Exception as e:
        return {"error": str(e)}


@app.post("/predict-demo/")
async def predict_single_image_self_demo(file: UploadFile):
    try:
        # Read the uploaded image
        image = await file.read()
        img = Image.open(io.BytesIO(image))

        # Model Params
        img_size = (256, 256)
        channels = 3  # either BGR or Grayscale
        color = "rgb"
        img_shape = (img_size[0], img_size[1], channels)

        # Preprocess the image to match the input requirements of your model
        img = img.resize((256, 256))  # Adjust the size as needed
        img = np.asarray(img)
        img = preprocess_input(img)

        # Make a prediction using your pre-trained model
        predictions = model_demo.predict(np.expand_dims(img, axis=0))

        # Convert prediction indices to class labels
        predicted_class_index = np.argmax(predictions[0])
        predicted_class_label = class_demo_mapping.get(predicted_class_index, "Unknown")

        matched_rows = df[df["Plant Name"] == predicted_class_label.lower()]
        matched_rows_json = matched_rows.to_json(orient="records")

        if not matched_rows.empty:
            matched_rows_json = matched_rows.to_json(orient="records")

            result = {
                "predicted_class": predicted_class_label,
                "class_probabilities": predictions.tolist(),
                "details": matched_rows_json,
            }
        else:
            result = {
                "predicted_class": predicted_class_label,
                "class_probabilities": predictions.tolist(),
            }

        return result

    except Exception as e:
        return {"error": str(e)}


def hi():
    return "si"


if __name__ == "__main__":
    uvicorn.run(app, port=8000)