using Plots, Statistics, LinearAlgebra;
using Flux, StochasticDiffEq, DiffEqBase.EnsembleAnalysis, DifferentialEquations;
using FileIO, Serialization;
using Distributions;
using Optim;
using StatsBase;
using Test;
include("learningProb.jl");


function QREvolution(prob::LearningProblem)

    # TODO, clean this up. I put this here as a hack so I don't have to change much, but these are redundant.
    # further, I've consolidated some of this functionality...it'd be nice to rewrite this using standard utility functions
    minimizer = prob.params;
    tspan = (prob.tsteps[1], prob.tsteps[end]);
    numDims, numSteps, numTrajs = size(prob.validationData);

    function fullSystemDrift(u,p,t)
        vgtPiece = prob.modelDrift_(u,p,t);

        numTrajs = size(u)[2];
        tetradPiece = zeros(9,numTrajs);
        for i in 1:numTrajs
            M = reshape(u[10:18,i], (3,3));
            ρ = reshape(u[1:9,i], (3,3));
            tetradPiece[:,i] = reshape(M'*ρ, 9);
        end

        return vcat(tetradPiece, vgtPiece);
    end

    W = WienerProcess(0.0, prob.validationData[:,1,:]);
    
    sdeProb = SDEProblem(fullSystemDrift,
                         prob.modelDiff_,
                         reshape(prob.validationData[:,1,1], (numDims,1)),
                         tspan,
                         p = prob.params,
                         noise = W,
                         noise_rate_prototype=zeros(18,18));
    initLimit = 7.0;
    finalLimit = 7.0;
    solver = SOSRI();

    
    
    #first, let's figure out binning to create the histograms using the simple validationData
    numTrajectories = size(prob.validationData)[3];
    u0 = prob.validationData[:,1,:];
    
    R_0 = zeros(numTrajectories);
    Q_0 = zeros(numTrajectories);
    R_f = zeros(numTrajectories);
    Q_f = zeros(numTrajectories);
    
    for i in 1:numTrajectories
        M_0 = reshape(u0[10:18,i],(3,3));
	M_f = reshape(prob.validationData[10:18,end,i],(3,3));

	Q_0[i] = -tr(M_0^2)/2.0;
	R_0[i] = -tr(M_0^3)/3.0;

	Q_f[i] = -tr(M_f^2)/2.0;
	R_f[i] = -tr(M_f^3)/3.0;
    end

    #trim to fit in histogram
    initBins = 20;
    badIndx = findall(x -> abs(x) > initLimit, Q_0);
    deleteat!(Q_0, badIndx);
    deleteat!(R_0, badIndx);
    badIndx = findall(x -> abs(x) > initLimit, R_0);
    deleteat!(Q_0, badIndx);
    deleteat!(R_0, badIndx);

    finalBins = 20;
    badIndx = findall(x -> abs(x) > finalLimit, Q_f);
    deleteat!(Q_f, badIndx);
    deleteat!(R_f, badIndx);
    badIndx = findall(x -> abs(x) > finalLimit, R_f);
    deleteat!(Q_f, badIndx);
    deleteat!(R_f, badIndx);

    h_0 = fit(Histogram, (R_0, Q_0), nbins = initBins);
    h_f = fit(Histogram, (R_f, Q_f), nbins = finalBins);
    scaleLimits = (0.0, maximum(h_f.weights));
    diffScaleLimits = (-maximum(h_f.weights), maximum(h_f.weights));

    # p_0 = plot(h_0, title="Initial Condition");
    # p_f = plot(h_f, title="True evolution");
    p_0 = contourf(midpoints(h_0.edges[1]),
                  midpoints(h_0.edges[2]),
                  h_0.weights',
                  clims = scaleLimits);
    p_f = contourf(midpoints(h_f.edges[1]),
                  midpoints(h_f.edges[2]),
                  h_f.weights',
                  clims = scaleLimits);

    #now, take nSDE + weights, generate new trajectories, and plot
    sol = zeros(size(prob.validationData)[1],size(prob.validationData)[2],1);
    unstableIdx = [];
    firstTime = true;

    sol = solve(sdeProb,
                u0 = prob.validationData[:,1,:],
                EM(),
                tstops = prob.tsteps,
                p=minimizer,
                saveat= prob.tsteps);

    sol = permutedims(sol, [1, 3, 2]);
    display(size(sol));
    numLearnedTraj = size(sol)[3];
    
    R_learned = zeros(numLearnedTraj);
    Q_learned = zeros(numLearnedTraj);

    for i in 1:numLearnedTraj
        M = reshape(sol[10:18,end,i],(3,3));
	Q_learned[i] = -tr(M^2)/2.0;
	R_learned[i] = -tr(M^3)/3.0;
    end

    #trim to same size as above:
    badIndx = findall(x -> abs(x) > finalLimit, Q_learned);
    deleteat!(Q_learned, badIndx);
    deleteat!(R_learned, badIndx);
    badIndx = findall(x -> abs(x) > finalLimit, R_learned);
    deleteat!(Q_learned, badIndx);
    deleteat!(R_learned, badIndx);

    display(size(R_learned));
    display(size(R_f));

    h_learned = fit(Histogram, (R_learned, Q_learned), h_f.edges);
    #use same axes as the true evolution for comparison
    p_learned = contourf(midpoints(h_learned.edges[1]),
                         midpoints(h_learned.edges[2]),
                         h_learned.weights',
                         clims = scaleLimits);

    h_diff = Histogram(h_learned.edges, h_f.weights .- h_learned.weights);
    p_diff = contourf(midpoints(h_diff.edges[1]),
                      midpoints(h_diff.edges[2]),
                      h_diff.weights',
                      clims = diffScaleLimits);
    
    #now calculate divergence
    trueDivArr = zeros(length(prob.tsteps), numTrajectories);
    learnedDivArr = zeros(length(prob.tsteps), numTrajectories);
    
    for i in 1:length(prob.tsteps)
        for j in 1:numLearnedTraj
            M_true = reshape(prob.validationData[10:18,i,j], (3,3));
            M_learned = reshape(sol[10:18,i,j], (3,3));
            
            trueDivArr[i,j] = tr(M_true);
            learnedDivArr[i,j] = tr(M_learned);
        end
    end
    p_divOverTime = scatter(1:length(prob.tsteps),
                         i -> mean(trueDivArr[i,:]),
                         label="avg div of truth");
    scatter!(1:length(prob.tsteps), i -> mean(learnedDivArr[i,:]), label="avg div of learned");
    scatter!(1:length(prob.tsteps), i -> mean(abs, learnedDivArr[i,:]), label="avg(abs(div)) of learned");
    scatter!(1:length(prob.tsteps), i -> var(learnedDivArr[i,:]),  label="var div of learned");
    plot!(legend=:topleft);
    
    p_1 = plot(p_0, p_f, p_learned, layout = (3,1));
    p_2 = plot(p_diff, p_divOverTime, layout = (2,1));
    p = plot(p_1,p_2,size = (1920,1080))

    return p;
end

function qrEvolution(prob::LearningProblem)

    # TODO, clean this up. I put this here as a hack so I don't have to change much, but these are redundant.
    # further, I've consolidated some of this functionality...it'd be nice to rewrite this using standard utility functions
    minimizer = prob.params;
    tspan = (prob.tsteps[1], prob.tsteps[end]);
    numDims, numSteps, numTrajs = size(prob.validationData);

    function fullSystemDrift(u,p,t)
        vgtPiece = prob.modelDrift_(u,p,t);

        numTrajs = size(u)[2];
        tetradPiece = zeros(9,numTrajs);
        for i in 1:numTrajs
            M = reshape(u[10:18,i], (3,3));
            ρ = reshape(u[1:9,i], (3,3));
            tetradPiece[:,i] = reshape(M'*ρ, 9);
        end

        return vcat(tetradPiece, vgtPiece);
    end

    W = WienerProcess(0.0, prob.validationData[:,1,:]);
    
    sdeProb = SDEProblem(fullSystemDrift,
                         prob.modelDiff_,
                         reshape(prob.validationData[:,1,1], (numDims,1)),
                         tspan,
                         p = prob.params,
                         noise = W,
                         noise_rate_prototype=zeros(18,18));

    solver = SOSRI();
    
    #first, let's figure out binning to create the histograms using the simple validationData
    numTrajectories = size(prob.validationData)[3];
    u0 = prob.validationData[:,1,:];
    
    r_0 = zeros(numTrajectories);
    q_0 = zeros(numTrajectories);
    r_f = zeros(numTrajectories);
    q_f = zeros(numTrajectories);
    
    for i in 1:numTrajectories
        M_0 = reshape(u0[10:18,i],(3,3));
	M_f = reshape(prob.validationData[10:18,end,i],(3,3));

        b_0 = M_0/norm(M_0);
        b_f = M_f/norm(M_f);

	q_0[i] = -tr(b_0^2)/2.0;
	r_0[i] = -tr(b_0^3)/3.0;

	q_f[i] = -tr(b_f^2)/2.0;
	r_f[i] = -tr(b_f^3)/3.0;
    end

    numBins = 30;
    h_0 = fit(Histogram, (r_0, q_0), nbins = numBins);
    h_f = fit(Histogram, (r_f, q_f), nbins = numBins);
    scaleLimits = (0.0, maximum(h_f.weights));
    diffScaleLimits = (-100,100);
    # p_0 = plot(h_0, title="Initial Condition");
    # p_f = plot(h_f, title="True evolution");
    p_0 = contour(midpoints(h_0.edges[1]),
                  midpoints(h_0.edges[2]),
                  h_0.weights',
                  clims = scaleLimits);
    p_f = contour(midpoints(h_f.edges[1]),
                  midpoints(h_f.edges[2]),
                  h_f.weights',
                  clims = scaleLimits);

    sol = solve(sdeProb,
                u0 = prob.validationData[:,1,:],
                EM(),
                tstops = prob.tsteps,
                p=minimizer,
                saveat= prob.tsteps);

    sol = permutedims(sol, [1, 3, 2]);
    display(size(sol));
    numLearnedTraj = size(sol)[3];
    
    r_learned = zeros(numLearnedTraj);
    q_learned = zeros(numLearnedTraj);

    for i in 1:numLearnedTraj
        M = reshape(sol[10:18,end,i],(3,3));
        b = M/norm(M);
	q_learned[i] = -tr(b^2)/2.0;
	r_learned[i] = -tr(b^3)/3.0;
    end

    h_learned = fit(Histogram, (r_learned, q_learned), h_f.edges);
    #use same axes as the true evolution for comparison
    #p_learned = plot(h_learned, xlims=xlims(p_f), ylims=ylims(p_f), title="Learned evolution");
    p_learned = contour(midpoints(h_learned.edges[1]),
                        midpoints(h_learned.edges[2]),
                        h_learned.weights',
                        clims = scaleLimits);

    h_diff = Histogram(h_learned.edges, h_f.weights .- h_learned.weights);
    p_diff = contourf(midpoints(h_diff.edges[1]),
                      midpoints(h_diff.edges[2]),
                      h_diff.weights',
                      clims = diffScaleLimits);
        
    p = plot(p_0, p_f, p_learned, p_diff, layout=(2,2), size = (1920,1080))

    return p;
end

function rhoStatisticsEvolution(prob::LearningProblem)
    tspan = (prob.tsteps[1], prob.tsteps[end]);
    numDims, numSteps, numTrajs = size(prob.validationData);

    function fullSystemDrift(u,p,t)
        vgtPiece = prob.modelDrift_(u,p,t);

        numTrajs = size(u)[2];
        tetradPiece = zeros(9,numTrajs);
        for i in 1:numTrajs
            M = reshape(u[10:18,i], (3,3));
            ρ = reshape(u[1:9,i], (3,3));
            tetradPiece[:,i] = reshape(M'*ρ, 9);
        end

        return vcat(tetradPiece, vgtPiece);
    end

    sdeProb = SDEProblem(fullSystemDrift,
                         prob.modelDiff_,
                         reshape(prob.validationData[:,1,1], (numDims,1)),
                         tspan,
                         p = prob.params);
    solver = SOSRI();
    
    #first, let's figure out binning to create the histograms using the simple validationData
    numTrajectories = size(prob.validationData)[3];
    u0 = prob.validationData[:,1,:];
    
    a_0 = zeros(numTrajectories);
    b_0 = zeros(numTrajectories);
    c_0 = zeros(numTrajectories);
    a_f = zeros(numTrajectories);
    b_f = zeros(numTrajectories);
    c_f = zeros(numTrajectories);
    
    Threads.@threads for i in 1:numTrajectories
        rho_0 = reshape(prob.validationData[1:9,1,i], (3,3));
        rho_f = reshape(prob.validationData[1:9,end,i],(3,3));
        a_0[i], b_0[i], c_0[i] = eigen(rho_0'rho_0).values;
        a_f[i], b_f[i], c_f[i] = eigen(rho_f'rho_f).values;
    end

    numBins = 50;
    h_a0 = fit(Histogram, a_0, nbins=numBins);
    h_b0 = fit(Histogram, b_0, nbins=numBins);
    h_c0 = fit(Histogram, c_0, nbins=numBins);

    p_0 = plot(h_a0, label="a_0", fillalpha=0.4);
    plot!(p_0, h_b0, label="b_0", fillalpha=0.4);
    plot!(p_0, h_c0, label="c_0", fillalpha=0.4, title="Initial Distributions of Eigenvalues");

    h_af = fit(Histogram, a_f, nbins=numBins);
    h_bf = fit(Histogram, b_f, nbins=numBins);
    h_cf = fit(Histogram, c_f, nbins=numBins);

    p_f = plot(h_af, label="a_f", fillalpha=0.4);
    plot!(p_f, h_bf, label="b_f", fillalpha=0.4);
    plot!(p_f, h_cf, label="c_f", fillalpha=0.4, title="Final Distributions of Eigenvalues");

    ylim = maximum([h_af.weights; h_bf.weights; h_cf.weights]);
    xlim = max(h_af.edges[1][end], h_bf.edges[1][end], h_cf.edges[1][end]);

    sol = solve(sdeProb,
                u0 = prob.validationData[:,1,:],
                EM(),
                tstops = prob.tsteps,
                p=prob.params,
                saveat= prob.tsteps);

    sol = permutedims(sol, [1, 3, 2]);
    display(size(sol));
    numLearnedTraj = size(sol)[3];
    
    a_learned = zeros(numLearnedTraj);
    b_learned = zeros(numLearnedTraj);
    c_learned = zeros(numLearnedTraj);

    Threads.@threads for i in 1:numLearnedTraj
        rho = reshape(sol[1:9,end,i], (3,3));
        a_learned[i], b_learned[i], c_learned[i] = eigen(rho'rho).values;
    end

    h_al = fit(Histogram, a_learned, nbins=numBins);
    h_bl = fit(Histogram, b_learned, nbins=numBins);
    h_cl = fit(Histogram, c_learned, nbins=numBins);

    p_learned = plot(h_al, label="a_f", fillalpha=0.4);
    plot!(h_bl, label="b_f", fillalpha=0.4);
    plot!(h_cl, label="c_f", fillalpha=0.4, xlims=(0.0,xlim), ylims=(0.0,ylim));
        
    p = plot(p_0, p_f, p_learned, layout=(1,3), size = (1920,800))

    return p;
end

# arr is a tuple of things to delete, say (R,Q,weights)
function trim(arr,Rmin,Rmax,Qmin,Qmax)
    R = arr[1];
    Q = arr[2];
    
    badIndx = findall(x -> x>Rmax, R);
    for i in 1:length(arr)
        deleteat!(arr[i],badIndx);
    end
    
    badIndx = findall(x -> x<Rmin, R);
    for i in 1:length(arr)
        deleteat!(arr[i],badIndx);
    end
    
    badIndx = findall(x -> x>Qmax, Q);
    for i in 1:length(arr)
        deleteat!(arr[i],badIndx);
    end
    
    badIndx = findall(x -> x<Qmin, Q)
    for i in 1:length(arr)
        deleteat!(arr[i],badIndx);
    end
end
# data   - DNS, ODE data, size = (dims(18), numTimeSteps, numTraj)
# Rlimit - where to trim R
# Qlimit - where to trim Q
# bins  - 2 entry tuple
# generateWeights - function that takes data as input, and outputs R,Q,weights arrays.
    # function generateWeights(data)
    #     numTraj = size(data)[3];
    #     R = zeros(numTraj);
    #     Q = zeros(numTraj);
    #     w = zeros(numTraj);
    #     for i in 1:numTraj
    #         M = reshape(data[10:18,end,i], (3,3));
    #         Q[i] = -tr(M^2)/2.0;
    #         R[i] = -tr(M^3)/3.0;
    #         w[i] = 1.0; #equal weight to each, i.e, counting metric
    #     end        
    #     return R,Q;
    # end
    # function generateEnergyFluxDensity(data)
    #     numTraj = size(data)[3];
    #     R = zeros(numTraj);
    #     Q = zeros(numTraj);
    #     e = zeros(numTraj);
    
    #     for i in 1:numTraj
    #         M = reshape(data[10:18, 1, i], (3,3));
    #         Q[i] = -tr(M^2)/2.0;
    #         R[i] = -tr(M^3)/3.0;
    #         e[i] =  tr(M^2*M');
    #     end
    #     return R,Q,e;
    # end
function drawRQContour(data, Rlimit, Qlimit, numBins, generateWeights)

    R,Q,w = generateWeights(data);
    
    trim((R,Q,w), -Rlimit, Rlimit, -Qlimit, Qlimit);

    edgeQ = range(-Qlimit, Qlimit, length=numBins);
    edgeR = range(-Rlimit, Rlimit, length=numBins);
    h = fit(Histogram, (R,Q), weights(w), (edgeR, edgeQ), closed=:left);
    
    c = contour(midpoints(h.edges[1]),
                midpoints(h.edges[2]),
                h.weights');

    return c;
end

# function drawVectorField(data, Rlimits, Qlimits, tangentFuncs)
#     for i in 1:length(tangentFuncs)
#         R = zeros(size(data)[3]);
#         Q = zeros(size(data)[3]);
#         for i in 1:size(data)[3]
#             M = reshape(data[10:18,1,i],(3,3));
#             R[i] = -tr(M^3)/3.0;
#             Q[i] = -tr(M^2)/2.0;
#         end
#         Rmin = Rlimits[1];
#         Rmax = Rlimits[2];
#         Qmin = Qlimits[1];
#         Qmax = Qlimits[2];
#         trim(R,Q,Rlimits[1],Rlimits[2],Qlimits[1],Qlimits[2]);
        
#     end
# end

#TODO having this around is useful on a variety of levels, perhaps export from generateData.jl
function trueDrift(u, p, t)
    ρ = reshape(u[1:9],(3,3));
    M = reshape(u[10:18],(3,3));
    k = inv(ρ);
    Π = (k*k')/tr(k*k');
    placeHolder = [(M'*ρ); (-M^2 + tr(M^2)*Π)];
    return [reshape(placeHolder[1:3,1:3],9); reshape(placeHolder[4:6,1:3],9)] #du/dt = [ dρdt; dMdt]
end

#tangent functions is an array of functions that map from \R^18 → \R^2 and return [dR/dt dQ/dt]
#function drawVectorField(data, Rrange, Qrange, tangentFuncs, dt, timepoint; titles = nothing)
function drawVectorField(prob::LearningProblem)

    #---------------------------TODO CLEAN ALL THIS UP----------------------#
    #TODO this is really ugly
    function learnedRQEvolution(u)
        M = reshape(u[10:18], (3,3));
        dM = reshape(prob.drift_(u, prob.trainingResults.minimizer, 0.0)[10:18], (3,3));
        dQ = -tr(M*dM);
        dR = -tr(M^2*dM)/3.0;
        return [dR; dQ];
    end
    function trueRQEvolution(u)
        dM = reshape(trueDrift(u,0.0,0.0)[10:18], (3,3));
        M = reshape(u[10:18], (3,3));
        dQ = -tr(M*dM);
        dR = -tr(M^2*dM)/3.0;
        return [dR; dQ];
    end
    function diff(u)
        return learnedRQEvolution(u) - trueRQEvolution(u);
    end
    
    # TODO, clean this up. I put this here as a hack so I don't have to change much, but these are redundant.
    data = prob.validationData; #pretty sure for this Eularian view, I only want first timestep...TODO resolve
    Rrange = range(-15.0, 15.0, step=1.0);
    Qrange = range(-15.0, 15.0, step=1.0);
    tangentFuncs = [trueRQEvolution, learnedRQEvolution, diff]
    dt = prob.tsteps[2] - prob.tsteps[1];
    timepoint = 1;
    titles = ["True RQ Evolution" "Learned RQ Evolution" "Difference"]
    #---------------------------TODO CLEAN ALL THIS UP END----------------------#

    
    rMid = midpoints(Rrange);
    qMid = midpoints(Qrange);
    pltArr = [];
    seperatix_r = range(Rrange[1], Rrange[end], length = 1000);
    seperatix_q = -((27.0/4.0).*(seperatix_r.^2)).^(1/3);
    #M should be length 9 array
    function calcRQFromSample(Marr)
        M = reshape(Marr, (3,3));
        R = -tr(M^3)/3.0;
        Q = -tr(M^2)/2.0;
        return R,Q;
    end

    function inRange(M, Rrange, Qrange)
        R,Q = calcRQFromSample(M);
        return (Rrange[1] <= R < Rrange[2]) & (Qrange[1] <= Q < Qrange[2]);
    end

    #M_vals should be 9xnumTraj array
    function find(M_vals, Rrange, Qrange)
        numTraj = size(M_vals)[2];
        trueIdx = [];
        for i in 1:numTraj
            if (inRange(M_vals[:,i], Rrange, Qrange))
                append!(trueIdx, i);
            end
        end
        return Array{Int64, 1}(trueIdx);
    end
    
    for k in 1:length(tangentFuncs)
        df = tangentFuncs[k];
        #I do a reallocation per loop because I mess up the arrays at the end. Not clean but...
        R = zeros(length(rMid),length(qMid));
        Q = zeros(length(rMid),length(qMid));
        Δ = zeros(length(rMid),length(qMid),2); #a trailing 2 here so each point (Q,R) has 2d vector

        for i in 1:length(qMid)
            for j in 1:length(rMid)
                R[i,j] = rMid[j];
                Q[i,j] = qMid[i];
                Rmin,Rmax = [Rrange[j], Rrange[j+1]];
                Qmin,Qmax = [Qrange[i], Qrange[i+1]];
                samples = find(data[10:18,timepoint,:], (Rmin,Rmax), (Qmin,Qmax));
                for l in 1:length(samples)
                    Δ[i,j,:] += dt*df(data[:,timepoint,samples[l]]);
                end
                Δ[i,j,:] /= length(samples);
            end
        end
        R = reshape(R, prod(size(R)));
        Q = reshape(Q, prod(size(Q)));
        Δ = flatten(Δ);
        dR = Δ[:,1];
        dQ = Δ[:,2];
        push!(pltArr, quiver(R,Q,quiver=(dR,dQ)));
        plot!(pltArr[k], seperatix_r, seperatix_q);
        if(titles != nothing && length(titles) == length(tangentFuncs))
            plot!(pltArr[k], title=titles[k]);
        end
    end


    return pltArr;
end

# Notice here that qr vs QR isn't specified. If the ground truth data is normalized, (i.e., prob.validationData = [ρ; b] with norm(b)=1), then you'll obtain qr.
# TODO ^this probably isn't the best way to do this...
function drawConditionalMeanTangents(prob::LearningProblem, Rrange::StepRangeLen, Qrange::StepRangeLen)
    dt = prob.tsteps[2] - prob.tsteps[1];
    timepoint = 1;
    timestep = 1;
    tspan = (prob.tsteps[timepoint], prob.tsteps[timepoint+timestep]);

    W = WienerProcess(0.0, prob.validationData[:,1,:]);
    sdeProb = SDEProblem(prob.fullSystemDrift_,
                         prob.modelDiff_,
                         prob.validationData[:,1,1],
                         tspan,
                         p = prob.params,
                         noise = W,
                         noise_rate_prototype=zeros(18,18));
    predicted = Array(solve(sdeProb,
                            u0 = prob.validationData[:,1,:],
                            EM(),
                            tstops = prob.tsteps,
                            p=prob.params,
                            saveat = prob.tsteps));
    predicted = permutedims(predicted, [1,3,2]);
    

    rMid = Array(midpoints(Rrange));
    qMid = Array(midpoints(Qrange));
    pltArr = [];
    seperatix_r = range(Rrange[1], Rrange[end], length=1000);
    seperatix_q = -((27.0/4.0).*(seperatix_r.^2)).^(1/3);

    function calcRQFromSample(Marr)
        @assert length(Marr) == 9;
        local M = reshape(Marr, (3,3));
        local R = -tr(M^3)/3.0;
        local Q = -tr(M^2)/2.0;
        return [R; Q];
    end

    function inRange(M, R_bounds, Q_bounds)
        local R,Q = calcRQFromSample(M);
        return (R_bounds[1] <= R < R_bounds[2]) & (Q_bounds[1] <= Q < Q_bounds[2]);
    end

    function find(M_vals, R_bounds, Q_bounds)
        numTrajs = size(M_vals)[2];
        trueIdx = [inRange(M_vals[:,i], R_bounds, Q_bounds) for i in 1:numTrajs];
        return findall(trueIdx);
    end

    R = zeros(length(rMid), length(qMid));
    Q = zeros(length(rMid), length(qMid));
    groundTruthΔ = zeros(length(rMid), length(qMid),2);
    modelDriftΔ = zeros(length(rMid), length(qMid),2);

    function calcRQTangentFrom_dM(dM,M) #dM,M have size (9)
        dm = reshape(dM, (3,3));
        m = reshape(M, (3,3));
        dQ = -tr(dm*m);
        dR = -tr(dm*m*m);
        return [dR; dQ];
    end

    for i in 1:length(rMid)
        for j in 1:length(qMid)
            R[i,j] = rMid[i];
            Q[i,j] = qMid[j];
        end
    end
    
    Threads.@threads for i in 1:length(rMid)
        for j in 1:length(qMid)
            samples = find(prob.validationData[10:18,timepoint,:],
                           (Rrange[i], Rrange[i+1]),
                           (Qrange[j], Qrange[j+1]));
            for l in 1:length(samples)
                intermediate1 = calcRQFromSample(prob.validationData[10:18,timepoint+timestep,samples[l]]);
                intermediate2 = calcRQFromSample( prob.validationData[10:18,timepoint,samples[l]]);
                groundTruthΔ[i,j,:] .+= (intermediate1 .- intermediate2)/(dt*timestep);
                dM = (predicted[10:18,end,samples[l]] - predicted[10:18, 1, samples[l]])/(dt*timestep);
                modelDriftΔ[i,j,:] += calcRQTangentFrom_dM(dM,predicted[10:18,1,samples[l]]);
            end
            groundTruthΔ[i,j,:] /= length(samples);
            modelDriftΔ[i,j,:] /= length(samples);
        end
    end
    R = reshape(R, prod(size(R)));
    Q = reshape(Q, prod(size(Q)));
    modelDriftΔ = dt*Flux.flatten(modelDriftΔ);
    groundTruthΔ = dt*Flux.flatten(groundTruthΔ);

    plt = plot(seperatix_r, seperatix_q, label="separatix");
    quiver!(plt, R,Q,quiver=(groundTruthΔ[:,1],groundTruthΔ[:,2]),label="groundTruth",color=:orange);
    quiver!(plt, R,Q,quiver=(modelDriftΔ[:,1], modelDriftΔ[:,2]),label="predicted",color=:green);
    annotate!(plt, :bottomright, "Predicted", :green);
    annotate!(plt, :bottomright, "\nGround Truth", :orange);    

    return plt;
end

function strainVorticityAlignmentPDFs(Aij::Array{Float64,2}; nbins = 15)
   stats_e1 = zeros(numTrajs);
   stats_e2 = zeros(numTrajs);
   stats_e3 = zeros(numTrajs);

   function costheta(e_i, omega)
      return (dot(e_i,omega))/(norm(e_i)*norm(omega));
   end

   for i in 1:numTrajs
      M = reshape(Aij, (3,3));
      S = 0.5*(M+M');
      W = 0.5*(M-M');
      omega = 2*[ W[3,2]; W[1,3]; W[2,1] ];
      vecs = eigen(S).vectors;
      e1 = vecs[:,3];
      e2 = vecs[:,2];
      e3 = vecs[:,1];
      stats_e1[i] = costheta(e1,omega);
      stats_e2[i] = costheta(e2,omega);
      stats_e3[i] = costheta(e3,omega);
   end
   h1 = fit(Histogram, stats_e1, nbins=nbins);
   h2 = fit(Histogram, stats_e2, nbins=nbins);
   h3 = fit(Histogram, stats_e3, nbins=nbins);

   return h1, h2, h3;
end

function generateForwardPrediction(prob::LearningProblem, startTimeInd, endTimeInd)
    dt = prob.tsteps[2] - prob.tsteps[1];
    tspan = (prob.tsteps[startTimeInd], prob.tsteps[endTimeInd]);

    W = WienerProcess(0.0, prob.validationData[:,1,:]);
    sdeProb = SDEProblem(prob.fullSystemDrift_,
                         prob.modelDiff_,
                         prob.validationData[:,1,1],
                         tspan,
                         p = prob.params,
                         noise = W,
                         noise_rate_prototype=zeros(18,18));
    predicted = Array(solve(sdeProb,
                            u0 = prob.validationData[:,1,:],
                            EM(),
                            tstops = prob.tsteps,
                            p=prob.params,
                            saveat = prob.tsteps));
    predicted = permutedims(predicted, [1,3,2]);

    return predicted;
end

function plotStrainRateVorticityAlignments(prob::LearningProblem; timepoint=1; nbins=15)
    predicted = generateForwardPrediction(prob, 1, timepoint);

    groundTruth_e1PDF, groundTruth_e2PDF, groundTruth_e3PDF = strainVorticityAlignmentPDFs(prob.validationData[10:18,timepoint,:],nbins=nbins);
    predicted_e1PDF, predicted_e2PDF, predicted_e3PDF = strainVorticityAlignmentPDFs(predicted[10:18,end,:],nbins=nbins);

    p = plot(midpoints(groundTruth_e1PDF.edges[1]), groundTruth_e1PDF.weights);
    plot!(p, midpoints(groundTruth_e2PDF.edges[1]), groundTruth_e2PDF.weights);
    plot!(p, midpoints(groundTruth_e3PDF.edges[1]), groundTruth_e3PDF.weights);

    plot!(p, midpoints(predicted_e1PDF.edges[1]), predicted_e1PDF.weights);
    plot!(p, midpoints(predicted_e2PDF.edges[1]), predicted_e2PDF.weights);
    plot!(p, midpoints(predicted_e3PDF.edges[1]), predicted_e3PDF.weights);

    return p;
end

function sStarPDF(Aij::Array{Float64,2}; nbins = 15)
   stats_sStar = zeros(numTrajs);

   Threads.@threads for i in 1:numTrajs
      M = reshape(Aij, (3,3));
      S = 0.5*(M+M');
      values = eigen(S).values;
      alpha = values[3];
      beta  = values[2];
      gamma = values[1];
      stats_sStar[i] = -(3*sqrt(6)*alpha*beta*gamma)/((alpha^2+beta^2+gamma^2)^(3/2));
   end
   h = fit(Histogram, stats_sStar, nbins=nbins);

   return h;
end

function niceLossPlot(prob::LearningProblem)
    p = plot(prob.lossArray, color=:blue, title="Loss", xlabel="epoch", label="observed loss", size=(1920,1080));
    return p;
end

function postProcess(prob::LearningProblem)
    p = plot(prob.lossArray,
             title="Loss",
             xlabel="epoch");
    plts = drawVectorField(prob);
    push!(plts, QREvolution(prob));
    push!(plts, plot(prob.lossArray,
                       yaxis=:log,
                       title="Loss for minibatch size of $(prob.miniBatchSize), ODE",
                       xlabel="epoch"));

    return plts;
end
