import * as faceapi from 'face-api.js';
import { FaceMatcher } from 'face-api.js';
import { Doc, DocListCast } from '../../../fields/Doc';
import { DocData } from '../../../fields/DocSymbols';
import { List } from '../../../fields/List';
import { listSpec } from '../../../fields/Schema';
import { Cast, ImageCast, NumCast, StrCast } from '../../../fields/Types';
import { ImageField } from '../../../fields/URLField';
import { DocumentType } from '../../documents/DocumentTypes';
import { Docs } from '../../documents/Documents';
import { DocumentManager } from '../../util/DocumentManager';

/**
 * A singleton class that handles face recognition and manages face Doc collections for each face found.
 * Displaying an image doc anywhere will trigger this class to test if the image contains any faces.
 * If it does, each recognized face will be compared to a stored, global set of faces (each face is represented
 * as a face collection Doc). If the face matches a face collection Doc, then it will be added to that
 * collection along with the numerical representation of the face, its face descriptor.
 *
 * Image Doc's that are added to one or more face collection Docs will be given these metadata fields:
 *     <image data field>_faceDescriptors - list of all the numerical face representations found in the image.
 *     <image data field>_faces - list of unique face Docs corresponding to recognized faces in the image.
 *
 * unique face Doc's are created for each person identified and are stored in the Dashboard's myUniqueFaces field
 *
 * Each unique face Doc represents a unique face and collects all matching face images for that person.  It has these fields:
 *     face - a string label for the person that was recognized                                     (TODO: currently it's just a 'face#')
 *     face_descriptors - a list of all the face descriptors for different images of the person
 *     face_docList - a list of all image Docs that contain a face for the person
 */
export class FaceRecognitionHandler {
    static _instance: FaceRecognitionHandler;
    private _apiModelReady = false;
    private _pendingAPIModelReadyDocs: Doc[] = [];

    public static get Instance() {
        return FaceRecognitionHandler._instance ?? new FaceRecognitionHandler();
    }

    /**
     * Loads an image
     */
    private static loadImage = (imgUrl: ImageField): Promise<HTMLImageElement> => {
        const [name, type] = imgUrl.url.href.split('.');
        const imageURL = `${name}_o.${type}`;

        return new Promise((resolve, reject) => {
            const img = new Image();
            img.crossOrigin = 'anonymous';
            img.onload = () => resolve(img);
            img.onerror = err => reject(err);
            img.src = imageURL;
        });
    };

    /**
     * return the metadata field name where unique face Docs are stored
     * @param imgDoc image with faces
     * @returns name of field
     */
    private static ImageDocFaceField = (imgDoc: Doc) => `${Doc.LayoutFieldKey(imgDoc)}_faces`;

    /**
     * Returns an array of faceDocs for each face recognized in the image
     * @param imgDoc image with faces
     * @returns faceDoc array
     */
    private static ImageDocFaces = (imgDoc: Doc) => DocListCast(imgDoc[`${Doc.LayoutFieldKey(imgDoc)}_faces`]);

    /**
     * initializes an image with an empty list of face descriptors
     * @param imgDoc image to initialize
     */
    private static initImageDocFaceDescriptors = (imgDoc: Doc) => {
        imgDoc[DocData][`${Doc.LayoutFieldKey(imgDoc)}_faceDescriptors`] = new List<List<number>>();
    };
    /**
     * returns the face descriptors for each face found on an image Doc
     * @param imgDoc
     * @returns list of face descriptors
     */
    public static ImageDocFaceDescriptors = (imgDoc: Doc) => imgDoc[DocData][`${Doc.LayoutFieldKey(imgDoc)}_faceDescriptors`] as List<List<number>>;

    /**
     * Adds a face descriptor for a face found in an image
     * @param imgDoc image Doc with face
     * @param faceDescriptor descriptor of a face
     */
    public static ImageDocAddFaceDescriptor = (imgDoc: Doc, faceDescriptor: List<number>) => {
        Cast(imgDoc[DocData][`${Doc.LayoutFieldKey(imgDoc)}_faceDescriptors`], listSpec('number'), null).push(faceDescriptor as unknown as number);
    };

    /**
     * returns a list of all face collection Docs on the current dashboard
     * @returns face collection Doc list
     */
    public static UniqueFaces = () => DocListCast(Doc.ActiveDashboard?.[DocData].myUniqueFaces);

