from typing import List

import numpy as np
import plotly.express as px
import plotly.graph_objects as go
import polars as pl
import streamlit as st

from pixel_patrol.default_tabs import DefaultTabs
from pixel_patrol.utils.utils import create_sprite_image
from pixel_patrol.widgets.widget_interface import ImagePrevalidationWidget


class ImageMosaikWidget(ImagePrevalidationWidget):

    @property
    def tab(self) -> str:
        return DefaultTabs.VISUALIZATION.value

    @property
    def name(self) -> str:
        return "Image Mosaic"

    def required_columns(self) -> List[str]:
        return ["*"]

    def options(self, data_frame: pl.DataFrame) -> dict:
        sort_column = st.selectbox("Sort mosaik by..", data_frame.columns)
        return {
            "sort_column": sort_column,
        }

    def render(self, files_dataframe: pl.DataFrame, sort_column: str = None):
        """
        Display a mosaic of random images for each selected folder.
        """
        if sort_column:
            files_dataframe = files_dataframe.sort(sort_column)

        mosaic_image = create_sprite_image(files_dataframe, border=True, border_size=5)
        unique_folders = files_dataframe.unique(subset=["selected_folder_short"])
        color_map = {row["selected_folder_short"]: row["color"] for row in unique_folders.to_dicts()}

        st.markdown("""### Thumbnails

The little pictures printed here are 2D representations of each file in the selected folders. They are generated by picking the center slice of all dimensions other than X and Y
 and scaling them into a square, so they do not represent the images accurately. It allows gaining quick access over the dataset though, and sorting it by any of the existing 
 measurement points (like min or mean) empowers us to visualize relations between images easier.  
""")

        fig = px.imshow(np.array(mosaic_image), binary_string=True)

        # Create dummy traces (no actual data) for the legend
        dummy_traces = [
            go.Scatter(
                x=[None],
                y=[None],
                mode='markers',
                marker=dict(size=10, color=color_map[label]),
                name=label,
                showlegend=True
            )
            for label in color_map
        ]

        # Add the dummy traces to your figure
        for trace in dummy_traces:
            fig.add_trace(trace)

        fig.update_layout(
            autosize=True,
            height=min(1000, mosaic_image.height),
            xaxis=dict(visible=False),
            yaxis=dict(visible=False),
            margin=dict(l=0, r=0, t=0, b=0)
        )

        # Step 6: Display in Streamlit
        st.plotly_chart(fig, use_container_width=True)
