import { clamp, round } from './rng'

/**
 * Directed power flow over the twin topology (4.2).
 *
 * **What this is, stated plainly, because the distinction matters:** this is a *demand
 * allocation* over the topology, not an AC load-flow solution. It does not solve the
 * power-flow equations, it does not iterate to a voltage/angle fixed point, and it does
 * not model impedance. What it does is answer the question the control room actually
 * asks of a one-line diagram — *where is the energy going, how much of it, and which
 * paths are close to their limit* — from the same telemetry the rest of the platform
 * runs on.
 *
 * Every result carries `method: 'demand_allocation'` so a downstream surface cannot
 * present it as a solved load flow. Replacing it with a Newton–Raphson solver later
 * changes this file and nothing else.
 *
 * The rules:
 *
 * - Generation injects, load withdraws, and storage does either depending on the sign of
 *   its loading — a battery at −40 % is charging, and charging is a load.
 * - An edge carries the demand of everything downstream of it. Downstream demand is
 *   accumulated by a depth-first walk with an explicit visited set, so a ring topology
 *   terminates rather than recursing forever.
 * - **A transit asset's own meter governs the path feeding it, and nothing further.** A
 *   transformer is not a load, so accumulating sink demand alone would leave its
 *   throughput invisible — and a strategy that transfers 190 MW off it would change no
 *   path in the picture. So the edge *into* a node carries `max(demand below, its own
 *   meter)`. But that maximum is deliberately not what its parent accumulates: the parent
 *   adds sink demand only. Letting a meter propagate upward adds the same megawatt again
 *   at every level between the load and the generator, and a substation ends up reporting
 *   more demand than its whole region draws.
 * - **A load fed from two directions is not counted twice.** Each parent takes the share
 *   of a child's demand that its own path capacity represents, so a meshed feeder pair
 *   splits the load between them instead of each carrying all of it.
 * - **Flow reverses when the far end injects more than it consumes.** That is the whole
 *   point of §10: a feeder into a substation whose battery has begun discharging is
 *   carrying power the other way, and the animation has to say so.
 * - Current comes from the flow and the *sending* node's own measured voltage and power
 *   factor, not from a nominal figure — an undervolt bus carries more amps for the same
 *   megawatts, which is exactly when a line matters most.
 */

// ─────────────────────────────────────────────────────────────────────────────
// Inputs
// ─────────────────────────────────────────────────────────────────────────────

/** Asset types that put power into the network rather than taking it out. */
const GENERATION = new Set(['power_plant', 'solar_farm', 'wind_farm'])

/** Asset types that only take power out. */
const CONSUMPTION = new Set(['industrial_load', 'ev_hub'])

/** Storage, which does both, decided by the sign of its loading. */
const STORAGE = new Set(['battery_storage'])

export type FlowRole = 'source' | 'sink' | 'transit' | 'storage'

export interface FlowNodeInput {
  code: string
  assetTypeKey: string
  capacityMw: number
  /** Signed utilisation. Negative on storage means charging. */
  loadPct: number
  loadMw: number
  voltageKv: number
  powerFactor: number
  isOutage: boolean
}

export interface FlowEdgeInput {
  from: string
  to: string
  kind: string
  /** Thermal rating of the path in MW. Zero means "unrated" and is treated as unknown. */
  capacityMw: number
}

export type FlowDirection = 'forward' | 'reverse' | 'idle'

export interface FlowEdgeResult {
  from: string
  to: string
  kind: string
  /**
   * The end power is actually leaving, and the end it is arriving at. For a forward flow
   * these equal `from` and `to`; on a reversal they are swapped, so a renderer can animate
   * `sourceCode → sinkCode` without knowing anything about the topology's declared
   * direction.
   */
  sourceCode: string
  sinkCode: string
  direction: FlowDirection
  /** Always positive. The direction fields carry the sign. */
  activePowerMw: number
  reactivePowerMvar: number
  apparentPowerMva: number
  currentA: number
  powerFactor: number
  voltageKv: number
  capacityMw: number
  utilisationPct: number
  /** Resistive loss estimate, rising with the square of utilisation. */
  lossMw: number
  /** 0–1, for line width and animation speed. Never a constant. */
  intensity: number
  state: 'normal' | 'watch' | 'warning' | 'critical'
  isBottleneck: boolean
}

