# temp allocation
abstract type ODEAllocation end

struct TractionRateAllocFFTConv{T, U, P<:FFTW.Plan} <: ODEAllocation
    dÏ„_dt::T # traction rate of interest
    relv::T # relative velocity including zero-padding
    relvnp::T # relative velocity excluding zero-padding area
    dÏ„_dt_dft::U # stress rate in discrete fourier domain
    relv_dft::U # relative velocity in discrete fourier domain
    dÏ„_dt_buffer::T # stress rate including zero-padding zone for fft
    pf::P # real-value-FFT forward operator
end

struct StressRateAllocMatrix{T} <: ODEAllocation
    reldÏµ::T
end

function gen_alloc(::Val{:BEMFault}, nx::I, nÎ¾::I; T=Float64, fftw_flags::UInt32=FFTW.PATIENT) where {I<:Integer}
    x1 = Matrix{T}(undef, 2 * nx - 1, nÎ¾)
    p1 = plan_rfft(x1, 1; flags=fftw_flags)

    return TractionRateAllocFFTConv(
        Matrix{T}(undef, nx, nÎ¾),
        zeros(T, 2nx-1, nÎ¾), zeros(T, nx, nÎ¾), # for relative velocity, including zero
        [Matrix{Complex{T}}(undef, nx, nÎ¾) for _ âˆˆ 1: 2]...,
        Matrix{T}(undef, 2nx-1, nÎ¾),
        p1)
end

function gen_alloc(::Val{:BEMMantle}, n::Integer; T=Float64)
    StressRateAllocMatrix(Matrix{T}(undef, n, 6))
end

# ode parts
@inline function relative_velocity!(alloc::TractionRateAllocFFTConv, vpl::T, v::AbstractMatrix{T}) where T
    @batch for j âˆˆ axes(v, 2)
        for i âˆˆ axes(v, 1)
            alloc.relvnp[i,j] = v[i,j] - vpl # no zero paddings, used for `LinearAlgebra.BLAS`
            alloc.relv[i,j] = alloc.relvnp[i,j] # there exists zero paddings in `alloc.relv`
        end
    end
end

@inline function dÏ„_dt!(gf::AbstractArray{T, 3}, alloc::TractionRateAllocFFTConv) where {T<:Complex}
    mul!(alloc.relv_dft, alloc.pf, alloc.relv)
    fill!(alloc.dÏ„_dt_dft, zero(T))
    # potential simplification by using https://github.com/mcabbott/Tullio.jl
    @batch for j âˆˆ axes(gf, 2)
        for l âˆˆ axes(gf, 3)
            for i âˆˆ axes(gf, 1)
                alloc.dÏ„_dt_dft[i,j] += gf[i,j,l] * alloc.relv_dft[i,l]
            end
        end
    end
    ldiv!(alloc.dÏ„_dt_buffer, alloc.pf, alloc.dÏ„_dt_dft)
    @batch for j âˆˆ axes(alloc.dÏ„_dt, 2)
        for i âˆˆ axes(alloc.dÏ„_dt, 1)
            alloc.dÏ„_dt[i,j] = alloc.dÏ„_dt_buffer[i,j]
        end
    end
end

# build ode
function assemble(
    gf::AbstractArray,
    p::RateStateQuasiDynamicProperty,
    u0::ArrayPartition, tspan::NTuple{2};
    se::StateEvolutionLaw=DieterichStateLaw(), kwargs...)

    alloc = gen_alloc(Val(:BEMFault), size(u0.x[1], 1), size(u0.x[1], 2); kwargs...)
    return ODEProblem{true}(ode, u0, tspan, (p, alloc, gf, se))
end

function assemble(
    gf::AbstractArray,
    p::RateStateQuasiDynamicProperty,
    dila::DilatancyProperty,
    u0::ArrayPartition, tspan::NTuple{2};
    se::StateEvolutionLaw=DieterichStateLaw(), kwargs...)

    alloc = gen_alloc(Val(:BEMFault), size(u0.x[1], 1), size(u0.x[1], 2); kwargs...)
    return ODEProblem{true}(ode, u0, tspan, (p, dila, alloc, gf, se))
end

function assemble(
    gfâ‚�â‚�::AbstractArray,
    gfâ‚�â‚‚::AbstractMatrix,
    gfâ‚‚â‚�::AbstractMatrix,
    gfâ‚‚â‚‚::AbstractMatrix,
    pf::RateStateQuasiDynamicProperty,
    pa::ViscosityProperty,
    u0::ArrayPartition, tspan::NTuple{2};
    se::StateEvolutionLaw=DieterichStateLaw(), kwargs...)

    allocâ‚� = gen_alloc(Val(:BEMFault), size(u0.x[1])...; kwargs...)
    allocâ‚‚ = gen_alloc(Val(:BEMMantle), size(u0.x[3], 1))
    return ODEProblem{true}(ode, u0, tspan, (pf, pa, allocâ‚�, allocâ‚‚, gfâ‚�â‚�, gfâ‚�â‚‚, gfâ‚‚â‚�, gfâ‚‚â‚‚, se))
