import sys, os, shutil, imp, warnings, cStringIO, re

import IPython
from IPython.Shell import MatplotlibShell

try:
    from hashlib import md5
except ImportError:
    from md5 import md5

from docutils.parsers.rst import directives
import sphinx


sphinx_version = sphinx.__version__.split(".")
# The split is necessary for sphinx beta versions where the string is
# '6b1'
sphinx_version = tuple([int(re.split('[a-z]', x)[0])
                        for x in sphinx_version[:2]])



COMMENT, INPUT, OUTPUT =  range(3)
rgxin = re.compile('In \[(\d+)\]:\s?(.*)\s*')
rgxout = re.compile('Out\[(\d+)\]:\s?(.*)\s*')
fmtin = 'In [%d]:'
fmtout = 'Out[%d]:'

def block_parser(part):
    """
    part is a string of ipython text, comprised of at most one
    input, one ouput, comments, and blank lines.  The block parser
    parses the text into a list of::

      blocks = [ (TOKEN0, data0), (TOKEN1, data1), ...]
    where TOKEN is one of [COMMENT | INPUT | OUTPUT ] and
    data is, depending on the type of token::

      COMMENT : the comment string

      INPUT: the (DECORATOR, INPUT_LINE, REST) where
         DECORATOR: the input decorator (or None)
         INPUT_LINE: the input as string (possibly multi-line)
         REST : any stdout generated by the input line (not OUTPUT)


      OUTPUT: the output string, possibly multi-line

    """

    block = []
    lines = part.split('\n')
    #print 'PARSE', lines
    N = len(lines)
    i = 0
    decorator = None
    while 1:

        if i==N:
            # nothing left to parse -- the last line
            break

        line = lines[i]
        i += 1
        line_stripped = line.strip()
        if line_stripped.startswith('#'):
            block.append((COMMENT, line))
            continue


        if line_stripped.startswith('@'):
            # we're assuming at most one decorator -- may need to
            # rethink
            decorator = line_stripped
            continue

        # does this look like an input line?
        matchin = rgxin.match(line)
        if matchin:
            lineno, inputline = int(matchin.group(1)), matchin.group(2)

            # the ....: continuation string
            continuation = '   %s:'%''.join(['.']*(len(str(lineno))+2))
            Nc = len(continuation)
            # input lines can continue on for more than one line, if
            # we have a '\' line continuation char or a function call
            # echo line 'print'.  The input line can only be
            # terminated by the end of the block or an output line, so
            # we parse out the rest of the input line if it is
            # multiline as well as any echo text

            rest = []
            while i<N:

                # look ahead; if the next line is blank, or a comment, or
                # an output line, we're done

                nextline = lines[i]
                matchout = rgxout.match(nextline)
                #print "nextline=%s, continuation=%s, starts=%s"%(nextline, continuation, nextline.startswith(continuation))
                if matchout or nextline.startswith('#'):
                    break
                elif nextline.startswith(continuation):
                    inputline += '\n' + nextline[Nc:]
                else:
                    rest.append(nextline)
                i+= 1

            block.append((INPUT, (decorator, inputline, '\n'.join(rest))))
            continue

        # if it looks like an output line grab all the text to the end
        # of the block
        matchout = rgxout.match(line)
        if matchout:
            lineno, output = int(matchout.group(1)), matchout.group(2)
            if i<N-1:
                output = '\n'.join([output] + lines[i:])

            #print 'OUTPUT', output
            block.append((OUTPUT, output))
            break

    #print 'returning block', block
    return block



import matplotlib
matplotlib.use('Agg')


