"Implementa a análise estrutural."
module FEA

using LinearAlgebra,
      SparseArrays,
      NearestNeighbors

export No, 
       Material, 
       Elemento, 
       Estrutura, 
       comprimento, 
       num_gls_estrutura, 
       vetor_forcas, 
       volume, 
       deslocamentos, 
       gls, 
       k, 
       tem_sobreposicao,
       carregar_apoiar_no!

"Número de graus de liberdade de um nó."
const num_gls_no = 2

# =======================================================================
"""
    No(id::Integer, 
       x::Real, 
       y::Real; 
       fx::Real = 0, 
       fy::Real = 0, 
       rx::Bool = false, 
       ry::Bool = false)

Implementa as propriedades de um nó.

# Argumentos
- `id::Integer`: Identificação do nó.
- `x::Real`: Coordenada x.
- `y::Real`: Coordenada y.
- `fx::Real`: Força na direção x.
- `fy::Real`: Força na direção y.
- `rx::Bool`: Apoio na direção x. true para impedido e 'false' para livre.
- `ry::Bool`: Apoio na direção y. true para impedido e 'false' para livre.
"""
mutable struct No
    id::Int64
    x::Float64
    y::Float64
    fx::Float64
    fy::Float64
    rx::Bool
    ry::Bool

    function No(id::Integer, 
                x::Real, 
                y::Real; 
                fx::Real = 0, 
                fy::Real = 0, 
                rx::Bool = false, 
                ry::Bool = false)

        if id < 1
            throw(ArgumentError("O número de identificação do nó deve ser maior que 1!"))
        end
        new(id, x, y, fx, fy, rx, ry)
    end
end

"""
    distancia(no₁::No, no₂::No)::Float64

Retorna a distância entre `no₁` e `no₂`.
"""
distancia(no₁::No, no₂::No)::Float64 = norm([(no₂.x - no₁.x), (no₂.y - no₁.y)])

"""
    gls(no::No)::Vector{Int64}

Retorna os graus de liberdade do nó.
"""
gls(no::No)::Vector{Int64} = [num_gls_no * no.id - 1, num_gls_no * no.id]

"""
    vetor_forcas(no::No)::Vector{Float64}

Retorna as forças nodais.
"""
vetor_forcas(no::No)::Vector{Float64} = [no.fx, no.fy]

"""
    vetor_apoios(no::No)::Vector{Bool}

    Retorna os apoios nodais.
"""
vetor_apoios(no::No)::Vector{Bool} = [no.rx, no.ry]

"""
    gls_livres(no::No)::Vector{Int64}

Retorna os graus de liberdade livres do nó.
"""
gls_livres(no::No)::Vector{Int64} = [i for (i, j) in zip(gls(no), vetor_apoios(no)) if j == false]

"""
    sao_colineares(no₁::No, no₂::No, no₃::No)::Bool

Verifica se três nós são colineares.

É feito o cálculo da área do triângulo formado pelos três nós. Caso essa área seja nula, 
os nós são colineares.
"""
function sao_colineares(no₁::No, no₂::No, no₃::No)::Bool
    m = [no₁.x no₁.y 1
        no₂.x no₂.y 1
        no₃.x no₃.y 1]

    return isapprox(det(m), 0; atol=1e-9)
end

# =====================================================================================
"""
Define as propriedades do material que compõe os elementos.

# Argumentos
- `v::Integer`: Coeficiente de Poisson.
- `E::AbstractFloat`: Módulo de Elasticidade.
"""
struct Material
    id::Int64
    E::Float64

    function Material(id::Integer, E::AbstractFloat)
        if id < 1
            throw(ArgumentError("O número de identificação do material não pode ser menor que 1!"))
        end
        if !(E > 0)
            throw(ArgumentError("O Módulo de Elasticidade deve ser maior que 0!"))
        end
        new(id, E)
    end
end

