#include "diffuse.hpp"
#include <algorithm>
#include <cmath>
#include <fstream>
#include <locale>
#include <new>
#include <numbers>
#include <utility>

namespace {
using lantern::Vec3;
Vec3 add(Vec3 a, Vec3 b) { return {a.x+b.x,a.y+b.y,a.z+b.z}; }
Vec3 scale(Vec3 v, double s) { return {v.x*s,v.y*s,v.z*s}; }
Vec3 product(Vec3 a, Vec3 b) { return {a.x*b.x,a.y*b.y,a.z*b.z}; }
double dot(Vec3 a, Vec3 b) { return a.x*b.x+a.y*b.y+a.z*b.z; }
Vec3 cross(Vec3 a, Vec3 b) { return {a.y*b.z-a.z*b.y,a.z*b.x-a.x*b.z,a.x*b.y-a.y*b.x}; }
bool finite(Vec3 v) { return std::isfinite(v.x)&&std::isfinite(v.y)&&std::isfinite(v.z); }
bool channels(Vec3 v, double upper) { return finite(v)&&v.x>=0&&v.y>=0&&v.z>=0&&v.x<=upper&&v.y<=upper&&v.z<=upper; }
std::expected<Vec3, lantern::Error> unit(Vec3 v) {
    const auto length=std::hypot(v.x,v.y,v.z);
    if (!finite(v)||!std::isfinite(length)||length<1e-12)
        return std::unexpected(lantern::Error{"Direction is outside the supported nonzero domain."});
    return Vec3{v.x/length,v.y/length,v.z/length};
}
}

