import numpy as np

from pysit import *

from pysit.gallery.gallery_base import PrecomputedGalleryModel

__all__ = ['MarmousiModel2', 'marmousi2']

class MarmousiModel2(PrecomputedGalleryModel):

    """ Marmousi II community velocity model.
    """

    # Names
    model_name = "Marmousi2"

    # A sanitized name for filesystem work
    fs_full_name = "marmousi2"
    fs_short_name = "marm2"

    # Available data
    supported_physics = ['acoustic', 'variable-density-acoustic', 'elastic']
    supported_physical_parameters = ['density', 'vp', 'vs']

    # Descriptive data
    valid_dimensions = (2,)
    @property #read only
    def dimension(self):
        return 2

    # File information
    _local_parameter_filenames = { 'vp' : 'vp_marmousi-ii.segy.gz',
                                   'vs' : 'vs_marmousi-ii.segy.gz',
                                   'density' : 'density_marmousi-ii.segy.gz'}

    _parameter_scale_factor = { 'vp' : 1000.0,
                                'vs' : 1000.0,
                                'density' : 1.0}

    _vp_file_sources = ['http://www.agl.uh.edu/downloads/vp_marmousi-ii.segy.gz']
    _vs_file_sources = ['http://www.agl.uh.edu/downloads/vs_marmousi-ii.segy.gz']
    _density_file_sources = ['http://www.agl.uh.edu/downloads/density_marmousi-ii.segy.gz']

    _remote_file_sources = {'vp' : _vp_file_sources,
                            'vs' : _vs_file_sources,
                            'density' : _density_file_sources}

    _model_transposed = False

    # Model specification
    base_physical_origin =  np.array([0.0, 0.0])
    base_physical_size = np.array([17000.0, 3500.0])
    base_physical_dimensions_units = ('m', 'm')

    base_pixels = np.array([13601, 2801])
    base_pixel_scale = np.array([1.25, 1.25])
    base_pixel_units = ('m', 'm')


    # Water properties specify the way the water masking is handled
    # (None, ) indicates no water
    # ('depth', <depth: float>) specifies that a preset depth from the base_physical_origin is to be used
    # ('mask', ) indicates that the specified mask name should be used
    water_properties = ('depth', 473.0) #m, see Marmousi2, an elasticc upgrade for Marmousi

    _initial_configs = {'smooth_width': {'sigma':1000.0},
                        'smooth_low_pass': {'freq':1./1000.},
                        'constant': {'velocity': 3000.0},
                        'gradient': {'min':1500.0, 'max':3000}}

    _scale_map = {'full':   np.array([ 1.25,  1.25]),
                  'large':  np.array([ 2.50,  2.50]),
                  'medium': np.array([ 5.00,  5.00]),
                  'small':  np.array([10.00, 10.00]),
                  'mini':   np.array([20.00, 20.00]),}

    # old marmousi mini-square had dx=24m
    patches = { 'original' : {'origin': np.array([3950, 440.0]),
                              'size': np.array([9200.0,3000.0])},

                'mini-square' : {'origin': np.array([3950+195*24.0, 440.0]),
                                 'size': np.array([128*24.0, 128*24.0])},

                'left-layers-square' : {'origin': np.array([1500.0, 0.0]),
                                        'size': np.array([3000.0, 3000.0])},

                'right-layers-square' : {'origin': np.array([11630.0, 0.0]),
                                         'size': np.array([3000.0, 3000.0])},

                'fracture-square' : {'origin': np.array([8630.0, 0.0]),
                                     'size': np.array([3000.0, 3000.0])},
              }

def download(parameter='all'):

    if parameter == 'all':
        for p in MarmousiModel2.supported_physical_parameters:
            MarmousiModel2._download_and_prepare(p)
        for p in MarmousiModel2.supported_masks:
            MarmousiModel2._download_and_prepare(p)
    else:
        MarmousiModel2._download_and_prepare(parameter)

def marmousi2(patch=None, compact=False, **kwargs):

    model_config = dict(physics='acoustic',
                        origin=None,
                        size=None,
                        pixel_scale='mini',
                        pixels=None,
                        initial_model_style='smooth_low_pass',
                        initial_config={},
                        fix_water_layer=True,)

    # Make any changes
    model_config.update(kwargs)

    if patch in MarmousiModel2.patches:
        model_config.update(MarmousiModel2.patches[patch])

    C, C0, m, d = MarmousiModel2(**model_config).get_setup()
    if compact:
        # we have to specify in the PML that we want to use the compact version of Helmholtz
        # operator (without auxiliary fields)
        m.x.lbc.domain_bc.compact = True

    return C, C0, m, d

if __name__ == '__main__':

    import matplotlib.pyplot as plt

    M = MarmousiModel2(initial_model_style='smooth_low_pass', pixel_scale='mini')

    C, C0, m, d = M.get_setup()
#   C, C0, m, d = marmousi2(patch='mini-square')

    fig = plt.figure()
    fig.add_subplot(2,1,1)
    vis.plot(C, m)
    fig.add_subplot(2,1,2)
    vis.plot(C0, m)
    plt.show()

    M = MarmousiModel2(initial_model_style='smooth_low_pass', pixel_scale='medium')

    C, C0, m, d = M.get_setup()
#   C, C0, m, d = marmousi2(patch='mini-square')

    fig = plt.figure()
    fig.add_subplot(2,1,1)
    vis.plot(C, m)
    fig.add_subplot(2,1,2)
    vis.plot(C0, m)
    plt.show()
numpy.array