from numpy.core.umath import *
from scipy import *
from numpy import limits, display_test
import sys

_have_wx = 0
if not display_test.have_x11() or display_test.try_XOpenDisplay():
    try:
        import wxPython
        _have_wx = 1
    except ImportError,msg:
        print __file__,msg

if _have_wx:
    import wxplt
    import plot_objects
    import gui_thread
    plot_module = wxplt
    plot_class = gui_thread.register(plot_module.plot_frame)

_figure = []
_active = None


def figure(which_one = None):
    global _figure; global _active
    if which_one is None:
        title ='Figure %d' % len(_figure)
        _figure.append(plot_class(title=title))
        _active = _figure[-1]
    elif (type(which_one) == type(1)) or (type(which_one) == type(1.)):
        try:
            _active = _figure[int(which_one)]
            _active.Raise()
        except IndexError:
            msg = "There are currently only %d active figures" % len(_figure)
            raise IndexError, msg
    elif which_one in _figure:
        _active = which_one
        _active.Raise()
    else:
        try:
            if which_one.__type_hack__ == "plot_canvas":
                _active = which_one
                _figure.append(_active)
                _active.Raise()
            else:
                raise ValueError, "The specified figure or index is not not known"
        except (AtrributeError):
            pass
    fig = current()
    return fig


def validate_active():
    global _active
    if _active is None: figure()
    try:
        if not _active.proxy_object_alive:
            _active = None
            figure()
    except:
        pass

def current():
    return _active

def redraw():
    validate_active()
    _active.redraw()

def close(which_one = None):
    global _figure; global _active
    if which_one is None:
        try:
            _active.Close()
            _figure.remove(_active)
        except ValueError:
            pass

        try:
            # should make sure the new plot window really exist
            set_new_active()
        except IndexError: _active = None
    elif which_one == 'all':
        for fig in _figure: fig.Close()
        _active = None
    else:
        raise NotImplementedError, "currently close only works with"\
                                   " _active window or 'all'"
        #try:
        #   _figure.remove(which_one)
        #   which_one.close()
        #except ValueError:
        #   which_one.close()

def set_new_active():
    # should validate new active here
    try:
        _active = _figure[-1]
    except IndexError:
        _active = None


def _auto_all():
    validate_active()
    _active.x_axis.bounds = ['auto','auto']
    _active.y_axis.bounds = ['auto','auto']
    _active.x_axis.tick_interval = 'auto'
    _active.y_axis.tick_interval = 'auto'
def autoscale():
    validate_active()
    _auto_all()
    _active.update()

def _an_axis(ax,setting):
    ticks = ax.ticks
    interval = ax.ticks[1]- ax.ticks[0]
    if setting in ['normal','auto']:
        ax.bounds = ['auto','auto']
    elif setting == 'freeze':
        ax.bounds = [axes[0],axes[1]]
        ax.tick_interval = interval
    elif setting in ['tight','fit']:
        ax.bounds = ['fit','fit']
        ax.tick_interval = 'auto'
    else:
        ax.bounds = [setting[0],setting[1]]
        if len(setting) > 2:
            ax.tick_interval = setting[2]

def xaxis(rng):
    validate_active()
    _an_axis(_active.x_axis,rng)
    _active.update()

def yaxis(rng):
    validate_active()
    _an_axis(_active.y_axis,rng)
    _active.update()

def title(name):
    validate_active()
    _active.title.text = name
    _active.update()

def xtitle(name):
    validate_active()
    _active.x_title.text = name
    _active.update()

def ytitle(name):
    validate_active()
    _active.y_title.text = name
    _active.update()

on = 'on'
off = 'off'
def grid(state=None):
    validate_active()
    if state is None:
        if _active.x_axis.grid_visible in ['on','yes']:
            _active.x_axis.grid_visible = 'off'
            _active.y_axis.grid_visible = 'off'
        else:
            _active.x_axis.grid_visible = 'on'
            _active.y_axis.grid_visible = 'on'
    elif state in ['on','off','yes','no']:
        _active.x_axis.grid_visible = state
        _active.y_axis.grid_visible = state
    else:
        raise ValueError, 'grid argument can be "on","off",'\
                          '"yes","no". Not ' + state
    _active.update()

def hold(state):
    validate_active()
    if state in ['on','off','yes','no']:
        _active.hold = state
    else:
        raise ValueError, 'holds argument can be "on","off",'\
                          '"yes","no". Not ' + state

