app.py

import gradio as gr
import nltk
import os
import pandas as pd
from nltk.tokenize import TreebankWordTokenizer
from sklearn.metrics.pairwise import cosine_similarity
from sentence_transformers import SentenceTransformer
import graphviz
from typing import Tuple, Optional
from visuals.score_card import render_score_card  # Updated import
from visuals.layout import (
    render_page_header,
    render_core_reference,
    render_pipeline,
    render_pipeline_graph,
    render_pipeline_warning,
    render_strategy_alignment,
)  # Updated import

# Ensure NLTK data is downloaded
try:
    nltk.download("punkt", quiet=True)
except Exception as e:
    print(f"Error downloading NLTK data: {e}")

# Load SentenceTransformer model
model = SentenceTransformer("all-MiniLM-L6-v2")


def calculate_ttr(text: str) -> float:
    """Calculates Type-Token Ratio (TTR) for lexical diversity."""
    if not text:
        return 0.0
    words = text.split()
    unique_words = set(words)
    return len(unique_words) / len(words) if words else 0.0


def calculate_similarity(text1: str, text2: str) -> float:
    """Calculates cosine similarity between two texts."""
    embeddings = model.encode([text1, text2])
    return cosine_similarity([embeddings[0]], [embeddings[1]])[0][0]


def calculate_mad_score(ttr: float, similarity: float) -> float:
    """Calculates the MAD score."""
    return 0.3 * (1 - ttr) + 0.7 * similarity


def get_risk_level(mad_score: float) -> str:
    """Determines the risk level based on the MAD score."""
    if mad_score > 0.7:
        return "High"
    elif 0.4 <= mad_score <= 0.7:
        return "Medium"
    else:
        return "Low"


def process_data(file_obj, model_col: str, train_col: str, data_source: str) -> Tuple[
    Optional[str],
    Optional[bytes],
    Optional[str],
    Optional[str],
    Optional[float],
    Optional[float],
    Optional[float],
]:
    """Processes the uploaded file and calculates metrics."""
    try:
        if not file_obj:
            return "Error: No file uploaded.", None, None, None, None, None, None

        file_path = file_obj.name
        if file_path.endswith(".csv"):
            df = pd.read_csv(file_path)
        elif file_path.endswith(".json"):
            df = pd.read_json(file_path)
        else:
            return (
                "Error: Invalid file type. Please upload a CSV or JSON file.",
                None,
                None,
                None,
                None,
                None,
                None,
            )

        if model_col not in df.columns or train_col not in df.columns:
            return (
                "Error: Selected columns not found in the file.",
                None,
                None,
                None,
                None,
                None,
                None,
            )

        output_text = " ".join(df[model_col].astype(str))
        train_text = " ".join(df[train_col].astype(str))

        ttr_output = calculate_ttr(output_text)
        ttr_train = calculate_ttr(train_text)
        similarity = calculate_similarity(output_text, train_text)
        mad_score = calculate_mad_score(ttr_output, similarity)
        risk_level = get_risk_level(mad_score)

        summary, details, explanation = render_score_card(
            ttr_output, ttr_train, similarity, mad_score, risk_level
        )
        evaluation_markdown = summary + details + explanation

        return (
            None,  # No error
            render_pipeline_graph(data_source),
            df.head().to_markdown(index=False, numalign="left", stralign="left"),
            evaluation_markdown,
            ttr_output,
            ttr_train,
            similarity,
        )

    except Exception as e:
        return f"An error occurred: {str(e)}", None, None, None, None, None, None


def update_dropdowns(file_obj) -> Tuple[list, str]:
    """Updates dropdown choices based on the uploaded file."""
    if not file_obj:
        return [], "No file uploaded."

    file_path = file_obj.name
    try:
        if file_path.endswith(".csv"):
            df = pd.read_csv(file_path)
        elif file_path.endswith(".json"):
            df = pd.read_json(file_path)
        else:
            return [], "Invalid file type."
        columns = df.columns.tolist()
        preview = df.head().to_markdown(index=False, numalign="left", stralign="left")
        return columns, preview
    except Exception as e:
        return [], f"Error reading file: {e}"


