package au.id.zancanaro;

import java.io.OutputStream;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.function.Function;
import java.util.function.Predicate;

public class RoseTree<T> {
    private final T value;
    private final Iterable<RoseTree<T>> children;

    public RoseTree(T value, Iterable<RoseTree<T>> children) {
        this.value = value;
        this.children = children;
    }

    public T getValue() {
        return value;
    }

    public Iterator<RoseTree<T>> getChildren() {
        return children.iterator();
    }

    public static <T> RoseTree<T> pure(T value) {
        return new RoseTree<>(value, Collections.emptyList());
    }

    public static <T> RoseTree<T> join(RoseTree<RoseTree<T>> tree) {
        return new RoseTree<>(
                tree.getValue().getValue(),
                () -> Iterators.concat(
                        Iterators.mappingIterator(RoseTree::join, tree.children.iterator()),
                        tree.getValue().children.iterator()));
    }

    public static <T> Iterator<RoseTree<T>[]> permutations(RoseTree<T>[] trees) {
        return Iterators.flatten(Iterators.rangeIterator(trees.length, index ->
                        Iterators.mappingIterator(child -> {
                            @SuppressWarnings("unchecked")
                            RoseTree<T>[] result = (RoseTree<T>[]) new RoseTree[trees.length];
                            for (int i = 0; i < trees.length; ++i) {
                                result[i] = (i == index ? child : trees[i]);
                            }
                            return result;
                        }, trees[index].getChildren())
        ));
    }

    public static <T, R> RoseTree<R> zip(Function<T[], R> fn, RoseTree<T>[] trees) {
        @SuppressWarnings("unchecked")
        T[] heads = (T[]) new Object[trees.length];
        for (int i = 0; i < trees.length; ++i) {
            heads[i] = trees[i].getValue();
        }
        return new RoseTree<>(
                fn.apply(heads),
                () -> Iterators.mappingIterator(
                        roses -> RoseTree.zip(fn, roses),
                        RoseTree.permutations(trees)));
    }

    public <R> RoseTree<R> fmap(Function<T, R> f) {
        return new RoseTree<>(
                f.apply(this.value),
                () -> Iterators.mappingIterator(tree -> tree.fmap(f), this.children.iterator()));
    }

    public <R> RoseTree<R> flatmap(Function<T, RoseTree<R>> f) {
        return RoseTree.join(this.fmap(f));
    }

    public RoseTree<T> filter(Predicate<T> predicate) {
        if (predicate.test(this.getValue())) {
            return new RoseTree<>(
                    this.getValue(),
                    () -> Iterators.mappingIterator(tree -> tree.filter(predicate),
                            Iterators.filteringIterator(
                                    tree -> predicate.test(tree.getValue()),
                                    this.getChildren())));
        } else {
            throw new IllegalArgumentException("Current value doesn't match predicate: whoops!");
        }
    }
}