def axis(setting):
    validate_active()
    x_ticks = _active.x_axis.ticks
    print type(x_ticks), dir(x_ticks)
    x_interval = x_ticks[1]- x_ticks[0]
    y_ticks = _active.y_axis.ticks
    x_interval = x_ticks[1]- y_ticks[0]
    axes = array((x_ticks[0],x_ticks[-1],y_ticks[0],y_ticks[-1]),Float)
    # had to use client below cause of __setattr__ troubles in plot_frame
    if setting == 'normal':
        _active.client.aspect_ratio = setting
        _auto_all()
    elif setting == 'equal':
        _active.client.aspect_ratio = setting
    elif setting == 'freeze':
        _active.x_axis.bounds = [axes[0],axes[1]]
        _active.y_axis.bounds = [axes[2],axes[3]]
        _active.x_axis.tick_interval = x_interval
        _active.x_axis.tick_interval = y_interval
    elif setting in ['tight','fit']:
        _active.x_axis.bounds = ['fit','fit']
        _active.y_axis.bounds = ['fit','fit']
        _active.x_axis.tick_interval = 'auto'
        _active.x_axis.tick_interval = 'auto'
    else:
        _active.x_axis.bounds = [setting[0],setting[1]]
        _active.y_axis.bounds = [setting[2],setting[3]]
    _active.update()

def save(file_name,format='png'):
    _active.save(file_name,format)


##########################################################
#----------------- plotting machinery -------------------#
##########################################################

#---- array utilities ------------

def is1D(a):
    as = shape(a)
    if(len(as) == 1):
        return 1
    if(as[0] == 1 or as[1]==1):
        return 1
    return 0

def row(a):
    return reshape(asarray(a),[1,-1])
def col(a):
    return reshape(asarray(a),[-1,1])

SizeMismatch = 'SizeMismatch'
SizeError = 'SizeError'
NotImplemented = 'NotImplemented'

#------------ Numerical constants ----------------

# really should do better than this...
BIG = 1e20
SMALL = 1e-20

#------------ plot group parsing -----------------
from types import *

def plot_groups(data):
    remains = data; groups = []
    while len(remains):
        group,remains = get_plot_group(remains)
        groups.append(group)
    return groups

def get_plot_group(data):
    group = ()
    remains = data
    state = 0
    finished = 0
    while(len(remains) > 0 and not finished):
        el = remains[0]
        if(state == 0):
            el = asarray(el)
            state = 1
        elif(state == 1):
            if(type(el) == StringType):
                finished = 1
            else:
                el = asarray(el)
            state = 2
        elif(state == 2):
            finished = 1
            if(type(el) != StringType):
                break
        try:
            if el.typecode() == 'D':
                print 'warning plotting magnitude of complex values'
                el = abs(el)
        except:
            pass
        group = group + (el,)
        remains = remains[1:]
    return group, remains

def hstack(tup):
        #horizontal stack (column wise)
    return concatenate(tup,1)

def lines_from_group(group):
    lines = []
    plotinfo = ''
    x = group[0]
    ar_num = 1
    if len(group) > 1:
        if type(group[1]) == StringType:
            plotinfo = group[1]
        else:
            ar_num = 2
            y = group[1]
    if len(group) == 3:
        plotinfo = group[2]
    #force 1D arrays to 2D columns
    if is1D(x):
        x = col(x)
    if ar_num == 2 and is1D(y):
        y = col(y)

    xs = shape(x)
    if ar_num == 2:  ys = shape(y)
    #test that x and y have compatible shapes
    if ar_num == 2:
        #check that each array has the same number of rows
        if(xs[0] != ys[0] ):
            raise SizeMismatch, ('rows', xs, ys)
        #check that x.cols = y.cols
        #no error x has 1 column
        if(xs[1] > 1 and xs[1] != ys[1]):
            raise SizeMismatch, ('cols', xs, ys)

    #plot x against index
    if(ar_num == 1):
        for y_data in transpose(x):
            index = arange(len(y_data))
            pts = hstack(( col(index), col(y_data) ))
            pts = remove_bad_vals(pts)
            line = plot_module.line_object(pts)
            lines.append(line)
    #plot x vs y
    elif(ar_num ==2):
        #x is effectively 1D
        if(xs[1] == 1):
            for y_data in transpose(y):
                pts = hstack(( col(x), col(y_data) ))
                pts = remove_bad_vals(pts)
                line = plot_module.line_object(pts)
                lines.append(line)
        #x is 2D
        else:
            x = transpose(x); y = transpose(y)
            for i in range(len(x)):
                pts = hstack(( col(x[i]), col(y[i]) ))
                pts = remove_bad_vals(pts)
                line = plot_module.line_object(pts)
                lines.append(line)
    color,marker,line_type = process_format(plotinfo)
    #print color,marker,line_type
    for line in lines:
        if color != 'auto':
            line.color = 'custom'
            line.set_color(color)
            #print color
        if not marker:
            line.marker_type = 'custom'
            line.markers.visible = 'no'
            #print marker
        elif marker != 'auto':
            line.marker_type = 'custom'
            line.markers.symbol = marker
            line.markers.visible = 'yes'
            #print marker
        if not line_type:
            line.line_type = 'custom'
            line.line.visible = 'no'
        elif line_type != 'auto':
            line.line_type = 'custom'
            line.line.visible = 'yes'
            line.line.style = line_type
            #print line_type
        #print line.markers.visible,    line.line.visible,
    return lines

