import * as grok from 'datagrok-api/grok';
import * as ui from 'datagrok-api/ui';
import * as DG from 'datagrok-api/dg';

import $ from 'cash-dom';
import {GridNeighbor} from '@datagrok-libraries/gridext/src/ui/GridNeighbor';
import {TreeHelper} from '../../src/utils/tree-helper';
import {GridTreeRendererBase} from './tree-renderers/grid-tree-renderer-base';
import {LeafRangeGridTreeRenderer} from '../../src/viewers/tree-renderers/grid-tree-renderer';
import {DistanceMetric, NodeCuttedType, NodeType} from '@datagrok-libraries/bio/src/trees';
import {TreeCutOptions} from '@datagrok-libraries/bio/src/trees/dendrogram';
import {markupNode, MarkupNodeType} from './tree-renderers/markup';
import {attachDivToGrid} from '../utils';
import {
  PROPS as _D_PROPS,
  PROPS_CATS as _D_PROPS_CATS,
} from './dendrogram';
import {RectangleTreeHoverType} from './tree-renderers/rectangle-tree-placer';
import {GridTreePlacer} from './tree-renderers/grid-tree-placer';
import {Unsubscribable} from 'rxjs';
import {ITreeHelper} from '@datagrok-libraries/bio/src/trees/tree-helper';

//@ts-ignore
import '../css/injected-dendrogram.css';

/** Options enabling the "mark cluster representatives (medoids)" feature in the Assign Clusters
 * dialog. `colNames` are the feature columns the tree was built on; `distance` is the metric used. */
export type MedoidOptions = {colNames: string[], distance: DistanceMetric};

