import ast
import itertools
from collections import Counter
import pandas as pd
import plotly.express as px
from plotly.graph_objs import Figure
import streamlit as st
from typing import Any


def get_elements(x: str, key: str = "elements") -> Any:
    try:
        result = ast.literal_eval(x)[key]
        print(type(result), result)
        if isinstance(result, str):
            return [result]
        return result
    except:
        return []


def get_most_common_elements(df: pd.DataFrame, column: str,  key: str = "elements") -> list:
    built_elements = df[column].apply(get_elements, args=(key,))
    most_common_elements = Counter(itertools.chain.from_iterable(built_elements)).most_common()
    return most_common_elements


def get_plot_from_most_common_elements(df: pd.DataFrame, column: str, key: str = "elements") -> Figure:
    most_common_elements = get_most_common_elements(df, column, key)
    most_common_elements = pd.DataFrame(most_common_elements, columns=[column, "count"])
    return px.bar(
        most_common_elements,
        x=column,
        y="count",
        labels={"count": "# Objects", column: "Built Elements"},
    )


def aggregate_fauna_elements(df: pd.DataFrame) -> dict:
    fauna_elements = {}
    fauna = df["fauna_identification"].apply(lambda x: ast.literal_eval(x)["fauna"])

    for i in range(len(fauna)):
        for element in fauna[i]:
            if element["type"] in fauna_elements:
                fauna_elements[element["type"]] += element["count"]
            else:
                fauna_elements[element["type"]] = element["count"]
    return fauna_elements