"""
Streamlit app để trực quan hóa embedding 3 chiều (PCA / UMAP / t-SNE) + phân cụm.

Chạy:
	streamlit run visual_data.py --server.port 8501

Yêu cầu cài đặt (một lần):
	pip install streamlit plotly scikit-learn umap-learn numpy pandas

Tính năng:
 - Load file JSON lớn chứa các object {"filepath": ..., "embedding": [...]} hoặc định dạng JSON lines.
 - Tùy chọn sample n phần tử (random) để tăng tốc.
 - Chọn thuật toán giảm chiều: PCA, UMAP, t-SNE.
 - Tham số điều chỉnh: n_neighbors, min_dist (UMAP); perplexity (t-SNE); n_components=3.
 - KMeans clustering (tuỳ chọn) để tô màu điểm; hoặc tô màu theo regex/substring trong tên file.
 - Lọc theo từ khóa trong đường dẫn.
 - Tải xuống toạ độ 3D + nhãn cluster.

File embedding quan sát được có thể không phải JSON array chuẩn; script sẽ thử:
 1. Parse như JSON array.
 2. Parse như JSON lines (mỗi dòng 1 object).
 3. Parse thủ công bằng cách tìm pattern {"filepath": ... , "embedding": [ ... ]}.

Nếu kích thước > ~1e6 bytes, dùng đọc streaming để giảm RAM.
"""

from __future__ import annotations

import json
import os
import re
import math
import random
from dataclasses import dataclass
from typing import Iterable, List, Optional, Tuple

import numpy as np
import pandas as pd
import streamlit as st
import colorsys
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score

try:
	import umap  # type: ignore
except ImportError:  # pragma: no cover
	umap = None  # handled later


@dataclass
class EmbeddingRecord:
	filepath: str
	embedding: List[float]


def _smart_json_object_stream(raw_text: str) -> Iterable[str]:
	"""Yield JSON object strings from a large raw buffer.

	Heuristic: find balanced braces starting with {"filepath": ...}.
	This is a fallback when content is not standard array / jsonlines.
	"""
	brace = 0
	buf = []
	in_obj = False
	for ch in raw_text:
		if ch == '{':
			if not in_obj:
				in_obj = True
				buf = ['{']
				brace = 1
			else:
				brace += 1
				buf.append(ch)
		elif ch == '}':
			if in_obj:
				brace -= 1
				buf.append('}')
				if brace == 0:
					yield ''.join(buf)
					in_obj = False
			else:
				# stray closing
				continue
		else:
			if in_obj:
				buf.append(ch)


def load_embeddings(
	path: str,
	sample_size: Optional[int] = None,
	sampling_seed: int = 42,
	max_objects: Optional[int] = None,
) -> List[EmbeddingRecord]:
	"""Load embeddings from a possibly large JSON / JSONL / raw file.

	Args:
		path: file path
		sample_size: random sample (after load) if provided
		sampling_seed: RNG seed
		max_objects: hard cap to stop early (for speed)
	"""
	# size = os.path.getsize(path)  # kích thước có thể dùng sau nếu muốn tối ưu đọc streaming
	# First attempt: JSON array
	records: List[EmbeddingRecord] = []
	def to_rec(obj) -> Optional[EmbeddingRecord]:
		if not isinstance(obj, dict):
			return None
		if 'embedding' in obj:
			fp = str(obj.get('filepath') or obj.get('file_path') or obj.get('path') or '')
			emb = obj['embedding']
			if isinstance(emb, list) and fp:
				return EmbeddingRecord(fp, emb)
		return None

	try:
		with open(path, 'r', encoding='utf-8') as f:
			text = f.read()
		text_stripped = text.strip()
		if text_stripped.startswith('[') and text_stripped.endswith(']'):
			arr = json.loads(text_stripped)
			for obj in arr:
				rec = to_rec(obj)
				if rec:
					records.append(rec)
					if max_objects and len(records) >= max_objects:
						break
		else:
			raise ValueError('Not a JSON array')
	except Exception:
		# Retry as JSON lines
		records = []
		try:
			with open(path, 'r', encoding='utf-8') as f:
				for line in f:
					line_strip = line.strip().rstrip(',')
					if not line_strip:
						continue
					if not line_strip.startswith('{'):
						continue
					try:
						obj = json.loads(line_strip)
						rec = to_rec(obj)
						if rec:
							records.append(rec)
							if max_objects and len(records) >= max_objects:
								break
					except json.JSONDecodeError:
						continue
			if not records:
				raise ValueError('No JSONL records')
		except Exception:
			# Fallback: heuristic extraction
			records = []
			with open(path, 'r', encoding='utf-8') as f:
				raw = f.read()
			for obj_str in _smart_json_object_stream(raw):
				if 'embedding' not in obj_str:
					continue
				# Clean possible trailing ',"
				try:
					# Attempt to fix malformed numbers like '1.2\n421875' (broken newline) by removing stray newlines inside arrays
					fixed = re.sub(r"(\d)\n(\d)", r"\1\2", obj_str)
					obj = json.loads(fixed)
				except Exception:
					continue
				rec = to_rec(obj)
				if rec:
					records.append(rec)
					if max_objects and len(records) >= max_objects:
						break
	if not records:
		raise RuntimeError("Không load được embedding nào từ file.")

	# Random sample if needed
	if sample_size and sample_size < len(records):
		random.seed(sampling_seed)
		records = random.sample(records, sample_size)
	return records


