package com.github.sh0nk.matplotlib4j;

import com.google.common.base.Preconditions;
import com.google.common.collect.ContiguousSet;
import com.google.common.collect.DiscreteDomain;
import com.google.common.collect.Range;

import java.util.List;
import java.util.function.BiFunction;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

public class NumpyUtils {

    // TODO: more options from numpy
    public static List<Double> linspace(double start, double end, int num) {
        Preconditions.checkArgument(num >= 0);
        return ContiguousSet.create(Range.closedOpen(0, num), DiscreteDomain.integers())
                .stream().map(x -> (x  * (end - start)) / (num - 1) + start).collect(Collectors.toList());
    }

    public static List<Double> arange(double start, double end, double step) {
        double scaledStart = start / step;
        double scaledEnd = end / step;
        double floorGap = scaledStart - (int) scaledStart;
        return ContiguousSet.create(Range.closed((int) scaledStart, (int) scaledEnd), DiscreteDomain.integers())
                .stream().map(x -> (x + floorGap) * step).collect(Collectors.toList());
    }

    public static <T extends Number> Grid<T> meshgrid(List<T> x, List<T> y) {
        Grid<T> grid = new Grid<>();
        grid.x = IntStream.range(0, y.size()).mapToObj(i -> x).collect(Collectors.toList());
        grid.y = y.stream().map(t -> (IntStream.range(0, x.size()).mapToObj(i -> t).collect(Collectors.toList()))).collect(Collectors.toList());
        return grid;
    }

    public static class Grid<T extends Number> {
        public List<List<T>> x;
        public List<List<T>> y;

        public <R> List<List<R>> calcZ(BiFunction<T, T, R> biFunction) {
            return IntStream.range(0, x.size()).mapToObj(i ->
                IntStream.range(0, x.get(i).size()).mapToObj(j ->
                    biFunction.apply(x.get(i).get(j), y.get(i).get(j))
                ).collect(Collectors.toList())
            ).collect(Collectors.toList());
        }
    }

}