    /**
     * Removes a unique face from the set of recognized unique faces
     * @param faceDoc unique face Doc
     * @returns
     */
    public static DeleteUniqueFace = (faceDoc: Doc) => Doc.ActiveDashboard && Doc.RemoveDocFromList(Doc.ActiveDashboard[DocData], 'myUniqueFaces', faceDoc);

    /**
     * returns the labels associated with a face collection Doc
     * @param faceDoc unique face Doc
     * @returns label string
     */
    public static UniqueFaceLabel = (faceDoc: Doc) => StrCast(faceDoc[DocData].face);

    public static SetUniqueFaceLabel = (faceDoc: Doc, value: string) => (faceDoc[DocData].face = value);
    /**
     * Returns all the face descriptors associated with a unique face Doc
     * @param faceDoc unique face Doc
     * @returns face descriptors
     */
    public static UniqueFaceDescriptors = (faceDoc: Doc) => faceDoc[DocData].face_descriptors as List<List<number>>;

    /**
     * Returns a list of all face image Docs associated with a unique face Doc
     * @param faceDoc unique face Doc
     * @returns image Docs
     */
    public static UniqueFaceImages = (faceDoc: Doc) => DocListCast(faceDoc[DocData].face_images);

    /**
     * Adds a face image to a unique face Doc, adds the unique face Doc to the images list of reognized faces,
     * and updates the unique face's set of face image descriptors
     * @param img - image with faces to add to a face collection Doc
     * @param faceDescriptor - the face descriptor for the face in the image to add
     * @param faceDoc - unique face Doc
     */
    public static UniqueFaceAddFaceImage = (img: Doc, faceDescriptor: List<number>, faceDoc: Doc) => {
        Doc.AddDocToList(faceDoc, 'face_images', img);
        Cast(faceDoc.face_descriptors, listSpec('number'), null).push(faceDescriptor as unknown as number); // items are lists of numbers, not numbers, but type system can't handle that
        Doc.AddDocToList(img[DocData], FaceRecognitionHandler.ImageDocFaceField(img), faceDoc);
    };

    /**
     * Removes a face from a unique Face Doc, and updates the unique face's set of face image descriptors
     * @param img - image with faces to remove
     * @param faceDoc - unique face Doc
     */
    public static UniqueFaceRemoveFaceImage = (img: Doc, faceDoc: Doc) => {
        Doc.RemoveDocFromList(faceDoc[DocData], 'face_images', img);
        const descriptorsEqual = (a: List<number>, b: List<number>) => (a === b ? true : a.length === b.length ? a.every((element, index) => element === b[index]) : false);
        faceDoc[DocData].face_descriptors = new List<List<number>>(FaceRecognitionHandler.UniqueFaceDescriptors(faceDoc).filter(fd => !FaceRecognitionHandler.ImageDocFaceDescriptors(img).some(desc => descriptorsEqual(fd, desc))));
        Doc.RemoveDocFromList(img[DocData], FaceRecognitionHandler.ImageDocFaceField(img), faceDoc);
    };

    constructor() {
        FaceRecognitionHandler._instance = this;
        this.loadAPIModels().then(() => this._pendingAPIModelReadyDocs.forEach(this.classifyFacesInImage));
        DocumentManager.Instance.AddAnyViewRenderedCB(dv => FaceRecognitionHandler.Instance.classifyFacesInImage(dv.Document));
    }

    /**
     * Loads the face detection models.
     */
    private loadAPIModels = async () => {
        const MODEL_URL = `/models`;
        await faceapi.loadFaceDetectionModel(MODEL_URL);
        await faceapi.loadFaceLandmarkModel(MODEL_URL);
        await faceapi.loadFaceRecognitionModel(MODEL_URL);
        this._apiModelReady = true;
    };

