package net.ranides.assira.collection.sets;

import net.ranides.assira.collection.maps.AVLTreeMultiMap;
import net.ranides.assira.collection.maps.MultiMap;
import net.ranides.assira.generic.CompareUtils;
import net.ranides.assira.generic.Pair;

import java.util.Comparator;

public class AVLCrossSet<K1, K2> extends AVLTreeSet<Pair<K1, K2>> implements CrossSet<K1, K2> {

    private final Comparator<? super K1> comparator1;
    private final Comparator<? super K2> comparator2;

    public AVLCrossSet() {
        super();
        this.comparator1 = null;
        this.comparator2 = null;
    }

    public AVLCrossSet(Comparator<? super K1> comparator1, Comparator<? super K2> comparator2) {
        super((a,b) -> {
            int cmp1 = CompareUtils.cmp(comparator1, a.getFirst(), b.getFirst());
            return cmp1 != 0 ? cmp1 : CompareUtils.cmp(comparator2, a.getSecond(), b.getSecond());
        });
        this.comparator1 = comparator1;
        this.comparator2 = comparator2;
    }

    protected final int compare1(K1 key1, K1 key2) {
        try {
            return CompareUtils.cmp(comparator1, key1, key2);
        } catch(ClassCastException ex){
            // "unchecked" generic methods could pass something strange here
            return -1;
        }
    }

    protected final int compare2(K2 key1, K2 key2) {
        try {
            return CompareUtils.cmp(comparator2, key1, key2);
        } catch(ClassCastException ex){
            // "unchecked" generic methods could pass something strange here
            return -1;
        }
    }

    @Override
    public boolean contains(K1 key1, K2 key2) {
        return contains(Pair.of(key1, key2));
    }

    @Override
    public boolean add(K1 key1, K2 key2) {
        return add(Pair.of(key1, key2));
    }

    @Override
    public boolean remove(K1 key1, K2 key2) {
        return remove(Pair.of(key1, key2));
    }

    @Override
    public MultiSet<K2> reduceFirst() {
        MultiSet<K2> result = new AVLTreeMultiSet<>(this::compare2);
        forEach((key1, key2) -> result.add(key2));
        return result;
    }

    @Override
    public MultiSet<K1> reduceSecond() {
        MultiSet<K1> result = new AVLTreeMultiSet<>(this::compare1);
        forEach((key1, key2) -> result.add(key1));
        return result;
    }

    @Override
    public MultiMap<K1, K2> asMap() {
        MultiMap<K1, K2> result = new AVLTreeMultiMap<>(this::compare1);
        forEach((key1, key2) -> result.put(key1, key2));
        return result;
    }
}
