# Scientific Computing and FEM
# Project C: mug simulation with ngsolve


from ngsolve import *
# from netgen.csg import *
from netgen.geom2d import SplineGeometry
import matplotlib.pyplot as plt
import time

""" It's not possible to generate a 3D model of the mug. I tried several ways.
One can build the geometry and have a look on it. However, the computer is not able to generate a
mesh from it. It got stuck. Using the jupyterhub from university does not change the situation.
So we will continue solving this example in 2D instead of 3D. """

# # generating the mesh 3D
# geo = CSGeometry()

# # parameter
# h = 105     # height
# r_bottom = 55   # radius bottom
# r_top = 83  # radius top
# thickness = 5   # thickness of cup
# fill_height = 66

# # outer cone (mug outside)
# outer = Cone(Pnt(0,0,0), Pnt(0,0,h), r_bottom, r_top)

# # inner cone
# inner = Cone(Pnt(0,0,thickness), Pnt(0,0,h), r_bottom - thickness, r_top - thickness)

# # ceramic (=difference)
# #cup = outer - inner

# # cutting liquid / air
# cyl_liquid = Cylinder(Pnt(0,0,0), Pnt(0,0,fill_height), r_top+10)
# cyl_air = Cylinder(Pnt(0,0,fill_height), Pnt(0,0,h), r_top+10)

# inside = inner
# # liquid (lower part)
# liquid = inside * cyl_liquid

# # air (upper part)
# air = inside * cyl_air

# cup = outer - (liquid + air)

# # whole geometry + material
# geo.Add(cup.mat("ceramic"))
# geo.Add(liquid.mat("liquid"))
# geo.Add(air.mat("air"))

# mesh = Mesh(geo.GenerateMesh(maxh=40.5))
# print("Elements:", mesh.ne)

# Draw(mesh, mesh.Materials())
# geo.Draw()

# -----------------------------------

# generating the 2D model and the corresponding mesh
geo = SplineGeometry()

midpoint = 22+66/99*(83/2-25)
points_mm = [(25,0), (83/2,105), (-83/2,105), (-25,0),     # outer points of ceramic
            (22,6), (midpoint,72), (-midpoint,72), (-22,6), # points of liquid
            (83/2-3,105), (-83/2+3,105)] # upper points of air
factor = 10**-3  # transformation from mm to m
points = [[x * factor, y * factor] for x, y in points_mm]

p0,p1,p2,p3,p4,p5,p6,p7,p8,p9 = [geo.AppendPoint(*pnt) for pnt in points]

curves = [[["line",p0,p1],"right",1,0],
            [["line",p1,p8],"top_mug",1,0],
            [["line",p8,p5],"rest",1,3],
            [["line",p5,p4],"rest",1,2],
            [["line",p4,p7],"rest",1,2],
            [["line",p7,p6],"rest",1,2],
            [["line",p6,p9],"rest",1,3],
            [["line",p9,p2],"top_mug",1,0],
            [["line",p2,p3],"left",1,0],
            [["line",p3,p0],"bottom",1,0],
            [["line",p5,p6],"rest",2,3],
            [["line",p8,p9],"top_air",3,0]]

for geom_part, name, left, right in curves:
    if name != "rest":  	# only on the boundary we need a name for bc
        geo.Append(geom_part, leftdomain=left, rightdomain=right, bc=name)
    else:
        geo.Append(geom_part, leftdomain=left, rightdomain=right)

geo.SetMaterial (1, "ceramic")
geo.SetMaterial (2, "liquid")
geo.SetMaterial (3, "air")

mesh = Mesh(geo.GenerateMesh(maxh=0.0007))
#print (f"segments of the boundary: {mesh.GetBoundaries()}")
#print (f"name of the materials: {mesh.GetMaterials()}")

#Draw(mesh) # plotting the mesh