class EmbeddedSphinxShell:

    def __init__(self):

        self.cout = cStringIO.StringIO()

        IPython.Shell.Term.cout = self.cout
        IPython.Shell.Term.cerr = self.cout
        argv = ['-autocall', '0']
        self.user_ns = {}
        self.user_glocal_ns = {}

        self.IP = IPython.ipmaker.make_IPython(
            argv, self.user_ns, self.user_glocal_ns, embedded=True,
            #shell_class=IPython.Shell.InteractiveShell,
            shell_class=MatplotlibShell,
            rc_override=dict(colors = 'NoColor'))

        self.input = ''
        self.output = ''

        self.is_verbatim = False
        self.is_doctest = False
        self.is_suppress = False

        # on the first call to the savefig decorator, we'll import
        # pyplot as plt so we can make a call to the plt.gcf().savefig
        self._pyplot_imported = False

        # we need bookmark the current dir first so we can save
        # relative to it
        self.process_input('bookmark ipy_basedir')
        self.cout.seek(0)
        self.cout.truncate(0)

    def process_input(self, line):
        'process the input, capturing stdout'
        #print "input='%s'"%self.input
        stdout = sys.stdout
        sys.stdout = self.cout
        #self.IP.resetbuffer()
        self.IP.push(self.IP.prefilter(line, 0))
        #self.IP.runlines(line)
        sys.stdout = stdout


    def process_block(self, block):
        """
        process block from the block_parser and return a list of processed lines
        """

        #print 'BLOCK', block
        ret = []

        output = None
        input_lines = None

        m = rgxin.match(str(self.IP.outputcache.prompt1).strip())
        lineno = int(m.group(1))

        input_prompt = fmtin%lineno
        output_prompt = fmtout%lineno
        image_file = None
        image_directive = None
        for token, data in block:

            if token==COMMENT:
                if not self.is_suppress:
                    ret.append(data)

            elif token==INPUT:

                decorator, input, rest = data
                #print 'INPUT:', data
                is_verbatim = decorator=='@verbatim' or self.is_verbatim
                is_doctest = decorator=='@doctest' or self.is_doctest
                is_suppress = decorator=='@suppress' or self.is_suppress
                is_savefig = decorator is not None and decorator.startswith('@savefig')
                #print 'is_verbatim=%s, is_doctest=%s, is_suppress=%s, is_savefig=%s'%(is_verbatim, is_doctest, is_suppress, is_savefig)
                input_lines = input.split('\n')


                continuation = '   %s:'%''.join(['.']*(len(str(lineno))+2))
                Nc = len(continuation)

                if is_savefig:
                    saveargs = decorator.split(' ')
                    filename = saveargs[1]
                    outfile = os.path.join('_static/%s'%filename)
                    # build out an image directive like
                    # .. image:: somefile.png
                    #    :width 4in
                    #
                    # from an input like
                    # savefig somefile.png width=4in
                    imagerows = ['.. image:: %s'%outfile]

                    for kwarg in saveargs[2:]:
                        arg, val = kwarg.split('=')
                        arg = arg.strip()
                        val = val.strip()
                        imagerows.append('   :%s: %s'%(arg, val))


                    image_file = outfile
                    image_directive = '\n'.join(imagerows)



                # TODO: can we get "rest" from ipython
                #self.process_input('\n'.join(input_lines))


                is_semicolon = False
                for i, line in enumerate(input_lines):
                    if line.endswith(';'):
                        is_semicolon = True

                    if i==0:
                        # process the first input line
                        if is_verbatim:
                            self.process_input('')
                        else:
                            # only submit the line in non-verbatim mode
                            self.process_input(line)
                        formatted_line = '%s %s'%(input_prompt, line)
                    else:
                        # process a continuation line
                        if not is_verbatim:
                            self.process_input(line)

                        formatted_line = '%s %s'%(continuation, line)


                    if not is_suppress:
                        ret.append(formatted_line)

                if not is_suppress:
                    if len(rest.strip()):
                        if is_verbatim:
                            # the "rest" is the standard output of the
                            # input, which needs to be added in
                            # verbatim mode
                            ret.append("%s"%rest)
                            ret.append('')

                self.cout.seek(0)
                output = self.cout.read()
                if not is_suppress and not is_semicolon and not is_verbatim:
                    ret.append(output)

                self.cout.truncate(0)




            elif token==OUTPUT:
                #print 'token==OUTPUT is_verbatim=%s'%is_verbatim
                if is_verbatim:
                    # construct a mock output prompt
                    output = '%s %s\n'%(fmtout%lineno, data)
                    ret.append(output)

                #print 'token==OUTPUT', output
                if is_doctest:
                    submitted = data.strip()
                    found = output
                    if found is not None:
                        ind = found.find(output_prompt)
                        if ind<0:
                            raise RuntimeError('output prompt="%s" does not match out line=%s'%(output_prompt, found))
                        found = found[len(output_prompt):].strip()

                        if found!=submitted:
                            raise RuntimeError('doctest failure for input_lines="%s" with found_output="%s" and submitted output="%s"'%(input_lines, found, submitted))
                        #print 'doctest PASSED for input_lines="%s" with found_output="%s" and submitted output="%s"'%(input_lines, found, submitted)


        if image_file is not None:
            self.insure_pyplot()
            command = 'plt.gcf().savefig("%s")'%image_file
            #print 'SAVEFIG', command
            self.process_input('bookmark ipy_thisdir')
            self.process_input('cd -b ipy_basedir')
            self.process_input(command)
            self.process_input('cd -b ipy_thisdir')
            self.cout.seek(0)
            self.cout.truncate(0)

        #print 'returning', ret, figure
        return ret, image_directive


    def insure_pyplot(self):
        if self._pyplot_imported:
            return
        self.process_input('import matplotlib.pyplot as plt')



shell = EmbeddedSphinxShell()


