SupplyAgent/web/src/types/heatTree.ts

import type { CategoryNode, NodeCounts, NodeWeights, WeightDimKey } from './category'
import { buildWeightScale, formatWeightScore, scoreHeatT } from './category'

export interface PreparedNode {
  id: number
  name: string
  level: number | null
  description: string | null
  hung_word_count?: number
  weights?: Partial<NodeWeights>
  counts?: Partial<NodeCounts>
  parent: PreparedNode | null
  path: string[]
  children: PreparedNode[]
  leafCount: number
}

export interface PreparedTree {
  roots: PreparedNode[]
  flat: PreparedNode[]
  maxDepth: number
  byId: Map<number, PreparedNode>
}

export const FULL_TREE_TAB = 'full' as const
export type HeatTabKey = typeof FULL_TREE_TAB | WeightDimKey

export const HEAT_DIM_TABS: { key: WeightDimKey; label: string }[] = [
  { key: 'total_score', label: '全局热度' },
  { key: 'ext_pop', label: '外部热度' },
  { key: 'plat_sust_pop', label: '平台持续热度' },
  { key: 'plat_ly_pop', label: '去年同期热度' },
  { key: 'recent_pop', label: '近期热度' },
]

function toPreparedNode(node: CategoryNode): Omit<PreparedNode, 'parent' | 'path' | 'leafCount'> {
  return {
    id: node.id,
    name: node.name || '(未命名)',
    level: node.level,
    description: node.description,
    hung_word_count: node.hung_word_count,
    weights: node.weights,
    counts: node.counts,
    children: [],
  }
}

function computeLeafCount(node: PreparedNode): number {
  if (!node.children.length) {
    node.leafCount = 1
    return 1
  }
  node.leafCount = node.children.reduce((sum, child) => sum + computeLeafCount(child), 0)
  return node.leafCount
}

export function prepareTree(nodes: CategoryNode[]): PreparedTree {
  const flat: PreparedNode[] = []
  const byId = new Map<number, PreparedNode>()
  let maxDepth = 0

  function walk(
    source: CategoryNode[],
    parent: PreparedNode | null,
    path: string[],
  ): PreparedNode[] {
    const out: PreparedNode[] = []
    for (const raw of source) {
      const base = toPreparedNode(raw)
      const node: PreparedNode = {
        ...base,
        parent,
        path: [...path, base.name],
        leafCount: 1,
        children: [],
      }
      node.children = walk(raw.children ?? [], node, node.path)
      flat.push(node)
      byId.set(node.id, node)
      maxDepth = Math.max(maxDepth, node.path.length - 1)
      out.push(node)
    }
    return out
  }

  const roots = walk(nodes, null, [])
  roots.forEach(computeLeafCount)

  return { roots, flat, maxDepth, byId }
}

export function descendantsAtDepth(roots: PreparedNode[], targetDepth: number): PreparedNode[] {
  const results: PreparedNode[] = []
  const walk = (node: PreparedNode) => {
    const depth = node.path.length - 1
    if (depth === targetDepth) {
      results.push(node)
      return
    }
    if (depth < targetDepth) node.children.forEach(walk)
  }
  roots.forEach(walk)
  return results
}

export function activeBaseDepth(startDepth: number, focus: PreparedNode | null): number {
  const focusDepth = focus ? focus.path.length - 1 : 0
  return Math.max(focusDepth, startDepth)
}

export function viewRoots(
  roots: PreparedNode[],
  focus: PreparedNode | null,
  startDepth: number,
): PreparedNode[] {
  const scopeRoots = focus ? [focus] : roots
  const baseDepth = activeBaseDepth(startDepth, focus)
  const view = descendantsAtDepth(scopeRoots, baseDepth)
  return view.length ? view : scopeRoots
}

export function relativeMaxDepth(nodes: PreparedNode[], baseDepth: number): number {
  let maximum = 0
  const walk = (node: PreparedNode) => {
    maximum = Math.max(maximum, node.path.length - 1 - baseDepth)
    node.children.forEach(walk)
  }
  nodes.forEach(walk)
  return maximum
}

export interface IcicleRect {
  node: PreparedNode
  x: number
  y: number
  width: number
  height: number
}

/** Fixed column width per level (px). */
export const COLUMN_WIDTH = 200

