import numpy as np

from sfepy.base.conf import ProblemConf
from sfepy.discrete.problem import Problem
from sfepy.mesh.mesh_generators import gen_block_mesh

from sfepy.postprocess.viewer import Viewer
from scipy.interpolate import interpn

import matplotlib.pyplot as plt



from matplotlib import ticker
    

    



class Block2d:
    
    """ 
    Small wrapper for sfepy problem class.
    
    Constructs a rectangle of specificed dimensions with a rectangular FEM mesh.
    Rectangle is gripped at one end and pulled away at the other end, by a specified 
    distance. FEM mesh sixe can be specified. Centre of rectangle can be specified. 
    
    Having done this, this class exposes the stretch amount and the distribution 
    of the stiffness tensor D as paramters that can be set, and allows querying 
    of the displacement at a set of points in the domain. 
    
    rectangle is stretched parallel to the y-axis. short side of rectangle is x-axis. 
    
    Also drawing. """
    
    def __init__(self, dims, center_location, cell_sizes=np.array([2,2]),
                 prob_file="C:\\Users\\wbald\\sfepythings\\blocks\\fem\\prob_desc_2d.py", 
                 put_mesh="C:\\Users\\wbald\\sfepythings"):
        
        """ 
        dims:               array, dimensions of rectangle [x,y] 
        center_location:    array, centre of rectangle [x,y]
        cell sizes:         array, x and y side length of FEM rectangular elements
        stretch:            default distance by which to displace the upper y axis edge. 
        prob_file:          problem description file
        put_mesh:           where to save the mesh file
        """
        
        assert(dims.shape[0] == 2)
        assert(cell_sizes.shape[0] == 2)
        
        # assume linear elasticity. Fix strain of rectangle to 0.001 and query
        # at different strains by scaling linearly
        self.dims = dims
        self.prob_file = prob_file        
        self.FEM_model_strain = 0.001
        self.elongation= self.FEM_model_strain * self.dims[1]
        
        

        nums = np.divide(dims, cell_sizes)
        nums = np.around(nums).astype(int) + np.array([1,1])
        
        blockmesh = gen_block_mesh(dims, nums, center_location)
        blockmesh.write(put_mesh + '\\mesh.vtk')
        
        conf = ProblemConf.from_file(prob_file)
        
        # 'region_ylower__1' holds the edge to be fixed
        # 'region_yupper__2' holds the edge to be displaced
        conf.regions['region_ylower__1'].select = 'vertices in (y < ' + str(0.01) + ')'
        conf.regions['region_yupper__2'].select = 'vertices in (y > ' + str(dims[1] - 0.01) + ')'
        
        conf.ebcs['ebc_Displaced__1'].dofs['u.1'] = self.elongation
        
        self.prob = 
Problem.from_conf(conf)