namespace lantern {
std::expected<Frame, Error> frame_from_normal(Vec3 normal) {
    if (!finite(normal)||std::abs(dot(normal,normal)-1)>1e-10)
        return std::unexpected(Error{"A sampling frame requires a finite unit normal."});
    // Choose a helper away from parallel. The branch avoids a pole singularity.
    const Vec3 helper=std::abs(normal.z)<0.9 ? Vec3{0,0,1} : Vec3{0,1,0};
    const auto tangent=unit(cross(helper,normal));
    if (!tangent) return std::unexpected(tangent.error());
    return Frame{*tangent,cross(normal,*tangent),normal};
}
Vec3 to_world(const Frame& f, Vec3 local) noexcept {
    return add(add(scale(f.tangent,local.x),scale(f.bitangent,local.y)),scale(f.normal,local.z));
}
std::expected<DiffuseScene, Error> DiffuseScene::build(std::span<const Surface> surfaces, Vec3 environment) {
    if (surfaces.size()>max_spheres||!channels(environment,16))
        return std::unexpected(Error{"Too many surfaces or invalid constant environment."});
    std::array<Sphere,max_spheres> geometry{};
    for (std::size_t i=0;i<surfaces.size();++i) {
        const auto& s=surfaces[i]; const auto c=s.sphere.center;
        if (!channels(s.reflectance,1)||!channels(s.emission,16)||!finite(c)||
            std::abs(c.x)>32||std::abs(c.y)>32||std::abs(c.z)>32||
            !std::isfinite(s.sphere.radius)||s.sphere.radius<0.1||s.sphere.radius>32)
            return std::unexpected(Error{"Transport requires bounded geometry, reflectance [0,1], and emission [0,16]."});
        geometry[i]=s.sphere;
    }
    auto index=SceneIndex::build(std::span{geometry}.first(surfaces.size()));
    if (!index) return std::unexpected(index.error());
    try {
        DiffuseScene result;
        result.surfaces_.assign(surfaces.begin(),surfaces.end());
        result.index_=std::move(*index);
        result.environment_=environment;
        return result;
    } catch (const std::bad_alloc&) { return std::unexpected(Error{"Surface ownership allocation failed."}); }
}
std::expected<PathTrace, Error> DiffuseScene::trace(Ray ray, Rng& rng, unsigned max_bounces) const {
    if (max_bounces>8) return std::unexpected(Error{"At most eight diffuse scatters are supported."});
    PathTrace trace;
    Vec3 beta{1,1,1};
    for (unsigned bounce=0;bounce<=max_bounces;++bounce) {
        const auto hit=index_.closest(ray); // Also validates the initial/unit ray.
        if (!hit) return std::unexpected(hit.error());
        auto& step=trace.steps[trace.count++];
        step.throughput=beta;
        if (!*hit) {
            trace.radiance=add(trace.radiance,product(beta,environment_));
            step.radiance=trace.radiance;
            trace.end=PathEnd::environment;
            return trace;
        }
        const auto found=std::find_if(surfaces_.begin(),surfaces_.end(),[&](const Surface& s){return s.sphere.id==(*hit)->id;});
        if (found==surfaces_.end()) return std::unexpected(Error{"Scene index/material identity invariant failed."});
        step.object=found->sphere.id;
        step.point=add(ray.origin,scale(ray.direction,(*hit)->distance));
        const auto outward=unit(add(step.point,scale(found->sphere.center,-1)));
        if (!outward) return std::unexpected(outward.error());
        step.normal=dot(ray.direction,*outward)>0 ? scale(*outward,-1) : *outward;
        trace.radiance=add(trace.radiance,product(beta,found->emission));
        step.radiance=trace.radiance;
        if (bounce==max_bounces) { trace.end=PathEnd::depth_limit; return trace; }
        if (found->reflectance.x==0&&found->reflectance.y==0&&found->reflectance.z==0) {
            trace.end=PathEnd::absorbed; return trace;
        }
        const double u=rng.uniform(); // Draw dimensions in a specified order.
        const double v=rng.uniform();
        const auto sample=cosine_hemisphere(u,v);
        if (!sample) return std::unexpected(sample.error());
        const auto frame=frame_from_normal(step.normal);
        if (!frame) return std::unexpected(frame.error());
        const auto direction=unit(to_world(*frame,sample->direction));
        if (!direction) return std::unexpected(direction.error());
        step.local_direction=sample->direction;
        step.cosine=sample->direction.z;
        step.pdf=sample->pdf; // Density per steradian in an orthonormal frame.
        const auto brdf=scale(found->reflectance,1/std::numbers::pi);
        beta=product(beta,scale(brdf,step.cosine/step.pdf));
        step.throughput=beta;
        // Deliberate small-scene approximation, not an error-bounded ray offset.
        ray={add(step.point,scale(step.normal,1e-4)),*direction};
    }
    return std::unexpected(Error{"Path bound invariant failed."});
}
std::expected<DiffuseScene, Error> study_scene() {
    const std::array<Surface,4> surfaces{{
        {{1,{-0.9,0,0},0.8},{0.8,0.2,0.1},{}},
        {{2,{0.9,0,0},0.8},{0.15,0.4,0.8},{}},
        {{3,{0,-8.85,0},8},{0.65,0.65,0.65},{}},
        {{4,{0,2.5,-0.5},0.55},{0,0,0},{4,3,2}}
    }};
    return DiffuseScene::build(surfaces,{0.35,0.45,0.65});
}
std::expected<Image, Error> render_diffuse(const DiffuseScene& scene, RenderSettings settings) {
    if (settings.width==0||settings.height==0||settings.width>512||settings.height>512||
        settings.samples==0||settings.samples>128||settings.max_bounces>8||
        settings.width*settings.height>2'097'152/settings.samples)
        return std::unexpected(Error{"Render dimensions, samples or total camera-ray budget are unsupported."});
    try {
        Image image{settings.width,settings.height,std::vector<Vec3>(settings.width*settings.height)};
        const double aspect=static_cast<double>(settings.width)/static_cast<double>(settings.height);
        for (std::size_t y=0;y<settings.height;++y) for (std::size_t x=0;x<settings.width;++x) {
            const auto pixel=y*settings.width+x;
            Rng rng{settings.seed+static_cast<std::uint64_t>(pixel)};
            Vec3 sum{};
            for (std::size_t sample=0;sample<settings.samples;++sample) {
                const double u=rng.uniform(); const double v=rng.uniform();
                const auto direction=unit({(2*(static_cast<double>(x)+u)/static_cast<double>(settings.width)-1)*aspect*0.5,
                    (1-2*(static_cast<double>(y)+v)/static_cast<double>(settings.height))*0.5,-1});
                if (!direction) return std::unexpected(direction.error());
                const auto traced=scene.trace({{0,0.4,5},*direction},rng,settings.max_bounces);
                if (!traced) return std::unexpected(traced.error());
                sum=add(sum,traced->radiance);
            }
            image.linear[pixel]=scale(sum,1/static_cast<double>(settings.samples));
        }
        return image;
    } catch (const std::bad_alloc&) { return std::unexpected(Error{"Image allocation failed."}); }
}
std::expected<void, Error> write_ppm(const Image& image, const std::filesystem::path& path) {
    if (image.width==0||image.height==0||image.width>512||image.height>512||image.linear.size()!=image.width*image.height)
        return std::unexpected(Error{"Image storage and extent disagree."});
    for (const auto color:image.linear) if (!channels(color,1e6))
        return std::unexpected(Error{"Image contains invalid linear radiance."});
    std::ofstream output{path,std::ios::binary|std::ios::trunc};
    output.imbue(std::locale::classic());
    if (!output) return std::unexpected(Error{"Could not open the PPM destination."});
    output<<"P6\n"<<image.width<<' '<<image.height<<"\n255\n";
    for (const auto color:image.linear) {
        // An explicitly approximate gamma-2 display curve, not exact sRGB.
        const auto encode=[](double v){return static_cast<unsigned char>(std::lround(255*std::sqrt(std::clamp(v,0.0,1.0))));};
        const std::array<unsigned char,3> pixel{encode(color.x),encode(color.y),encode(color.z)};
        output.write(reinterpret_cast<const char*>(pixel.data()),3);
    }
    output.close();
    if (!output) return std::unexpected(Error{"PPM write failed; the destination may be incomplete."});
    return {};
}
} // namespace lantern