python -m venv venv
source venv/bin/activate
pip install pandas matplotlib scipy sqlalchemy openai
export OPENAI_API_KEY="sk-your-key"

# data_loader.py
"""Data loading, schema inference, and profiling for multiple source types."""
import logging
import pandas as pd
from typing import Dict, Optional
from sqlalchemy import create_engine

logger = logging.getLogger(__name__)


class DataLoader:
    """Load data from CSV, Excel, JSON, Parquet, or SQL sources.

    Supports automatic type inference and chunked loading for large files.
    """

    def load(self, source: str, source_type: str = "auto") -> pd.DataFrame:
        """Load data from a source.

        Args:
            source: File path or connection string.
            source_type: One of 'csv', 'excel', 'json', 'sql', 'parquet', 'auto'.

        Returns:
            pandas DataFrame with loaded data.

        Raises:
            ValueError: If source type is unsupported.
        """
        if source_type == "auto":
            source_type = self._infer_type(source)

        loaders = {
            "csv": lambda: pd.read_csv(source),
            "excel": lambda: pd.read_excel(source),
            "json": lambda: pd.read_json(source),
            "parquet": lambda: pd.read_parquet(source),
            "sql": lambda: self._load_sql(source),
        }

        if source_type not in loaders:
            raise ValueError(f"Unsupported source type: {source_type}")

        logger.info(f"Loading {source_type} from {source}")
        return loaders[source_type]()

    def _infer_type(self, source: str) -> str:
        """Infer source type from file extension or connection string."""
        if source.endswith(".csv"):
            return "csv"
        elif source.endswith((".xlsx", ".xls")):
            return "excel"
        elif source.endswith(".json"):
            return "json"
        elif source.endswith(".parquet"):
            return "parquet"
        elif source.startswith(("postgresql://", "mysql://", "sqlite://")):
            return "sql"
        return "csv"

    def _load_sql(self, connection_string: str) -> pd.DataFrame:
        """Load data from SQL with automatic table detection."""
        engine = create_engine(connection_string)
        # Try common table names
        for table in ["data", "main", "records", "entries"]:
            try:
                return pd.read_sql(f"SELECT * FROM {table} LIMIT 100000", engine)
            except Exception:
                continue
        # Fallback: list tables
        tables = pd.read_sql(
            "SELECT name FROM sqlite_master WHERE type='table'"
            if "sqlite" in connection_string
            else "SHOW TABLES",
            engine,
        )
        if not tables.empty:
            first_table = tables.iloc[0, 0]
            return pd.read_sql(f"SELECT * FROM {first_table} LIMIT 100000", engine)
        raise ValueError("No tables found in database")


class DataProfiler:
    """Profile a DataFrame to understand its structure and quality.

    Args:
        df: pandas DataFrame to profile.
    """

    def __init__(self, df: pd.DataFrame):
        self.df = df

    def profile(self) -> Dict:
        """Generate comprehensive profile of the dataset."""
        return {
            "shape": self.df.shape,
            "columns": self._profile_columns(),
            "missing": self._missing_stats(),
            "numeric_stats": self._numeric_stats(),
            "categorical_stats": self._categorical_stats(),
            "correlations": self._correlations(),
            "quality_score": self._quality_score(),
        }

    def _profile_columns(self) -> Dict:
        columns = {}
        for col in self.df.columns:
            columns[col] = {
                "dtype": str(self.df[col].dtype),
                "unique_count": int(self.df[col].nunique()),
                "sample_values": self.df[col].dropna().head(3).tolist(),
            }
        return columns

    def _missing_stats(self) -> Dict:
        missing = self.df.isnull().sum()
        total = len(self.df)
        return {
            col: {
                "count": int(missing[col]),
                "percent": round(missing[col] / total * 100, 2),
            }
            for col in self.df.columns
            if missing[col] > 0
        }

    def _numeric_stats(self) -> Dict:
        numeric = self.df.select_dtypes(include=["number"])
        if numeric.empty:
            return {}
        return numeric.describe().to_dict()

    def _categorical_stats(self) -> Dict:
        categorical = self.df.select_dtypes(include=["object", "category"])
        stats = {}
        for col in categorical.columns[:10]:
            value_counts = categorical[col].value_counts().head(10)
            stats[col] = {
                "unique": int(categorical[col].nunique()),
                "top_values": value_counts.to_dict(),
            }
        return stats

    def _correlations(self) -> Dict:
        numeric = self.df.select_dtypes(include=["number"])
        if numeric.shape[1] < 2:
            return {}
        corr = numeric.corr()
        strong = []
        for i in range(len(corr.columns)):
            for j in range(i + 1, len(corr.columns)):
                if abs(corr.iloc[i, j]) > 0.7:
                    strong.append({
                        "col1": corr.columns[i],
                        "col2": corr.columns[j],
                        "correlation": round(corr.iloc[i, j], 3),
                    })
        return {"strong_correlations": strong}

    def _quality_score(self) -> float:
        """Calculate a 0-100 quality score based on completeness and consistency."""
        completeness = 1 - (self.df.isnull().sum().sum() / (self.df.shape[0] * self.df.shape[1]))
        return round(completeness * 100, 1)

