package eu.jonahbauer.raytracing;

import eu.jonahbauer.raytracing.math.AABB;
import eu.jonahbauer.raytracing.math.Vec3;
import eu.jonahbauer.raytracing.render.texture.CheckerTexture;
import eu.jonahbauer.raytracing.render.texture.Color;
import eu.jonahbauer.raytracing.render.camera.SimpleCamera;
import eu.jonahbauer.raytracing.render.material.*;
import eu.jonahbauer.raytracing.render.texture.ImageTexture;
import eu.jonahbauer.raytracing.render.texture.PerlinTexture;
import eu.jonahbauer.raytracing.scene.Hittable;
import eu.jonahbauer.raytracing.scene.Scene;
import eu.jonahbauer.raytracing.scene.SkyBox;
import eu.jonahbauer.raytracing.scene.hittable2d.Parallelogram;
import eu.jonahbauer.raytracing.scene.hittable3d.Box;
import eu.jonahbauer.raytracing.scene.hittable3d.ConstantMedium;
import eu.jonahbauer.raytracing.scene.hittable3d.Sphere;
import eu.jonahbauer.raytracing.scene.util.HittableBinaryTree;
import org.jetbrains.annotations.NotNull;

import java.util.*;
import java.util.function.IntFunction;

public class Examples {
    private static final Map<String, IntFunction<Example>> REGISTRY = new HashMap<>();

    private static void register(@NotNull String name, @NotNull IntFunction<Example> example) {
        REGISTRY.put(name, example);
    }

    static {
        register("SIMPLE", Examples::getSimpleScene);
        register("SPHERES", Examples::getSpheres);
        register("SQUARES", Examples::getSquares);
        register("LIGHT", Examples::getLight);
        register("CORNELL", Examples::getCornellBox);
        register("CORNELL_SMOKE", Examples::getCornellBoxSmoke);
        register("DIAGRAMM", Examples::getDiagramm);
        register("EARTH", Examples::getEarth);
        register("PERLIN", Examples::getPerlin);
        register("FINAL", Examples::getFinal);
    }

    public static @NotNull IntFunction<Example> getByName(@NotNull String name) {
        var out = REGISTRY.get(name);
        if (out == null) throw new IllegalArgumentException("unknown example " + name + ", expected one of " + REGISTRY.keySet());
        return out;
    }

    public static @NotNull Example getSimpleScene(int height) {
        if (height <= 0) height = 675;
        return new Example(
                new Scene(
                        getSkyBox(),
                        new Sphere(new Vec3(0, -100.5, -1.0), 100.0, new LambertianMaterial(new Color(0.8, 0.8, 0.0))),
                        new Sphere(new Vec3(0, 0, -1.2), 0.5, new LambertianMaterial(new Color(0.1, 0.2, 0.5))),
                        new Sphere(new Vec3(-1.0, 0, -1.2), 0.5, new DielectricMaterial(1.5)),
                        new Sphere(new Vec3(-1.0, 0, -1.2), 0.4, new DielectricMaterial(1 / 1.5)),
                        new Sphere(new Vec3(1.0, 0, -1.2), 0.5, new MetallicMaterial(new Color(0.8, 0.6, 0.2), 0.0))
                ),
                SimpleCamera.builder()
                        .withImage(height * 16 / 9, height)
                        .build()
        );
    }

    public static @NotNull Example getSpheres(int height) {
        if (height <= 0) height = 675;

        var rng = new Random(1);
        var objects = new ArrayList<Hittable>();
        objects.add(new Sphere(
                new Vec3(0, -1000, 0), 1000,
                new LambertianMaterial(new CheckerTexture(0.32, new Color(.2, .3, .1), new Color(.9, .9, .9)))
        ));

        for (int a = -11; a < 11; a++) {
            for (int b = -11; b < 11; b++) {
                var center = new Vec3(a + 0.9 * rng.nextDouble(), 0.2, b + 0.9 * rng.nextDouble());
                if (Vec3.distance(center, new Vec3(4, 0.2, 0)) <= 0.9) continue;

                Material material;
                var rnd = rng.nextDouble();
                if (rnd < 0.8) {
                    // diffuse
                    var albedo = Color.multiply(Color.random(rng), Color.random(rng));
                    material = new LambertianMaterial(albedo);
                } else if (rnd < 0.95) {
                    // metal
                    var albedo = Color.random(rng, 0.5, 1.0);
                    var fuzz = rng.nextDouble() * 0.5;
                    material = new MetallicMaterial(albedo, fuzz);
                } else {
                    // glass
                    material = new DielectricMaterial(1.5);
                }

                objects.add(new Sphere(center, 0.2, material));
            }
        }

        objects.add(new Sphere(new Vec3(0, 1, 0), 1.0, new DielectricMaterial(1.5)));
        objects.add(new Sphere(new Vec3(-4, 1, 0), 1.0, new LambertianMaterial(new Color(0.4, 0.2, 0.1))));
        objects.add(new Sphere(new Vec3(4, 1, 0), 1.0, new MetallicMaterial(new Color(0.7, 0.6, 0.5))));

        var camera = SimpleCamera.builder()
                .withImage(height * 16 / 9, height)
                .withPosition(new Vec3(13, 2, 3))
                .withTarget(new Vec3(0, 0, 0))
                .withFieldOfView(Math.toRadians(20))
                .withFocusDistance(10.0)
                .withBlurAngle(Math.toRadians(0.6))
                .build();

        return new Example(new Scene(getSkyBox(), objects), camera);
    }

