export Climatology

#area of a grid box rectangular in latitude and longitude
# colatitude Î¸ âˆˆ [0,Ï€]
# longitude Ï• âˆˆ [0,2Ï€]
cellarea(r, Î”Ï•, Î¸â‚�, Î¸â‚‚) = (r^2)*abs(Î”Ï•*(cos(Î¸â‚�) - cos(Î¸â‚‚)))

checktranspose(A::Matrix)::Matrix = (size(A,1) > size(A,2)) ? collect(transpose(A)) : A

function readgrid(fn, v)::Matrix
    #read temperature file
    X = ncread(fn, v)
    #squash third dimension if necessary
    d = ndims(X)
    if d == 3
        return checktranspose(dropdims(X, dims=3))
    elseif d == 2
        return checktranspose(X)
    end
    error("unusual number of dimensions ($d) found in file $fn, variable $v")
end

struct Climatology{ğ�’¯}
    mask::BitMatrix #land mask (1 for land, 0 otherwise)
    r::Matrix{ğ�’¯} #cell runoff [m/s]
    T::Matrix{ğ�’¯} #cell temperature [K]
    A::Matrix{ğ�’¯} #cell area [m^2]
    f::Matrix{ğ�’¯} #cell land fraction [-]
    lat::Vector{ğ�’¯} #cell latitudes
    n::Int64 #number of rows/latitudes
    m::Int64 #number of columns/longitudes
end

function Base.show(io::IO, ğ�’¸::Climatology{ğ�’¯}) where {ğ�’¯}
    @unpack mask, r, T, A, f, n, m = ğ�’¸
    println(io, "$n x $m Climatology{$ğ�’¯}")
    Tmax = round(maximum(T[mask]), sigdigits=4)
    Tmin = round(minimum(T[mask]), sigdigits=4)
    println(io, "  temperature âˆˆ [$Tmin, $Tmax] K")
    rmax = round(maximum(r[mask]), sigdigits=4)
    rmin = round(minimum(r[mask]), sigdigits=4)
    println(io, "  runoff âˆˆ [$rmin, $rmax] m/s")
    println(io, "  land fraction = $(landfraction(ğ�’¸))")
    N = sum(mask)
    print(io, "  $N land cells, $(n*m) total cells")
end

Base.size(ğ�’¸::Climatology)::NTuple{2,Int64} = (ğ�’¸.n, ğ�’¸.m)

function Base.size(ğ�’¸::Climatology, dim::Int)::Int64
    @assert 1 <= dim <= 2 "Climatology has only two dimensions"
    size(ğ�’¸)[dim]
end

function Climatology(fnr::String, #runoff file name
                     vr::String,  #runoff variable name
                     nullr::Real, #runoff empty/fill value
                     convr::Real, #runoff conversion factor
                     fnT::String, #temperature file name
                     vT::String,  #temperature variable name
                     fnf::String, #land fraction file name
                     vf::String; #land fraction variable name
                     fnlat::String="", #empty will use runoff file
                     latname::String="lat")
    #read runoff grid
    r = readgrid(fnr, vr)
    #nullify null values
    @. r[r â‰ˆ convert(eltype(r), nullr)] = NaN
    #replace negatives
    @. r[r < 0] = 0
    #apply conversion factor
    r .*= convr

    #get cell coordinates
    fnlat = isempty(fnlat) ? fnr : fnlat
    lat = collect(eltype(r), ncread(fnlat, latname))

    #read temperature grid
    T = readgrid(fnT, vT)

    #read cell land fractions
    f = readgrid(fnf, vf)
    #discard nonsense
    @. f[(f < 0) | (f > 1)] = 0

    #insist on the same types for grids
    ğ�’¯ = promote_type(eltype.((r, T, f))...)
    r, T, f = convert.(Matrix{ğ�’¯}, (r, T, f))

    #demand everything is the same size
    @assert size(r) == size(T) == size(f) "Climatology arrays must have the same size/shape"
    n, m = size(r)
    @assert length(lat) == n "latitude vector length ($(length(lat))) must match number of Climatology rows ($n)"

    #make a mask from the non-NaN runoff values
    mask = @. r |> isnan |> !

    #check that the mask agrees with land fractions
    @inbounds for i âˆˆ 1:n, j âˆˆ 1:m
        if mask[i,j]
            #should have land in the cell
            @assert 0 < f[i,j] <= 1 "Climatology mask indicates a land cell, but land fraction (f) is zero"
        else
            #should have ocean in the cell
            @assert f[i,j] == 0 "Climatology mask indicates an ocean cell, but land fraction (f) is nonzero"
        end
    end

    #cell areas
    A = zeros(ğ�’¯, n, m)
    Î”Ï• = 2Ï€/m
    Î¸â‚˜ = (Ï€/180)*(lat[2:end] .+ lat[1:end-1])/2 .+ Ï€/2
    aâ‚� = cellarea(ğ��‘â‚‘, Î”Ï•, 0, Î¸â‚˜[1])
    aâ‚™ = cellarea(ğ��‘â‚‘, Î”Ï•, Î¸â‚˜[end], Ï€)
    for j âˆˆ 1:m
        A[1,j] = aâ‚�
        A[n,j] = aâ‚™
    end
    for i âˆˆ 2:n-1
        aáµ¢ = cellarea(ğ��‘â‚‘, Î”Ï•, Î¸â‚˜[i-1], Î¸â‚˜[i])
        for j âˆˆ 1:m
            A[i,j] = aáµ¢
        end
    end

    #construct
    Climatology{ğ�’¯}(mask, r, T, A, f, lat, n, m)
