"""extract_features handler — per-cluster feature ranking + ORM persistence."""
import sys
import traceback

import numpy as np
import polars as pl

from ..session_store import SessionStore, _generate_id
from ..profile_util import profile
from analysis.constants import RANDOM_SEED, UMAP_BATCH_SIZE, UMAP_TRAIN_SAMPLE
import logging
logger = logging.getLogger(__name__)


async def _handle_extract_features(
    dataset_id: str,
    cluster_result_id: str,
    top_k: int = 10,
    method: str = 'zscore',
    save_to_db: bool = True,
    run_id: int | None = None,
) -> dict:
    """Compute per-cluster distinguishing features and optionally persist."""
    store = SessionStore()
    dataset_entry = store.get_dataset(dataset_id)
    if dataset_entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}

    cluster_entry = store.get_cluster_result(cluster_result_id)
    if cluster_entry is None:
        return {'error': f'Cluster result not found: {cluster_result_id}',
                'truncated': False}

    lf: pl.LazyFrame = dataset_entry['lazyframe']
    labels = cluster_entry['labels']

    # D5: If clustering downsampled the data, filter matching rows
    _ds_idx = cluster_entry.get('params', {}).get('_ds_idx')
    if _ds_idx is not None:
        df = (lf.with_row_index('__ds_idx')
              .filter(pl.col('__ds_idx').is_in(_ds_idx))
              .drop('__ds_idx')
              .collect(streaming=True))
    else:
        df = lf.collect(streaming=True)

    # Truncate labels if df was filtered but labels weren't
    if len(labels) != len(df):
        labels = labels[:len(df)]

    # Only use numeric columns for feature extraction
    numeric_cols = [c for c in df.columns
                    if df[c].dtype in (pl.Float32, pl.Float64,
                                        pl.Int32, pl.Int64,
                                        pl.UInt32, pl.UInt64)]
    if len(labels) != len(df):
        labels = labels[:len(df)]

    # Global statistics per column
    global_stats = {}
    for col in numeric_cols:
        s = df[col].drop_nulls()
        if len(s) > 0:
            global_stats[col] = {
                'mean': float(s.mean()),
                'std': float(s.std()) if s.std() is not None else 1.0,
            }

    unique_labels = sorted(set(labels))
    all_features = []

    for label in unique_labels:
        mask = np.array(labels) == label
        if mask.sum() < 2:
            continue

        cluster_df = df.filter(pl.Series('__mask__', mask))
        cluster_features = []

        for col in numeric_cols:
            c_series = cluster_df[col].drop_nulls()
            gs = global_stats.get(col)
            if gs is None or gs['std'] == 0 or len(c_series) == 0:
                continue

            c_mean = float(c_series.mean())
            c_std = float(c_series.std()) if c_series.std() is not None else 0.0

            if method == 'zscore':
                score = (c_mean - gs['mean']) / max(gs['std'], 1e-10)
            else:
                score = abs(c_mean - gs['mean']) / max(c_std + gs['std'], 1e-10)

            cluster_features.append({
                'feature_name': col,
                'cluster_label': int(label),
                'mean': round(c_mean, 4),
                'std': round(c_std, 4),
                'global_mean': round(gs['mean'], 4),
                'global_std': round(gs['std'], 4),
                'distinguishing_score': round(float(score), 4),
                'distinguishing_method': method,
            })

        cluster_features.sort(key=lambda x: abs(x['distinguishing_score']), reverse=True)
        all_features.extend(cluster_features[:top_k])

    # ── Compute UMAP-2D embedding ──────────────────────────────────
    if len(numeric_cols) >= 2 and len(df) > 2:
        try:
            from sklearn.preprocessing import StandardScaler
            import umap
            MAX_UMAP_TRAIN = UMAP_TRAIN_SAMPLE
            BATCH_SIZE = UMAP_BATCH_SIZE
            data_matrix = df.select(numeric_cols).to_numpy()
            if len(df) > MAX_UMAP_TRAIN:
                idx_sample = np.random.RandomState(RANDOM_SEED).choice(
                    len(df), size=MAX_UMAP_TRAIN, replace=False)
                mat_sample = data_matrix[idx_sample]
                scaler = StandardScaler().fit(mat_sample)
                mat_sample_scaled = np.nan_to_num(scaler.transform(mat_sample),
                                                  nan=0.0, posinf=0.0, neginf=0.0)
                reducer = umap.UMAP(n_components=2, random_state=RANDOM_SEED,
                                    n_neighbors=15, min_dist=0.1,
                                    metric='euclidean')
                reducer.fit(mat_sample_scaled)
                coords_list = []
                for start in range(0, len(df), BATCH_SIZE):
                    end = min(start + BATCH_SIZE, len(df))
                    mat_batch = scaler.transform(data_matrix[start:end])
                    mat_batch = np.nan_to_num(mat_batch, nan=0.0, posinf=0.0, neginf=0.0)
                    coords_list.append(reducer.transform(mat_batch))
                coords = np.vstack(coords_list)
            else:
                data_scaled = StandardScaler().fit_transform(data_matrix)
                data_scaled = np.nan_to_num(data_scaled, nan=0.0, posinf=0.0, neginf=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(data_scaled)
            non_numeric = [c for c in df.columns if c not in numeric_cols and not c.startswith('_')]
            entity_col_umap = non_numeric[0] if non_numeric else df.columns[0]
            labels_list = cluster_entry.get('labels', [])
            from analysis.models import EntityProfile as EP
            from analysis.models import ClusterResult as CR
            from analysis.models import AnalysisRun as AR
            from asgiref.sync import sync_to_async
            csv_glob = dataset_entry.get('metadata', {}).get('csv_glob', 'manual')
            run_obj = await sync_to_async(
                AR.objects.filter(csv_glob=csv_glob).order_by('-id').first,
                thread_sensitive=True
            )()
            if run_obj:
                run_id_for_ep = run_obj.id
                profiles = []
                cr_cache = {}
                for i, row in enumerate(df.iter_rows(named=True)):
                    if i >= len(coords): break
                    ev = str(row.get(entity_col_umap, ''))
                    if not ev: continue
                    lbl = labels_list[i] if i < len(labels_list) else -1
                    cache_key = f'{run_id_for_ep}_{lbl}'
                    if cache_key not in cr_cache:
                        cr_cache[cache_key] = CR.objects.filter(run_id=run_id_for_ep, cluster_label=lbl).first()
                    profiles.append(EP(
                        entity_value=ev, run_id=run_id_for_ep, cluster_label=lbl,
                        cluster=cr_cache[cache_key],
                        embedding_x=float(coords[i, 0]),
                        embedding_y=float(coords[i, 1]),
                    ))
                if profiles:
                    EP.objects.bulk_create(profiles, ignore_conflicts=True, batch_size=1000)
        except Exception:
            print(f'[UMAP_ERR] {traceback.format_exc()}', file=sys.stderr)

    # Persist to Django ORM
    db_saved = False
    if save_to_db:
        try:
            from asgiref.sync import sync_to_async
            from analysis.models import AnalysisRun as AnalysisRunModel
            ce = cluster_entry
            async def _do_save(ce_inner):
                if run_id is not None:
                    run = await sync_to_async(AnalysisRunModel.objects.get)(id=run_id)
                else:
                    csv_glob_fallback = dataset_entry.get('metadata', {}).get('csv_glob', 'manual')
                    run = await sync_to_async(
                        AnalysisRunModel.objects.filter(csv_glob=csv_glob_fallback).order_by('-id').first,
                        thread_sensitive=True
                    )()
                    if run is None:
                        run = await sync_to_async(AnalysisRunModel.objects.create)(
                            csv_glob=csv_glob_fallback, status='running', total_flows=0,
                        )
                run.status = 'completed'
                run.cluster_count = ce_inner['n_clusters']
                await sync_to_async(run.save)()
                return await sync_to_async(_save_features_to_db)(cluster_result_id, all_features, run_id=run.id)
            db_saved = await _do_save(ce)
        except Exception as e:
            print(f'[DB_SAVE_ERROR] {e}\n{traceback.format_exc()}', file=sys.stderr)
            db_saved = False


    result_id = _generate_id('feat')
    store.store_feature_result(
        result_id=result_id,
        features=all_features,
        parent_cluster_result_id=cluster_result_id,
    )

    return {
        'feature_result_id': result_id,
        'features': all_features,
        'n_features': len(all_features),
        'db_saved': db_saved,
    }


@profile
def _save_features_to_db(
    cluster_result_id: str,
    features: list[dict],
    run_id: int | None = None,
) -> bool:
    """Save extracted features to the Django ORM ClusterFeature model."""
    if run_id is not None:
        from analysis.models import AnalysisRun as AnalysisRun2
        run = AnalysisRun2.objects.filter(id=run_id).first()
        if run is None:
            return False
    else:
        store = SessionStore()
        cluster_entry = store.get_cluster_result(cluster_result_id)
        if cluster_entry is None:
            return False

        parent_ds_id = cluster_entry.get('parent_dataset_id', '')
        from analysis.models import AnalysisRun
        ds_entry = store.get_dataset(parent_ds_id)
        if ds_entry is None:
            return False

        csv_glob = ds_entry.get('metadata', {}).get('csv_glob', '')
        unique_global = sorted(set(f['cluster_label'] for f in features))
        n_clusters_global = len([l for l in unique_global if l >= 0])
        try:
            run = AnalysisRun.objects.filter(csv_glob=csv_glob).order_by('-created_at').first()
            if run is None:
                run = AnalysisRun.objects.create(
                    csv_glob=csv_glob,
                    status='completed',
                    cluster_count=n_clusters_global,
                )
        except Exception:
            logger.error("unknown failed: {}".format(traceback.format_exc()))
            return False

    from analysis.models import ClusterResult as ClusterResultModel

    store = SessionStore()
    cluster_entry = store.get_cluster_result(cluster_result_id)
    label_counts = {}
    if cluster_entry:
        full_labels = cluster_entry.get('labels', [])
        for lbl in full_labels:
            label_counts[lbl] = label_counts.get(lbl, 0) + 1

    unique_labels = sorted(set(
        f['cluster_label'] for f in features
    ))
    n_clusters = len([l for l in unique_labels if l >= 0])
    for label in unique_labels:
        if label == -1:
            continue
        label_features = [f for f in features if f['cluster_label'] == label]
        if not label_features:
            continue

        cluster_db, _ = ClusterResultModel.objects.get_or_create(
            run=run,
            cluster_label=label,
            defaults={'size': label_counts.get(label, 1)},
        )
        if label_counts.get(label) is not None:
            ClusterResultModel.objects.filter(
                run=run, cluster_label=label
            ).update(size=label_counts[label])

        from analysis.models import ClusterFeature as ClusterFeatureModel
        for feat in label_features:
            try:
                ClusterFeatureModel.objects.get_or_create(
                    cluster=cluster_db,
                    feature_name=feat['feature_name'],
                    defaults={
                        'mean': feat.get('mean'),
                        'std': feat.get('std'),
                        'median': None,
                        'p25': None,
                        'p75': None,
                        'missing_rate': None,
                        'distinguishing_score': feat.get('distinguishing_score'),
                        'distinguishing_method': feat.get('distinguishing_method', 'zscore'),
                    }
                )
            except Exception:
                logger.error("unknown failed: {}".format(traceback.format_exc()))
                continue

    run.cluster_count = n_clusters
    run.save(update_fields=['cluster_count'])
    return True
