module Otm

using ..FEA,
      LinearAlgebra

export Otimizador, 
       otimizar!, 
       gerar_trelica_hiperconectada

"""
    gerar_trelica_hiperconectada(Lx::Real, 
                                 Ly::Real, 
                                 nx::Integer,
                                 ny::Integer, 
                                 material::Material, 
                                 nivel::Integer=2, 
                                 area::Real=1)::Estrutura

Gera uma treliça hiperconectada de domínio estendido retangular.

# Argumentos
- `Lx::Real`: Comprimento ao longo de x.
- `Ly::Real`: Comprimento ao longo de y.
- `nx::Integer`: Número de nós na direção x.
- `ny::Integer`: Número de nós na direção y.
- `material::Material`: Material de todos os elementos.
- `nivel::Integer=2`: Nível de conectividade dos nós.
- `area::Real=1`: Área inicial das seções dos elementos.
"""
function gerar_trelica_hiperconectada(Lx::Real, 
                                      Ly::Real, 
                                      nx::Integer,
                                      ny::Integer, 
                                      material::Material, 
                                      nivel::Integer=2, 
                                      area::Real=1)::Estrutura
    # Distância entre nós na direção x.
    lx = Lx/(nx - 1)
    # Distância entre nós na direção y.
    ly = Ly/(ny - 1)
    # Número total de nós na estrutura.
    n = nx * ny
    
    # Criação da matriz que contém as posições de cada nó.
    m = reshape(1:n, (nx, ny))'
    # Criação dos nós da estrutura.
    nos::Vector{No} = []
    k = 1
    for i=1:ny
        y = Ly - (i - 1)*ly
        for j=1:nx
            x = (j - 1)*lx
            push!(nos, No(k, x, y))
            k += 1
        end
    end

    # Criação da treliça hiperconectada sem que haja sobreposição de barras colineares.
    elementos::Vector{Elemento} = []
    c = 1
    for i=1:ny
        for j=1:nx
            no_ref = nos[m[i, j]]
            @info "Gerando treliça: $(m[i, j]) / $n"
            for ii=max(1, i - nivel):min(ny, i + nivel)
                for jj=max(1, j - nivel):min(nx, j + nivel)
                    if !(ii == i && jj == j)
                        no_viz = nos[m[ii, jj]]
                        el_valido = true

                        el_tmp = Elemento(c, no_ref, no_viz, area, material)

                        for el in elementos
                            if tem_sobreposicao(el_tmp, el)
                                el_valido = false
                                break
                            end
                        end
                        if el_valido
                            push!(elementos, el_tmp)
                            c += 1
                        end
                    end
                end
            end
        end
    end
    return Estrutura(nos, elementos)
end

"""
    Otimizador(volume_max::Real, estrutura:: Estrutura)

Implementa as propriedades do modelo

volume_max: Volume de material disponível para o problema de otimização.
estrutura: Estrutura que será otimizada.
"""
mutable struct Otimizador
    volume_max::Float64
    estrutura::Estrutura
    x_min::Float64
    x_max::Float64
    x_ini::Float64
    move::Float64
    η::Float64
    x::Vector
    u::Vector

    function Otimizador(volume_max::Real, estrutura:: Estrutura)
        Ls = [comprimento(el) for el in estrutura.elementos]
        x_max = 1e3 * volume_max / minimum(Ls)
        x_ini = volume_max / sum(Ls)
        x_min = 0
        move = x_ini * 1e4
        η = 0.5
        x = fill(x_ini, length(Ls))
        u = Vector{Float64}(undef, num_gls_estrutura(estrutura))

        new(volume_max, estrutura, x_min, x_max, x_ini, move, η, x, u)
    end
end

"""
    fo(otimizador::Otimizador)::Float64

Retorna a função objetivo.
"""
function fo(otimizador::Otimizador)::Float64
    return vetor_forcas(otimizador.estrutura)' * otimizador.u
end

"""
    ∇fo(otimizador::Otimizador)::Float64

Retorna a sensibilidade da função objetivo.
"""
function ∇fo(otimizador::Otimizador)::Vector{Float64}
    els = otimizador.estrutura.elementos

    return [begin
        uₑ = otimizador.u[gls(els[i])]
        -uₑ' * k(els[i]) * uₑ
    end for i=eachindex(els)]
end

"""
    oc!(otimizador::Otimizador)

Atualiza as variáveis de projeto com o OC.
"""
function oc!(otimizador::Otimizador)
    n = length(otimizador.estrutura.elementos)

    sens_vol = [comprimento(el) for el in otimizador.estrutura.elementos]
    sens_fo = ∇fo(otimizador)
    vol_atual = volume(otimizador.estrutura)

    η = otimizador.η
    bm = -sens_fo ./ sens_vol
    
    l₁ = 0
    l₂ = 1.2 * maximum(bm)
    x_novo = Array{Float64}(undef, n)

    c = 1
    while (l₂ - l₁) > 1e-10 * (l₂ + 1) && c ≤ 1000
        lₘ = (l₂ + l₁) / 2

        # Correção de erro numérico
        be = max.(0, bm / lₘ)
        t = @. otimizador.x_min + (otimizador.x - otimizador.x_min) * be^η
        x_novo = @fastmath @. max(max(min(min(t, otimizador.x + otimizador.move), otimizador.x_max), 
                                  otimizador.x - otimizador.move), otimizador.x_min)
        
        if (vol_atual - otimizador.volume_max) + sens_vol' * (x_novo - otimizador.x) > 0
            l₁ = lₘ
        else
            l₂ = lₘ
        end
        c += 1
    end
    otimizador.x = x_novo
end

"""
    otimizar!(otimizador::Otimizador; max_it=1000, ϵ=1e-9) 

Minimiza a compliance do problema.
"""
function otimizar!(otimizador::Otimizador; max_it=1000, ϵ=1e-9)
    num_elems = length(otimizador.estrutura.elementos)
    erro = Inf
    
    iter = 0
    while erro > ϵ && iter ≤ max_it
        norm₀ = norm(otimizador.x)

        for i=1:num_elems
            otimizador.estrutura.elementos[i].area = otimizador.x[i]
        end

        otimizador.u = deslocamentos(otimizador.estrutura)

        oc!(otimizador)
        norm₁ = norm(otimizador.x)

        if iter > 1
            erro = abs(norm₁ - norm₀) / norm₀
        end

        iter += 1
        @info "It: $iter \tCompliance: $(fo(otimizador))\t Erro: $erro \tVolume: $(volume(otimizador.estrutura))"
    end
end

end