"""
Data Cleaning & Visualization Project
A Gradio-based interactive dashboard for cleaning raw datasets and generating insights.
"""

import gradio as gr
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from io import StringIO
import warnings

warnings.filterwarnings("ignore")
sns.set_theme(style="whitegrid", palette="muted")
plt.rcParams["figure.figsize"] = (10, 5)
plt.rcParams["figure.dpi"] = 100


# ---------------------------------------------------------------------------
# Helper functions
# ---------------------------------------------------------------------------

def load_dataframe(file):
    """Load CSV or Excel file into a DataFrame."""
    if file is None:
        return None, "Please upload a file first."
    try:
        path = file.name if hasattr(file, "name") else file
        if str(path).lower().endswith((".xlsx", ".xls")):
            df = pd.read_excel(path)
        else:
            df = pd.read_csv(path)
        return df, f"Loaded successfully • Shape: {df.shape[0]} rows × {df.shape[1]} columns"
    except Exception as e:
        return None, f"Error loading file: {str(e)}"


def get_overview(df):
    """Return text overview of the dataset."""
    if df is None:
        return "No data loaded."

    buffer = StringIO()
    df.info(buf=buffer)
    info_str = buffer.getvalue()

    missing = df.isnull().sum()
    missing_pct = (missing / len(df) * 100).round(2)
    missing_df = pd.DataFrame({"Missing Count": missing, "Missing %": missing_pct})
    missing_df = missing_df[missing_df["Missing Count"] > 0].sort_values(
        "Missing Count", ascending=False
    )

    overview = f"""
### Dataset Overview
- **Rows:** {df.shape[0]:,}
- **Columns:** {df.shape[1]}
- **Memory Usage:** {df.memory_usage(deep=True).sum() / 1024**2:.2f} MB
- **Duplicate Rows:** {df.duplicated().sum():,}
- **Numeric Columns:** {len(df.select_dtypes(include=np.number).columns)}
- **Categorical Columns:** {len(df.select_dtypes(include=["object", "category"]).columns)}

### Missing Values
"""
    if missing_df.empty:
        overview += "No missing values found.\n"
    else:
        overview += missing_df.to_markdown() + "\n"

    overview += "\n### Column Data Types\n"
    overview += df.dtypes.to_frame("dtype").to_markdown()

    overview += "\n\n### First 5 Rows\n"
    overview += df.head().to_markdown(index=False)

    return overview


def handle_missing(df, strategy, fill_value=None):
    """Handle missing values according to chosen strategy."""
    if df is None:
        return None, "No data loaded."

    df = df.copy()
    before = df.isnull().sum().sum()

    if strategy == "Drop rows with any missing":
        df = df.dropna()
    elif strategy == "Drop columns with >50% missing":
        thresh = len(df) * 0.5
        df = df.dropna(axis=1, thresh=thresh)
    elif strategy == "Fill numeric with mean":
        num_cols = df.select_dtypes(include=np.number).columns
        df[num_cols] = df[num_cols].fillna(df[num_cols].mean())
    elif strategy == "Fill numeric with median":
        num_cols = df.select_dtypes(include=np.number).columns
        df[num_cols] = df[num_cols].fillna(df[num_cols].median())
    elif strategy == "Fill categorical with mode":
        cat_cols = df.select_dtypes(include=["object", "category"]).columns
        for col in cat_cols:
            mode_val = df[col].mode()
            if not mode_val.empty:
                df[col] = df[col].fillna(mode_val[0])
    elif strategy == "Fill with custom value":
        if fill_value is None or fill_value == "":
            return df, "Please provide a custom fill value."
        try:
            # try numeric conversion
            val = float(fill_value)
            df = df.fillna(val)
        except ValueError:
            df = df.fillna(fill_value)
    elif strategy == "Forward fill":
        df = df.ffill()
    elif strategy == "Backward fill":
        df = df.bfill()

    after = df.isnull().sum().sum()
    msg = f"Missing values: {before} → {after}  |  Rows remaining: {len(df):,}"
    return df, msg


