/**
 * @fileoverview Visualization for span-based annotations.
 *
 * This provides LIT modules for sparse span annotations using the
 * MultiSegmentAnnotations type. Annotations are rendered in-line as highlight
 * spans in running text, which is well-suited for tasks like QA or entity
 * recognition which have a small number of spans over a longer passage.
 *
 * We provide two module classes:
 * - AnnotatedTextGoldModule for gold annotations (in the input data)
 * - AnnotatedTextModule for model predictions
 */

// tslint:disable:no-new-decorators
import '../elements/annotated_text_vis';

import {html} from 'lit';
import {customElement} from 'lit/decorators.js';
import {makeObservable, observable} from 'mobx';

import {LitModule} from '../core/lit_module';
import {type AnnotationGroups, TextSegments} from '../elements/annotated_text_vis';
import {MultiSegmentAnnotations, TextSegment} from '../lib/lit_types';
import {styles as sharedStyles} from '../lib/shared_styles.css';
import {type IndexedInput, ModelInfoMap, Spec} from '../lib/types';
import {doesOutputSpecContain, filterToKeys, findSpecKeys} from '../lib/utils';

/** LIT module for model output. */
@customElement('annotated-text-gold-module')
export class AnnotatedTextGoldModule extends LitModule {
  static override title = 'Annotated Text (gold)';
  static override duplicateForExampleComparison = true;
  static override duplicateForModelComparison = false;
  static override numCols = 4;
  static override template =
      (model: string, selectionServiceIndex: number, shouldReact: number) => {
        return html`
      <annotated-text-gold-module model=${model} .shouldReact=${shouldReact}
        selectionServiceIndex=${selectionServiceIndex}>
      </annotated-text-gold-module>`;
      };

  static override get styles() {
    return sharedStyles;
  }

  override renderImpl() {
    const input = this.selectionService.primarySelectedInputData;
    if (!input) return null;

    const dataSpec = this.appState.currentDatasetSpec;

    // Text segment fields
    const segmentNames = findSpecKeys(dataSpec, TextSegment);
    const segments: TextSegments = filterToKeys(input.data, segmentNames);
    const segmentSpec = filterToKeys(dataSpec, segmentNames);

    // Annotation fields
    const annotationNames = findSpecKeys(dataSpec, MultiSegmentAnnotations);
    const annotations: AnnotationGroups =
        filterToKeys(input.data, annotationNames);
    const annotationSpec = filterToKeys(dataSpec, annotationNames);

    // If more than one model is selected, AnnotatedTextModule will be offset
    // vertically due to the model name header, while this one won't be.
    // So, add an offset so that the content still aligns when there is a
    // AnnotatedTextGoldModule and a AnnotatedTextModule side-by-side.
    const offsetForHeader = !this.appState.compareExamplesEnabled &&
        this.appState.currentModels.length > 1;

    // clang-format off
    return html`
      ${offsetForHeader? html`<div class='offset-for-module-header'></div>` : null}
      <annotated-text-vis .segments=${segments}
                          .segmentSpec=${segmentSpec}
                          .annotations=${annotations}
                          .annotationSpec=${annotationSpec}>
      </annotated-text-vis>`;
    // clang-format on
  }

  static override shouldDisplayModule(
      modelSpecs: ModelInfoMap, datasetSpec: Spec) {
    return findSpecKeys(datasetSpec, MultiSegmentAnnotations).length > 0;
  }
}

/** LIT module for model output. */
@customElement('annotated-text-module')
export class AnnotatedTextModule extends LitModule {
  static override title = 'Annotated Text (predicted)';
  static override duplicateForExampleComparison = true;
  static override duplicateForModelComparison = true;
  static override numCols = 4;
  static override template =
      (model: string, selectionServiceIndex: number, shouldReact: number) => {
        return html`
      <annotated-text-module model=${model} .shouldReact=${shouldReact}
        selectionServiceIndex=${selectionServiceIndex}>
      </annotated-text-module>`;
      };

  static override get styles() {
    return sharedStyles;
  }

  @observable private currentData?: IndexedInput = undefined;
  @observable private currentPreds: AnnotationGroups = {};

  constructor() {
    super();
    // TODO(b/319249148): Assess whether these need to be @observable vs @state.
    makeObservable(this);
  }

  override firstUpdated() {
    const getPrimarySelectedInputData = () =>
        this.selectionService.primarySelectedInputData;
    this.reactImmediately(getPrimarySelectedInputData, data => {
      this.updateToSelection(data);
    });
  }

  private async updateToSelection(input: IndexedInput|null) {
    if (input == null) {
      this.currentData = undefined;
      this.currentPreds = {};
      return;
    }
    // Before waiting for the backend call, update data and clear annotations.
    this.currentData = input;
    this.currentPreds = {};  // empty preds will render as (no data)

    const promise = this.apiService.getPreds(
        [input], this.model, this.appState.currentDataset,
        [MultiSegmentAnnotations], [], 'Retrieving annotations');
    const results = await this.loadLatest('answers', promise);
    if (results === null) return;

    // Update data again, in case selection changed rapidly.
    this.currentData = input;
    this.currentPreds = results[0] as AnnotationGroups;
  }

  override renderImpl() {
    if (!this.currentData) return null;

    const segmentNames =
        findSpecKeys(this.appState.currentDatasetSpec, TextSegment);
    const segments: TextSegments =
        filterToKeys(this.currentData.data, segmentNames);
    const segmentSpec =
        filterToKeys(this.appState.currentDatasetSpec, segmentNames);

    const outputSpec = this.appState.getModelSpec(this.model).output;
    const annotationSpec = filterToKeys(
        outputSpec, findSpecKeys(outputSpec, MultiSegmentAnnotations));
    // clang-format off
    return html`
      <annotated-text-vis .segments=${segments}
                          .segmentSpec=${segmentSpec}
                          .annotations=${this.currentPreds}
                          .annotationSpec=${annotationSpec}>
      </annotated-text-vis>`;
    // clang-format on
  }

  static override shouldDisplayModule(
      modelSpecs: ModelInfoMap, datasetSpec: Spec) {
    return doesOutputSpecContain(modelSpecs, MultiSegmentAnnotations);
  }
}

declare global {
  interface HTMLElementTagNameMap {
    'annotated-text-module': AnnotatedTextModule;
    'annotated-text-gold-module': AnnotatedTextGoldModule;
  }
}
