#include "directional_shadow.h"

#include "dx_exception.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <limits>

#include "render_type_def.h"

using namespace DirectX;

namespace {
// Direct3D 的 NDC 深度范围为 [0, 1]，不是 OpenGL 的 [-1, 1]。
std::array<XMVECTOR, 8> BuildFrustumCornersWorld(FXMMATRIX view, CXMMATRIX projection) {

    XMVECTOR determinant{};

    const XMMATRIX inverseViewProjection =
            XMMatrixInverse(&determinant, XMMatrixMultiply(view, projection));

    constexpr float ndcX[2] = {-1.0f, 1.0f};
    constexpr float ndcY[2] = {-1.0f, 1.0f};
    constexpr float ndcZ[2] = {0.0f, 1.0f};

    std::array<XMVECTOR, 8> corners{};

    std::size_t index = 0;

    for (const float z: ndcZ) {
        for (const float y: ndcY) {
            for (const float x: ndcX) {
                const XMVECTOR cornerNDC = XMVectorSet(x, y, z, 1.0f);

                corners[index++] = XMVector3TransformCoord(cornerNDC, inverseViewProjection);
            }
        }
    }

    return corners;
}
}

//= public
void DirectionalShadow::Initialize(ID3D11Device *device, const Settings &settings) {
    m_settings = settings;

    m_shadowMap.Initialize(device, settings.resolution);

    D3D11_BUFFER_DESC bufferDesc{};
    bufferDesc.ByteWidth = static_cast<UINT>(sizeof(ShadowCBData));
    bufferDesc.Usage = D3D11_USAGE_DYNAMIC;
    bufferDesc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
    bufferDesc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;

    ThrowIfFailed(device->CreateBuffer(
        &bufferDesc, nullptr, m_shadowCB.GetAddressOf()));

    D3D11_RASTERIZER_DESC rasterizerDesc{};
    rasterizerDesc.FillMode = D3D11_FILL_SOLID;
    rasterizerDesc.CullMode = D3D11_CULL_BACK;
    rasterizerDesc.FrontCounterClockwise = settings.frontCounterClockwise;

    // 写入阴影深度时进行偏移，减轻 Shadow Acne。
    rasterizerDesc.DepthBias = settings.rasterDepthBias;
    rasterizerDesc.SlopeScaledDepthBias = settings.rasterSlopeBias;
    rasterizerDesc.DepthBiasClamp = settings.rasterBiasClamp;

    rasterizerDesc.DepthClipEnable = TRUE;
    rasterizerDesc.ScissorEnable = FALSE;
    rasterizerDesc.MultisampleEnable = FALSE;
    rasterizerDesc.AntialiasedLineEnable = FALSE;

    ThrowIfFailed(device->CreateRasterizerState(
        &rasterizerDesc, m_rasterizerState.GetAddressOf()));

    D3D11_DEPTH_STENCIL_DESC depthDesc{};
    depthDesc.DepthEnable = TRUE;
    depthDesc.DepthWriteMask = D3D11_DEPTH_WRITE_MASK_ALL;
    depthDesc.DepthFunc = D3D11_COMPARISON_LESS_EQUAL;
    depthDesc.StencilEnable = FALSE;

    ThrowIfFailed(device->CreateDepthStencilState(
        &depthDesc, m_depthStencilState.GetAddressOf()));
}

