path: root/src/client/views/collections/collectionFreeForm/ImageLabelBox.tsx

import { FontAwesomeIcon } from '@fortawesome/react-fontawesome';
import { Colors, IconButton } from '@dash/components';
import similarity from 'compute-cosine-similarity';
import { ring } from 'ldrs';
import 'ldrs/ring';
import { action, computed, makeObservable, observable, reaction } from 'mobx';
import { observer } from 'mobx-react';
import React from 'react';
import { imageUrlToBase64 } from '../../../../ClientUtils';
import { Utils, numberRange } from '../../../../Utils';
import { Doc, NumListCast, Opt } from '../../../../fields/Doc';
import { List } from '../../../../fields/List';
import { ImageCastToNameType, ImageCastWithSuffix } from '../../../../fields/Types';
import { gptGetEmbedding, gptImageLabel } from '../../../apis/gpt/GPT';
import { DocumentType } from '../../../documents/DocumentTypes';
import { Docs } from '../../../documents/Documents';
import { DragManager } from '../../../util/DragManager';
import { SettingsManager } from '../../../util/SettingsManager';
import { SnappingManager } from '../../../util/SnappingManager';
import { ViewBoxBaseComponent } from '../../DocComponent';
import { MainView } from '../../MainView';
import { DocumentView } from '../../nodes/DocumentView';
import { FieldView, FieldViewProps } from '../../nodes/FieldView';
import { OpenWhere } from '../../nodes/OpenWhere';
import './ImageLabelBox.scss';
import { MarqueeOptionsMenu } from './MarqueeOptionsMenu';

export class ImageInformationItem {}

export class ImageLabelBoxData {
    // eslint-disable-next-line no-use-before-define
    static _instance: ImageLabelBoxData;
    @observable _docs: Doc[] = [];
    @observable _labelGroups: string[] = [];

    constructor() {
        makeObservable(this);
        ImageLabelBoxData._instance = this;
    }
    public static get Instance() {
        return ImageLabelBoxData._instance ?? new ImageLabelBoxData();
    }

    @action
    public setData = (docs: Doc[]) => {
        this._docs = docs;
    };

    @action
    addLabel = (labelIn: string) => {
        const label = labelIn.toUpperCase().trim();
        if (label.length > 0) {
            if (!this._labelGroups.includes(label)) {
                this._labelGroups = [...this._labelGroups, label.startsWith('#') ? label : '#' + label];
            }
        }
    };

    @action
    removeLabel = (label: string) => {
        const labelUp = label.toUpperCase();
        this._labelGroups = this._labelGroups.filter(group => group !== labelUp);
    };
}

@observer
export class ImageLabelBox extends ViewBoxBaseComponent<FieldViewProps>() {
    public static LayoutString(fieldKey: string) {
        return FieldView.LayoutString(ImageLabelBox, fieldKey);
    }
    // eslint-disable-next-line no-use-before-define
    public static Instance: ImageLabelBox;

    private _dropDisposer?: DragManager.DragDropDisposer;
    private _inputRef = React.createRef<HTMLInputElement>();
    @observable _loading: boolean = false;
    private _currentLabel: string = '';

    protected createDropTarget = (ele: HTMLDivElement) => {
        this._dropDisposer?.();
        ele && (this._dropDisposer = DragManager.MakeDropTarget(ele, this.onInternalDrop.bind(this), this.layoutDoc));
    };

    protected onInternalDrop(e: Event, de: DragManager.DropEvent): boolean {
        const { docDragData } = de.complete;
        if (docDragData) {
            ImageLabelBoxData.Instance.setData(ImageLabelBoxData.Instance._docs.concat(docDragData.droppedDocuments));
            return false;
        }
        return false;
    }

    @computed get _labelGroups() {
        return ImageLabelBoxData.Instance._labelGroups;
    }

    @computed get _selectedImages() {
        // return DocListCast(this.dataDoc.data);
        return ImageLabelBoxData.Instance._docs;
    }
    @observable _displayImageInformation: boolean = false;

    constructor(props: FieldViewProps) {
        super(props);
        makeObservable(this);
        ring.register();
        ImageLabelBox.Instance = this;
    }

    // ImageLabelBox.Instance.setData()
    /**
     * This method is called when the SearchBox component is first mounted. When the user opens
     * the search panel, the search input box is automatically selected. This allows the user to
     * type in the search input box immediately, without needing clicking on it first.
     */
    componentDidMount() {
        this.classifyImagesInBox();
        reaction(
            () => this._selectedImages,
            () => this.classifyImagesInBox()
        );
    }