# =====================================================================================
"""
    Elemento(id::Integer, no₁::No, no₂::No, area::Real, material::Material)

Implementa as propriedades de uma elemento de treliça.

# Argumentos
- `id::Integer`: Número de identificação do elemento.
- `no₁::No`: Nó inicial.
- `no₂::No`: Nó final.
- `area::Real`: Área da seção transversal.
- `material::Material`: Material utilizado no elemento.
"""
mutable struct Elemento
    id::Int64
    no₁::No
    no₂::No
    area::Float64
    material::Material

    function Elemento(id::Integer, no₁::No, no₂::No, area::Real, material::Material)
        if id < 1
            throw(ArgumentError("O número de identificação do elemento não pode ser menor que 1!"))
        end
        if area < 0
            throw(ArgumentError("A área da seção não pode ser menor que 0!"))
        end
        new(id, no₁, no₂, area, material)
    end
end

"""
    comprimento(elemento::Elemento)::Float64

Retorna o comprimento do elemento.
"""
comprimento(elemento::Elemento)::Float64 = distancia(elemento.no₁, elemento.no₂)

"""
    gls(elemento::Elemento)::Vector{Int64}

Retorna os graus de liberdade do elemento.
"""
gls(elemento::Elemento)::Vector{Int64} = vcat(gls(elemento.no₁), gls(elemento.no₂))

"""
    vetor_forcas(elemento::Elemento)::Vector{Float64}

Retorna o vetor de forças atuantes nos graus de liberdade do elemento.
"""
vetor_forcas(elemento::Elemento)::Vector{Float64} = vcat(vetor_forcas(elemento.no₁), vetor_forcas(elemento.no₂))

"""
    gls_livres(elemento::Elemento)::Vector{Int64}

Retorna os graus de liberdade livres do elemento.
"""
gls_livres(elemento::Elemento)::Vector{Int64} = vcat(gls_livres(elemento.no₁), gls_livres(elemento.no₂))

"""
    angulo(elemento::Elemento)::Float64

Retorna o ângulo de inclinação do elemento em relação a horizontal.

Segue as convenções do círculo trigonométrico.
"""
function angulo(elemento::Elemento)::Float64
    dx = elemento.no₂.x - elemento.no₁.x
    dy = elemento.no₂.y - elemento.no₁.y

    return atan(dy, dx)
end

"""
    matriz_rotacao(elemento::Elemento)::Matrix{Float64}

Retorna a matriz de rotação do elemento.
"""
function matriz_rotacao(elemento::Elemento)::Matrix{Float64}
    θ = angulo(elemento)
    s = sin(θ)
    c = cos(θ)

    return [c s 0 0
           -s c 0 0
            0 0 c s
            0 0 -s c]
end

"""
    k_local(elemento::Elemento)::Matrix{Float64}

Retorna a matriz de rigidez local do elemento.
"""
function k_local(elemento::Elemento)::Matrix{Float64}
    L = comprimento(elemento)
    E = elemento.material.E
    A = elemento.area

    return (E * A / L) * [1 0 -1 0
                          0 0 0 0
                         -1 0 1 0
                          0 0 0 0]
end

"""
    k(elemento::Elemento)::Matrix{Float64}

Retorna a matriz de rigidez do elemento rotacionada no sistema global.
"""
function k(elemento::Elemento)::Matrix{Float64}
    k = k_local(elemento)
    r = matriz_rotacao(elemento)

    return r' * k * r
end

"""
    sao_colineares(elemento₁::Elemento, elemento₂::Elemento)::Bool

Verifica se dois elementos são colineares.
"""
function sao_colineares(elemento₁::Elemento, elemento₂::Elemento)::Bool
    el1 = elemento₁
    el2 = elemento₂
    if sao_colineares(el1.no₁, el1.no₂, el2.no₁) && sao_colineares(el1.no₁, el1.no₂, el2.no₂)
        return true
    else
        return false
    end
end

