import { generate } from "shortid";
import Ngram from "./Ngram";
import Sequence from "./Sequence";
import Stimulus from "./Stimulus";
import Tree from "./Tree";

const MAX_DEPTH = 10;

export default class Model  {

    get id(): string {
        return this._id;
    }

    get depth(): number {
        return this._depth;
    }

    get interpolationfactor(): number {
        return this._interpolationfactor;
    }

    get stimulusOccurence(): number {
        return this._stimulusOccurence;
    }

    private _id: string;
    private readonly _depth: number;
    private readonly _interpolationfactor: number;
    private _sessionIdList: string[];
    private _treeMap: Map <string, Tree>;
    private _stimulusOccurence: number;

    constructor(depth: number, interpolationfactor: number= 2, id?: string) {
        if (id) {
            this._id = id;
        } else {
            this._id = generate();
        }

        if (depth <= 0 || depth > MAX_DEPTH) {
            throw new Error("cannot create model, depth between 1 and " + MAX_DEPTH);
        }
        this._depth = depth;

        this._interpolationfactor = interpolationfactor;

        this._sessionIdList = [];
        this._treeMap = new Map();
        this._stimulusOccurence = 0;
    }

    // For testing purpose
    public getTreeByStimulus(stimulus: Stimulus): Tree {
        return this._treeMap.get(stimulus.value);
    }

    public linkWithSession(sessionId: string): void {
        if (!this._sessionIdList.includes(sessionId)) {
            this._sessionIdList.push(sessionId);
        }
    }

    public getLinkedSessionIdList(): string[] {
        return this._sessionIdList.slice(0, this._sessionIdList.length);
    }

    public learnSequence(sequence: Sequence): void {
        while (sequence.length > 0) {
            const context = this.fitContextToDepth(sequence.getContext());
            if (context.length === 0) {
                return;
            }
            const [prefix, last] = sequence.cloneAndPop();
            if (last instanceof Note) {
                this.addNoteKnowingContext(last, this.fitContextToDepth(prefix.getContext()));
            } else if (last instanceof Stimulus) {
                this.contextOccurs(context);
                this.addStimulusKnowingContext(last, this.fitContextToDepth(prefix.getContext()));
            }
            sequence = prefix;
        }
    }

    public getNoteDistributionListMap(sequence: Sequence): Map<string, NoteDistribution[]> {
        if (sequence.length === 0) {
            throw new Error("cannot getNoteDistributionListMap, sequence's context is empty");
        }
        let context = sequence.getContext();
        const lastStimulusdOfContext = context[context.length - 1];

        const distributionListMap: Map<string, NoteDistribution[]> = new Map();
        if (!this._treeMap.has(lastStimulusdOfContext.value)) {
            return distributionListMap;
        }

        while (context.length > 0) {
            const tree = this._treeMap.get(lastStimulusdOfContext.value);
            const distributionForContext: Map<string, NoteDistribution> = tree.getNoteDistributionsFromPrefixes(context);

            distributionForContext.forEach((distribution, note) => {
                if (!distributionListMap.has(note)) {
                    distributionListMap.set(note, [distribution]);
                } else {
                    distributionListMap.get(note).push(distribution);
                }
            });
            context = context.slice(1, context.length);
        }
        return distributionListMap;
    }

    public getStimulusProbabilityMap(context: Stimulus[]): Map<string, number> {
        if (context.length === 0) {
            throw new Error("cannot getStimulusProbabilityMap, sequence's context is empty");
        }
        const lastStimulusdOfContext = context[context.length - 1];
        if (! this._treeMap.has(lastStimulusdOfContext.value)) {
            return new Map();
        } else {
            const tree = this._treeMap.get(lastStimulusdOfContext.value);
            return tree.getStimulusInterpolatedProbabilityMap(context, this._interpolationfactor);
        }
    }

    /*
    Starting from a context, build sequences for each possible next event, and return the built sequences and their probabilities
    */
    public getSequenceWithProbability(context: Stimulus[]): Array<{sequence: Sequence, probability: number}> {
        const sequenceWithProbabilities = [];
        if (context.length === 0) {
            throw new Error("cannot getStimulusProbabilityMapForExactContext, sequence's context is empty");
        }
        const lastStimulusdOfContext = context[context.length - 1];
        if (!this._treeMap.has(lastStimulusdOfContext.value)) {
            return sequenceWithProbabilities;
        } else {
            const tree = this._treeMap.get(lastStimulusdOfContext.value);
            const successors = tree.successorStimulusSet;

            for (const successor of successors) {
                const probabilities = tree.getStimulusProbabilityList(successor, context);
                for (let contextSize = 0; contextSize < probabilities.length; ++contextSize) {
                    const probability = probabilities[contextSize];
                    const stimuli = context.slice(context.length - contextSize, context.length);
                    stimuli.push(successor);
                    sequenceWithProbabilities.push({
                        sequence: new Sequence(stimuli),
                        probability,
                    });
                }
            }
            return sequenceWithProbabilities;
        }
    }

