"""
init-mater-project

Copyright (C) 2025 [Lauranne Sarribouette] <lauranne.sarribouette@univ-grenoble-alpes.fr>

SPDX-License-Identifier: LGPL-3.0-or-later
"""

import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import streamlit as st
from typing import Dict, Optional, List

from src.app import data, components


def display_time_series_tab(
    df_full: pd.DataFrame, dimension_filters: Dict[str, List[str]]
):
    """Display time series with multi-line charts"""
    st.markdown("## 📊 Time Series Evolution")

    time_range = components.render_time_filter_for_tab(
        df_full, mode="range", key_suffix="timeseries"
    )

    df_filtered = data.filter_data(
        df_full, time_range, dimension_filters, aggregate_all=False
    )

    if df_filtered.empty:
        st.warning("⚠️ No data matches current filters.")
        return

    num_series = data.count_series(df_filtered)

    if num_series > 50:
        st.warning(
            f"⚠️ Too many series ({num_series}). "
            "Chart may be hard to read. Consider filtering dimensions."
        )

    plot_df = df_filtered.reset_index()

    if "time" not in plot_df.columns:
        st.info("No time dimension found in data")
        return

    dimension_cols = [col for col in plot_df.columns if col not in ["time", "value"]]

    if dimension_cols:
        plot_df["series"] = plot_df[dimension_cols].apply(
            lambda row: " ".join(str(v) for v in row), axis=1
        )
    else:
        plot_df["series"] = "Total"

    fig = px.line(
        plot_df,
        x="time",
        y="value",
        color="series",
        labels={"value": "Value", "time": "Time"},
    )

    fig.update_layout(
        height=500,
        showlegend=True,
        hovermode="x unified",
        legend=dict(orientation="h", yanchor="bottom", y=-0.5, xanchor="center", x=0.5),
    )
    st.plotly_chart(
        fig, use_container_width=True
    )  # TODO: soon to be deprecated, use `width="stretch" instead`


def display_area_tab(df_full: pd.DataFrame, dimension_filters: Dict[str, List[str]]):
    """Display stacked area charts"""
    st.markdown("## 📈 Area (Stacked)")

    dimensions = [name for name in df_full.index.names if name not in ["time"]]

    if not dimensions:
        st.info("No dimensions available for stacking")
        return

    col1, col2 = st.columns([2, 1])

    with col1:
        stack_by = st.selectbox("Stack by dimension", dimensions, key="area_stack_by")

    with col2:
        mode = st.radio(
            "Mode", ["Absolute", "Percentage"], horizontal=True, key="area_mode"
        )

    time_range = components.render_time_filter_for_tab(
        df_full, mode="range", key_suffix="area"
    )

    df_filtered = data.filter_data(
        df_full, time_range, dimension_filters, aggregate_all=False
    )

    if df_filtered.empty:
        st.warning("⚠️ No data available")
        return

    plot_df = df_filtered.reset_index()

    if "time" not in plot_df.columns:
        st.info("No time dimension found")
        return

    if stack_by in plot_df.columns:
        grouped = plot_df.groupby(["time", stack_by])["value"].sum().reset_index()

        if mode == "Percentage":
            totals = grouped.groupby("time")["value"].sum().reset_index()
            totals.columns = ["time", "total"]
            grouped = grouped.merge(totals, on="time")
            grouped["value"] = (grouped["value"] / grouped["total"]) * 100

        fig = px.area(
            grouped,
            x="time",
            y="value",
            color=stack_by,
            title=f"Stacked Area ({'%' if mode == 'Percentage' else 'Absolute'})",
            labels={"value": "Percentage" if mode == "Percentage" else "Value"},
        )

        fig.update_layout(
            height=500,
            showlegend=True,
            hovermode="x unified",
            legend=dict(
                orientation="h", yanchor="bottom", y=-0.5, xanchor="center", x=0.5
            ),
        )
        st.plotly_chart(
            fig, use_container_width=True
        )  # TODO: soon to be deprecated, use `width="stretch" instead`


