# [[file:~/Dropbox/Research/TetradTurbulence/MLCoarseGrainedVGT/docs/Article/article.org::*Problem Struct Definition][Problem Struct Definition:1]]
# unfortunately order does matter here, so be careful when modifying
mutable struct LearningProblem

    #--------------set by operator--------------#
    initialConditions;#::Array{Float64, 2}; #shape=(18,trajectories)
    trainingData;#::Array{Float64, 3}; #shape = (18, length(tsteps), trajectories)
    
    x;#::Array{Float64, 2}; #shape=(numdims, numTrajs*(numSteps-1)
    y;#::Array{Float64, 2}; #shape=(numdims, numTrajs*(numSteps-1)
    
    validationData;#::Array{Float64, 3};
    tsteps::Array{Float64, 1};
    miniBatchSize::Int;
    maxiters::Int;
    optimizer; #e.g., GradientDescent(), SGD(), etc.

    trueDrift_; #should take du, u, p, t as input
    trueDiff_;  # and modify du, return nothing

    encode_; #takes data of size(18,numTsteps,numTrajs) and outputs size(reducedDims, numTsteps, numTrajs)
    decode_; #inverse of encode
    
    modelDrift_; #should take u,p,t as params
            # and return tangent of length(u)

    modelDiff_; #should take u,p,t as params
           # and return noise of length(u)

    fullSystemDrift_; #encodes, applies drift, decodes. Mainly for post-processing.
    
    loss_;
    lossAndGrads_; #should take x,y,p as inputs
           # and return the loss, and gradient w.r.t. parameters p
    
    params::Array{Float64, 1};

    #parameter lengths
    driftLength;
    diffLength;

    #functions called by modelDrift, modelDiff_ respectively
    parameterizedDriftClosure_;
    parameterizedPrecisionMat_;

    #--------------set by training--------------#
    lossArray::Array{Float64, 1};


    #--------------set by post-processing--------------#
    plts; #dictionary

end

#--------------constructor for use by operator--------------#
function LearningProblem(initialConditions,
                         trainingData,#::Array{Float64, 3},
                         x,#::Array{Float64, 2},
                         y,#::Array{Float64, 2},
                         validationData,#::Array{Float64, 3},
                         tsteps::Array{Float64, 1},
                         miniBatchSize::Int,
                         maxiters::Int,
                         optimizer,
                         trueDrift_,
                         trueDiff_,
                         encode_,
                         decode_,
                         modelDrift_,
                         modelDiff_,
                         fullSystemDrift_,
                         loss_,
                         lossAndGrads_,
                         initialParams::Array{Float64, 1},
                         driftLength,
                         diffLength,
                         parameterizedDriftClosure_,
                         parameterizedPrecisionMat_)
    
    return LearningProblem(initialConditions,
                           trainingData,
                           x,
                           y,
                           validationData,
                           tsteps,
                           miniBatchSize,
                           maxiters,
                           optimizer,
                           trueDrift_,
                           trueDiff_,
                           encode_,
                           decode_,
                           modelDrift_,
                           modelDiff_,
                           fullSystemDrift_,
                           loss_,
                           lossAndGrads_,
                           initialParams,
                           driftLength,
                           diffLength,
                           parameterizedDriftClosure_,
                           parameterizedPrecisionMat_,
                           zeros(maxiters), #lossArray
                           Dict{String, Plots.Plot{Plots.GRBackend}}()); #plots
end

#delete large allocations before saving problem to disk.
function myserialize(path, l::LearningProblem)
    l.initialConditions = nothing;
    l.trainingData = nothing;
    l.x = nothing;
    l.y = nothing;
    serialize(path, l);
end

# used inside of problem setup to condense along time dimension to provide training data
# of size (numDims, numSamples=numTsteps*numTrajs) and to generate ground truth tangents
# to compare learning against
function parseTrainingData(data, dt)
    numDims, numSteps, numTrajs = size(data);
    
    x = zeros(numDims, numTrajs*(numSteps-1)); #the last time-step is not a valid sample
    y = [(data[:,j+1,i] - data[:,j,i])/dt for i in 1:numTrajs for j in 1:(numSteps-1)]; #generate array of array of tangents
    y = hcat(y...); #y is 9900 arrays of 18 each. We want size(y)=(18,9900)

    stride = numSteps-1;
    for i in 1:numTrajs
        for j in 1:(numSteps-1)
            x[:,stride*(i-1)+(j-1)+1] = data[:,j,i]; #the perils of 1-based indexing.
            #I suppose i,j could start at 0, and then add one on RHS.
        end
    end

    return x,y;
end
# Problem Struct Definition:1 ends here