def remove_duplicates(df):
    """Remove duplicate rows."""
    if df is None:
        return None, "No data loaded."
    before = len(df)
    df = df.drop_duplicates()
    after = len(df)
    return df, f"Duplicates removed: {before - after}  |  Rows remaining: {after:,}"


def detect_outliers_iqr(df, columns=None):
    """Detect outliers using IQR method. Returns boolean mask of outlier rows."""
    if df is None:
        return None
    num_cols = df.select_dtypes(include=np.number).columns.tolist()
    if columns:
        num_cols = [c for c in columns if c in num_cols]
    if not num_cols:
        return pd.Series([False] * len(df), index=df.index)

    outlier_mask = pd.Series([False] * len(df), index=df.index)
    for col in num_cols:
        Q1 = df[col].quantile(0.25)
        Q3 = df[col].quantile(0.75)
        IQR = Q3 - Q1
        lower = Q1 - 1.5 * IQR
        upper = Q3 + 1.5 * IQR
        outlier_mask |= (df[col] < lower) | (df[col] > upper)
    return outlier_mask


def handle_outliers(df, method, columns=None):
    """Handle outliers: remove or cap them."""
    if df is None:
        return None, "No data loaded."

    df = df.copy()
    num_cols = df.select_dtypes(include=np.number).columns.tolist()
    if columns:
        selected = [c.strip() for c in columns.split(",") if c.strip()]
        num_cols = [c for c in selected if c in num_cols]
    if not num_cols:
        return df, "No numeric columns selected for outlier handling."

    mask = detect_outliers_iqr(df, num_cols)
    n_outliers = mask.sum()

    if method == "Remove outlier rows":
        df = df[~mask]
        msg = f"Removed {n_outliers} outlier rows  |  Rows remaining: {len(df):,}"
    elif method == "Cap outliers (winsorize)":
        for col in num_cols:
            Q1 = df[col].quantile(0.25)
            Q3 = df[col].quantile(0.75)
            IQR = Q3 - Q1
            lower = Q1 - 1.5 * IQR
            upper = Q3 + 1.5 * IQR
            df[col] = df[col].clip(lower=lower, upper=upper)
        msg = f"Capped outliers in columns: {', '.join(num_cols)}  |  Outliers affected: {n_outliers}"
    else:
        msg = "No action taken."
    return df, msg


def plot_missing_heatmap(df):
    """Visualize missing values."""
    if df is None or df.empty:
        return None
    fig, ax = plt.subplots(figsize=(12, 5))
    sns.heatmap(df.isnull(), cbar=False, yticklabels=False, cmap="YlOrRd", ax=ax)
    ax.set_title("Missing Values Heatmap (yellow/red = missing)")
    plt.tight_layout()
    return fig


def plot_distributions(df, columns=None):
    """Plot histograms / KDE for numeric columns."""
    if df is None:
        return None
    num_cols = df.select_dtypes(include=np.number).columns.tolist()
    if columns:
        selected = [c.strip() for c in columns.split(",") if c.strip()]
        num_cols = [c for c in selected if c in num_cols]
    if not num_cols:
        return None

    n = min(len(num_cols), 6)
    cols = num_cols[:n]
    fig, axes = plt.subplots(1, n, figsize=(4 * n, 4))
    if n == 1:
        axes = [axes]
    for ax, col in zip(axes, cols):
        sns.histplot(df[col].dropna(), kde=True, ax=ax, color="#4C72B0")
        ax.set_title(col)
        ax.set_xlabel("")
    plt.suptitle("Numeric Distributions", y=1.02)
    plt.tight_layout()
    return fig