# printing the domains of the mug
print_domains = False
if print_domains == True:
    cf = CoefficientFunction([
        1 if mat=="ceramic" else
        2 if mat=="liquid" else
        3 if mat=="air" else 0
        for mat in mesh.GetMaterials()
    ])

    Draw(cf, mesh, "domains")

# --------------------------------------

# defining the parameters
# heat storage capacity (J/(m^3*K))
c_domain_values = {'ceramic': 1240*2.3*10**3, 'liquid': 4.138*10**6, 'air': 1.21*10**3}
c = mesh.MaterialCF(c_domain_values)
# thermal conductivity (W/(m*K))
lam_domain_values = {'ceramic': 5, 'liquid': 0.6, 'air': 0.026}
lam = mesh.MaterialCF(lam_domain_values)
# heat transfer coefficient (W/(m^2*K))
alpha = 8
# outside temperature
u_out = 18

# -----------------------------------

# setting up the FEM
# H1-conforming finite element space
fes = H1(mesh, order=1, autoupdate=True)

# define trial- and test-functions
u = fes.TrialFunction()
v = fes.TestFunction()

# define the bilinear form (stiffness + Robin part)
a = BilinearForm(fes, symmetric=True)
a += (lam*grad(u)*grad(v))*dx
a += (alpha*u*v)*ds

# define the bilinear form (mass)
m = BilinearForm(fes, symmetric=True)
m += c*u*v*dx

# define the right-hand side (Robin part)
f = LinearForm(fes)
f += (alpha*u_out*v)*ds
    
# assemble the system of linear equation
a.Assemble()
m.Assemble()
f.Assemble()
minv = m.mat.Inverse(fes.FreeDofs(), inverse="pardiso")

# solution field
gfu = GridFunction(fes)

# ------------------------------------

# stationary problem (no time aspect)
# gfu.vec.data = a.mat.Inverse(fes.FreeDofs(), inverse="pardiso")*f.vec.data
# Draw(gfu)

# -----------------------------------

# functions that will be helpful

# calculatin the average temperatur in a given subdomain
def mean_temp(gfu,mesh,domain):
    temp_int = Integrate(gfu, mesh, definedon=mesh.Materials(domain))
    vol = Integrate(1, mesh, definedon=mesh.Materials(domain))

    return temp_int/vol

# -----------------------------------

# Q1
run_Q1 = False

if run_Q1 == True:

    # initial data for Q1
    init_u_liquid = 80
    init_u_domain_values = {'ceramic': u_out, 'liquid': init_u_liquid, 'air': u_out}
    init_u = mesh.MaterialCF(init_u_domain_values)
    gfu.Set(init_u)
    Draw(gfu)

    time.sleep(7)
    #Redraw(blocking=True)

    # setting up the FEM environment
    time_max = 10*60    # breaking condition for simulation
    tau = 0.5  # duration of 1 timestep
    temp_opt = 67

    theta = 0.5 # paramater for time discretization
    # 0 (explicit Euler), 0.5 (Crank Nicolson scheme), 1 (implicit Euler)

    mstar = m.mat.CreateMatrix()
    mstar.AsVector().data = m.mat.AsVector() + theta*tau*a.mat.AsVector()
    mstarinv = mstar.Inverse(fes.FreeDofs(), inverse="pardiso")
    r = gfu.vec.CreateVector()

    time_act = 0    # starting time =0 for loop condition
    temp = 0    # here just for starting condition in while loop
    data_q1 = [[time_act,u_out]]

    # calculating T_warm, temperature when ceramic has optimal temperature temp_opt 
    while temp<temp_opt and time_act<time_max:
        time_act = time_act + tau   # actual time point
        r.data = (m.mat - (1 - theta)*tau*a.mat)*gfu.vec + tau*f.vec
        gfu.vec.data = mstarinv * r
        temp = mean_temp(gfu,mesh,'ceramic')
        data_q1.append([time_act,temp])
        if time_act % 5 == 0:
            Draw(gfu, min=u_out,max=init_u_liquid,autoscale=False)
        if time_act == 451: # time of maximal temperature (ceramic); for u_0 = 80 !!!
            gfu.Save("solution_Q1.vol") 
    # finding the maximal temperature of the ceramic cup
    time_temp_max, temp_ceramic_max = max(data_q1, key=lambda x: x[1])
    print(f"u_0^liquid = {init_u_liquid} °C.")
    print(f"After {time_temp_max} s the ceramic part reaches its maximal temperature of {temp_ceramic_max:.3f} °C.")


