import type { Dataset, BenchmarkRecord } from '../types'

/**
 * Typed normalization + analysis layer for the kernel scaling explorer.
 *
 * Raw JSONL records (see ../types.ts) are normalized into measurements with
 * numeric shape axes. All aggregation is explicit: repeated measurements of
 * one configuration are collapsed with the median (the benchmark suite's
 * declared aggregation policy), and every analysis counts each unique
 * configuration once.
 */

/** The benchmark suite reports medians; repeats of a config are median-aggregated. */
export const AGGREGATION_POLICY = 'median' as const

/** Above this many series the scaling chart demands an explicit selection. */
export const MAX_SCALING_SERIES = 12

/** Default exploratory display gates (coverage rules, not significance claims). */
export const DEFAULT_GATES = { minConfigs: 12, minValues: 3 }

/** One raw benchmark run, normalized. */
export interface NormalizedMeasurement {
  /** Benchmark identity (the dataset / file name, e.g. "matmul"). */
  kernel: string
  /** Git commit id (`rev`, or legacy `commit`); 'unknown' when absent. */
  rev: string
  date: string
  host: string
  hardware: string
  /** Non-axis identity tags: kernel impl, dtype/quantization, opt_level, … */
  tags: Record<string, string>
  /** Shape dimensions as numbers, e.g. { m: 4096, k: 4096, n: 8192 }. */
  axes: Record<string, number>
  metrics: Record<string, number>
}

export interface KernelModel {
  kernel: string
  measurements: NormalizedMeasurement[]
  /** Axis names in recorded order (m, k, n / bs, seq, heads, d / …). */
  axisNames: string[]
  metricNames: string[]
}

/** Multi-select run filters; empty selection = no restriction. */
export interface RunFilters {
  hosts: string[]
  revs: string[]
  hardware: string[]
  dateFrom: string
  dateTo: string
}

/** Axis value selection: OR within one axis, AND across axes, [] = All. */
export type AxisFilters = Record<string, number[]>

export interface ExplorerFilters {
  run: RunFilters
  axes: AxisFilters
}

export function emptyRunFilters(): RunFilters {
  return { hosts: [], revs: [], hardware: [], dateFrom: '', dateTo: '' }
}

export function emptyFilters(): ExplorerFilters {
  return { run: emptyRunFilters(), axes: {} }
}

// ---------------------------------------------------------------------------
// Normalization
// ---------------------------------------------------------------------------

const AXIS_PREFIX = 'shape.'

/** Normalize one parsed record (tags already contain shape.<dim> strings). */
export function normalizeMeasurement(kernel: string, r: BenchmarkRecord): NormalizedMeasurement | null {
  const axes: Record<string, number> = {}
  for (const [k, v] of Object.entries(r.tags)) {
    if (!k.startsWith(AXIS_PREFIX)) continue
    const n = Number(v)
    if (!Number.isFinite(n)) return null
    axes[k.slice(AXIS_PREFIX.length)] = n
  }
  if (Object.keys(axes).length === 0) return null
  const tags: Record<string, string> = {}
  for (const [k, v] of Object.entries(r.tags)) {
    // rev/commit are identity fields (m.rev), not execution settings.
    if (!k.startsWith(AXIS_PREFIX) && k !== 'rev' && k !== 'commit') tags[k] = v
  }
  return {
    kernel,
    rev: r.tags.rev ?? r.tags.commit ?? 'unknown',
    date: r.date,
    host: r.host,
    hardware: r.hardware,
    tags,
    axes,
    metrics: { ...r.metrics },
  }
}

/** Build the kernel model for a dataset; null when it holds no shape axes. */
export function buildKernelModel(ds: Dataset): KernelModel | null {
  const measurements: NormalizedMeasurement[] = []
  const axisNames: string[] = []
  const metricNames: string[] = []
  for (const r of ds.records) {
    const m = normalizeMeasurement(ds.name, r)
    if (!m) continue
    measurements.push(m)
    for (const a of Object.keys(m.axes)) if (!axisNames.includes(a)) axisNames.push(a)
    for (const k of Object.keys(m.metrics)) if (!metricNames.includes(k)) metricNames.push(k)
  }
  if (measurements.length === 0 || axisNames.length === 0) return null
  return { kernel: ds.name, measurements, axisNames, metricNames }
}

