
#pragma once

#include <mitsuba/core/matrix.h>
#include <mitsuba/core/vector.h>
#include <mitsuba/bidir/HG.hpp>

#include <algorithm>

MTS_NAMESPACE_BEGIN

struct CoherenceProp {
    Matrix3x3 Theta{0,0,0,0,0,0,0,0,0}, Q{0,0,0,0,0,0,0,0,0};
    Spectrum hg[HGCoefs][HGCoefs] = {};
    Vector3 s{0,0,0};

    operator bool() const { return s.length()>0 && !Q.isZero() && !Theta.isZero(); }   // Checks if valid CSD
    
    auto eval(unsigned n, unsigned m, Vector3 xi) const {
        xi = Q*xi / s.length();
        
        Spectrum ret;
        for (int i=0; i<SPECTRUM_SAMPLES; i++) 
            ret[i] = hg[n][m][i]!=.0f ? hg[n][m][i] * HG(n,m,xi*k(i),Theta) : .0f;
        return ret;
    }
    auto eval(Vector3 xi) const {
        xi = Q*xi / s.length();

        Spectrum v{ .0f };
        ForEachHGCoeff(n,m)
        for (int i=0; i<SPECTRUM_SAMPLES; i++)
            v[i] += hg[n][m][i]!=.0f ? hg[n][m][i] * HG(n,m,xi*k(i),Theta) : .0f;
        return v;
    }
    
    static Float lambda(const int bin) noexcept {
        return .5f*(Spectrum{}.getBinCoverage(bin).first+Spectrum{}.getBinCoverage(bin).second)*1e-8f;
    }
    static Float k(const int bin) noexcept {
        return 2*M_PI/lambda(bin);
    }
};

inline void hg_init_coefs(CoherenceProp &csd, const Float (&src)[HGCoefs][HGCoefs]) {
    ForEachHGCoeff(n,m)
        csd.hg[n][m] = Spectrum{ src[n][m] };
}

inline auto make_csd_spherical_source(const Matrix3x3 &Q, Float r, Float l) noexcept {
    CoherenceProp csd{};
    
    csd.Theta.setZero();
    csd.Theta(0,0) = 1. / (r*r);
    csd.Theta(1,1) = 1. / (r*r);
    csd.Theta(2,2) = 1. / (l*l);
    csd.Q = Q;

    hg_init_coefs(csd, hg_coeffs_sphere);

    return csd;
}

inline auto make_csd_cylindrical_source(const Matrix3x3 &Q, Float a, Float r, Float l) noexcept {
    CoherenceProp csd{};
    
    csd.Theta.setZero();
    csd.Theta(0,0) = 1. / (a*a);
    csd.Theta(1,1) = 1. / (r*r);
    csd.Theta(2,2) = 1. / (l*l);
    csd.Q = Q;

    hg_init_coefs(csd, hg_coeffs_cylinder);

    return csd;
}

MTS_NAMESPACE_END