def reduce_embeddings(
	X: np.ndarray,
	method: str,
	random_state: int = 42,
	umap_neighbors: int = 15,
	umap_min_dist: float = 0.1,
	tsne_perplexity: int = 30,
	tsne_learning_rate: str | float = 'auto',
) -> Tuple[np.ndarray, dict]:
	"""Project high-dim embeddings to 3D.

	Returns (coords (n,3), meta_info)
	"""
	meta = {"method": method}
	if method == 'PCA':
		pca = PCA(n_components=3, random_state=random_state)
		coords = pca.fit_transform(X)
		meta['explained_variance_ratio'] = pca.explained_variance_ratio_.tolist()
		return coords, meta
	if method == 'UMAP':
		if umap is None:
			raise RuntimeError("Chưa cài umap-learn: pip install umap-learn")
		reducer = umap.UMAP(
			n_components=3,
			n_neighbors=umap_neighbors,
			min_dist=umap_min_dist,
			metric='cosine',
			random_state=random_state,
		)
		coords = reducer.fit_transform(X)
		meta['umap_graph_connectivity'] = float(reducer.graph_.getnnz())
		return coords, meta
	if method == 't-SNE':
		perplexity = min(tsne_perplexity, max(5, (X.shape[0] - 1) // 3))
		tsne = TSNE(
			n_components=3,
			perplexity=perplexity,
			learning_rate=tsne_learning_rate,
			init='pca',
			random_state=random_state,
			n_iter=1000,
			verbose=0,
		)
		coords = tsne.fit_transform(X)
		meta['effective_perplexity'] = perplexity
		return coords, meta
	raise ValueError(f"Unknown method {method}")


def kmeans_cluster(coords: np.ndarray, n_clusters: int, seed: int = 42) -> Tuple[np.ndarray, float]:
	if n_clusters <= 1:
		return np.zeros(coords.shape[0], dtype=int), float('nan')
	km = KMeans(n_clusters=n_clusters, n_init='auto', random_state=seed)
	labels = km.fit_predict(coords)
	score = float('nan')
	if len(set(labels)) > 1 and coords.shape[0] >= n_clusters * 5:
		try:
			score = silhouette_score(coords, labels)
		except Exception:
			pass
	return labels, score


def build_dataframe(recs: List[EmbeddingRecord]) -> pd.DataFrame:
	return pd.DataFrame({
		'filepath': [r.filepath for r in recs],
		'embedding': [r.embedding for r in recs],
	})


def load_cluster_file(
	path: str,
	expected_n: int,
	noise_label: int = -1,
) -> np.ndarray:
	"""Load cluster labels from a JSON result file.

	Supports formats:
	- {"results": [ ... ]}
	- [ ... ]
	Each item may contain one of: cluster, cluster_id, label, is_noise, filepath.
	If only is_noise exists: non-noise -> 0, noise -> noise_label.
	If filepath present, mapping is done by filepath, otherwise by index order.
	"""
	try:
		with open(path, 'r', encoding='utf-8') as f:
			content = json.load(f)
	except Exception as e:
		raise RuntimeError(f'Không đọc được file cluster: {e}')

	if isinstance(content, dict) and 'results' in content:
		items = content['results']
	elif isinstance(content, list):
		items = content
	else:
		raise RuntimeError('Định dạng file cluster không hợp lệ (cần list hoặc có key "results").')

	# Detect if filepath-based mapping
	use_filepath = any(isinstance(it, dict) and 'filepath' in it for it in items)

	labels = np.full(expected_n, noise_label, dtype=int)
	if use_filepath:
		# Build path->label
		mapping = {}
		for it in items:
			if not isinstance(it, dict):
				continue
			fp = it.get('filepath')
			if not fp:
				continue
			if 'cluster' in it:
				val = it['cluster']
			elif 'cluster_id' in it:
				val = it['cluster_id']
			elif 'label' in it:
				val = it['label']
			elif 'is_noise' in it:
				val = (0 if not it.get('is_noise') else noise_label)
			else:
				val = 0
			try:
				mapping[str(fp)] = int(val)
			except Exception:
				continue
		return labels, mapping  # second value used later to map onto df

	# Index-based mapping
	collected = []
	for it in items:
		if not isinstance(it, dict):
			# accept raw int labels
			if isinstance(it, int):
				collected.append(int(it))
			continue
		if 'cluster' in it:
			val = it['cluster']
		elif 'cluster_id' in it:
			val = it['cluster_id']
		elif 'label' in it:
			val = it['label']
		elif 'is_noise' in it:
			val = (0 if not it.get('is_noise') else noise_label)
		else:
			val = 0
		try:
			collected.append(int(val))
		except Exception:
			collected.append(noise_label)

	for i in range(min(expected_n, len(collected))):
		labels[i] = collected[i]
	return labels, None


def main():  # pragma: no cover - Streamlit entry
	st.set_page_config(page_title="Embedding 3D Viewer", layout="wide")
	st.title("🔍 Embedding 3D Viewer")
	st.caption("Trực quan hóa tương quan embedding hóa đơn (Qwen2-VL).")

	default_path = 'embeddings_factures_osteopathie_1k_qwen.json'
	path = st.text_input('Đường dẫn file embedding', value=default_path)
	col_top = st.columns(4)
	sample_size = col_top[0].number_input('Sample (0 = all)', min_value=0, value=1000, step=100)
	max_objects = col_top[1].number_input('Max objects đọc (0 = no limit)', min_value=0, value=0, step=500)
	seed = col_top[2].number_input('Seed', min_value=0, value=42, step=1)
	show_raw = col_top[3].checkbox('Hiện bảng raw', value=False)

	algo = st.selectbox('Thuật toán giảm chiều', ['UMAP', 'PCA', 't-SNE'], index=0)
	with st.expander('Tham số giảm chiều'):
		if algo == 'UMAP':
			umap_neighbors = st.slider('n_neighbors', 5, 100, 15, 1)
			umap_min_dist = st.slider('min_dist', 0.0, 1.0, 0.1, 0.01)
			tsne_perplexity = 30
		elif algo == 't-SNE':
			tsne_perplexity = st.slider('perplexity', 5, 100, 30, 1)
			umap_neighbors = 15
			umap_min_dist = 0.1
		else:
			umap_neighbors = 15
			umap_min_dist = 0.1
			tsne_perplexity = 30

	with st.expander('Phân cụm & màu sắc'):
		cluster_source = st.radio('Nguồn cluster', ['KMeans','Load file','None'], index=0, horizontal=True)
		n_clusters = st.slider('Số cluster KMeans', 2, 100, 10, 1, disabled=(cluster_source != 'KMeans'))
		cluster_file_path = st.text_input('Đường dẫn file cluster (JSON)', value='cluster/dbscan_results.json', disabled=(cluster_source != 'Load file'))
		noise_label = st.number_input('Giá trị label cho noise', value=-1, step=1, disabled=(cluster_source != 'Load file'))
		path_filter = st.text_input('Filter substring (lọc filepath, comma = OR)', value='')
		color_by_substring = st.text_input('Color theo substring (vd: osteo, facture)', value='')
		palette_name = st.selectbox('Bảng màu', ['auto','okabe-ito','tol','plotly','bold','dark24','set3','d3','hcl'], index=0, help='Chọn bảng màu dễ phân biệt hơn (okabe-ito, tol = thân thiện cho người mù màu)')
		marker_size = st.slider('Kích thước điểm', 2, 15, 5, 1)
		st.caption('Nếu file cluster chỉ chứa is_noise: noise -> noise_label, còn lại -> 0.')

	load_btn = st.button('🚀 Load & Giảm chiều', type='primary')

	if load_btn:
		if not os.path.isfile(path):
			st.error(f'File không tồn tại: {path}')
			return
		with st.spinner('Đang load embeddings...'):
			recs = load_embeddings(
				path,
				sample_size=sample_size or None,
				sampling_seed=int(seed),
				max_objects=max_objects or None,
			)
		st.success(f'Loaded {len(recs)} embeddings')

		df = build_dataframe(recs)
		dim = len(df['embedding'].iloc[0])
		st.write(f'Chiều gốc: {dim}')

		# Filter
		if path_filter.strip():
			tokens = [t.strip() for t in path_filter.split(',') if t.strip()]
			if tokens:
				mask = df['filepath'].apply(lambda p: any(tok.lower() in p.lower() for tok in tokens))
				df = df[mask].reset_index(drop=True)
				st.info(f'Filter còn {len(df)} bản ghi.')
				if df.empty:
					return

		X = np.vstack(df['embedding'].values).astype(np.float32)
		with st.spinner('Đang giảm chiều...'):
			coords, meta = reduce_embeddings(
				X,
				algo,
				random_state=int(seed),
				umap_neighbors=umap_neighbors,
				umap_min_dist=umap_min_dist,
				tsne_perplexity=tsne_perplexity,
			)
		df[['x', 'y', 'z']] = coords
		st.write('Meta:', meta)

		# Clustering (three modes)
		if cluster_source == 'KMeans':
			with st.spinner('KMeans clustering...'):
				labels, sil = kmeans_cluster(coords, n_clusters, seed)
			df['cluster'] = labels
			if not math.isnan(sil):
				st.write(f'Silhouette: {sil:.4f}')
			else:
				st.write('Silhouette: N/A')
		elif cluster_source == 'Load file':
			if not os.path.isfile(cluster_file_path):
				st.error(f'Không tìm thấy file cluster: {cluster_file_path}')
				return
			with st.spinner('Đang nạp label từ file cluster...'):
				loaded_labels, mapping = load_cluster_file(cluster_file_path, len(df), noise_label=int(noise_label))
			if mapping is not None:
				# map by filepath
				labels = []
				miss = 0
				for fp in df['filepath']:
					lab = mapping.get(fp)
					if lab is None:
						lab = int(noise_label)
						miss += 1
					labels.append(lab)
				labels = np.array(labels, dtype=int)
				if miss:
					st.warning(f'{miss} filepath không tìm thấy trong file cluster – gán noise.')
			else:
				labels = loaded_labels
			df['cluster'] = labels
			# Try silhouette when >1 cluster and not only noise
			uniq = set(labels)
			if len([u for u in uniq if u != int(noise_label)]) > 1:
				try:
					mask = labels != int(noise_label)
					sil = silhouette_score(coords[mask], labels[mask])
					st.write(f'Silhouette (exclude noise): {sil:.4f}')
				except Exception:
					st.write('Silhouette: N/A')
			else:
				st.write('Silhouette: N/A')
		else:  # None
			df['cluster'] = -1
			st.write('Không áp dụng clustering.')

		# Color scheme
		if color_by_substring.strip():
			subs = [s.strip() for s in color_by_substring.split(',') if s.strip()]
			def color_from_sub(p: str) -> str:
				for i, ssub in enumerate(subs):
					if ssub.lower() in p.lower():
						return ssub
				return 'other'
			df['color_group'] = df['filepath'].apply(color_from_sub)
			color_col = 'color_group'
		else:
			color_col = 'cluster'

		# --- Palette handling -------------------------------------------------
		def get_base_palette(name: str) -> List[str]:
			name = name.lower()
			if name == 'okabe-ito':  # 8 colors, colorblind-safe
				return ["#000000","#E69F00","#56B4E9","#009E73","#F0E442","#0072B2","#D55E00","#CC79A7"]
			if name == 'tol':  # Paul Tol (12)
				return ["#4477AA","#66CCEE","#228833","#CCBB44","#EE6677","#AA3377","#BBBBBB","#000000","#EEDD88","#FFAABB","#99DDFF","#44BB99"]
			if name == 'plotly':
				from plotly.colors import qualitative as q
				return q.Plotly
			if name == 'bold':
				from plotly.colors import qualitative as q
				return q.Bold
			if name == 'dark24':
				from plotly.colors import qualitative as q
				return q.Dark24
			if name == 'set3':
				from plotly.colors import qualitative as q
				return q.Set3
			if name == 'd3':
				from plotly.colors import qualitative as q
				return q.D3
			return []

		def generate_hcl_like(n: int) -> List[str]:
			# Simple evenly spaced hues in HSV then adjust to look more balanced
			cols = []
			for i in range(n):
				h = (i / n) % 1.0
				s = 0.55 + 0.35 * ((i * 37) % 2)  # alternate saturation
				v = 0.85 if (i % 3) else 0.98
				r, g, b = colorsys.hsv_to_rgb(h, s, v)
				cols.append('#%02X%02X%02X' % (int(r*255), int(g*255), int(b*255)))
			return cols

		def build_palette(name: str, k: int) -> List[str]:
			if name == 'auto':
				return []  # let plotly decide
			if name == 'hcl':
				return generate_hcl_like(k)
			base = get_base_palette(name)
			if k <= len(base):
				return base[:k]
			# extend by generating extra colors using HSV golden ratio
			cols = list(base)
			gold = 0.61803398875
			h = 0.1
			while len(cols) < k:
				h = (h + gold) % 1.0
				r, g, b = colorsys.hsv_to_rgb(h, 0.6, 0.95)
				newc = '#%02X%02X%02X' % (int(r*255), int(g*255), int(b*255))
				if newc not in cols:
					cols.append(newc)
			return cols

		color_display_col = color_col + '_display'
		# ensure discrete (string) for numeric clusters
		if np.issubdtype(df[color_col].dtype, np.number):
			df[color_display_col] = df[color_col].astype(int).astype(str)
		else:
			df[color_display_col] = df[color_col].astype(str)
		groups = df[color_display_col].unique()
		palette_seq = build_palette(palette_name, len(groups)) if palette_name else []

		import plotly.express as px
		fig = px.scatter_3d(
			df,
			x='x', y='y', z='z',
			color=color_display_col,
			color_discrete_sequence=palette_seq if palette_seq else None,
			hover_data={'filepath': True, 'cluster': True, 'x': ':.2f', 'y': ':.2f', 'z': ':.2f'},
			title=f'Embedding 3D ({algo})',
			opacity=0.9,
			height=800,
		)
		fig.update_traces(marker={'size': int(marker_size)})
		fig.update_layout(legend=dict(title='Nhóm', itemsizing='constant'))
		st.plotly_chart(fig, use_container_width=True)

		if show_raw:
			st.dataframe(df.head(100))

		# Download
		out_csv = df[['filepath', 'x', 'y', 'z', 'cluster']].to_csv(index=False).encode('utf-8')
		st.download_button('⬇️ Tải toạ độ 3D (CSV)', out_csv, file_name='embedding_3d.csv', mime='text/csv')

		st.caption('Hoàn tất.')


if __name__ == '__main__':  # pragma: no cover
	# Khi chạy bằng 'streamlit run', sys.argv chỉ chứa tên file => ta luôn gọi main()
	# Nếu muốn test nhanh CLI, có thể thêm arg '--cli-test'
	if '--cli-test' in os.sys.argv:
		test_path = 'embeddings_factures_osteopathie_1k_qwen.json'
		if os.path.exists(test_path):
			recs = load_embeddings(test_path, sample_size=5)
			print(f'[CLI TEST] Loaded {len(recs)} embeddings dim={len(recs[0].embedding)}')
		else:
			print('[CLI TEST] Không tìm thấy file test.')
	else:
		main()

