import type { ToolsDataset, ToolsMetricMap, ToolsScatterRowId } from '$lib/types/tools-data';
import { adaptBucketRanking } from './adapt-ranking';
import type {
	BucketLatencyMs,
	BucketMetricMap,
	BucketScatterRowId,
	BucketVisualizationSnapshot
} from './types';

function modelKeyById(dataset: ReturnType<typeof adaptBucketRanking>) {
	return new Map(
		dataset.models.map((model) => [model.modelId ?? model.name, model.modelId ?? model.name])
	);
}

function groupKeyById(snapshot: BucketVisualizationSnapshot) {
	return new Map(snapshot.catalog.groups.map((group) => [group.group_id, group.label.en]));
}

function modelKey(id: string, keys: ReadonlyMap<string, string>) {
	return keys.get(id) ?? id;
}

function groupKey(id: string, keys: ReadonlyMap<string, string>) {
	return keys.get(id) ?? id;
}

function latencyMs(latency: BucketLatencyMs) {
	return {
		p50: latency.p50,
		p95: latency.p95,
		p99: latency.p99
	};
}

function remapMetricMap(
	map: BucketMetricMap,
	rowKey: (id: string) => string,
	modelKeyForId: (id: string) => string
): ToolsMetricMap {
	return Object.fromEntries(
		Object.entries(map).map(([rowId, row]) => [
			rowKey(rowId),
			Object.fromEntries(
				Object.entries(row).map(([modelId, value]) => [modelKeyForId(modelId), value])
			)
		])
	);
}

function remapScatterValues(
	values: Record<BucketScatterRowId, Record<string, number | null>>,
	modelKeyForId: (id: string) => string
) {
	return Object.fromEntries(
		Object.entries(values).map(([rowId, row]) => [
			rowId,
			Object.fromEntries(
				Object.entries(row).map(([modelId, value]) => [modelKeyForId(modelId), value])
			)
		])
	) as Record<ToolsScatterRowId, Record<string, number | null>>;
}

export function adaptBucketTools(snapshot: BucketVisualizationSnapshot): ToolsDataset {
	const ranking = adaptBucketRanking(snapshot);
	const modelKeys = modelKeyById(ranking);
	const groupKeys = groupKeyById(snapshot);
	const modelKeyForId = (id: string) => modelKey(id, modelKeys);
	const groupKeyForId = (id: string) => groupKey(id, groupKeys);
	const datasetRefs = snapshot.catalog.datasets.map((dataset) => ({
		id: dataset.dataset_id,
		groupId: dataset.group_id,
		groupKey: groupKeyForId(dataset.group_id),
		label: dataset.label,
		language: dataset.language
	}));

	return {
		...ranking,
		datasets: datasetRefs,
		visualizations: {
			radar: {
				title: snapshot.radar.title,
				subtitle: snapshot.radar.subtitle,
				modelKeys: snapshot.radar.model_ids.map(modelKeyForId),
				groupKeys: snapshot.radar.group_ids.map(groupKeyForId),
				defaultModelKeys: snapshot.radar.default_model_ids.map(modelKeyForId),
				values: remapMetricMap(snapshot.radar.values, groupKeyForId, modelKeyForId),
				f1Values: remapMetricMap(snapshot.radar.f1_values, groupKeyForId, modelKeyForId)
			},
			scatter: {
				title: snapshot.scatter.title,
				xAxis: snapshot.scatter.x_axis,
				yAxis: snapshot.scatter.y_axis,
				points: snapshot.scatter.points.map((point) => ({
					modelKey: modelKeyForId(point.model_id),
					fpr: point.fpr,
					fnr: point.fnr,
					tooltip: point.tooltip
				})),
				rowIds: snapshot.scatter.row_ids,
				values: remapScatterValues(snapshot.scatter.values, modelKeyForId)
			},
			heatmap: {
				title: snapshot.heatmap.title,
				subtitle: snapshot.heatmap.subtitle,
				metric: snapshot.heatmap.metric,
				rows: datasetRefs.filter((dataset) => snapshot.heatmap.row_ids.includes(dataset.id)),
				columnModelKeys: snapshot.heatmap.column_ids.map(modelKeyForId),
				values: remapMetricMap(snapshot.heatmap.values, (id) => id, modelKeyForId)
			},
			groupedBars: {
				title: snapshot.grouped_bars.title,
				subtitle: snapshot.grouped_bars.subtitle,
				groupKeys: snapshot.grouped_bars.group_ids.map(groupKeyForId),
				models: snapshot.grouped_bars.models.map((model) => ({
					modelKey: modelKeyForId(model.model_id),
					groups: Object.fromEntries(
						Object.entries(model.groups).map(([groupId, metrics]) => [
							groupKeyForId(groupId),
							metrics
						])
					)
				}))
			},
			pareto: {
				title: snapshot.pareto.title,
				subtitle: snapshot.pareto.subtitle,
				xAxis: snapshot.pareto.x_axis,
				yAxis: snapshot.pareto.y_axis,
				byModelKey: Object.fromEntries(
					Object.entries(snapshot.pareto.by_model_id).map(([modelId, item]) => [
						modelKeyForId(modelId),
						{
							latencyMs: latencyMs(item.latency_ms),
							integral: item.integral,
							fpr: item.fpr
						}
					])
				)
			},
			performance: {
				title: snapshot.performance.title,
				rows: snapshot.performance.rows.map((row) => ({
					modelKey: modelKeyForId(row.model_id),
					rank: row.rank,
					latencyMs: latencyMs(row.latency_ms),
					errorRate: row.error_rate,
					integral: row.integral
				}))
			},
			robustness: {
				title: snapshot.robustness.title,
				rows: snapshot.robustness.rows.map((row) => ({
					modelKey: modelKeyForId(row.model_id),
					scoreReal: row.score_real,
					scoreRobust: row.score_robust,
					deltaScore: row.delta_score,
					fnrReal: row.fnr_real,
					fnrRobust: row.fnr_robust,
					deltaFnr: row.delta_fnr,
					fprReal: row.fpr_real,
					fprRobust: row.fpr_robust,
					deltaFpr: row.delta_fpr,
					categoryLabel: row.category_label
				}))
			}
		},
		drilldownIndex: {
			heatmap: snapshot.drilldownIndex.heatmap.map((row) => ({
				modelKey: modelKeyForId(row.model_id),
				groupKey: groupKeyForId(row.group_id),
				datasetIds: row.dataset_ids,
				selectedMetric: row.selected_metric,
				metrics: row.metrics,
				description: row.description,
				promptViewerAvailable: row.prompt_viewer_available,
				promptViewerPath: row.prompt_viewer_path
			}))
		}
	};
}