#! /usr/bin/env python
"""
Define a system of measurement units to make it easy
to work with centimeters, grams, ergs, et cetera.

Version of 2010-03-02.

Example of use, employing defined-in units:
>>> import units
>>> m = units.meter
>>> s = units.second
>>> g = 9.8 * m / s**2
>>> t = 3 * s
>>> distance = 0.5 * g * t**2
>>> print distance
44.1 meter

Example of use, defining one's own units:
>>> import units
>>> dollars = units.Unit( 1, dollar=1 )
>>> widgets = units.Unit( 1, widget=1 )
>>> pallets = units.Unit( 1, pallet=1 )
>>> unit_price = 3.14 * dollars / widgets
>>> pallet_price = unit_price * 24 * widgets / pallets
>>> print pallet_price
75.36 dollar /pallet


Revision history:
       2006-02-20: initial version.
       2008-11-13: Cleaned up __str__, fixing a bug.
       2009-01-01: Massive overhaul: base on a dictionary.
       2009-01-04: Fix __radd__ bug; add exception.
       2010-03-02: Minor cleanup.  More thought into null or zero units.
"""

import exceptions
import math


class MismatchedUnits( exceptions.Exception ):
  pass


def sqrt( u ):
  """Return the square root of a unit-bearing number."""

  a = {}
  for k,v in u.units.iteritems():
    if v % 2 != 0:
      raise RuntimeError( "sqrt only allowed if units have even powers." )
    a[k] = v/2

  return Unit( math.sqrt( u.value ), **a )


def number_or_unit( value, **units ):
  """If there are not units, or their exponents are all zero,
  return just the value; otherwise, build a Unit and return it.
  """

  if all( v == 0 for _,v in units.iteritems() ):
    return value
  return Unit( value, **units )


def _raise_exception_if_mismatched( a, b ):
  """Raise an exception if b is not a Unit, or if the units
  of a and b differ.
  """
  
  if not isinstance( b, Unit ):
    raise MismatchedUnits, "combining with object not of Unit class"
  if a.units != b.units:
    raise MismatchedUnits, "combining with object with different units"


class Unit( object ):
  """A number, with attached units."""
  
  def __init__( self, value, **args ):
    """Example: g = Unit( 9.8, meter=1, second=-2 )

    The names of the units are whatever you want (except Python
    reserved words), and the integers are the exponents.

    You are not prevented from creating objects with no exponents or
    with exponents of zero, but arithmetic operations on such an object
    are likely to return objects of the underlying value's type.
    """
    self.units = {}
    for k,v in args.iteritems():
      if v != 0:
        self.units[k] = v
    self.value = value

  def __add__( self, other ):
    _raise_exception_if_mismatched( self, other )
    return Unit( self.value + other.value, **self.units )

  def __radd__( self, other ):
    return self + other

  def __sub__( self, other ):
    _raise_exception_if_mismatched( self, other )
    return Unit( self.value - other.value, **self.units )

  def __rsub__( self, other ):
    _raise_exception_if_mismatched( self, other )
    return Unit( other.value - self.value, **self.units )


  def __neg__( self ):
    return Unit( -self.value, **self.units )


  def __mul__( self, other ):
    if isinstance( other, Unit ):
      a = self.units.copy()
      for k,v in other.units.iteritems():
        a[k] = v + a.get( k, 0 )
      return number_or_unit( self.value * other.value, **a )
    else:
      return number_or_unit( self.value * other, **self.units )


  def __rmul__( self, other ):
    return self * other


  def __div__( self, other ):
    """Return self/other."""

    if isinstance( other, Unit ):
      a = self.units.copy()
      for k,v in other.units.iteritems():
        a[k] = - v + a.get( k, 0 )
      return number_or_unit( self.value / other.value, **a )
    else:
      return number_or_unit( self.value / other, **self.units )


  def __rdiv__( self, other ):
    """Return other/self."""
    if isinstance( other, Unit ):
      return __div__( other, self )
    else:
      a = {}
      for k,v in self.units.iteritems():
        a[k] = -v
      return number_or_unit( other / self.value, **a )


  def __str__( self ):
    numerator = ""
    denominator = ""
    for k,v in self.units.iteritems():
      if v > 0:
        numerator += " %s" % k
        if v > 1: numerator += "^%d" % v

      elif v < 0:
        denominator += " /%s" % k
        if v < -1: denominator += "^%d" % -v

    return str( self.value ) + numerator + denominator


  def __eq__( self, other ):
    _raise_exception_if_mismatched( self, other )
    return self.value == other.value


  def __pow__( self, power ):
    """Returns self**power"""
    assert isinstance(power, int)
    a = {}
    for k,v in self.units.iteritems():
      a[k] = v * power
    return Unit( self.value ** power, **a )


  def has_units( self, other ):
    """Tests if two things have the same units."""

    if isinstance( other, Unit ):
      return self.units == other.units
    else:
      return all( v == 0 for _,v in self.units.iteritems() )

    

    