// ---------------------------------------------------------------------------
// Filtering
// ---------------------------------------------------------------------------

/** Matches one measurement against run + axis filters. */
export function matchesFilters(m: NormalizedMeasurement, f: ExplorerFilters): boolean {
  const r = f.run
  if (r.hosts.length > 0 && !r.hosts.includes(m.host)) return false
  if (r.revs.length > 0 && !r.revs.includes(m.rev)) return false
  if (r.hardware.length > 0 && !r.hardware.includes(m.hardware)) return false
  if (r.dateFrom && m.date < r.dateFrom) return false
  if (r.dateTo && m.date > r.dateTo) return false
  for (const [axis, values] of Object.entries(f.axes)) {
    if (!values || values.length === 0) continue
    const v = m.axes[axis]
    if (v === undefined || !values.includes(v)) return false
  }
  return true
}

export function applyExplorerFilters(model: KernelModel, f: ExplorerFilters): NormalizedMeasurement[] {
  return model.measurements.filter((m) => matchesFilters(m, f))
}

// ---------------------------------------------------------------------------
// Configuration identity and repeat aggregation
// ---------------------------------------------------------------------------

/**
 * Configuration identity for measurement matching: kernel, full shape,
 * host/device and every execution setting tag — deliberately WITHOUT rev or
 * date (revision varies for comparison, dates are repeats).
 */
export function configKey(m: NormalizedMeasurement): string {
  const axes = Object.keys(m.axes)
    .sort()
    .map((k) => `${k}=${m.axes[k]}`)
    .join(' ')
  const tags = Object.keys(m.tags)
    .sort()
    .map((k) => `${k}=${m.tags[k]}`)
    .join(' ')
  return [m.kernel, axes, m.host, m.hardware, tags].join('|')
}

/**
 * Everything but axes/rev/date — one homogeneous comparison group inside
 * which coefficients are never mixed across hosts or revisions.
 */
export function groupKey(m: NormalizedMeasurement): string {
  const tags = Object.keys(m.tags)
    .sort()
    .map((k) => `${k}=${m.tags[k]}`)
    .join(' ')
  return [m.kernel, m.host, m.hardware, m.rev, tags].join('|')
}

export function groupLabel(m: NormalizedMeasurement): string {
  const shortRev = m.rev.length > 12 ? m.rev.slice(0, 8) : m.rev
  const tags = Object.keys(m.tags)
    .sort()
    .map((k) => `${k}=${m.tags[k]}`)
    .join(' · ')
  return [m.host, shortRev, tags].filter(Boolean).join(' · ')
}

/** Median of a numeric sample (the declared aggregation policy). */
export function median(values: number[]): number {
  const sorted = [...values].sort((a, b) => a - b)
  const mid = sorted.length >> 1
  return sorted.length % 2 === 1 ? sorted[mid] : (sorted[mid - 1] + sorted[mid]) / 2
}

/** One unique configuration: repeats aggregated per metric, one point per rev. */
export interface ConfigPoint {
  key: string
  axes: Record<string, number>
  host: string
  hardware: string
  tags: Record<string, string>
  rev: string
  /** Median of the repeats per metric. */
  metrics: Record<string, number>
  /** How many raw measurements were aggregated. */
  repeats: number
  dates: string[]
}

/**
 * Aggregate measurements into unique configurations (per rev): each
 * configuration counts once regardless of how many repeats it has.
 */
export function aggregateConfigs(ms: NormalizedMeasurement[]): ConfigPoint[] {
  const buckets = new Map<string, NormalizedMeasurement[]>()
  for (const m of ms) {
    const key = `${configKey(m)}|${m.rev}`
    const b = buckets.get(key)
    if (b) b.push(m)
    else buckets.set(key, [m])
  }
  return [...buckets.values()].map((group) => {
    const first = group[0]
    const metrics: Record<string, number> = {}
    for (const k of Object.keys(first.metrics)) {
      const values = group.map((m) => m.metrics[k]).filter((v) => v !== undefined)
      if (values.length > 0) metrics[k] = median(values)
    }
    return {
      key: `${configKey(first)}|${first.rev}`,
      axes: first.axes,
      host: first.host,
      hardware: first.hardware,
      tags: first.tags,
      rev: first.rev,
      metrics,
      repeats: group.length,
      dates: [...new Set(group.map((m) => m.date))].sort(),
    }
  })
}

