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

# loaders/csv_loader.py
import pandas as pd
from pathlib import Path
from typing import Dict

class CSVLoader:
    def __init__(self):
        self.datasets: Dict[str, pd.DataFrame] = {}

    def load(self, file_path: str, name: str = None) -> str:
        path = Path(file_path)
        if not path.exists():
            return f"Error: File {file_path} not found"
        df = pd.read_csv(path)
        dataset_name = name or path.stem
        self.datasets[dataset_name] = df
        return f"Loaded {len(df)} rows, {len(df.columns)} columns as '{dataset_name}'"

    def get_schema(self, name: str) -> str:
        if name not in self.datasets:
            return f"Dataset '{name}' not found"
        df = self.datasets[name]
        schema = []
        for col in df.columns:
            dtype = str(df[col].dtype)
            non_null = df[col].count()
            sample = df[col].dropna().iloc[0] if not df[col].dropna().empty else "N/A"
            schema.append(f"  {col}: {dtype} ({non_null} non-null, sample: {sample})")
        return f"Table: {name}\nColumns:\n" + "\n".join(schema)

    def get_sample(self, name: str, n: int = 5) -> str:
        if name not in self.datasets:
            return f"Dataset '{name}' not found"
        return self.datasets[name].head(n).to_string()

    def get_datasets(self) -> list:
        return list(self.datasets.keys())

# loaders/sql_loader.py
import pandas as pd
from sqlalchemy import create_engine, text
from typing import Dict

class SQLLoader:
    def __init__(self, connection_string: str):
        self.engine = create_engine(connection_string)
        self.tables: Dict[str, pd.DataFrame] = {}

    def load_table(self, table_name: str) -> str:
        try:
            df = pd.read_sql_table(table_name, self.engine)
            self.tables[table_name] = df
            return f"Loaded {len(df)} rows from {table_name}"
        except Exception as e:
            return f"Error loading {table_name}: {str(e)}"

    def execute_query(self, query: str) -> pd.DataFrame:
        with self.engine.connect() as conn:
            return pd.read_sql(text(query), conn)

    def get_schema(self) -> str:
        schemas = []
        for name, df in self.tables.items():
            cols = ", ".join(f"{c} ({df[c].dtype})" for c in df.columns)
            schemas.append(f"{name}: {cols}")
        return "\n".join(schemas)

