from flask import Flask, request, jsonify
import mlflow
import pandas as pd
from pycaret.classification import *
from pycaret.regression import *
from pycaret.clustering import *
import json
import os
import groq

app = Flask(__name__)

# Initialize GROQ client
groq_client = groq.Client(api_key=os.getenv("GROQ_API_KEY"))

# MLflow Configuration
mlflow.set_tracking_uri("http://127.0.0.1:5000")
mlflow.set_experiment("Neural-Vision Enhanced")

@app.route('/analyze', methods=['POST'])
def analyze():
    try:
        data = request.json
        prompt = data.get('prompt')
        context = json.loads(data.get('context'))
        metrics = data.get('metrics', {})
        
        # Create GROQ prompt with context
        system_prompt = f"""
        You are a data science assistant analyzing model metrics and data.
        Context: {json.dumps(context, indent=2)}
        Metrics: {json.dumps(metrics, indent=2)}
        """
        
        # Get GROQ response
        response = groq_client.chat.completions.create(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": prompt}
            ],
            model="mixtral-8x7b-32768",
            temperature=0.7,
            max_tokens=1024
        )
        
        return jsonify({
            "analysis": response.choices[0].message.content
        })
    
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/train', methods=['POST'])
def train_model():
    try:
        data = request.json
        df = pd.DataFrame(data['data'])
        problem_type = data['problem_type']
        target = data.get('target')
        
        if problem_type == "Classification":
            setup(df, target=target, session_id=42)
        elif problem_type == "Regression":
            setup(df, target=target, session_id=42)
        else:
            setup(df, session_id=42)
        
        best_model = compare_models()
        metrics = pull().to_dict()
        
        # Log to MLflow
        with mlflow.start_run():
            mlflow.log_metrics(metrics)
            mlflow.sklearn.log_model(best_model, "model")
        
        return jsonify({
            "model": str(best_model),
            "metrics": metrics
        })
    
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(host='127.0.0.1', port=5001)