export interface ScopeCounts {
  /** Unique (shape × host × settings) configurations. */
  configurations: number
  /** Raw measurements (all repeats). */
  samples: number
}

export function countScope(ms: NormalizedMeasurement[]): ScopeCounts {
  return { configurations: new Set(ms.map(configKey)).size, samples: ms.length }
}

// ---------------------------------------------------------------------------
// Scaling chart model
// ---------------------------------------------------------------------------

export interface ScalingPoint {
  x: number
  /** Median metric over the repeats at this x of this series. */
  y: number
  repeats: number
  date: string
  rev: string
  /** Baseline value at the same configuration, or null when unavailable. */
  baselineY: number | null
}

export interface ScalingSeries {
  id: string
  label: string
  rev: string
  /** Values of every dimension except x that identify this curve. */
  fixed: Record<string, number>
  /** Measured points only, sorted by x. Never interpolated. */
  points: ScalingPoint[]
  isBaseline: boolean
}

export interface ScalingModel {
  x: string
  metric: string
  series: ScalingSeries[]
  seriesCount: number
  baselineRev: string | null
  /** False when any series holds a non-positive x value (blocks log scale). */
  positiveXOnly: boolean
  /** True when only one distinct x value remains across the scope. */
  needsMoreX: boolean
}

/**
 * Curves for the scaling chart: one series per (rev × host × settings ×
 * remaining axis values). Points are measured values only — segments between
 * them are drawn as straight connectors and labelled as such in the UI.
 * Baseline values are attached per point; missing baselines stay null
 * ("unavailable"), never zero.
 */
export function buildScalingModel(
  ms: NormalizedMeasurement[],
  x: string,
  metric: string,
  baselineRev: string | null
): ScalingModel {
  interface Bucket {
    rev: string
    host: string
    hardware: string
    tags: Record<string, string>
    fixed: Record<string, number>
    ys: number[]
    xs: number[]
    repeats: number
    date: string
  }
  const buckets = new Map<string, Bucket>()
  // Baseline lookup: same fixed configuration at the baseline rev.
  const baselineIndex = new Map<string, Map<number, number[]>>()

  const seriesId = (rev: string, host: string, hardware: string, tags: Record<string, string>, fixed: Record<string, number>) =>
    [
      rev,
      host,
      hardware,
      Object.keys(tags).sort().map((k) => `${k}=${tags[k]}`).join(' '),
      Object.keys(fixed).sort().map((k) => `${k}=${fixed[k]}`).join(' '),
    ].join('|')

  for (const m of ms) {
    const xv = m.axes[x]
    const yv = m.metrics[metric]
    if (xv === undefined || yv === undefined) continue
    const fixed: Record<string, number> = {}
    for (const [k, v] of Object.entries(m.axes)) if (k !== x) fixed[k] = v
    const id = seriesId(m.rev, m.host, m.hardware, m.tags, fixed)
    let b = buckets.get(id)
    if (!b) {
      b = { rev: m.rev, host: m.host, hardware: m.hardware, tags: m.tags, fixed, ys: [], xs: [], repeats: 0, date: m.date }
      buckets.set(id, b)
    }
    b.xs.push(xv)
    b.ys.push(yv)
    b.repeats += 1
    if (m.date > b.date) b.date = m.date

    if (baselineRev && m.rev === baselineRev) {
      const fk = seriesId('', m.host, m.hardware, m.tags, fixed).slice(1) // identity without rev
      let byX = baselineIndex.get(fk)
      if (!byX) baselineIndex.set(fk, (byX = new Map()))
      const arr = byX.get(xv)
      if (arr) arr.push(yv)
      else byX.set(xv, [yv])
    }
  }

  const series: ScalingSeries[] = []
  for (const [id, b] of buckets) {
    // Aggregate repeats at identical x within the series.
    const byX = new Map<number, number[]>()
    for (let i = 0; i < b.xs.length; i++) {
      const arr = byX.get(b.xs[i])
      if (arr) arr.push(b.ys[i])
      else byX.set(b.xs[i], [b.ys[i]])
    }
    const fk = seriesId('', b.host, b.hardware, b.tags, b.fixed).slice(1)
    const points: ScalingPoint[] = [...byX.entries()]
      .map(([xv, ys]) => {
        let baselineY: number | null = null
        if (baselineRev && b.rev !== baselineRev) {
          const base = baselineIndex.get(fk)?.get(xv)
          if (base && base.length > 0) baselineY = median(base)
        } else if (baselineRev && b.rev === baselineRev) {
          baselineY = median(ys)
        }
        return {
          x: xv,
          y: median(ys),
          repeats: ys.length,
          date: b.date,
          rev: b.rev,
          baselineY,
        }
      })
      .sort((a, c) => a.x - c.x)
    const labelParts = [
      b.rev.length > 12 ? b.rev.slice(0, 8) : b.rev,
      ...Object.keys(b.fixed).sort().map((k) => `${k}=${formatAxisValue(b.fixed[k])}`),
    ]
    series.push({
      id,
      label: labelParts.join(' · '),
      rev: b.rev,
      fixed: b.fixed,
      points,
      isBaseline: baselineRev !== null && b.rev === baselineRev,
    })
  }
  series.sort((a, b) => a.label.localeCompare(b.label))

  const xs = series.flatMap((s) => s.points.map((p) => p.x))
  return {
    x,
    metric,
    series,
    seriesCount: series.length,
    baselineRev,
    positiveXOnly: xs.every((v) => v > 0),
    needsMoreX: new Set(xs).size < 2,
  }
}

