import random
import statistics
import time

import matplotlib.pyplot as plt
import pandas as pd
import plotly.graph_objects as go
import streamlit as st

import algorithms as alg

st.title("Sorting Algorithm Visualizer")


def render_metrics(m):
    if not m:
        return
    ms = m.get("seconds", 0.0) * 1000.0
    comps = m.get("comparisons", 0)
    moves = m.get("moves", 0)
    c1, c2, c3 = st.columns(3)
    with c1:
        st.metric("Time (ms)", f"{ms:.2f}")
    with c2:
        st.metric("Comparisons", f"{comps:,}")
    with c3:
        st.metric("Moves", f"{moves:,}")


st.subheader("Input Configuration")
length = st.slider("List length", 5, 20, 8)
# Input type selector
input_type = st.selectbox(
    "Input type", ["Random", "Sorted", "Reversed", "Few unique"], index=0
)
# Generate data based on input type
if input_type == "Random":
    data = random.sample(range(1, 30), length)  # unique values
elif input_type == "Sorted":
    data = sorted(random.sample(range(1, 30), length))
elif input_type == "Reversed":
    data = sorted(random.sample(range(1, 30), length), reverse=True)
else:  # Few unique
    pool = random.sample(range(1, 30), k=min(3, max(1, length // 3)))
    data = [random.choice(pool) for _ in range(length)]
    random.shuffle(data)

st.write(f"Input array: {data}")

st.subheader("Scaling Benchmark (n -> time)")
sizes_str = st.text_input("Sizes", value="10, 20, 40, 80, 160")
try:
    sizes = [int(x.strip()) for x in sizes_str.split(",") if x.strip()]
    sizes = [s for s in sizes if s > 0]
except Exception:
    sizes = []

runs_scale = st.slider("Runs per size", 3, 50, 10, step=1)

algo_options = [
    "Insertion Sort",
    "Merge Sort",
    "Quick Sort",
    "Counting Sort",
    "Radix Sort (LSD)",
    "Heap Sort",
    "Shell Sort",
    "Bucket Sort",
]
selected_algos = st.multiselect(
    "Algorithms to include",
    options=algo_options,
    default=[
        "Insertion Sort",
        "Merge Sort",
        "Quick Sort",
        "Counting Sort",
        "Radix Sort (LSD)",
        "Heap Sort",
        "Shell Sort",
        "Bucket Sort",
    ],
)


def make_data(n: int, input_type: str):
    if input_type == "Random":
        return [random.randint(1, max(30, n * 3)) for _ in range(n)]
    elif input_type == "Sorted":
        arr = [random.randint(1, max(30, n * 3)) for _ in range(n)]
        return sorted(arr)
    elif input_type == "Reversed":
        arr = [random.randint(1, max(30, n * 3)) for _ in range(n)]
        return sorted(arr, reverse=True)
    else:
        pool_size = max(2, min(10, n // 5))
        pool = [random.randint(1, max(30, n * 3)) for _ in range(pool_size)]
        return [random.choice(pool) for _ in range(n)]


## cal algo
def get_algo_fn(name: str):
    if name == "Insertion Sort":
        return lambda arr: alg.insertion_sort(arr, record_steps=False)
    if name == "Merge Sort":
        return lambda arr: alg.merge_sort(arr, record_steps=False)
    if name == "Quick Sort":
        return lambda arr: alg.quick_sort(arr, record_steps=False)
    if name == "Counting Sort":
        return lambda arr: alg.counting_sort(arr, record_steps=False)
    if name == "Radix Sort (LSD)":
        return lambda arr: alg.radix_sort_lsd(arr, base=10, record_steps=False)
    if name == "Heap Sort":
        return lambda arr: alg.heap_sort(arr, record_steps=False)
    if name == "Shell Sort":
        return lambda arr: alg.shell_sort(arr, record_steps=False)
    if name == "Bucket Sort":
        return lambda arr: alg.bucket_sort(arr, record_steps=False)
    raise ValueError(name)


if st.button("Run Scaling Benchmark"):
    if not sizes:
        st.error("Please enter at least one valid size (e.g., 10, 20, 40).")
        st.stop()

    rows = []
    for n in sizes:
        for algo_name in selected_algos:
            fn = get_algo_fn(algo_name)
            times = []
            for _ in range(runs_scale):
                arr = make_data(n, input_type)
                _, m = fn(arr)
                times.append(m["seconds"] * 1000.0)  # ms

            avg_ms = statistics.mean(times)
            std_ms = statistics.pstdev(times) if len(times) > 1 else 0.0

            rows.append(
                {
                    "n": n,
                    "Algorithm": algo_name,
                    "Average Time (ms)": avg_ms,
                    "Std Dev (ms)": std_ms,
                }
            )

    df_scale = pd.DataFrame(rows)
    fig_scale = go.Figure()

    for algo_name in selected_algos:
        sub = df_scale[df_scale["Algorithm"] == algo_name].sort_values("n")
        fig_scale.add_trace(
            go.Scatter(
                x=sub["n"],
                y=sub["Average Time (ms)"],
                mode="lines+markers",
                name=algo_name,
            )
        )
    fig_scale.update_layout(
        title=f"Scaling Benchmark (input_type = {input_type}, runs = {runs_scale})",
        xaxis_title="n (input size)",
        yaxis_title="Average Time (ms)",
        height=480,
        width=1000,
    )
    st.plotly_chart(fig_scale, use_container_width=True)

    st.dataframe(
        df_scale.sort_values(["Algorithm", "n"]).style.format(
            {
                "Average Time (ms)": "{:.3f}",
                "Std Dev (ms)": "{:.3f}",
            }
        ),
        use_container_width=True,
    )

    st.download_button(
        "Download Scaling CSV",
        data=df_scale.to_csv(index=False).encode("utf-8"),
        file_name="scaling_benchmark.csv",
        mime="text/csv",
    )

if st.button("Run Comparison"):

    data_insertion = data.copy()
    data_merge = data.copy()
    data_quick = data.copy()
    data_counting = data.copy()
    data_radix = data.copy()
    data_heap = data.copy()
    data_shell = data.copy()
    data_bucket = data.copy()

    steps_insertion, metrics_insertion = alg.insertion_sort(data_insertion)
    steps_merge, metrics_merge = alg.merge_sort(data_merge)
    steps_quick, metrics_quick = alg.quick_sort(data_quick)
    steps_counting, metrics_counting = alg.counting_sort(data_counting)
    steps_radix, metrics_radix = alg.radix_sort_lsd(data_radix, base=10)
    steps_heap, metrics_heap = alg.heap_sort(data_heap)
    steps_shell, metrics_shell = alg.shell_sort(data_shell)
    steps_bucket, metrics_bucket = alg.bucket_sort(data_bucket)

    def create_animation(steps, title, color_fn):
        if not steps:
            return go.Figure(
                layout=go.Layout(
                    width=900,
                    height=420,
                    title=title,
                    xaxis=dict(visible=False),
                    yaxis=dict(visible=False),
                    annotations=[dict(text="No steps to display", showarrow=False)],
                )
            )
        frames = []
        for i, step in enumerate(steps):
            array = step["array"]
            active_index = step.get("active_index", -1)
            sorted_boundary = step.get("sorted_boundary", -1)

            colors = color_fn(len(array), active_index, sorted_boundary)

            frames.append(
                go.Frame(
                    data=[
                        go.Scatter(
                            x=list(range(len(array))),
                            y=array,
                            mode="markers+text",
                            marker=dict(size=28, color=colors),
                            text=array,
                            textposition="middle center",
                        )
                    ],
                    name=f"Step {i+1}",
                )
            )

        initial = steps[0]
        initial_colors = color_fn(
            len(initial["array"]),
            initial.get("active_index", -1),
            initial.get("sorted_boundary", -1),
        )

        fig = go.Figure(
            data=[
                go.Scatter(
                    x=list(range(len(initial["array"]))),
                    y=initial["array"],
                    mode="markers+text",
                    marker=dict(size=28, color=initial_colors),
                    text=initial["array"],
                    textposition="middle center",
                )
            ],
            layout=go.Layout(
                width=900,
                height=420,
                title=title,
                xaxis=dict(range=[-0.5, len(initial["array"]) - 0.5]),
                yaxis=dict(range=[0, max(max(s["array"]) for s in steps) + 5]),
                updatemenus=[
                    dict(
                        type="buttons",
                        buttons=[dict(label="Play", method="animate", args=[None])],
                        showactive=False,
                    )
                ],
                sliders=[
                    {
                        "steps": [
                            {
                                "args": [
                                    [f"Step {i+1}"],
                                    {"frame": {"duration": 500, "redraw": True}},
                                ],
                                "label": f"{i+1}",
                                "method": "animate",
                            }
                            for i in range(len(frames))
                        ],
                        "transition": {"duration": 0},
                        "x": 0,
                        "y": -0.1,
                        "currentvalue": {"prefix": "Step: "},
                    }
                ],
            ),
            frames=frames,
        )
        return fig

    def insertion_colors(length, active_index, sorted_boundary):
        return [
            (
                "red"
                if j == active_index
                else "green" if j <= sorted_boundary else "gray"
            )
            for j in range(length)
        ]

    def merge_colors(length, active_index, sorted_boundary):
        return [
            (
                "purple"
                if j == active_index
                else "blue" if j <= sorted_boundary else "gray"
            )
            for j in range(length)
        ]

    def quick_colors(length, active_index, sorted_boundary):
        return [
            (
                "orange"
                if j == active_index
                else "green" if j == sorted_boundary else "gray"
            )
            for j in range(length)
        ]

    def counting_colors(length, active_index, sorted_boundary):
        return [
            (
                "purple"
                if j == active_index
                else "green" if j == sorted_boundary else "gray"
            )
            for j in range(length)
        ]

    def radix_colors(length, active_index, sorted_boundary):
        return [
            (
                "purple"
                if j == active_index
                else "green" if j <= sorted_boundary else "gray"
            )
            for j in range(length)
        ]

    def heap_colors(length, active_index, sorted_boundary):
        return [
            (
                "orange"
                if j == active_index
                else (
                    "green"
                    if (sorted_boundary != -1 and j >= sorted_boundary)
                    else "gray"
                )
            )
            for j in range(length)
        ]

    def shell_colors(length, active_index, sorted_boundary):
        return [("orange" if j == active_index else "gray") for j in range(length)]

    def bucket_colors(length, active_index, sorted_boundary):
        return [
            (
                "purple"
                if j == active_index
                else "green" if j <= sorted_boundary else "gray"
            )
            for j in range(length)
        ]

    (
        tab_ins,
        tab_mer,
        tab_quick,
        tab_count,
        tab_radix,
        tab_heap,
        tab_shell,
        tab_bucket,
    ) = st.tabs(
        [
            "Insertion",
            "Merge",
            "Quick",
            "Counting",
            "Radix (LSD)",
            "Heap",
            "Shell",
            "Bucket",
        ]
    )

    with tab_ins:
        st.plotly_chart(
            create_animation(steps_insertion, "Insertion Sort", insertion_colors),
            use_container_width=True,
        )
    with tab_mer:
        st.plotly_chart(
            create_animation(steps_merge, "Merge Sort", merge_colors),
            use_container_width=True,
        )
    with tab_quick:
        st.plotly_chart(
            create_animation(steps_quick, "Quick Sort", quick_colors),
            use_container_width=True,
        )
    with tab_count:
        st.plotly_chart(
            create_animation(steps_counting, "Counting Sort", counting_colors),
            use_container_width=True,
        )
    with tab_radix:
        st.plotly_chart(
            create_animation(steps_radix, "Radix Sort (LSD)", radix_colors),
            use_container_width=True,
        )
    with tab_heap:
        st.plotly_chart(
            create_animation(steps_heap, "Heap Sort", heap_colors),
            use_container_width=True,
        )
    with tab_shell:
        st.plotly_chart(
            create_animation(steps_shell, "Shell Sort", shell_colors),
            use_container_width=True,
        )
    with tab_bucket:
        st.plotly_chart(
            create_animation(steps_bucket, "Bucket Sort", bucket_colors),
            use_container_width=True,
        )

    df = pd.DataFrame(
        [
            {
                "Algorithm": "Insertion Sort",
                "Time (ms)": metrics_insertion["seconds"] * 1000,
                "Comparisons": metrics_insertion["comparisons"],
                "Moves": metrics_insertion["moves"],
                "Frames": len(steps_insertion),
                "Sorted OK": steps_insertion[-1]["array"] == sorted(data),
            },
            {
                "Algorithm": "Merge Sort",
                "Time (ms)": metrics_merge["seconds"] * 1000,
                "Comparisons": metrics_merge["comparisons"],
                "Moves": metrics_merge["moves"],
                "Frames": len(steps_merge),
                "Sorted OK": steps_merge[-1]["array"] == sorted(data),
            },
            {
                "Algorithm": "Quick Sort",
                "Time (ms)": metrics_quick["seconds"] * 1000,
                "Comparisons": metrics_quick["comparisons"],
                "Moves": metrics_quick["moves"],
                "Frames": len(steps_quick),
                "Sorted OK": steps_quick[-1]["array"] == sorted(data),
            },
            {
                "Algorithm": "Counting Sort",
                "Time (ms)": metrics_counting["seconds"] * 1000,
                "Comparisons": metrics_counting["comparisons"],
                "Moves": metrics_counting["moves"],
                "Frames": len(steps_counting),
                "Sorted OK": steps_counting[-1]["array"] == sorted(data),
            },
            {
                "Algorithm": "Radix Sort(LSD)",
                "Time (ms)": metrics_radix["seconds"] * 1000,
                "Comparisons": metrics_radix["comparisons"],
                "Moves": metrics_radix["moves"],
                "Frames": len(steps_radix),
                "Sorted OK": steps_radix[-1]["array"] == sorted(data),
            },
            {
                "Algorithm": "Heap Sort",
                "Time (ms)": metrics_heap["seconds"] * 1000,
                "Comparisons": metrics_heap["comparisons"],
                "Moves": metrics_heap["moves"],
                "Frames": len(steps_heap),
                "Sorted OK": steps_heap[-1]["array"] == sorted(data),
            },
            {
                "Algorithm": "Shell Sort",
                "Time (ms)": metrics_shell["seconds"] * 1000,
                "Comparisons": metrics_shell["comparisons"],
                "Moves": metrics_shell["moves"],
                "Frames": len(steps_shell),
                "Sorted OK": steps_shell[-1]["array"] == sorted(data),
            },
            {
                "Algorithm": "Bucket Sort",
                "Time (ms)": metrics_bucket["seconds"] * 1000,
                "Comparisons": metrics_bucket["comparisons"],
                "Moves": metrics_bucket["moves"],
                "Frames": len(steps_bucket),
                "Sorted OK": steps_bucket[-1]["array"] == sorted(data),
            },
        ]
    )
    st.subheader("Summary Table")
    st.dataframe(df.style.format({"Time (ms)": "{:.2f}"}), use_container_width=True)

    csv = df.to_csv(index=False).encode("utf-8")
    st.download_button(
        "Download CSV", data=csv, file_name="sorting_summary.csv", mime="text/csv"
    )