from Cuboid_1 import *
from sfepy.mechanics.matcoefs import stiffness_from_youngpoisson
from sfepy.postprocess.probes_vtk import Probe

import os

def stress_strain(out, pb, state, extend=False):
    """
    Calculate and output strain and stress for given displacements.
    """
    from sfepy.base.base import Struct
    import matplotlib.pyplot as plt
    import matplotlib.font_manager as fm

    ev = pb.evaluate
    strain = ev('ev_cauchy_strain.3.Omega(u)', mode='el_avg')
    stress = ev('ev_cauchy_stress.3.Omega(Steel.D, u)', mode='el_avg')

    out['cauchy_strain'] = Struct(name='output_data', mode='cell',
                                  data=strain, dofs=3)
    out['cauchy_stress'] = Struct(name='output_data', mode='cell',
                                  data=stress, dofs=3)

    probe = Probe(out, pb.domain.mesh, probe_view=True)

    ps0 = [[0.0,  0.0, 0.0], [ 0.0,  0.0, 0.0],[ 0.0,  0.0, 0.0]]
    ps1 = [[5.0,  0.0, 0.0], [ 0.0, 1.0, 0.0], [ 0.0,  0.0, 1.0]]
    n_point = 1000

    labels = ['%s -> %s' % (p0, p1) for p0, p1 in zip(ps0, ps1)]
    probes = []
    for ip in xrange(len(ps0)):
        p0, p1 = ps0[ip], ps1[ip]
        probes.append('line%d' % ip)
        probe.add_line_probe('line%d' % ip, p0, p1, n_point)

    for ip, label in zip(probes, labels):
        fig = plt.figure()
        plt.clf()
        fig.subplots_adjust(hspace=0.4)
        plt.subplot(311)
        pars, vals = probe(ip, 'u')
        for ic in range(vals.shape[1] - 1):
            plt.plot(pars, vals[:,ic], label=r'$u_{%d}$' % (ic + 1),
                     lw=1, ls='-', marker='+', ms=3)
        plt.ylabel('displacements')
        plt.xlabel('probe %s' % label, fontsize=8)
        plt.legend(loc='best', prop=fm.FontProperties(size=10))

        sym_indices = [0, 4, 1]
        sym_labels = ['11', '22', '12']

        plt.subplot(312)
        pars, vals = probe(ip, 'cauchy_strain')
        for ii, ic in enumerate(sym_indices):
            plt.plot(pars, vals[:,ic], label=r'$e_{%s}$' % sym_labels[ii],
                     lw=1, ls='-', marker='+', ms=3)
        plt.ylabel('Cauchy strain')
        plt.xlabel('probe %s' % label, fontsize=8)
        plt.legend(loc='best', prop=fm.FontProperties(size=8))

        plt.subplot(313)
        pars, vals = probe(ip, 'cauchy_stress')
        for ii, ic in enumerate(sym_indices):
            plt.plot(pars, vals[:,ic], label=r'$\sigma_{%s}$' % sym_labels[ii],
                     lw=1, ls='-', marker='+', ms=3)
        plt.ylabel('Cauchy stress')
        plt.xlabel('probe %s' % label, fontsize=8)
        plt.legend(loc='best', prop=fm.FontProperties(size=8))

        opts = pb.conf.options
        filename_results = os.path.join(opts.get('output_dir'),
                                        'Cuboid_probe_%s.png' % ip)

        fig.savefig(filename_results)

    return out

materials['Steel'][0].update({'D' : stiffness_from_youngpoisson(2, young, poisson)})

options.update({
    'post_process_hook' : 'stress_strain',
})