#!/usr/bin/env python

# Jacob Joseph
# 4 June 2012

# Custom scales for matplotlib.  In particular, an exponential scale.
# Mostly from: http://stackoverflow.com/questions/4899516/how-to-produce-an-exponentially-scaled-axis

# Use:
# mscale.register_scale(CustomScale)
# pyplot.xscale('exp10')

from numpy import log10, log2
import matplotlib.pyplot as plt

from matplotlib import scale as mscale
from matplotlib import transforms as mtransforms

class ExpScale(mscale.ScaleBase):
    name = 'exp10'

    def __init__(self, axis, **kwargs):
        mscale.ScaleBase.__init__(self)
        self.thresh = None #thresh

    def get_transform(self):
        return self.CustomTransform(self.thresh)

    def set_default_locators_and_formatters(self, axis):
        pass

    class CustomTransform(mtransforms.Transform):
        input_dims = 1
        output_dims = 1
        is_separable = True

        def __init__(self, thresh):
            mtransforms.Transform.__init__(self)
            self.thresh = thresh

        def transform(self, a):
            print a, a/44000, 10**(a/44000/10)
            return 10**(a/44000/10)

        def inverted(self):
            return CustomScale.InvertedCustomTransform(self.thresh)

    class InvertedCustomTransform(mtransforms.Transform):
        input_dims = 1
        output_dims = 1
        is_separable = True

        def __init__(self, thresh):
            mtransforms.Transform.__init__(self)
            self.thresh = thresh

        def transform(self, a):
            print "inv", a, log2(a)*(10/log2(10)), 10**(a/10)
            return log2(a)*(10/log2(10)), 10**(a/10)

        def inverted(self):
            return CustomScale.CustomTransform(self.thresh)