end

function ode(du::T, u::T, p::Tuple{P, AL, A, SE}, t::U
    ) where {T, U, P<:AbstractProperty, AL<:ODEAllocation, A, SE<:StateEvolutionLaw}

    v, Î¸, _ = u.x
    dv, dÎ¸, dÎ´ = du.x
    prop, alloc, gf, se = p

    relative_velocity!(alloc, prop.vpl, v)
    dÏ„_dt!(gf, alloc)
    update_fault!(prop, alloc, v, Î¸, dv, dÎ¸, dÎ´, se)
end

function ode(du::T, u::T, p::Tuple{P1, P2, AL, A, SE}, t::U
    ) where {
        T, U,
        P1<:RateStateQuasiDynamicProperty, P2<:DilatancyProperty,
        AL<:TractionRateAllocFFTConv,
        A, SE<:StateEvolutionLaw
    }

    v, Î¸, _, ð�“… = u.x
    dv, dÎ¸, dÎ´, dð�“… = du.x
    prop, dila, alloc, gf, se = p

    relative_velocity!(alloc, prop.vpl, v)
    dÏ„_dt!(gf, alloc)
    update_fault_with_dilatancy!(prop, dila, alloc, v, Î¸, ð�“…, dv, dÎ¸, dÎ´, dð�“…, se)
end

function ode(du::T, u::T, p::Tuple{P1, P2, AL1, AL2, A, U, U, U, SE}, t::V
    ) where {
        T, U, V, A,
        SE<:StateEvolutionLaw,
        P1<:RateStateQuasiDynamicProperty, P2<:ViscosityProperty,
        AL1<:TractionRateAllocFFTConv, AL2<:StressRateAllocMatrix
    }

    v, Î¸, _, Ïƒ, _ = u.x
    dv, dÎ¸, dÏµ, dÏƒ, dÎ´ = du.x
    pf, pa, alloc1, alloc2, gfâ‚�â‚�, gfâ‚�â‚‚, gfâ‚‚â‚�, gfâ‚‚â‚‚, se = p

    relative_velocity!(alloc1, pf.vpl, v)
    update_strain_rate!(pa, Ïƒ, dÏµ)
    relative_strain_rate!(alloc2, dÏµ, pa.dÏµâ‚€)
    dÏ„_dt!(gfâ‚�â‚�, alloc1) # fault - fault
    matvecmul!(vec(alloc1.dÏ„_dt), gfâ‚‚â‚�, vec(alloc2.reldÏµ), true, true) # mantle - fault
    matvecmul!(vec(dÏƒ), gfâ‚�â‚‚, vec(alloc1.relvnp)) # fault - mantle
    matvecmul!(vec(dÏƒ), gfâ‚‚â‚‚, vec(alloc2.reldÏµ), true, true) # mantle - mantle
    update_fault!(pf, alloc1, v, Î¸, dv, dÎ¸, dÎ´, se)
end

@inline function update_strain_rate!(p::ViscosityProperty, Ïƒ::T, dÏµ::T) where T
    @batch for i âˆˆ axes(Ïƒ, 1)
        Ïƒkk = (Ïƒ[i,1] + Ïƒ[i,4] + Ïƒ[i,6]) / 3
        Ïƒxx = Ïƒ[i,1] - Ïƒkk
        Ïƒyy = Ïƒ[i,4] - Ïƒkk
        Ïƒzz = Ïƒ[i,6] - Ïƒkk
        Ïƒxy, Ïƒxz, Ïƒyz = Ïƒ[i,2], Ïƒ[i,3], Ïƒ[i,5]
        Ï„norm = sqrt(Ïƒxx^2 + Ïƒyy^2 + Ïƒzz^2 + 2 * (Ïƒxy^2 + Ïƒxz^2 + Ïƒyz^2))
        dÏµ[i,1] = dÏµ_dt(p, i, Ïƒxx, Ï„norm)
        dÏµ[i,2] = dÏµ_dt(p, i, Ïƒxy, Ï„norm)
        dÏµ[i,3] = dÏµ_dt(p, i, Ïƒxz, Ï„norm)
        dÏµ[i,4] = dÏµ_dt(p, i, Ïƒyy, Ï„norm)
        dÏµ[i,5] = dÏµ_dt(p, i, Ïƒyz, Ï„norm)
        dÏµ[i,6] = dÏµ_dt(p, i, Ïƒzz, Ï„norm)
    end
