import {AdditionalNode, GraphNode, GraphNodeWithIds, isAdditionalNode} from "./types";
import {Dict} from "../types";
import deepcopy from "deepcopy";
import {isNonNullable, shuffleInplace, sortByProperty, uniquify} from "../utils";
import seedrandom from "seedrandom";

function generateIds<T extends Dict<unknown>>(node: GraphNode<T>, nodeUid: () => string): GraphNodeWithIds<T> {
  return {
    ...node,
    ___uid: `${node.name.replace(/[^a-z]/gi, '').toLowerCase().substring(0, 3)}_${nodeUid()}`,
    ___uidInputs: [],
    ___uidOutputs: []
  }
}

function generateAdditional<T extends Dict<unknown>>(uid: string, type: 'i'|'h'|'o', nodeUid: () => string): GraphNodeWithIds<AdditionalNode> {
  return {
    inputs: type !== 'i' ? [uid] : [],
    outputs: type !== 'o' ? [uid] : [],
    name: uid,
    ___type: type,
    ___uid: `${uid.replace(/[^a-z]/gi, '').toLowerCase().substring(0, 3)}_${nodeUid()}_${type}`,
    ___uidInputs: [],
    ___uidOutputs: []
  }
}

function splitIntoRows<T extends Dict<unknown>>(inputs: string[], nodes: GraphNode<T>[], nodeUid: () => string): [string[][], Dict<GraphNodeWithIds<T>>] {
  const nodesWithId: Dict<GraphNodeWithIds<T>> = {}
  const available = new Set(inputs)
  let queue = [...nodes]
  const rows: string[][] = []
  while (queue.length) {
    const amount = queue.length
    const availableOfRow: string[] = []
    rows.push([])

    queue = queue.filter((node) => {
      const isPlaceable = node.inputs.every(input => available.has(input))
      if (isPlaceable) {
        const nodeWithId: GraphNodeWithIds<T> = generateIds(node, nodeUid)
        rows[rows.length - 1].push(nodeWithId.___uid)
        nodesWithId[nodeWithId.___uid] = nodeWithId
        availableOfRow.push(...node.outputs)
      }
      return !isPlaceable
    })
    availableOfRow.map(uid => available.add(uid))

    if (amount === queue.length) {
      console.warn("Loop detected! Left over:", queue)
      rows.pop()
      break
    }
  }
  return [rows, nodesWithId]
}

function addAdditionalNodes<T extends Dict<unknown>>(rowsRaw: string[][], nodesRaw: Dict<GraphNodeWithIds<T>>, inputs: string[], outputs: string[], nodeUid: () => string): [string[][], Dict<GraphNodeWithIds<T|AdditionalNode>>] {
  const addedInputs: Dict<number> = {}
  const res: string[][] = deepcopy([[], ...rowsRaw, []])
  const resNodes: Dict<GraphNodeWithIds<T|AdditionalNode>> = deepcopy(nodesRaw)
  for (let i = 0; i < rowsRaw.length; i++) {
    const rowInputs = uniquify(rowsRaw[i].flatMap(row => resNodes[row].inputs))
    for (let rowInput of rowInputs) {
      if (rowInput in addedInputs) {
        for (let j = addedInputs[rowInput]+1; j < i+1; j++) {
          const nodeWithId: GraphNodeWithIds<AdditionalNode> = generateAdditional(rowInput, 'h', nodeUid)
          res[j].push(nodeWithId.___uid)
          resNodes[nodeWithId.___uid] = nodeWithId
        }
        addedInputs[rowInput] = i
      } else if (inputs.includes(rowInput)) {
        const nodeWithId: GraphNodeWithIds<AdditionalNode> = generateAdditional(rowInput, 'i', nodeUid)
        res[i].push(nodeWithId.___uid)
        resNodes[nodeWithId.___uid] = nodeWithId
        addedInputs[rowInput] = i
      }
    }
    const rowOutputs = uniquify(rowsRaw[i].flatMap(row => resNodes[row].outputs))
    for (let rowOutput of rowOutputs) {
      if (outputs?.includes(rowOutput)) {
        const nodeWithId: GraphNodeWithIds<AdditionalNode> = generateAdditional(rowOutput, 'o', nodeUid)
        res[i+2].push(nodeWithId.___uid)
        resNodes[nodeWithId.___uid] = nodeWithId
      }
      addedInputs[rowOutput] = i+1
    }
  }
  if (res[res.length-1].length === 0) res.splice(res.length-1, 1)
  return [res, resNodes]
}

