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