module FEM 
export Node, Element, stiffness, tri_tiles, gradient, center, normalize, p

include("./Parameters.jl")

using GeometryBasics
using GeometryBasics.LinearAlgebra

struct Node
  x::Float64
  y::Float64
  index::UInt
end

p(n::Node) = Point2f(n.x, n.y)

struct Element
  nodes::Vector{Node}
  a::Vector{Float64}
  b::Vector{Float64}
  c::Vector{Float64}
  Δ::Float64

  function Element(n::Vector{Node})
    xs = map(nd->nd.x,n)
    ys = map(nd->nd.y,n)
    a = cross(xs, ys)
    b = cross(ys, ones(3))
    c = cross(ones(3), xs)
    Δ = dot(xs,b) / 2

    new(n,a,b,c,Δ)
  end
end

function stiffness(e::Element)::Matrix{Float64} 
  #  tensor product
  Hₑ  = e.b .* e.b'
  Hₑ += e.c .* e.c'
  Hₑ *= Parameters.hz*Parameters.k/4e.Δ

  return Hₑ
end

function tri_tiles(L::Float64, divisions::Int, trapezoidal::Bool=false, biased::Bool=false, ring::Bool=false)::Tuple{Vector{Node}, Vector{Element}}
  nodes = Matrix{Node}(undef, divisions+1, divisions+1)
  
  i = 1
  for y in 0:divisions
    for x in 0:divisions
      xₑ = L*x/divisions
      yₑ = L*y/divisions

      if biased
        B = yₑ/2Parameters.L
        xₑ = xₑ*(xₑ*B/Parameters.L - B + 1)
      end

      if trapezoidal
        xₑ *= 1 - 0.5*(y/divisions)
      end

      if ring
        r = Parameters.L*(1 + x/divisions)
        θ = (π/4)*(y/divisions)

        xₑ = r*cos(θ)
        yₑ = r*sin(θ)
        xₑ -= Parameters.L
      end

      nodes[x+1,y+1] = Node(xₑ,yₑ,i)
      i += 1
    end
  end

  elements = []

  for y in 1:divisions
    for x in 1:divisions
      # lower/upper triangle
      push!(elements, Element([nodes[x,y], nodes[x+1,y], nodes[x,y+1]]))
      push!(elements, Element([nodes[x+1,y+1], nodes[x,y+1], nodes[x+1,y]]))
    end
  end

  return vec(nodes), elements
end

function gradient(e::Element, T::Vector{Float64})::Vec2f
  return Vec2f([
    e.b[1] e.b[2] e.b[3] ;
    e.c[1] e.c[2] e.c[3] 
  ] * [
    T[e.nodes[1].index]
    T[e.nodes[2].index]
    T[e.nodes[3].index]
  ] / 2e.Δ)
end

center(e::Element)::Point2f = Point2f(sum([n.x for n in e.nodes])/3.0, sum([n.y for n in e.nodes])/3.0)
normalize(v::Vec2f)::Vec2f = v / sqrt(v[1]^2 + v[2]^2)

end