from __future__ import division, print_function, unicode_literals
import decimal
import math
import numpy as np
from matplotlib.ticker import Formatter


# Proposed "enhancement" of the EngFormatter
class EnhancedEngFormatter(Formatter):
    """
    Formats axis values using engineering prefixes to represent powers of 1000,
    plus a specified unit, e.g., 10 MHz instead of 1e7.
    """

    # the unicode for -6 is the greek letter mu
    # commeted here due to bug in pep8
    # (https://github.com/jcrocholl/pep8/issues/271)

    # The SI engineering prefixes
    ENG_PREFIXES = {
        -24: "y",
        -21: "z",
        -18: "a",
        -15: "f",
        -12: "p",
         -9: "n",
         -6: "\u03bc",
         -3: "m",
          0: "",
          3: "k",
          6: "M",
          9: "G",
         12: "T",
         15: "P",
         18: "E",
         21: "Z",
         24: "Y"
    }

    def __init__(self, unit="", places=None, space_sep=True):
        """ Parameters
            ----------
            unit: str (default: "")
                Unit symbol to use.

            places: int (default: None)
                Number of digits after the decimal point.
                If it is None, falls back to the floating point format '%g'.

            space_sep: boolean (default: True)
                If True, a (single) space is used between the value and the
                prefix/unit, else the prefix/unit is directly appended to the
                value.
        """
        self.unit = unit
        self.places = places
        # <ENH>
        if space_sep is True:
            self.sep = " "  # 1 space
        else:
            self.sep = ""  # no space
        # <\ENH>

    def __call__(self, x, pos=None):
        s = "%s%s" % (self.format_eng(x), self.unit)
        return self.fix_minus(s)

    def format_eng(self, num):
        """ Formats a number in engineering notation, appending a letter
        representing the power of 1000 of the original number. Some examples:

        >>> format_eng(0)       # for self.places = 0
        '0'

        >>> format_eng(1000000) # for self.places = 1
        '1.0 M'

        >>> format_eng("-1e-6") # for self.places = 2
        u'-1.00 \u03bc'

        @param num: the value to represent
        @type num: either a numeric value or a string that can be converted to
                   a numeric value (as per decimal.Decimal constructor)

        @return: engineering formatted string
        """

        dnum = decimal.Decimal(str(num))

        sign = 1

        if dnum < 0:
            sign = -1
            dnum = -dnum

        if dnum != 0:
            pow10 = decimal.Decimal(int(math.floor(dnum.log10() / 3) * 3))
        else:
            pow10 = decimal.Decimal(0)

        pow10 = pow10.min(max(self.ENG_PREFIXES.keys()))
        pow10 = pow10.max(min(self.ENG_PREFIXES.keys()))

        prefix = self.ENG_PREFIXES[int(pow10)]

        mant = sign * dnum / (10 ** pow10)

        # <ENH>
        if self.places is None:
            format_str = "%g{sep:s}%s".format(sep=self.sep)
        elif self.places == 0:
            format_str = "%i{sep:s}%s".format(sep=self.sep)
        elif self.places > 0:
            format_str = "%.{p:i}f{sep:s}%s".format(p=self.places,
                                                    sep=self.sep)
        # <\ENH>

        formatted = format_str % (mant, prefix)

        formatted = formatted.strip()
        if (self.unit != "") and (prefix == self.ENG_PREFIXES[0]):
            # <ENH>
            formatted = formatted + self.sep
            # <\ENH>

        return formatted


# DEMO
def demo_formatter(**kwargs):
    """ Print the strings produced by the EnhancedEngFormatter for a list of
        arbitrary test values.
    """
    TEST_VALUES = [1.23456789e-6, 0.1, 1, 999.9, 1001]
    unit = kwargs.get('unit', "")
    space_sep = kwargs.get('space_sep', True)
    formatter = EnhancedEngFormatter(**kwargs)

    print("\n[Case: unit='{u:s}' & space_sep={s}]".format(u=unit, s=space_sep))
    print(*["{tst};".format(tst=formatter(value)) for value in TEST_VALUES])

if __name__ == '__main__':
    """ Matplotlib current behavior (w/ space separator) """
    demo_formatter(unit="s", space_sep=True)
    # >> 1.23457 μs; 100 ms; 1 s; 999.9 s; 1.001 ks;
    demo_formatter(unit="", space_sep=True)
    # >> 1.23457 μ; 100 m; 1; 999.9; 1.001 k;

    """ New possibility (w/o space separator) """
    demo_formatter(unit="s", space_sep=False)
    # >> 1.23457μs; 100ms; 1s; 999.9s; 1.001ks;
    demo_formatter(unit="", space_sep=False)
    # >> 1.23457μ; 100m; 1; 999.9; 1.001k;