end

#--------------------------------------

export meanlandtemperature
export meanlandrunoff, totallandrunoff

checkcut(cut) = @assert cut >= 0 "latitude cutoff must be positive"

checksize(ğ�’¸, X) = @assert size(X) == size(ğ�’¸) "size mismatch between array and Climatology"

#already exported
function landfraction(ğ�’¸::Climatology{ğ�’¯}, cut::Real=Inf) where {ğ�’¯}
    @unpack mask, A, f, lat, n, m = ğ�’¸
    checkcut(cut)
    @multiassign num, den = zero(ğ�’¯)
    @inbounds for i âˆˆ 1:n, j âˆˆ 1:m
        #don't check the mask here b/c denominator must accumulate area
        if -cut <= lat[i] <= cut
            num += A[i,j]*f[i,j]
            den += A[i,j]
        end
    end
    return iszero(den) ? zero(ğ�’¯) : num/den
end

function landmean(X::AbstractMatrix{ğ�’¯}, ğ�’¸::Climatology{ğ�’¯}, cut::Real=Inf) where {ğ�’¯}
    @unpack mask, A, f, lat, n, m = ğ�’¸
    checksize(ğ�’¸, X)
    checkcut(cut)
    @multiassign num, den = zero(ğ�’¯)
    @inbounds for i âˆˆ 1:n, j âˆˆ 1:m
        if mask[i,j] & (-cut <= lat[i] <= cut)
            #land area of cell
            LA = A[i,j]*f[i,j]
            #contributions to averaging
            num += LA*X[i,j]
            den += LA
        end
    end
    return iszero(den) ? zero(ğ�’¯) : num/den
end

function landsum(X::AbstractMatrix{ğ�’¯}, ğ�’¸::Climatology{ğ�’¯}, cut::Real=Inf) where {ğ�’¯}
    @unpack mask, A, f, lat, n, m = ğ�’¸
    checksize(ğ�’¸, X)
    checkcut(cut)
    s = zero(ğ�’¯)
    @inbounds for i âˆˆ 1:n, j âˆˆ 1:m
        if mask[i,j] & (-cut <= lat[i] <= cut)
            #land area of cell
            LA = A[i,j]*f[i,j]
            #contribution to sum
            s += LA*X[i,j]
        end
    end
    return s
end

meanlandtemperature(ğ�’¸::Climatology, cut::Real=Inf) = landmean(ğ�’¸.T, ğ�’¸, cut)

meanlandrunoff(ğ�’¸::Climatology, cut::Real=Inf) = landmean(ğ�’¸.r, ğ�’¸, cut)

totallandrunoff(ğ�’¸::Climatology, cut::Real=Inf) = landsum(ğ�’¸.r, ğ�’¸, cut)

