include("./fem.jl")
include("./Parameters.jl")
include("./print_HTP.jl")

using .FEM
using GLMakie
using Makie.GeometryBasics

@enum Variation BASIC V1 V2 V3 V4_div V4_mul

function solve(name::String, variation::Variation)
  div = 9
  H = zeros(Float64,(div+1)^2, (div+1)^2)

  (nodes, elements) = tri_tiles(Parameters.L, div, variation==V1, variation==V2, variation==V3)
  
  for (index, el) in enumerate(elements)
    Hₑ = stiffness(el)

    if variation == V4_div && index in Parameters.elₘ 
      Hₑ /= Parameters.c 
    elseif variation == V4_mul && index in Parameters.elₘ 
      Hₑ *= Parameters.c 
    end

    for i in 1:3
      for j in 1:3
        H[el.nodes[i].index, el.nodes[j].index] += Hₑ[i,j]
      end
    end
  end
  
  P = zeros(Float64, 90)
  
  # impose neumann conditions on top edge
  # adapted from https://mathoverflow.net/questions/5085/how-to-apply-neuman-boundary-condition-to-finite-element-method-problems
  # maybe check if this actually makes sense
  
  nᵧ = -Parameters.q / Parameters.k
  # skip=2 since only every second element has a edge along the bottom: ◺◹
  for ∂_el in elements[1:2:18]
    n1 = ∂_el.nodes[1] 
    n2 = ∂_el.nodes[2]
    l = abs(n2.x - n1.x)
    P[n1.index] += nᵧ*l/2
    P[n2.index] += nᵧ*l/2
  end
  
  rhs = P - H[1:90,91:100]*fill(Parameters.T, 10,1)
  T = vec(H[1:90,1:90]\rhs)
  append!(T, fill(Parameters.T, 10))
  
  reaction_forces = H[91:100,:]*T
  
  centers = center.(elements)

  gradients = map(el -> gradient(el, T), elements)
  flux = gradients .* -1

  norm = maximum(map(g-> sqrt(g[1]^2 + g[2]^2), gradients))/(Parameters.L/div)*2.5
  
  gradients ./= norm
  flux ./= norm
  
  set_theme!(theme_black())
  f = Figure(resolution = (1536, 1024))
  
  tris = map(e->Polygon([p(e.nodes[1]), p(e.nodes[2]), p(e.nodes[3])]), elements)
  poly(f[1, 1], tris, color=:transparent, linestyle=:solid, strokewidth=0.8, strokecolor=:white, transparency=true)

  xs = map(n->n.x, nodes)
  ys = map(n->n.y, nodes)

  xs = reshape(xs, (div+1, div+1))
  ys = reshape(ys, (div+1, div+1))
  Ts = reshape(T,  (div+1, div+1))

  surface(f[2, 1], xs, ys, Ts, colormap=:matter, axis=(type=Axis3,))

  contour(f[1:2,2:3], map(n->n.x,nodes), map(n->n.y,nodes), T, levels=16, colormap=:matter)
  arrows!(f[1:2,2:3], centers, gradients, arrowcolor=:red, linecolor=:red)
  arrows!(f[1:2,2:3], centers, flux, arrowcolor=:blue, linecolor=:blue)
  
  save("plots/$name.png", current_figure())

  print_HTP(H, T, P, "txt/htp_$name.txt")
end

solve("baseline", BASIC)
solve("variation1", V1)
solve("variation2", V2)
solve("variation3", V3)
solve("variation4_div", V4_div)
solve("variation4_mul", V4_mul)