def display_snapshot_tab(
    df_full: pd.DataFrame, dimension_filters: Dict[str, List[str]]
):
    """Display snapshot views: Pie, Treemap, Bar"""
    st.markdown("## 🎯 Snapshot")

    dimensions = [name for name in df_full.index.names if name not in ["time"]]

    if not dimensions:
        st.info("No dimensions available")
        return

    time_points = []
    if "time" in df_full.index.names:
        time_points = sorted(df_full.index.get_level_values("time").unique())

    col1, col2 = st.columns(2)

    with col1:
        viz_type = st.radio(
            "Visualization type",
            ["Pie", "Treemap", "Bar"],
            horizontal=True,
            key="snapshot_viz_type",
            label_visibility="collapsed",
        )

    selected_time = None
    if time_points:
        with col2:
            time_range = components.render_time_filter_for_tab(
                df_full, mode="single", key_suffix="snapshot"
            )
            selected_time = time_range[0] if time_range else None

    primary_dim = st.selectbox("Dimension", dimensions, key="snapshot_primary_dim")

    df_filtered = data.filter_data(
        df_full,
        (selected_time, selected_time) if selected_time else None,
        dimension_filters,
        aggregate_all=False,
    )

    if df_filtered.empty:
        st.warning("⚠️ No data available")
        return

    if "time" in df_filtered.index.names:
        df_time = df_filtered.xs(selected_time, level="time", drop_level=True)
    else:
        df_time = df_filtered

    if viz_type == "Pie":
        if primary_dim in df_time.index.names:
            grouped = df_time.groupby(level=primary_dim).sum().reset_index()

            fig = px.pie(
                grouped,
                values="value",
                names=primary_dim,
                title=f"Distribution by {primary_dim}",
            )
            fig.update_layout(
                height=500,
                showlegend=True,
                hovermode="x unified",
                legend=dict(
                    orientation="h", yanchor="bottom", y=-0.5, xanchor="center", x=0.5
                ),
            )
            st.plotly_chart(
                fig, use_container_width=True
            )  # TODO: soon to be deprecated, use `width="stretch" instead`

    elif viz_type == "Treemap":
        treemap_data = data.prepare_treemap_data(df_time, primary_dim)

        if not treemap_data.empty:
            fig = go.Figure(
                go.Treemap(
                    labels=treemap_data["labels"],
                    parents=treemap_data["parents"],
                    values=treemap_data["values"],
                    ids=treemap_data["ids"],
                    texttemplate="<b>%{label}</b><br>%{percentParent}",
                    hovertemplate="<b>%{label}</b><br>Value: %{value}<br>Percent: %{percentParent}<extra></extra>",
                )
            )

            fig.update_layout(
                height=500,
                title=f"Treemap: {primary_dim}",
            )
            st.plotly_chart(
                fig, use_container_width=True
            )  # TODO: soon to be deprecated, use `width="stretch" instead`

    elif viz_type == "Bar":
        if primary_dim in df_time.index.names:
            grouped = df_time.groupby(level=primary_dim).sum().reset_index()

            fig = px.bar(
                grouped,
                x=primary_dim,
                y="value",
                title=f"Distribution by {primary_dim}",
                labels={"value": "Value"},
            )
            fig.update_layout(height=500)
            st.plotly_chart(
                fig, use_container_width=True
            )  # TODO: soon to be deprecated, use `width="stretch" instead`


def display_heatmap_tab(df_full: pd.DataFrame, dimension_filters: Dict[str, List[str]]):
    """Display heatmap"""
    st.markdown("## 🔥 Heatmap")

    dimensions = [name for name in df_full.index.names if name not in ["time"]]

    if not dimensions:
        st.info("No dimensions available")
        return

    row_dim = st.selectbox("Dimension (rows)", dimensions, key="heatmap_row_dim")

    df_filtered = data.filter_data(
        df_full,
        None,  # No time filter for heatmap
        dimension_filters,
        aggregate_all=False,
    )

    if df_filtered.empty:
        st.warning("⚠️ No data available")
        return

    if "time" not in df_filtered.index.names or row_dim not in df_filtered.index.names:
        st.error(f"Required dimensions not found: time or {row_dim}")
        return

    heatmap_data = df_filtered.groupby(level=[row_dim, "time"]).sum()
    pivot = heatmap_data.unstack(level="time", fill_value=0)

    if pivot.empty:
        st.warning("No data to display")
        return

    fig = px.imshow(
        pivot.values,
        labels=dict(x="Time", y=row_dim, color="Value"),
        x=[
            col.strftime("%Y") if hasattr(col, "strftime") else str(col)
            for col in pivot.columns.get_level_values(0)
        ],
        y=pivot.index.tolist(),
        aspect="auto",
        color_continuous_scale="YlOrRd",
    )

    fig.update_layout(height=max(500, len(pivot) * 30), title=f"{row_dim} × time")

    st.plotly_chart(
        fig, use_container_width=True
    )  # TODO: soon to be deprecated, use `width="stretch" instead`


def display_raw_data_tab(
    df_full: pd.DataFrame,
    dimension_filters: Dict[str, List[str]],
    simulation_name: Optional[str],
):
    """Display raw data with export"""
    st.markdown("## 📋 Raw Data")

    time_range = components.render_time_filter_for_tab(
        df_full, mode="range", key_suffix="rawdata"
    )

    df_filtered = data.filter_data(
        df_full, time_range, dimension_filters, aggregate_all=False
    )

    if df_filtered.empty:
        st.warning("⚠️ No data matches current filters")
        return

    st.data_editor(
        df_filtered.reset_index(),
        use_container_width=True,
        num_rows="dynamic",  # TODO: soon to be deprecated, use `width="stretch" instead`
    )

    csv = df_filtered.to_csv().encode("utf-8")
    filename = f"{simulation_name}_data.csv" if simulation_name else "data.csv"

    st.download_button(
        label="📥 Download CSV",
        data=csv,
        file_name=filename,
        mime="text/csv",
    )