export interface FlowNodeResult {
  code: string
  role: FlowRole
  /** Positive when the node puts power into the network. */
  netInjectionMw: number
  inflowMw: number
  outflowMw: number
  /** Everything drawn by this node and everything below it. */
  downstreamDemandMw: number
  downstreamNodes: number
  imbalanceMw: number
}

export interface PowerFlowResult {
  method: 'demand_allocation'
  edges: FlowEdgeResult[]
  nodes: FlowNodeResult[]
  totalGenerationMw: number
  totalDemandMw: number
  totalLossMw: number
  /** Largest single edge flow, the figure every intensity is normalised against. */
  peakFlowMw: number
  reversedEdges: number
  bottlenecks: string[]
  /**
   * Generation minus demand minus losses. Never forced to zero: a topology that does not
   * balance is telling you something, and hiding it behind a slack bus would throw that
   * away.
   */
  imbalanceMw: number
}

// ─────────────────────────────────────────────────────────────────────────────
// Roles
// ─────────────────────────────────────────────────────────────────────────────

export function flowRole(node: FlowNodeInput): FlowRole {
  if (STORAGE.has(node.assetTypeKey)) return 'storage'
  if (GENERATION.has(node.assetTypeKey)) return 'source'
  if (CONSUMPTION.has(node.assetTypeKey)) return 'sink'
  return 'transit'
}

/**
 * What a node puts into the network, in MW. Negative means it takes power out.
 *
 * An asset out of service contributes nothing in either direction — not a small residual,
 * and not its last known value.
 */
export function netInjectionMw(node: FlowNodeInput): number {
  if (node.isOutage) return 0
  const magnitude = Math.abs(node.loadMw)

  switch (flowRole(node)) {
    case 'source':
      return magnitude
    case 'sink':
      return -magnitude
    case 'storage':
      // The sign of loading is the whole answer: negative loading is charging, and a
      // charging battery is a load like any other.
      return node.loadPct < 0 ? -magnitude : magnitude
    default:
      // Transformers, substations and lines carry power rather than making or using it.
      // Their losses are attributed to the edges, not to the node.
      return 0
  }
}

// ─────────────────────────────────────────────────────────────────────────────
// Solve
// ─────────────────────────────────────────────────────────────────────────────

const LOSS_COEFFICIENT: Record<string, number> = {
  // Fraction of throughput lost at full rating. Overhead line, cable, transformer and
  // busbar all differ by roughly an order of magnitude in the same direction real plant
  // does; the numbers are indicative, and the field is named an estimate everywhere it
  // surfaces.
  supplies: 0.022,
  feeds: 0.018,
  transmits: 0.031,
  cable: 0.014,
  busbar: 0.004,
  transforms: 0.009,
}

const DEFAULT_LOSS = 0.02

function stateFor(utilisationPct: number): FlowEdgeResult['state'] {
  if (utilisationPct >= 100) return 'critical'
  if (utilisationPct >= 90) return 'warning'
  if (utilisationPct >= 75) return 'watch'
  return 'normal'
}

