#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Feb 24 14:36:40 2020

@author: bleyerj
"""
from dolfin import *
import numpy as np
import ufl
from mfront_wrapper.utils import *

class Gradient:
    def __init__(self, variable, expression, name=None):
        self.variable = variable
        self.expression = expression
        shape = ufl.shape(self.expression)
        if len(shape)==1:
            self.shape = shape[0]
        elif shape==():
            self.shape = 0
        else:
            self.shape = shape
        if name is None:
            self.name = "Gradient-{}".format(np.random.randint(1e6))
        else:
            self.name = name
    def __call__(self, v):
        return ufl.replace(self.expression, {self.variable:v})
    def variation(self, v):
        """ Directional derivative in direction v """
        return ufl.derivative(self.expression, self.variable, v)

    def initialize_function(self, mesh, quadrature_degree):
        We = get_quadrature_element(mesh.ufl_cell(), quadrature_degree, self.shape)
        self.function_space = FunctionSpace(mesh, We)
        self.function = Function(self.function_space, name=self.name)

    def update(self, x):
        self.function.vector().set_local(x)

class Grad(Gradient):
    """ Gradient operator : grad(u) """
    def __init__(self, variable, name=None):
        Gradient.__init__(self, variable,
                          nonsymmetric_tensor_to_vector(grad(variable)), name)

class GradT(Gradient):
    """ Transposed gradient operator : grad(u).T """
    def __init__(self, variable, name=None):
        Gradient.__init__(self, variable,
                          nonsymmetric_tensor_to_vector(grad(variable).T), name)

class SymGrad(Gradient):
    """ Symmetric gradient operator : sym(grad(u)) """
    def __init__(self, variable, name=None):
        Gradient.__init__(self, variable,
                          symmetric_tensor_to_vector(sym(grad(variable))), name)

class Grad_Id(Gradient):
    """ Transformation gradient operator : grad(u) + Id(d) """
    def __init__(self, variable, name=None):
        Gradient.__init__(self, variable,
                          symmetric_tensor_to_vector(grad(variable) +
                          Identity(variable.geometric_dimension()), T22=1), name)
class Var(Gradient):
    """ The variable itself : u """
    def __init__(self, variable, name=None):
        Gradient.__init__(self, variable, variable, name)


class Flux:
    def __init__(self, gradients, shape=None, dual_gradient=None, name=None):
        if type(gradients)==list:
            self.gradients = gradients
        else:
            self.gradients = [gradients]
        self.num_gradients = len(self.gradients)
        self.gradient_shapes = [g.shape for g in self.gradients]
        if shape is None:
            assert len(set(self.gradient_shapes))==1, "Cannot deduce flux shape from non-unique gradient shapes"
            self.shape = self.gradient_shapes[0]
        else:
            self.shape = shape
        if dual_gradient is None:
            self.dual_gradient = self.gradients[0]
        else:
            self.dual_gradient = dual_gradient
        if name is None:
            self.name = "Flux-{}".format(np.random.randint(1e6))
        else:
            self.name = name

    def initialize_functions(self, mesh, quadrature_degree):
        We = get_quadrature_element(mesh.ufl_cell(), quadrature_degree, self.shape)
        W = FunctionSpace(mesh, We)
        self.function = Function(W, name=self.name)
        self.tangents = [Function(FunctionSpace(mesh,
                          get_quadrature_element(mesh.ufl_cell(),
                          quadrature_degree, dim=(self.shape, g.shape))), name="d{}_d{}".format(self.name, g.name))
                          for g in self.gradients]
        for g in self.gradients:
            g.initialize_function(mesh, quadrature_degree)

    def update(self, x):
        self.function.vector().set_local(x)



class Stress(Flux):
    def __init__(self, variable, name="sigma"):
        Flux.__init__(self, SymGrad(variable, name="eps"), name=name)


class ThermalFlux(Flux):
    def __init__(self, variable, name="j"):
        Flux.__init__(self, [Grad(variable, name="grad(T)"), Var(variable, name="T")],
                             variable.geometric_dimension(), name=name)

# mesh = UnitSquareMesh(10, 10)
# V = VectorFunctionSpace(mesh, "CG", 1)
# VT = FunctionSpace(mesh, "CG", 1)

# u = Function(V, name="u")
# T = Function(VT, name="Temperature")
# v = Function(V, name="v")
# g = Grad_Id(u, name="F")
# sig = Stress(u)
# j = ThermalFlux(T)
# j.initialize_functions(mesh, 2)
# print([ufl.shape(tang) for tang in j.tangents])
# print([tang.name() for tang in j.tangents])