end

@inline function relative_strain_rate!(alloc::StressRateAllocMatrix, dÏµ::AbstractMatrix, dÏµâ‚€::AbstractVector)
    @batch for i âˆˆ axes(dÏµ, 1)
        for j âˆˆ axes(dÏµ, 2)
            alloc.reldÏµ[i,j] = dÏµ[i,j] - dÏµâ‚€[j]
        end
    end
end

# rate-and-state update
@inline function update_fault!(
    p::RateStateQuasiDynamicProperty, alloc::TractionRateAllocFFTConv,
    v::T, Î¸::T, dv::T, dÎ¸::T, dÎ´::T, se::StateEvolutionLaw) where T

    @batch for i âˆˆ eachindex(v)
        Ïˆ1 = exp((p.fâ‚€ + p.b[i] * log(p.vâ‚€ * max(zero(eltype(Î¸)), Î¸[i]) / p.L[i])) / p.a[i]) / 2p.vâ‚€
        Ïˆ2 = p.Ïƒ[i] * Ïˆ1 / hypot(1, v[i] * Ïˆ1)
        dÎ¼_dv = p.a[i] * Ïˆ2
        dÎ¼_dÎ¸ = p.b[i] / Î¸[i] * v[i] * Ïˆ2
        dÎ¸[i] = dÎ¸_dt(se, v[i], Î¸[i], p.L[i])
        dv[i] = (alloc.dÏ„_dt[i] - dÎ¼_dÎ¸ * dÎ¸[i]) / (dÎ¼_dv + p.Î·)
        dÎ´[i] = v[i]
    end
end

@inline function update_fault_with_dilatancy!(
    p::RateStateQuasiDynamicProperty, dila::DilatancyProperty,
    alloc::TractionRateAllocFFTConv,
    v::T, Î¸::T, ð�“…::T, dv::T, dÎ¸::T, dÎ´::T, dð�“…::T, se::StateEvolutionLaw) where T

    @batch for i âˆˆ eachindex(v)
        dÎ¸[i] = dÎ¸_dt(se, v[i], Î¸[i], p.L[i])
        dð�“…[i] = dð�“…_dt(dila, i, ð�“…[i], Î¸[i], dÎ¸[i])

        aá¶  = p.a[i] / p.fâ‚€
        bá¶  = p.b[i] / p.fâ‚€
        vá¶  = max(zero(eltype(v)), v[i] / p.vâ‚€)
        Î¸á¶  = max(zero(eltype(Î¸)), Î¸[i] * p.vâ‚€ / p.L[i])
        vá¶ áµƒâ�»Â¹ = vá¶  ^ (aá¶  - 1)
        Î¸á¶ áµ‡â�»Â¹ = Î¸á¶  ^ (bá¶  - 1)
        vá¶ áµƒ = vá¶  ^ aá¶ 
        Î¸á¶ áµ‡ = Î¸á¶  ^ bá¶ 

        dv[i] = (
            alloc.dÏ„_dt[i] +
            p.fâ‚€ * dð�“…[i] * vá¶ áµƒ * Î¸á¶ áµ‡ -
            p.fâ‚€ * (p.Ïƒ[i] - ð�“…[i]) * vá¶ áµƒ * Î¸á¶ áµ‡â�»Â¹ * bá¶  * p.vâ‚€ / p.L[i] * dÎ¸[i]
        ) / (
            p.fâ‚€ * (p.Ïƒ[i] - ð�“…[i]) * vá¶ áµƒâ�»Â¹ * Î¸á¶ áµ‡ * aá¶  / p.vâ‚€
        )
        dÎ´[i] = v[i]
    end
end

# evolution law
@inline dÎ¸_dt(::DieterichStateLaw, v::T, Î¸::T, L::T) where T = 1 - v * Î¸ / L

# dilantancy law
@inline dð�“…_dt(dila::DilatancyProperty, i::Int, ð�“…::T, Î¸::T, dÎ¸::T) where T = -(ð�“… - dila.pâ‚€[i]) / dila.tâ‚š[i] + dila.Ïµ[i] / dila.Î²[i] / Î¸ * dÎ¸

# viscosity law
@inline dÏµ_dt(p::PowerLawViscosityProperty, i::Int, Ïƒ::T, Ï„norm::T) where T = p.Î³[i] * Ïƒ * Ï„norm ^ p.n[i]
@inline function dÏµ_dt(p::CompositePowerLawViscosityProperty, i::Int, Ïƒ::T, Ï„norm::T) where T
    ans = zero(T)
    for pp âˆˆ p.piter
        ans += dÏµ_dt(pp, i, Ïƒ, Ï„norm)
    end
    ans
end