"""
    tem_sobreposicao(elemento₁::Elemento, elemento₂::Elemento)::Bool

Verifica se há sobreposição entre dois elementos.
"""
function tem_sobreposicao(elemento₁::Elemento, elemento₂::Elemento)::Bool
    # Verificar se os nós dos elementos são colineares.
    el1 = elemento₁
    el2 = elemento₂
    L1 = comprimento(el1)
    L2 = comprimento(el2)
    # Se os dois elementos possuírem os mesmos nós, há sobreposição.
    if el1.no₁ in [el2.no₁, el2.no₂] && el1.no₂ in [el2.no₁, el2.no₂]
        return true
    else
        if sao_colineares(el1, el2)
            ls = [distancia(el1.no₁, el2.no₁),
                  distancia(el1.no₁, el2.no₂),
                  distancia(el1.no₂, el2.no₁),
                  distancia(el1.no₂, el2.no₂)]
            
            # Distância entre os nós extremos dos dois elementos de referência.
            dmax = maximum(ls)

            if dmax ≤ (L1 + L2) && !isapprox(dmax, L1 + L2; atol=1e-9)
                return true
            else
                return false
            end
        else
            return false
        end
    end
end

# =====================================================================================
"""
    Estrutura(nos::Vector{No}, elementos::Vector{Elemento}, λ::AbstractFloat=1e-9)

Implementa as propriedades de uma Estrutura.

# Argumentos
- `nos::Array{No}`: Nós que compõem a estrutura.
- `elementos::Array{Elemento}`: Elementos que compõem a estrutura.
- `λ::Float64`: Parâmetro de Tikhonov.
"""
struct Estrutura
    nos::Vector{No}
    elementos::Vector{Elemento}
    λ::Float64

    function Estrutura(nos::Vector{No}, elementos::Vector{Elemento}, λ::AbstractFloat=1e-9)
        new(nos, elementos, λ)
    end
end

"""
    carregar_apoiar_no!(nos::Vector{No}, x, y; fx=0, fy=0, rx=false, ry=false)

Adiciona as cargas `fx` e `fy` ao nó mais próximo do ponto de coordenadas `x` e `y`.
"""
function carregar_apoiar_no!(nos::Vector{No}, x, y; fx=0, fy=0, rx=false, ry=false)
    # Matriz contendo os dados de todos os nós.
    m = Matrix{Float64}(undef, (2, length(nos)))
    for i=eachindex(nos)
        m[1, i] = nos[i].x
        m[2, i] = nos[i].y
    end
    kd = KDTree(m)
    no_p = knn(kd, [x, y], 1)[1][1]

    nos[no_p].fx = fx
    nos[no_p].fy = fy
    nos[no_p].rx = rx
    nos[no_p].ry = ry
end

"""
    num_gls_estrutura(estrutura::Estrutura)::Int64 = num_gls_no * length(estrutura.nos)

Retorna o número de graus de liberdade da estrutura.
"""
num_gls_estrutura(estrutura::Estrutura)::Int64 = num_gls_no * length(estrutura.nos)

"""
    gls_livres(estrutura::Estrutura)::Vector{Int64}

Retorna os graus de liberdade livres da estrutura.
"""
gls_livres(estrutura::Estrutura)::Vector{Int64} = vcat([gls_livres(el) for el in estrutura.elementos]...) |> unique |> sort

"""
    volume(estrutura::Estrutura)::Float64

Retorna o volume da estrutura.
"""
volume(estrutura::Estrutura)::Float64 = sum([el.area * comprimento(el) for el in estrutura.elementos])

"""
    vetor_forcas(estrutura::Estrutura, incluir_gls_impedidos::Bool = true)::Vector{Float64}

Retorna o vetor de forças da estrutura.
"""
function vetor_forcas(estrutura::Estrutura, incluir_gls_impedidos::Bool = true)::Vector{Float64}
    forcas = Array{Float64}(undef, num_gls_estrutura(estrutura))

    for no in estrutura.nos
        forcas[gls(no)] = [no.fx, no.fy]
    end

    if incluir_gls_impedidos
        return forcas
    else
        return forcas[gls_livres(estrutura)]
    end
