from matplotlib import cm
from matplotlib import pyplot
from matplotlib import rcParams
from matplotlib import ticker
from matplotlib.colorbar import ColorbarBase
from matplotlib.colors import Normalize
from matplotlib.collections import PatchCollection
from matplotlib.patches import Rectangle

from scipy.io import mmio

from fipy.tools import numerix

__all__ = ["MatplotlibSparseMatrixViewer"]

class SignedLogFormatter(ticker.LogFormatter):
    """Format signed values for log axis

    formatted values range from large positive to small positive to small negative to large negative.

    if attribute decadeOnly is True, only the decades will be labelled.
    """
    def __init__(self, base=10.0, labelOnlyBase=True, threshold=0.):
        """
        base is used to locate the decade tick,
        which will be the only one to be labeled if labelOnlyBase
        is False
        """
        self._base = base+0.0
        self.labelOnlyBase=labelOnlyBase
        self.threshold = threshold

    def __call__(self, x, pos=None):
        'Return the format for tick val x at position pos'
        vmin, vmax = self.axis.get_view_interval()
        d = abs(vmax - vmin)
        b=self._base
        # only label the decades
        sgn = numerix.sign(x)
        x = abs(x)
        x += self.threshold
        isDecade = ticker.is_decade(x)
        bx = b**x
        sgnbx = sgn * bx
        if not isDecade and self.labelOnlyBase: s = ''
        elif x > 4: s= '%+1.0e' % sgnbx
        elif x < 0: s =  '%+1.0e' % sgnbx
        else        : s =  self.pprint_val(sgnbx, d)
        return s

    def pprint_val(self, x, d):
        #if the number is not too big and it's an int, format it as an
        #int
        if abs(x)<1e4 and x==int(x): return '%+d' % x

        d = 10**(d / 2.)

        if d < 1e-2: fmt = '%+1.3e'
        elif d < 1e-1: fmt = '%+1.3f'
        elif d > 1e5: fmt = '%+1.1e'
        elif d > 10 : fmt = '%+1.1f'
        elif d > 1 : fmt = '%+1.2f'
        else: fmt = '%+1.3f'
        s =  fmt % x
        #print d, x, fmt, s
        tup = s.split('e')
        if len(tup)==2:
            mantissa = tup[0].rstrip('0').rstrip('.')
            sgn = tup[1][0].replace('+', '')
            exponent = tup[1][1:].lstrip('0')
            s = '%se%s%s' %(mantissa, sgn, exponent)
        else:
            s = s.rstrip('0').rstrip('.')

        return s

class SignedLogLocator(ticker.LogLocator):
    """Determine the tick locations for "signed" log axes

    Locate ticks from large positive to small positive to small negative to large negative
    """

    def __init__(self, base=10.0, subs=[1.0], threshold=0.):
        """
        place ticks on the location= base**i*subs[j]
        """
        ticker.LogLocator.__init__(self, base, subs)
        self.numticks = 7
        self.threshold = threshold

    def _set_numticks(self):
        self.numticks = 7  # todo; be smart here; this is just for dev

    def __call__(self):
        'Return the locations of the ticks'
        b=self._base

        vmin, vmax = self.axis.get_view_interval()

        if vmax<vmin:
            vmin, vmax = vmax, vmin

        if vmax * vmin > 0:
            raise ValueError("The interval must range from negative to positive values")

        ticklocs = []

        for limit, sgn in zip([vmax, -vmin], [1, -1]):
            numdec = numerix.floor(limit+self.threshold)-numerix.ceil(self.threshold)

            if self._subs is None: # autosub
                if numdec>10:
                    subs = numerix.array([1.0])
                elif numdec>6:
                    subs = numerix.arange(2.0, b, 2.0)
                else:
                    subs = numerix.arange(2.0, b)
                subs = numerix.log(subs) / numerix.log(b)
            else:
                subs = self._subs
                if numdec == 0 and len(subs) == 1:
                    subs = numerix.array(list(subs) + list(numerix.log(numerix.arange(2.0, b)) / numerix.log(b)))

            stride = 1
            while numdec/stride+1 > self.numticks:
                stride += 1

            for decadeStart in numerix.arange(numerix.floor(self.threshold),
                                              numerix.ceil(limit + self.threshold)+stride,
                                              stride):
                ticks = subs + decadeStart - self.threshold
                ticklocs.extend( sgn * ticks.compress(ticks > 0) )

        return numerix.array(ticklocs)

    def autoscale(self):
        'Try to choose the view limits intelligently'
        self.verify_intervals()
        vmin, vmax = self.dataInterval.get_bounds()

        if vmax<vmin:
            vmin, vmax = vmax, vmin

        if vmax * vmin > 0:
            raise ValueError("The interval must range from negative to positive values")

        if vmax == self.threshold:
            vmax += 1
        if vmin == -self.threshold:
            vmin -= 1

        exponent, remainder = divmod(numerix.log10(vmax - vmin), 1)

        if remainder < 0.5:
            exponent -= 1
        scale = 10**(-exponent)
        vmin = numerix.floor(scale*vmin)/scale
        vmax = numerix.ceil(scale*vmax)/scale

        return nonsingular(vmin, vmax)