def plot_boxplots(df, columns=None):
    """Boxplots to show outliers."""
    if df is None:
        return None
    num_cols = df.select_dtypes(include=np.number).columns.tolist()
    if columns:
        selected = [c.strip() for c in columns.split(",") if c.strip()]
        num_cols = [c for c in selected if c in num_cols]
    if not num_cols:
        return None

    n = min(len(num_cols), 6)
    cols = num_cols[:n]
    fig, axes = plt.subplots(1, n, figsize=(3.5 * n, 4))
    if n == 1:
        axes = [axes]
    for ax, col in zip(axes, cols):
        sns.boxplot(y=df[col], ax=ax, color="#55A868")
        ax.set_title(col)
    plt.suptitle("Boxplots (Outlier Detection)", y=1.02)
    plt.tight_layout()
    return fig


def plot_correlation(df):
    """Correlation heatmap for numeric features."""
    if df is None:
        return None
    num_df = df.select_dtypes(include=np.number)
    if num_df.shape[1] < 2:
        return None
    corr = num_df.corr()
    fig, ax = plt.subplots(figsize=(10, 8))
    sns.heatmap(
        corr,
        annot=True,
        fmt=".2f",
        cmap="coolwarm",
        center=0,
        square=True,
        linewidths=0.5,
        ax=ax,
        cbar_kws={"shrink": 0.8},
    )
    ax.set_title("Correlation Matrix")
    plt.tight_layout()
    return fig


def plot_categorical(df, column=None):
    """Bar chart for a categorical column."""
    if df is None:
        return None
    cat_cols = df.select_dtypes(include=["object", "category"]).columns.tolist()
    if not cat_cols:
        return None
    col = column if column and column in cat_cols else cat_cols[0]
    counts = df[col].value_counts().head(15)
    fig, ax = plt.subplots(figsize=(10, 5))
    sns.barplot(x=counts.values, y=counts.index, ax=ax, palette="viridis")
    ax.set_title(f"Top categories in '{col}'")
    ax.set_xlabel("Count")
    plt.tight_layout()
    return fig


def generate_summary_report(df):
    """Generate a storytelling-style text report of key findings."""
    if df is None:
        return "No data available."

    lines = ["# Data Insights Report\n"]
    lines.append(f"**Dataset size after cleaning:** {df.shape[0]:,} rows × {df.shape[1]} columns\n")

    # Missing
    missing_total = df.isnull().sum().sum()
    lines.append(f"- Remaining missing values: **{missing_total}**\n")

    # Duplicates
    lines.append(f"- Duplicate rows remaining: **{df.duplicated().sum()}**\n")

    # Numeric summary
    num_df = df.select_dtypes(include=np.number)
    if not num_df.empty:
        lines.append("\n## Numeric Features Summary\n")
        desc = num_df.describe().T[["mean", "std", "min", "50%", "max"]].round(2)
        lines.append(desc.to_markdown())
        lines.append("\n")

        # Strongest correlations
        if num_df.shape[1] >= 2:
            corr = num_df.corr().abs()
            # zero out diagonal without mutating original
            corr_no_diag = corr.where(~np.eye(corr.shape[0], dtype=bool), 0)
            if corr_no_diag.max().max() > 0:
                max_pair = corr_no_diag.stack().idxmax()
                lines.append(
                    f"\n**Strongest correlation:** `{max_pair[0]}` ↔ `{max_pair[1]}` "
                    f"(r = {num_df.corr().loc[max_pair]:.3f})\n"
                )

    # Categorical
    cat_cols = df.select_dtypes(include=["object", "category"]).columns
    if len(cat_cols) > 0:
        lines.append("\n## Categorical Features\n")
        for col in cat_cols[:5]:
            top = df[col].value_counts().head(3)
            lines.append(f"- **{col}** top values: " + ", ".join(f"`{k}` ({v})" for k, v in top.items()))
        lines.append("\n")

    lines.append("\n---\n*Report generated by the Data Cleaning & Visualization dashboard.*")
    return "\n".join(lines)


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