    public crossEntropy(context: Stimulus[]): number {
        const probaOfUnknown = 1e-6;

        if (context.length === 0) {
            return probaOfUnknown;
        }

        let probabilitySum = 0;
        for (let index = context.length - 1; index > 0; index--) {
            const lastStimulus = context[index];
            const contextSequence = context.slice(0, index);
            const probabilityMap = this.getStimulusProbabilityMap(contextSequence);

            const modelProba = probabilityMap.get(lastStimulus.value);
            let proba;
            if (modelProba === null || modelProba === undefined || modelProba === 0) {
                proba = probaOfUnknown;
            } else {
                proba = modelProba * (1 - probaOfUnknown);
            }
            probabilitySum = probabilitySum + Math.log2(proba);
        }
        return -(probabilitySum / (context.length - 1));
    }

    public getAllNgram(): Ngram[] {
        const ngramSet = [];
        for (const tree of this._treeMap.values()) {
            tree.getAllNgram().forEach( (ngram) => {
                ngramSet.push(ngram);
            });
        }
        return ngramSet;
    }

    public sortNgram(ngrams): Ngram[][] {
        const ngramsPerDepth = [];
        for (let i = 0; i < this.depth; ++i) {
            ngramsPerDepth.push([]);
        }
        for (const ngram of ngrams) {
            ngramsPerDepth[ngram.n - 1].push(ngram);
        }
        return ngramsPerDepth;
    }

    public profileBasedCoverage(otherModel: Model): any {
        const coverages: any = {};
        const ngrams = this.getAllNgram();
        const ngramsToCover = otherModel.getAllNgram();

        const ngramsSorted: Ngram[][] = this.sortNgram(ngrams);
        const otherNgramsSorted: Ngram[][] = otherModel.sortNgram(ngramsToCover);

        for (let i = 0; i < otherModel.depth; i++) {
            if (otherNgramsSorted[i].length > 0) {
                coverages[(i + 1).toString()] =  profileBasedCoverageForN(ngramsSorted[i], otherNgramsSorted[i]);
            }
        }
        coverages.all = profileBasedCoverageForN(ngrams, ngramsToCover);
        return coverages;
    }

    private contextOccurs(context: Stimulus[]): void {
        const lastStimulusOfContext = context[context.length - 1];
        const tree = this.findTreeByStimulus(lastStimulusOfContext);
        tree.contextOccurs(context);
    }

    private addStimulusKnowingContext(stimulus: Stimulus, context: Stimulus[]): void {

        if (context.length === 0) {
            const tree = this.findTreeByStimulus(stimulus);
            tree.learnStimulusKnowingContext(undefined, []);
        } else {
            const lastStimulusOfSequence = context[context.length - 1];
            const tree = this.findTreeByStimulus(lastStimulusOfSequence);
            tree.learnStimulusKnowingContext(stimulus, context);
        }
    }

    private addNoteKnowingContext(note: Note, context: Stimulus[]): void {
        if (context.length === 0) { return; }
        const lastStimulusOfSequence = context[context.length - 1];
        const tree = this.findTreeByStimulus(lastStimulusOfSequence);
        tree.learnNoteKnowingContext(note, context);
    }

    private findTreeByStimulus(stimulus: Stimulus): Tree {
        if (!this._treeMap.has(stimulus.value)) {
            this._treeMap.set(stimulus.value, new Tree(stimulus));
        }
        return this._treeMap.get(stimulus.value);
    }

    private fitContextToDepth(context: Stimulus[]): Stimulus[] {
        if (context.length < (this.depth)) {
            return context;
        } else {
            return context.slice(context.length - this.depth, context.length);
        }
    }

}

function profileBasedCoverageForN(ngrams: Ngram[], ngramsToCover: Ngram[]): {} {
    let totalOccurencesOther = 0;
    for (const ngram of ngramsToCover) {
        totalOccurencesOther += ngram.occurence;
    }

    const coverProbabilityMap = new Map();
    for (const ngram of ngramsToCover) {
        const probabilityOfSequence = ngram.occurence / totalOccurencesOther;
        coverProbabilityMap.set(ngram.key, probabilityOfSequence);
    }

    let coverage = 0;
    for (const ngram of ngrams) {
        if (coverProbabilityMap.has(ngram.key)) {
            coverage += coverProbabilityMap.get(ngram.key);
        }
    }
    const coverageRate = parseFloat(coverage.toFixed(2));
    const notCovered = ngramsToCover
        .filter((other) => !ngrams.some((ngram) => other.key === ngram.key))
        .map((ngram) => ({
            occurence: ngram.occurence,
            actions: ngram.key,
            probability: coverProbabilityMap.get(ngram.key),
        }))
        .sort((a, b) => b.probability - a.probability);
    const totalToCover = ngramsToCover.length;
    const covered = totalToCover - notCovered.length;

    return {
        coverageRate,
        notCovered,
        totalToCover,
        covered,
    };
}
