#!/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, symmetric=None):
        self.variable = variable
        if symmetric is None:
            self.expression = expression
        elif symmetric:
            if self.variable.geometric_dimension()==2:
                self.expression = as_vector([symmetric_tensor_to_vector(expression)[i] for i in range(4)])
            else:
                self.expression = symmetric_tensor_to_vector(expression)
        else:
            if self.variable.geometric_dimension()==2:
                self.expression = as_vector([nonsymmetric_tensor_to_vector(expression)[i] for i in range(5)])
            else:
                self.expression = nonsymmetric_tensor_to_vector(expression)
        shape = ufl.shape(self.expression)
        if len(shape)==1:
            self.shape = shape[0]
        elif shape==():
            self.shape = 0
            #self.expression = as_vector([self.expression])
        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, name, shape, gradients=[]):
        self.shape = shape
        self.name = name
        self.gradients = gradients

    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)
        values = [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]
        keys = [g.name for g in self.gradients]
        self.tangent_blocks = dict(zip(keys, values))

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

    def previous(self):
        try:
            self.previous = Function(self.function.function_space(), name=self.name+"_previous")
            self.previous.assign(self.function)
            return self.previous
        except:
            raise ValueError("Function must be initialized first.")

# 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])