    /**
     * Creates a new, empty unique face Doc
     * @returns a unique face Doc
     */
    private createUniqueFaceDoc = (dashboard: Doc) => {
        const faceDocNum = NumCast(dashboard[DocData].myUniqueFaces_count) + 1;
        dashboard[DocData].myUniqueFaces_count = faceDocNum; // TODO: improve to a better name

        const uniqueFaceDoc = Docs.Create.UniqeFaceDocument({
            title: `Face ${faceDocNum}`,
            _layout_reflowHorizontal: true,
            _layout_reflowVertical: true,
            _layout_nativeDimEditable: true,
            _layout_borderRounding: '20px',
            _layout_fitWidth: true,
            _layout_autoHeight: true,
            _face_showImages: true,
            _width: 400,
            _height: 100,
        });
        const uface = uniqueFaceDoc[DocData];
        uface.face = `Face${faceDocNum}`;
        uface.face_images = new List<Doc>();
        uface.face_descriptors = new List<List<number>>();
        Doc.SetContainer(uniqueFaceDoc, Doc.MyFaceCollection);

        Doc.ActiveDashboard && Doc.AddDocToList(Doc.ActiveDashboard[DocData], 'myUniqueFaces', uniqueFaceDoc);
        return uniqueFaceDoc;
    };

    /**
     * Finds the most similar matching Face Document to a face descriptor
     * @param faceDescriptor face descriptor number list
     * @returns face Doc
     */
    private findMatchingFaceDoc = (faceDescriptor: Float32Array) => {
        if (!Doc.ActiveDashboard || FaceRecognitionHandler.UniqueFaces().length < 1) {
            return undefined;
        }

        const faceDescriptors = FaceRecognitionHandler.UniqueFaces().map(faceDoc => {
            const float32Array = FaceRecognitionHandler.UniqueFaceDescriptors(faceDoc).map(fd => new Float32Array(Array.from(fd)));
            return new faceapi.LabeledFaceDescriptors(FaceRecognitionHandler.UniqueFaceLabel(faceDoc), float32Array);
        });
        const faceMatcher = new FaceMatcher(faceDescriptors, 0.6);
        const match = faceMatcher.findBestMatch(faceDescriptor);
        if (match.label !== 'unknown') {
            for (const faceDoc of FaceRecognitionHandler.UniqueFaces()) {
                if (FaceRecognitionHandler.UniqueFaceLabel(faceDoc) === match.label) {
                    return faceDoc;
                }
            }
        }
        return undefined;
    };

    /**
     * When a document is added, this finds faces in the images and tries to
     * match them to existing unique faces, otherwise new unique face(s) are created.
     * @param imgDoc The document being analyzed.
     */
    private classifyFacesInImage = async (imgDoc: Doc) => {
        if (!Doc.UserDoc().recognizeFaceImages) return;
        const activeDashboard = Doc.ActiveDashboard;
        if (!this._apiModelReady || !activeDashboard) {
            this._pendingAPIModelReadyDocs.push(imgDoc);
        } else if (imgDoc.type === DocumentType.LOADING && !imgDoc.loadingError) {
            setTimeout(() => this.classifyFacesInImage(imgDoc), 1000);
        } else {
            const imgUrl = ImageCast(imgDoc[Doc.LayoutFieldKey(imgDoc)]);
            if (imgUrl && !DocListCast(Doc.MyFaceCollection.examinedFaceDocs).includes(imgDoc[DocData])) { // only examine Docs that have an image and that haven't already been examined.
                Doc.AddDocToList(Doc.MyFaceCollection, 'examinedFaceDocs', imgDoc[DocData]);
                FaceRecognitionHandler.initImageDocFaceDescriptors(imgDoc);
                FaceRecognitionHandler.loadImage(imgUrl).then( // load image and analyze faces
                img => faceapi.detectAllFaces(img).withFaceLandmarks().withFaceDescriptors()
                    .then(imgDocFaceDescriptions => { // For each face detected, find a match.
                        for (const fd of imgDocFaceDescriptions) {
                            const faceDescriptor = new List<number>(Array.from(fd.descriptor));
                            FaceRecognitionHandler.ImageDocAddFaceDescriptor(imgDoc, faceDescriptor); // add face descriptor to image's list of descriptors
                            const matchedUniqueFace = this.findMatchingFaceDoc(fd.descriptor) ?? this.createUniqueFaceDoc(activeDashboard);
                            FaceRecognitionHandler.UniqueFaceAddFaceImage(imgDoc, faceDescriptor, matchedUniqueFace); // add image/faceDescriptor to matched unique face
                        } //
                        return imgDocFaceDescriptions;
                    }) 
                );
            } // prettier-ignore
        }
    };
}