import time, types, itertools, os.path
from scipy.interpolate import interp1d
from scipy.integrate import odeint
from scipy.optimize import fmin
from collections import OrderedDict

class DiffEq(object):
    """designed to make numerically solving diff eqs more readable
    
    setattr and getattr have been modified to provide the following effects:
        >>> self._foo=2.#setting an _under var makes it still returned by...
        >>> self.foo
        2.0
        >>> self._foo=2 #but only if it's a float
        >>> self.foo
        AttributeError: self has no attribute 'foo'
        >>> # Now suppose _foo is a class method and bar is in _flags.
        >>> self._foo = lambda self: self.bar
        >>> self.flags.append('bar')
        >>> self.bar = 3.14
        >>> self.foo    #finding the object having not found it before:
                        #self._foo() is called and stored to self.__foo
        3.14            #self.__foo is returned
        
        >>> self.foo    #finding the object having asked for it before:
        3.14            #self.__foo is returned directly
        
        >>> self.bar=2  #setting bar resets all self.__foos
        
        >>> self.foo    #self._foo() is called again, stored to self.__foo
        2               #the new self.__foo is returned
    
    Note that with this behavior you should treat all variable names, be they
    normal, _under, or __dunder, as in the same namespace. Don't even try to
    differentiate variable names by underscores.
    
    The benefit of this system is that every normal float will be optimized,
    while every _under float will be ignored completely, and flagging a var
    with _under will not change the behavior of any functions using it.
    Additionally, if you suddenly realize that foo should be dependent on
    bar and baz, you simply define:
        def _foo(self):
            return do_stuff(self.bar, self.baz)
    and bar and baz will be optimized while foo will be ignored -- and no
    function code using foo needs to be rewritten.
    """
    def __init__(self):
        self._time = time.time()
        
        # All floats lacking an initial _underscore will be optimized.
        self.opt_vars = OrderedDict()
        for key, value in vars(self).items():
            if key[0] == "_": continue
            if type(value) != float: continue
            self.opt_vars[key] = value
    def gettime(self): return time.time() - self.time
    
    def __getattribute__(self, attr):
        """Provides a few getattr psuedonyms:
        
        self.foo tries to return:
            self.foo    if foo exists
            self._foo   if _foo is a float
            self._foo() if _foo is a method
        see DiffEq's documentation for examples
        """
        sgetattr = super(DiffEq, self).__getattribute__
        try: return sgetattr(attr)
        except AttributeError as err:
            if attr[0] == "_": raise err
            value = sgetattr("_" + attr)
            if type(value) == types.MethodType:
                try: return sgetattr("__" + attr)
                except AttributeError:
                    value = value()
                    setattr(self, "__" + attr, value)
                    return value
            if type(value) == float: return value
            raise err
    
    _flags = []
    def __setattribute__(self, attr):
        """Clears dunder updates when a flag variable is changed."""
        sgetattr = super(DiffEq, self).__getattribute__
        ssetattr = super(DiffEq, self).__setattribute__
        if attr in self.flags:
            for key, value in vars(self).items():
                if key[:2] != "__": continue
                if type(sgetattr(key[2:])) != types.MethodType: continue
                vars(self).remove(key)
        ssetattr(attr)
    
    def solve_diff_eq(self):
        """wrapper for odeint and self.diff_eq
        
        requires the following:
            self.diff_eq
            self.init_vars
            self.timespace
        """
        time0 = self.gettime()
        self.time = numpy.linspace(*self.timespace)
        self.soln = odeint(self.diff_eq, self.init_vars, self.time)
        print "  ::  took %s seconds" % (self.gettime() - time0)
    
    def opt_init(self, guesses=None):
        if guesses == None:
            for var in self.opt_vars:
                self.opt_vars[var] = getattr(self, var)
        else:
            for var, guess in itertools.izip(self.opt_vars, guesses):
                setattr(self, var, guess)
    
    #_outfile = tempfile.mktemp(dir='data', prefix=str(self.gettime()))
    outfile = lambda: os.path.join('data', str(self.gettime()))
    def opt_fun(self, guesses, f, func, init_vars):
        """Optimize the free variables for the diff eq."""
        self.opt_init()
        map(setattr, itertools.repeat(self), self.opt_vars.iterkeys(), guesses)
        self.__dict__.update(dict(zip(self.opt_vars.iterkeys(), guesses)))
        outfile = self.__dict__.get('outfile')
        
        time, soln = self.solve_diff_eq()
        
        phi, phi1, v_tan, tension, moment, momentum, max_arg = \
            interpret(soln)
        phi, phi1, v_tan, tension, moment, momentum, time = \
            map(lambda x: numpy.ndarray.__getslice__(x, 0, max_arg+1),
                [phi, phi1, v_tan, tension, moment, momentum, time])
        
        time, (phi, phi1, v_tan, tension, moment, momentum, max_arg) = \
            solver(func, init_vars, returns=['time', 'slices'])
        
        # Parameter values.
        f.write("%s, " % r_roll)
        f.write("%s, " % length)
        f.write("%s, " % r0)
        f.write("%s, " % m)
        f.write("%s, " % Mass)
        
        # Values at the max velocity point.
        f.write("%s, " %     phi[-1])
        f.write("%s, " %    phi1[-1])
        f.write("%s, " % tension[-1])
        f.write("%s\n" %   v_tan[-1])
        
        time, soln = self.solve_func(
        
        self.interpret()
        
        return self.optimize_value()