'''
Created on Feb 11, 2018

@author: rch
'''

import os
from ibvpy.mathkit.tensor import EPS

# from mayavi import mlab
# from tvtk.api import tvtk
from ibvpy.view.plot3d.viz3d import Vis3D, Viz3D

import numpy as np
import traits.api as tr


class Vis3DLattice(Vis3D):
    r'''Record the response triggered by the hist object the time increments
    upon a record event.
    '''
    var = tr.Str('<unnamed>')

    extension = '.npz'

    def setup(self):
        self.new_dir()
        ts = self.tstep
        fe_domain = ts.fe_domain
        X_Ia_list = []
        I_Li_list = []
        Xm_Lia_list = []
        # Consistently with other adaptors, this one allows for
        # a loop over several lattice domains.
        for domain in fe_domain:
            xdomain = domain.xmodel
            var_function = domain.tmodel.var_dict.get(self.var, None)
            if var_function == None or xdomain.hidden:
                continue
            DELTA_x_ab = xdomain.vtk_expand_operator
            X_Ia = np.einsum(
                'Ia,ab->Ib',
                xdomain.X_Ia, DELTA_x_ab
            )
            Xm_Lia = np.einsum(
                'Lib,ba->Lia',
                xdomain.Xm_Lia, DELTA_x_ab
            )
            X_Ia_list.append(X_Ia)
            Xm_Lia_list.append(Xm_Lia)
            I_Li_list.append(xdomain.I_Li)
        if len(I_Li_list) == 0:
            raise ValueError(
                'Empty output for field variable %s in model %s' %
                (self.var, domain.tmodel)
            )
        X_Ia = np.hstack(X_Ia_list)
        I_Li = np.hstack(I_Li_list)
        fname = 'lattice_tessellation'
        target_file = os.path.join(
            self.dir, fname.replace('.', '_')
        ) + '.' + self.extension
        np.savez(target_file,
                 X_Ia=X_Ia, I_Li=I_Li, Xm_Lia=Xm_Lia)
        self.lattice_file = target_file

    def update(self):
        ts = self.tstep
        U = ts.U_k
        t = ts.t_n1
        U_Ia_list = []
        U_Lipa_list = []
        state_fields = []
        for domain in ts.fe_domain:
            xdomain = domain.xmodel
            state_field = domain.state_n.get(self.var, None)
            if (xdomain.hidden):
                # If the state variable not present in the domain, skip
                continue
            if (state_field is not None):
                state_fields.append(state_field.flatten())
            DELTA_x_ab = xdomain.vtk_expand_operator
            U_Ib = U[xdomain.o_Ipa][..., 0, :]
            U_Ia_list.append(np.einsum(
                '...a,ab->...b', U_Ib, DELTA_x_ab
            ))
            U_Lipa = U[xdomain.o_Lipa]
            U_Lipa_list.append(np.einsum(
                '...a,ab->...b', U_Lipa, DELTA_x_ab
            ))
            if len(U_Lipa) == 0:
                raise ValueError('no fields for variable %s' % self.var)
        # Save the displacements
        U_Ia = np.vstack(U_Ia_list)
        U_Lipa = np.vstack(U_Lipa_list)
 #       state_vec = np.hstack(state_fields)
        fname = 'U_step_%008.4f' % t
        target_file = os.path.join(
            self.dir, fname.replace('.', '_')
        ) + '.' + self.extension
        # , state_vec=state_vec)
        np.savez(target_file, U_Ia=U_Ia, U_Lipa=U_Lipa)
        self.add_file(target_file)