# -----------------------

# Q2
run_Q2 = True

if run_Q2 == True:

    # initial data for Q2
    init_u_liquid = 85  # now we have coffee (nearly the same behavior as water)
    gfu_old = GridFunction(fes)
    gfu_old.Load("solution_Q1.vol")
    chi_liquid = mesh.MaterialCF({'liquid': 1}, default=0)
    cf_init = gfu_old * (1 - chi_liquid) + init_u_liquid * chi_liquid
    gfu.Set(cf_init)
    Draw(gfu, min=u_out,max=init_u_liquid,autoscale=False)
    time.sleep(7)

    # setting up the FEM environment
    time_max = 10*60*7    # breaking condition for simulation
    tau = 2.5  # duration of 1 timestep
    temp_opt = 50   # coffee should cool down

    theta = 0.5 # paramater for time discretization
    # 0 (explicit Euler), 0.5 (Crank Nicolson scheme), 1 (implicit Euler)

    mstar = m.mat.CreateMatrix()
    mstar.AsVector().data = m.mat.AsVector() + theta*tau*a.mat.AsVector()
    mstarinv = mstar.Inverse(fes.FreeDofs(), inverse="pardiso")
    r = gfu.vec.CreateVector()

    time_act = 0    # starting time =0 for loop condition
    temp = init_u_liquid    # here just for starting condition in while loop
    data_q2 = [[time_act,init_u_liquid]]

    # calculating T_cof, when coffee drops to temperature temp_opt 
    while temp>temp_opt and time_act<time_max:
        time_act = time_act + tau   # actual time point
        r.data = (m.mat - (1 - theta)*tau*a.mat)*gfu.vec + tau*f.vec
        gfu.vec.data = mstarinv * r
        temp = mean_temp(gfu,mesh,'liquid')
        data_q2.append([time_act,temp])
        if time_act % 25 == 0:
            Draw(gfu, min=u_out,max=init_u_liquid,autoscale=False)
        if time_act % 100 == 0:
            print(f'time: {time_act}; temperature: {temp:.3f}') 
        
    print(f'end-time: {time_act}; temperature: {temp:.3f}')

# -----------------------

# -----------------------

# plotting of the temperature curves

# Q1: heating the ceramic
plot_1 = False
if plot_1 == True:
    x,y = zip(*data_q1)
    plt.figure()
    plt.plot(x, y)
    plt.grid(True)
    plt.axhline(y=temp_ceramic_max, color='red', linestyle='--', label=f'T_max={temp_ceramic_max:.3f}') # horizontal line
    plt.axvline(x=time_temp_max, color='black', linestyle='--', label=f't={time_temp_max}') # vertical line
    plt.title('temperature curve of ceramic')
    plt.xlabel('time (in s)')
    plt.ylabel('temperature (°C)')
    plt.legend()
    plt.savefig('temperature_ceramic_q1.png', dpi=600)
    plt.show()
    plt.close()

# Q2: coffee cool down
plot_2 = False
if plot_2 == True:
    x,y = zip(*data_q2)
    plt.figure()
    plt.plot(x, y)
    plt.grid(True)
    plt.title('temperature curve of coffee')
    plt.xlabel('time (in s)')
    plt.ylabel('temperature (°C)')
    plt.savefig('temperature_coffee_q2.png', dpi=600)
    plt.show()
    plt.close()