#!/usr/bin/env python
# 12.01.2007, c 
import os.path as op
import shutil
from optparse import OptionParser

import sfepy
from sfepy.base.base import *
from sfepy.base.conf import ProblemConf, get_standard_keywords
from sfepy.fem import ProblemDefinition
from sfepy.fem.evaluate import assemble_by_blocks
from sfepy.homogenization.phono import transform_plot_data, plot_logs, \
     plot_gaps, detect_band_gaps, compute_cat, compute_polarization_angles
from sfepy.homogenization.engine import HomogenizationEngine
from sfepy.applications import SimpleApp
from sfepy.solvers import Solver, eig
from sfepy.base.plotutils import plt

def make_save_hook( base_name, post_process_hook = None, file_per_var = None ):
    def save_phono_correctors( state, problem, ir, ic ):
        problem.save_state( (base_name % (ir, ic)) + '.vtk', state,
                            post_process_hook = post_process_hook,
                            file_per_var = file_per_var )
    return save_phono_correctors

def try_set_defaults( obj, attr, defaults ):
    try:
        values = getattr( obj, attr )
        set_defaults( values, defaults )
    except:
        values = defaults
    return values
    
def report_iw_cat( iw_dir, christoffel ):
    output( 'incident wave direction:' )
    output( iw_dir )
    output( 'Christoffel acoustic tensor:' )
    output( christoffel )

