/***********************************************/
/*           Copyright (c) 2025 Belmu          */
/*             All Rights Reserved             */
/***********************************************/

/*
    [Credits]:
        JNNGL - help with understanding focus logic (https://github.com/JNNGL)

    [References]:
        Pharr, M., Jakob, W., & Humphreys, G. (2016). Physically Based Rendering: From theory to implementation (3rd ed.). Morgan Kaufmann. Section 6.4, “Realistic Cameras.” Retrieved July 11, 2025, from https://pbr-book.org/3ed‑2018/Camera_Models/Realistic_Cameras
        Wikipedia. (2025). Ray transfer matrix analysis. https://en.wikipedia.org/wiki/Ray_transfer_matrix_analysis
*/

mat2 computeRayTransferMatrix(float wavelength) {
    mat2 rayTransferMatrix = mat2(1.0);

    for (int i = 0; i < lensElementsCount; i++) {
        LensElement element = lensSystem[i];

        bool isApertureStop = element.radius == 0.0;

        mat2 propagationMatrix = mat2(
            1.0              , 0.0, 
            element.thickness, 1.0
        );

        if (!isApertureStop) {
            float n1 = i == 0 ? airIOR : sellmeier(lensSystem[i - 1].coefficients, wavelength);
            float n2 = sellmeier(element.coefficients, wavelength);

            mat2 refractionMatrix = mat2(
                1.0, (n1 - n2) / (element.radius * n2), 
                0.0, n1 / n2
            );

            rayTransferMatrix = refractionMatrix * rayTransferMatrix;
        }

        if (i != lensElementsCount - 1)
            rayTransferMatrix = propagationMatrix * rayTransferMatrix;
    }

    return rayTransferMatrix;
}

float getFocus(mat2 rayTransferMatrix, float focusDistance) {
    const float     x = 0.001;
    float       theta = atan(x, max0(focusDistance));

    vec2 paraxialRay = rayTransferMatrix * vec2(x, theta);
    
    return -paraxialRay.x / tan(paraxialRay.y);
}

float estimateLensFocusShift(float wavelength, float focusDistance) {
    float focusShift = 0.0;
    float frontZ     = getFrontLensElementZ();

    mat2 rayTransferMatrix = computeRayTransferMatrix(wavelength);

    for (int i = 0; i < LENS_FOCUS_ITERATIONS; i++) {
        float distanceToSensor = focusDistance + frontZ - focusShift;
        float sensorZ          = getFocus(rayTransferMatrix, distanceToSensor);

        focusShift = max0(sensorZ) - rearLensElement.thickness;
    }

    return focusShift;
}