#already exported
meanlandlatitude(ğ�’¸::Climatology) = landmean(repeat(ğ�’¸.lat, 1, ğ�’¸.m), ğ�’¸)

#already exported
meanabslandlatitude(ğ�’¸::Climatology) = landmean(repeat(abs.(ğ�’¸.lat), 1, ğ�’¸.m), ğ�’¸)

#--------------------------------------
export meanoceandistance

function sph2cart(Î¸::T, Ï•::T) where {T}
    sâ‚œ, câ‚œ = sincos(Î¸)
    sâ‚š, câ‚š = sincos(Ï•)
    return SVector{3,T}(sâ‚œ*câ‚š, sâ‚œ*sâ‚š, câ‚œ)
end

function arclength(câ‚�::SVector{3,T}, câ‚‚::SVector{3,T}) where {T}
    d = câ‚� â‹… câ‚‚
    if d > one(T)
        return zero(T)
    elseif d < -one(T)
        return convert(T,Ï€)
    end
    acos(d)
end

function meanoceandistance(ğ�’¸::Climatology{ğ�’¯}, cut::Real=Inf, R::Real=ğ��‘â‚‘) where {ğ�’¯}
    @unpack mask, lat, n, m = ğ�’¸
    checkcut(cut)
    @assert any(mask .== 0) "no ocean cells, can't compute distances to ocean"
    #we assume longitude values cells are evenly spaced, as they ought to be
    Î”Ï• = 2Ï€/m
    Ï• = collect(ğ�’¯, LinRange(Î”Ï•/2, 2Ï€ - Î”Ï•/2, m))
    #convert the latitude values to radians
    Î¸ = collect(ğ�’¯, lat*(Ï€/180) .+ Ï€/2)
    #create a grid of cartesian coordinates for each cell
    C = Matrix{SVector{3,ğ�’¯}}(undef, n, m)
    for i âˆˆ 1:n, j âˆˆ 1:m
        C[i,j] = sph2cart(Î¸[i], Ï•[j])
    end
    #mean arclength from land cells to ocean cells
    â„’ = zero(ğ�’¯)
    count::Int64 = 0
    @inbounds for i âˆˆ 1:n, j âˆˆ 1:m
        #check if its a land cell withing the desired latitude band
        if mask[i,j] & (-cut <= lat[i] <= cut)
            #find the minimum distance to the ocean for cell i,j
            ğ�“� = Inf
            cáµ¢â±¼ = C[i,j]
            for k âˆˆ 1:n, l âˆˆ 1:m
                #cell k,l must be ocean
                if !mask[k,l] & ((i != k) | (j != l)) #don't check a cell against itself
                    ğ�“�â‚–â‚— = arclength(cáµ¢â±¼, C[k,l])
                    if ğ�“�â‚–â‚— < ğ�“�
                        ğ�“� = ğ�“�â‚–â‚—
                    end
                end
            end
            count += 1
            â„’ += ğ�“�
        end
    end
    #normalize by cell count and convert from radians to meters
    R*â„’/count
end

#--------------------------------------
export perimeter

function perimeter(mask::BitMatrix, lat::Vector{ğ�’¯}, R::Real=ğ��‘â‚‘) where {ğ�’¯}
    @assert size(mask,1) == length(lat)
    n, m = size(mask)
    B = CircularArray(mask)
    Î”Ï• = convert(ğ�’¯, 2Ï€/m)
    Î”Î¸ = convert(ğ�’¯, Ï€/n)
    c = cos.(LinRange(-Ï€/2, Ï€/2, n+1))
    p = zero(ğ�’¯)
    for i âˆˆ 1:n, j âˆˆ 1:m
        if B[i,j]
            #look left and right
            for k âˆˆ (j-1,j+1)
                if !B[i,k]
                    p += Î”Î¸*R
                end
            end
            #look "up" 
            if !B[i-1,j]
                p += Î”Ï•*R*c[i]
            end
            #look "down" 
            if !B[i+1,j]
                p += Î”Ï•*R*c[i+1]
            end
        end
    end
    return p
end

perimeter(ğ�’¸::Climatology, R::Real=ğ��‘â‚‘) = perimeter(ğ�’¸.mask, ğ�’¸.lat, R)