# Define familiar units of mass:
                   
kilogram =       Unit( 1, kilogram=1 )
gram = kilogram / 1000.



# Define familiar units of length:

meter = Unit( 1, meter=1 )
centimeter = meter / 100.
inch = 2.54 * centimeter
foot = 12 * inch
mile = 5280 * foot



# Define familiar units of time:

second =     Unit( 1, second=1 )
minute = 60 * second
hour = 60 * minute
day = 86400 * second
week = 7 * day
year = 365.24 * day



# Define familiar units of charge:

coulomb = Unit( 1, coulomb=1 )



# Define familiar units of temperature:

kelvin = Unit( 1, kelvin=1 )
fahrenheit = kelvin * 9. / 5



# Define familiar units of force:

dyne = gram*centimeter/second**2
newton = kilogram*meter/second**2



# Define familiar units of energy:

erg = centimeter * dyne
joule = newton * meter



# Define familiar units of power:

watt = joule/second


if __name__ == "__main__":

  assert isinstance( number_or_unit( 1, meter=1 ), Unit )
  assert isinstance( number_or_unit( 1, meter=0 ), int )

  g = Unit( 9.5, meter=1, second=-2 )
  assert isinstance( g, Unit )
  assert g.units[ "meter"  ] == 1
  assert g.units[ "second" ] == -2
  assert len( g.units ) == 2
  assert g.value == 9.5

  halfg = g / 2.
  assert halfg.units[ "meter"  ] == 1
  assert halfg.units[ "second" ] == -2
  assert len( halfg.units ) == 2
  assert halfg.value == 4.75
  
  threehalvesg = g + halfg
  assert threehalvesg == Unit( 9.5 + 4.75, meter=1, second=-2 )
  assert threehalvesg - halfg == g
  assert halfg - threehalvesg == -g
  assert threehalvesg == 1.5 * g
  assert threehalvesg == g * 1.5
  assert threehalvesg / g == 1.5
  assert threehalvesg / 1.5 == g
  assert 1 / g == Unit( 1/9.5, meter=-1, second=2 )

  g2 = g**2
  assert g2 == Unit( 9.5**2, meter=2, second=-4 )
  assert sqrt( g2 ) == g

  assert g.has_units( halfg )
  assert not g.has_units( g2 )
  assert not g.has_units( 1 )
  assert Unit( 1 ).has_units( 1 )
  assert Unit( 1, meter=0 ).has_units( 1 )

  t = 3 * second
  v = g * t
  d = 0.5 * g * t**2
  assert str( g ) == "9.5 meter /second^2"
  assert str( t ) == "3 second"
  assert str( v ) == "28.5 meter /second"
  assert str( d ) == "42.75 meter"


  # Confirm that when units cancel, a pure number results:
  
  miles = d / mile
  assert isinstance( miles, float )

  assert g / ( meter / second**2 ) == 9.5

  assert Unit( 2, meter=1 ) * Unit( 3, meter=-1 ) == 6
  assert Unit( 2., meter=1 ) / Unit( 4, meter=1 ) == 0.5



  # Ensure that __rmul__ isn't an infinite loop:

  assert 1 * g == g * 1


  # Demonstrate another use of has_units:

  assert g.has_units( inch / day**2 )


  # Verify that mismatched units cause exceptions:

  c = 3e8 * meter / second

  for abomination in ( "c == g",
                       "c == 0",
                       "0 == c",
                       "c + g",
                       "1 + g",
                       "g + 1",
                       "c - g",
                       "1 - g",
                       "g - 1",
                       "0 == g - g" ):
    try:
      t = eval( abomination )
    except MismatchedUnits:
      pass
    else:
      print "*** %s failed to raise an exception! ***" % abomination
      
  

  print "Done with tests."