# # Utility functions used in `demo_trsp.ipynb`

using MeshArrays, FortranFiles

# #### Function that reads files that were generated by `trsp_prep`

function trsp_read(mygrid::gcmgrid,dirIn::String)
    #read data from files generated by trsp_prep
    TrspX=mygrid.read(dirIn*"TrspX.bin",MeshArray(mygrid,Float32))
    TrspY=mygrid.read(dirIn*"TrspY.bin",MeshArray(mygrid,Float32))
    TauX=mygrid.read(dirIn*"TauX.bin",MeshArray(mygrid,Float32))
    TauY=mygrid.read(dirIn*"TauY.bin",MeshArray(mygrid,Float32))
    SSH=mygrid.read(dirIn*"SSH.bin",MeshArray(mygrid,Float32))

    return TrspX, TrspY, TauX, TauY, SSH
end

# #### Function that generates small binary files (2D) from large netcdf ones (4D)

function trsp_prep(mygrid::gcmgrid,GridVariables::Dict,dirOut::String="")

    #wind stress
    fileName="nctiles_climatology/oceTAUX/oceTAUX"
    oceTAUX=Main.read_nctiles(fileName,"oceTAUX",mygrid)
    fileName="nctiles_climatology/oceTAUY/oceTAUY"
    oceTAUY=Main.read_nctiles(fileName,"oceTAUY",mygrid)
    oceTAUX=mask(oceTAUX,0.0)
    oceTAUY=mask(oceTAUY,0.0)

    #sea surface height anomaly
    fileName="nctiles_climatology/ETAN/ETAN"
    ETAN=Main.read_nctiles(fileName,"ETAN",mygrid)
    fileName="nctiles_climatology/sIceLoad/sIceLoad"
    sIceLoad=Main.read_nctiles(fileName,"sIceLoad",mygrid)
    rhoconst=1029.0
    myssh=(ETAN+sIceLoad./rhoconst)
    myssh=mask(myssh,0.0)

    #seawater transports
    fileName="nctiles_climatology/UVELMASS/UVELMASS"
    U=Main.read_nctiles(fileName,"UVELMASS",mygrid)
    fileName="nctiles_climatology/VVELMASS/VVELMASS"
    V=Main.read_nctiles(fileName,"VVELMASS",mygrid)
    U=mask(U,0.0)
    V=mask(V,0.0)

    #time averaging and vertical integration
    TrspX=similar(GridVariables["DXC"])
    TrspY=similar(GridVariables["DYC"])
    TauX=similar(GridVariables["DXC"])
    TauY=similar(GridVariables["DYC"])
    SSH=similar(GridVariables["XC"])

    for i=1:mygrid.nFaces
        tmpX=mean(U.f[i],dims=4)
        tmpY=mean(V.f[i],dims=4)
        for k=1:length(GridVariables["RC"])
            tmpX[:,:,k]=tmpX[:,:,k].*GridVariables["DYG"].f[i]
            tmpX[:,:,k]=tmpX[:,:,k].*GridVariables["DRF"][k]
            tmpY[:,:,k]=tmpY[:,:,k].*GridVariables["DXG"].f[i]
            tmpY[:,:,k]=tmpY[:,:,k].*GridVariables["DRF"][k]
        end
        TrspX.f[i]=dropdims(sum(tmpX,dims=3),dims=(3,4))
        TrspY.f[i]=dropdims(sum(tmpY,dims=3),dims=(3,4))
        TauX.f[i]=dropdims(mean(oceTAUX.f[i],dims=3),dims=3)
        TauY.f[i]=dropdims(mean(oceTAUY.f[i],dims=3),dims=3)
        SSH.f[i]=dropdims(mean(myssh.f[i],dims=3),dims=3)
    end

    if !isempty(dirOut)
        write_bin(TrspX,dirOut*"TrspX.bin")
        write_bin(TrspY,dirOut*"TrspY.bin")
        write_bin(TauX,dirOut*"TauX.bin")
        write_bin(TauY,dirOut*"TauY.bin")
        write_bin(SSH,dirOut*"SSH.bin")
    end

    return TrspX, TrspY, TauX, TauY, SSH
end

# #### Function that writes a `MeshArray` object to a binary file using `FortranFile`

function write_bin(inFLD::MeshArray,filOut::String)
    recl=prod(inFLD.grid.ioSize)*4
    tmp=Float32.(convert2gcmfaces(inFLD))
    println("saving to file: "*filOut)
    f =  FortranFile(filOut,"w",access="direct",recl=recl,convert="big-endian")
    write(f,rec=1,tmp)
    close(f)
end
