import json
import math
import sqlite3

import streamlit as st
import torch
import torchvision
from PIL import Image
from huggingface_hub import hf_hub_download
from torchvision import transforms
from transformers import AutoModelForImageClassification, AutoConfig

# Set the page title
st.title("Global Bird Classification App")

# Input latitude and longitude (optional)
latitude = st.number_input("Enter latitude (optional)", value=None, format="%f")
longitude = st.number_input("Enter longitude (optional)", value=None, format="%f")

st.text('Please fill the coordinates before upload image.')

# Upload an image
uploaded_file = st.file_uploader("Please select an image", type=["jpg", "jpeg", "png"])

lang = st.selectbox(
    "Result Language",
    options=[2, 1, 0],
    format_func=lambda x: {
        2: "Latina (Nomen Scientificum)",
        1: "English (IOC 10.1)",
        0: "中文 （中国大陆）",
    }[x]
)


classify_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


# crop and classification
def classify_objects(classification_model, image, species_list):
    input_tensor = classify_transforms(image).unsqueeze(0).to(device)
    with torch.no_grad():
        logits = classification_model(input_tensor)[0]
        filtered = get_filtered_predictions(logits, species_list)
        return softmax(filtered)


def softmax(tuples):
    # `torch.nn.functional.softmax` requires the input to be `Tensor`, so I implemented it myself
    values = [t[1] for t in tuples]
    exp_values = [math.exp(v) for v in values]
    sum_exp_values = sum(exp_values)
    softmax_values = [ev / sum_exp_values for ev in exp_values]
    updated_tuples = [(t[0], softmax_values[i]) for i, t in enumerate(tuples)]
    updated_tuples.sort(key=lambda t: t[1], reverse=True)

    return updated_tuples


def get_filtered_predictions(predictions: list[float], species_list: list[int]) -> list[tuple[int, float]]:
    original = {index: value for index, value in enumerate(predictions)}

    if species_list:
        filtered_predictions = [(key, value) for key, value in original.items() if key in species_list]
    else:
        filtered_predictions = [(key, value) for key, value in original.items()]

    return filtered_predictions


class DistributionDB:
    def __init__(self, db_path):
        self.con = sqlite3.connect(db_path)
        self.cur = self.con.cursor()


    def get_list(self, lat, lng) -> list:
        self.cur.execute(f'''
SELECT m.cls
FROM distributions AS d
LEFT OUTER JOIN places AS p
  ON p.worldid = d.worldid
LEFT OUTER JOIN sp_cls_map AS m
  ON d.species = m.species
WHERE p.south <= {lat}
  AND p.north >= {lat}
  AND p.east >= {lng}
  AND p.west <= {lng}
GROUP BY d.species, m.cls;
''')

        return [row[0] for row in self.cur]

    def close(self):
        self.cur.close()
        self.con.close()


# If the user uploads an image
if uploaded_file is not None:
    try:
        label_map_path = hf_hub_download(repo_id='sunjiao/osea', filename='bird_info.json')
        st.success(f"Successfully downloaded labels from Hugging Face Hub!")
    except Exception as e:
        st.error(f"Failed to download the file: {e}")
        st.stop()

    with open(label_map_path, 'r') as f:
        data = f.read()

    bird_info = json.loads(data)

    species_list = None
    if latitude and longitude:
        try:
            sqlite_path = hf_hub_download(repo_id='sunjiao/osea', filename='avonet.db')
            st.success(f"Successfully downloaded distribution database from Hugging Face Hub!")
        except Exception as e:
            st.error(f"Failed to download the file: {e}")
            st.stop()

        db = DistributionDB(sqlite_path)
        species_list = db.get_list(latitude, longitude)
        db.close()

    # Open the image
    image = Image.open(uploaded_file)

    # Display the uploaded image
    st.image(image, caption="Uploaded Image", use_container_width=True)

    try:
        weight_dict = hf_hub_download(repo_id='sunjiao/osea', filename='pytorch_model.bin')
        st.success(f"Successfully downloaded weight dict from Hugging Face Hub!")
    except Exception as e:
        st.error(f"Failed to download the file: {e}")
        st.stop()

    model = torchvision.models.resnet34(num_classes=11000)
    model.load_state_dict(torch.load(weight_dict, map_location=device))
    model.eval()

    results = classify_objects(model, image, species_list)

    top3_results = results[:3]

    # Display the top 3 results and their probabilities
    st.subheader("Classification Results (Top 3):")
    for result in top3_results:
        st.write(f"{bird_info[result[0]][lang]}: {result[1]:.4f}")

# Display latitude and longitude if provided
if latitude is not None and longitude is not None:
    st.write(f"Entered Latitude: {latitude}, Longitude: {longitude}")