function linkNodes<T extends Dict<unknown>>(rowsWithInOut: string[][], nodesWithInOut: Dict<GraphNodeWithIds<T|AdditionalNode>>): Dict<GraphNodeWithIds<T|AdditionalNode>> {
  type Store = {input: Dict<string[]>[], output: Dict<string[]>[]}

  const store: Store = {
    input: Array.from(rowsWithInOut, Object),
    output: Array.from(rowsWithInOut, Object)
  }
  const nodesPremapping = (node: GraphNodeWithIds<T|AdditionalNode>, rowIdx: number) => {
    node.inputs.forEach(input => store.input[rowIdx][input] = [...(store.input[rowIdx][input] ?? []), node.___uid])
    node.outputs.forEach(output => store.output[rowIdx][output] = [...(store.output[rowIdx][output] ?? []), node.___uid])
  };
  const linkInOut = (node: GraphNodeWithIds<T|AdditionalNode>, rowIdx: number) => {
    if (rowIdx > 0)
      node.___uidInputs = uniquify(node.inputs.flatMap(input => store.output[rowIdx - 1][input]).filter(isNonNullable))
    if (rowIdx < rowsWithInOut.length-1)
      node.___uidOutputs = uniquify(node.outputs.flatMap(output => store.input[rowIdx + 1][output]).filter(isNonNullable))
  };

  rowsWithInOut.forEach((row, rowIdx) => row.forEach(uid => nodesPremapping(nodesWithInOut[uid], rowIdx)))
  rowsWithInOut.forEach((row, rowIdx) => row.forEach(uid => linkInOut(nodesWithInOut[uid], rowIdx)))
  return nodesWithInOut
}

function crossingsOf<T extends Dict<unknown>>(fixed: string[], flex: string[], nodes: Dict<GraphNodeWithIds<T|AdditionalNode>>, isDown: boolean): [number[][], number] {
  const result = Array.from(flex, () => Array.from(flex, () => 9999999))
  let score = 0
  for (let i = 0; i < flex.length-1; i++) {
    for (let j = i+1; j < flex.length; j++) {
      const inputsI = new Set(nodes[flex[i]][isDown ? '___uidInputs' : '___uidOutputs'])
      const inputsJ = new Set(nodes[flex[j]][isDown ? '___uidInputs' : '___uidOutputs'])
      const diff = fixed.reduce(([size, acc], curr) => {
        inputsI.has(curr) && size--
        return [size, acc + (inputsJ.has(curr) ? size : 0)]
      }, [inputsI.size, 0])[1] - fixed.reduce(([size, acc], curr) => {
        inputsJ.has(curr) && size--
        return [size, acc + (inputsI.has(curr) ? size : 0)]
      }, [inputsJ.size, 0])[1]
      result[i][j] = diff
      result[j][i] = -diff
      score += diff
    }
  }
  return [result, score]
}