    public static @NotNull Example getSquares(int height) {
        if (height <= 0) height = 600;
        return new Example(
                new Scene(
                        getSkyBox(),
                        new Parallelogram(new Vec3(-3, -2, 5), new Vec3(0, 0, -4), new Vec3(0, 4, 0), new LambertianMaterial(new Color(1.0, 0.2, 0.2))),
                        new Parallelogram(new Vec3(-2, -2, 0), new Vec3(4, 0, 0), new Vec3(0, 4, 0), new LambertianMaterial(new Color(0.2, 1.0, 0.2))),
                        new Parallelogram(new Vec3(3, -2, 1), new Vec3(0, 0, 4), new Vec3(0, 4, 0), new LambertianMaterial(new Color(0.2, 0.2, 1.0))),
                        new Parallelogram(new Vec3(-2, 3, 1), new Vec3(4, 0, 0), new Vec3(0, 0, 4), new LambertianMaterial(new Color(1.0, 0.5, 0.0))),
                        new Parallelogram(new Vec3(-2, -3, 5), new Vec3(4, 0, 0), new Vec3(0, 0, -4), new LambertianMaterial(new Color(0.2, 0.8, 0.8)))
                ),
                SimpleCamera.builder()
                        .withImage(height, height)
                        .withFieldOfView(Math.toRadians(80))
                        .withPosition(new Vec3(0, 0, 9))
                        .withTarget(new Vec3(0, 0, 0))
                        .build()
        );
    }

    public static @NotNull Example getLight(int height) {
        if (height <= 0) height = 225;
        return new Example(
                new Scene(
                        new Sphere(new Vec3(0, -1000, 0), 1000, new LambertianMaterial(new Color(0.2, 0.2, 0.2))),
                        new Sphere(new Vec3(0, 2, 0), 2, new LambertianMaterial(new Color(0.2, 0.2, 0.2))),
                        new Parallelogram(new Vec3(3, 1, -2), new Vec3(2, 0, 0), new Vec3(0, 2, 0), new DiffuseLight(new Color(4.0, 4.0, 4.0))),
                        new Sphere(new Vec3(0, 7, 0), 2, new DiffuseLight(new Color(4.0, 4.0, 4.0)))
                ),
                SimpleCamera.builder()
                        .withImage(height * 16 / 9, height)
                        .withFieldOfView(Math.toRadians(20))
                        .withPosition(new Vec3(26, 3, 6))
                        .withTarget(new Vec3(0, 2, 0))
                        .build()
        );
    }

    public static @NotNull Example getCornellBox(int height) {
        if (height <= 0) height = 600;

        var red = new LambertianMaterial(new Color(.65, .05, .05));
        var white = new LambertianMaterial(new Color(.73, .73, .73));
        var green = new LambertianMaterial(new Color(.12, .45, .15));
        var light = new DiffuseLight(new Color(15.0, 15.0, 15.0));

        return new Example(
                new Scene(
                        new Box(
                                new AABB(new Vec3(0, 0, 0), new Vec3(555, 555, 555)),
                                white, white, red, green, white, null
                        ),
                        new Parallelogram(new Vec3(343, 554, 332), new Vec3(-130, 0, 0), new Vec3(0, 0, -105), light),
                        new Box(new Vec3(0, 0, 0), new Vec3(165, 330, 165), white).rotateY(Math.toRadians(15)).translate(new Vec3(265, 0, 295)),
                        new Box(new Vec3(0, 0, 0), new Vec3(165, 165, 165), white).rotateY(Math.toRadians(-18)).translate(new Vec3(130, 0, 65))
                ),
                SimpleCamera.builder()
                        .withImage(height, height)
                        .withFieldOfView(Math.toRadians(40))
                        .withPosition(new Vec3(278, 278, -800))
                        .withTarget(new Vec3(278, 278, 0))
                        .build()
        );
    }