export function solvePowerFlow(input: {
  nodes: FlowNodeInput[]
  edges: FlowEdgeInput[]
}): PowerFlowResult {
  const byCode = new Map(input.nodes.map((node) => [node.code, node]))

  // Only edges whose both ends are in the scene. An edge to a node that is not being
  // rendered would animate power arriving from nowhere.
  const edges = input.edges.filter((edge) => byCode.has(edge.from) && byCode.has(edge.to))

  const outgoing = new Map<string, FlowEdgeInput[]>()
  const incoming = new Map<string, FlowEdgeInput[]>()
  for (const edge of edges) {
    const out = outgoing.get(edge.from)
    if (out) out.push(edge)
    else outgoing.set(edge.from, [edge])

    const into = incoming.get(edge.to)
    if (into) into.push(edge)
    else incoming.set(edge.to, [edge])
  }

  /**
   * How much of a node's demand each of its feeders is responsible for.
   *
   * A load with one feeder gives that feeder all of it. A load with two gives each the
   * share its capacity represents — which is both the physically sensible split and the
   * only way a meshed network does not report every load several times over.
   */
  const inShare = new Map<string, number>()
  for (const [code, feeders] of incoming) {
    const total = feeders.reduce((sum, edge) => sum + Math.max(0, edge.capacityMw), 0)
    for (const edge of feeders) {
      inShare.set(
        `${edge.from}\u0000${edge.to}`,
        total > 0 ? Math.max(0, edge.capacityMw) / total : 1 / feeders.length,
      )
    }
    void code
  }

  const shareOf = (edge: FlowEdgeInput) => inShare.get(`${edge.from}\u0000${edge.to}`) ?? 1

  const injection = new Map<string, number>()
  const metering = new Map<string, number>()
  for (const node of input.nodes) {
    injection.set(node.code, netInjectionMw(node))
    // What a transit asset's own instrument says is going through it. Out-of-service
    // assets carry nothing, and generation and load are already handled as injection.
    metering.set(
      node.code,
      !node.isOutage && flowRole(node) === 'transit' ? Math.abs(node.loadMw) : 0,
    )
  }

  // ── Downstream demand, memoised, cycle-guarded ────────────────────────────
  //
  // `mw` is sink demand and is what a parent accumulates. `throughput` is what the edge
  // arriving at this node carries, and stops here.
  const demandCache = new Map<string, { mw: number; throughput: number; count: number }>()

  function downstream(code: string, seen: Set<string>): { mw: number; throughput: number; count: number } {
    const cached = demandCache.get(code)
    if (cached) return cached

    // A ring feeder is a legitimate topology, so a repeat visit stops the walk rather
    // than failing it. The demand it would have contributed is already counted on the
    // path that reached it first.
    if (seen.has(code)) return { mw: 0, throughput: 0, count: 0 }
    seen.add(code)

    const own = injection.get(code) ?? 0
    // Only withdrawal counts as demand. Injection is handled as a reversal, below.
    let mw = own < 0 ? -own : 0
    let count = 1

    for (const edge of outgoing.get(code) ?? []) {
      const child = downstream(edge.to, seen)
      // Only this path's share of what the child draws. Without it, a load reachable two
      // ways is carried in full by both.
      mw += child.mw * shareOf(edge)
      count += child.count
    }

    seen.delete(code)

    // The meter, where there is one. A transformer reading 190 MW is carrying 190 MW
    // whether or not every load behind it is in this scene — and when a strategy takes
    // load off it, this is the number that falls. It shapes the path arriving here and
    // is deliberately excluded from `mw`, which is what travels upward.
    const metered = metering.get(code) ?? 0
    const result = {
      mw: round(mw, 4),
      throughput: round(Math.max(mw, metered), 4),
      count,
    }
    // Memoise only acyclic results: inside a cycle the value depends on the path taken
    // to reach it, and caching that would leak one traversal's answer into another's.
    if (seen.size === 0) demandCache.set(code, result)
    return result
  }

  // ── Per-edge flow ─────────────────────────────────────────────────────────
  const results: FlowEdgeResult[] = []
  const inflow = new Map<string, number>()
  const outflow = new Map<string, number>()
  let totalLossMw = 0
  let reversedEdges = 0

  for (const edge of edges) {
    const fromNode = byCode.get(edge.from)!
    const toNode = byCode.get(edge.to)!

    const far = downstream(edge.to, new Set([edge.from]))
    const farInjection = injection.get(edge.to) ?? 0
    const share = shareOf(edge)

    // The far end's own generation serves its own subtree first. What is left over is
    // either demand this edge must supply, or surplus this edge must carry back — and
    // this path carries its share of it, not all of it.
    const netFar = (far.throughput - Math.max(0, farInjection)) * share
    const magnitude = Math.abs(netFar)

    const reversed = netFar < -1e-6
    if (reversed) reversedEdges += 1

    const outOfService = fromNode.isOutage || toNode.isOutage
    const activePowerMw = outOfService ? 0 : round(magnitude, 3)

    const direction: FlowDirection =
      activePowerMw < 1e-6 ? 'idle' : reversed ? 'reverse' : 'forward'

    // The sending end supplies the voltage and power factor. On a reversal that is the
    // node the topology calls the destination.
    const sender = reversed ? toNode : fromNode
    const powerFactor = clamp(sender.powerFactor || 0.95, 0.5, 1)
    const voltageKv = sender.voltageKv > 0.01 ? sender.voltageKv : Math.max(fromNode.voltageKv, toNode.voltageKv)

    const apparentPowerMva = powerFactor > 0 ? activePowerMw / powerFactor : 0
    const reactivePowerMvar = round(
      Math.sqrt(Math.max(0, apparentPowerMva * apparentPowerMva - activePowerMw * activePowerMw)),
      3,
    )

    const currentA =
      voltageKv > 0.01 && activePowerMw > 0
        ? round((apparentPowerMva * 1_000_000) / (Math.sqrt(3) * voltageKv * 1_000), 1)
        : 0

    // An unrated path falls back to the sending node's capacity: better an approximate
    // denominator than a utilisation of infinity.
    const capacityMw = edge.capacityMw > 0 ? edge.capacityMw : Math.max(1, sender.capacityMw)
    const utilisationPct = round((activePowerMw / capacityMw) * 100, 2)

    const coefficient = LOSS_COEFFICIENT[edge.kind] ?? DEFAULT_LOSS
    const loadingPu = capacityMw > 0 ? activePowerMw / capacityMw : 0
    const lossMw = round(coefficient * capacityMw * loadingPu * loadingPu, 4)
    totalLossMw += lossMw

    const sourceCode = reversed ? edge.to : edge.from
    const sinkCode = reversed ? edge.from : edge.to

    outflow.set(sourceCode, (outflow.get(sourceCode) ?? 0) + activePowerMw)
    inflow.set(sinkCode, (inflow.get(sinkCode) ?? 0) + activePowerMw)

    results.push({
      from: edge.from,
      to: edge.to,
      kind: edge.kind,
      sourceCode,
      sinkCode,
      direction,
      activePowerMw,
      reactivePowerMvar,
      apparentPowerMva: round(apparentPowerMva, 3),
      currentA,
      powerFactor: round(powerFactor, 4),
      voltageKv: round(voltageKv, 3),
      capacityMw: round(capacityMw, 2),
      utilisationPct,
      lossMw,
      intensity: 0,
      state: stateFor(utilisationPct),
      isBottleneck: utilisationPct >= 90,
    })
  }

  // ── Normalise intensity ───────────────────────────────────────────────────
  // Against the busiest edge in the scene, so the animation reads as a comparison within
  // this view rather than against an absolute nobody can see.
  const peakFlowMw = results.reduce((peak, edge) => Math.max(peak, edge.activePowerMw), 0)
  for (const edge of results) {
    edge.intensity =
      peakFlowMw > 0
        ? round(clamp(Math.sqrt(edge.activePowerMw / peakFlowMw), 0, 1), 4)
        : 0
  }

  // ── Node summaries ────────────────────────────────────────────────────────
  const nodes: FlowNodeResult[] = input.nodes.map((node) => {
    const own = injection.get(node.code) ?? 0
    const inMw = round(inflow.get(node.code) ?? 0, 3)
    const outMw = round(outflow.get(node.code) ?? 0, 3)
    const far = downstream(node.code, new Set())
    return {
      code: node.code,
      role: flowRole(node),
      netInjectionMw: round(own, 3),
      inflowMw: inMw,
      outflowMw: outMw,
      downstreamDemandMw: far.mw,
      downstreamNodes: Math.max(0, far.count - 1),
      imbalanceMw: round(inMw + Math.max(0, own) - outMw - Math.max(0, -own), 3),
    }
  })

  const totalGenerationMw = round(
    input.nodes.reduce((sum, node) => sum + Math.max(0, netInjectionMw(node)), 0),
    3,
  )
  const totalDemandMw = round(
    input.nodes.reduce((sum, node) => sum + Math.max(0, -netInjectionMw(node)), 0),
    3,
  )

  return {
    method: 'demand_allocation',
    edges: results,
    nodes,
    totalGenerationMw,
    totalDemandMw,
    totalLossMw: round(totalLossMw, 3),
    peakFlowMw: round(peakFlowMw, 3),
    reversedEdges,
    bottlenecks: results.filter((edge) => edge.isBottleneck).map((edge) => `${edge.from}→${edge.to}`),
    imbalanceMw: round(totalGenerationMw - totalDemandMw - totalLossMw, 3),
  }
}

