tianxuan/analysis/views/clustering.py

"""Cluster views: overview, detail, globe flows, and the clustering pipeline."""
import asyncio
import json
import traceback
import logging

from django.shortcuts import render, get_object_or_404
from django.http import Http404

from analysis.models import AnalysisRun, ClusterResult, EntityProfile
from analysis import nl_describe
from analysis.constants import RANDOM_SEED, UMAP_BATCH_SIZE, UMAP_TRAIN_SAMPLE
from analysis.profile_util import profile

from .helpers import _extract_lat, _extract_lon

logger = logging.getLogger(__name__)


def cluster_overview(request, display_id):
    """Overview of all clusters for a run, with UMAP scatter data and geo scatter data."""
    run = get_object_or_404(AnalysisRun, display_id=display_id)
    clusters = run.clusters.exclude(cluster_label=-1).order_by('-size').prefetch_related('features')
    noise = run.clusters.filter(cluster_label=-1).first()

    # Build UMAP scatter data for Canvas / Three.js 3D
    entities = run.entities.exclude(embedding_x=None).exclude(embedding_y=None).only(
        'embedding_x', 'embedding_y', 'embedding_z', 'cluster_label', 'entity_value'
    )
    has_z = any(e.embedding_z is not None for e in entities)
    scatter_data = [
        {'x': e.embedding_x, 'y': e.embedding_y, 'z': e.embedding_z if has_z else 0,
         'label': e.cluster_label if e.cluster_label is not None else -1, 'entity': e.entity_value}
        for e in entities
    ]

    # Noise cluster: features + entity detail
    noise_features = []
    noise_summary = ''
    noise_entity_count = 0
    if noise:
        noise_features_qs = noise.features.all().order_by('-distinguishing_score')[:15]
        noise_features = [
            {'feature_name': f.feature_name, 'score': f.distinguishing_score, 'mean': f.mean,
             'std': f.std, 'p25': f.p25, 'p75': f.p75}
            for f in noise_features_qs
        ]
        if noise_features:
            noise_summary = nl_describe.describe_cluster(noise_features)
        noise_entity_count = noise.size

    # LLM workflow timeline for auto runs
    llm_timeline_json = ''
    if run.run_type == 'auto':
        llm_thinking = run.llm_thinking or ''
        tool_calls = run.tool_calls_json or []
        if tool_calls or llm_thinking:
            llm_timeline_json = json.dumps({
                'thinking': llm_thinking,
                'tool_calls': tool_calls,
            })

    # Build geo scatter data from feature_json (lat/lon)
    geo_data = []
    geo_skipped = 0
    for e in run.entities.all().only('feature_json', 'cluster_label', 'entity_value'):
        fj = e.feature_json
        if not fj:
            continue
        lat = _extract_lat(fj.get('avg_latitude'))
        lon = _extract_lon(fj.get('avg_longitude'))
        if lat is not None and lon is not None:
            geo_data.append({
                'x': lon,
                'y': lat,
                'label': e.cluster_label if e.cluster_label is not None else -1,
                'entity': e.entity_value,
            })
        else:
            geo_skipped += 1

    # Cluster sizes for bar chart
    cluster_sizes = [{'label': c.cluster_label, 'size': c.size, 'silhouette': c.silhouette_score}
                     for c in clusters]

    # Top features per cluster from ClusterFeature + NL summaries
    top_features = {}
    for c in clusters:
        features_qs = c.features.all().order_by('-distinguishing_score')[:10]
        feats = [
            {'feature_name': f.feature_name, 'score': f.distinguishing_score, 'mean': f.mean,
             'std': f.std}
            for f in features_qs
        ]
        top_features[str(c.cluster_label)] = feats
        # Attach features + NL summary directly to cluster object for template
        c._top_features = feats
        c.nl_summary = nl_describe.describe_cluster(feats) if feats else ''

    return render(request, 'analysis/cluster_overview.html', {
        'run': run,
        'clusters': clusters,
        'noise': noise,
        'scatter_data_json': json.dumps(scatter_data),
        'has_z': has_z,
        'geo_data_json': json.dumps(geo_data),
        'geo_skipped': geo_skipped,
        'geo_count': len(geo_data),
        'cluster_sizes_json': json.dumps(cluster_sizes),
        'top_features_json': json.dumps(top_features),
        'noise_features': noise_features,
        'noise_summary': noise_summary,
        'noise_entity_count': noise_entity_count,
        'llm_timeline_json': llm_timeline_json,
        # Globe flow data for the sidebar
        'globe_flows_json': json.dumps(_get_globe_flows(run)),
        'selected_ds_id': f'upload_{run.display_id}',
    })