    public static @NotNull Example getCornellBoxSmoke(int height) {
        if (height <= 0) height = 600;
        var red = new LambertianMaterial(new Color(.65, .05, .05));
        var white = new LambertianMaterial(new Color(.73, .73, .73));
        var green = new LambertianMaterial(new Color(.12, .45, .15));
        var light = new DiffuseLight(new Color(7.0, 7.0, 7.0));

        return new Example(
                new Scene(
                        new Box(
                                new AABB(new Vec3(0, 0, 0), new Vec3(555, 555, 555)),
                                white, white, red, green, white, null
                        ),
                        new Parallelogram(new Vec3(113, 554, 127), new Vec3(330, 0, 0), new Vec3(0, 0, 305), light),
                        new ConstantMedium(
                                new Box(new Vec3(0, 0, 0), new Vec3(165, 330, 165), white).rotateY(Math.toRadians(15)).translate(new Vec3(265, 0, 295)),
                                0.01, new IsotropicMaterial(Color.BLACK)
                        ),
                        new ConstantMedium(
                                new Box(new Vec3(0, 0, 0), new Vec3(165, 165, 165), white).rotateY(Math.toRadians(-18)).translate(new Vec3(130, 0, 65)),
                                0.01, new IsotropicMaterial(Color.WHITE)
                        )
                ),
                SimpleCamera.builder()
                        .withImage(height, height)
                        .withFieldOfView(Math.toRadians(40))
                        .withPosition(new Vec3(278, 278, -800))
                        .withTarget(new Vec3(278, 278, 0))
                        .build()
        );
    }

    public static @NotNull Example getDiagramm(int height) {
        if (height <= 0) height = 450;

        record Partei(String name, Color color, double stimmen) { }
        var data = List.of(
                new Partei("CDU", new Color(0x00, 0x4B, 0x76), 18.9),
                new Partei("SPD", new Color(0xC0, 0x00, 0x3C), 25.7),
                new Partei("AfD", new Color(0x80, 0xCD, 0xEC), 10.3),
                new Partei("FDP", new Color(0xF7, 0xBB, 0x3D), 11.5),
                new Partei("DIE LINKE", new Color(0x5F, 0x31, 0x6E), 4.9),
                new Partei("GRÜNE", new Color(0x00, 0x85, 0x4A), 14.8),
                new Partei("CSU", new Color(0x00, 0x77, 0xB6), 5.2)
        );
        var white = new LambertianMaterial(new Color(.99, .99, .99));

        var count = data.size();
        var size = 75d;
        var spacing = 50d;

        var x = (count + 1) * spacing + count * size + 1000;
        var y = 500 + 1000;
        var z = 2 * spacing + size + 1000;

        var objects = new ArrayList<Hittable>();
        objects.add(new Parallelogram(new Vec3(0, 0, 0), new Vec3(x, 0, 0), new Vec3(0, y, 0), white));
        objects.add(new Parallelogram(new Vec3(0, 0, 0), new Vec3(0, 0, z), new Vec3(x, 0, 0), white));
        objects.add(new Parallelogram(new Vec3(0, 0, 0), new Vec3(0, y, 0), new Vec3(0, 0, z), white));

        for (int i = 0; i < data.size(); i++) {
            var partei = data.get(i);
            objects.add(new Box(
                    new Vec3((i + 1) * spacing +       i * size, 0, spacing),
                    new Vec3((i + 1) * spacing + (i + 1) * size, partei.stimmen() * 15, spacing + size),
                    new DielectricMaterial(1.5, partei.color())
            ));
        }

        return new Example(
                new Scene(new Color(1.25, 1.25, 1.25), objects),
                SimpleCamera.builder()
                        .withImage(height * 16 / 9, height)
                        .withPosition(new Vec3(700, 250, 800))
                        .withTarget(new Vec3(500, 200, 0))
                        .withFieldOfView(Math.toRadians(40))
                        .build()
        );
    }

