promptToCode.py

import pandas as pd
import textwrap
import re 
import textwrap
# 🧹 COMMON TYPO NORMALIZER
def normalize_prompt(prompt: str) -> str:
    fixes = {
        "histogramm": "histogram",
        "frequennt": "frequent",
        "statisstics": "statistics",
        "desscending": "descending",
        "commparing": "comparing",
        "hheatmap": "heatmap",
        "z--score": "z-score",
        "  ": " ", 
    }
    p = prompt.lower().strip()
    for wrong, right in fixes.items():
        p = p.replace(wrong, right)
    return p


def prompt_to_code(prompt: str, df: pd.DataFrame):
    """
    Convert known prompt templates into executable python code strings.
    If unrecognized, return None so Gemini can take over.
    """
    p = normalize_prompt(prompt)

    # ----------------------------------------
    # 1️⃣ SUMMARY
    # ----------------------------------------
  

    if p.startswith("summarize the dataset"):
        return textwrap.dedent("""
        summary = {
            "Metric": [
                "Total Rows",
                "Total Columns",
                "Numeric Columns",
                "Categorical Columns",
                "Missing Value Columns",
            ],
            "Value": [
                len(df),
                len(df.columns),
                len(df.select_dtypes(include='number').columns),
                len(df.select_dtypes(include='object').columns),
                df.isnull().sum().astype(bool).sum(),
            ]
        }

        result = pd.DataFrame(summary)
    """)


    # ----------------------------------------
    # 2️⃣ FREQUENT VALUES (Top 10)
    # ----------------------------------------
    if "top 10 most frequent values" in p:
        col = re.findall(r"'([^']+)'", prompt)
        if col:
            col = col[0]
            return textwrap.dedent(f"""
                result = df['{col}'].value_counts(dropna=False).head(10).reset_index()
                result.columns = ['value','count']
            """)

    # ----------------------------------------
    # 3️⃣ SUMMARY STATISTICS
    # ----------------------------------------
    if "summary statistics" in p or "describe" in p:
        return "result = df.select_dtypes(include=['number']).describe().T"

    # ----------------------------------------
    # 4️⃣ HISTOGRAM
    # ----------------------------------------
    if "histogram" in p:
        col = re.findall(r"'([^']+)'", prompt)
        if col:
            col = col[0]
            return textwrap.dedent(f"""
                plt.figure(figsize=(6,4))
                df['{col}'].dropna().astype(float).hist(bins=30)
                plt.title('Histogram of {col}')
                plt.xlabel('{col}')
                plt.ylabel('count')
                result_img_path = None
            """)

    # ----------------------------------------
    # 5️⃣ SORT DESC TOP 10
    # ----------------------------------------
    if "sort dataset by" in p and "descending" in p:
        col = re.findall(r"'([^']+)'", prompt)
        if col:
            col = col[0]
            return textwrap.dedent(f"""
                result = df.sort_values('{col}', ascending=False).head(10).reset_index(drop=True)
            """)

    # ----------------------------------------
    # 6️⃣ SCATTER PLOT
    # ----------------------------------------
    if "scatter" in p and "comparing" in p:
        cols = re.findall(r"'([^']+)'", prompt)
        if len(cols) >= 2:
            x, y = cols[0], cols[1]
            return textwrap.dedent(f"""
                plt.figure(figsize=(6,4))
                plt.scatter(df['{x}'], df['{y}'])
                plt.xlabel('{x}')
                plt.ylabel('{y}')
                plt.title('{x} vs {y}')
                result_img_path = None
            """)

    # ----------------------------------------
    # 7️⃣ CORRELATION HEATMAP
    # ----------------------------------------
    if "correlation heatmap" in p:
        return textwrap.dedent("""
            corr = df.select_dtypes(include=['number']).corr()
            plt.figure(figsize=(6,5))
            plt.imshow(corr, cmap='viridis', aspect='auto')
            plt.colorbar()
            plt.xticks(range(len(corr)), corr.columns, rotation=90)
            plt.yticks(range(len(corr)), corr.columns)
            plt.title('Correlation heatmap')
            result_img_path = None
        """)

    # ----------------------------------------
    # 8️⃣ OUTLIERS (Z-SCORE)
    # ----------------------------------------
    if "z-score" in p or "outlier" in p:
        return textwrap.dedent("""
            from scipy import stats
            num = df.select_dtypes(include=['number'])
            if num.shape[1] == 0:
                result = pd.DataFrame()
            else:
                z = abs(stats.zscore(num))
                mask = (z > 3).any(axis=1)
                result = df[mask].head(20).reset_index(drop=True)
        """)

    # ❌ UNKNOWN → SEND TO GEMINI
    return None