import io
import math
import base64
import datetime as dt
from dataclasses import dataclass
from typing import Optional, Tuple, List

import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import gradio as gr

# -----------------------------
# Helper: generate a fun demo dataset if user doesn't upload one
# -----------------------------

def make_demo_data(n_days: int = 180, seed: int = 7) -> pd.DataFrame:
    rng = np.random.default_rng(seed)
    start = dt.date.today() - dt.timedelta(days=n_days)
    dates = pd.date_range(start, periods=n_days, freq="D")

    categories = np.array(["Gadgets", "Gear", "Gifts"])  # product category
    regions = np.array(["NA", "EMEA", "APAC"])          # sales region

    cat = rng.choice(categories, size=n_days)
    region = rng.choice(regions, size=n_days)

    base = (
        200
        + 20 * np.sin(np.linspace(0, 6 * math.pi, n_days))  # seasonality
        + rng.normal(0, 15, n_days)
    )

    # Category and region effects
    cat_bump = np.where(cat == "Gadgets", 25, np.where(cat == "Gear", 10, -5))
    reg_bump = np.where(region == "NA", 15, np.where(region == "EMEA", 5, 0))

    sales = np.clip(base + cat_bump + reg_bump, 10, None)
    profit = np.clip(sales * (0.15 + rng.normal(0.0, 0.03, n_days)), 0, None)
    customers = np.clip((sales / 10) + rng.normal(0, 2, n_days), 1, None)

    df = pd.DataFrame(
        {
            "date": dates,
            "category": cat,
            "region": region,
            "sales": sales.round(2),
            "profit": profit.round(2),
            "customers": customers.round(0).astype(int),
        }
    )
    return df

# -----------------------------
# Core plotting logic
# -----------------------------

PLOT_TYPES = [
    "line",
    "scatter",
    "bar",
    "box",
    "violin",
    "heatmap (corr)",
]

COLOR_SCALES = [
    "Turbo",
    "Viridis",
    "Cividis",
    "Plasma",
    "Inferno",
    "Magma",
    "IceFire",
    "Tealrose",
    "Bluered",
]

@dataclass
class PlotConfig:
    plot_type: str
    x: Optional[str]
    y: Optional[str]
    color: Optional[str]
    facet_col: Optional[str]
    agg: str
    smooth: bool
    smooth_window: int
    add_trend: bool
    trend_degree: int
    points: bool
    theme: str
    color_scale: str


def _pick_numeric_cols(df: pd.DataFrame) -> List[str]:
    return [c for c in df.columns if pd.api.types.is_numeric_dtype(df[c])]


def _pick_categorical_cols(df: pd.DataFrame) -> List[str]:
    return [
        c
        for c in df.columns
        if not pd.api.types.is_numeric_dtype(df[c]) and df[c].nunique() < len(df) / 2
    ]


def _rolling_series(y: pd.Series, window: int) -> pd.Series:
    window = max(2, int(window))
    return y.rolling(window=window, min_periods=1, center=False).mean()


def _add_poly_trend(fig: go.Figure, xvals, yvals, degree: int, name: str = "Trend") -> None:
    # Safeguard: only add if enough points and numeric
    if len(xvals) < 3:
        return
    # Convert datetimes to ordinal for regression if needed
    if np.issubdtype(np.array(xvals).dtype, np.datetime64):
        x_numeric = pd.to_datetime(xvals).map(pd.Timestamp.toordinal).to_numpy()
    else:
        x_numeric = pd.Series(xvals).astype(float).to_numpy()

    y_numeric = pd.Series(yvals).astype(float).to_numpy()
    mask = np.isfinite(x_numeric) & np.isfinite(y_numeric)
    x_numeric = x_numeric[mask]
    y_numeric = y_numeric[mask]
    if len(x_numeric) < max(5, degree + 2):
        return

    coeffs = np.polyfit(x_numeric, y_numeric, deg=degree)
    xp = np.linspace(x_numeric.min(), x_numeric.max(), 200)
    yp = np.polyval(coeffs, xp)

    # Convert back datetime if original was datetime
    if np.issubdtype(np.array(xvals).dtype, np.datetime64):
        xp_plot = pd.to_datetime([dt.date.fromordinal(int(v)) for v in xp])
    else:
        xp_plot = xp

    fig.add_trace(
        go.Scatter(
            x=xp_plot,
            y=yp,
            mode="lines",
            name=name,
            line=dict(color="#222", width=3, dash="dash"),
            hovertemplate="%{x}<br>trend=%{y:.2f}<extra></extra>",
        )
    )


