# [[file:~/Dropbox/Research/TetradTurbulence/MLCoarseGrainedVGT/docs/Article/article.org::*Evaluation Criteria][Evaluation Criteria:1]]
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 createRollup(prob::LearningProblem; startTimeInd = 1, endTimeInd = nothing)
   numDims, numTsteps, numTrajs = size(prob.validationData);
   if (endTimeInd == nothing)
      endTimeInd = numTsteps;
   end
   results = Dict{String, Vector{Any}};
   predicted = generateForwardPrediction(prob, startTimeInd, endTimeInd);

   #--------------Prediction Plots--------------#
   #skewness/flatness of velocity gradient tensor
   stats_predicted = skewnessKurtosisVGT(predicted[10:18,endTimeInd,:]);
   stats_pltPredicted = scatter(stats_predicted);
   results = merge!(results, Dict("skew/kurt_predicted"=>[stats_pltPredicted]));

   #conditional mean tangents
    cmt_pltPredicted = qrConditionalMeanTangents(predicted[10:18,[1;2],:], dt = 1e-1);
    plot!(cmt_pltPredicted, xlims=(-1.0,1.0), ylims=(-1.0,1.0));
   results = merge!(results, Dict("qrCMT_predicted"=>[cmt_pltPredicted]));

   #evolved q-r pdfs
   qrPDF_predicted = qrHistogram(predicted[10:18,endTimeInd,:],nbins=30);
   qrPDF_predictedPlot = plot(qrPDF_predicted);
   results = merge!(results, Dict("qrPDF_predicted"=>[qrPDF_predictedPlot]));

   #pdfs of cosine angle between vorticity and strain-rate eigenvectors
   alignmentHists_predicted = strainVorticityAlignmentPDFs(predicted[10:18,endTimeInd,:]);
   results = merge!(results, Dict("alignmentPDFs_predicted"=>[plot(alignmentHists_predicted[i]) for i in 1:length(alignmentHists_predicted)]))

   #pdfs of s^*
   sStarHist_predicted = sStarPDF(predicted[10:18,endTimeInd,:]);
   results = merge!(results, Dict("sStarPDF_predicted"=>[plot(sStarHist_predicted)]));

   #pdf of enstrophy
   enstrophy_predicted = normalizedEnstrophyPDF(predicted[10:18,endTimeInd,:]);
   results = merge!(results, Dict("normalizedEnstrophyPDF_predicted"=>[plot(enstrophy_predicted)]));

   #pdf of dissipation
   dissipation_predicted = normalizedDissipationPDF(predicted[10:18,endTimeInd,:]);
   results = merge!(results, Dict("normalizedDissipationPDF_predicted"=>[plot(dissipation_predicted)]));

   #loss
   results = merge!(results, Dict("loss"=>[plot(prob.lossArray)]));

   #--------------Ground Truth Plots--------------#
   #skewness/flatness of velocity gradient tensor
   stats_groundTruth = skewnessKurtosisVGT(prob.validationData[10:18,endTimeInd,:]);
   stats_pltGroundTruth = scatter(stats_groundTruth);
   results = merge!(results, Dict("skew/kurt_groundTruth"=>[stats_pltGroundTruth]));   

   #conditional mean tangents
   cmt_pltGroundTruth = qrConditionalMeanTangents(prob.validationData[10:18,[1;2],:], dt = 1e-1, plotToOverlay=results["qrCMT_predicted"][1]);
    results = merge!(results, Dict("qrCMT_groundTruth"=>[cmt_pltGroundTruth]));
    

   #evolved q-r pdfs
   qrPDF_initial = qrHistogram(prob.validationData[10:18,startTimeInd,:],nbins=30);
   qrPDF_initialPlot = plot(qrPDF_initial);
   results = merge!(results, Dict("qrPDF_initial"=>[qrPDF_initialPlot]));

   qrPDF_groundTruth = qrHistogram(prob.validationData[10:18,endTimeInd,:]);
   qrPDF_groundTruthPlot = plot(qrPDF_groundTruth);
   results = merge!(results, Dict("qrPDF_groundTruth"=>[qrPDF_groundTruthPlot]));

   #pdfs of cosine angle between vorticity and strain-rate eigenvectors
   alignmentHists_groundTruth = strainVorticityAlignmentPDFs(prob.validationData[10:18,endTimeInd,:]);
   results = merge!(results, Dict("alignmentPDFs_groundTruth"=>[plot(alignmentHists_groundTruth[i]) for i in 1:length(alignmentHists_groundTruth)]))

   #pdfs of s^*
   sStarHist_groundTruth = sStarPDF(prob.validationData[10:18,endTimeInd,:]);
   results = merge!(results, Dict("sStarPDF_groundTruth"=>[plot(sStarHist_groundTruth)]));

   #pdf of enstrophy
   enstrophy_groundTruth = normalizedEnstrophyPDF(prob.validationData[10:18,endTimeInd,:]);
   results = merge!(results, Dict("normalizedEnstrophyPDF_groundTruth"=>[plot(enstrophy_groundTruth)]));

   #pdf of dissipation
   dissipation_groundTruth = normalizedDissipationPDF(prob.validationData[10:18,endTimeInd,:]);
   results = merge!(results, Dict("normalizedDissipationPDF_groundTruth"=>[plot(dissipation_groundTruth)]));

    
   return results;