class AcousticBandGapsApp( SimpleApp ):

    def process_options( options ):
        """Application options setup. Sets default values for missing
        non-compulsory options."""
        get = options.get_default_attr
        
        clear_cache = get( 'clear_cache', {} )
        eigensolver = get( 'eigensolver', 'eig.sgscipy' )
        eig_problem = get( 'eig_problem', 'simple' )
        schur = get( 'schur', None )

        elasticity_contrast = get( 'elasticity_contrast', 1.0 )
        scale_epsilon = get( 'scale_epsilon', 1.0 )

        incident_wave_dir = get( 'incident_wave_dir', None )
        dispersion = get( 'dispersion', 'simple' )
        dispersion_conf = get( 'dispersion_conf', None )
        homogeneous = get( 'homogeneous', False )

        save = get( 'save_eig_vectors', (0, 0) )
        eig_range = get( 'eig_range', None )

        freq_margins = get( 'freq_margins', (5, 5) )
        # Given in per cent.
        freq_margins = 0.01 * nm.array( freq_margins, dtype = nm.float64 )

        fixed_eig_range = get( 'fixed_eig_range', None )

        # Given in per cent.
        freq_step = 0.01 * get( 'freq_step', 5 )

        feps = get( 'feps', 1e-8 )
        zeps = get( 'zeps', 1e-8 )
        teps = get( 'teps', 1e-4 )
        teps_rel = get( 'teps_rel', True )

        eig_vector_transform = get( 'eig_vector_transform', None )
        plot_transform = get( 'plot_transform', None )
        plot_transform_wave = get( 'plot_transform_wave', None )
        plot_transform_angle = get( 'plot_transform_angle', None )

        plot_options = get( 'plot_options', {'show' : True,'legend' : False,} )

        fig_name = get( 'fig_name', None )
        fig_name_wave = get( 'fig_name_wave', None )
        fig_name_angle = get( 'fig_name_angle', None )

        aux = {
            'resonance' : 'eigenfrequencies',
            'masked' : 'masked eigenfrequencies',
            'eig_min' : 'min eig($M^*$)',
            'eig_mid' : 'mid eig($M^*$)',
            'eig_max' : 'max eig($M^*$)',
            'y_axis' : 'eigenvalues of mass matrix $M^*$',
        }
        plot_labels = try_set_defaults( options, 'plot_labels', aux )

        aux = {
            'resonance' : 'eigenfrequencies',
            'masked' : 'masked eigenfrequencies',
            'eig_min' : r'$\kappa$(min)',
            'eig_mid' : r'$\kappa$(mid)',
            'eig_max' : r'$\kappa$(max)',
            'y_axis' : 'polarization angles',
        }
        plot_labels_angle = try_set_defaults( options, 'plot_labels_angle', aux )

        aux = {
            'resonance' : 'eigenfrequencies',
            'masked' : 'masked eigenfrequencies',
            'eig_min' : r'wave number (min)',
            'eig_mid' : r'wave number (mid)',
            'eig_max' : r'wave number (max)',
            'y_axis' : 'wave numbers',
        }
        plot_labels_wave = try_set_defaults( options, 'plot_labels_wave', aux )

        plot_rsc =  {
            'resonance' : {'linewidth' : 0.5, 'color' : 'r', 'linestyle' : '-' },
            'masked' : {'linewidth' : 0.5, 'color' : 'r', 'linestyle' : ':' },
            'x_axis' : {'linewidth' : 0.5, 'color' : 'k', 'linestyle' : '--' },
            'eig_min' : {'linewidth' : 0.5, 'color' : 'b', 'linestyle' : '--' },
            'eig_mid' : {'linewidth' : 0.5, 'color' : 'b', 'linestyle' : '-.' },
            'eig_max' : {'linewidth' : 0.5, 'color' : 'b', 'linestyle' : '-' },
            'strong_gap' : {'linewidth' : 0, 'facecolor' : (1, 1, 0.5) },
            'weak_gap' : {'linewidth' : 0, 'facecolor' : (1, 1, 1) },
            'propagation' : {'linewidth' : 0, 'facecolor' : (0.5, 1, 0.5) },
            'params' : {'axes.labelsize': 'large',
                        'text.fontsize': 'large',
                        'legend.fontsize': 'large',
                        'xtick.labelsize': 'large',
                        'ytick.labelsize': 'large',
                        'text.usetex': False},
        }
        plot_rsc = try_set_defaults( options, 'plot_rsc', plot_rsc )

        eigenmomentum = get( 'eigenmomentum', None,
                             'missing "eigenmomentum" in options!' )
        
        region_to_material = get( 'region_to_material', None,
                                  'missing "region_to_material" in options!' )

        tensor_names = get( 'tensor_names', None,
                            'missing "tensor_names" in options!' )

        volume = get( 'volume', None, 'missing "volume" in options!' )

        if eig_problem == 'simple_liquid':
            liquid_region = get('liquid_region', None,
                                'missing "liquid_region" in options!')
        else:
            liquid_region = None

        return Struct( **locals() )
    process_options = staticmethod( process_options )

    def process_options_pv( options ):
        """Application options setup for phase velocity computation. Sets
        default values for missing non-compulsory options."""
        get = options.get_default_attr

        clear_cache = get( 'clear_cache', {} )
        eigensolver = get( 'eigensolver', 'eig.sgscipy' )

        incident_wave_dir = get( 'incident_wave_dir', None )
        dispersion = get( 'dispersion', 'simple' )
        dispersion_conf = get( 'dispersion_conf', None )
        homogeneous = get( 'homogeneous', False )
        fig_suffix = get( 'fig_suffix', '.pdf' )

        region_to_material = get( 'region_to_material', None,
                                  'missing "region_to_material" in options!' )

        tensor_names = get( 'tensor_names', None,
                            'missing "tensor_names" in options!' )

        volume = get( 'volume', None, 'missing "volume" in options!' )

        return Struct( **locals() )
    process_options_pv = staticmethod( process_options_pv )

    def __init__( self, conf, options, output_prefix, **kwargs ):
        SimpleApp.__init__( self, conf, options, output_prefix,
                            init_equations = False )

        self.setup_options()
        self.cached_coefs = None
        self.cached_iw_dir = None
        self.cached_christoffel = None
        self.cached_evp = None

        output_dir = self.problem.output_dir
        shutil.copyfile( conf._filename,
                         op.join( output_dir, op.basename( conf._filename ) ) )

    def setup_options( self ):
        SimpleApp.setup_options( self )
        
        if self.options.phase_velocity:
            process_options = AcousticBandGapsApp.process_options_pv
        else:
            process_options = AcousticBandGapsApp.process_options
        self.app_options += process_options( self.conf.options )
    
    def call( self ):
        """In parametric runs, cached data (homogenized coefficients,
        Christoffel acoustic tensor and eigenvalue problem solution) are
        cleared according to 'clear_cache' aplication options.

        Example:

        clear_cache = {'cached_christoffel' : True, 'cached_evp' : True}
        """
        options = self.options

        for key, val in self.app_options.clear_cache.iteritems():
            if val and key.startswith('cached_'):
                setattr(self, key, None)

        if options.phase_velocity:
            # No band gaps in this case.
            return self.compute_phase_velocity()

        evp = self.solve_eigen_problem()

        self.fix_eig_range( evp.eigs.shape[0] )

        if options.detect_band_gaps:
            bg = detect_band_gaps( self.problem, evp.kind,
                                   evp.eigs_rescaled, evp.eig_vectors,
                                   self.app_options, self.conf.funmod )

            if options.plot:
                plot_range, teigs = transform_plot_data( bg.logs.eigs,
                                                         bg.opts.plot_transform,
                                                         self.conf.funmod )

                plot_rsc = bg.opts.plot_rsc
                plot_opts =  bg.opts.plot_options
                plot_labels =  bg.opts.plot_labels
                
                plt.rcParams.update( plot_rsc['params'] )

                fig = plot_gaps( 1, plot_rsc, bg.gaps, bg.kinds,
                                 bg.freq_range_margins, plot_range,
                                 clear = True )
                fig = plot_logs( 1, plot_rsc, plot_labels, bg.logs.freqs, teigs,
                                 bg.valid[bg.eig_range],
                                 bg.freq_range_initial,
                                 plot_range, False,
                                 show_legend = plot_opts['legend'],
                                 new_axes = True )

                fig_name = bg.opts.fig_name
                if fig_name is not None:
                    fig.savefig( fig_name )
                if plot_opts['show']:
                    plt.show()

        elif options.analyze_dispersion:
            christoffel, iw_dir = self.compute_cat(ret_iw_dir=True)

            bg = detect_band_gaps( self.problem, evp.kind,
                                   evp.eigs_rescaled, evp.eig_vectors,
                                   self.app_options, self.conf.funmod,
                                   christoffel = christoffel )

            output( 'computing polarization angles...' )
            pas = compute_polarization_angles( iw_dir, bg.logs.eig_vectors )
            output( '...done' )

            bg.polarization_angles = pas

            output( 'computing phase velocity...' )
            bg.phase_velocity = self.compute_phase_velocity()
            output( '...done' )

            if options.plot:
                plot_rsc = bg.opts.plot_rsc
                plot_opts =  bg.opts.plot_options
                plt.rcParams.update( plot_rsc['params'] )

                aux = transform_plot_data( pas,
                                           bg.opts.plot_transform_angle,
                                           self.conf.funmod )
                plot_range, pas = aux
                
                plot_labels =  bg.opts.plot_labels_angle

                fig = plot_gaps( 1, plot_rsc, bg.gaps, bg.kinds,
                                 bg.freq_range_margins, plot_range,
                                 clear = True )
                fig = plot_logs( 1, plot_rsc, plot_labels, bg.logs.freqs, pas,
                                 bg.valid[bg.eig_range],
                                 bg.freq_range_initial,
                                 plot_range, False,
                                 show_legend = plot_opts['legend'],
                                 new_axes = True )

                fig_name = bg.opts.fig_name_angle
                if fig_name is not None:
                    fig.savefig( fig_name )

                aux = transform_plot_data( bg.logs.eigs,
                                           bg.opts.plot_transform_wave,
                                           self.conf.funmod )
                plot_range, teigs = aux

                plot_labels =  bg.opts.plot_labels_wave

                fig = plot_gaps( 2, plot_rsc, bg.gaps, bg.kinds,
                                 bg.freq_range_margins, plot_range,
                                 clear = True )
                fig = plot_logs( 2, plot_rsc, plot_labels, bg.logs.freqs, teigs,
                                 bg.valid[bg.eig_range],
                                 bg.freq_range_initial,
                                 plot_range, False,
                                 show_legend = plot_opts['legend'],
                                 new_axes = True )

                fig_name = bg.opts.fig_name_wave
                if fig_name is not None:
                    fig.savefig( fig_name )
                if plot_opts['show']:
                    
plt.show()