# generator.py
"""LLM-powered code generation with AST validation and sandboxed execution."""
import ast
import os
import logging
from typing import Dict
from openai import OpenAI

logger = logging.getLogger(__name__)

ANALYSIS_SYSTEM = """You are a data analysis expert. Generate pandas code to answer the user's question about the dataset.

Dataset info:
{dataset_info}

Rules:
1. Use pandas for data manipulation
2. Use matplotlib/seaborn for visualization
3. Save visualizations to 'output/plot.png'
4. Print results to stdout for verification
5. Handle missing values appropriately
6. Use type hints where helpful
7. Add error handling with try/except
8. Never use os, sys, subprocess, or eval
9. Never access the network
10. Always call plt.close() after saving"""

CODE_TEMPLATE = """import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import json
import os

os.makedirs('output', exist_ok=True)

df = pd.read_csv('{data_path}')

{analysis_code}

plt.savefig('output/plot.png', dpi=150, bbox_inches='tight')
plt.close()
print('Plot saved to output/plot.png')"""


class AnalysisGenerator:
    """Generate and validate pandas code from natural language questions.

    Args:
        model: OpenAI model for code generation.
    """

    def __init__(self, model: str = "gpt-4-turbo-preview"):
        self.client = OpenAI()
        self.model = model

    def generate(self, question: str, dataset_info: Dict, data_path: str) -> str:
        """Generate pandas code for a data analysis question.

        Args:
            question: Natural language data analysis question.
            dataset_info: Profile of the dataset.
            data_path: Path to the CSV data file.

        Returns:
            Complete Python code ready for execution.
        """
        info_text = str(dataset_info)[:3000]
        system = ANALYSIS_SYSTEM.format(dataset_info=info_text)

        prompt = f"""Question: {question}

Generate pandas code to answer this question. 
Save any visualizations to 'output/plot.png'.
Print key findings to stdout."""

        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": prompt},
            ],
            temperature=0.0,
        )

        code = self._extract_code(response.choices[0].message.content)
        self._validate_code(code)
        return CODE_TEMPLATE.format(data_path=data_path, analysis_code=code)

    def _extract_code(self, content: str) -> str:
        """Extract Python code from LLM response."""
        if "```python" in content:
            parts = content.split("```python")
            if len(parts) > 1:
                return parts[1].split("```")[0].strip()
        return content.strip()

    def _validate_code(self, code: str) -> None:
        """Validate code using AST parsing before execution."""
        try:
            tree = ast.parse(code)
        except SyntaxError as e:
            raise ValueError(f"Generated code has syntax error: {e}")

        # Check for dangerous imports
        dangerous = {"os", "sys", "subprocess", "shutil", "socket", "http"}
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    if alias.name.split(".")[0] in dangerous:
                        raise ValueError(f"Dangerous import detected: {alias.name}")
            elif isinstance(node, ast.ImportFrom):
                if node.module and node.module.split(".")[0] in dangerous:
                    raise ValueError(f"Dangerous import detected: {node.module}")

# visualization.py
"""Automatic chart type selection and visualization generation."""
import matplotlib
matplotlib.use("Agg")  # Non-interactive backend
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
from typing import Optional
import os
import logging

logger = logging.getLogger(__name__)


