package eu.jonahbauer.raytracing.math;

import org.jetbrains.annotations.NotNull;

import java.util.Optional;
import java.util.random.RandomGenerator;

public record Vec3(double x, double y, double z) {
    public static final Vec3 ZERO = new Vec3(0, 0, 0);
    public static final Vec3 MAX = new Vec3(Double.MAX_VALUE, Double.MAX_VALUE, Double.MAX_VALUE);
    public static final Vec3 MIN = new Vec3(-Double.MAX_VALUE, -Double.MAX_VALUE, -Double.MAX_VALUE);
    public static final Vec3 UNIT_X = new Vec3(1, 0, 0);
    public static final Vec3 UNIT_Y = new Vec3(0, 1, 0);
    public static final Vec3 UNIT_Z = new Vec3(0, 0, 1);

    public Vec3 {
        assert Double.isFinite(x) && Double.isFinite(y) && Double.isFinite(z) : "x, y and z must be finite";
    }

    /**
     * {@return a uniformly random vector with components in the range [-1, 1)}
     */
    public static @NotNull Vec3 random(@NotNull RandomGenerator random) {
        return new Vec3(
                Math.fma(2, random.nextDouble(), -1),
                Math.fma(2, random.nextDouble(), -1),
                Math.fma(2, random.nextDouble(), -1)
        );
    }

    /**
     * {@return a uniformly random unit vector}
     */
    public static @NotNull Vec3 random(@NotNull RandomGenerator random, boolean unit) {
        if (!unit) return random(random);
        Vec3 vec;
        double squared;
        do {
            vec = random(random);
            squared = vec.squared();
        } while (squared > 1);
        return vec.div(Math.sqrt(squared));
    }

    public static @NotNull Vec3 reflect(@NotNull Vec3 vec, @NotNull Vec3 normal) {
        var factor = - 2 * normal.times(vec);
        return new Vec3(
                Math.fma(factor, normal.x(), vec.x()),
                Math.fma(factor, normal.y(), vec.y()),
                Math.fma(factor, normal.z(), vec.z())
        );
    }

    public static @NotNull Optional<Vec3> refract(@NotNull Vec3 vec, @NotNull Vec3 normal, double ri) {
        vec = vec.unit();
        var cosTheta = Math.min(- vec.times(normal), 1.0);
        var sinTheta = Math.sqrt(1 - cosTheta * cosTheta);
        if (ri * sinTheta > 1) return Optional.empty();

        var rOutPerp = vec.plus(normal.times(cosTheta)).times(ri);
        var rOutParallel = normal.times(- Math.sqrt(Math.abs(1 - rOutPerp.squared())));
        return Optional.of(rOutPerp.plus(rOutParallel));
    }

    public static @NotNull Vec3 rotate(@NotNull Vec3 vec, @NotNull Vec3 axis, double angle) {
        Vec3 vxp = axis.cross(vec);
        Vec3 vxvxp = axis.cross(vxp);
        return vec.plus(vxp.times(Math.sin(angle))).plus(vxvxp.times(1 - Math.cos(angle)));
    }

    public static double distance(@NotNull Vec3 a, @NotNull Vec3 b) {
        return a.minus(b).length();
    }

    public static @NotNull Vec3 average(@NotNull Vec3 current, @NotNull Vec3 next, int index) {
        var factor = 1d / index;
        return new Vec3(
                Math.fma(factor, next.x() - current.x(), current.x()),
                Math.fma(factor, next.y() - current.y(), current.y()),
                Math.fma(factor, next.z() - current.z(), current.z())
        );
    }

    public static @NotNull Vec3 max(@NotNull Vec3 a, @NotNull Vec3 b) {
        return new Vec3(
                Math.max(a.x(), b.x()),
                Math.max(a.y(), b.y()),
                Math.max(a.z(), b.z())
        );
    }

    public static @NotNull Vec3 min(@NotNull Vec3 a, @NotNull Vec3 b) {
        return new Vec3(
                Math.min(a.x(), b.x()),
                Math.min(a.y(), b.y()),
                Math.min(a.z(), b.z())
        );
    }

    public @NotNull Vec3 plus(double x, double y, double z) {
        return new Vec3(this.x + x, this.y + y, this.z + z);
    }

    public @NotNull Vec3 minus(double x, double y, double z) {
        return new Vec3(this.x - x, this.y - y, this.z - z);
    }

    public @NotNull Vec3 plus(@NotNull Vec3 b) {
        return new Vec3(this.x + b.x, this.y + b.y, this.z + b.z);
    }
    
    public @NotNull Vec3 minus(@NotNull Vec3 b) {
        return new Vec3(this.x - b.x, this.y - b.y, this.z - b.z);
    }
    
    public double times(@NotNull Vec3 b) {
        return this.x * b.x + this.y * b.y + this.z * b.z;
    }
    
    public @NotNull Vec3 times(double b) {
        return new Vec3(this.x * b, this.y * b, this.z * b);
    }

    public @NotNull Vec3 neg() {
        return new Vec3(-x, -y, -z);
    }

    public @NotNull Vec3 inv() {
        return new Vec3(1 / x, 1 / y, 1 / z);
    }

    public @NotNull Vec3 cross(@NotNull Vec3 b) {
        return new Vec3(
                this.y() * b.z() - b.y() * this.z(),
                this.z() * b.x() - b.z() * this.x(),
                this.x() * b.y() - b.x() * this.y()
        );
    }

    public @NotNull Vec3 div(double b) {
        return new Vec3(this.x / b, this.y / b, this.z / b);
    }

    public double squared() {
        return this.x * this.x + this.y * this.y + this.z * this.z;
    }
    
    public double length() {
        return Math.sqrt(squared());
    }

    public boolean isNearZero() {
        var s = 1e-8;
        return Math.abs(x) < s && Math.abs(y) < s && Math.abs(z) < s;
    }
    
    public @NotNull Vec3 unit() {
        return div(length());
    }

    public double get(int axis) {
        return switch (axis) {
            case 0 -> x;
            case 1 -> y;
            case 2 -> z;
            default -> throw new IndexOutOfBoundsException(axis);
        };
    }

    public @NotNull Vec3 withX(double x) {
        return new Vec3(x, y, z);
    }

    public @NotNull Vec3 withY(double y) {
        return new Vec3(x, y, z);
    }

    public @NotNull Vec3 withZ(double z) {
        return new Vec3(x, y, z);
    }
}