/** Compact dimension values for labels: 4096 -> 4k. */
export function formatAxisValue(v: number): string {
  if (v >= 1000 && v % 1024 === 0) return `${v / 1024}k`
  if (v >= 1000 && v % 1000 === 0) return `${v / 1000}k`
  return String(v)
}

// ---------------------------------------------------------------------------
// Correlation analysis
// ---------------------------------------------------------------------------

export type CorrelationScope = 'filtered' | 'sweep'

export interface CorrelationGroup {
  key: string
  label: string
  rev: string
  host: string
  hardware: string
  tags: Record<string, string>
  configurations: number
}

export type CellReason = 'constant-axis' | 'constant-metric' | 'coverage' | 'missing'

export interface CorrelationCell {
  axis: string
  metric: string
  /** Spearman rho (average ranks for ties); null when not computable. */
  rho: number | null
  /** Why rho is unavailable, when it is. */
  reason?: CellReason
  /** Unique configurations that contributed. */
  n: number
  distinctValues: number
  /** Exploratory display gate result (coverage rule, not significance). */
  passed: boolean
}

export interface CorrelationModel {
  scope: CorrelationScope
  group: CorrelationGroup | null
  axes: string[]
  metrics: string[]
  cells: CorrelationCell[]
  configurations: number
  gates: { minConfigs: number; minValues: number }
}

/** Average ranks for ties (each tie group gets the mean of its rank span). */
export function rankAverage(values: number[]): number[] {
  const indexed = values.map((v, i) => ({ v, i }))
  indexed.sort((a, b) => a.v - b.v)
  const ranks = new Array<number>(values.length)
  let i = 0
  while (i < indexed.length) {
    let j = i
    while (j + 1 < indexed.length && indexed[j + 1].v === indexed[i].v) j++
    const avg = (i + j + 2) / 2 // 1-based ranks i+1..j+1 averaged
    for (let k = i; k <= j; k++) ranks[indexed[k].i] = avg
    i = j + 1
  }
  return ranks
}

/**
 * Spearman rank correlation. Returns null when undefined (fewer than two
 * points or a constant input) — never a misleading zero.
 */
export function spearman(xs: number[], ys: number[]): number | null {
  if (xs.length !== ys.length || xs.length < 2) return null
  const rx = rankAverage(xs)
  const ry = rankAverage(ys)
  const mx = rx.reduce((a, b) => a + b, 0) / rx.length
  const my = ry.reduce((a, b) => a + b, 0) / ry.length
  let num = 0
  let dx = 0
  let dy = 0
  for (let i = 0; i < rx.length; i++) {
    num += (rx[i] - mx) * (ry[i] - my)
    dx += (rx[i] - mx) ** 2
    dy += (ry[i] - my) ** 2
  }
  if (dx === 0 || dy === 0) return null
  return num / Math.sqrt(dx * dy)
}

