using DifferentialEquations
using StochasticDiffEq

include("../src/utils.jl")

function generateNoisyTetradData(u0, tsteps)
    numTrajs = size(u0)[2];
    tspan = (tsteps[1], tsteps[end]);
    c = reshape(createNoiseRateForcingHIT(15,15), (9,9));

    function tetradDrift(u, p, t)
        numDims, numTrajs = size(u);
        rhoPiece = zeros(9,numTrajs);
        vgtPiece = zeros(9,numTrajs);
        for i in 1:numTrajs
            ρ = reshape(u[1:9],(3,3));
            M = reshape(u[10:18],(3,3));
            k = inv(ρ);
            Π = (k*k')/tr(k*k');
            rhoPiece[:,i] = reshape(M'*ρ, 9);
            vgtPiece[:,i] = reshape(-M^2 + tr(M^2)*Π, 9);
        end
        return cat(rhoPiece, vgtPiece, dims=1);
    end

    function hitDiff(u,p,t)
        return [Matrix(I,9,9) zeros(9,9)
                zeros(9,9)    c];
    end
    W = WienerProcess(0.0, u0);
    
    sdeProb = SDEProblem(tetradDrift,
                         hitDiff,
                         u0[:,1],
                         tspan,
                         noise=W,
                         noise_rate_prototype=zeros(18,18));
    
    sol = solve(sdeProb,
                u0 = u0,
                EM(),
                tstops = tsteps,
                saveat = tsteps);

    return permutedims(Array(sol), [1,3,2]);
end