// ─────────────────────────────────────────────────────────────────────────────
// Tracing (§54, §55)
// ─────────────────────────────────────────────────────────────────────────────

/**
 * Walk back from a load to whatever is feeding it.
 *
 * Follows the *solved* flow rather than the declared topology, so a path that has
 * reversed is traced the way power is actually moving — which is the only version of the
 * answer worth showing.
 */
export function tracePowerSource(
  flow: PowerFlowResult,
  loadCode: string,
  maxDepth = 12,
): { path: string[]; edges: FlowEdgeResult[]; sources: string[] } {
  const incoming = new Map<string, FlowEdgeResult[]>()
  for (const edge of flow.edges) {
    if (edge.direction === 'idle') continue
    const list = incoming.get(edge.sinkCode)
    if (list) list.push(edge)
    else incoming.set(edge.sinkCode, [edge])
  }

  const path: string[] = [loadCode]
  const used: FlowEdgeResult[] = []
  const seen = new Set([loadCode])
  let current = loadCode

  for (let depth = 0; depth < maxDepth; depth += 1) {
    const feeders = incoming.get(current)
    if (!feeders || feeders.length === 0) break

    // The dominant feeder. A load fed from two directions has a majority path, and
    // showing every branch at once is how a trace becomes unreadable.
    const strongest = feeders
      .filter((edge) => !seen.has(edge.sourceCode))
      .sort((a, b) => b.activePowerMw - a.activePowerMw)[0]
    if (!strongest) break

    used.push(strongest)
    current = strongest.sourceCode
    seen.add(current)
    path.push(current)
  }

  const roles = new Map(flow.nodes.map((node) => [node.code, node.role]))
  return {
    path,
    edges: used,
    sources: path.filter((code) => roles.get(code) === 'source' || roles.get(code) === 'storage'),
  }
}