def _get_globe_flows(run, max_rows=500):
    """Extract flow data for 3D globe visualization from a run."""
    from analysis.data_loader import load_csv_directory, load_from_db
    try:
        lf = None
        if run.sqlite_table:
            lf = load_from_db(run.sqlite_table)
        if lf is None and run.csv_glob:
            lf, _, _, _, _ = load_csv_directory(run.csv_glob)
        if lf is None:
            return []
        df = lf.head(max_rows).collect()
        from analysis.type_classifier import TLS_HEX_MAP
        from analysis.geoip import lookup as geo_lookup
        flows = []
        src_ip_col = next((c for c in df.columns if 'src_ip' in c.lower() or ':ips' in c.lower()), None)
        dst_ip_col = next((c for c in df.columns if 'dst_ip' in c.lower() or ':ipd' in c.lower()), None)
        tls_col = next((c for c in df.columns if c.lower() in ('0ver', 'tls_version', '_tlsver', 'version')), None)
        bytes_col = next((c for c in df.columns if c.lower() in ('8ack', '8pak', '8byt', 'bytes_sent', 'bytes')), None)
        for row in df.iter_rows(named=True):
            sip = str(row.get(src_ip_col, '')) if src_ip_col else ''
            dip = str(row.get(dst_ip_col, '')) if dst_ip_col else ''
            if not sip or not dip:
                continue
            sg = geo_lookup(sip)
            dg = geo_lookup(dip)
            if not sg or not dg:
                continue
            tls_val = str(row.get(tls_col, '') or '') if tls_col else ''
            tls_ver = TLS_HEX_MAP.get(tls_val.replace(' ', ''), tls_val)
            flows.append({
                'slat': sg['lat'], 'slon': sg['lon'],
                'dlat': dg['lat'], 'dlon': dg['lon'],
                'tls': tls_ver,
                'bytes': float(row.get(bytes_col, 0) or 0) if bytes_col else 0,
            })
        return flows
    except Exception:
        logger.error("unknown failed: {}".format(traceback.format_exc()))
        return []


def cluster_detail(request, display_id, cluster_label):
    """Detailed view of a single cluster with raw data rows, SVD features, and NL summary."""
    import json
    import polars as pl
    run = get_object_or_404(AnalysisRun, display_id=display_id)
    try:
        cluster_label = int(cluster_label)
    except (ValueError, TypeError):
        raise Http404("Invalid cluster label")
    cluster = get_object_or_404(ClusterResult, run=run, cluster_label=cluster_label)
    features = cluster.features.all().order_by('-distinguishing_score')[:50]

    # Load raw data rows for this cluster from EP row indices
    row_data = []
    try:
        from analysis.data_loader import load_csv_directory, load_from_db
        lf = None
        if run.sqlite_table:
            lf = load_from_db(run.sqlite_table)
        if lf is None and run.csv_glob:
            lf, _, _, _, _ = load_csv_directory(run.csv_glob)
        if lf is not None:
            # Get row indices assigned to this cluster from EntityProfile
            eps = cluster.entities.all().order_by('entity_value')[:100]
            indices = []
            for ep in eps:
                try:
                    idx = int(ep.entity_value.replace('row_', ''))
                    indices.append(idx)
                except (ValueError, AttributeError):
                    continue
            if indices:
                df = lf.collect(streaming=True)
                df_filtered = df[indices]
                # Convert to dict list for template (limit columns for display)
                cols = df_filtered.columns[:20]  # show first 20 columns
                for row_idx, row in enumerate(df_filtered.select(cols).iter_rows(named=True)):
                    entry = {'index': indices[row_idx] if row_idx < len(indices) else row_idx}
                    for k, v in row.items():
                        entry[k] = v
                    row_data.append(entry)
    except Exception as exc:
        logger.warning(f'Failed to load raw data for cluster detail: {exc}')

    nl_summary = ''
    feats_list = [{'feature_name': f.feature_name, 'score': f.distinguishing_score,
                   'mean': f.mean, 'std': f.std} for f in features[:10]]
    if feats_list:
        nl_summary = nl_describe.describe_cluster(feats_list)

    return render(request, 'analysis/cluster_detail.html', {
        'run': run,
        'cluster': cluster,
        'features': features,
        'row_data': row_data,
        'nl_summary': nl_summary,
    })


@profile
def _run_clustering_pipeline(run, store, ds_id, feature_columns, algorithm,
                              min_cluster_size, run_umap=True, head=None):
    """Clustering pipeline on raw rows. Delegates to service layer."""
    from analysis.services.clustering import run_clustering_pipeline
    run_clustering_pipeline(run, store, ds_id, feature_columns, algorithm,
                            min_cluster_size, run_umap=run_umap, head=head)