def ipython_directive(name, arguments, options, content, lineno,
                      content_offset, block_text, state, state_machine,
                      ):

    debug = ipython_directive.DEBUG
    shell.is_suppress = options.has_key('suppress')
    shell.is_doctest = options.has_key('doctest')
    shell.is_verbatim = options.has_key('verbatim')

    #print 'ipy', shell.is_suppress, options
    parts = '\n'.join(content).split('\n\n')
    lines = ['.. sourcecode:: ipython', '']

    figures = []
    for part in parts:
        block = block_parser(part)

        if len(block):
            rows, figure = shell.process_block(block)
            for row in rows:
                lines.extend(['    %s'%line for line in row.split('\n')])

            if figure is not None:
                figures.append(figure)

    for figure in figures:
        lines.append('')
        lines.extend(figure.split('\n'))
        lines.append('')

    #print lines
    if len(lines)>2:
        if debug:
            print '\n'.join(lines)
        else:
            #print 'INSERTING %d lines'%len(lines)
            state_machine.insert_input(
                lines, state_machine.input_lines.source(0))

    return []

ipython_directive.DEBUG = False

def setup(app):
    setup.app = app
    options = {
        'suppress': directives.flag,
        'doctest': directives.flag,
        'verbatim': directives.flag,
        }


    app.add_directive('ipython', ipython_directive, True, (0, 2, 0), **options)


def test():

    examples = [
        r"""
In [9]: pwd
Out[9]: '/home/jdhunter/py4science/book'

In [10]: cd bookdata/
/home/jdhunter/py4science/book/bookdata

In [2]: from pylab import *

In [2]: ion()

In [3]: im = imread('stinkbug.png')

@savefig mystinkbug.png width=4in
In [4]: imshow(im)
Out[4]: <matplotlib.image.AxesImage object at 0x39ea850>

""",
        r"""

In [1]: x = 'hello world'

# string methods can be
# used to alter the string
@doctest
In [2]: x.upper()
Out[2]: 'HELLO WORLD'

@verbatim
In [3]: x.st<TAB>
x.startswith  x.strip
""",
    r"""

In [130]: url = 'http://ichart.finance.yahoo.com/table.csv?s=CROX\
   .....: &d=9&e=22&f=2009&g=d&a=1&br=8&c=2006&ignore=.csv'

In [131]: print url.split('&')
--------> print(url.split('&'))
['http://ichart.finance.yahoo.com/table.csv?s=CROX', 'd=9', 'e=22', 'f=2009', 'g=d', 'a=1', 'b=8', 'c=2006', 'ignore=.csv']

In [60]: import urllib

""",
    r"""\

In [133]: import numpy.random

@suppress
In [134]: numpy.random.seed(2358)

@doctest
In [135]: np.random.rand(10,2)
Out[135]:
array([[ 0.64524308,  0.59943846],
       [ 0.47102322,  0.8715456 ],
       [ 0.29370834,  0.74776844],
       [ 0.99539577,  0.1313423 ],
       [ 0.16250302,  0.21103583],
       [ 0.81626524,  0.1312433 ],
       [ 0.67338089,  0.72302393],
       [ 0.7566368 ,  0.07033696],
       [ 0.22591016,  0.77731835],
       [ 0.0072729 ,  0.34273127]])

""",

    r"""
In [106]: print x
--------> print(x)
jdh

In [109]: for i in range(10):
   .....:     print i
   .....:
   .....:
0
1
2
3
4
5
6
7
8
9


""",

        r"""

In [144]: from pylab import *

In [145]: ion()

# use a semicolon to suppress the output
@savefig test_hist.png width=4in
In [151]: hist(np.random.randn(10000), 100);


@savefig test_plot.png width=4in
In [151]: plot(np.random.randn(10000), 'o');
   """,

        r"""
# use a semicolon to suppress the output
In [151]: plt.clf()

@savefig plot_simple.png width=4in
In [151]: plot([1,2,3])

@savefig hist_simple.png width=4in
In [151]: hist(np.random.randn(10000), 100);

""",
     r"""
# update the current fig
In [151]: ylabel('number')

In [152]: title('normal distribution')


@savefig hist_with_text.png
In [153]: grid(True)

        """,


        r"""

In [239]: 1/2
@verbatim
Out[239]: 0

In [240]: 1.0/2.0
Out[240]: 0.5
""",

        r"""
@verbatim
In [6]: pwd
Out[6]: '/home/jdhunter/mypy'
""",

        r"""
@verbatim
In [151]: myfile.upper?
Type:           builtin_function_or_method
Base Class:     <type 'builtin_function_or_method'>
String Form:    <built-in method upper of str object at 0x980e2f0>
Namespace:      Interactive
Docstring:
    S.upper() -> string
    Return a copy of the string S converted to uppercase.
 """
    ]



    ipython_directive.DEBUG = True
    #options = dict(suppress=True)
    options = dict()
    for example in examples:
        content = example.split('\n')
        ipython_directive('debug', arguments=None, options=options,
                          content=content, lineno=0,
                          content_offset=None, block_text=None,
                          state=None, state_machine=None,
                          )


if __name__=='__main__':
    test()