state_df = gr.State(value=None)


def on_upload(file, current_df):
    df, msg = load_dataframe(file)
    if df is not None:
        return df, msg, get_overview(df), df.head(20)
    return current_df, msg, "Upload failed.", None


def on_clean_missing(df, strategy, fill_value):
    if df is None:
        return None, "Upload data first.", None, None
    new_df, msg = handle_missing(df, strategy, fill_value)
    return new_df, msg, get_overview(new_df), new_df.head(20)


def on_remove_dupes(df):
    if df is None:
        return None, "Upload data first.", None, None
    new_df, msg = remove_duplicates(df)
    return new_df, msg, get_overview(new_df), new_df.head(20)


def on_outliers(df, method, columns):
    if df is None:
        return None, "Upload data first.", None, None
    new_df, msg = handle_outliers(df, method, columns)
    return new_df, msg, get_overview(new_df), new_df.head(20)


def on_visualize(df, dist_cols, box_cols, cat_col):
    if df is None:
        return None, None, None, None, "Upload & clean data first."
    fig_miss = plot_missing_heatmap(df)
    fig_dist = plot_distributions(df, dist_cols)
    fig_box = plot_boxplots(df, box_cols)
    fig_corr = plot_correlation(df)
    fig_cat = plot_categorical(df, cat_col)
    report = generate_summary_report(df)
    return fig_miss, fig_dist, fig_box, fig_corr, fig_cat, report


def download_cleaned(df):
    if df is None:
        return None
    path = "/tmp/cleaned_dataset.csv"
    df.to_csv(path, index=False)
    return path


# ---------------------------------------------------------------------------
# Build the Gradio interface
# ---------------------------------------------------------------------------

CUSTOM_CSS = """
.gradio-container { max-width: 1200px !important; }
.header-text { text-align: center; margin-bottom: 1rem; }
"""