export function injectTreeForGridUI2(
  grid: DG.Grid, treeRoot: NodeType | null, leafColName?: string, neighborWidth: number = 100,
  cut?: TreeCutOptions, medoidOptions?: MedoidOptions
): GridNeighbor {
  const th: ITreeHelper = new TreeHelper();
  const treeNb: GridNeighbor = attachDivToGrid(grid, neighborWidth);
  if (treeRoot) markupNode(treeRoot);
  const totalLength: number = treeRoot ? (treeRoot as MarkupNodeType).subtreeLength! : 1;
  if (Number.isNaN(totalLength))
    throw new Error('Can not calculate totalLength for the tree.');

  const placer: GridTreePlacer<MarkupNodeType> = new GridTreePlacer<MarkupNodeType>(grid, totalLength);

  const renderer: GridTreeRendererBase<MarkupNodeType> =
    LeafRangeGridTreeRenderer.create(grid, treeRoot, placer);
  renderer.attach(treeNb.root!);

  renderer.onAfterRender.subscribe(({target, context, lengthRatio}) => {
    if (cutSlider) {
      const tgt = target as GridTreeRendererBase<MarkupNodeType>;

      cutSlider.root.style.left = `${0}px`;
      cutSlider.root.style.width = `${tgt.view!.clientWidth}px`;
      cutSlider.root.style.height = `${grid.colHeaderHeight}px`;

      const posX = cutSlider.value! * lengthRatio + tgt.leftPadding * window.devicePixelRatio;

      context.strokeStyle = '#A00000';
      context.moveTo(posX, 0);
      context.lineTo(posX, context.canvas.height);
      context.stroke();
    }
  });

  function assingClusters() {
    if (!treeRoot)
      return;
    showClusterAsignmentDialog(treeRoot, th, grid.dataFrame, renderer, medoidOptions);
  }

  treeNb.root?.addEventListener('contextmenu', (ev) => {
    ev.stopImmediatePropagation();
    ev.preventDefault();
    const menu = DG.Menu.popup();
    menu.item('Assign Clusters', () => { assingClusters(); });
    menu.item('Reset Zoom', () => renderer.onResetZoom());
    menu.show({x: ev.clientX, y: ev.clientY, causedBy: ev});
  });

  const assignClustersButton = ui.iconFA('magic', assingClusters, 'Assign Clusters');
  assignClustersButton.classList.add('dendrogram-assign-clusters-bttn');
  treeNb.root?.appendChild(assignClustersButton);


  let cutSlider: DG.InputBase<number | null> | null = null;
  if (cut) {
    // TODO: Get max from tree height
    //@ts-ignore
    const treeHeight: number = (treeRoot as MarkupNodeType).subtreeLength;
    console.log(treeHeight);
    cutSlider = ui.input.slider('', {value: 0, min: 0, max: treeHeight});
    $(cutSlider.root).find('input').each((_, el) => {
      el.setAttribute('step', '0.01');
      el.style.width = '100%';
      el.style.height = `100%`;
    });
    cutSlider.root.style.position = 'absolute';
    cutSlider.root.style.top = `${0}px`;
    cutSlider.root.style.backgroundColor = '#FFF0F0';

    treeNb.root!.appendChild(cutSlider.root);

    cutSlider.onChanged.subscribe((value) => {
      // console.debug('Dendrogram: injectTreeToGrid() cutSlider.onChanged() ' + `${cutSlider!.value}`);
      // th.cutTreeToGrid(newickRoot, cutSlider!.value!, grid.dataFrame, leafColName, 'Cluster');

      const t1 = Date.now();
      th.treeCutAsTree(treeRoot, value!, true);
      const t2 = Date.now();
      console.debug('Dendrogram: injectTreeForGrid() cutSlider.onChanged() treeCutAsTree() ' +
        `ET: ${((t2 - t1) / 1000).toString()}`);

      const newickRootCopy = JSON.parse(JSON.stringify(treeRoot));
      const newickRootCutted = th.treeCutAsTree(newickRootCopy, value!);
      th.markClusters(newickRootCutted as NodeCuttedType, cut.dataDf, leafColName ?? null, cut.clusterColName);
      th.buildClusters(newickRootCutted as NodeCuttedType, cut.clusterDf, cut.clusterColName, leafColName);

      markupNode(newickRootCutted!);
      renderer.treeRoot = newickRootCutted as MarkupNodeType;
    });
  }

  function alignGridWithTree(): void {
    // grid.dataFrame is null once the table view is closed; deferred callers
    // (onFilterChanged / onRowsResized via setTimeout) may run after that.
    if (!grid.dataFrame)
      return;
    const [viewedRoot] = th.setGridOrder(treeRoot, grid, leafColName);
    if (viewedRoot) markupNode(viewedRoot);

    renderer.treeRoot = viewedRoot as MarkupNodeType;
  }

  // initial alignment tree with grid
  alignGridWithTree();

  // -- Handling events --

  // When a node is clicked, all of the leaf nodes come from it should become selected
  function rendererOnCurrentChanged() {
    window.setTimeout(() => {
      if (!renderer || !placer) return;

      // Reset the current row index
      // grid.dataFrame.currentRowIdx = -1;
      const selectionIndexes = new Set<number>();
      // only one current becomes first of sub leaves
      const selectionLeafs: NodeType[] | null = renderer.current ? th.getLeafList(renderer.current.node) : null;
      if (selectionLeafs && selectionLeafs.length > 0) {
        const selectionLeafNames = new Set<string>(selectionLeafs.map((leaf) => leaf.name));
        if (leafColName) {
          const leafCol: DG.Column = grid.dataFrame.getCol(leafColName);
          const rowCount = grid.dataFrame.rowCount;
          for (let rowI: number = 0; rowI < rowCount; rowI++) {
            const rowLeafName: string = leafCol.get(rowI);
            if (selectionLeafNames.has(rowLeafName)) {
              selectionIndexes.add(rowI);
              break;
            }
          }
        } else {
          selectionLeafs.forEach((leaf) => {
            selectionIndexes.add(parseInt(leaf.name));
          });
        }
      }
      if (selectionIndexes.size === 1) {
        grid.dataFrame.currentRowIdx = selectionIndexes.values().next().value!;
      } else {
        grid.dataFrame.selection.init((i) => selectionIndexes.has(i));
        grid.dataFrame.selection.fireChanged();
      }
      grid.invalidate(); // fixing stall current on changed currentRowIdx to -1
    });
  }

  function rendererOnMouseOverChanged() {
    window.setTimeout(() => {
      if (!renderer || !placer) return;

      const oldMouseOverRowIdx: number = grid.dataFrame.mouseOverRowIdx;
      // only one mouseOver becomes first of sub leaves
      const mouseOverLeaf: NodeType | null = renderer.mouseOver ? th.getLeafList(renderer.mouseOver.node)[0] : null;
      let newMouseOverRowIdx: number = -1;
      if (mouseOverLeaf) {
        if (leafColName) {
          const leafCol: DG.Column = grid.dataFrame.getCol(leafColName);
          const rowCount = grid.dataFrame.rowCount;
          for (let rowI = 0; rowI < rowCount; rowI++) {
            const rowLeafName: string = leafCol.get(rowI);
            if (rowLeafName == mouseOverLeaf.name) {
              newMouseOverRowIdx = rowI;
              break;
            }
          }
        } else {
          newMouseOverRowIdx = parseInt(mouseOverLeaf.name);
        }
      }
      if (newMouseOverRowIdx != oldMouseOverRowIdx)
        grid.dataFrame.mouseOverRowIdx = newMouseOverRowIdx;
    });
  }

  function rendererOnSelectionChanged() {
    window.setTimeout(() => {
      if (!renderer || !placer) return;

      /* Here we get selected rows from dataFrame (leaves only).
       * Some of selected nodes can be in subtree of renderer.selections nodes.
       * We need to merge nodes to form selections object.
       * Nodes subs can be selected or deselected.
       */

      const oldSelection: DG.BitSet = grid.dataFrame.selection.clone();
      if (renderer.selections.length == 0) {
        grid.dataFrame.selection.init((_) => { return false; }, false);
      } else {
        const leafCol: DG.Column | null = !!leafColName ? grid.dataFrame.getCol(leafColName) : null;
        const nodeNameSet = new Set(
          renderer.selectedNodes
            .map((sn) => th.getNodeList(sn).map((n) => n.name))
            .flat());
        // console.debug('Dendrogram: Dendrogram.rendererOnSelectionChanged(), ' +
        //   `nodeNameSet = ${JSON.stringify([...nodeNameSet])}`);

        grid.dataFrame.selection.init(
          (rowI) => {
            const nodeName = !!leafCol ? leafCol.get(rowI) : `${rowI}`;
            return nodeNameSet.has(nodeName);
          },
          false);
      }

      const newSelection: DG.BitSet = grid.dataFrame.selection;
      let selectionChanged: boolean = oldSelection.length !== newSelection.length || // != -> true (changed)
        oldSelection.trueCount !== newSelection.trueCount;
      if (!selectionChanged) {
        for (let rowI: number = 0; rowI < oldSelection.length; rowI++) {
          if (oldSelection.get(rowI) !== newSelection.get(rowI)) {
            selectionChanged = true;
            break;
          }
        }
      }
      if (selectionChanged)
        grid.dataFrame.selection.fireChanged();
    }, 0 /* next event cycle*/);
  }

  function dataFrameOnCurrentRowChanged(_value: any) {
    const leafCol: DG.Column | null = !!leafColName ? grid.dataFrame.getCol(leafColName) : null;
    const idx: number = grid.dataFrame.currentRowIdx;
    const currentLeafName: string | null = idx == -1 ? null : !!leafCol ? leafCol.get(idx) : `${idx}`;

    const th: ITreeHelper = new TreeHelper();
    const currentLeaf: MarkupNodeType | null = !renderer.treeRoot ? null : th.getNodeList(renderer.treeRoot)
      .find((leaf) => currentLeafName == leaf.name) ?? null;
    const current: RectangleTreeHoverType<MarkupNodeType> | null = currentLeaf ? {
      node: currentLeaf,
      nodeHeight: placer!.getNodeHeight(renderer.treeRoot, currentLeaf)!,
    } : null;

    renderer.current = current;
  }

  function dataFrameOnMouseOverRowChanged(_value: any) {
    if (!renderer || !placer) return;

    const leafCol: DG.Column | null = !!leafColName ? grid.dataFrame.getCol(leafColName) : null;
    const idx: number = grid.dataFrame.mouseOverRowIdx;
    const mouseOverLeafName: string | null = idx == -1 ? null : !!leafCol ? leafCol.get(idx) : `${idx}`;

    const th: ITreeHelper = new TreeHelper();
    const mouseOverLeaf: MarkupNodeType | null = th.getLeafList(renderer.treeRoot)
      .find((leaf) => mouseOverLeafName == leaf.name) ?? null;
    const mouseOver: RectangleTreeHoverType<MarkupNodeType> | null = mouseOverLeaf ? {
      node: mouseOverLeaf,
      nodeHeight: placer!.getNodeHeight(renderer.treeRoot, mouseOverLeaf)!,
    } : null;

    renderer.mouseOver = mouseOver;
  }

  function dataFrameOnSelectionChanged(_value: any) {
    if (!renderer || !placer) return;

    const leafList: MarkupNodeType[] = th.getLeafList(renderer.treeRoot);
    const leafDict: { [name: string]: MarkupNodeType } = {};
    for (const node of leafList) {
      if (node.name in leafDict)
        throw new Error('Non unique key tree node name');
      leafDict[node.name] = node;
    }

    const rowDict: { [name: string]: number } = {};
    const leafCol: DG.Column | null = leafColName ? grid.dataFrame.getCol(leafColName) : null;
    const rowCount: number = grid.dataFrame.rowCount;
    for (let rowI = 0; rowI < rowCount; rowI++) {
      const leafName: string = leafCol ? leafCol.get(rowI) : `${rowI}`;
      rowDict[leafName] = rowI;
    }

    const selLeafNames: { [name: string]: number } = {};
    const selIndexes = grid.dataFrame.selection.getSelectedIndexes();
    for (const selRowI of selIndexes) {
      const leafName: string = leafCol ? leafCol.get(selRowI) : `${selRowI}`;
      selLeafNames[leafName] = selRowI;
    }
    const selNodeList: MarkupNodeType[] = th.getNodesByLeaves<MarkupNodeType>(renderer.treeRoot, selLeafNames);

    const selections: RectangleTreeHoverType<MarkupNodeType>[] = [];
    for (const selNode of selNodeList)
      selections.push({node: selNode, nodeHeight: placer.getNodeHeight(renderer.treeRoot, selNode)!});

    renderer.selections = selections;
  }

  // Variable to track if filter is changed and prevent the sorting change event
  let filterChangeCounter = 0;

  function dataFrameOnFilterChanged(_value: any) {
    // TODO: Filter newick tree
    console.debug('Dendrogram: injectTreeForGridUI2() grid.dataFrame.onFilterChanged()');
    filterChangeCounter += 1;
    // to prevent nested fire event in event handler
    window.setTimeout(() => {
      alignGridWithTree();
    }, 0);
  }

  function dfOnSortingChanged(_value?: any) {
    // If the reordering is caused by the filter change, return
    if (filterChangeCounter > 0) {
      filterChangeCounter -= 1;
      return;
    }
    const treeOverlay = ui.div();
    treeOverlay.style.width = treeNb.root!.style.width;
    treeOverlay.style.height = treeNb.root!.style.height;
    treeOverlay.classList.add('dendrogram-overlay');

    const sortInfoDiv = ui.div('Revert columns sort order to see Dendrogram Tree');
    const realignButton = ui.button('Revert sort', () => {
      alignGridWithTree();

      treeNb?.root?.removeChild(treeOverlay);
      sortingSub = grid.onRowsSorted.subscribe(dfOnSortingChanged);
    });

    const infoContainer = ui.divV(
      [sortInfoDiv, realignButton]
    );

    treeOverlay.appendChild(infoContainer);
    treeNb.root?.appendChild(treeOverlay);
    sortingSub.unsubscribe();
  }

  function treeNeighborOnClosed() {
    for (const sub of subs)
      sub.unsubscribe();
  }

  const subs: Unsubscribable[] = [];
  subs.push(treeNb.onClosed.subscribe(treeNeighborOnClosed));
  subs.push(renderer.onCurrentChanged.subscribe(rendererOnCurrentChanged));
  subs.push(renderer.onMouseOverChanged.subscribe(rendererOnMouseOverChanged));
  subs.push(renderer.onSelectionChanged.subscribe(rendererOnSelectionChanged));

  let sortingSub = grid.onRowsSorted.subscribe(dfOnSortingChanged);
  subs.push(sortingSub);

  subs.push(grid.onRowsResized.subscribe(dataFrameOnFilterChanged));
  subs.push(grid.dataFrame.onCurrentRowChanged.subscribe(dataFrameOnCurrentRowChanged));
  subs.push(grid.dataFrame.onMouseOverRowChanged.subscribe(dataFrameOnMouseOverRowChanged));
  subs.push(grid.dataFrame.onSelectionChanged.subscribe(dataFrameOnSelectionChanged));
  subs.push(grid.dataFrame.onFilterChanged.subscribe(dataFrameOnFilterChanged));

  return treeNb;
}