end

"""
    vetor_apoios(estrutura::Estrutura)::Vector{Int64}

Retorna os graus de liberdade impedidos.
"""
function vetor_apoios(estrutura::Estrutura)::Vector{Int64}
    apoios::Array{Int64} = []

    for no in estrutura.nos
        vaps = vetor_apoios(no)

        if any(vaps)
            gl = gls(no)
            append!(apoios, [gl[i] for (i, j) in enumerate(vaps) if j])
        end
    end
    return apoios
end

"""
    k(estrutura::Estrutura)::SparseMatrixCSC

Retorna a matriz de rigidez esparsa da estrutura.

Procedimento de montagem com 3 vetores separados.
"""
function k(estrutura::Estrutura)::SparseMatrixCSC
    # Número de termos em cada matriz de rigidez de elementos.
    num_gle = (num_gls_no * 2)
    apoios = vetor_apoios(estrutura)

    num_it = length(estrutura.elementos) * num_gle^2

    linhas = ones(Int64, num_it)
    colunas = ones(Int64, num_it)
    termos = zeros(num_it)

    c = 1

    for elem in estrutura.elementos
        gle = gls(elem)
        k_glob = k(elem)

        # Posições livres do elemento.
        pos_livres = [i for (i, gl) in enumerate(gle) if gl ∉ apoios]

        if length(pos_livres) > 0
            for i in pos_livres
                lin = gle[i] - count(x -> x < gle[i], apoios)
                for j in pos_livres
                    linhas[c] = lin
                    colunas[c] = gle[j] - count(x -> x < gle[j], apoios)
                    termos[c] = i != j ? k_glob[i, j] : k_glob[i, j] + estrutura.λ
                    c += 1
                end
            end
        end
    end
    ks = sparse(linhas, colunas, termos)
    dropzeros!(ks)

    return ks
end

"""
    deslocamentos(estrutura::Estrutura, incluir_gls_impedidos::Bool = true)::Array{Float64}

Retorna os deslocamentos nodais.
"""
function deslocamentos(estrutura::Estrutura, incluir_gls_impedidos::Bool = true)::Array{Float64}
    u = k(estrutura) \ vetor_forcas(estrutura, false)
    if !incluir_gls_impedidos
        return u
    else
        uf = zeros(num_gls_estrutura(estrutura))
        uf[gls_livres(estrutura)] = u
        return uf
    end
end

"""
    nos_deslocados(estrutura::Estrutura, escala::Int64 = 1)::Vector{No}

Retorna um vetor contendo os nós em suas novas posições com a estrutura deformada.
"""
function nos_deslocados(estrutura::Estrutura, escala::Int64 = 1)::Vector{No}
    num_nos = length(estrutura.nos)
    u_pontos = reshape(deslocamentos(estrutura), (2, num_nos))'
    nos_desloc = Array{No}(undef, length(estrutura.nos))

    for i in 1:num_nos
        nos_desloc[i] = No(i, estrutura.nos[i].x + escala * u_pontos[i, 1],
            estrutura.nos[i].y + escala * u_pontos[i, 2])
    end
    return nos_desloc
end

"""
    forcas_internas(estrutura::Estrutura)::Vector{Float64}

Retorna os valores dos esforços normais atuantes em cada elemento.
"""
function forcas_internas(estrutura::Estrutura)::Vector{Float64}
    forcas = Array{Float64}(undef, length(estrutura.elementos))
    nos_def = nos_deslocados(estrutura)

    for (i, elem) in enumerate(estrutura.elementos)
        l₀ = comprimento(elem)
        l₁ = distancia(nos_def[elem.no₁.id], nos_def[elem.no₂.id])
        deform = (l₁ - l₀) / l₀
        forcas[i] = elem.area * elem.material.E * deform
    end

    return forcas
end

end