def build_plot(df: pd.DataFrame, cfg: PlotConfig) -> go.Figure:
    # Basic validation
    if df is None or df.empty:
        return go.Figure()

    # Auto-fill axes if missing
    numeric_cols = _pick_numeric_cols(df)
    cat_cols = _pick_categorical_cols(df)

    x = cfg.x or ("date" if "date" in df.columns else (cat_cols[0] if cat_cols else (numeric_cols[0] if numeric_cols else None)))
    y = cfg.y or (numeric_cols[0] if numeric_cols else None)

    if cfg.plot_type == "heatmap (corr)":
        num = df[numeric_cols].copy() if numeric_cols else pd.DataFrame()
        if num.empty:
            fig = go.Figure()
            fig.add_annotation(text="No numeric columns for correlation heatmap.", showarrow=False)
            return fig
        corr = num.corr(numeric_only=True)
        fig = px.imshow(
            corr,
            text_auto=True,
            color_continuous_scale=cfg.color_scale,
            aspect="auto",
            title="Correlation Heatmap",
        )
        fig.update_layout(template=cfg.theme)
        return fig

    if x is None or y is None:
        fig = go.Figure()
        fig.add_annotation(text="Select X and Y columns to plot.", showarrow=False)
        return fig

    df_plot = df.copy()

    # Optional aggregation for line/bar: group by X (and color if provided)
    if cfg.agg != "none" and cfg.plot_type in {"line", "bar"}:
        group_keys = [x] + ([cfg.color] if cfg.color else [])
        if cfg.agg == "sum":
            df_plot = df_plot.groupby(group_keys, dropna=False)[y].sum().reset_index()
        elif cfg.agg == "mean":
            df_plot = df_plot.groupby(group_keys, dropna=False)[y].mean().reset_index()
        elif cfg.agg == "median":
            df_plot = df_plot.groupby(group_keys, dropna=False)[y].median().reset_index()

    # Build the base figure with Plotly Express
    if cfg.plot_type == "line":
        fig = px.line(
            df_plot,
            x=x,
            y=y,
            color=cfg.color,
            facet_col=cfg.facet_col,
            template=cfg.theme,
            color_continuous_scale=cfg.color_scale,
        )
    elif cfg.plot_type == "scatter":
        fig = px.scatter(
            df_plot,
            x=x,
            y=y,
            color=cfg.color,
            facet_col=cfg.facet_col,
            template=cfg.theme,
            color_continuous_scale=cfg.color_scale,
            render_mode="auto",
        )
        if cfg.points:
            fig.update_traces(marker=dict(size=9, opacity=0.7))
    elif cfg.plot_type == "bar":
        fig = px.bar(
            df_plot,
            x=x,
            y=y,
            color=cfg.color,
            facet_col=cfg.facet_col,
            template=cfg.theme,
            color_continuous_scale=cfg.color_scale,
        )
    elif cfg.plot_type == "box":
        fig = px.box(
            df_plot,
            x=cfg.color if cfg.color else None,
            y=y,
            points="all" if cfg.points else False,
            color=cfg.color,
            template=cfg.theme,
            color_discrete_sequence=px.colors.qualitative.Set2,
        )
    elif cfg.plot_type == "violin":
        fig = px.violin(
            df_plot,
            x=cfg.color if cfg.color else None,
            y=y,
            box=True,
            points="all" if cfg.points else False,
            color=cfg.color,
            template=cfg.theme,
            color_discrete_sequence=px.colors.qualitative.Pastel,
        )
    else:
        fig = go.Figure()

    # Optional smoothing (rolling mean) for line/scatter
    if cfg.smooth and cfg.plot_type in {"line", "scatter"}:
        try:
            if cfg.color and cfg.color in df_plot.columns:
                for key, grp in df_plot.groupby(cfg.color):
                    sm = _rolling_series(grp[y].reset_index(drop=True), cfg.smooth_window)
                    fig.add_trace(
                        go.Scatter(
                            x=grp[x],
                            y=sm,
                            mode="lines",
                            name=f"{key} (rolling {cfg.smooth_window})",
                            line=dict(width=3),
                            opacity=0.85,
                        )
                    )
            else:
                sm = _rolling_series(df_plot[y], cfg.smooth_window)
                fig.add_trace(
                    go.Scatter(
                        x=df_plot[x],
                        y=sm,
                        mode="lines",
                        name=f"rolling {cfg.smooth_window}",
                        line=dict(width=3),
                        opacity=0.85,
                    )
                )
        except Exception:
            pass

    # Optional polynomial trend for scatter/line (degree 1..3)
    if cfg.add_trend and cfg.plot_type in {"line", "scatter"}:
        if cfg.color and cfg.color in df_plot.columns:
            for key, grp in df_plot.groupby(cfg.color):
                _add_poly_trend(fig, grp[x], grp[y], cfg.trend_degree, name=f"{key} trend")
        else:
            _add_poly_trend(fig, df_plot[x], df_plot[y], cfg.trend_degree)

    # Decorate
    fig.update_layout(
        title=f"{cfg.plot_type.title()} — {y} vs {x}",
        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
        hovermode="x unified" if cfg.plot_type == "line" else "closest",
        margin=dict(l=40, r=20, t=60, b=40),
    )

    # Fun: subtle background stripes for readability
    fig.update_xaxes(showgrid=True, gridcolor="rgba(0,0,0,0.06)")
    fig.update_yaxes(showgrid=True, gridcolor="rgba(0,0,0,0.06)")

    return fig

