/*
 * 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.sequence;

import com.milaboratory.primitivio.PrimitivI;
import com.milaboratory.primitivio.PrimitivO;
import com.milaboratory.primitivio.annotations.Serializable;
import com.milaboratory.primitivio.blocks.PrimitivIOBlocksUtil;
import com.milaboratory.util.io.ByteArrayDataOutput;
import com.milaboratory.util.io.IOUtil;
import gnu.trove.iterator.TLongIterator;
import gnu.trove.list.array.TLongArrayList;
import gnu.trove.set.hash.TLongHashSet;
import net.jpountz.lz4.LZ4Compressor;
import net.jpountz.lz4.LZ4FastDecompressor;
import org.apache.commons.io.IOUtils;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.util.AbstractCollection;
import java.util.Arrays;
import java.util.Iterator;
import java.util.Objects;

@Serializable(by = IO.ShortSequenceSetSerializer.class)
public final class ShortSequenceSet extends AbstractCollection<NucleotideSequence> {
    private final TLongHashSet set;

    public ShortSequenceSet() {
        this(new TLongHashSet());
    }

    private ShortSequenceSet(TLongHashSet set) {
        this.set = set;
    }

    public boolean add(NucleotideSequence seq) {
        return set.add(toLong(seq));
    }

    public boolean contains(NucleotideSequence seq) {
        return !seq.containsWildcards() && set.contains(toLong(seq));
    }

    @Override
    public int size() {
        return set.size();
    }

    @Override
    public Iterator<NucleotideSequence> iterator() {
        TLongIterator it = set.iterator();
        return new Iterator<NucleotideSequence>() {
            @Override
            public boolean hasNext() {
                return it.hasNext();
            }

            @Override
            public NucleotideSequence next() {
                return fromLong(it.next());
            }
        };
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (!(o instanceof ShortSequenceSet)) return false;
        ShortSequenceSet that = (ShortSequenceSet) o;
        return set.equals(that.set);
    }

    @Override
    public int hashCode() {
        return Objects.hash(set);
    }

    public static long toLong(NucleotideSequence seq) {
        if (seq.size() > 29)
            throw new IllegalArgumentException("Only sequences shorter than 30 nucleotides are supported.");
        if (seq.containsWildcards())
            throw new IllegalArgumentException("Sequences containing wildcards are not supported.");
        long ret = 0;
        for (int i = 0; i < seq.size(); i++) {
            ret <<= 2;
            ret |= seq.codeAt(i);
        }
        ret |= ((long) seq.size()) << 58;
        return ret;
    }

    public static NucleotideSequence fromLong(long seq) {
        int size = (int) (seq >>> 58);
        SequenceBuilder<NucleotideSequence> seqBuilder = NucleotideSequence.ALPHABET.createBuilder()
                .ensureCapacity(size);
        for (int i = 0; i < size; i++)
            seqBuilder.append((byte) (3 & (seq >>> (2 * (size - i - 1)))));
        return seqBuilder.createAndDestroy();
    }

    public static final String MAGIC = "MiSSS_V1";

    public byte[] serialize() {
        TLongArrayList list = new TLongArrayList(set);
        list.sort();
        ByteArrayDataOutput bos = new ByteArrayDataOutput();
        try (PrimitivO o = new PrimitivO(bos)) {
            long p = 0, c;
            for (int i = 0; i < list.size(); i++) {
                c = list.get(i);
                o.writeVarLong(c - p);
                p = c;
            }
        }
        byte[] buffer = bos.getBuffer();
        LZ4Compressor compressor = PrimitivIOBlocksUtil.highLZ4Compressor();
        byte[] compressed = compressor.compress(buffer);
        byte[] magicBytes = MAGIC.getBytes(StandardCharsets.US_ASCII);
        byte[] output = new byte[magicBytes.length + 12 + compressed.length];
        System.arraycopy(magicBytes, 0, output, 0, magicBytes.length);
        IOUtil.writeIntBE(compressed.length, output, magicBytes.length);
        IOUtil.writeIntBE(buffer.length, output, magicBytes.length + 4);
        IOUtil.writeIntBE(list.size(), output, magicBytes.length + 8);
        System.arraycopy(compressed, 0, output, magicBytes.length + 12, compressed.length);
        return output;
    }

    public static boolean isShortSequenceSet(InputStream is) throws IOException {
        byte[] magicBytes = MAGIC.getBytes(StandardCharsets.US_ASCII);
        is.mark(magicBytes.length);
        try {
            byte[] header = new byte[magicBytes.length];
            if (IOUtils.read(is, header) != header.length)
                return false;
            return Arrays.equals(header, magicBytes);
        } finally {
            is.reset();
        }
    }

    public static ShortSequenceSet deserialize(InputStream is) throws IOException {
        byte[] magicBytes = MAGIC.getBytes(StandardCharsets.US_ASCII);
        byte[] header = new byte[magicBytes.length + 12];
        IOUtils.readFully(is, header);
        if (!Arrays.equals(magicBytes, Arrays.copyOf(header, magicBytes.length)))
            throw new IllegalArgumentException("Malformed sequence set");
        int compressedSize = IOUtil.readIntBE(header, magicBytes.length);
        int uncompressedSize = IOUtil.readIntBE(header, magicBytes.length + 4);
        int numberOfRecords = IOUtil.readIntBE(header, magicBytes.length + 8);
        byte[] compressed = new byte[compressedSize];
        IOUtils.readFully(is, compressed);
        LZ4FastDecompressor decompressor = PrimitivIOBlocksUtil.defaultLZ4Decompressor();
        byte[] uncompressed = new byte[uncompressedSize];
        decompressor.decompress(compressed, uncompressed);
        TLongHashSet set = new TLongHashSet();
        try (PrimitivI i = new PrimitivI(new ByteArrayInputStream(uncompressed))) {
            long v = 0;
            for (int n = 0; n < numberOfRecords; n++)
                set.add(v += i.readVarLong());
        }
        return new ShortSequenceSet(set);
    }
}