class Viz3DLattice(Viz3D):

    warp_factor = tr.Float(1.0, auto_set=False, enter_set=True)

    def setup(self):
        m = mlab

        npzfile = np.load(self.vis3d.lattice_file)
        X_Ia = npzfile['X_Ia']
        I_Li = npzfile['I_Li']
        Xm_Lia = npzfile['Xm_Lia']

        self.d = tvtk.PolyData(points=X_Ia, lines=I_Li)
        self.warp_vector = m.pipeline.warp_vector(self.d)
        self.warp_vector.filter.scale_factor = self.warp_factor

        tube = m.pipeline.tube(self.warp_vector,
                               tube_radius=self.tube_radius)
        self.lines = m.pipeline.surface(tube, color=(0.1, 0.1, 0.1))

        X_aiL = np.einsum('Lia->aiL', X_Ia[I_Li])
        Xm_aiL = np.einsum('Lia->aiL', Xm_Lia)
        XIm_aiL = np.concatenate(
            [
                np.einsum('iaL->aiL',
                          np.array([X_aiL[:, 0, ...], Xm_aiL[:, 0, ...]])),
                np.einsum('iaL->aiL',
                          np.array([X_aiL[:, 1, ...], Xm_aiL[:, 0, ...]]))
            ], axis=-1
        )
        XIm_aLi = np.einsum('aiL->aLi', XIm_aiL)
        XIm_aJ = XIm_aLi.reshape(3, -1)
        XIm_Ja = np.einsum('aJ->Ja', XIm_aJ)
        J_Li = np.arange(XIm_aJ.shape[1]).reshape(-1, 2)
        self.d_J = tvtk.PolyData(points=XIm_Ja, lines=J_Li)
        self.warp_vector_J = m.pipeline.warp_vector(self.d_J)
        self.warp_vector_J.filter.scale_factor = self.warp_factor

        tube_J = m.pipeline.tube(self.warp_vector_J,
                                 tube_radius=self.tube_radius)
        self.lines_J = m.pipeline.surface(tube_J, color=(0.1, 0.1, 0.1))
        glyph = m.pipeline.glyph(self.warp_vector_J)

        glyph.glyph.glyph_source.glyph_source = glyph.glyph.glyph_source.glyph_dict[
            'cube_source']
        glyph.glyph.glyph.range = np.array([0., 0.1])
        glyph.glyph.scale_mode = 'scale_by_vector'
        glyph.glyph.color_mode = 'color_by_vector'

    tube_radius = tr.Float(0.01)

    visible = tr.Property(tr.Bool)

    def _set_visible(self, visible):
        return
        self.d.visible = visible

    def _get_visible(self):
        return True  # self.d.visible

    def plot(self, vot):
        fe_domain = self.vis3d.tstep.fe_domain
        xdomain = fe_domain[0].xmodel
        dXm_Lia = xdomain.dXm_Lia

        idx = self.vis3d.tstep.hist.get_time_idx(vot)
        fname = self.vis3d.file_list[idx]
        npzfile = np.load(fname)
        U_Ia = npzfile['U_Ia']
        self.d.point_data.vectors = U_Ia

        U_Lipa = npzfile['U_Lipa']
        Um_trans_Lia, Phi_Lia = np.einsum('Lipa->pLia', U_Lipa)
        Um_rot_Lia = np.einsum('abc,...b,...c->...a',
                               EPS, Phi_Lia, dXm_Lia)
        Um_Lia = Um_trans_Lia + Um_rot_Lia
        Um_trans_aiL = np.einsum('Lia->aiL', Um_trans_Lia)
        Um_aiL = np.einsum('Lia->aiL', Um_Lia)
        UIm_aiL = np.concatenate(
            [
                np.einsum('iaL->aiL',
                          np.array([Um_trans_aiL[:, 0, ...], Um_aiL[:, 0, ...]])),
                np.einsum('iaL->aiL',
                          np.array([Um_trans_aiL[:, 1, ...], Um_aiL[:, 1, ...]]))
            ], axis=-1
        )
        UIm_aLi = np.einsum('aiL->aLi', UIm_aiL)
        UIm_aJ = UIm_aLi.reshape(3, -1)
        UIm_Ja = np.einsum('aJ->Ja', UIm_aJ)
        self.d_J.point_data.vectors = UIm_Ja
