"""Anomaly detection + visualization handlers."""
import traceback
import numpy as np
from analysis.constants import MAX_DISTRIBUTION_SAMPLE, RANDOM_SEED
import polars as pl

from ..session_store import SessionStore
from ._helpers import _resolve_dataset, NUMERIC_TYPE_NAMES
import logging
logger = logging.getLogger(__name__)


async def _handle_detect_anomalies(dataset_id: str, contamination: float = 0.05) -> dict:
    """Run Isolation Forest on scored entity data."""
    from sklearn.ensemble import IsolationForest
    entry, err = _resolve_dataset(dataset_id)
    if err:
        return err

    lf: pl.LazyFrame = entry['lazyframe']
    schema = entry.get('schema', {})
    num_cols = [c for c in schema if schema[c].split('(')[0].strip() in NUMERIC_TYPE_NAMES and not c.startswith('_')]
    if not num_cols:
        return {'error': 'No numeric columns for anomaly detection', 'truncated': False}

    df = lf.select(num_cols).collect(streaming=True)
    mat = df.to_numpy()
    mat = np.nan_to_num(mat, nan=0.0)

    n = len(mat)
    chunk_size = min(MAX_SAMPLE_ROWS, n)
    labels = np.zeros(n, dtype=int)
    for start in range(0, n, chunk_size):
        end = min(start + chunk_size, n)
        iso = IsolationForest(n_estimators=100, max_samples=min(256, chunk_size),
                              contamination=contamination, random_state=RANDOM_SEED, n_jobs=-1)
        labels[start:end] = iso.fit_predict(mat[start:end])

    anomaly_count = int((labels == -1).sum())
    return {
        'status': 'completed',
        'total_entities': n,
        'anomaly_count': anomaly_count,
        'anomaly_rate': round(anomaly_count / max(n, 1), 4),
        'dataset_id': dataset_id,
    }


async def _handle_visualize_anomalies(dataset_id: str) -> dict:
    """Generate UMAP-2D embedding for anomaly visualization."""
    from sklearn.preprocessing import StandardScaler
    import umap
    entry, err = _resolve_dataset(dataset_id)
    if err:
        return err

    lf: pl.LazyFrame = entry['lazyframe']
    schema = entry.get('schema', {})
    num_cols = [c for c in schema if schema[c].split('(')[0].strip() in NUMERIC_TYPE_NAMES and not c.startswith('_')]
    if len(num_cols) < 2:
        return {'error': 'Need at least 2 numeric columns for UMAP', 'truncated': False}

    try:
        n_total = lf.select(pl.len()).collect(streaming=True).item()
    except Exception:
        logger.error("_handle_visualize_anomalies failed: {}".format(traceback.format_exc()))
        n_total = 0
    sample_lf = lf if n_total <= MAX_DISTRIBUTION_SAMPLE else lf.collect(streaming=True).sample(n=MAX_DISTRIBUTION_SAMPLE, seed=RANDOM_SEED).lazy()
    df = sample_lf.select(num_cols).collect(streaming=True)
    mat = StandardScaler().fit_transform(df.to_numpy())
    mat = np.nan_to_num(mat, nan=0.0)
    reducer = umap.UMAP(n_components=2, random_state=RANDOM_SEED,
                        n_neighbors=15, min_dist=0.1, metric='euclidean')
    coords = reducer.fit_transform(mat)

    risk_col = next((c for c in schema if 'risk' in c.lower()), None)
    score_col = next((c for c in schema if 'proxy_score' in c.lower() or 'threat' in c.lower()), None)

    non_num = [c for c in df.columns if c not in num_cols and not c.startswith('_')]
    ent_col = non_num[0] if non_num else df.columns[0]
    entities = df.select(ent_col).to_series().to_list() if ent_col in df.columns else []

    scatter = [
        {'x': float(coords[i, 0]), 'y': float(coords[i, 1]),
         'entity': str(entities[i]) if i < len(entities) else '',
         'score': round(float(mat[i].mean()), 4)}
        for i in range(len(coords))
    ]
    return {
        'status': 'completed',
        'scatter_points': len(scatter),
        'scatter_data': scatter[:200],
        'dataset_id': dataset_id,
    }