# query/nl_to_sql.py
from openai import OpenAI
from typing import Dict

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

    def translate(
        self, question: str, schema: str, tables: Dict[str, any]
    ) -> str:
        table_samples = []
        for name, df in list(tables.items())[:3]:
            table_samples.append(f"Table: {name}\n{df.head(3).to_string()}")
        samples = "\n\n".join(table_samples)
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": f"""You are a SQL expert. 
                Generate SQL queries based on natural language questions.
                
                Schema:
                {schema}

                Sample data:
                {samples}

                Rules:
                1. Use only SELECT statements (no INSERT, UPDATE, DELETE)
                2. Use table aliases for readability
                3. Include LIMIT clause
                4. Use proper JOINs when needed"""},
                {"role": "user", "content": question},
            ],
            temperature=0.0,
        )
        sql = response.choices[0].message.content.strip()
        if sql.startswith("```sql"):
            sql = sql[7:]
        if sql.endswith("```"):
            sql = sql[:-3]
        return sql.strip()

# analysis/stats.py
import pandas as pd
import numpy as np
from typing import Dict, Any

class StatsAnalyzer:
    def __init__(self):
        self.client = None

    def summarize(self, df: pd.DataFrame) -> Dict[str, Any]:
        summary = {
            "shape": df.shape,
            "columns": list(df.columns),
            "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()},
            "missing": df.isnull().sum().to_dict(),
            "numeric_stats": {},
            "categorical_stats": {},
        }
        for col in df.select_dtypes(include=[np.number]).columns:
            summary["numeric_stats"][col] = {
                "mean": float(df[col].mean()),
                "std": float(df[col].std()),
                "min": float(df[col].min()),
                "max": float(df[col].max()),
                "median": float(df[col].median()),
            }
        for col in df.select_dtypes(include=["object", "category"]).columns:
            value_counts = df[col].value_counts().head(5)
            summary["categorical_stats"][col] = {
                "unique": int(df[col].nunique()),
                "top_values": value_counts.to_dict(),
            }
        return summary

    def auto_insights(self, df: pd.DataFrame) -> str:
        insights = []
        numeric_cols = df.select_dtypes(include=[np.number]).columns
        for col in numeric_cols:
            skew = df[col].skew()
            if abs(skew) > 1:
                insights.append(f"'{col}' is heavily skewed (skew={skew:.2f})")
            q1, q3 = df[col].quantile(0.25), df[col].quantile(0.75)
            iqr = q3 - q1
            outliers = ((df[col] < q1 - 1.5 * iqr) | (df[col] > q3 + 1.5 * iqr)).sum()
            if outliers > 0:
                insights.append(f"'{col}' has {outliers} potential outliers")
        if len(numeric_cols) >= 2:
            corr = df[numeric_cols].corr()
            for i in range(len(numeric_cols)):
                for j in range(i + 1, len(numeric_cols)):
                    if abs(corr.iloc[i, j]) > 0.7:
                        insights.append(
                            f"Strong correlation ({corr.iloc[i, j]:.2f}) between "
                            f"'{numeric_cols[i]}' and '{numeric_cols[j]}'"
                        )
        return "\n".join(insights) if insights else "No significant patterns detected"

# analysis/visualizer.py
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path

class Visualizer:
    def __init__(self, output_dir: str = "./plots"):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(exist_ok=True)
        plt.style.use("seaborn-v0_8-whitegrid")

    def create_chart(
        self, df: pd.DataFrame, chart_type: str, x: str, y: str = None, title: str = ""
    ) -> str:
        fig, ax = plt.subplots(figsize=(10, 6))
        if chart_type == "bar":
            if y:
                df.plot(kind="bar", x=x, y=y, ax=ax)
            else:
                df[x].value_counts().head(10).plot(kind="bar", ax=ax)
        elif chart_type == "line":
            df.plot(kind="line", x=x, y=y, ax=ax)
        elif chart_type == "scatter":
            df.plot(kind="scatter", x=x, y=y, ax=ax)
        elif chart_type == "histogram":
            df[x].hist(ax=ax, bins=30)
        elif chart_type == "box":
            df[[x]].boxplot(ax=ax)
        ax.set_title(title or f"{chart_type.title()} of {x}")
        plt.tight_layout()
        filename = f"{chart_type}_{x}.png"
        filepath = self.output_dir / filename
        fig.savefig(filepath)
        plt.close()
        return str(filepath)

    def auto_visualize(self, df: pd.DataFrame) -> list:
        charts = []
        numeric_cols = df.select_dtypes(include=["number"]).columns
        if len(numeric_cols) >= 2:
            charts.append(self.create_chart(
                df, "scatter", numeric_cols[0], numeric_cols[1],
                f"{numeric_cols[0]} vs {numeric_cols[1]}"
            ))
        for col in numeric_cols[:3]:
            charts.append(self.create_chart(
                df, "histogram", col, title=f"Distribution of {col}"
            ))
        cat_cols = df.select_dtypes(include=["object"]).columns
        for col in cat_cols[:2]:
            charts.append(self.create_chart(
                df, "bar", col, title=f"Counts of {col}"
            ))
        return charts

# agent.py
from loaders.csv_loader import CSVLoader
from query.nl_to_sql import NLToSQL
from analysis.stats import StatsAnalyzer
from analysis.visualizer import Visualizer
from typing import Dict

class DataAnalysisAgent:
    def __init__(self, model: str = "gpt-4-turbo-preview"):
        self.csv_loader = CSVLoader()
        self.query_gen = NLToSQL(model)
        self.stats = StatsAnalyzer()
        self.viz = Visualizer()
        self.datasets: Dict[str, any] = {}

    def load_csv(self, path: str, name: str = None) -> str:
        result = self.csv_loader.load(path, name)
        dataset_name = name or path.split("/")[-1].split("\\")[-1].replace(".csv", "")
        if dataset_name in self.csv_loader.datasets:
            self.datasets[dataset_name] = self.csv_loader.datasets[dataset_name]
        return result

    def ask(self, question: str) -> Dict:
        if not self.datasets:
            return {"error": "No datasets loaded"}
        schema = "\n".join(
            self.csv_loader.get_schema(name)
            for name in self.datasets.keys()
        )
        sql = self.query_gen.translate(question, schema, self.datasets)
        try:
            import sqlite3
            conn = sqlite3.connect(":memory:")
            for name, df in self.datasets.items():
                df.to_sql(name, conn, index=False, if_exists="replace")
            import pandas as pd
            result_df = pd.read_sql_query(sql, conn)
            conn.close()
            insights = self.stats.auto_insights(result_df)
            charts = self.viz.auto_visualize(result_df)
            return {
                "question": question,
                "sql": sql,
                "result": result_df.head(20).to_dict("records"),
                "row_count": len(result_df),
                "insights": insights,
                "charts": charts,
                "summary": self.stats.summarize(result_df),
            }
        except Exception as e:
            return {
                "question": question,
                "sql": sql,
                "error": str(e),
            }

import pytest
from agent import DataAnalysisAgent

@pytest.fixture
def agent():
    return DataAnalysisAgent()

def test_load_csv(agent):
    result = agent.load_csv("test_data.csv", "test")
    assert "Loaded" in result

def test_ask(agent):
    agent.load_csv("test_data.csv", "test")
    result = agent.ask("How many rows are there?")
    assert "sql" in result
    assert "result" in result or "error" in result

def test_stats():
    import pandas as pd
    analyzer = StatsAnalyzer()
    df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
    summary = analyzer.summarize(df)
    assert "shape" in summary

# main.py
from agent import DataAnalysisAgent

def main():
    agent = DataAnalysisAgent()
    print("Data Analysis Agent. Type 'quit' to exit.\n")
    while True:
        cmd = input("Command (load/ask/quit): ").strip()
        if cmd == "quit":
            break
        elif cmd == "load":
            path = input("CSV path: ").strip()
            print(agent.load_csv(path))
        elif cmd == "ask":
            question = input("Your question: ").strip()
            result = agent.ask(question)
            if "error" in result:
                print(f"Error: {result['error']}")
            else:
                print(f"SQL: {result['sql']}")
                print(f"Rows: {result['row_count']}")
                print(f"Insights: {result['insights']}")

if __name__ == "__main__":
    main()