from flask import Flask, request, jsonify
from flask_cors import CORS
import pandas as pd
from openai import OpenAI
import re
import plotly.io as pio
import json
import plotly.express as px
import os
from dotenv import load_dotenv
import logging

# Load environment variables from .env file
load_dotenv()

# Get the API key from the environment variable
api_key = os.getenv('OPENAI_API_KEY')
pio.renderers.default = 'json'

app = Flask(__name__)
CORS(app, resources={r"/analyze": {"origins": "https://viz.zhanglearning.com"}})

logging.basicConfig(level=logging.INFO)

@app.route('/analyze', methods=['POST'])
def analyze():
    app.logger.info("Analyze function called")
    if 'file' not in request.files:
        app.logger.error("No file uploaded")
        return jsonify({'error': '没有文件上传'}), 400
    
    file = request.files['file']
    prompt = request.form.get('prompt', '')
    
    if file.filename == '':
        app.logger.error("No file selected")
        return jsonify({'error': '没有选择文件'}), 400
    
    if file and file.filename.endswith('.csv'):
        try:
            df = pd.read_csv(file)
            app.logger.info(f"CSV file read successfully. Shape: {df.shape}")
            data_info = df.dtypes.to_dict()
            
            app.logger.info("Initializing OpenAI client")
            client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com")
            
            messages = [
                {"role": "system", "content": """
                你是数据分析、可视化和 Jupyter Notebook 开发方面的专家，专注于 Python 库，如 pandas、matplotlib、seaborn 和 numpy。
                **关键原则：**
                - 用准确的 Python 示例写出简洁的技术回复。
                - 在数据分析工作流中优先考虑可读性和可重复性。
                - 在适当的时候使用函数式编程；避免不必要的类。
                - 优先使用向量化操作而不是显式循环以获得更好的性能。
                - 使用描述性的变量名以反映它们所包含的数据。
                - 遵循 Python 代码的 PEP 8 风格指南。

                **数据分析和操作：**
                - 使用 pandas 进行数据操作和分析。
                - 在可能的情况下，优先使用方法链进行数据转换。
                - 使用 loc 和 iloc 进行明确的数据选择。
                - 利用 groupby 操作进行高效的数据聚合。

                **可视化：**
                - 使用 matplotlib 进行低级绘图控制和自定义。
                - 使用 seaborn 进行统计可视化和美观的默认设置。
                - 创建带有适当标签、标题和图例的信息丰富且视觉上吸引人的图。
                - 使用适当的配色方案并考虑色盲可访问性。

                **Jupyter Notebook 最佳实践：**
                - 使用 Markdown 单元格以清晰的部分结构笔记本。
                - 使用有意义的单元格执行顺序以确保可重复性。
                - 在 Markdown 单元格中包含解释性文本以记录分析步骤。
                - 保持代码单元格专注且模块化，以便于理解和调试。
                - 使用诸如 %matplotlib inline 之类的魔术命令进行内联绘图。

                **错误处理和数据验证：**
                - 在分析开始时实施数据质量检查。
                - 适当地处理缺失数据（插补、删除或标记）。
                - 对于容易出错的操作使用 try-except 块，尤其是在读取外部数据时。
                - 验证数据类型和范围以确保数据完整性。

                **依赖项：**
                - pandas
                - numpy
                - matplotlib
                - seaborn
                - jupyter
                - scikit-learn（用于机器学习任务）

                **关键约定：**
                1. 以数据探索和汇总统计开始分析。
                2. 创建可重用的绘图函数以实现一致的可视化。
                3. heatmap已经从plotly.express迁移到plotly.graph_objects中
                参考 pandas、matplotlib 和 Jupyter 的官方文档以获取最佳实践和最新的 API。
                """},
                {"role": "user", "content": f"根据接收的数据字段和类型：{data_info},{prompt},注意：我已经安装好了所有依赖;请确保在代码中使用 'df' 变量来引用数据框;直接给我最终代码即可，不要写注释;请确保在代码中使用 'df' 变量来引用数据框。使用 plotly.express 进行可视化，并使用 'px' 作为别名。不要使用 df_filtered 变量，所有的过滤操作都应该直接在 df 上进行。"}
            ]
            
            app.logger.info("Sending request to OpenAI")
            response = client.chat.completions.create(
                model="deepseek-coder",
                messages=messages,
                stream=False
            )
            
            app.logger.info("Received response from OpenAI")
            response_code = response.choices[0].message.content
            code_blocks = re.findall(r'```(.*?)```', response_code, re.DOTALL)
            cleaned_code_blocks = [code.replace("python\n", "") for code in code_blocks]
            
            results = []
            for code in cleaned_code_blocks:
                try:
                    app.logger.info(f"Executing code block: {code[:100]}...")  # Log first 100 chars of code
                    local_vars = {'df': df, 'px': px}
                    exec(code, globals(), local_vars)
                    for var_name, var_value in local_vars.items():
                        if var_name.startswith('fig'):
                            results.append(pio.to_json(var_value))
                    app.logger.info(f"Code block executed successfully. Results: {len(results)}")
                except Exception as e:
                    app.logger.error(f"Error executing code: {str(e)}")
            
            app.logger.info(f"Returning {len(results)} plots")
            return jsonify({'plots': results})
        except Exception as e:
            app.logger.error(f"Error in analyze function: {str(e)}")
            return jsonify({'error': str(e)}), 500
    
    app.logger.error("Unsupported file type")
    return jsonify({'error': '不支持的文件类型'}), 400

if __name__ == '__main__':
    app.run(debug=True)