export function layoutIcicle(
  roots: PreparedNode[],
  baseDepth: number,
  columnWidth: number,
  canvasHeight: number,
): IcicleRect[] {
  if (!roots.length || columnWidth <= 0 || canvasHeight <= 0) return []

  const totalLeaves = roots.reduce((sum, node) => sum + node.leafCount, 0)
  if (totalLeaves <= 0) return []

  const rows: IcicleRect[] = []
  let yCursor = 0

  const place = (node: PreparedNode, y: number, height: number) => {
    const relativeDepth = node.path.length - 1 - baseDepth
    rows.push({
      node,
      x: relativeDepth * columnWidth,
      y,
      width: columnWidth,
      height,
    })
    let childY = y
    for (const child of node.children) {
      const childHeight = height * (child.leafCount / node.leafCount)
      place(child, childY, childHeight)
      childY += childHeight
    }
  }

  for (const root of roots) {
    const height = canvasHeight * (root.leafCount / totalLeaves)
    place(root, yCursor, height)
    yCursor += height
  }

  return rows
}

/** Depth-based neutral palette (全局树). */
const DEPTH_COLORS = [
  '#dbeafe',
  '#bfdbfe',
  '#93c5fd',
  '#7dd3fc',
  '#a5f3fc',
  '#99f6e4',
  '#bbf7d0',
  '#d9f99d',
  '#fde68a',
  '#fed7aa',
  '#fecaca',
  '#e9d5ff',
]

export function nodeStructureFill(node: PreparedNode): string {
  // Absolute depth so focused/level views keep the same colors as the global tree.
  const absoluteDepth = Math.max(0, node.path.length - 1)
  return DEPTH_COLORS[absoluteDepth % DEPTH_COLORS.length] ?? '#f1f5f9'
}

/**
 * Map-like heat: white = no score; cool blue → teal → amber → red.
 * t is percentile in [0, 1] (全局热度 tab uses min-max linear instead).
 */
export function mapHeatFill(t: number | null): string {
  if (t == null) return '#ffffff'
  const clamp = Math.min(1, Math.max(0, t))
  const stops: [number, [number, number, number]][] = [
    [0, [219, 234, 254]],
    [0.25, [125, 211, 252]],
    [0.5, [52, 211, 153]],
    [0.75, [251, 191, 36]],
    [1, [239, 68, 68]],
  ]
  let i = 0
  while (i < stops.length - 1 && clamp > stops[i + 1][0]) i += 1
  const [t0, c0] = stops[i]
  const [t1, c1] = stops[Math.min(i + 1, stops.length - 1)]
  const u = t1 === t0 ? 0 : (clamp - t0) / (t1 - t0)
  const r = Math.round(c0[0] + (c1[0] - c0[0]) * u)
  const g = Math.round(c0[1] + (c1[1] - c0[1]) * u)
  const b = Math.round(c0[2] + (c1[2] - c0[2]) * u)
  return `rgb(${r}, ${g}, ${b})`
}

export function mapHeatText(t: number | null): string {
  if (t == null) return '#64748b'
  if (t >= 0.72) return '#ffffff'
  return '#0f172a'
}

export function nodeDimScore(
  node: PreparedNode,
  dim: WeightDimKey,
): { weight: number | null; count: number } {
  const count = node.counts?.[dim] ?? 0
  const weight = node.weights?.[dim]
  if (count <= 0 || weight == null || Number.isNaN(weight)) {
    return { weight: null, count }
  }
  return { weight, count }
}

export function collectDimScale(nodes: PreparedNode[], dim: WeightDimKey): number[] {
  const values: number[] = []
  for (const node of nodes) {
    const { weight, count } = nodeDimScore(node, dim)
    if (weight != null && count > 0) values.push(weight)
  }
  return buildWeightScale(values)
}

export function nodeHeatT(
  node: PreparedNode,
  dim: WeightDimKey,
  scale: number[],
): number | null {
  const { weight, count } = nodeDimScore(node, dim)
  if (weight == null || count <= 0) return null
  return scoreHeatT(dim, weight, scale)
}

export function formatNodeDimScore(node: PreparedNode, dim: WeightDimKey): string {
  const { weight, count } = nodeDimScore(node, dim)
  return formatWeightScore(dim, weight, count)
}

export function nodeFillColor(node: PreparedNode, _baseDepth?: number): string {
  return nodeStructureFill(node)
}

export function nodeTextColor(): string {
  return '#1e293b'
}