    @action
    groupImages = () => {
        this.groupImagesInBox();
    };

    @action
    startLoading = () => {
        this._loading = true;
    };

    @action
    endLoading = () => {
        this._loading = false;
    };

    @action
    toggleDisplayInformation = () => {
        this._displayImageInformation = !this._displayImageInformation;
        if (this._displayImageInformation) {
            this._selectedImages.forEach(doc => (doc._layout_showTags = true));
        } else {
            this._selectedImages.forEach(doc => (doc._layout_showTags = false));
        }
    };

    @action
    submitLabel = () => {
        const input = document.getElementById('new-label') as HTMLInputElement;
        ImageLabelBoxData.Instance.addLabel(this._currentLabel);
        this._currentLabel = '';
        input.value = '';
    };

    onInputChange = action((e: React.ChangeEvent<HTMLInputElement>) => {
        this._currentLabel = e.target.value;
    });

    classifyImagesInBox = async () => {
        this.startLoading();

        // Converts the images into a Base64 format, afterwhich the information is sent to GPT to label them.

        const imageInfos = this._selectedImages.map(async doc => {
            if (!doc.$tags_chat) {
                const url = ImageCastWithSuffix(doc[Doc.LayoutDataKey(doc)], '_o') ?? '';
                return imageUrlToBase64(url).then(hrefBase64 =>
                !hrefBase64 ? undefined : 
                    gptImageLabel(hrefBase64,'Give three labels to describe this image.').then(labels => 
                        ({ doc, labels }))) ; // prettier-ignore
            }
        });

        (await Promise.all(imageInfos)).forEach(imageInfo => {
            if (imageInfo) {
                imageInfo.doc.$tags_chat = (imageInfo.doc.$tags_chat as List<string>) ?? new List<string>();

                const labels = imageInfo.labels.split('\n');
                labels.forEach(label => {
                    const hashLabel =
                        '#' +
                        label
                            .replace(/^\d+\.\s*|-|f\*/, '')
                            .replace(/^#/, '')
                            .trim();
                    (imageInfo.doc.$tags_chat as List<string>).push(hashLabel);
                });
            }
        });

        this.endLoading();
    };

    /**
     * Groups images to most similar labels.
     */
    groupImagesInBox = action(async () => {
        this.startLoading();

        await Promise.all(
            this._selectedImages
                .map(doc => ({ doc, labels: doc.$tags_chat as List<string> }))
                .map(({ doc, labels }) => labels.map((label, index) => gptGetEmbedding(label).then(embedding => (doc[`$tags_embedding_${index + 1}`] = new List<number>(embedding)))))
        );

        const labelToEmbedding = new Map<string, number[]>();
        // Create embeddings for the labels.
        await Promise.all(this._labelGroups.map(async label => gptGetEmbedding(label).then(labelEmbedding => labelToEmbedding.set(label, labelEmbedding))));

        // For each image, loop through the labels, and calculate similarity. Associate it with the
        // most similar one.
        this._selectedImages.forEach(doc => {
            const embedLists = numberRange((doc.$tags_chat as List<string>).length).map(n => Array.from(NumListCast(doc[`$tags_embedding_${n + 1}`])));
            const bestEmbedScore = (embedding: Opt<number[]>) => Math.max(...embedLists.map(l => (embedding && similarity(Array.from(embedding), l)!) || 0));
            const {label: mostSimilarLabelCollect} = 
                this._labelGroups.map(label => ({ label, similarityScore: bestEmbedScore(labelToEmbedding.get(label)) }))
                           .reduce((prev, cur) => cur.similarityScore < 0.3 || cur.similarityScore <= prev.similarityScore ? prev: cur,
                                   { label: '', similarityScore: 0, }); // prettier-ignore
            doc.$data_label = mostSimilarLabelCollect; // The label most similar to the image's contents.
        });

        this.endLoading();

        if (this._selectedImages) {
            MarqueeOptionsMenu.Instance.groupImages();
        }

        MainView.Instance.closeFlyout();
    });

    render() {
        if (this._loading) {
            return (
                <div className="image-box-container" style={{ pointerEvents: 'all', color: SnappingManager.userColor, background: SnappingManager.userBackgroundColor }}>
                    <l-ring size="60" color="white" />
                </div>
            );
        }

        if (this._selectedImages.length === 0) {
            return (
                <div className="searchBox-container" style={{ pointerEvents: 'all', color: SnappingManager.userColor, background: SnappingManager.userBackgroundColor }} ref={ele => this.createDropTarget(ele!)}>
                    <p style={{ fontSize: 'large' }}>In order to classify and sort images, marquee select the desired images and press the &apos;Classify and Sort Images&apos; button. Then, add the desired groups for the images to be put in.</p>
                </div>
            );
        }

        return (
            <div className="searchBox-container" style={{ pointerEvents: 'all', color: SnappingManager.userColor, background: SnappingManager.userBackgroundColor }} ref={ele => this.createDropTarget(ele!)}>
                <div className="searchBox-bar" style={{ pointerEvents: 'all', color: SnappingManager.userColor, background: SnappingManager.userBackgroundColor }}>
                    <IconButton
                        tooltip={'See image information'}
                        onPointerDown={this.toggleDisplayInformation}
                        icon={this._displayImageInformation ? <FontAwesomeIcon icon="caret-up" /> : <FontAwesomeIcon icon="caret-down" />}
                        color={MarqueeOptionsMenu.Instance.userColor}
                        style={{ width: '19px' }}
                    />
                    <input
                        defaultValue=""
                        autoComplete="off"
                        onChange={this.onInputChange}
                        onKeyDown={e => {
                            e.key === 'Enter' ? this.submitLabel() : null;
                            e.stopPropagation();
                        }}
                        type="text"
                        placeholder="Input groups for images to be put into..."
                        aria-label="label-input"
                        id="new-label"
                        className="searchBox-input"
                        style={{ width: '100%', borderRadius: '5px' }}
                        ref={this._inputRef}
                    />
                    <IconButton
                        tooltip={'Add a label'}
                        onPointerDown={() => {
                            const input = document.getElementById('new-label') as HTMLInputElement;
                            ImageLabelBoxData.Instance.addLabel(this._currentLabel);
                            this._currentLabel = '';
                            input.value = '';
                        }}
                        icon={<FontAwesomeIcon icon="plus" />}
                        color={MarqueeOptionsMenu.Instance.userColor}
                        style={{ width: '19px' }}
                    />
                    {this._labelGroups.length > 0 ? <IconButton tooltip={'Group Images'} onPointerDown={this.groupImages} icon={<FontAwesomeIcon icon="object-group" />} color={Colors.MEDIUM_BLUE} style={{ width: '19px' }} /> : <div></div>}
                </div>
                <div>
                    <div className="image-label-list">
                        {this._labelGroups.map(group => {
                            return (
                                <div key={Utils.GenerateGuid()}>
                                    <p style={{ color: MarqueeOptionsMenu.Instance.userColor }}>{group}</p>
                                    <IconButton
                                        tooltip={'Remove Label'}
                                        onPointerDown={() => {
                                            ImageLabelBoxData.Instance.removeLabel(group);
                                        }}
                                        icon={'x'}
                                        color={MarqueeOptionsMenu.Instance.userColor}
                                        style={{ width: '8px' }}
                                    />
                                </div>
                            );
                        })}
                    </div>
                </div>
                {this._displayImageInformation ? (
                    <div className="image-information-list">
                        {this._selectedImages.map(doc => {
                            const [name, type] = ImageCastToNameType(doc[Doc.LayoutDataKey(doc)]);
                            return (
                                <div className="image-information" style={{ borderColor: SettingsManager.userColor }} key={Utils.GenerateGuid()}>
                                    <img
                                        src={`${name}_o.${type}`}
                                        onClick={async () => {
                                            await DocumentView.showDocument(doc, { willZoomCentered: true });
                                        }}></img>
                                    <div className="image-information-labels" onClick={() => this._props.addDocTab(doc, OpenWhere.addRightKeyvalue)}>
                                        {(doc.$tags_chat as List<string>).map(label => {
                                            return (
                                                <div key={Utils.GenerateGuid()} className="image-label" style={{ backgroundColor: SettingsManager.userVariantColor, borderColor: SettingsManager.userColor }}>
                                                    {label}
                                                </div>
                                            );
                                        })}
                                    </div>
                                </div>
                            );
                        })}
                    </div>
                ) : (
                    <div></div>
                )}
            </div>
        );
    }
}

Docs.Prototypes.TemplateMap.set(DocumentType.IMAGEGROUPER, {
    layout: { view: ImageLabelBox, dataField: 'data' },
    options: { acl: '', _width: 400 },
});