function optimizeOrder<T extends Dict<unknown>>(rowsWithInOut: string[][], nodesWithInOut: Dict<GraphNodeWithIds<T|AdditionalNode>>): [string[][], number] {

  function addCosts(costsUp: [number[][], number], costsDown: [number[][], number]): [number[][], number] {
    return [
      costsUp[0].map((row, rowIdx) => row.map((col, colIdx) => col+costsDown[0][rowIdx][colIdx])),
      costsUp[1] + costsDown[1]
    ]
  }

  function improveRow(top: string[]|undefined, mid: string[], bot: string[]|undefined) {
    const costsUp = top ? crossingsOf(top, mid, nodesWithInOut, true) : undefined
    const costsDown = bot ? crossingsOf(bot, mid, nodesWithInOut, false) : undefined
    const [costs, scoreConst] = costsUp && costsDown ? addCosts(costsUp, costsDown) : costsDown ?? costsUp ?? [undefined, 0]
    let score = scoreConst
    if (!costs) return score
    let improvementInRow = true
    while (improvementInRow) {
      improvementInRow = false
      const colBest = costs
        .map((_, idx) => costs
          .slice(0, idx)
          .map(r => r[idx])
          .reverse()
          .reduce(([sum, best], curr, i) => {
            return sum + curr > best.sum
              ? [sum + curr, {sum: sum+curr, move: -i-1, idx}] as const
              : [sum + curr, best] as const
          }, [0, {sum: -10000, move: 0, idx: 0}] as readonly [number, Reduce])[1]
        )
        .filter(isNonNullable)
      const rowBest = costs
        .map((_, idx) => costs[idx]
          .slice(idx+1)
          .reduce(([sum, best], curr, i) => {
            return sum + curr > best.sum
              ? [sum + curr, {sum: sum+curr, move: i+1, idx}] as const
              : [sum + curr, best] as const
          }, [0, {sum: -10000, move: 0, idx: 0}] as readonly [number, Reduce])[1]
        )
        .filter(isNonNullable)
      const replacement = [...colBest, ...rowBest].sort(sortByProperty(red => -red.sum))[0]
      if (replacement.sum > 0) {
        score -= 2*replacement.sum
        improvement = true
        improvementInRow = true
        costs.splice(replacement.move+replacement.idx, 0, ...costs.splice(replacement.idx, 1))
        costs.map(col => col.splice(replacement.move+replacement.idx, 0, ...col.splice(replacement.idx, 1)))
        mid.splice(replacement.move+replacement.idx, 0, ...mid.splice(replacement.idx, 1))
      }
    }
    return score
  }

  type Reduce = {sum: number, move: number, idx: number}
  let improvement = true
  let scores: number[] = []
  while (improvement) {
    improvement = false
    rowsWithInOut.reduce((top, mid, idx) => {
      scores[idx] = improveRow(top, mid, rowsWithInOut[idx+1])
      return mid
    }, undefined as string[]|undefined)
    rowsWithInOut.reduceRight((bot, mid, idx) => {
      scores[idx] = improveRow(rowsWithInOut[idx-1], mid, bot)
      return mid
    }, undefined as string[]|undefined)
  }
  return [rowsWithInOut, scores.reduce((a, b) => a+b)]
}

export function findBest<T extends Dict<unknown>>(rowsWithInOut: string[][], nodesWithInOut: Dict<GraphNodeWithIds<T|AdditionalNode>>, timeLimit: number) {
  let bestScore = Infinity
  let bestRows = deepcopy(rowsWithInOut)
  let limit = Date.now() + timeLimit
  let iterSinceImprovements = 0
  let rng = seedrandom.alea("We are SEEED, ya!")
  while (true) {
    if (Date.now() > limit) break
    if (iterSinceImprovements > 100) break
    const [rowsOptimized, score] = optimizeOrder(rowsWithInOut, nodesWithInOut)
    if (score < bestScore) {
      bestScore = score
      bestRows = deepcopy(rowsOptimized)
      iterSinceImprovements = 0
    } else {
      iterSinceImprovements++
    }
    rowsOptimized.forEach((row, ) => shuffleInplace(row, rng))
  }
  console.log(bestScore)
  return bestRows
}

export function graphUntangled<T extends Dict<unknown>>(nodes: GraphNode<T>[], inputs: string[], outputs: string[], timeLimit = 0): [string[][], Dict<GraphNodeWithIds<T|AdditionalNode>>] {
  const nodeSeed = seedrandom.alea("node")
  const nodeUid = () => Math.abs(nodeSeed.int32()).toString(16).substring(0, 3)
  //console.log('---------------')
  const [rowsRaw, nodesRaw] = splitIntoRows(inputs, nodes, nodeUid)
  //console.log("Step 1")
  //console.table(rowsRaw)
  //console.table(nodesRaw)
  const [rowsWithInOut, nodesWithInOut] = addAdditionalNodes(rowsRaw, nodesRaw, inputs, outputs, nodeUid)
  //console.log("Step 2")
  //console.table(rowsWithInOut)
  //console.table(nodesWithInOut)
  const nodesLinked = linkNodes(rowsWithInOut, nodesWithInOut)
  //console.log("Step 3")
  //console.table(rowsWithInOut)
  //console.table(nodesLinked)
  const bestRows = findBest(rowsWithInOut, nodesWithInOut, timeLimit)
  return [bestRows, nodesLinked]
}