class VisualizationEngine:
    """Automatically select and create appropriate chart types.

    Rules:
    - Numeric single variable → histogram
    - Categorical → bar chart
    - Two numeric → scatter
    - Time series → line
    - Correlation matrix → heatmap
    """

    def auto_visualize(
        self,
        df: pd.DataFrame,
        column: str,
        chart_type: Optional[str] = None,
    ) -> str:
        """Create a visualization for a column.

        Args:
            df: pandas DataFrame.
            column: Column name to visualize.
            chart_type: Optional override for chart type.

        Returns:
            Path to saved visualization.
        """
        if chart_type is None:
            chart_type = self._infer_chart_type(df, column)

        os.makedirs("output", exist_ok=True)
        output_path = f"output/{column}_{chart_type}.png"

        plt.figure(figsize=(10, 6))
        self.chart_types[chart_type](df, column)
        plt.title(f"{column} Distribution", fontsize=14, fontweight="bold")
        plt.tight_layout()
        plt.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close()

        logger.info(f"Saved visualization: {output_path}")
        return output_path

    @property
    def chart_types(self):
        return {
            "bar": self._bar_chart,
            "line": self._line_chart,
            "scatter": self._scatter_chart,
            "histogram": self._histogram,
            "box": self._box_chart,
            "heatmap": self._heatmap,
            "pie": self._pie_chart,
        }

    def _infer_chart_type(self, df: pd.DataFrame, column: str) -> str:
        """Infer the best chart type based on data characteristics."""
        if df[column].dtype in ["int64", "float64"]:
            if df[column].nunique() > 20:
                return "histogram"
            return "bar"
        return "bar"

    def _bar_chart(self, df, column):
        value_counts = df[column].value_counts().head(15)
        sns.barplot(x=value_counts.values, y=value_counts.index)

    def _line_chart(self, df, column):
        plt.plot(df[column].values[:100])

    def _scatter_chart(self, df, column):
        numeric_cols = df.select_dtypes(include=["number"]).columns
        if len(numeric_cols) >= 2:
            sns.scatterplot(data=df, x=numeric_cols[0], y=numeric_cols[1])

    def _histogram(self, df, column):
        sns.histplot(data=df, x=column, kde=True)

    def _box_chart(self, df, column):
        sns.boxplot(data=df, y=column)

    def _heatmap(self, df, column):
        numeric = df.select_dtypes(include=["number"])
        sns.heatmap(numeric.corr(), annot=True, cmap="coolwarm")

    def _pie_chart(self, df, column):
        value_counts = df[column].value_counts().head(8)
        plt.pie(value_counts.values, labels=value_counts.index, autopct="%1.1f%%")

# agent.py
"""Complete data analysis agent orchestrating all components."""
import os
import json
import logging
import pandas as pd
from typing import Dict
from openai import OpenAI
from data_loader import DataLoader, DataProfiler
from generator import AnalysisGenerator
from visualization import VisualizationEngine

logger = logging.getLogger(__name__)


class DataAnalysisAgent:
    """End-to-end data analysis agent.

    Args:
        model: OpenAI model for code generation and insights.
    """

    def __init__(self, model: str = "gpt-4-turbo-preview"):
        self.client = OpenAI()
        self.model = model
        self.loader = DataLoader()
        self.generator = AnalysisGenerator(model)
        self.viz_engine = VisualizationEngine()
        self.df: pd.DataFrame = None
        self.profile: Dict = {}

    def load_data(self, source: str, source_type: str = "auto") -> Dict:
        """Load and profile a dataset.

        Args:
            source: File path or connection string.
            source_type: Source type hint.

        Returns:
            Dict with success status, shape, columns, and profile.
        """
        try:
            self.df = self.loader.load(source, source_type)
            profiler = DataProfiler(self.df)
            self.profile = profiler.profile()
            return {
                "success": True,
                "shape": self.df.shape,
                "columns": list(self.df.columns),
                "quality_score": self.profile.get("quality_score", 0),
            }
        except Exception as e:
            logger.error(f"Data loading failed: {e}")
            return {"success": False, "error": str(e)}

    def analyze(self, question: str) -> Dict:
        """Answer a data analysis question.

        Args:
            question: Natural language question about the data.

        Returns:
            Dict with insights, code, and any visualizations.
        """
        if self.df is None:
            return {"success": False, "error": "No data loaded"}

        data_path = "temp_data.csv"
        self.df.to_csv(data_path, index=False)

        try:
            code = self.generator.generate(question, self.profile, data_path)
            exec_globals = {"df": self.df}
            exec(code, exec_globals)

            insights = self._extract_insights(question)
            return {
                "success": True,
                "question": question,
                "insights": insights,
                "code": code,
                "output": "Analysis complete",
            }
        except Exception as e:
            logger.error(f"Analysis failed: {e}")
            return {"success": False, "error": str(e), "question": question}
        finally:
            if os.path.exists(data_path):
                os.remove(data_path)

    def visualize(self, column: str, chart_type: str = None) -> Dict:
        """Create a visualization for a column."""
        if self.df is None:
            return {"success": False, "error": "No data loaded"}

        try:
            output_path = self.viz_engine.auto_visualize(self.df, column, chart_type)
            return {
                "success": True,
                "chart_type": chart_type or "auto",
                "output_path": output_path,
            }
        except Exception as e:
            return {"success": False, "error": str(e)}

    def _extract_insights(self, question: str) -> str:
        """Use LLM to interpret analysis results in natural language."""
        prompt = f"""Based on this dataset profile:
{json.dumps(self.profile, indent=2, default=str)[:2000]}

Question: {question}

Provide key insights and findings in 2-3 sentences. Be specific with numbers."""

        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "You are a data analyst."},
                {"role": "user", "content": prompt},
            ],
            temperature=0.3,
        )
        return response.choices[0].message.content