using Pkg;
using Flux;
using Zygote;
using LinearAlgebra;
using Optim;
using Plots;
using Serialization;
using PyCall;
using Statistics;
using Printf;
#Pkg.activate("./");
#Pkg.instantiate();

function calcSymmetricTensorBasis!(dest, M,rho)
    numBasisElements, numSamples = size(dest);
    @assert (numBasisElements == 42)
    @assert (numSamples == size(M)[end] == size(rho)[end]);

    Threads.@threads for i in 1:numSamples
        id = Matrix(I,3,3);
        S = 0.5*(M[:,:,i]+M[:,:,i]');
        W = 0.5*(M[:,:,i]-M[:,:,i]');
        g = rho[:,:,i]'*rho[:,:,i];

        dest[1,i]  = S                    
        dest[2,i]  = g                    
        dest[3,i]  = S^2                  
        dest[4,i]  = W^2                  
        dest[5,i]  = g^2                  
        dest[6,i]  = S*W-W*S              
        dest[7,i]  = S*g + g*S            
        dest[8,i]  = g*W - W*g            
        dest[9,i]  = S*W^2 + W^2*S        
        dest[10,i] = S*g^2 + g^2*S        
        dest[11,i] = S^2*W - W*S^2        
        dest[12,i] = g^2*W - W*g^2        
        dest[13,i] = g*S^2 + S^2*g        
        dest[14,i] = g*W^2 + W^2*g        
        dest[15,i] = S*W*g - g*W*S        
        dest[16,i] = S*g*W - W*g*S        
        dest[17,i] = g*S*W - W*S*g        
        dest[18,i] = S^2*W^2 + W^2*S^2    
        dest[19,i] = S^2*g^2 + g^2*S^2    
        dest[20,i] = W^2*g^2 + g^2*W^2    
        dest[21,i] = S^2*W*g - g*W*S^2    
        dest[22,i] = S^2*g*W - W*g*S^2    
        dest[23,i] = W^2*g*S + S*g*W^2    
        dest[24,i] = g^2*S*W - W*S*g^2    
        dest[25,i] = g^2*W*S - S*W*g^2    
        dest[26,i] = S*W^2*g + g*W^2*S    
        dest[27,i] = S*g^2*W - W*g^2*S    
        dest[28,i] = W*S^2*g - g*S^2*W    
        dest[29,i] = g*S*W^2 + W^2*S*g    
        dest[30,i] = S*g^2*S^2 + S^2*g^2*S
        dest[31,i] = W*S^2*W^2 - W^2*S^2*W
        dest[32,i] = W*g^2*W^2 - W^2*g^2*W
        dest[33,i] = g*S^2*g^2 + g^2*S^2*g      
        dest[34,i] = g*W^2*g^2 + g^2*W^2*g      
        dest[35,i] = S^2*W^2*g + g*W^2*S^2
        dest[36,i] = S^2*g^2*W - W*g^2*S^2
        dest[37,i] = W^2*g^2*S + S*g^2*W^2
        dest[38,i] = g^2*S^2*W - W*S^2*g^2
        dest[39,i] = S*W^2*g^2 + g^2*W^2*S
        dest[40,i] = g*S^2*W^2 + W^2*S^2*g
        dest[41,i] = id;
    end

    return;
end

function calcYifengSymmetricTensorBasis!(dest, M,rho)
    numBasisElements, numSamples = size(dest);
    @assert (numSamples == size(M)[end] == size(rho)[end]);

    Threads.@threads for i in 1:numSamples
        id = Matrix(I,3,3);
        S = 0.5*(M[:,:,i]+M[:,:,i]');
        W = 0.5*(M[:,:,i]-M[:,:,i]');
        g = rho[:,:,i]'*rho[:,:,i];

        dest[1,i]  = S # TB#1                   
        dest[2,i]  = S*W-W*S               # TB#2
        dest[3,i]  = S^2 - (tr(S^2)/3)*id; # TB#3                 
        dest[4,i]  = W^2 - (tr(W^2)/3)*id; # TB#4                
        dest[5,i]  = S^2*W - W*S^2                       # TB#5
        dest[6,i]  = S*W^2 + W^2*S - (2*tr(S*W^2)/3)*id; # TB#6
        dest[7,i]  = W*S*W^2 - W^2*S*W # TB#7 -- justification needed!
        dest[8,i]  = S*W*S^2 - S^2*W*S # TB#8 -- justification needed!
        dest[9,i]  = S^2*W^2 + W^2*S^2 - (2*tr(S^2*W^2)/3)*id; # TB#9
        dest[10,i] = W*S^2*W^2 - W^2*S^2*W  # TB#10

    end

    return;
end

function calcSkewSymmetricTensorBasis!(dest, M,rho)
    numBasisElements, numSamples = size(dest);
    @assert (numBasisElements == 17)
    @assert (numSamples == size(M)[end] == size(rho)[end]);

    Threads.@threads for i in 1:numSamples
        id = Matrix(I,3,3);
        S = 0.5*(M[:,:,i]+M[:,:,i]');
        W = 0.5*(M[:,:,i]-M[:,:,i]');
        g = rho[:,:,i]'*rho[:,:,i];

        dest[1,i]  = g * W + W * g                       
        dest[2,i]  = S^2 * W + W * S^2                   
        dest[3,i]  = g^2 * W + W * g^2                   
        dest[4,i]  = S * W*g + g * W * S           
        dest[5,i]  = S * g*W + W * g * S           
        dest[6,i]  = g * S*W + W * S * g           
        dest[7,i]  = S^2 * W * g + g * W*S^2       
        dest[8,i]  = S^2 * g + W * g * S^2            
        dest[9,i]  = S^2 * g * W + W * g * S^2     
        dest[10,i] = g^2 * S * W + W * S * g^2     
        dest[11,i] = g^2 * W * S + S * W * g^2     
        dest[12,i] = S * g^2 * W + W * g^2 * S     
        dest[13,i] = W * S^2 * g + g * S^2 * W     
        dest[14,i] = W * S^2 * W^2 + W^2 * S^2 * W 
        dest[15,i] = W * g^2 * W^2 + W^2 * g^2 * W 
        dest[16,i] = S^2 * g^2 * W + W * g^2 * S^2 
        dest[17,i] = g^2 * S^2 * W + W * S^2 * g^2 
    end

    return;
end

function calcTensorBasis!(mij,roij,symTB,skewSymTB)
    calcSymmetricTensorBasis!(symTB, mij, roij);
    calcSkewSymmetricTensorBasis!(skewSymTB, mij, roij);
end

function calcInvariants!(dest,S,W,g)
    numSamples = size(S)[end];

    #TODO: compare with calcYifengInvariants! and merge changes
    Threads.@threads for i in 1:numSamples
        dest[1,i]  = tr(S[i])               
        dest[2,i]  = tr(S[i]^2)             
        dest[3,i]  = tr(g[i])               
        dest[4,i]  = tr(S[i] * g[i])        
        dest[5,i]  = tr(g[i]^2)             
        dest[6,i]  = tr(S[i] * g[i]^2)      
        dest[7,i]  = tr(W[i])               
        dest[8,i]  = tr(S[i] * W[i])        
        dest[9,i]  = tr(g[i] * W[i])        
        dest[10,i] = tr(S[i] * g[i]*W[i])   
        dest[11,i] = tr(W[i]^2)             
        dest[12,i] = tr(S[i] * W[i]^2)      
        dest[13,i] = tr(g[i] * W[i]^2)      
        dest[14,i] = tr(S[i] * g[i]*W[i]^2) 
    end
end

function calcYifengInvariants!(dest,mij,roij)
    numInvariants, numSamples = size(dest);
    
    Threads.@threads for i in 1:numSamples
        S = 0.5*(mij[:,:,i]+mij[:,:,i]');
        W = 0.5*(mij[:,:,i]-mij[:,:,i]');
        dest[1,i]  = tr(S^2)      #λ_1
        dest[2,i]  = tr(W^2)      #λ_2
        dest[3,i]  = tr(S^3);     #λ_3
        dest[4,i] = tr(S * W^2)   #λ_4
        dest[5,i] = tr(W^2 * S^2) #λ_5   
    end
end


#function TBNN_LearningPressure(miniBatchSize, maxiters)
    #--------------------begin problem setup--------------------#
    basepath = "/groups/chertkov/data/tetradData/InertialRangeData/";
    vgtDataPath = string(basepath,"aij_1024_gaus8_tetrad.bin");
    rhoDataPath = string(basepath,"roij_1024_gaus8_tetrad.bin");
    pressureDataPath = string(basepath, "pij_1024_gaus8_tetrad.bin");
    viscousDataPath = string(basepath, "vis_1024_gaus8_tetrad.bin");

    numTrajs = 32768;
    numTsteps = 1500;
    numTstepsToConsider = 200;
    numTstepsToSkip = 10;
    dt = 3e-4;

    numSymBasis = 42;
    numSkewSymBasis = 17;
    numInvariants = 14;

    if (wantToGenerateData == true)
        py"""
        import numpy as np
        mij = np.fromfile(file=$vgtDataPath, dtype=float)
        mij = mij.reshape([$numTrajs, $numTsteps, 3, 3])
          
        roij = np.fromfile(file=$rhoDataPath, dtype=float)
        roij = roij.reshape([$numTrajs, $numTsteps, 3, 3])

        pij = np.fromfile(file=$pressureDataPath, dtype=float)
        pij = pij.reshape([$numTrajs, $numTsteps, 3,3]);

        vij = np.fromfile(file=$viscousDataPath, dtype=float)
        vij = vij.reshape([$numTrajs, $numTsteps, 3,3]);

        #dmdt = [[(mij[i,j,:,:] - mij[i,j-1,:,:])/$dt for i in range($numTrajs)] for j in range(1,$numTsteps)];
        #subgrid = np.array([[dmdt[j][i] - (-pij[i,j,:,:] - np.matmul(mij[i,j,:,:],mij[i,j,:,:]) + vij[i,j,:,:]) for i in range($numTrajs)] for j in range(0:($numTsteps-1))])

        tsteps = range(0,$numTstepsToConsider,$numTstepsToSkip);

        mij = mij[:, tsteps, :, :];
        roij = roij[:, tsteps, :,:];
        pij = pij[:, tsteps, :, :];
        vij = vij[:, tsteps, :, :];
        #subgrid = subgrid[:, tsteps, :, :];
        """
        
        roij = PyArray(py"roij"o);
        mij = PyArray(py"mij"o);
        pij = PyArray(py"pij"o);
        visij = PyArray(py"vij"o);
        #subgridij = PyArray(py"subgrid"o); # Π_ij = dm/dt-(-p_{ij} - m_{ik}m_{kj} + v_{ij})

        roij = permutedims(roij, [3,4,2,1]);
        mij = permutedims(mij, [3,4,2,1]);
        pij = permutedims(pij, [3,4,2,1]);
        visij = permutedims(visij, [3,4,2,1]);
        #subgridij = permutedims(subgridij, [3,4,2,1]);

        numTsteps = size(mij)[3];

        mij = reshape(mij, (3,3,numTsteps*numTrajs));
        roij = reshape(roij, (3,3,numTsteps*numTrajs));
        pij = reshape(pij, (3,3, numTsteps*numTrajs));
        visij = reshape(visij, (3,3,numTsteps*numTrajs));
        #subgridij = reshape(subgridij, [3,3,numTsteps*numTrajs]);

        #generate learning inputs
        symTB = fill(zeros(3,3), (numSymBasis, numTsteps*numTrajs));
        skewSymTB = fill(zeros(3,3), (numSkewSymBasis, numTsteps*numTrajs));
        invariants = zeros(numInvariants, numTsteps*numTrajs);
        calcYifengSymmetricTensorBasis!(symTB, mij, roij);
        calcSkewSymmetricTensorBasis!(skewSymTB, mij, roij);
        # TODO this is garbage, skewSymTB[1,:] != W calcYifengInvariants!(invariants, symTB[1,:], skewSymTB[1,:], symTB[2,:]); #S,W,g

        x = invariants;
        numTrajs = size(x)[end];
        
        #learning "labels" are pressure hessian at each timestep/traj
        y = pij;

        normalizations = hcat(mean(x[1:numInvariants,:],dims=2),
                              var(x[1:numInvariants,:],dims=2));
    end

    # tsteps = Array{Float64, 1}(range(0.0, step=dt*numTstepsToSkip, length=numTsteps));
    # dt *= numTstepsToSkip;

    # miniBatchSize = 200;
    # maxiters = 20; 
    #optimizer = Flux.Optimise.Optimiser(ExpDecay(0.001, 0.5, 1000, 1e-7), ADAM(0.001));
    optimizer = ADAM(0.0001);
    # see https://fluxml.ai/Flux.jl/stable/training/optimisers/

hiddenDim = 30;
θ, dnn_pressure = Flux.destructure(Chain(Dense(numInvariants,numSymBasis)));
#    θ, dnn_pressure = Flux.destructure(Chain(Dense(numInvariants,hiddenDim,tanh),
#                                             Dense(hiddenDim, hiddenDim, tanh),
#                                             Dense(hiddenDim,numSymBasis)));
    
    #u will be vgt (and g=ρ^Tρ eventually) in order to calculate tensor basis
    # we then transform to invariants, and pass that in to NN, recombining
    # just before we predict
    function predictPressureHessian(u,p,sampleIndxs)
        numSamples = size(u)[end];

        predCoefs = dnn_pressure(p)(u);

        intermediate = [[predCoefs[i,j]*symTB[i,sampleIndxs[j]] for i in 1:numSymBasis] for j in 1:numSamples];
        predictedPressureHessian = [sum(intermediate[j]) for j in 1:numSamples];

        return predictedPressureHessian;
    end

    function loss(features, labels, ps, sampleIndxs)
        numSamples = size(features)[end]
        pred = predictPressureHessian(features,
                                      ps,
                                      sampleIndxs);
        loss = [norm(y[:,:,i] - pred[i])^2 for i in 1:numSamples];
        return sum(loss)/numSamples;
    end

    function ψ(x)
       return log(1+x+x^2/2.0);
    end

    function dψdx(x)
       return (2(x+1))/(x^2 + 2x + 2);
    end

    function lossAndGrads(batchX, batchY, p, sampleIndxs)
        numSamples = size(batchX)[end];
        predCoefs = dnn_pressure(p)(batchX);
        intermediate = [[predCoefs[i,j]*symTB[i,sampleIndxs[j]] for i in 1:numSymBasis] for j in 1:numSamples];
        predictedPressureHessian = [sum(intermediate[j]) for j in 1:numSamples];
        δ = [(batchY[:,:,i] - predictedPressureHessian[i]) for i in 1:numSamples];
        contraction = hcat([[sum(δ[i].*symTB[j,i]) for i in 1:numSamples] for j in 1:numSymBasis]...);
        dnn_dθ(x) = pullback(ps->dnn_pressure(ps)(x), p)[2](1);
        grads = (2/numSamples)*Flux.pullback(ps->dnn_pressure(ps)(batchX),p)[2](contraction')[1]/numSamples;
        loss = sum([norm(δ[i])^2 for i in 1:numSamples])/numSamples;
        return loss,grads;
    end

    function truncatedLossAndGrads(batchX, batchY, p, sampleIndxs)
        numSamples = size(batchX)[end];
        predCoefs = dnn_pressure(p)(batchX);
        intermediate = [[predCoefs[i,j]*symTB[i,sampleIndxs[j]] for i in 1:numSymBasis] for j in 1:numSamples];
        predictedPressureHessian = [sum(intermediate[j]) for j in 1:numSamples];
        δ = [(batchY[:,:,i] - predictedPressureHessian[i]) for i in 1:numSamples];
        contraction = hcat([[-dψdx(sum(δ[i].*symTB[j,i])) for i in 1:numSamples] for j in 1:numSymBasis]...);
        grads = (2/numSamples)*pullback(ps->dnn_pressure(ps)(batchX),p)[2](contraction')[1]/numSamples;
        loss = sum([ψ(norm(δ[i])^2) for i in 1:numSamples])/numSamples;
        return loss,grads;
    end

    
    println("Setup complete, beginning learning");

function train(train_x,train_y, maxiters, miniBatchSize; seed = nothing)
    lossArray = zeros(maxiters);
    timeSinceLastUpdate = time();

    function customTrain!(ps, train_x, train_y, opt, sampleIdxs)
        l,gs = lossAndGrads(train_x, train_y, ps, sampleIdxs);
        ps += 0.1*gs;
        #Flux.update!(opt, ps, gs);
        return l;
    end

    if (seed != nothing)
        ps = seed;
    else
        ps = θ;
    end
    debugIters = 1;

    sampleIdxs = rand(1:numTrajs, miniBatchSize);
    for i in 1:maxiters
        lossArray[i] = customTrain!(ps, train_x[:,sampleIdxs], train_y[:,:,sampleIdxs], optimizer, sampleIdxs);
        if (i % debugIters == 0)
            @printf("%e   %d, %e\n", mean(lossArray[i-(debugIters-1):i]), i, (time()-timeSinceLastUpdate)/debugIters);
            timeSinceLastUpdate=time();
        end
    end
    return lossArray;
end