import re
color_re = re.compile('[ymcrgbwk]')
color_trans = {'y':'yellow','m':'magenta','c':'cyan','r':'red','g':'green',
               'b':'blue', 'w':'white','k':'black'}
# this one isn't quite right
marker_re = re.compile('[ox+s^v]|(?:[^-])[.]')
marker_trans = {'.':'dot','o':'circle','x':'cross','+':'plus','s':'square',
                '^':'triangle','v':'down_triangle'}

line_re = re.compile('--|-\.|[-:]')
line_trans = {'-':'solid',':':'dot','-.':'dot dash','--':'dash'}
def process_format(format):
    if format == '':
        return 'auto','auto','auto'
    color,marker,line = 'auto',None,None
    m = color_re.findall(format)
    if len(m): color = color_trans[m[0]]
    m = marker_re.findall(format)
    # the -1 takes care of 'r.', etc
    if len(m): marker = marker_trans[m[0][-1]]
    m = line_re.findall(format)
    if len(m): line = line_trans[m[0]]
    return color,marker,line

def remove_bad_vals(x):
    # !! Fix axis order when interface changed.
    # mapping:
    #    NaN -> 0
    #    Inf -> limits.double_max
    #   -Inf -> limits.double_min
    y = nan_to_num(x)
    big = limits.double_max / 10
    small = limits.double_min / 10
    y = clip(y,small,big)
    return y

def stem(*data):
    if len(data) == 1:
        n = arange(len(data[0]))
        x = data[0]
        ltype = ['b-','mo']
    if len(data) == 2:
        if type(data[1]) is types.StringType:
            ltype = [data[1],'mo']
            n = arange(len(data[0]))
            x = data[0]
        elif type(data[1]) in [types.ListType, types.TupleType]:
            n = arange(len(data[0]))
            x = data[0]
            ltype = data[1][:2]
        else:
            n = data[0]
            x = data[1]
            ltype = ['b-','mo']
    elif len(data) > 2:
        n = data[0]
        x = data[1]
        ltype = data[2]
        if type(ltype) is types.StringType:
            ltype = [ltype,'mo']
    else:
        raise ValueError, "Invalid input arguments."

    if len(n) != len(x):
        raise SizeMismatch, ('lengths', len(n), len(x))
    # line at zero:
    newdata = []
    newdata.extend([[n[0],n[-1]],[0,0],ltype[0]])

    # stems
    for k in range(len(x)):
        newdata.extend([[n[k],n[k]],[0,x[k]],ltype[0]])

    # circles
    newdata.extend([n,x,ltype[1]])
    keywds = {'fill_style': 'transparent'}
    return plot(*newdata,**keywds)



def plot(*data,**keywds):
    groups = plot_groups(data)
    lines = []
    for group in groups:
        lines.extend(lines_from_group(group))
        #default to markers being invisible
        #lines[-1].markers.visible = 'no'
    # check for hold here
    for name in plot_objects.poly_marker._attributes.keys():
        value = keywds.get(name)
        if value is not None:
            for k in range(len(lines)):
                exec('lines[k].markers.%s = value' % name)
    validate_active()
    if not _active.hold in ['on','yes']:
        _active.line_list.data = [] # clear it out
        _active.image_list.data = [] # clear it out
    for i in lines:
        _active.line_list.append(i)
    _active.update()
    return _active

def markers(visible=None):
    pass

#-------------------------------------------------------------------#
#--------------------------- image ---------------------------------#
#-------------------------------------------------------------------#

def image(img,x=None,y=None,colormap = 'grey',scale='no'):
    """Colormap should really default to the current colormap..."""
    # check for hold here
    validate_active()
    image = wxplt.image_object(img,x,y,colormap=colormap,scale=scale)
    if not _active.hold in ['on','yes']:
        _active.line_list.data = [] # clear it out
        _active.image_list.data = [] # clear it out
        _active.image_list.append(image)
        try:
            axis('equal')
        except AttributeError:
            # cluge to handle case where ticks didn't exist when
            # calling axis()
            _active.client.layout_all()
            axis('equal')
    else:
        _active.image_list.append(image)
        _active.update()
    return _active

def imagesc(img,x=None,y=None,colormap = 'grey'):
    image(img,x,y,colormap,scale='yes')

#matlab equivalence
xlabel = xtitle
ylabel = ytitle

def speed_test():
    p = plot([1,2,3],'r:o')
    s1 = (200,200)
    s2 = (400,400)
    p.SetSize(s1)
    for i in range(20):
        if p.GetSizeTuple()[0] == 200:
            p.SetSize(s2)
        else:
            p.SetSize(s1)