# -----------------------------
# Gradio UI
# -----------------------------

def infer_options(df: pd.DataFrame) -> Tuple[List[str], List[str]]:
    num_cols = _pick_numeric_cols(df)
    cat_cols = _pick_categorical_cols(df)
    return num_cols, cat_cols


def read_csv(file: Optional[gr.File]) -> pd.DataFrame:
    if file is None:
        return make_demo_data()
    try:
        df = pd.read_csv(file.name)
    except Exception:
        # try excel
        try:
            df = pd.read_excel(file.name)
        except Exception as e:
            raise ValueError(f"Failed to read file: {e}")
    return df


def ui_refresh_columns(file: Optional[gr.File]):
    df = read_csv(file)
    num_cols, cat_cols = infer_options(df)
    all_cols = list(df.columns)
    # reasonable defaults
    x_default = "date" if "date" in all_cols else (cat_cols[0] if cat_cols else (num_cols[0] if num_cols else None))
    y_default = num_cols[0] if num_cols else None

    return (
        gr.Dropdown.update(choices=all_cols, value=x_default),
        gr.Dropdown.update(choices=num_cols, value=y_default),
        gr.Dropdown.update(choices=all_cols + [None], value=None),
        gr.Dropdown.update(choices=cat_cols + [None], value=None),
        df.head(10),
    )


def ui_plot(file, plot_type, x, y, color, facet_col, agg, smooth, smooth_window, add_trend, trend_degree, points, theme, color_scale):
    df = read_csv(file)
    cfg = PlotConfig(
        plot_type=plot_type,
        x=x,
        y=y,
        color=color if color != "None" else None,
        facet_col=facet_col if facet_col != "None" else None,
        agg=agg,
        smooth=smooth,
        smooth_window=int(smooth_window or 5),
        add_trend=add_trend,
        trend_degree=int(trend_degree or 1),
        points=bool(points),
        theme=theme,
        color_scale=color_scale,
    )
    fig = build_plot(df, cfg)
    return fig

