import cv2
import numpy as np
import streamlit as st
from datasets import load_dataset
from PIL import Image
from matplotlib.pyplot import imshow
from transformers import pipeline
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
import json

def load_dataset_from_Hugging_Face():
  reloaded_dataset_thermal = load_dataset("gokulraja17/rice-thermal-demo")
  reloaded_dataset_rgb = load_dataset("gokulraja17/rice-rgb-demo2")
  return reloaded_dataset_thermal,reloaded_dataset_rgb

def load_model_from_Hugging_Face(type):
    if(type=="thermal"):
        feature_extractor = AutoFeatureExtractor.from_pretrained("flagship/convnext-fine-tuned-thermal-new_demo")
        model = AutoModelForImageClassification.from_pretrained("flagship/convnext-fine-tuned-thermal-new_demo")      
    else:
        feature_extractor = AutoFeatureExtractor.from_pretrained("gokulraja17/convnext-fine-tuned-rgb-demo2")
        model = AutoModelForImageClassification.from_pretrained("gokulraja17/convnext-fine-tuned-rgb-demo2")
    return feature_extractor,model

def get_image_and_label(reloaded_dataset_thermal,reloaded_dataset_rgb,index):
  url_thermal = reloaded_dataset_thermal["test"][index]
  url_rgb = reloaded_dataset_rgb["test"][index]
  return url_thermal["image"],url_thermal["label"],url_rgb["image"],url_rgb["label"]

def create_pipeline(_feature_extractor,_model):

  pipe = pipeline("image-classification",
                model=_model,
                feature_extractor=_feature_extractor)
  return pipe

def get_output_label(result):
  max = -999
  label = ""
  for i in result:
    if(max<i['score']):
      max = i['score']
      label = i["label"]
  return label

def consolidate_score(thermal_result,rgb_result):
  consol_score_1=[]
  consol_score_2=[]
  consol_score=[]
  for i in range(0,len(thermal_result)):
    consol_score_1.append(list(thermal_result[i].values()))
    consol_score_2.append(list(rgb_result[i].values()))
  temp_label=""
  for i in range(0,len(thermal_result)):
    temp_label=consol_score_1[i][1]
    for j in range(0,len(rgb_result)):
      if(temp_label==consol_score_2[j][1]):
        consol_score.append([temp_label,consol_score_1[i][0],consol_score_2[j][0]])
  output_consol=[]
  for i in consol_score:
    output_consol.append({'score':(((i[2]/100)*35)+((i[1]/100)*65)),'label':i[0]})
  return output_consol
thermal_check = st.checkbox("Thermal Input")
rgb_check  = st.checkbox("RGB Input")
both_check  = st.checkbox("Both")
if(thermal_check or rgb_check):
    if(thermal_check):
    
        st.write("Thermal Image file uploader...")
        uploaded_file = st.file_uploader("Choose a thermal image file", type=["jpg","png"])

        if uploaded_file is not None:
            input_image = Image.open(uploaded_file)
            st.write("Input Image...")
            st.image(input_image)
            dataset_thermal,dataset_rgb = load_dataset_from_Hugging_Face()
            labels=dataset_thermal["train"].features["label"].names
            id2label={k:v for k,v in enumerate(labels)}
            label2id = {v:k for k,v in enumerate(labels)}

            feature_extractor_thermal,model_thermal = load_model_from_Hugging_Face("thermal")
            pipe_thermal = create_pipeline(feature_extractor_thermal,model_thermal)
            # opencv_image = np.array(opencv_image)
            thermal_result=pipe_thermal(input_image)
        
        
            consolidated_result=get_output_label(thermal_result)
            if(consolidated_result!="RiceLeafs_Healthy"):
                f = open("remedy.json")
                data = json.load(f)
                for i in data[consolidated_result.split("_")[0]]:
                    if(i["disease_name"]==consolidated_result.split("_")[1]):
                        for key in i:
                            st.write(key," : ",i[key])
            else:
                st.write(consolidated_result.split("_")[0]," is ",consolidated_result.split("_")[1])
        
    if(rgb_check):
            
        st.write("RGB Image file uploader...")
        uploaded_file = st.file_uploader("Choose a rgb image file", type=["jpg","png"])

        if uploaded_file is not None:
            input_image = Image.open(uploaded_file)
            st.write("Input Image...")
            st.image(input_image)
            dataset_thermal,dataset_rgb = load_dataset_from_Hugging_Face()
            labels=dataset_thermal["train"].features["label"].names
            id2label={k:v for k,v in enumerate(labels)}
            label2id = {v:k for k,v in enumerate(labels)}
            feature_extractor_rgb,model_rgb = load_model_from_Hugging_Face("rgb")
            pipe_rgb = create_pipeline(feature_extractor_rgb,model_rgb)
            rgb_result=pipe_rgb(input_image)

            consolidated_result=get_output_label(rgb_result)
            if(consolidated_result!="RiceLeafs_Healthy"):
                f = open("remedy.json")
                data = json.load(f)
                for i in data[consolidated_result.split("_")[0]]:
                    if(i["disease_name"]==consolidated_result.split("_")[1]):
                        for key in i:
                            st.write(key," : ",i[key])
            else:
                st.write(consolidated_result.split("_")[0]," is ",consolidated_result.split("_")[1])

if(both_check):
    st.write("Thermal Image file uploader...")
    uploaded_file = st.file_uploader("Choose a thermal image file", type=["jpg","png"])

    if uploaded_file is not None:
        input_image = Image.open(uploaded_file)
        st.write("Input Image...")
        st.image(input_image)
        dataset_thermal,dataset_rgb = load_dataset_from_Hugging_Face()
        labels=dataset_thermal["train"].features["label"].names
        id2label={k:v for k,v in enumerate(labels)}
        label2id = {v:k for k,v in enumerate(labels)}

        feature_extractor_thermal,model_thermal = load_model_from_Hugging_Face("thermal")
        pipe_thermal = create_pipeline(feature_extractor_thermal,model_thermal)
        # opencv_image = np.array(opencv_image)
        thermal_result=pipe_thermal(input_image)
    st.write("RGB Image file uploader...")
    uploaded_file = st.file_uploader("Choose a rgb image file", type=["jpg","png"])

    if uploaded_file is not None:
        input_image = Image.open(uploaded_file)
        st.write("Input Image...")
        st.image(input_image)
        dataset_thermal,dataset_rgb = load_dataset_from_Hugging_Face()
        labels=dataset_thermal["train"].features["label"].names
        id2label={k:v for k,v in enumerate(labels)}
        label2id = {v:k for k,v in enumerate(labels)}
        feature_extractor_rgb,model_rgb = load_model_from_Hugging_Face("rgb")
        pipe_rgb = create_pipeline(feature_extractor_rgb,model_rgb)
        rgb_result=pipe_rgb(input_image)
        
        consolidated_result=get_output_label(consolidate_score(thermal_result,rgb_result))
        if(consolidated_result!="RiceLeafs_Healthy"):
                f = open("remedy.json")
                data = json.load(f)
                for i in data[consolidated_result.split("_")[0]]:
                    if(i["disease_name"]==consolidated_result.split("_")[1]):
                        for key in i:
                            st.write(key," : ",i[key])
        else:
            st.write(consolidated_result.split("_")[0]," is ",consolidated_result.split("_")[1])