# 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.

# flow ninja: the purpose of this is to manage collections of vector estimates and perform
#             assimilations of local flow conditions.

# NOTE: this does not handle 3D flows presently, this only handles 2D flows (az and mag),
#       but . . the flows can have 3D context (e.g., a z coordinate). When there is a 3D
#       flow sensor available, I will update this to work in 3D.

import os
import sys
import copy
from math import *
import numpy as np
import pandas as pd
from threading import Lock

from weight_engine import *
from particle_ninja import *

# import number ninja (you need to specify this yourself!)
if os.name == 'posix':
    try:
        sys.path.append (os.path.abspath ('../packages/'))
    except:
        try:
            sys.path.append (os.path.abspath ('./packages/'))
        except:
            pass
        
else:
    sys.path.append ('C://Users//tom//Dropbox//python')

from number_ninja import *

class flow_ninja:
    def __init__ (self, weight_engine_in = None, model = 'invdist', space_wt = 0.75, space_zerodist = 100.0,
                  time_zerodist = 3000.0, space_exponent = 0.2, time_exponent = 0.2, is_wind = False,
                  threadsafe = True):
        '''
        constructor creates the number ninjas to handle az and mag data
        
        weight_engine_in = pre-defined weight engine, if unsupplied, takes the weight engine
                           parameters defined in other arguments
        model = the model to use for flow weighting
        space_wt = space weight for weight engine, the time weight is 1 - space_wt
        space_zerodist = the distance whereby space influence is zero
        time_zerodist = the distance whereby time influence is zero
        space_exponent = the space weighting exponent
        time_exponent = the time weighting exponent
        is_wind = config to denote this is a wind measurement (and thus the direction of flow
               is input as the direction the flow is coming from, not the direction that it
               is going!). This will be stored internally as the direction it is going (as per
               the semantic concept of flow).
        threadsafe = if this is false, the threadlock is replaced with a fake lock, which allows easier
                     pickling (but obviously is not threadsafe).
        '''
        
        self.is_wind = is_wind
        self.az = number_ninja (n = None, is_angle = True, log_time = True, always_return = True,
                                threadsafe = threadsafe)
        self.mag = number_ninja (n = None, is_angle = False, log_time = True, always_return = True,
                                 threadsafe = threadsafe)
        
        # either manually construct weight engine, or use pre-defined one
        if not weight_engine_in is None:
            self.weight_engine = weight_engine_in
        else:
            self.weight_engine = weight_engine (model, space_wt, space_zerodist, time_zerodist,
                                                space_exponent, time_exponent)
        
        return
    
    def add_data_np (self, azs, mags, xs = None, ys = None, zs = None, ts = None):
        '''
        method to add numpy arrays of vals, xs, ys, zs, ts. Note that this uses the
        add_data method incrementally to guard against issues with the data and to
        honor the pre-existing setup of the number ninja. This is by no means efficient,
        and direct assignment of the arrays is another option.
            
        azs = the azimuth values (not optional!)
        mags = the magnitude values (not optional!)
        xs = x positions (optional)
        ys = y positions (optional)
        zs = z positions (optional)
        ts = time positions as a datetime object (optional, defaults to now)
        '''
        
        # set defaults
        x = None
        y = None
        z = None
        t = None

        # corece to numpy arrays if these are pandas series (pandas series index from 'index', which is not safe
        # to trust as the dataframe could be sliced or the index could be something random)
        if type (azs) is pd.core.series.Series:
            azs = np.array (azs)
        if type (mags) is pd.core.series.Series:
            mags = np.array (mags)
        if type (xs) is pd.core.series.Series:
            xs = np.array (xs)
        if type (ys) is pd.core.series.Series:
            ys = np.array (ys)
        if type (zs) is pd.core.series.Series:
            zs = np.array (zs)
        if type (ts) is pd.core.series.Series:
            ts = np.array (ts)

        # incrementally loop over the data and append to the flow ninja
        for i in range (0, azs.shape[0]):
            if not np.isnan (azs[i]) and not np.isnan (mags[i]):
                if not xs is None:
                    x = xs[i]
                if not ys is None:
                    y = ys[i]
                if not zs is None:
                    z = zs[i]
                if not ts is None:
                    t = ts[i]
                
                self.add_data (az = azs[i], mag = mags[i], x = x, y = y, z = z, t = t)
            
        return    
    
    def add_data (self, az, mag, x = None, y = None, z = None, t = None):
        '''
        the method adds values to the list, with a noted x, y, z, time position
        
        az = the azimuth of flow (in the case of wind, the direction it is coming from) (deg)
        mag = the magnitude of flow
        x = x position (optional)
        y = y position (optional)
        z = z position (optional)
        t = time position as a datetime object (optional, defaults to now)
        '''
        # flip az if this is wind
        if self.is_wind:
            az = (180.0 + az) % 360.0
                
        # add data to underlying number ninjas
        self.az.add_data (az, x = x, y = y, z = z, t = t)
        self.mag.add_data (mag, x = x, y = y, z = z, t = t)
        return
    
    def sample (self, test_x, test_y, test_t, test_z = None, return_index = False):
        '''
        method to perform weighted sample of flow conditions at a given location in
        space and time. This can be used for stochastic simulations of lagrangian
        style advection (among other things).
        
        test_x = x location to sample from
        test_y = y location to sample from
        test_z = z location to sample from (optional)
        test_t = time location to sample from
        return_index = also return the sampled index
        returns az and mag (and index if specified)
        '''
        
        if len (self.az.vals) > 0:
            # get the weighted sample of the az, then use index to get mag
            az, index = self.az.spacetime_sample (test_x = test_x, test_y = test_y, test_t = test_t,
                                                  weight_engine = self.weight_engine, test_z = test_z)
            mag = float (self.mag.vals [index])

            if return_index:
                return (az, mag, index)
            else:
                return (az, mag)

        else:
            if return_index:
                return (np.nan, np.nan, np.nan)
            else:
                return (np.nan, np.nan)
    
    def calc_mean (self, test_x, test_y, test_t, test_z = None, scalar_mean = True):
        '''
        method to perform weighted mean in space for a given location
        
        test_x = x location to sample from
        test_y = y location to sample from
        test_z = z location to sample from (optional)
        test_t = time location to sample from
        scalar_mean = calculate the scalar mean, if false, return the vector mean
        returns az and mag
        '''
        # if we are returning the scalar mean, calculate the means separately
        if scalar_mean:
            az, influence = self.az.spacetime_mean (test_x = test_x, test_y = test_y, test_t = test_t,
                                                    weight_engine = self.weight_engine, test_z = test_z, circle_ranging = True)
            mag, influence = self.mag.spacetime_mean (test_x = test_x, test_y = test_y, test_t = test_t,
                                                      weight_engine = self.weight_engine, test_z = test_z)
        
        else:
            # use the mag ninja to get the influence array and manually create vector mean (discard mean mag)
            temp, influence = self.mag.spacetime_mean (test_x = test_x, test_y = test_y, test_t = test_t,
                                                       weight_engine = self.weight_engine, test_z = test_z)
            
            # create weighted vector mean
            az = np.nan
            mag = np.nan

            # redefine from np.nan if we are able to perform the calculation
            if influence.shape[0] > 0:
                influence_sum = influence.sum ()            # calc the sum of influence
                if not np.isnan (influence_sum):
                    if influence_sum > 0.0:
                        na_mask = ~np.logical_or (np.isnan (self.az.vals), np.isnan (influence))
                        mask = np.logical_and ((influence > 0.0), na_mask)
                        az_masked = self.az.vals [mask]
                        mag_masked = self.mag.vals [mask]
                        influence_masked = influence [mask]
                        az_masked_radians = az_masked * np.pi / 180.0
                        xcp = mag_masked * np.sin (az_masked_radians)
                        ycp = mag_masked * np.cos (az_masked_radians)
                        xcp_mean = np.average (xcp, weights = influence_masked)
                        ycp_mean = np.average (ycp, weights = influence_masked)
                        az = degrees (atan2 (xcp_mean, ycp_mean))
                        az = az % 360.0
                        mag = sqrt (xcp_mean**2.0 + ycp_mean**2.0)

        return (az, mag)
    
    def assimilate (self, locs, scalar_mean = True, az_colname = 'az', mag_colname = 'mag'):
        '''
        method to assimilate for a pandas dataframe containing coordinates in space and time.
        This is a useful method for making assimilations to view in paraview or qgis as flow glyphs
        
        locs = pandas dataframe containing points where we need to create mean estimates for
               must have 'x', 'y', 'z', 't', cols. Can be missing 'z' column if no z coords available.
        scalar_mean = calculate the scalar mean (instead of the vector mean) - these two numbers
                      are rather different and you may choose one vs another for a given flow
                      assimilation task
        az_colname = the aximuth column name (for appending to existing dataframes)
        mag_colname = the magnitude column name (for appending to existing dataframes)
        returns the dataframe with az and mag columns
        '''
        locs [az_colname] = np.nan
        locs [mag_colname] = np.nan
        
        # calculate weighted means
        for i in range (0, locs.shape[0]):
            x = locs.loc [i, 'x']
            y = locs.loc [i, 'y']
            t = locs.loc [i, 't']
            
            if 'z' in locs.columns:
                z = locs.loc [i, 'z']
            else:
                z = None
            
            # perform calculation and append to dataframe
            az, mag = self.calc_mean (test_x = x, test_y = y, test_t = t, test_z = z, 
                                      scalar_mean = scalar_mean)
            locs.loc[i, az_colname] = az
            locs.loc[i, mag_colname] = mag
    
        return (locs)
    
    def repeat_sample (self, locs, az_colname = 'az', mag_colname = 'mag'):
        '''
        method to repeatedly pull samples from the flow ninja
        
        locs = pandas dataframe containing points where we need to create mean estimates for
               must have 'x', 'y', 'z', 't', cols. Can be missing 'z' column if no z coords available.
        az_colname = the aximuth column name (for appending to existing dataframes)
        mag_colname = the magnitude column name (for appending to existing dataframes)
        returns the dataframe with az and mag columns
        '''
        
        locs [az_colname] = np.nan
        locs [mag_colname] = np.nan
        
        # calculate the sample
        for i in range (0, locs.shape[0]):
            x = locs.loc [i, 'x']
            y = locs.loc [i, 'y']
            t = locs.loc [i, 't']
            
            if 'z' in locs.columns:
                z = locs.loc [i, 'z']
            else:
                z = None
            
            # make a sample
            az, mag = self.sample (test_x = x, test_y = y, test_t = t, test_z = z)
            locs.loc[i, az_colname] = az
            locs.loc[i, mag_colname] = mag
    
        return (locs)
    
    def timeslice (self, min_time, max_time):
        '''
        method to timeslice the number ninjas to limit the volume of numbers in the ninja
        
        min_time = the minimum time to keep
        max_time = the maximum time to keep
        '''
        self.az.timeslice (min_time, max_time)
        self.mag.timeslice (min_time, max_time)
        return
    
    def add_data_file (self, filename, az_col = 'az', mag_col = 'mag'):
        '''
        method to add data from a file. Do note that this could be optimized for huge flows
        as it presently reads in the file twice as it calls the underlying number ninja add_data_file
        method. Also note all the important gotchas from the number_ninja add_data_file method. In
        particular it is vital to note that this only appends data into the ninjas. It does not clear
        the ninjas and replace the data.
        
        filename = the filename to add data from
        az_col = the azimuth column (defaults to az)
        mag_col = the magnitude column (defaults to mag)
        '''
        self.az.add_data_file (filename, val_col = az_col)
        self.mag.add_data_file (filename, val_col = mag_col)
        return

    def sort_in_time (self):
        '''
        method to sort the underlying ninjas in time. Note that this does an 'in place' sort, and will wipe
        any context implied by ordering.
        '''
        self.az.sort_in_time ()
        self.mag.sort_in_time ()
        return

    def write (self, filename):
        '''
        method to write out the record as a pandas dataframe
        filename = filename to write out to
        '''
        
        if self.az.vals.shape[0] > 0:
            df = pd.DataFrame (self.az.vals, columns = ['az'])
            df['mag'] = self.mag.vals
            
            if len (self.az.x) > 0:
                df['x'] = self.az.x
            if len (self.az.y) > 0:
                df['y'] = self.az.y
            if len (self.az.z) > 0:
                df['z'] = self.az.z
            if len (self.az.t) > 0:
                df['t'] = self.az.t
            
            df.to_csv (filename, index = False)
        
        return
    
    def deepcopy (self):
        '''
        method to deepcopy the individual number ninjas
        returns a copy
        '''
        selfcopy = copy.copy (self)
        selfcopy.az = self.az.deepcopy ()
        selfcopy.mag = self.mag.deepcopy ()
        return (selfcopy)
  
    