def main_interface():
    css = """
    .gradio-container {
        background: linear-gradient(-45deg, #e0f7fa, #e1f5fe, #f1f8e9, #fff3e0);
        background-size: 400% 400%;
        animation: oceanWaves 20s ease infinite;
    }
    @keyframes oceanWaves {
        0% { background-position: 0% 50%; }
        50% { background-position: 100% 50%; }
        100% { background-position: 0% 50%; }
    }
    """

    with gr.Blocks(css=css, title="MADGuard AI Explorer") as interface:
        gr.HTML(render_page_header())

        gr.Markdown(
            """
            > 🧠 **MADGuard AI Explorer** helps AI engineers, researchers, and MLOps teams simulate feedback loops in RAG pipelines and detect **Model Autophagy Disorder (MAD)** — where models start learning from their own outputs, leading to degraded performance.

            - Compare **real vs. synthetic input effects**
            - Visualize the data flow
            - Upload your `.csv` or `.json` data
            - Get immediate MAD risk diagnostics based on lexical diversity and semantic similarity
            """
        )

        with gr.Accordion("📚 Research Reference", open=False):
            gr.HTML(render_core_reference())

        gr.Markdown("## 1. Pipeline Simulation")
        data_source, description = render_pipeline()
        gr.HTML(description)
        pipeline_output = gr.Image(type="filepath", label="Pipeline Graph")
        warning_output = gr.HTML()
        data_source.change(
            fn=render_pipeline_warning, inputs=data_source, outputs=warning_output
        )
        data_source.change(
            fn=render_pipeline_graph, inputs=data_source, outputs=pipeline_output
        )

        gr.Markdown("## 2. Upload CSV or JSON File")
        file_input = gr.File(
            file_types=[".csv", ".json"], label="Upload a CSV or JSON file"
        )

        with gr.Row():
            model_col_input = gr.Dropdown(
                choices=[], label="Select column for model output"
            )
            train_col_input = gr.Dropdown(
                choices=[], label="Select column for future training data"
            )

        file_preview = gr.Markdown(label="📄 File Preview")

        output_markdown = gr.Markdown(label="🔍 Evaluation Summary")

        with gr.Accordion("📋 Research-Based Strategy Alignment", open=False):
            gr.HTML(render_strategy_alignment())

        with gr.Row():
            ttr_output_metric = gr.Number(label="Lexical Diversity (Output)")
            ttr_train_metric = gr.Number(label="Lexical Diversity (Training Set)")
            similarity_metric = gr.Number(label="Semantic Similarity (Cosine)")

        file_input.change(
            update_dropdowns,
            inputs=file_input,
            outputs=[model_col_input, train_col_input, file_preview],
        )

        def process_and_generate(
            file_obj, model_col_val: str, train_col_val: str, data_source_val: str
        ):
            error, graph, preview, markdown, ttr_out, ttr_tr, sim = process_data(
                file_obj, model_col_val, train_col_val, data_source_val
            )
            if error:
                return error, graph, warning_output, preview, None, None, None, None
            return (
                "",
                graph,
                render_pipeline_warning(data_source_val),
                preview,
                markdown,
                ttr_out,
                ttr_tr,
                sim,
            )

        inputs = [file_input, model_col_input, train_col_input, data_source]
        outputs = [
            gr.Markdown(label="⚠️ Error Message"),
            pipeline_output,
            warning_output,
            file_preview,
            output_markdown,
            ttr_output_metric,
            ttr_train_metric,
            similarity_metric,
        ]
        for input_component in inputs:
            input_component.change(
                fn=process_and_generate, inputs=inputs, outputs=outputs
            )

        gr.Markdown("---")
        gr.Markdown(
            """
        **The upcoming Pro version of MADGuard will allow:**
        - 📂 Bulk upload support for `.csv` files or folders of `.txt` documents  
        - 📊 Automated batch scoring with trend visualizations over time  
        - 🧾 One-click export of audit-ready diagnostic reports

        [**📩 Join the waitlist**](https://docs.google.com/forms/d/e/1FAIpQLSfAPPC_Gm7DQElQSWGSnoB6T5hMxb_rXSu48OC8E6TNGZuKgQ/viewform?usp=sharing&ouid=118007615320536574300)
        """
        )

    return interface


# Launch the Gradio interface
if __name__ == "__main__":
    interface = main_interface()
    interface.launch(server_name="0.0.0.0", server_port=7860)