#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Feb  1 08:49:11 2019

@author: bleyerj
"""
from dolfin import *
import mfront_wrapper as mf
import numpy as np
import meshio

fname = "../meshes/rod3D.msh"
msh = meshio.read("../meshes/rod3D.msh")
for cell in msh.cells:
    if cell.type == "triangle":
        triangle_cells = cell.data
    elif  cell.type == "tetra":
        tetra_cells = cell.data

for key in msh.cell_data_dict["gmsh:physical"].keys():
    if key == "triangle":
        triangle_data = msh.cell_data_dict["gmsh:physical"][key]
    elif key == "tetra":
        tetra_data = msh.cell_data_dict["gmsh:physical"][key]
tetra_mesh = meshio.Mesh(points=msh.points, cells={"tetra": tetra_cells},
                         cell_data={"name_to_read":[tetra_data]})
# triangle_mesh =meshio.Mesh(points=msh.points,
#                            cells=[("triangle", triangle_cells)],
#                            cell_data={"name_to_read":[triangle_data]})
triangle_mesh =meshio.Mesh(points=msh.points,
                           cells=[("triangle", triangle_cells)],
                           cell_data={"name_to_read":[triangle_data]})
meshio.write("mesh.xdmf", tetra_mesh)
meshio.write("mf.xdmf", triangle_mesh)

# remove blank spaces in field_date keys
tags = dict([(key.strip(), value) for (key, value) in msh.field_data.items()])


mesh = Mesh()
mvc = MeshValueCollection("size_t", mesh, 3)
with XDMFFile("mesh.xdmf") as infile:
    infile.read(mesh)
    infile.read(mvc, "name_to_read")
cf = cpp.mesh.MeshFunctionSizet(mesh, mvc)
mvc = MeshValueCollection("size_t", mesh, 2)
with XDMFFile("mf.xdmf") as infile:
    infile.read(mvc, "name_to_read")
mf = MeshFunction("size_t", mesh, 2)

def get_tag(tag):
    return tags[tag][0]

dx = Measure("dx", domain=mesh, subdomain_data=cf)
ds = Measure("ds", domain=mesh, subdomain_data=mf)

V = FunctionSpace(mesh, "CG", 2)
T = Function(V, name="Temperature")

Text = 1e3
v = Function(V)
bc = DirichletBC(V, Constant(Tl), mf, get_tag("00SOO"))
bc[0].apply(v.vector())
facets = MeshFunction("size_t", mesh, 1)
ds = Measure("ds", subdomain_data=facets)

dt = Constant(1e-3)
theta = Constant(0.5)

material = mf.MFrontNonlinearMaterial("../materials/src/libBehaviour.so",
                                      "HeatTransfer",
                                      hypothesis="plane_strain")
problem = mf.MFrontNonlinearProblem(T, material, quadrature_degree=0)
problem.bc = bc
j = mf.ThermalFlux(T)
j.initialize_functions(mesh, 0)
problem.flux = j
H = mf.Flux(mf.Var(T), name="Enthalpy")
H.initialize_functions(mesh, 0)
problem.state_variables.append(H)

H0 = H.previous()
j0 = j.previous()
problem.L = (problem.u_*(H.function - H0)-dt*(theta*dot(grad(problem.u_), j.function)+
                 -(1-theta)*dot(grad(problem.u_), j0)))*problem.dx
problem.a = (problem.u_*H.tangents[0]*problem.du -
             dt*theta*dot(grad(problem.u_), j.tangents[0]*grad(problem.du) +
                          j.tangents[1]*problem.du))*problem.dx

print(problem.material.data_manager.K.shape)

T.interpolate(Constant(Tl))

problem.solve(T.vector())

x = np.linspace(0, length)
import matplotlib.pyplot as plt
plt.plot(x, np.array([T(xi, width/2) for xi in x]))
#plt.figure()
#plot(project(flux.function[0], FunctionSpace(mesh, "DG",0)))