/**
 * Everything downstream of an asset — what stops being fed if it stops.
 *
 * The counterpart of the trace above, and the honest answer to "what does this transformer
 * matter to".
 */
export function traceDependentLoads(
  flow: PowerFlowResult,
  assetCode: string,
  maxDepth = 12,
): { codes: string[]; edges: FlowEdgeResult[]; demandMw: number } {
  const forward = new Map<string, FlowEdgeResult[]>()
  for (const edge of flow.edges) {
    if (edge.direction === 'idle') continue
    const list = forward.get(edge.sourceCode)
    if (list) list.push(edge)
    else forward.set(edge.sourceCode, [edge])
  }

  const codes: string[] = []
  const used: FlowEdgeResult[] = []
  const seen = new Set([assetCode])
  let frontier = [assetCode]

  for (let depth = 0; depth < maxDepth && frontier.length > 0; depth += 1) {
    const next: string[] = []
    for (const code of frontier) {
      for (const edge of forward.get(code) ?? []) {
        if (seen.has(edge.sinkCode)) continue
        seen.add(edge.sinkCode)
        codes.push(edge.sinkCode)
        used.push(edge)
        next.push(edge.sinkCode)
      }
    }
    frontier = next
  }

  const byCode = new Map(flow.nodes.map((node) => [node.code, node]))
  const demandMw = round(
    codes.reduce((sum, code) => {
      const node = byCode.get(code)
      return sum + (node && node.netInjectionMw < 0 ? -node.netInjectionMw : 0)
    }, 0),
    3,
  )

  return { codes, edges: used, demandMw }
}
