app/services/data.py

from typing import Optional, Dict, Any
from bson import ObjectId
from app.schemas.auth import CreateUser
from fastapi.exceptions import HTTPException
from app.models.mongodb.metadata import Metadata
from motor.motor_asyncio import AsyncIOMotorCollection
from app.database.mongodb import get_async_collection
from fastapi import UploadFile
import pandas as pd
from io import BytesIO, StringIO
from app.ml.preprocessing import preprocess
from bson import ObjectId
from app.utils.utils import generate_operation_id
from app.utils.s3_utils import async_upload_file_to_s3, s3_key_for_upload, get_s3_url
from app.services.models import ModelService
from typing import List
from app.core.constant import SAFE_GLOBALS
from scipy.stats import chi2_contingency, pointbiserialr, f_oneway
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import base64
import aiohttp

class DataService:
    """Operations for data service including metadata handling"""

    def __init__(self):
        self.metadata_collection_name = "metadata"
        self.preprocessed_data_collection_name = "preprocessed_data"
        self.model_service = ModelService()

    @property
    def metadata(self) -> AsyncIOMotorCollection:
        """Lazy-load async collection"""
        return get_async_collection(self.metadata_collection_name)
    
    @property
    def preprocessed_data(self) -> AsyncIOMotorCollection:
        """Lazy-load async preprocessed_data collection"""
        return get_async_collection(self.preprocessed_data_collection_name)
    
    async def read_file(self, url: dict):
        if url.lower().endswith(".csv"):
            df = pd.read_csv(url)
            df.columns = [col.replace('.', '_') for col in df.columns]
            df = df.fillna('null')
            return {"data": df.to_dict(orient="records")}
        elif url.lower().endswith(".json"):
            df = pd.read_json(url)
            df = df.fillna('null')
            return {"data": df.to_dict(orient="records")[0]}
        else:
            raise ValueError("URL must point to a CSV or JSON file")
    
    async def save_metadata(self, metadata: Dict[str, Any]):
        """
        Save metadata info (columns + target_columns) into MongoDB.
        Example:
        {
            "columns": {"name": "string", "age": "int"},
            "target_columns": ["price", "rating"]
        }
        """
        document = {
            "columns": metadata.get("columns", {}),
            "target_columns": metadata.get("target_columns", []),
            "columns_to_use_in_reasoning": metadata.get("columns_to_use_in_reasoning", []),
            "ground_truth_column_for_reasoning": metadata.get("ground_truth_column_for_reasoning", []),
            "language": metadata.get("language", "en"),
            "file_url": metadata.get("file_url", "")
        }
        result = await self.metadata.insert_one(Metadata(**document).model_dump())
        return result.inserted_id
    
    async def update_metadata(self, metadata: Dict[str, Any]):
        """
        Update metadata info by metadata_id.
        """
        metadata_id = metadata.get("metadata_id")
        document = {
            "columns": metadata.get("columns", {}),
        }
        result = await self.metadata.update_one(
            {"_id": ObjectId(metadata_id)},
            {"$set": document}
        )
        return result.modified_count > 0
    
    async def get_metadata(self, metadata_id: str):
        """
        Fetch metadata document by its ID.
        """
        result = await self.metadata.find_one({"_id": ObjectId(metadata_id)})
        result["_id"] = str(result["_id"])
        return result if result else None
    
    async def get_feature_types(self, file: UploadFile):
        contents = await file.read()
        if file.filename.endswith(".csv"):
            df = pd.read_csv(BytesIO(contents))
        elif file.filename.endswith(".json"):
            df = pd.read_json(BytesIO(contents))
        else:
            raise ValueError("File must be a CSV or JSON file")
        
        feature_types = {}
        for col in df.columns:
            if pd.api.types.is_numeric_dtype(df[col]):
                feature_types[col] = "numerical"
            elif pd.api.types.is_categorical_dtype(df[col]) or df[col].nunique() / len(df) < 0.05:
                feature_types[col] = "categorical"
            else:
                feature_types[col] = "text"
        
        file_id = generate_operation_id()
        csv_buffer = StringIO()
        df.to_csv(csv_buffer, index=False)
        s3_key = s3_key_for_upload(f"{file_id}-{file.filename}")
        
        await async_upload_file_to_s3(csv_buffer.getvalue().encode("utf-8"), s3_key)
        print(f"File uploaded to S3: {s3_key}")
        return {
            "feature_types": feature_types,
            "file_url": get_s3_url(s3_key)
        } 
    
    async def preprocessing(self, metadata_id: str = None, custom_code: UploadFile = None):
        metadata = await self.get_metadata(metadata_id) if metadata_id else None
        file_url = metadata.get("file_url", "")
        async with aiohttp.ClientSession() as session:
            async with session.get(file_url) as resp:
                if resp.status != 200:
                    raise ValueError(f"Failed to download file. HTTP {resp.status}")
                content = await resp.read()

        lower_url = file_url.lower()
        if lower_url.endswith(".csv"):
            df = pd.read_csv(BytesIO(content))
        elif lower_url.endswith(".json"):
            df = pd.read_json(BytesIO(content))
        else:
            raise ValueError("Only CSV or JSON files are supported")
        
        if custom_code:
            print("Custom preprocessing detected. Running uploaded script...")
            if not custom_code.filename.endswith(".py"):
                raise HTTPException(400, "Custom preprocessing file must be a .py script")
            custom_script = (await custom_code.read()).decode("utf-8")
            df_cleaned = self._execute_custom_preprocessor(custom_script, df, metadata)
        else:
            df_cleaned = preprocess(df, metadata)
        
        file_id = generate_operation_id()
        csv_buffer = StringIO()
        df_cleaned.to_csv(csv_buffer, index=False)
        s3_key = s3_key_for_upload(f"{file_id}-preprocessed.csv")
        
        await async_upload_file_to_s3(csv_buffer.getvalue().encode("utf-8"), s3_key)
        print(f"File uploaded to S3: {s3_key}")

        result = await self.preprocessed_data.insert_one({
            "metadata_id": metadata_id,
            "s3_path": s3_key,
            "custom_preprocessing": bool(custom_code)
        })

        return {
            "metadata_id": metadata_id,
            "preprocess_id": str(result.inserted_id),
            "s3_url": get_s3_url(s3_key),
            "custom": bool(custom_code)
        }
        
    
    async def get_preprocessing_data(self, preprocess_id: str = None):
        result = await self.preprocessed_data.find_one({"_id": ObjectId(preprocess_id)})
        print(f"[DATA SERVICE] Fetched preprocessed data: {result}")
        if result:
            result["_id"] = str(result["_id"])
            print(f"[DATA SERVICE] Fetched preprocessed data: {result}")
            return result
        return None
    
    async def get_model_metadata(self, user_id: str, model_name: str, version: int):
        """
        Fetch metadata document by its ID.
        """
        model = await self.model_service.get_by_version(user_id, model_name, version)
        metadata_id = model.get("metadata_id") if model else None
        if not metadata_id:
            return None
        result = await self.get_metadata(metadata_id)
        return result if result else None
    
    async def compare_metadata(self, metadata_id: str, columns: Dict[str, str], target_columns: List[str], columns_to_use_in_reasoning: Optional[List[str]] = None, ground_truth_column_for_reasoning: Optional[List[str]] = None):
        """
        Compare two metadata documents and return differences.
        """
        metadata_stored = await self.get_metadata(metadata_id)
        metadata_incoming = {
            "columns": columns,
            "target_columns": target_columns,
            "columns_to_use_in_reasoning": columns_to_use_in_reasoning,
            "ground_truth_column_for_reasoning": ground_truth_column_for_reasoning
        }
        
        if not metadata_stored:
            raise ValueError("First metadata not found")
        
        differences = {
            "columns_added": [],
            "columns_removed": [],
            "target_columns_added": [],
            "target_columns_removed": [],
            "columns_to_use_in_reasoning_added": [],
            "columns_to_use_in_reasoning_removed": [],
            "ground_truth_column_for_reasoning_added": [],
            "ground_truth_column_for_reasoning_removed": []
        }
        
        # Compare columns
        cols_stored = set(metadata_stored["columns"].keys())
        cols_incoming = set(metadata_incoming["columns"].keys())
        
        # Compare target columns
        targets_stored = set(metadata_stored["target_columns"])
        targets_incoming = set(metadata_incoming["target_columns"])
        
        # compare reasoning columns
        reasoning_stored = set(metadata_stored.get("columns_to_use_in_reasoning", []))
        reasoning_incoming = set(metadata_incoming.get("columns_to_use_in_reasoning", []))

        # Compare ground truth columns for reasoning
        ground_truth_stored = set(metadata_stored.get("ground_truth_column_for_reasoning", []))
        ground_truth_incoming = set(metadata_incoming.get("ground_truth_column_for_reasoning", []))

        feature_columns_stored = cols_stored - targets_stored - reasoning_stored - ground_truth_stored
        feature_column_incoming = cols_incoming - targets_incoming - reasoning_incoming - ground_truth_incoming
        
        differences["feature_columns_added"] = list(feature_column_incoming - feature_columns_stored)
        differences["feature_columns_removed"] = list(feature_columns_stored - feature_column_incoming)
        
        differences["columns_to_use_in_reasoning_added"] = list(reasoning_incoming - reasoning_stored)
        differences["columns_to_use_in_reasoning_removed"] = list(reasoning_stored - reasoning_incoming)

        differences["ground_truth_column_for_reasoning_added"] = list(ground_truth_incoming - ground_truth_stored)
        differences["ground_truth_column_for_reasoning_removed"] = list(ground_truth_stored - ground_truth_incoming)

        differences["target_columns_added"] = list(targets_incoming - targets_stored)
        differences["target_columns_removed"] = list(targets_stored - targets_incoming)
        
        return differences
    
    def _execute_custom_preprocessor(self, code: str, df, metadata):
        local_env = {}
        exec(code, SAFE_GLOBALS, local_env)

        if "custom_preprocess" not in local_env:
            raise ValueError("Uploaded script must define function `custom_preprocess(df, metadata)`.")

        custom_fn = local_env["custom_preprocess"]
        return custom_fn(df.copy(), metadata)
    

    def cramers_v(self, x, y):
        confusion_matrix = pd.crosstab(x, y)
        n = confusion_matrix.values.sum()
        if confusion_matrix.empty:
            return 0.0
        if n == 0:
            return 0.0
        if confusion_matrix.shape[0] < 2 or confusion_matrix.shape[1] < 2:
            return 0.0
        
        chi2 = chi2_contingency(confusion_matrix, correction=False)[0]
        r, k = confusion_matrix.shape
        return np.sqrt(chi2 / (n * (min(r - 1, k - 1)))) if n > 0 and min(r, k) > 1 else 0.0

    def correlation_ratio(self, categories, values):
        fcat = np.array(categories)
        values = np.array(values, dtype=float)
        categories_unique = np.unique(fcat[~pd.isnull(fcat)])
        y_avg_total = np.nanmean(values)
        ss_between, ss_within = 0, 0
        for cat in categories_unique:
            mask = (fcat == cat)
            group = values[mask]
            n = len(group)
            if n > 0:
                y_avg_group = np.nanmean(group)
                ss_between += n * (y_avg_group - y_avg_total) ** 2
                ss_within += np.nansum((group - y_avg_group) ** 2)
        return np.sqrt(ss_between / (ss_between + ss_within)) if (ss_between + ss_within) > 0 else 0.0

    async def get_analysis_result(self, metadata_id) -> Dict:
        metadata = await self.get_metadata(metadata_id)
        columns = metadata.get("columns", {})
        target_columns = metadata.get("target_columns", [])
        file_url = metadata.get("file_url", "")
        async with aiohttp.ClientSession() as session:
            async with session.get(file_url) as resp:
                if resp.status != 200:
                    raise ValueError(f"Failed to download file. HTTP {resp.status}")
                content = await resp.read()

        lower_url = file_url.lower()
        if lower_url.endswith(".csv"):
            df = pd.read_csv(BytesIO(content))
        elif lower_url.endswith(".json"):
            df = pd.read_json(BytesIO(content))
        else:
            raise ValueError("Only CSV or JSON files are supported")

        for col, dtype in columns.items():
            if dtype == 'numerical':
                df[col] = pd.to_numeric(df[col], errors='coerce')
            elif dtype == 'categorical':
                df[col] = df[col].astype('category')
            elif dtype == 'text':
                df[col] = df[col].astype(str)

        features = [c for c in df.columns if c not in target_columns]
        numeric_feats = [c for c in features if columns[c] == 'numerical']
        categorical_feats = [c for c in features if columns[c] == 'categorical']
        text_feats = [c for c in features if columns[c] == 'text']

        result = {
            "feature_target_correlation": {},
            "feature_correlation_matrix": {},
            "vif": {},
            "heatmap_base64": None
        }

        for target in target_columns:
            result["feature_target_correlation"][target] = {}
            for feat in numeric_feats:
                if columns[target] == 'numerical':
                    corr_val = df[[feat, target]].corr().iloc[0, 1]
                else:
                    corr_val = self.correlation_ratio(df[target], df[feat])
                result["feature_target_correlation"][target][feat] = float(corr_val)

            for feat in categorical_feats:
                if columns[target] == 'numerical':
                    corr_val = self.correlation_ratio(df[feat], df[target])
                else:
                    corr_val = self.cramers_v(df[target], df[feat])
                result["feature_target_correlation"][target][feat] = float(corr_val)

        # Custom correlation matrix using Pearson, correlation_ratio, and cramers_v
        all_feats = numeric_feats + categorical_feats
        corr_matrix = pd.DataFrame(index=all_feats, columns=all_feats, dtype=float)

        for i in all_feats:
            for j in all_feats:
                if i == j:
                    corr_matrix.loc[i, j] = 1.0
                elif pd.isna(corr_matrix.loc[i, j]):
                    type_i = columns[i]
                    type_j = columns[j]
                    if type_i == 'numerical' and type_j == 'numerical':
                        val = df[[i, j]].corr().iloc[0, 1]
                    elif type_i == 'numerical' and type_j == 'categorical':
                        val = self.correlation_ratio(df[j], df[i])
                    elif type_i == 'categorical' and type_j == 'numerical':
                        val = self.correlation_ratio(df[i], df[j])
                    elif type_i == 'categorical' and type_j == 'categorical':
                        val = self.cramers_v(df[i], df[j])
                    else:
                        val = np.nan
                    corr_matrix.loc[i, j] = val
                    corr_matrix.loc[j, i] = val

        result["feature_correlation_matrix"] = corr_matrix.fillna(0).round(3).to_dict()
        plt.figure(figsize=(12, 10))
        sns.heatmap(corr_matrix, annot=False, cmap="coolwarm")
        plt.title("Combined Feature Correlation Heatmap")
        buf = BytesIO()
        plt.tight_layout()
        plt.savefig(buf, format="png")
        plt.savefig("heatmap.png", format="png", dpi=300, bbox_inches="tight")
        plt.close()
        buf.seek(0)
        # Upload to S3
        heatmap_filename = f"{generate_operation_id()}_heatmap.png"
        heatmap_s3_key = s3_key_for_upload(heatmap_filename)
        
        try:
            await async_upload_file_to_s3(buf.getvalue(), heatmap_s3_key)
            heatmap_url = get_s3_url(heatmap_s3_key)
            print(f"[ANALYSIS] Heatmap uploaded to S3: {heatmap_url}")
            result["heatmap_url"] = heatmap_url
        except Exception as e:
            print(f"[ANALYSIS] Failed to upload heatmap to S3: {e}")
            result["heatmap_url"] = None
        
        return result