class MatplotlibSparseMatrixViewer:
    def __init__(self, title="Sparsity"):
        self.title = title

        self.L_width = 0.8
        self.margin = (1. - self.L_width) / 2
        self.b_width = self.margin
        self.c_width = self.margin / 3
        self.buffer = 1.5 * self.margin
        self.aspect = (self.margin + self.L_width                   # M
                       + self.buffer + self.c_width                 # colorbar
                       + self.buffer + self.b_width + self.margin)  # b

        pyplot.ion()

        fig = pyplot.figure(figsize=pyplot.figaspect(1. / self.aspect))
        self.id = fig.number

    def plot(self, matrix, RHSvector, log='auto'):
        import tempfile
        import os

        if "print" in os.environ['FIPY_DISPLAY_MATRIX'].lower().split():
            print "-"*75
            print self.title
            print "-"*75
            print "L:"
            print matrix
            print "b:", RHSvector

        (f, mtxName) = tempfile.mkstemp(suffix='.mtx')
        matrix.exportMmf(mtxName)
        mtx = mmio.mmread(mtxName)
        os.remove(mtxName)

        pyplot.ion()

        c = mtx.tocoo()
        y = c.row
        x = c.col
        z = c.data
        N = matrix._shape[0]

        b = RHSvector
        if numerix.shape(b) == ():
            b = numerix.zeros((N,), 'l')

        if len(z) == 0:
            y = numerix.zeros((1,), 'l')
            x = numerix.zeros((1,), 'l')
            z = numerix.zeros((1,), 'l')

        def signed_to_logs(v):
            return (numerix.where(v > 0, numerix.log10(v), numerix.nan),
                    numerix.where(v < 0, numerix.log10(-v), numerix.nan))

        def logs_to_signed(v, plus, minus):
            v = numerix.where(v > 0, plus, -minus)
            v = numerix.where(numerix.isnan(v), 0., v)

            return v

        zPlus, zMinus = signed_to_logs(z)
        bPlus, bMinus = signed_to_logs(b)

        logs = (zPlus, zMinus, bPlus, bMinus)

        log = ((log == True)
               or (log == 'auto'
                   and (numerix.nanmax(numerix.concatenate(logs))
                        - numerix.nanmin(numerix.concatenate(logs)) > 2)))

        if log:
            zMin = numerix.nanmin(numerix.concatenate(logs))
            zMax = numerix.nanmax(numerix.concatenate(logs))

            zMin -= 0.5

            numdec = numerix.floor(zMax) - numerix.ceil(zMin)
            if numdec < 0:
                zMax += 0.5

            for v in logs:
                v -= zMin

            zRange = zMax - zMin

            if zRange == 0:
                zRange = numerix.nanmax(zPlus) + 1

            z = logs_to_signed(z, zPlus, zMinus)
            b = logs_to_signed(b, bPlus, bMinus)

            fmt = SignedLogFormatter(threshold=zMin)
            loc = SignedLogLocator(threshold=zMin)

        else:
            zRange = max(abs(numerix.concatenate((z, b))))

            if zRange == 0:
                zRange = 1

            fmt = None
            loc = None


        pyplot.ioff()

        fig = pyplot.figure(self.id)
        fig.clf()

        usetex = rcParams['text.usetex']
        rcParams['text.usetex'] = False

        cmap = cm.RdBu

        norm = Normalize(vmin=-zRange, vmax=zRange)

        x0 = self.margin
        L_ax = fig.add_axes([x0 / self.aspect, self.margin, self.L_width / self.aspect, self.L_width])
        L_ax.text(0.5, -0.1, "L",
                  transform=L_ax.transAxes, horizontalalignment='center', verticalalignment='baseline')

        x0 += self.L_width + self.buffer
        c_ax = fig.add_axes([x0 / self.aspect, self.margin, self.c_width / self.aspect, self.L_width])

        x0 += self.c_width + self.buffer
        b_ax = fig.add_axes([x0 / self.aspect, self.margin, self.b_width / self.aspect, self.L_width],
                            sharey=L_ax)
        b_ax.text(0.5, -0.1, "b",
                  transform=b_ax.transAxes, horizontalalignment='center', verticalalignment='baseline')



        def scatterRectangles(x, y, z, norm=None, cmap=None):
            patches = [Rectangle(numerix.array([X - 0.5, Y - 0.5]), 1., 1.,
                                 edgecolor='none') for X, Y in zip(x, y)]

            collection = PatchCollection(patches, norm=norm, cmap=cmap,
                                         edgecolors='none')
            collection.set_array(z)

            return collection

        L_ax.add_collection(scatterRectangles(x=x, y=y, z=z,
                                              norm=norm, cmap=cmap))

        b_ax.add_collection(scatterRectangles(x=numerix.zeros((N,), 'l'), y=numerix.arange(N), z=b,
                                              norm=norm, cmap=cmap))

        ColorbarBase(ax=c_ax, cmap=cmap, norm=norm, orientation='vertical',
                     format=fmt, ticks=loc)

        pyplot.setp((b_ax.get_xticklabels(),
                     b_ax.get_yticklabels(),
                     b_ax.get_xticklines(),
                     b_ax.get_yticklines()), visible=False)

        L_ax.set_xlim(xmin=-0.5, xmax=N-0.5)
        L_ax.set_ylim(ymax=-0.5, ymin=N-0.5)

        b_ax.set_xlim(xmin=-0.5, xmax=0.5)
        b_ax.set_ylim(ymax=-0.5, ymin=N-0.5)

        fig.suptitle(self.title, x=0.5, y=0.95, fontsize=14)

        pyplot.draw()

        rcParams['text.usetex'] = usetex