    public static @NotNull Example getEarth(int height) {
        if (height <= 0) height = 450;

        return new Example(
                new Scene(
                        getSkyBox(),
                        new Sphere(Vec3.ZERO, 2, new LambertianMaterial(new ImageTexture("/earthmap.jpg")))
                ),
                SimpleCamera.builder()
                        .withImage(height * 16 / 9, height)
                        .withFieldOfView(Math.toRadians(20))
                        .withPosition(new Vec3(12, 0, 0))
                        .withTarget(Vec3.ZERO)
                        .build()
        );
    }

    public static @NotNull Example getPerlin(int height) {
        if (height <= 0) height = 450;

        var material = new LambertianMaterial(new PerlinTexture(4));

        return new Example(
                new Scene(
                        getSkyBox(),
                        new Sphere(new Vec3(0, -1000, 0), 1000, material),
                        new Sphere(new Vec3(0, 2, 0), 2, material)
                ),
                SimpleCamera.builder()
                        .withImage(height * 16 / 9, height)
                        .withFieldOfView(Math.toRadians(20))
                        .withPosition(new Vec3(13, 2, 3))
                        .withTarget(Vec3.ZERO)
                        .build()
        );
    }

    public static @NotNull Example getFinal(int height) {
        if (height <= 0) height = 400;

        var objects = new ArrayList<Hittable>();
        var random = new Random(1);

        // boxes
        var boxes = new ArrayList<Hittable>();
        var ground = new LambertianMaterial(new Color(0.48, 0.83, 0.53));
        for (int i = 0; i < 20; i++) {
            for (int j = 0; j < 20; j++) {
                var w = 100.0;
                var x0 = -1000.0 + i * w;
                var z0 = -1000.0 + j * w;
                var y0 = 0.0;
                var x1 = x0 + w;
                var y1 = random.nextInt(1, 101);
                var z1 = z0 + w;
                boxes.add(new Box(new Vec3(x0, y0, z0), new Vec3(x1, y1, z1), ground));
            }
        }
        objects.add(new HittableBinaryTree(boxes));

        // light
        objects.add(new Parallelogram(
                new Vec3(123, 554, 147), new Vec3(300, 0, 0), new Vec3(0, 0, 265),
                new DiffuseLight(new Color(7., 7., 7.))
        ));

        // spheres with different materials
        objects.add(new Sphere(new Vec3(400, 400, 200), 50, new LambertianMaterial(new Color(0.7, 0.3, 0.1))));
        objects.add(new Sphere(new Vec3(260, 150, 45), 50, new DielectricMaterial(1.5)));
        objects.add(new Sphere(new Vec3(0, 150, 145), 50, new MetallicMaterial(new Color(0.8, 0.8, 0.9), 1.0)));

        // glass sphere filled with gas
        var boundary = new Sphere(new Vec3(360, 150, 145), 70, new DielectricMaterial(1.5));
        objects.add(boundary);
        objects.add(new ConstantMedium(boundary, 0.2, new IsotropicMaterial(new Color(0.2, 0.4, 0.9))));

        // put the world in a glass sphere
        objects.add(new ConstantMedium(
                new Sphere(new Vec3(0, 0, 0), 5000, new DielectricMaterial(1.5)),
                0.0001, new IsotropicMaterial(new Color(1., 1., 1.))
        ));

        // textures spheres
        objects.add(new Sphere(new Vec3(400, 200, 400), 100, new LambertianMaterial(new ImageTexture("/earthmap.jpg"))));
        objects.add(new Sphere(new Vec3(220, 280, 300), 80, new LambertianMaterial(new PerlinTexture(0.2))));

        // box from spheres
        var white = new LambertianMaterial(new Color(.73, .73, .73));
        var spheres = new ArrayList<Hittable>();
        for (int j = 0; j < 1000; j++) {
            spheres.add(new Sphere(new Vec3(random.nextDouble(165), random.nextDouble(165), random.nextDouble(165)), 10, white));
        }
        objects.add(new HittableBinaryTree(spheres).rotateY(Math.toRadians(15)).translate(new Vec3(-100, 270, 395)));

        return new Example(
                new Scene(objects),
                SimpleCamera.builder()
                        .withImage(height, height)
                        .withFieldOfView(Math.toRadians(40))
                        .withPosition(new Vec3(478, 278, -600))
                        .withTarget(new Vec3(278, 278, 0))
                        .build()
        );
    }

    private static @NotNull SkyBox getSkyBox() {
        return SkyBox.gradient(new Color(0.5, 0.7, 1.0), Color.WHITE);
    }
}
