package eu.jonahbauer.raytracing.scene.util;

import eu.jonahbauer.raytracing.math.*;
import eu.jonahbauer.raytracing.scene.Hittable;
import org.jetbrains.annotations.NotNull;

import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;

public final class HittableOctree extends HittableCollection {
    private final @NotNull Octree<Hittable> objects;
    private final @NotNull Optional<BoundingBox> bbox;

    public HittableOctree(@NotNull List<? extends @NotNull Hittable> objects) {
        var result = newOctree(objects);
        this.objects = result.getKey();
        this.bbox = Optional.of(result.getValue());
    }

    public HittableOctree(@NotNull Hittable @NotNull... objects) {
        this(List.of(objects));
    }

    @Override
    public void hit(@NotNull Ray ray, @NotNull State state) {
        objects.hit(ray, object -> hit(state, ray, object));
    }

    @Override
    public @NotNull Optional<BoundingBox> getBoundingBox() {
        return bbox;
    }

    private static @NotNull Entry<@NotNull Octree<Hittable>, @NotNull BoundingBox> newOctree(@NotNull List<? extends Hittable> objects) {
        Vec3 center = Vec3.ZERO, max = Vec3.MIN, min = Vec3.MAX;

        int i = 1;
        for (var object : objects) {
            var bbox = object.getBoundingBox().orElseThrow();
            center = Vec3.average(center, bbox.center(), i++);
            max = Vec3.max(max, bbox.max());
            min = Vec3.min(min, bbox.min());
        }

        var dimension = Arrays.stream(new double[] {
                Math.abs(max.x() - center.x()),
                Math.abs(max.y() - center.y()),
                Math.abs(max.z() - center.z()),
                Math.abs(min.x() - center.x()),
                Math.abs(min.y() - center.y()),
                Math.abs(min.z() - center.z())
        }).max().orElse(10);

        var out = new Octree<Hittable>(center, dimension);
        objects.forEach(object -> out.add(object.getBoundingBox().get(), object));
        return Map.entry(out, new BoundingBox(min, max));
    }
}
