package pathtracing;

import basics.math.algebra.Basis;
import basics.math.algebra.Ray;
import basics.math.algebra.Vector;
import geometry.scene.Hit;
import geometry.scene.Scene;
import optics.light.Color;
import optics.light.LightSource;
import optics.view.Camera;
import optics.view.PinholeCamera;
import renderer.Renderer;

import java.util.Optional;

public class Pathtracer implements Renderer {

    private final Camera camera = new PinholeCamera(new Vector(0,0,-4), Vector.origin(), 90, 1e-3, 1e3);
    private final Scene scene = Scene.generatePathtracingExampleScene();

    private double traceShadow(Vector point, Vector direction, double distance) {
        Ray shadowRay = new Ray(point, direction, 1e-3, distance);

        Optional<Hit> result = scene.intersect(shadowRay);

        if (result.isEmpty()) {
            return 1.0;
        }

        return 0.0;
    }

    private double castShadow(LightSource target, Vector point, Vector surfaceNormal) {
        Vector lightDirection = target.getDirection(point);
        double distance = target.getDistance(point);
        double shadow = traceShadow(point, lightDirection, distance);
        return Math.max(surfaceNormal.dot(lightDirection), 0.0) * shadow;
    }

    private Color traceDirect(Ray directRay, int depth) {
        if (depth == 8) {
            return Color.BLACK;
        }

        Optional<Hit> result = scene.intersect(directRay);

        if (result.isEmpty()) {
            // Background color
            return Color.BLACK;
        }

        Vector intersection = directRay.travel(result.get().getDistance());
        Vector normal = result.get().getMesh().normalAt(intersection, directRay.getDirection());

        BSDF material = (BSDF) result.get().getMesh().getMaterial();

        var diffuse = diffuseCosineWeightedBRDF(intersection, normal, material, depth);

        return diffuse.add(material.getEmission());
    }

    private Color diffuseCosineWeightedBRDF(Vector intersection, Vector normal, BSDF bsdf, int depth) {
        var next = new Ray(intersection, sampleNormalHemisphere(normal, bsdf.getRoughness()), 1e-3, 1e3);

        var cosTheta = next.getDirection().dot(normal);

        var pdf = cosTheta / Math.PI;

        var indirectLight = traceDirect(next, depth + 1).scale(pdf);

        // direct light

        Color directLight = Color.BLACK;

        for (LightSource light : scene.getLights()) {
            double directInfluence = castShadow(light, intersection, normal);

            Color incoming = bsdf.getReflectance().scale(directInfluence).mul(light.getColor());

            directLight = directLight.add(incoming);
        }

        return bsdf.getReflectance().mul(indirectLight.add(directLight));
    }

    private Vector cosineHemisphere() {
        double u1 = Math.random();
        double u2 = Math.random();

        double r = Math.sqrt(u1);
        double theta = 2 * Math.PI * u2;

        double x = r * Math.cos(theta);
        double y = r * Math.sin(theta);

        return new Vector(x, y, Math.sqrt(1 - u1));
    }

    private Vector sampleNormalHemisphere(Vector normal, double roughness) {
        Vector hemisphereSample = cosineHemisphere();

        var orthnormalBasis = Basis.constructOrthonormalBasis(normal);

        return orthnormalBasis.vectors[0].scale(hemisphereSample.z)
                .add(orthnormalBasis.vectors[1].scale(hemisphereSample.x))
                .add(orthnormalBasis.vectors[2].scale(hemisphereSample.y));
    }

    @Override
    public Color traceCameraRay(double u, double v) {
        var cameraRay = camera.generateViewRay(u, v);

        return traceDirect(cameraRay, 0);
    }

    @Override
    public Camera getCamera() {
        return camera;
    }
}
