package eu.jonahbauer.raytracing.render.texture;

import eu.jonahbauer.raytracing.math.Vec3;
import org.jetbrains.annotations.NotNull;

import java.util.Objects;
import java.util.Random;
import java.util.function.DoubleFunction;
import java.util.random.RandomGenerator;

public final class PerlinTexture implements Texture {
    private static final int POINT_COUNT = 256;
    private static final @NotNull Random RANDOM = new Random();
    private static final @NotNull DoubleFunction<Color> GREYSCALE = t -> new Color(t, t, t);

    private final double scale;
    private final int turbulence;
    private final @NotNull DoubleFunction<Color> color;

    private final int mask;
    private final Vec3[] randvec;
    private final int[] permX;
    private final int[] permY;
    private final int[] permZ;

    public PerlinTexture() {
        this(1.0);
    }

    public PerlinTexture(double scale) {
        this(scale, 7);
    }

    public PerlinTexture(double scale, int turbulence) {
        this(scale, turbulence, GREYSCALE);
    }

    public PerlinTexture(double scale, int turbulence, @NotNull DoubleFunction<Color> color) {
        this(scale, turbulence, color, POINT_COUNT, RANDOM);
    }

    public PerlinTexture(
            double scale, int turbulence, @NotNull DoubleFunction<Color> color,
            int count, @NotNull RandomGenerator random
    ) {
        if ((count & (count - 1)) != 0) throw new IllegalArgumentException("count must be a power of two");
        if (turbulence <= 0) throw new IllegalArgumentException("turbulence must be positive");

        this.scale = scale;
        this.turbulence = turbulence;
        this.color = Objects.requireNonNull(color, "color");

        this.mask = count - 1;
        this.randvec = new Vec3[count];
        for (int i = 0; i < count; i++) {
            this.randvec[i] = Vec3.random(random, true);
        }
        this.permX = generatePerm(count, random);
        this.permY = generatePerm(count, random);
        this.permZ = generatePerm(count, random);
    }

    private static int @NotNull[] generatePerm(int count, @NotNull RandomGenerator random) {
        int[] p = new int[count];
        for (int i = 0; i < count; i++) {
            p[i] = i;
        }
        permutate(p, random);
        return p;
    }

    private static void permutate(int @NotNull[] p, @NotNull RandomGenerator random) {
        for (int i = p.length - 1; i > 0; i--) {
            int target = random.nextInt(i);
            int tmp = p[i];
            p[i] = p[target];
            p[target] = tmp;
        }
    }

    public double getNoise(@NotNull Vec3 p) {
        var x = p.x() * scale;
        var y = p.y() * scale;
        var z = p.z() * scale;

        var u = x - Math.floor(x);
        var v = y - Math.floor(y);
        var w = z - Math.floor(z);

        int i = (int) Math.floor(x);
        int j = (int) Math.floor(y);
        int k = (int) Math.floor(z);

        var c = new Vec3[2][2][2];
        for (int di = 0; di < 2; di++) {
            for (int dj = 0; dj < 2; dj++) {
                for (int dk = 0; dk < 2; dk++) {
                    c[di][dj][dk] = randvec[permX[(i + di) & mask] ^ permY[(j + dj) & mask] ^ permZ[(k + dk) & mask]];
                }
            }
        }

        return interpolate(c, u, v, w);
    }

    public double getNoise(@NotNull Vec3 p, int depth) {
        var accum = 0.0;
        var temp = p;
        var weight = 1.0;

        for (int i = 0; i < depth; i++) {
            accum += weight * getNoise(temp);
            weight *= 0.5;
            temp = temp.times(2);
        }

        return accum;
    }

    @Override
    public @NotNull Color get(double u, double v, @NotNull Vec3 p) {
        var noise = getNoise(p, turbulence);
        var t = Math.fma(0.5, Math.sin(Math.PI * noise), 0.5);
        return color.apply(t);
    }

    private static double interpolate(Vec3[][][] c, double u, double v, double w) {
        var uu = u * u * (3 - 2 * u);
        var vv = v * v * (3 - 2 * v);
        var ww = w * w * (3 - 2 * w);

        var accum = 0.0;
        for (int i = 0; i < 2; i++) {
            for (int j = 0; j < 2; j++) {
                for (int k = 0; k < 2; k++) {
                    var weight = new Vec3(u - i, v - j, w - k);
                    accum += (i * uu + (1 - i) * (1 - uu))
                             * (j * vv + (1 - j) * (1 - vv))
                             * (k * ww + (1 - k) * (1 - ww))
                             * c[i][j][k].times(weight);
                }
            }
        }
        return accum;
    }
}
