/*
 * Copyright (c) 2022 MiLaboratories Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.milaboratory.core.motif;

import com.milaboratory.core.sequence.Sequence;

import static com.milaboratory.core.motif.LongArrayBitHelper.*;

abstract class BitapStateIterator {
    final BitapData data;
    final Sequence sequence;
    int errors;
    final long[][] R;
    final int to;
    int symbolsProcessed = 0;
    int current;
    boolean match;

    BitapStateIterator(BitapData data, Sequence sequence, int count, int from, int to) {
        if (sequence.getAlphabet().size() != data.patternMask.length)
            throw new IllegalArgumentException();
        this.data = data;
        this.sequence = sequence;
        this.R = LongArrayBitHelper.ones2D(count, data.segmentCount);
        for (int i = 1; i < count; ++i)
            leftShiftLittleEndian(R[i], i);
        this.to = to;
        this.current = from;
    }

    /**
     * Try to advance the state to the next position, returns false, if sequence boundary is reached.
     * {@link BitapStateIterator @match} shows whether tha match was found.
     */
    abstract boolean nextState();

    abstract int currentPosition();

    static final class ExactMatchStateIterator extends BitapStateIterator {
        public ExactMatchStateIterator(BitapData data, Sequence sequence, int from, int to) {
            super(data, sequence, 1, from, to);
            errors = 0;
        }

        @Override
        boolean nextState() {
            // Reset state
            match = false;

            // Check end of sequence
            if (current == to)
                return false;

            // Main part
            leftShiftLittleEndian(R[0], 1);
            orAssign(R[0], data.patternMask[sequence.codeAt(current)]);
            ++current;
            match = !getBitLittleEndian(R[0], data.size - 1);

            // Next state calculated
            return true;
        }

        @Override
        int currentPosition() {
            return current - data.size;
        }
    }

    static final class SubstitutionOnlyFirstStateIterator extends BitapStateIterator {
        public SubstitutionOnlyFirstStateIterator(BitapData data, Sequence sequence, int maxSubstitutions, int from, int to) {
            super(data, sequence, maxSubstitutions + 1, from, to);
        }

        @Override
        boolean nextState() {
            // Reset state
            match = false;

            // Check end of sequence
            if (current == to)
                return false;

            int d;
            long[]
                    swapTemp,
                    // These are the only two array allocations in this method
                    preMismatchTmp = new long[data.segmentCount],
                    mismatchTmp = new long[data.segmentCount];

            // Main part
            // long matchingMask = (1L << ());

            long[] currentPatternMask = data.patternMask[sequence.codeAt(current)];
            ++current;
            ++symbolsProcessed;

            // Exact match on the previous step == match with insertion on current step
            leftShiftLittleEndian(R[0], 1);
            assign(mismatchTmp, R[0]);
            orAssign(R[0], currentPatternMask);

            if (!getBitLittleEndian(R[0], data.size - 1)) {
                errors = 0;
                match = true;
            }

            for (d = 1; d < R.length; ++d) {
                leftShiftLittleEndian(R[d], 1);
                assign(preMismatchTmp, R[d]);
                orAssign(R[d], currentPatternMask);
                andAssign(R[d], mismatchTmp);
                if (!match && !getBitLittleEndian(R[d], data.size - 1) && symbolsProcessed >= data.size) {
                    errors = d;
                    match = true;
                }

                // Read as: mismatchTmp = preMismatchTmp
                // but preserving already allocated array for the next preMismatchTmp
                swapTemp = mismatchTmp;
                mismatchTmp = preMismatchTmp;
                preMismatchTmp = swapTemp;
            }

            return true;
        }

        @Override
        int currentPosition() {
            assert current >= data.size;
            return current - data.size;
        }
    }

    static abstract class IndelStateIterator extends BitapStateIterator {
        public IndelStateIterator(BitapData data, Sequence sequence, int maxErrors, int from, int to) {
            super(data, sequence, maxErrors + 1, from, to);
        }

        void updateState(long[] currentPatternMask) {
            // long matchingMask = (1L << ());
            // data.size - 1

            long[]
                    swapTmp,
                    preInsertionTmp = new long[data.segmentCount],
                    preMismatchTmp = new long[data.segmentCount],
                    insertionTmp = new long[data.segmentCount],
                    deletionTmp = new long[data.segmentCount],
                    mismatchTmp = new long[data.segmentCount];

            // Exact match on the previous step == match with insertion on current step
            assign(insertionTmp, R[0]);
            leftShiftLittleEndian(R[0], 1);
            assign(mismatchTmp, R[0]);
            orAssign(R[0], currentPatternMask);
            assign(deletionTmp, R[0]);

            if (!getBitLittleEndian(R[0], data.size - 1)) {
                errors = 0;
                match = true;
            }

            for (int d = 1; d < R.length; ++d) {
                assign(preInsertionTmp, R[d]);
                leftShiftLittleEndian(R[d], 1);
                assign(preMismatchTmp, R[d]);
                orAssign(R[d], currentPatternMask);

                leftShiftLittleEndian(deletionTmp, 1);
                andAssign(R[d], insertionTmp, mismatchTmp, deletionTmp);
                if (!match && !getBitLittleEndian(R[d], data.size - 1) && symbolsProcessed >= data.size - R.length + 1) {
                    errors = d;
                    match = true;
                }

                assign(deletionTmp, R[d]);

                // Read as:
                //   insertionTmp = preInsertionTmp;
                //   mismatchTmp = preMismatchTmp;
                // but preserving already allocated array for the next preMismatchTmp

                swapTmp = insertionTmp;
                insertionTmp = preInsertionTmp;
                preInsertionTmp = swapTmp;

                swapTmp = mismatchTmp;
                mismatchTmp = preMismatchTmp;
                preMismatchTmp = swapTmp;
            }
        }
    }

    static final class SubstitutionAndIndelLastStateIterator extends IndelStateIterator {
        public SubstitutionAndIndelLastStateIterator(BitapData data, Sequence sequence, int maxErrors, int from, int to) {
            super(data, sequence, maxErrors, from, to);
        }

        @Override
        boolean nextState() {
            // Reset state
            match = false;

            // Check end of sequence
            if (current == to)
                return false;

            // Main part
            ++symbolsProcessed;
            updateState(data.patternMask[sequence.codeAt(current++)]);

            return true;
        }

        @Override
        int currentPosition() {
            return current - 1;
        }
    }

    static final class SubstitutionAndIndelFirstStateIterator extends IndelStateIterator {
        public SubstitutionAndIndelFirstStateIterator(BitapData data, Sequence sequence, int maxErrors, int from, int to) {
            super(data, sequence, maxErrors, to - 1, from);
        }

        @Override
        boolean nextState() {
            // Reset state
            match = false;

            // Check end of sequence
            if (current == to - 1)
                return false;

            // Main part
            ++symbolsProcessed;
            updateState(data.reversePatternMask[sequence.codeAt(current--)]);

            return true;
        }

        @Override
        int currentPosition() {
            return current + 1;
        }
    }
}
