package eu.jonahbauer.raytracing.render.renderer;

import eu.jonahbauer.raytracing.math.Range;
import eu.jonahbauer.raytracing.math.Ray;
import eu.jonahbauer.raytracing.render.texture.Color;
import eu.jonahbauer.raytracing.render.camera.Camera;
import eu.jonahbauer.raytracing.render.canvas.Canvas;
import eu.jonahbauer.raytracing.scene.Scene;
import org.jetbrains.annotations.NotNull;

import java.util.Random;
import java.util.SplittableRandom;
import java.util.random.RandomGenerator;
import java.util.stream.IntStream;

public final class SimpleRenderer implements Renderer {
    private final int sqrtSamplesPerPixel;
    private final int maxDepth;
    private final double gamma;

    private final boolean parallel;
    private final boolean iterative;

    public static @NotNull Builder builder() {
        return new Builder();
    }

    public static @NotNull Renderer withDefaults() {
        return new Builder().build();
    }

    private SimpleRenderer(@NotNull Builder builder) {
        this.sqrtSamplesPerPixel = (int) Math.sqrt(builder.samplesPerPixel);
        this.maxDepth = builder.maxDepth;
        this.gamma = builder.gamma;

        this.parallel = builder.parallel;
        this.iterative = builder.iterative;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void render(@NotNull Camera camera, @NotNull Scene scene, @NotNull Canvas canvas) {
        if (canvas.getWidth() != camera.getWidth() || canvas.getHeight() != camera.getHeight()) {
            throw new IllegalArgumentException("sizes of camera and canvas are different");
        }

        if (iterative) {
            renderIterative(camera, scene, canvas);
        } else {
            renderNonIterative(camera, scene, canvas);
        }
    }

    /**
     * Renders the {@code scene} as seen by the {@code camera} to the {@code canvas}, taking one sample per pixel at
     * a time and updating the canvas after each sample.
     */
    private void renderIterative(@NotNull Camera camera, @NotNull Scene scene, @NotNull Canvas canvas) {
        var random = new Random();

        // render one sample after the other
        int i = 0;
        for (int sj = 0; sj < sqrtSamplesPerPixel; sj++) {
            for (int si = 0; si < sqrtSamplesPerPixel; si++) {
                var sample = ++i;
                var sif = si;
                var sjf = sj;
                getScanlineStream(camera.getHeight(), parallel).forEach(y -> {
                    for (int x = 0; x < camera.getWidth(); x++) {
                        var ray = camera.cast(x, y, sif, sjf, sqrtSamplesPerPixel, random);
                        var c = getColor(scene, ray, random);
                        canvas.set(x, y, Color.average(canvas.get(x, y), c, sample));
                    }
                });
            }
        }

        // apply gamma correction
        getScanlineStream(camera.getHeight(), parallel).forEach(y -> {
            for (int x = 0; x < camera.getWidth(); x++) {
                canvas.set(x, y, Color.gamma(canvas.get(x, y), gamma));
            }
        });
    }

    /**
     * Renders the {@code scene} as seen by the {@code camera} to the {@code canvas}, taking some amount of samples
     * per pixel and updating the canvas after each pixel.
     */
    private void renderNonIterative(@NotNull Camera camera, @NotNull Scene scene, @NotNull Canvas canvas) {
        var splittable = new SplittableRandom();
        // render one pixel after the other
        getScanlineStream(camera.getHeight(), parallel).forEach(y -> {
            var random = splittable.split();
            for (int x = 0; x < camera.getWidth(); x++) {
                var color = Color.BLACK;
                int i = 0;
                for (int sj = 0; sj < sqrtSamplesPerPixel; sj++) {
                    for (int si = 0; si < sqrtSamplesPerPixel; si++) {
                        var ray = camera.cast(x, y, si, sj, sqrtSamplesPerPixel, random);
                        var c = getColor(scene, ray, random);
                        color = Color.average(color, c, ++i);
                    }
                }
                canvas.set(x, y, Color.gamma(color, gamma));
            }
        });
    }

    /**
     * {@return the color of the given ray in the given scene}
     */
    private @NotNull Color getColor(@NotNull Scene scene, @NotNull Ray ray, @NotNull RandomGenerator random) {
        return getColor0(scene, ray, maxDepth, random);
    }

    private @NotNull Color getColor0(@NotNull Scene scene, @NotNull Ray ray, int depth, @NotNull RandomGenerator random) {
        var color = Color.BLACK;
        var attenuation = Color.WHITE;

        while (depth-- > 0) {
            var optional = scene.hit(ray, new Range(0.001, Double.POSITIVE_INFINITY));
            if (optional.isEmpty()) {
                color = Color.add(color, Color.multiply(attenuation, scene.getBackgroundColor(ray)));
                break;
            }

            var hit = optional.get();
            var material = hit.material();
            var emitted = material.emitted(hit);
            var scatter = material.scatter(ray, hit, random);
            color = Color.add(color, Color.multiply(attenuation, emitted));

            if (scatter.isEmpty()) break;
            attenuation = Color.multiply(attenuation, scatter.get().attenuation());
            ray = scatter.get().ray();
        }

        return color;
    }

    /**
     * {@return a stream of the pixels in a canvas with the given size} The pixels {@code x} and {@code y} coordinate
     * are encoded in the longs lower and upper 32 bits respectively.
     */
    private static @NotNull IntStream getScanlineStream(int height, boolean parallel) {
        var stream = IntStream.range(0, height);
        return parallel ? stream.parallel() : stream;
    }

    public static class Builder {
        private int samplesPerPixel = 100;
        private int maxDepth = 10;
        private double gamma = 2.0;
        private boolean parallel = true;
        private boolean iterative = false;

        public @NotNull Builder withSamplesPerPixel(int samples) {
            if (samples <= 0) throw new IllegalArgumentException("samples must be positive");
            this.samplesPerPixel = samples;
            return this;
        }

        public @NotNull Builder withMaxDepth(int depth) {
            if (depth <= 0) throw new IllegalArgumentException("depth must be positive");
            this.maxDepth = depth;
            return this;
        }

        public @NotNull Builder withGamma(double gamma) {
            if (gamma <= 0 || !Double.isFinite(gamma)) throw new IllegalArgumentException("gamma must be positive");
            this.gamma = gamma;
            return this;
        }

        public @NotNull Builder withParallel(boolean parallel) {
            this.parallel = parallel;
            return this;
        }

        public @NotNull Builder withIterative(boolean iterative) {
            this.iterative = iterative;
            return this;
        }

        public @NotNull SimpleRenderer build() {
            return new SimpleRenderer(this);
        }
    }
}
