# FLOW NINJA
# Copyright 2016-2017 Thomas E. Barchyn

# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
# associated documentation files (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge, publish, distribute, 
# sublicense, and/or sell copies of the Software, and to permit persons to whom the Software
# is furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all copies or
# substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
# BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

# particle ninja: a simple particle prototype class to be built upon

import sys
import os
from math import *
import numpy as np
import pandas as pd

from flow_ninja import *

# import number ninja
from number_ninja import *

class particle_ninja:
    def __init__ (self, spec_flow, timestep = 1.0, forward = True, topo = None):
        '''
        constructor takes a flow ninja as an argument, which allows sampling. Note that the logs
        are set up in the set_start method, not here.
        
        spec_flow = an initialized flow ninja that allows sampling
        timestep = timestep in seconds (note: this will be converted to integer
                   microseconds, so if you require more precision than 0.001 s
                   provide a np.timedelta64 object instead of the time in seconds)
        forward = forward (downwind) advection of backwards (upwind) advection, note that this
                  goes backwards in time (so timesteps must be negative)
        topo = a point_interpolator object (or object where object.interpolate(x_loc, y_loc)
               produces a value). This will be used to assign the z value.
        '''
        self.spec_flow = spec_flow
        self.forward = forward
        
        if topo is None:
            self.interpolate_topo = False
        else:
            self.topo = topo
            self.interpolate_topo = True
            
        # check to make sure the timestep is proper sign
        if self.forward and timestep <= 0.0:
            print ('WARNING: converting particle ninja timestep to positive (forward in time)')
            timestep = -1.0 * timestep
        
        if not self.forward and timestep >= 0.0:
            print ('WARNING: converting particle ninja timestep to negative (backwards in time)')
            timestep = -1.0 * timestep
        
        self.timestep = timestep                # this will be np.timedelta64 if using np datetime format
        self.timestep_s = timestep              # this will be seconds (float)
        return
    
    def set_start (self, x_start, y_start, t_start, z_start = 0.0):
        '''
        method to set the start for a particle
        
        x_start = x starting position
        y_start = y starting position
        t_start = t starting time position
        z_start = z starting position (optional)
        '''
        
        self.x_start = x_start
        self.y_start = y_start
        self.z_start = z_start
        
        try:
            t_start = np.datetime64 (t_start)
            self.float_time = False
        except:
            self.float_time = True
        
        self.t_start = t_start
        self.log_time = 0.0         # internal log time for particle
        
        # update the timesteps to proper type
        if not self.float_time:
            if not type (self.timestep) is np.timedelta64:
                # add the timedelta in microseconds, because timedelta64 can only handle
                # integers (due to the complexity in handling fractions of non base 10 things)
                # this gives us a little bit of extra precision and avoids truncation
                # record a floating point timestep by conversion if necessary
                # coerce this to a timedelta object
                self.timestep = np.timedelta64 (int(self.timestep * 1000), 'ms')
            
            # make the timestep s in seconds
            if type (self.timestep_s) is np.timedelta64:
                self.timestep_s = np.divide (self.timestep, np.timedelta64 (1, 's'))
        
        # set internal coordinates
        self.x = x_start
        self.y = y_start
        self.t = t_start
        self.az = None
        self.mag = None
        self.dist = 0.0
        
        # assign topo if provided, else see if we have an interpolator
        if not z_start == 0.0 or not self.interpolate_topo:
            self.z = z_start                # use the assigned value
        else:
            self.z = self.topo.interpolate (self.x, self.y)

        # make numpy arrays for logs (instead of super slow pandas dataframes)
        self.log_x = number_ninja (minimalist_config = 'float')
        self.log_y = number_ninja (minimalist_config = 'float')
        self.log_z = number_ninja (minimalist_config = 'float')
        self.log_dist = number_ninja (minimalist_config = 'float')
        self.log_elapsed_time = number_ninja (minimalist_config = 'float')
        
        if self.float_time:
            self.log_t = number_ninja (minimalist_config = 'float')
        else:
            self.log_t = number_ninja (minimalist_config = 'time')
        
        # log it to the local log
        self.log_data ()
        return
    
    def move (self, sample = True):
        '''
        method to move the particle one step (note that this doesn't work in 3D). This either samples (randomly)
        or calculates a local average with the weighting engine. The latter is better as a low budget method
        of calculating an ensemble of lagrangian particle walks. But note that the mean is not equivalent to
        the ensemble approach (and could diverge significantly in any situation with large changes in the
        flow through space or time). Note that this takes vector means for the means, when not taking samples.
        
        sample = use sampling vs taking a weighted mean
        '''

        if sample:
            # sample an azimuth and magnitude with random weighted sampling
            self.az, self.mag = self.spec_flow.sample (test_x = self.x, test_y = self.y, test_t = self.t, test_z = self.z)
        else:
            # sample an azimuth and magnitude, with a vector mean magnitude
            self.az, self.mag = self.spec_flow.calc_mean (test_x = self.x, test_y = self.y, test_t = self.t, test_z = self.z,
                                                          scalar_mean = False)

        # make the move in planimetric coordinates
        self.x = self.x + (self.timestep_s * self.mag * sin (self.az * pi / 180.0))
        self.y = self.y + (self.timestep_s * self.mag * cos (self.az * pi / 180.0))
        
        # interpolate z location if supplied an interpolator
        if self.interpolate_topo:
            self.z = self.topo.interpolate (x_loc = self.x, y_loc = self.y)
        
        # advance time with a timedelta (make change if we haven't done so already)
        self.t = self.t + self.timestep
        self.log_time = self.log_time + self.timestep_s         # use floating point timestep
        
        # advance the dist
        self.dist = self.dist + (self.timestep_s * self.mag)
        
        # log it to the local log
        self.log_data ()
        return
    
    def backup (self, backup_dist):
        '''
        method to correct the location of the last position by backing up. This corrects the last location
        rather than moving backwards and adding a new record. Possibly, if moving backwards is something
        that is actually of interest, we can do this. Note that the backup_dist must be less than the distance
        traveled in the last timestep, it doesn't trace backwards up the particle advection.
        
        backup_dist = the distance to backup (must be less than last step)
        '''
        if self.log_dist.i < 2:
            return                  # just return as we cannot backup without at least 1 timestep
        else:
            # calculate relevant indices
            now = self.log_dist.i - 1
            then = now - 1
            
            # calculate fraction of last step to backup
            backup_frac = backup_dist / (self.log_dist.vals[now] - self.log_dist.vals[then])
            
            if backup_frac > 1.0:
                print ('WARNING: particle ninja cannot backup more than 1 step')
                return
            
            # perform weighted averages to update the last record of the number ninjas
            self.log_x.vals[now] = (self.log_x.vals[now] * (1.0 - backup_frac)) + (self.log_x.vals[then] * backup_frac)
            self.log_y.vals[now] = (self.log_y.vals[now] * (1.0 - backup_frac)) + (self.log_y.vals[then] * backup_frac)
            self.log_z.vals[now] = (self.log_z.vals[now] * (1.0 - backup_frac)) + (self.log_z.vals[then] * backup_frac)
            self.log_dist.vals[now] = self.log_dist.vals[now] - backup_dist
            
            # address the time differences
            time_rewind_s = backup_frac * self.timestep_s
            
            self.log_elapsed_time.vals[now] = self.log_elapsed_time.vals[now] - time_rewind_s
            
            if self.float_time:
                self.log_t.vals[now] = self.log_t.vals[now] - time_rewind_s
            else:
                time_rewind = np.timedelta64 (int(time_rewind_s * 1000), 'ms')
                self.log_t.vals[now] = self.log_t.vals[now] - time_rewind
                
        return
        
    def log_data (self):
        '''
        method to log each call to move to a pandas dataframe. If this is the first
        call, it creates the log dataframe as well (to ensure types are set properly
        after the first call).
        '''
        
        self.log_x.add_data (val = self.x)
        self.log_y.add_data (val = self.y)
        self.log_z.add_data (val = self.z)
        self.log_elapsed_time.add_data (val = self.log_time)
        self.log_dist.add_data (val = self.dist)
        self.log_t.add_data (val = self.t)
        return
    
    def get_log (self):
        '''
        method to get the logs as a pandas dataframe. This returns the data
        as a pandas dataframe.
        '''
        
        # scrub the minimalist arrays
        self.log_x.scrub_minimalist ()
        self.log_y.scrub_minimalist ()
        self.log_z.scrub_minimalist ()
        self.log_elapsed_time.scrub_minimalist ()
        self.log_dist.scrub_minimalist ()
        self.log_t.scrub_minimalist ()
        
        # create pandas dataframe
        pd_log = pd.DataFrame (self.log_x.vals, columns = ['x'])
        pd_log['y'] = self.log_y.vals
        pd_log['z'] = self.log_z.vals
        pd_log['elapsed_time'] = self.log_elapsed_time.vals
        pd_log['dist'] = self.log_dist.vals
        pd_log['t'] = self.log_t.vals
        return (pd_log)

    def calc_sinuosity (self):
        '''
        method to calc the present sinuosity. Sinuosity is defined as the ratio of total distance
        along the line to the distance between the start and end points. This ranges from 1 (straight
        line), to infinity (very very sinuos). I'm sure there are more complicated ways to
        determine the sinuosity, but for now, this gets at the point.
        
        returns sinuosity
        '''
        now = self.log_dist.i - 1                       # ensure to use the logged value
        dist_from_start = sqrt ((self.log_x.vals[now] - self.log_x.vals[0])**2.0 + (self.log_y.vals[now] - self.log_y.vals[0])**2.0)   
        
        # ratio total distance to distance from start
        sinuosity = self.log_dist.vals[now] / dist_from_start
        return (sinuosity)
        
    def calc_init_flowrate (self):
        '''
        method to calculate the initial flow rate at the first timestep
        '''
        try:
            init_dist = self.log_dist.vals[1] - self.log_dist.vals[0]    
            init_rate = init_dist / self.timestep_s
        except:
            init_rate = 0.0
            
        return (init_rate)
            