with gr.Blocks(
    title="Data Cleaning & Visualization",
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="slate"),
    css=CUSTOM_CSS,
) as demo:

    gr.Markdown(
        """
        # Data Cleaning & Visualization Project
        **Upload a raw dataset → clean missing values, outliers & duplicates → explore insights with interactive charts.**
        """
    )

    current_df = gr.State(None)

    with gr.Tab("1. Upload & Overview"):
        with gr.Row():
            file_input = gr.File(
                label="Upload CSV or Excel file",
                file_types=[".csv", ".xlsx", ".xls"],
            )
            status_msg = gr.Textbox(label="Status", interactive=False)

        overview_md = gr.Markdown("Upload a file to see the dataset overview.")
        preview_table = gr.Dataframe(label="Data Preview (first 20 rows)", interactive=False)

        file_input.change(
            on_upload,
            inputs=[file_input, current_df],
            outputs=[current_df, status_msg, overview_md, preview_table],
        )

    with gr.Tab("2. Clean Data"):
        gr.Markdown("### Handle Missing Values")
        with gr.Row():
            missing_strategy = gr.Dropdown(
                choices=[
                    "Drop rows with any missing",
                    "Drop columns with >50% missing",
                    "Fill numeric with mean",
                    "Fill numeric with median",
                    "Fill categorical with mode",
                    "Fill with custom value",
                    "Forward fill",
                    "Backward fill",
                ],
                value="Fill numeric with median",
                label="Missing-value strategy",
            )
            custom_fill = gr.Textbox(
                label="Custom fill value (only for 'Fill with custom value')",
                placeholder="e.g. 0 or Unknown",
            )
        btn_missing = gr.Button("Apply Missing-Value Handling", variant="primary")

        gr.Markdown("### Remove Duplicates")
        btn_dupes = gr.Button("Remove Duplicate Rows", variant="secondary")

        gr.Markdown("### Handle Outliers (IQR method)")
        with gr.Row():
            outlier_method = gr.Dropdown(
                choices=["Remove outlier rows", "Cap outliers (winsorize)"],
                value="Cap outliers (winsorize)",
                label="Outlier handling method",
            )
            outlier_cols = gr.Textbox(
                label="Columns (comma-separated, leave blank for all numeric)",
                placeholder="e.g. age, salary, price",
            )
        btn_outliers = gr.Button("Apply Outlier Handling", variant="primary")

        clean_status = gr.Textbox(label="Cleaning Status", interactive=False)
        clean_overview = gr.Markdown()
        clean_preview = gr.Dataframe(label="Cleaned Preview", interactive=False)

        btn_missing.click(
            on_clean_missing,
            inputs=[current_df, missing_strategy, custom_fill],
            outputs=[current_df, clean_status, clean_overview, clean_preview],
        )
        btn_dupes.click(
            on_remove_dupes,
            inputs=[current_df],
            outputs=[current_df, clean_status, clean_overview, clean_preview],
        )
        btn_outliers.click(
            on_outliers,
            inputs=[current_df, outlier_method, outlier_cols],
            outputs=[current_df, clean_status, clean_overview, clean_preview],
        )

        gr.Markdown("---")
        download_btn = gr.Button("Download Cleaned CSV")
        download_file = gr.File(label="Download")
        download_btn.click(download_cleaned, inputs=[current_df], outputs=[download_file])

    with gr.Tab("3. Visualize & Insights"):
        gr.Markdown("Select columns (optional) and generate plots + an automatic insights report.")
        with gr.Row():
            dist_cols = gr.Textbox(
                label="Distribution columns (comma-separated)",
                placeholder="Leave blank for first 6 numeric columns",
            )
            box_cols = gr.Textbox(
                label="Boxplot columns (comma-separated)",
                placeholder="Leave blank for first 6 numeric columns",
            )
            cat_col = gr.Textbox(
                label="Categorical column for bar chart",
                placeholder="Leave blank for first categorical column",
            )
        btn_viz = gr.Button("Generate Visualizations & Report", variant="primary")

        with gr.Row():
            plot_miss = gr.Plot(label="Missing Values Heatmap")
            plot_corr = gr.Plot(label="Correlation Matrix")
        with gr.Row():
            plot_dist = gr.Plot(label="Distributions")
            plot_box = gr.Plot(label="Boxplots")
        plot_cat = gr.Plot(label="Categorical Breakdown")
        insights_md = gr.Markdown(label="Insights Report")

        btn_viz.click(
            on_visualize,
            inputs=[current_df, dist_cols, box_cols, cat_col],
            outputs=[plot_miss, plot_dist, plot_box, plot_corr, plot_cat, insights_md],
        )

    with gr.Tab("About"):
        gr.Markdown(
            """
            ## Project Goals
            This interactive dashboard demonstrates the complete workflow of a **Data Cleaning & Visualization** project:

            ### Key Features Implemented
            - **Handle missing values** – multiple strategies (drop, mean/median/mode fill, ffill/bfill, custom)
            - **Remove duplicates**
            - **Detect & treat outliers** using the IQR method (remove or winsorize/cap)
            - **Visualizations** with Matplotlib & Seaborn:
              - Missing-value heatmap
              - Histograms / KDE distributions
              - Boxplots for outlier inspection
              - Correlation matrix
              - Categorical bar charts
            - **Automated insights report** that tells a short data story

            ### Expected Learning Outcomes
            - Data preprocessing techniques
            - Exploratory data visualization
            - Storytelling with data

            ### How to use
            1. Upload any CSV or Excel file in **Tab 1**
            2. Clean the data step-by-step in **Tab 2**
            3. Generate charts and the insights report in **Tab 3**
            4. Download the cleaned dataset when ready

            Built with **Gradio**, **Pandas**, **Matplotlib** & **Seaborn**.
            """
        )

if __name__ == "__main__":
    demo.launch()