import Ngram from "./Ngram";
import Note from "./Note";
import NoteDistribution from "./NoteDistribution";
import Sequence from "./Sequence";
import Stimulus from "./Stimulus";
import Tree from "./Tree_CSP";
import Model from "./Model";

export default class CSP_Model extends Model  {

    private _treeMap: Map <string, Tree>;
    private readonly _interpolationfactor: number;

    constructor(depth: number, interpolationfactor: number= 2, id?: string) {
        super(depth, id)
        this._interpolationfactor = interpolationfactor;

        this._treeMap = new Map();
    }

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

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

    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(sequence: Sequence): Map<string, number> {
        const context = sequence.getContext()
        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);
        }
    }
    
    public getAllNgram(): Ngram[] {
        const ngramSet = [];
        for (const tree of this._treeMap.values()) {
            tree.getAllNgram().forEach( (ngram) => {
                ngramSet.push(ngram);
            });
        }
        return ngramSet;
    }

    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);
    }

}
