package common.expression;

import java.util.Arrays;
import java.util.List;
import java.util.OptionalDouble;
import java.util.function.DoubleBinaryOperator;
import java.util.function.DoubleUnaryOperator;
import java.util.function.Function;
import java.util.function.ToDoubleFunction;
import java.util.stream.DoubleStream;
import java.util.stream.IntStream;
import java.util.stream.Stream;

/**
 * Known expression operations.
 *
 * @author Georgiy Korneev (kgeorgiy@kgeorgiy.info)
 */
@SuppressWarnings("StaticMethodOnlyUsedInOneClass")
public enum Operations {
    ;

    public static final Operation ARITH = builder -> {
        builder.ops.alias("negate", "Negate");
        builder.ops.alias("+", "Add");
        builder.ops.alias("-", "Subtract");
        builder.ops.alias("*", "Multiply");
        builder.ops.alias("/", "Divide");
    };

    public static final Operation NARY_ARITH = builder -> {
        builder.ops.remove("negate", "+", "-", "*", "/");

        builder.ops.unary("negate", "Negate", a -> -a);

        builder.ops.any("+", "Add", arith(0, Double::sum, 0));
        builder.ops.any("-", "Subtract", arith(0, (a, b) -> a - b, 1));
        builder.ops.any("*", "Multiply", arith(1, (a, b) -> a * b, 0));
        builder.ops.any("/", "Divide", arith(1, (a, b) -> a / b, 1));
    };


    // === Common

    public static Operation constant(final String name, final double value) {
        return constant(name, name, value);
    }

    public static Operation constant(final String name, final String alias, final double value) {
        return builder -> builder.ops.constant(name, alias, value);
    }

    public static Operation unary(final String name, final String alias, final DoubleUnaryOperator op) {
        return builder -> builder.ops.unary(name, alias, op);
    }

    public static Operation binary(final String name, final String alias, final DoubleBinaryOperator op) {
        return builder -> builder.ops.binary(name, alias, op);
    }

    public static AnyOp arith(final double zero, final DoubleBinaryOperator f, final int minArity) {
        final ExprTester.Func func = args -> args.length == 0 ? zero
                : args.length == 1 ? f.applyAsDouble(zero, args[0])
                        : Arrays.stream(args).reduce(f).orElseThrow();
        return new AnyOp(func, minArity, minArity + 5, 2);
    }


    // === More common

    public record Op(String name, String alias, int minArity, int maxArity, ExprTester.Func f) {
        public Operation fix(final int arity) {
            assert minArity <= arity && arity <= maxArity;
            final String patched =
                    name.contains("₀") ? name.replace("₀", fix(arity, '₀')) :
                    name.contains("0") ? name.replace("0", fix(arity, '0')) :
                    name + fix(arity, name.charAt(0) > 0xff ? '₀' : '0');
            return fixed(patched, alias + arity, arity, f);
        }

        private static String fix(final int arity, final char c) {
            return "" + (char) (c + arity);
        }

        public Operation any(final int fixedArity) {
            return checker -> checker.ops.any(name, alias, new AnyOp(f, minArity, maxArity, fixedArity));
        }
    }

    public static Op op(final String name, final String alias, final int minArity, final ExprTester.Func f) {
        return new Op(name, alias, minArity, minArity + 5, f);
    }

    public static Op op1(final String alias, final int minArity, final ExprTester.Func f) {
        return new Op(Character.toLowerCase(alias.charAt(0)) + alias.substring(1), alias, minArity, minArity + 5, f);
    }

    public static Op opS(final String name, final String alias, final int minArity, final ToDoubleFunction<DoubleStream> f) {
        return op(name, alias, minArity, args -> f.applyAsDouble(Arrays.stream(args)));
    }

    public static Op opO(final String name, final String alias, final int minArity, final Function<DoubleStream, OptionalDouble> f) {
        return opS(name, alias, minArity, f.andThen(OptionalDouble::orElseThrow)::apply);
    }

    public static Operation fixed(final String name, final String alias, final int arity, final ExprTester.Func f) {
        return builder -> builder.ops.fixed(name, alias, arity, f);
    }

    @SuppressWarnings("SameParameterValue")
    public static Operation range(final int min, final int max, final Op... ops) {
        final List<Operation> operations = IntStream.rangeClosed(min, max)
                .mapToObj(i -> Arrays.stream(ops).map(op -> op.fix(i)))
                .flatMap(Function.identity())
                .toList();
        return builder -> operations.forEach(op -> op.accept(builder));
    }

    @SuppressWarnings("SameParameterValue")
    public static Operation any(final int fixed, final Op... ops) {
        final List<Operation> operations = Arrays.stream(ops).map(op -> op.any(fixed)).toList();
        return builder -> operations.forEach(op -> op.accept(builder));
    }

    public static Operation infix(
            final String name,
            final String alias,
            final int priority,
            final DoubleBinaryOperator op
    ) {
        return checker -> checker.ops.infix(name, alias, priority, op);
    }


    // === Variables
    public static final Operation VARIABLES = builder ->
            Stream.of("y", "z").forEach(builder.ops::variable);


    // === OneTwo
    public static final Operation ONE = constant("one", 1);
    public static final Operation TWO = constant("two", 2);
    public static final Operation THREE = constant("three", 3);


    // === Clamp, wrap
    public static final Operation CLAMP = fixed("clamp", "Clamp", 3, args ->
            args[1] <= args[2] ? Math.min(Math.max(args[0], args[1]), args[2]) : Double.NaN);
    public static final Operation WRAP = fixed("wrap", "Wrap", 3, args ->
            args[1] < args[2]
                    ? args[0] - Math.floor((args[0] - args[1]) / (args[2] - args[1])) * (args[2] - args[1])
                    : Double.NaN);

    // === ArcTan
    public static final Operation ATAN = unary("atan", "ArcTan", Math::atan);
    public static final Operation ATAN2 = binary("atan2", "ArcTan2", Math::atan2);

    // === ArgMin, ArgMax

    private static Op arg(final String name, final String alias, final Function<DoubleStream, OptionalDouble> f) {
        return opO(
                name, "Arg" + alias, 1, args -> {
                    final double[] values = args.toArray();
                    return f.apply(Arrays.stream(values)).stream()
                            .flatMap(value -> IntStream.range(0, values.length)
                                    .filter(i -> values[i] == value).asDoubleStream())
                            .findFirst();
                }
        );
    }

    public static final Op ARG_MIN = arg("argMin", "Min", DoubleStream::min);
    public static final Op ARG_MAX = arg("argMax", "Max", DoubleStream::max);


    // === SoftClamp
    public static final Operation SOFT_CLAMP = fixed("softClamp", "SoftClamp", 4, args ->
            args[1] <= args[2] && args[3] > 0
                    ? args[1] + (args[2] - args[1]) / (1 + Math.exp(args[3] * ((args[2] + args[1]) / 2 - args[0])))
                    : Double.NaN);


    // === SinCos
    public static final Operation SIN = unary("sin", "Sin", Math::sin);
    public static final Operation COS = unary("cos", "Cos", Math::cos);


    // === Pow, Log
    public static final Operation POW = binary("pow", "Power", Math::pow);
    public static final Operation LOG = binary("log", "Log", (a, b) -> Math.log(Math.abs(b)) / Math.log(Math.abs(a)));


    // === Sum
    public static final Op SUM = opS("sum", "Sum", 0, DoubleStream::sum);
}
