/FileExtraction/Java_unseen/GenSpectrum_LAPIS/server/src/main/java/ch/ethz/lapis/api/query/InsertionStore.java
package ch.ethz.lapis.api.query;

import java.util.*;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

public class InsertionStore {

    /** Position -> (inserted bases -> list of sequence ids) */
    private final Map<Short, Map<String, List<Integer>>> positionMap = new HashMap<>();

    /** Sequence "id" -> insertions of the sequence */
    private final Map<Integer, List<String>> sequenceMap = new HashMap<>();

    /**
     *
     * @param insertion In the format as generated by Nextclade, e.g., 1234:AATCCG or 2345:EPE (only upper-case)
     */
    public void putInsertions(int sequenceId, String insertion) {
        String[] split = insertion.split(":");
        short position = Short.parseShort(split[0]);
        String insertedBases = split[1];
        if (!positionMap.containsKey(position)) {
            positionMap.put(position, new HashMap<>());
        }
        var insertionMapAtPosition = positionMap.get(position);
        if (!insertionMapAtPosition.containsKey(insertedBases)) {
            insertionMapAtPosition.put(insertedBases, new ArrayList<>());
        }
        var sequencesWithInsertion = insertionMapAtPosition.get(insertedBases);
        sequencesWithInsertion.add(sequenceId);
        if (!sequenceMap.containsKey(sequenceId)) {
            sequenceMap.put(sequenceId, new ArrayList<>());
        }
        sequenceMap.get(sequenceId).add(insertion);
    }

    /**
     *
     * @param insertionQuery Examples: 1234:AATCGG, 1234:EPE, 1234:AAT?, 1234:AAT?G, 1234:E?
     * @return A list of sequence ids
     */
    public List<Integer> find(String insertionQuery) {
        insertionQuery = insertionQuery.toUpperCase();
        String[] split = insertionQuery.split(":");
        short position = Short.parseShort(split[0]);
        String basesQuery = split[1];

        Map<String, List<Integer>> insertionsAtPosition = positionMap.get(position);
        if (insertionsAtPosition == null) {
            return List.of();
        }
        if (!basesQuery.contains("?")) {
            // No wildcard
            return insertionsAtPosition.getOrDefault(basesQuery, List.of());
        } else {
            // With wildcard
            Pattern queryRegex = Pattern.compile(basesQuery.replace("?", ".*"));
            List<Integer> allSequences = new ArrayList<>();
            insertionsAtPosition.forEach((insertion, sequences) -> {
                if (queryRegex.matcher(insertion).matches()) {
                    allSequences.addAll(sequences);
                }
            });
            return allSequences;
        }
    }

    public List<InsertionCount> countInsertions(Collection<Integer> ids) {
        Map<String, int[]> insertionCounts = new HashMap<>();
        for (int id : ids) {
            List<String> insertions  = sequenceMap.get(id);
            if (insertions != null) {
                for (String insertion : insertions) {
                    insertionCounts.compute(insertion, (k, v) -> v == null ?
                        new int[] { 0 } : v)[0]++;
                }
            }
        }
        return insertionCounts.entrySet().stream()
            .map(e -> new InsertionCount(e.getKey(), e.getValue()[0]))
            .collect(Collectors.toList());
    }

    public record InsertionCount(String insertion, int count) {
    }
}