function showClusterAsignmentDialog(
  treeRoot: NodeType, th: ITreeHelper, dataFrame: DG.DataFrame, renderer: GridTreeRendererBase<MarkupNodeType>,
  medoidOptions?: MedoidOptions
) {
  const dialog = ui.dialog('Assign Clusters');
  const treeHeight = (treeRoot as MarkupNodeType).subtreeLength!;
  if (!treeHeight)
    throw new Error('Can not calculate tree height.');

  const cutSlider = ui.input.float('Threshold', {value: treeHeight / 2, min: 0, max: treeHeight,
    showSlider: true, tooltipText: 'Cutting threshold for clusters (height of dendrogram from the root)'});
  const clusterInput = ui.input.int('Clusters', {value: 1, min: 1, tooltipText: 'Number of clusters after cutting'});
  const medoidInput = ui.input.bool('Medoid columns', {value: true,
    tooltipText: 'Add per-cluster medoid rank (1 = most representative) and average distance-to-cluster columns'});
  let processing = false;

  const getClusterCount = (threshold: number) => {
    return th.treeCutAsLeaves(treeRoot, threshold).length;
  };

  const searchForThresholdGivenClusterNum = (clusterNum: number): number => {
    let low = 0;
    let high = treeHeight;
    let bestThreshold = low;
    let minDiff = Number.MAX_VALUE;

    for (let i = 0; i < 20; i++) {
      const mid = (low + high) / 2;
      const count = getClusterCount(mid);

      if (count === clusterNum) return mid;

      const diff = Math.abs(count - clusterNum);
      if (diff < minDiff) {
        minDiff = diff;
        bestThreshold = mid;
      }

      if (count < clusterNum)
        low = mid;
      else
        high = mid;
    }
    return bestThreshold;
  };

  const subs: Unsubscribable[] = [];

  subs.push(cutSlider.onChanged.subscribe((value) => {
    if (processing || value === null) return;
    processing = true;
    try {
      clusterInput.value = getClusterCount(value);
    } finally {
      processing = false;
    }
  }));

  subs.push(DG.debounce(cutSlider.onChanged, 20).subscribe(() => {
    renderer.render('Invalidate');
  }));

  subs.push(clusterInput.onChanged.subscribe((value) => {
    if (processing || value === null) return;
    processing = true;
    try {
      cutSlider.value = searchForThresholdGivenClusterNum(value);
    } finally {
      processing = false;
    }
  }));

  dialog.add(cutSlider.root);
  dialog.add(clusterInput.root);
  if (medoidOptions)
    dialog.add(medoidInput.root);
  cutSlider.fireChanged();

  dialog.addButton('Assign', async () => {
    const threshold = cutSlider.value ?? 0;
    const clusters = th.treeCutAsLeaves(treeRoot, threshold);
    const colName = dataFrame.columns.getUnusedName(`Cluster (${threshold.toFixed(2)})`);
    const clusterCol = DG.Column.fromType(DG.TYPE.STRING, colName, dataFrame.rowCount);

    // collect the row indices of each cluster while writing the cluster id column
    const clusterRowIndexes: number[][] = [];
    clusters.forEach((node, idx) => {
      const clusterId = (idx + 1).toString();
      const leaves = th.getLeafList(node);
      const rows: number[] = [];
      for (const leaf of leaves) {
        const row = parseInt(leaf.name);
        if (row !== undefined && !Number.isNaN(row)) {
          clusterCol.set(row, clusterId, false);
          rows.push(row);
        }
      }
      clusterRowIndexes.push(rows);
    });
    const categoryOrder = clusters.map((_, idx) => (idx + 1).toString());
    clusterCol.setCategoryOrder(categoryOrder);

    clusterCol.setTag(DG.TAGS.DESCRIPTION,
      `Cluster assignment from cutting the dendrogram at height ${threshold.toFixed(2)}.`);
    clusterCol.fireValuesChanged();
    dataFrame.columns.add(clusterCol);

    if (medoidOptions && medoidInput.value) {
      const pi = DG.TaskBarProgressIndicator.create('Finding cluster representatives ...');
      try {
        const treeHelper = new TreeHelper();
        const {rankByRow, avgDistByRow} = await treeHelper.calcMedoids(
          dataFrame, medoidOptions.colNames, clusterRowIndexes, medoidOptions.distance);
        const rankColName = dataFrame.columns.getUnusedName(`Medoid Rank (${threshold.toFixed(2)})`);
        const rankCol = DG.Column.fromList(DG.TYPE.INT, rankColName, rankByRow);
        rankCol.setTag(DG.TAGS.DESCRIPTION,
          'Representativeness rank within the cluster (1 = medoid, the most central member). ' +
          'Filter to rank ≤ N to extract the top-N representatives per cluster.');
        dataFrame.columns.add(rankCol);
        const distColName = dataFrame.columns.getUnusedName(`Avg Distance to Cluster (${threshold.toFixed(2)})`);
        const distCol = DG.Column.fromList(DG.TYPE.FLOAT, distColName, avgDistByRow);
        distCol.setTag(DG.TAGS.DESCRIPTION,
          'Mean distance from this row to the other members of its cluster (lower = more central).');
        dataFrame.columns.add(distCol);
      } catch (e) {
        grok.shell.error('Failed to compute cluster representatives. See console for details.');
        console.error(e);
      } finally {
        pi.close();
      }
    }
    dialog.close();
  });

  dialog.show();

  subs.push(renderer.onAfterRender.subscribe(({target, context, lengthRatio}) => {
    const tgt = target as GridTreeRendererBase<MarkupNodeType>;

    // const drawCanvasSize = context.canvas.width - (tgt.leftPadding + tgt.rightPadding) * window.devicePixelRatio;
    // const actRatio = renderer.treeRoot?.subtreeLength ? drawCanvasSize / renderer.treeRoot!.subtreeLength : lengthRatio;
    const posX = tgt.treeXToCanvasX(cutSlider.value!) * window.devicePixelRatio;

    context.strokeStyle = '#A00000';
    context.moveTo(posX, 0);
    context.lineTo(posX, context.canvas.height);
    context.stroke();
  }));

  const closeSub = dialog.onClose.subscribe(() => {
    for (const sub of subs)
      sub.unsubscribe();
    closeSub.unsubscribe();
  });
}