void DirectionalShadow::UpdateMatrices(
    FXMMATRIX cameraView, CXMMATRIX cameraProjection, const XMFLOAT3 &lightDirection) {

    const auto corners = BuildFrustumCornersWorld(cameraView, cameraProjection);

    // 计算相机视锥的世界空间中心。
    XMVECTOR center = XMVectorZero();

    for (const XMVECTOR corner: corners) {
        center = XMVectorAdd(center, corner);
    }

    center = XMVectorScale(center, 1.0f / 8.0f);

    // 使用包围球构建稳定的正方形正交投影。
    float radius = 0.0f;

    for (const XMVECTOR corner: corners) {
        const XMVECTOR offset = XMVectorSubtract(corner, center);

        const float distance = XMVectorGetX(XMVector3Length(offset));

        radius = std::max(radius, distance);
    }

    // 对半径进行轻微量化，减少相机旋转导致的投影尺寸变化。
    radius = std::ceil(radius * 16.0f) / 16.0f;
    radius += m_settings.xyPadding;

    XMVECTOR forward = XMVector3Normalize(XMLoadFloat3(&lightDirection));

    // 当光照方向接近世界 Y 轴时，改用 Z 轴作为参考 Up，
    // 避免叉乘结果接近零。
    const XMVECTOR worldUp = XMVectorSet(0.0f, 1.0f, 0.0f, 0.0f);

    const float parallel = std::abs(XMVectorGetX(XMVector3Dot(forward, worldUp)));

    const XMVECTOR upCandidate = parallel > 0.98f ?
                                     XMVectorSet(0.0f, 0.0f, 1.0f, 0.0f) : worldUp;

    const XMVECTOR right = XMVector3Normalize(XMVector3Cross(upCandidate, forward));

    const XMVECTOR up = XMVector3Normalize(XMVector3Cross(forward, right));

    // 将阴影投影中心对齐到阴影贴图 texel。
    // 这样相机缓慢移动时，阴影不会持续发生亚像素抖动。
    const float worldUnitsPerTexel = (radius * 2.0f) / static_cast<float>(m_settings.resolution);

    float centerX = XMVectorGetX(XMVector3Dot(center, right));

    float centerY = XMVectorGetX(XMVector3Dot(center, up));

    const float centerZ = XMVectorGetX(XMVector3Dot(center, forward));

    centerX = std::round(centerX / worldUnitsPerTexel) * worldUnitsPerTexel;

    centerY = std::round(centerY / worldUnitsPerTexel) * worldUnitsPerTexel;

    XMVECTOR snappedCenter = XMVectorAdd(XMVectorScale(right, centerX), XMVectorScale(up, centerY));

    snappedCenter = XMVectorAdd(snappedCenter, XMVectorScale(forward, centerZ));

    // 光源位于视锥中心的反方向，并朝 lightDirection 看。
    const float lightDistance = radius + m_settings.zPadding;

    const XMVECTOR lightPosition = XMVectorSubtract(
        snappedCenter, XMVectorScale(forward, lightDistance));

    const XMMATRIX lightView = XMMatrixLookAtLH(
        lightPosition, snappedCenter, up);

    // 计算光源空间中的深度范围。
    float minZ = std::numeric_limits<float>::max();
    float maxZ = std::numeric_limits<float>::lowest();

    for (const XMVECTOR corner: corners) {
        XMFLOAT3 cornerLightSpace{};

        XMStoreFloat3(&cornerLightSpace, XMVector3TransformCoord(corner, lightView));

        minZ = std::min(minZ, cornerLightSpace.z);
        maxZ = std::max(maxZ, cornerLightSpace.z);
    }

    // 扩大 Z 范围，使相机视锥之外的物体仍有机会投射阴影。
    const float nearZ = std::max(0.01f, minZ - m_settings.zPadding);

    const float farZ = std::max(nearZ + 1.0f, maxZ + m_settings.zPadding);

    const XMMATRIX lightProjection = XMMatrixOrthographicOffCenterLH(
        -radius, radius, -radius, radius, nearZ, farZ);

    // HLSL 使用 row_major 和 mul(vector, matrix)，因此不转置。
    XMStoreFloat4x4(&m_cbData.lightViewProjection,
                    XMMatrixMultiply(lightView, lightProjection));

    const float texelSize = 1.0f / static_cast<float>(m_settings.resolution);

    m_cbData.shadowParams = XMFLOAT4{texelSize, texelSize,
        m_settings.minBias, m_settings.slopeBias};
}

void DirectionalShadow::BeginDepthPass(ID3D11DeviceContext *context) {
    UploadConstantBuffer(context);

    UnbindShadowSRV(context);

    context->OMSetRenderTargets(0, nullptr, m_shadowMap.DSV());

    context->ClearDepthStencilView(m_shadowMap.DSV(), D3D11_CLEAR_DEPTH, 1.0f, 0);

    const D3D11_VIEWPORT &viewport = m_shadowMap.Viewport();

    context->RSSetViewports(1, &viewport);
    context->RSSetState(m_rasterizerState.Get());

    context->OMSetDepthStencilState(m_depthStencilState.Get(), 0);

    ID3D11Buffer *shadowCB = m_shadowCB.Get();

    context->VSSetConstantBuffers(Z::RenderBindings::ShadowCB, 1, &shadowCB);

    // 不透明物体的深度阶段不需要 Pixel Shader。
    context->PSSetShader(nullptr, nullptr, 0);
    context->GSSetShader(nullptr, nullptr, 0);
}

void DirectionalShadow::EndDepthPass(ID3D11DeviceContext *context) const {
    // 解除 DSV 绑定，之后才能把阴影贴图作为 SRV 使用。
    context->OMSetRenderTargets(0, nullptr, nullptr);
}

void DirectionalShadow::BindMainPass(ID3D11DeviceContext *context) const {
    ID3D11Buffer *shadowCB = m_shadowCB.Get();

    context->VSSetConstantBuffers(Z::RenderBindings::ShadowCB, 1, &shadowCB);

    context->PSSetConstantBuffers(Z::RenderBindings::ShadowCB, 1, &shadowCB);

    ID3D11ShaderResourceView *shadowSRV = m_shadowMap.SRV();

    context->PSSetShaderResources(Z::RenderBindings::ShadowTexture, 1, &shadowSRV);

    ID3D11SamplerState *shadowSampler = m_shadowMap.Sampler();

    context->PSSetSamplers(Z::RenderBindings::ShadowSampler, 1, &shadowSampler);
}

void DirectionalShadow::UnbindShadowSRV(ID3D11DeviceContext *context) const {
    // 同一资源不能同时作为 SRV 和 DSV 使用。
    ID3D11ShaderResourceView *nullSRV = nullptr;

    context->PSSetShaderResources(Z::RenderBindings::ShadowTexture, 1, &nullSRV);
}




//= private
void DirectionalShadow::UploadConstantBuffer(ID3D11DeviceContext *context) const {
    D3D11_MAPPED_SUBRESOURCE mapped{};

    ThrowIfFailed(context->Map(m_shadowCB.Get(), 0,
                               D3D11_MAP_WRITE_DISCARD, 0, &mapped));

    std::memcpy(mapped.pData, &m_cbData, sizeof(m_cbData));

    context->Unmap(m_shadowCB.Get(), 0);
}