end
# Evaluation Criteria:1 ends here

# [[file:~/Dropbox/Research/TetradTurbulence/MLCoarseGrainedVGT/docs/Article/article.org::*Skewness & Kurtosis of predicted velocity gradient tensor][Skewness & Kurtosis of predicted velocity gradient tensor:1]]
function vorticityElements(Aij::Array{Float64,2})
   numDims, numTrajs = size(Aij);
   @assert numDims==9;

   arr = zeros(3,numTrajs);
   for i in 1:numTrajs
      M = reshape(Aij[:,i],(3,3));
      W = 0.5(M-M');
      arr[:,i] .= [W[3,2]; W[1,3]; W[2,1]];
   end

   return reshape(arr, prod(size(arr)));
end
function onDiagonalElements(Aij::Array{Float64,2})
   numDims, numTrajs = size(Aij);
   @assert numDims==9;

   arr = zeros(3,numTrajs);
   for i in 1:numTrajs
      M = reshape(Aij[:,i],(3,3));
      arr[:,i] .= [M[1,1]; M[2,2]; M[3,3]];
   end
   return reshape(arr, prod(size(arr)));
end
function skewnessKurtosisVGT(Aij::Array{Float64,2})
   onDiagonal = onDiagonalElements(Aij);
   offDiagonal = vorticityElements(Aij);

   return Dict("onDiagonalSkewness"=>skewness(onDiagonal),
	       "onDiagonalKurtosis"=>kurtosis(onDiagonal),
	       "offDiagonalSkewness"=>skewness(offDiagonal),
	       "offDiagonalKurtosis"=>kurtosis(offDiagonal));
end
# Skewness & Kurtosis of predicted velocity gradient tensor:1 ends here

# [[file:~/Dropbox/Research/TetradTurbulence/MLCoarseGrainedVGT/docs/Article/article.org::*Conditional mean tangents][Conditional mean tangents:1]]
# Aij needs to be of size (9, 2, numTrajs), representing 2 timesteps (start/end)
function qrConditionalMeanTangents(Aij::Array{Float64,3}; 
				   nbins_r=15,
				   nbins_q=15,
				   dt=3e-4, 
				   Rrange=range(-sqrt(3)/9, sqrt(3)/9, length=nbins_r),
				   Qrange=range(-0.5, 0.5, length=nbins_q),
				   label = nothing,
				   pltToOverlay = nothing)

   numDims, numTsteps, numTrajs = size(Aij);
   @assert numDims == 9;
   rmid = Array(midpoints(Rrange));
   qmid = Array(midpoints(Qrange));
   seperatix_r = range(Rrange[1], Rrange[end], length=1000);
   seperatix_q = -((27.0/4.0).*(seperatix_r.^2)).^(1/3);

   function calc_rqFromSample(Marr)
      @assert length(Marr) == 9;
      local M = reshape(Marr, (3,3));
      local b = M/norm(M);
      local r = -tr(b^3)/3.0;
      local q = -tr(b^2)/2.0;
      return [r; q];
   end

   function inRange(M, R_bounds, Q_bounds)
      local r,q = calc_rqFromSample(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));
   delta = zeros(length(rmid), length(qmid), 2);

   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(Aij[:,1,:],
			(Rrange[i], Rrange[i+1]),
			(Qrange[j], Qrange[j+1]));
	 for l in 1:length(samples)
	    end_rq = calc_rqFromSample(Aij[:,end,samples[l]]);
	    start_rq = calc_rqFromSample(Aij[:,1,samples[l]]);
	    delta[i,j,:] .+= (end_rq .- start_rq)/dt;
	 end
	 delta[i,j,:] /= length(samples);
      end
   end

   r = reshape(r, prod(size(r)));
   q = reshape(q, prod(size(q)));
   delta = Flux.flatten(delta);

   if (pltToOverlay == nothing)
      plt = plot(seperatix_r, seperatix_q, label="seperatix");
      quiver!(plt, r,q, quiver=(delta[:,1], delta[:,2]), label=label);
      return plt;
   else
      quiver!(pltToOverlay, r,q, quiver=(delta[:,1], delta[:,2]),label=label);
      return pltToOverlay; 
   end