/** Comparison groups present in scope (never silently mixed in one cell). */
export function correlationGroups(
  model: KernelModel,
  f: ExplorerFilters,
  scope: CorrelationScope
): CorrelationGroup[] {
  const pool = scope === 'sweep' ? sweepMeasurements(model, f) : applyExplorerFilters(model, f)
  const groups = new Map<string, CorrelationGroup>()
  for (const m of pool) {
    const key = groupKey(m)
    let g = groups.get(key)
    if (!g) groups.set(key, (g = { key, label: groupLabel(m), rev: m.rev, host: m.host, hardware: m.hardware, tags: m.tags, configurations: 0 }))
  }
  for (const g of groups.values()) {
    g.configurations = new Set(pool.filter((m) => groupKey(m) === g.key).map(configKey)).size
  }
  return [...groups.values()].sort((a, b) => a.label.localeCompare(b.label))
}

/**
 * "Whole sweep": ignore ONLY the axis filters; keep kernel, host, revision,
 * period and execution settings from the current selection.
 */
function sweepMeasurements(model: KernelModel, f: ExplorerFilters): NormalizedMeasurement[] {
  return applyExplorerFilters(model, { run: f.run, axes: {} })
}

/**
 * Correlation matrix within ONE homogeneous group (host × rev × settings).
 * Each unique configuration contributes exactly one (axis, metric) pair —
 * repeats are median-aggregated and never counted as extra shapes.
 */
export function buildCorrelationModel(
  model: KernelModel,
  f: ExplorerFilters,
  scope: CorrelationScope,
  groupKeySelected: string | null,
  gates: { minConfigs: number; minValues: number } = DEFAULT_GATES
): CorrelationModel {
  const pool = scope === 'sweep' ? sweepMeasurements(model, f) : applyExplorerFilters(model, f)
  const groups = correlationGroups(model, f, scope)
  const group = groups.find((g) => g.key === groupKeySelected) ?? groups[0] ?? null

  const cells: CorrelationCell[] = []
  let configurations = 0
  if (group) {
    const inGroup = pool.filter((m) => groupKey(m) === group.key)
    const configs = aggregateConfigs(inGroup)
    configurations = configs.length
    for (const axis of model.axisNames) {
      for (const metric of model.metricNames) {
        cells.push(cellFor(configs, axis, metric, gates))
      }
    }
  }
  return { scope, group, axes: model.axisNames, metrics: model.metricNames, cells, configurations, gates }
}

function cellFor(
  configs: ConfigPoint[],
  axis: string,
  metric: string,
  gates: { minConfigs: number; minValues: number }
): CorrelationCell {
  const xs: number[] = []
  const ys: number[] = []
  for (const c of configs) {
    const xv = c.axes[axis]
    const yv = c.metrics[metric]
    if (xv === undefined || yv === undefined) continue
    xs.push(xv)
    ys.push(yv)
  }
  const distinctValues = new Set(xs).size
  const n = xs.length
  const base: CorrelationCell = { axis, metric, rho: null, n, distinctValues, passed: false }

  if (n === 0) return { ...base, reason: 'missing' }
  if (distinctValues < 2) return { ...base, reason: 'constant-axis' }
  if (new Set(ys).size < 2) return { ...base, reason: 'constant-metric' }

  const rho = spearman(xs, ys)
  if (rho === null) return { ...base, reason: 'missing' }

  const passed = n >= gates.minConfigs && distinctValues >= gates.minValues
  if (!passed) return { ...base, rho, reason: 'coverage', passed: false }
  return { ...base, rho, passed: true }
}

// ---------------------------------------------------------------------------
// Anchor configuration + "inspect scaling"
// ---------------------------------------------------------------------------

export interface Anchor {
  /** Values to fix for every axis except `x`. */
  fixed: Record<string, number>
  /** Measured x values at the anchor — proves the sweep is real. */
  xValues: number[]
  /** Raw measurements behind the anchor. */
  samples: number
}

/**
 * Pick the measured anchor for inspecting `x` scaling: the values of the
 * other axes with the most measured x coverage (ties: most samples, then
 * newest). The anchor always comes from actual measurements.
 */