with gr.Blocks(theme=gr.themes.Soft(primary="#2563eb", neutral="#111827")) as demo:
    gr.Markdown(
        """
        # 📈 Data Viz Studio — with Gradio + Plotly
        Upload a CSV (or use the built‑in demo), then craft an **interesting** plot.
        
        **Tips**
        - If your data has a `date` column, try a **Line** plot with a **rolling average** and a **trend**.
        - Use **Color by** + **Facet** to spot differences across categories.
        - Use **Heatmap (corr)** to quickly find strong relationships between numeric columns.
        """
    )

    with gr.Row():
        file = gr.File(label="Upload CSV or Excel (optional)")
        refresh = gr.Button("🔄 Load / Refresh Columns", variant="secondary")
        use_demo = gr.Button("🎲 Use Demo Data", variant="secondary")

    with gr.Accordion("Data Preview", open=False):
        preview = gr.Dataframe(row_count=(5, "dynamic"), wrap=True)

    with gr.Row():
        plot_type = gr.Dropdown(PLOT_TYPES, value="line", label="Plot Type")
        theme = gr.Dropdown(px.colors.named_themes(), value="plotly_white", label="Theme")
        color_scale = gr.Dropdown(COLOR_SCALES, value="Turbo", label="Color Scale (for continuous)")

    with gr.Row():
        x = gr.Dropdown(choices=[], label="X Axis")
        y = gr.Dropdown(choices=[], label="Y Axis (numeric)")

    with gr.Row():
        color = gr.Dropdown(choices=[], value=None, label="Color by (categorical or numeric)")
        facet_col = gr.Dropdown(choices=[], value=None, label="Facet column (categorical)")
        agg = gr.Dropdown(["none", "sum", "mean", "median"], value="none", label="Aggregate (line/bar)")

    with gr.Row():
        smooth = gr.Checkbox(value=True, label="Add Rolling Average")
        smooth_window = gr.Slider(2, 60, value=7, step=1, label="Rolling Window (periods)")
        add_trend = gr.Checkbox(value=False, label="Add Polynomial Trendline")
        trend_degree = gr.Slider(1, 3, value=1, step=1, label="Trend Degree (1-3)")
        points = gr.Checkbox(value=True, label="Show Points (scatter/box/violin)")

    out = gr.Plot(label="Your Plot")

    # Wire events
    def set_demo(_):
        df = make_demo_data()
        # Save to an in-memory CSV for preview (not strictly necessary for plotting)
        num_cols, cat_cols = infer_options(df)
        all_cols = list(df.columns)
        x_default = "date" if "date" in all_cols else (cat_cols[0] if cat_cols else (num_cols[0] if num_cols else None))
        y_default = num_cols[0] if num_cols else None
        return (
            None,
            gr.Dropdown.update(choices=all_cols, value=x_default),
            gr.Dropdown.update(choices=num_cols, value=y_default),
            gr.Dropdown.update(choices=all_cols + [None], value=None),
            gr.Dropdown.update(choices=cat_cols + [None], value=None),
            df.head(10),
        )

    refresh.click(ui_refresh_columns, inputs=[file], outputs=[x, y, color, facet_col, preview])
    use_demo.click(set_demo, inputs=[use_demo], outputs=[file, x, y, color, facet_col, preview])

    # Auto-refresh columns on file change
    file.change(ui_refresh_columns, inputs=[file], outputs=[x, y, color, facet_col, preview])

    # Draw plot when any control changes
    controls = [file, plot_type, x, y, color, facet_col, agg, smooth, smooth_window, add_trend, trend_degree, points, theme, color_scale]
    for c in controls:
        c.change(ui_plot, inputs=controls, outputs=out)

    # Initial state using demo data
    _ = set_demo(None)
    out.update(value=build_plot(make_demo_data(), PlotConfig(
        plot_type="line", x="date", y="sales", color="region", facet_col=None,
        agg="mean", smooth=True, smooth_window=7, add_trend=False, trend_degree=1,
        points=False, theme="plotly_white", color_scale="Turbo"
    )))

if __name__ == "__main__":
    demo.launch()  # You can set share=True when running locally to create a public link