import numpy
import error
import math

#System specifications:

#System classes are responsible for defining the differential operator, and 
#rendering the initial state to seed solvers

#System constructors convert data from human read/writeable form, to an 
#internal representation

#System state consists of two parts: the time component, which is a float,
#and the w vector, which has some shape, subject only to the restriction that
#the differential operator returns an array with the same shape, when given a
#w vector and a time

def optional_w(f):
    def f_opt(self, w=False):
        if w == False:
            w = self.initial[0]['w']
        return f(self, w)
    f_opt.__name__ = f.__name__
    f_opt.__doc__ = f.__doc__
    return f_opt

def optional_state(f):
    def f_opt(self, state=False):
        if state == False:
            state = self.initial[0]
        return f(self, state)
    f_opt.__name__ = f.__name__
    f_opt.__doc__ = f.__doc__
    return f_opt

def statify(f):
    def f_new(state=False):
        if state == False:
            return f()
        return f(state['w'])
    f_new.__name__ = f.__name__
    f_new.__doc__ = f.__doc__
    return f_new

class System(object):
    
    "Dummy implementation of the System class."
    
    w_type = NotImplemented
    initial = NotImplemented
    
    @optional_state
    def differential(self, state):
        raise NotImplementedError
    
    def flat_differential(self, t=None, w=None):
        if w is None or t is None:
            return self.differential().view(dtype=numpy.float64)
        else:
            return self.differential(
                {'t': t, 'w': numpy.rec.array(w, 
                    dtype=self.initial[0]['w'].dtype)}).view(
                        dtype=numpy.float64)

class Polynomial(System):
    def __init__(self, coefficients, t=0):
        self.coefficients = coefficients
        #self.initial = 
    
    def exact(self, t):
        return sum([coefficients[i] * t**i for i in range(len(coefficients))])
    
    @optional_state
    def differential(self, state):
        t = state['t']
        return sum(i * [coefficients[i] * t**(i-1) for i in range(len(coefficients))])

class Exponential(System):
    def __init__(self, initial=1, lambda_=-1, t=0):
        self.lambda_ = lambda_
        self.w_type = numpy.float
        self.array_type = numpy.dtype([('t', numpy.float, 1), 
                                       ('w', self.w_type, 1)])
        self.initial = numpy.empty(1, self.array_type)
        self.initial[0]['t'] = t
        self.initial[0]['w'] = initial
    
    def exact(self, t):
        return self.initial[0]['w'] * math.exp(self.lambda_ * (t - self.initial[0]['t']))
    
    @optional_state
    def differential(self, state):
        return state['w'] * self.lambda_

class Orbital(System):
    
    def __init__(self, *bodies, **kwargs):
        t = 0
        if 't' in kwargs:
            t = kwargs['t']
        self.G = 1
        if 'G' in kwargs:
            self.G = kwargs['G']
        self.size = len(bodies)
        if self.size == 0: #Degenerate case in which nothing happens
            return
        if self.size == 1: #Degenerate case in which all methods have no error
            return
        self.masses = numpy.asarray([[body['mass']] for body in bodies])
        self.mass = sum(self.masses)
        self.potential_constants = -self.G*numpy.ma.masked_array(self.masses.transpose()*self.masses, numpy.eye(self.size))
        self.D = len(bodies[0]['pos'])
        self.zero = numpy.zeros(self.D)
        if self.D != 2 and self.D != 3:
            return
        self.point_type = numpy.dtype([('pos', numpy.float, self.D), 
                                   ('vel', numpy.float, self.D)])
        self.w_type = numpy.dtype((self.point_type, self.size))
        self.array_type = numpy.dtype([('t', numpy.float, 1), 
                                       ('w', self.w_type, 1)])
        self.initial = numpy.empty(1, self.array_type)
        self.initial[0]['t'] = t
        for i in range(self.size):
            for j in range(self.D):
                self.initial[0]['w'][i]['pos'][j] = bodies[i]['pos'][j]
                self.initial[0]['w'][i]['vel'][j] = bodies[i]['vel'][j]
        
        self.errors = {
            "barycenter": error.Projection(statify(self.barycenter), 
                                           self.projected_barycenter,
                                           self.momentum()/sum(self.masses),
                                           t),
            "energy": error.Invariant(statify(self.mechanical_energy)),
            "momentum": error.InvariantCoordinate(statify(self.momentum)),
            "angular momentum": (
                self.D==3 and 
                error.InvariantCoordinate or 
                error.Invariant)(statify(self.angular_momentum))}
        
        if (self.momentum() == self.zero).all():
            print "WARNING: The system's momentum is zero. This is physically unlikely, and leads to degenerate behavior."
    
    @optional_state
    def pos(self, state):
        return state['w']['pos']
    
    @optional_state
    def vel(self, state):
        return state['w']['vel']
    
    @optional_state
    def tee(self, state):
        return state['t']
    
    @optional_w
    def momenta(self, w):
        return self.masses * w['vel']
    
    def momentum(self, w=False):
        return sum(self.momenta(w))
    
    #@optional_w
    def mechanical_energy(self, w=False):
        return self.kinetic_energy(w) + self.potential_energy(w)
    
    @optional_w
    def kinetic_energy(self, w):
        return 0.5 * sum(sum(self.masses * (w['vel'] * w['vel'])))
    
    @optional_w
    def barycenter(self, w):
        return sum(self.masses * w['pos']) / self.mass
    
    def projected_barycenter(self, t=None):
        diff = 0
        if t is not None:
            diff = t - self.initial[0]['t']
        return self.barycenter() + diff * self.momentum() / self.mass
    
    @optional_w
    def angular_momentum(self, w):
        return sum(numpy.cross(w['pos']-numpy.repeat([self.barycenter(w)], 
                                                     self.size, 0), 
                               self.momenta(w)))
        
    @optional_w
    def displacements(self, w):
        pos = numpy.repeat([w['pos']], self.size, 0)
        #print pos.shape
        return pos - pos.transpose(1, 0, 2)
    
    #@optional_w
    def distances_squared(self, w=False):
        return (self.displacements(w)**2).sum(2)
    
    #@optional_w
    def distances(self, w=False):
        return self.distances_squared(w)**0.5
        
    #@optional_w
    def potentials(self, w=False):
        return self.potential_constants / self.distances(w)
    
    #@optional_w
    def potential_energy(self, w=False):
        return 0.5*self.potentials(w).filled(0).sum()
        
    #@optional_w
    def forces(self, w=False):
        return ((numpy.repeat(self.potentials(w) / self.distances_squared(w), 
                              self.D).reshape((self.size, self.size, self.D)) * 
                 self.displacements(w)).filled(0)).sum(0)
    
    #@optional_w
    def accelerations(self, w=False):
        return self.forces(w)/self.masses
    
    @optional_state
    def differential(self, state):
        w = state['w']
        result = numpy.empty_like(w)
        result['pos'] = w['vel']
        result['vel'] = self.accelerations(w)
        return result