export function selectAnchor(ms: NormalizedMeasurement[], x: string): Anchor | null {
  interface Proj {
    fixed: Record<string, number>
    xs: Set<number>
    samples: number
    date: string
  }
  const projs = new Map<string, Proj>()
  for (const m of ms) {
    const xv = m.axes[x]
    if (xv === undefined) continue
    const fixed: Record<string, number> = {}
    for (const [k, v] of Object.entries(m.axes)) if (k !== x) fixed[k] = v
    const key = Object.keys(fixed).sort().map((k) => `${k}=${fixed[k]}`).join(' ')
    let p = projs.get(key)
    if (!p) projs.set(key, (p = { fixed, xs: new Set(), samples: 0, date: m.date }))
    p.xs.add(xv)
    p.samples += 1
    if (m.date > p.date) p.date = m.date
  }
  const list = [...projs.values()]
  if (list.length === 0) return null
  list.sort((a, b) => b.xs.size - a.xs.size || b.samples - a.samples || b.date.localeCompare(a.date))
  const best = list[0]
  return { fixed: best.fixed, xValues: [...best.xs].sort((a, b) => a - b), samples: best.samples }
}

export interface InspectPlan {
  x: string
  /** Filters with every other axis fixed to the anchor values. */
  filters: ExplorerFilters
  anchor: Anchor
}

/**
 * "Inspect [axis] scaling": set x, clear the axis' own filter and fix every
 * other axis to a valid measured anchor configuration. Run filters and the
 * period are preserved. The returned state is explicit — callers should
 * display which values were fixed.
 */
export function inspectAxis(model: KernelModel, f: ExplorerFilters, x: string): InspectPlan | null {
  const pool = applyExplorerFilters(model, f)
  const anchor = selectAnchor(pool.length > 0 ? pool : model.measurements, x)
  if (!anchor) return null
  const axes: AxisFilters = { ...f.axes, [x]: [] }
  for (const [k, v] of Object.entries(anchor.fixed)) axes[k] = [v]
  return { x, filters: { run: f.run, axes }, anchor }
}

/**
 * The axis with the most distinct measured values — the natural default
 * sweep for the starting view.
 */
export function dominantAxis(model: KernelModel): string {
  let best = model.axisNames[0]
  let bestCount = -1
  for (const a of model.axisNames) {
    const n = new Set(model.measurements.map((m) => m.axes[a]).filter((v) => v !== undefined)).size
    if (n > bestCount) {
      bestCount = n
      best = a
    }
  }
  return best
}

// ---------------------------------------------------------------------------
// Linked scatter inspection
// ---------------------------------------------------------------------------

export interface ScatterPoint {
  x: number
  y: number
  repeats: number
  dates: string[]
  rev: string
  /** Full configuration + measurement description for hover/focus. */
  label: string
}

/**
 * One point per unique configuration for a selected (axis, metric) pair,
 * within the selected homogeneous group and scope. Repeats are aggregated
 * (median) and reported as `repeats`.
 */
export function buildScatter(
  model: KernelModel,
  f: ExplorerFilters,
  scope: CorrelationScope,
  groupKeySelected: string | null,
  axis: string,
  metric: string
): ScatterPoint[] {
  const pool = scope === 'sweep' ? sweepMeasurements(model, f) : applyExplorerFilters(model, f)
  const groups = correlationGroups(model, f, scope)
  const group = groups.find((g) => g.key === groupKeySelected) ?? groups[0] ?? null
  if (!group) return []
  const configs = aggregateConfigs(pool.filter((m) => groupKey(m) === group.key))
  const points: ScatterPoint[] = []
  for (const c of configs) {
    const xv = c.axes[axis]
    const yv = c.metrics[metric]
    if (xv === undefined || yv === undefined) continue
    const axes = Object.keys(c.axes)
      .sort()
      .map((k) => `${k}=${formatAxisValue(c.axes[k])}`)
      .join(' ')
    const tags = Object.keys(c.tags)
      .sort()
      .map((k) => `${k}=${c.tags[k]}`)
      .join(' · ')
    points.push({
      x: xv,
      y: yv,
      repeats: c.repeats,
      dates: c.dates,
      rev: c.rev,
      label: `${axes}${tags ? ` · ${tags}` : ''} · ${c.repeats} run${c.repeats === 1 ? '' : 's'} (${c.dates.join(', ')})`,
    })
  }
  return points.sort((a, b) => a.x - b.x || a.y - b.y)
}