end
# Conditional mean tangents:1 ends here

# [[file:~/Dropbox/Research/TetradTurbulence/MLCoarseGrainedVGT/docs/Article/article.org::*Evolved $q\text{-}r$ PDFs][Evolved $q\text{-}r$ PDFs:1]]
# Aij has size = (9,numTrajs)
function qrHistogram(Aij::Array{Float64,2}; nbins = 15)

    numDims, numTrajs = size(Aij);
    @assert numDims == 9;

    r = zeros(numTrajs);
    q = zeros(numTrajs);

    Threads.@threads for i in 1:numTrajs
	M = reshape(Aij[:,i],(3,3));
	b = M/norm(M);
	q[i] = -tr(b^2)/2.0;
	r[i] = -tr(b^3)/3.0;
    end

    h = fit(Histogram, (r, q), nbins = nbins);
    return h;
end
# Evolved $q\text{-}r$ PDFs:1 ends here

# [[file:~/Dropbox/Research/TetradTurbulence/MLCoarseGrainedVGT/docs/Article/article.org::*PDFs of cosine of angle between vorticity and strain-rate eigenvectors][PDFs of cosine of angle between vorticity and strain-rate eigenvectors:1]]
#Aij is of size (9,numTrajs)
function strainVorticityAlignmentPDFs(Aij::Array{Float64,2}; nbins = 15)
   numDims, numTrajs = size(Aij);
   @assert numDims == 9;

   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

   Threads.@threads for i in 1:numTrajs
      M = reshape(Aij[:,i], (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
# PDFs of cosine of angle between vorticity and strain-rate eigenvectors:1 ends here

# [[file:~/Dropbox/Research/TetradTurbulence/MLCoarseGrainedVGT/docs/Article/article.org::*PDF of $s^*$][PDF of $s^*$:1]]
#Aij is of size (9,numTrajs)
function sStarPDF(Aij::Array{Float64,2}; nbins = 15)
   numDims, numTrajs = size(Aij);
   @assert numDims == 9;
   stats_sStar = zeros(numTrajs);

   Threads.@threads for i in 1:numTrajs
      M = reshape(Aij[:,i], (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
# PDF of $s^*$:1 ends here

# [[file:~/Dropbox/Research/TetradTurbulence/MLCoarseGrainedVGT/docs/Article/article.org::*PDF of enstrophy][PDF of enstrophy:1]]
function normalizedEnstrophyPDF(Aij::Array{Float64,2}; nbins::Int = 15)
   numDims, numTrajs = size(Aij);
   @assert numDims == 9;
   stats = [sum(abs2,Aij[:,i]) for i in 1:numTrajs];
   meanEnstrophy = mean(stats);
   stats /= meanEnstrophy;
   h = fit(Histogram, stats, nbins=nbins);
   return h;
end
# PDF of enstrophy:1 ends here

# [[file:~/Dropbox/Research/TetradTurbulence/MLCoarseGrainedVGT/docs/Article/article.org::*PDF of dissipation][PDF of dissipation:1]]
#Aij is of size (9,numTrajs)
function normalizedDissipationPDF(Aij::Array{Float64,2}; nu::Float64 = 1.1277e-3, nbins::Int =15)
   numDims, numTrajs = size(Aij);
   @assert numDims == 9;
   stats = zeros(numTrajs);

   Threads.@threads for i in 1:numTrajs
      M = reshape(Aij[:,i], (3,3));
      S = 0.5*(M+M');
      stats[i] = nu*(norm(S))^2
   end
   meanDissipation = mean(stats);
   stats /= meanDissipation;
   h = fit(Histogram, stats, nbins=nbins);

   return h;
end
# PDF of dissipation:1 ends here

# [[file:~/Dropbox/Research/TetradTurbulence/MLCoarseGrainedVGT/docs/Article/article.org::*Loss][Loss:1]]
function plotLoss(prob::LearningProblem)
   p = plot(prob.lossArray, yaxis=log);
   return p;
end
# Loss:1 ends here
