################################################################################
#  Copyright (c) 2014 Paul Scherrer Institute. All rights reserved.
################################################################################

import sys
import time
import math
import os.path

from array import array
import jarray

import java.lang.Class as Class
import java.lang.Object as Object
import java.beans.PropertyChangeListener
import java.util.concurrent.Callable
import java.util.List
import java.lang.reflect.Array

import org.python.core.PyArray as PyArray

import ch.psi.utils.Threading as Threading
import ch.psi.utils.State as State
import ch.psi.utils.Convert as Convert
import ch.psi.pshell.core.Controller.CommandSource as CommandSource
import ch.psi.pshell.data.PlotDescriptor as PlotDescriptor
import ch.psi.pshell.dev.Readable as Readable
import ch.psi.pshell.dev.Writable as Writable
import ch.psi.pshell.dev.DeviceAdapter as DeviceAdapter
import ch.psi.pshell.dev.DeviceListener as DeviceListener
import ch.psi.pshell.dev.MotorGroupBase.MoveMode as MoveMode
import ch.psi.pshell.epics.Epics as Epics
import ch.psi.pshell.epics.EpicsScan as EpicsScan
import ch.psi.pshell.scan.LineScan 
import ch.psi.pshell.scan.AreaScan
import ch.psi.pshell.scan.VectorScan
import ch.psi.pshell.scan.ManualScan


################################################################################
#Scan functions
################################################################################
def onScanBeforeReadout(scan):
    try:
        if (scan.before_read!=None):
            scan.before_read()
    except AttributeError:
        pass

def onScanAfterReadout(scan):
    try:
        if (scan.after_read!=None):
            scan.after_read()
    except AttributeError:
        pass

class LineScan(ch.psi.pshell.scan.LineScan):
    def onBeforeReadout(self):
        onScanBeforeReadout(self)

    def onAfterReadout(self):
        onScanAfterReadout(self)

class AreaScan(ch.psi.pshell.scan.AreaScan):
    def onBeforeReadout(self):
        onScanBeforeReadout(self)

    def onAfterReadout(self):
        onScanAfterReadout(self)

class VectorScan(ch.psi.pshell.scan.VectorScan):
    def onBeforeReadout(self):
        onScanBeforeReadout(self)

    def onAfterReadout(self):
        onScanAfterReadout(self)

class ManualScan (ch.psi.pshell.scan.ManualScan):
    def __init__(self, writables, readables, start = None, end = None, steps = None, relative = False):
        ch.psi.pshell.scan.ManualScan.__init__(self, writables, readables, start, end, steps, relative, controller)

    def append(self,setpoints, positions, values):
        ch.psi.pshell.scan.ManualScan.append(self, to_object_array(setpoints), to_object_array(positions), to_object_array(values))


def sleep(sec):
    """Sleep.

    Args:
        sec(float): Time in seconds.

    """ 
    time.sleep(sec)

def to_list(obj):
    if isinstance(obj,tuple):
        return list(obj)
    if not isinstance(obj,list):
        return [obj,]
    return obj

def is_list(obj):
    return isinstance(obj,tuple) or isinstance(obj,list)

class_types = {
    'b': "java.lang.Byte",
    'h': "java.lang.Short",
    'u': "java.lang.Integer",
    'i': "java.lang.Integer",
    'l': "java.lang.Long",
    'c': "java.lang.Character",
    'f': "java.lang.Float",
    'd': "java.lang.Double",
    's': "java.lang.String",
    'o': "java.lang.Object",

    '[b': "[B",          #byte
    '[h': "[S",          #short
    '[u': "[I",          #int
    '[i': "[I",          #int
    '[l': "[J",          #long
    '[c': "[C",          #char
    '[f': "[F",          #float
    '[d': "[D",          #double
    '[s': "[Ljava.lang.String;",
    '[o': "[Ljava.lang.Object;",
}


def to_array(obj, type):   
    """Convert Python list to Java array.

    Args:
        obj(list): Original data.
        type(str): array type 'b' = byte, 'h' = short, 'i' = int, 'l' = long,  'f' = float, 'd' = double, 
                              'c' = char, 's' = String,  'o' = Object 
    Returns:
        Java array.

    """ 
    if type[0] == '[':
        type = type[1:]    
    arrayType = class_types.get("["+type)    

    if obj is None:
        return None
    if isinstance(obj,java.util.List):
        obj = obj.toArray(java.lang.reflect.Array.newInstance(Class.forName(class_types.get(type)),0))    
        obj = Convert.wrapperArrayToPrimitiveArray(obj)
    if isinstance(obj,PyArray):
        return obj        
    if is_list(obj):
        if type=='o' or  type== 's':
            ret = java.lang.reflect.Array.newInstance(Class.forName(class_types.get(type)),len(obj))
            for i in range (len(obj)):
                if is_list(obj):
                    ret[i] = to_array(obj[i],type)
                elif type == 's':
                    ret[i] = str(obj[i])
                else:
                    ret[i] = obj[i]        
            return ret            

        if len(obj)>0 and is_list(obj[0]):
            if  len(obj[0])>0 and is_list(obj[0][0]):
                    ret = java.lang.reflect.Array.newInstance(Class.forName(arrayType),len(obj),len(obj[0]))                    
                    for i in range(len(obj)):
                        ret[i]=to_array(obj[i], type)            
                    return ret
            else:
                ret = java.lang.reflect.Array.newInstance(Class.forName(arrayType),len(obj))
                for i in range(len(obj)):
                    ret[i]=to_array(obj[i], type)            
                return ret     
        return jarray.array(obj,type)
    return obj

def to_double_array(obj):   
    return to_array(obj, 'd')

def to_float_array(obj):   
    return to_array(obj, 'f')

def to_long_array(obj):   
    return to_array(obj, 'l')

def to_int_array(obj):   
    return to_array(obj, 'i')

def to_short_array(obj):   
    return to_array(obj, 'h')

def to_byte_array(obj):   
    return to_array(obj, 'b')

def to_string_array(obj):   
    return to_array(obj, 's')

def to_object_array(obj):
    return to_array(obj, 'o')

def lscan(writables, readables, start, end, steps, latency=0.0, relative = False, context=None, before_read=None, after_read=None):
    """Line Scan: positioners change together, linearly from start to end positions.

    Args:
        writables(list of Writable): Positioners set on each step.
        readables(list of Readable): Sensors to be sampled on each step.
        start(list of float): start positions of writables.
        end(list of float): final positions of writables.
        steps(int or float or list of float): number of scan steps (int) or step size (float).
        relative (bool, optional): if true, start and end positions are relative to current at start of the scan
        latency(float, optional): sleep time in each step before readout, defaults to 0.0.
        context(str, optional): plotting context name.
        before_read (function): callback on each step, before each readout.
        after_read (function): callback on each step, after each readout.

    Returns:
        ArrayList of ScanRecord objects.

    """    
    latency_ms=int(latency*1000)
    writables=to_list(writables)
    readables=to_list(readables)
    start=to_list(start)
    end=to_list(end)    
    scan = LineScan(writables,readables, start, end , steps, relative, latency_ms, controller)
    scan.before_read=before_read
    scan.after_read=after_read
    scan.setPlotName(context)
    scan.start()
    return scan.getResult()

def vscan(writables, readables, vector, latency=0.0, relative = False, context=None, before_read=None, after_read=None):
    """Vector Scan: positioners change following values provided in a vector.

    Args:
        writables(list of Writable): Positioners set on each step.
        readables(list of Readable): Sensors to be sampled on each step.
        vector(list of list of float): table of positioner values.
        relative (bool, optional): if true, start and end positions are relative to current at start of the scan
        latency(float, optional): sleep time in each step before readout, defaults to 0.0.        
        context(str, optional): plotting context name.
        before_read (function): callback on each step, before each readout.
        after_read (function): callback on each step, after each readout.

    Returns:
        ArrayList of ScanRecord objects.

    """    
    latency_ms=int(latency*1000)
    writables=to_list(writables)
    readables=to_list(readables)
    scan = VectorScan(writables,readables, vector, relative, latency_ms, controller)
    scan.before_read=before_read
    scan.after_read=after_read
    scan.setPlotName(context)
    scan.start()
    return scan.getResult()

def ascan(writables, readables, start, end, steps, latency=0.0, relative = False, context=None, before_read=None, after_read=None):
    """Area Scan: multi-dimentional scan, each positioner is a dimention.

    Args:
        writables(list of Writable): Positioners set on each step.
        readables(list of Readable): Sensors to be sampled on each step.
        start(list of float): start positions of writables.
        end(list of float): final positions of writables.
        steps(list of int or list of float): number of scan steps (int) or step size (float).
        relative (bool, optional): if true, start and end positions are relative to current at start of the scan
        latency(float, optional): sleep time in each step before readout, defaults to 0.0.
        context(str, optional): plotting context name.
        before_read (function): callback on each step, before each readout.
        after_read (function): callback on each step, after each readout.

    Returns:
        ArrayList of ScanRecord objects.

    """    
    latency_ms=int(latency*1000)
    writables=to_list(writables)
    readables=to_list(readables)
    start=to_list(start)
    end=to_list(end)    
    scan = AreaScan(writables,readables, start, end , steps, relative, latency_ms, controller)
    scan.before_read=before_read
    scan.after_read=after_read
    scan.setPlotName(context)
    scan.start()
    return scan.getResult()    

def tscan(readables, points, interval, context=None, before_read=None, after_read=None):
    """Time Scan: sensors are sampled in fixed time intervals.

    Args:
        readables(list of Readable): Sensors to be sampled on each step.
        points(int): number of samples.
        interval(float): time interval between readouts.
        context(str, optional): plotting context name.
        before_read (function): callback on each step, before each readout.
        after_read (function): callback on each step, after each readout.

    Returns:
        ArrayList of ScanRecord objects.

    """      
    latency_ms=int(interval*1000)
    writables=[]
    readables=to_list(readables)
    start=[0,]
    end=[points,]
    steps=points
    scan = LineScan(writables,readables, start, end , steps, False, latency_ms, controller)
    scan.before_read=before_read
    scan.after_read=after_read
    scan.setPlotName(context)
    scan.start()
    return scan.getResult()

def escan(name, context=None):
    """Epics Scan: execute an Epics Scan Record.

    Args:
        name(str): Name of scan record.
        context(str, optional): plotting context name.

    Returns:
        ArrayList of ScanRecord objects.

    """      
    scan = EpicsScan(name, controller)
    scan.setPlotName(context)
    scan.start()
    return scan.getResult()

    
################################################################################
#Plot functions
################################################################################

def plot(data, name = None, xdata = None, ydata=None, context=None):    
    """Request one or multiple plots of user data (1d, 2d or 3d)

    Args:
        data: array or list of values. For multiple plots, array of arrays or lists of values.
        name(str or list of str, optional): Plot name or list of names (if multiple plots). 
        xdata: array or list of values. For multiple plots, array of arrays or lists of values.
        ydata: array or list of values. For multiple plots, array of arrays or lists of values.
        context(str, optional): plotting context name.

    Returns:
        ArrayList of Plot objects.

    """      
    if (name is not None) and is_list(name):
        if len(name)==0:
            name=None;
        plots = java.lang.reflect.Array.newInstance(Class.forName("ch.psi.pshell.data.PlotDescriptor"), len(data))                    
        for i in range (len(data)):         
            plotName = None if (name is None) else name[i]
            x = xdata            
            if is_list(x) and len(x)>0 and (is_list(x[0]) or isinstance(x[0] , java.util.List)):
                x = x[i]                
            y = ydata
            if is_list(y) and len(y)>0 and (is_list(y[0]) or isinstance(y[0] , java.util.List)):
                y = y[i]
            plots[i] =  PlotDescriptor(plotName , to_double_array(data[i]), to_double_array(x), to_double_array(y))            
        return controller.plot(plots,context)
    else:
        plot = PlotDescriptor(name, to_double_array(data), to_double_array(xdata), to_double_array(ydata))
        return controller.plot(plot,context)

################################################################################
#Data access functions
################################################################################

def load_data(path, page=0):    
    slice = controller.dataManager.getData(path, page)
    return slice.sliceData

def get_attributes(path):    
    return controller.dataManager.getAttributes(path)

def save_dataset(path, data):
    data = to_double_array(data)
    controller.dataManager.setDataset(path,data)

def create_dataset(path, type, unsigned=False, dimensions=None):
    type = Class.forName(class_types.get(type,type))
    controller.dataManager.createDataset(path, type, unsigned, dimensions)

def create_table(path, names, types=None, lengths=None, dims=None):
    if (types is not None):
        for i in range (len(types)):
            types[i] = Class.forName(class_types.get(types[i],types[i]))
    controller.dataManager.createDataset(path, names,types, lengths,dims)

def save_data_item(path, data, index=None):
    data = to_double_array(data)
    if index is None:
        controller.dataManager.appendItem(path, data)
    else:
        controller.dataManager.setItem(path, data, index)

def save_table_item(path, data):
    #data = to_object_array(data)
    if is_list(data):
        arr = java.lang.reflect.Array.newInstance(Class.forName("java.lang.Object"),len(data))
        for i in range (len(data)):
            if is_list(data[i]):
                arr[i] = to_double_array(data[i])
            else:
                arr[i] = data[i]        
        data=arr
    controller.dataManager.appendItem(path, data)

def flush_data():
    controller.dataManager.flush()

def set_attribute(path, name, value, unsigned = False):
    controller.dataManager.setAttribute(path,name, value, unsigned)

def get_current_data_group():
    return controller.dataManager.getCurrentGroup()

def set_flush_records(value):
    """Override, for the executing scripot, the configuration option to flush each record.
    """
    controller.dataManager.flushRecords = value

################################################################################
#Epics access functions
################################################################################

channel_types = {
    'b': "java.lang.Byte",
    'i': "java.lang.Short",
    'l': "java.lang.Integer",
    'f': "java.lang.Float",
    'd': "java.lang.Double",
    's': "java.lang.String",
    
    '[b': "[B",
    '[i': "[S",
    '[l': "[I",
    '[f': "[F",
    '[d': "[D",
    '[s': "[Ljava.lang.String;",
}

def caget(name, type='s', size=None):    
    """Reads an Epics PV.

    Args:
        name(str): PV name
        type(str, optional): type of PV, defaults 'd'.
            Scalar values: 'b', 'i', 'l', 'd', 's'.
            Array: values: '[b', '[i,', '[l', '[d', '[s'.
        size (int, optional): for arrays, number of elements to be read. Default read all.

    Returns:
        PV value

    """          
    return Epics.get(name, Class.forName(channel_types.get(type,type)),size)

def cawait(name, value, timeout=None, comparator=None, type='s', size=None):    
    """Wait for a PV to have a given value.

    Args:
        name(str): PV name
        value (obj): value to compare to
        timeout(float, optional): time in seconds to wait. If None, waits forever.
        comparator(java.util.Comparator, optional): if None, equality.
        type(str, optional): type of PV, defaults 'd'.
            Scalar values: 'b', 'i', 'l', 'd', 's'.
            Array: values: '[b', '[i,', '[l', '[d', '[s'.
         size (int, optional): for arrays, number of elements to be read. Default read all.
    """          
    if (timeout is not None):
        timeout = int(timeout*1000)
    value = adjust_value(value)
    Epics.waitValue(name, value, comparator, timeout, Class.forName(channel_types.get(type,type)),size)


def adjust_value(value, var_type=None):
    if (value is None):
        return value
    if (var_type is not None):
        if is_list(value):    
            var_type = var_type.replace(',','').replace('[','')
            ret = []
            for item in value:
               ret.append(adjust_value(item), var_type) 
            value = ret
        else:
            var_type = var_type.lower()
            if var_type=='b':
                value = byte(value)
            elif var_type=='i':
                value = short(value)
            elif var_type=='l':
                value = int(value)
            elif var_type=='f':
                value = float(value)
            elif var_type=='d':
                value = float(value)
            elif var_type=='s':
                value = str(value)

    if isinstance(value,tuple):
        value =  list(value)
    if isinstance(value,list):
        list_type = type(value[0])
        array_types = {
            int: "i",
            long: "l",
            float:"d",
            str:Class.forName("java.lang.String"),    
        }        
        array_type = array_types.get(type(value[0]),'d')
        array = PyArray(array_type)        
        array.fromlist(value)         
        value=array
    return value

def caput(name, value, timeout = None):   
    """Writes to an Epics PV.

    Args:
        name(str): PV name
        value(scalar, string or array): new PV value.
        timeout(int, optional): timeout in seconds to the write. If None waits forever to completion.
    """  
    value=adjust_value(value)
    if (timeout is not None):
        timeout = int(timeout*1000)
    return Epics.put(name, value, timeout)

def caputq(name, value):        
    """Writes to an Epics PV and does not wait.

    Args:
        name(str): PV name
        value(scalar, string or array): new PV value.
    """  
    value=adjust_value(value)
    return Epics.putq(name, value)

def create_channel_device(channelName, type = None, size = None, deviceName = None):
    if type is None:
        type = 's'
    type = Class.forName(channel_types.get(type,type))
    dev = Epics.newChannelDevice(deviceName, channelName, type)
    dev.initialize()
    if (size is not None):
        dev.setSize(size)
    return dev
    
def create_channel(name, type = None, size = None):        
    if type is None:
        type = 's'
    type = Class.forName(channel_types.get(type,type))
    return Epics.newChannel(name, type, size)

class Channel(java.beans.PropertyChangeListener):
    def __init__(self, name, type = None, size = None, callback=None):
        self.channel = create_channel(name, type, size)
        self.callback = callback
    
    def get_name(self):
        return self.channel.name

    def get_size(self):
        return self.channel.size

    def set_size(self, size):
        self.channel.size = size

    def is_connected(self):
        return self.channel.connected

    def is_monitored(self):
        return self.channel.monitored

    def set_monitored(self, value):
        self.channel.monitored = value
        if (value):
            self.channel.addPropertyChangeListener(self)
        else:
            self.channel.removePropertyChangeListener(self)

    def propertyChange(self, pce):
        if pce.getPropertyName() == "value":
            if self.callback is not None:
                self.callback(pce.getNewValue())

    def put(self, value, timeout=None):
        if (timeout==None):
            self.channel.setValue(value)
        else:
            self.channel.setValueAsync(value).get(int(timeout*1000), java.util.concurrent.TimeUnit.MILLISECONDS);

    def putq(self, value):
        self.channel.setValueNoWait(value)

    def get(self):
        return self.channel.getValue()

    def wait_for_value(self, value, timeout=None, comparator=None):
        if comparator is None:
            if timeout is None:
                self.channel.waitForValue(value)
            else:
                self.channel.waitForValue(value, int(timeout*1000))
        else:
            if timeout is None:
                self.channel.waitForValue(value, comparator)
            else:
                self.channel.waitForValue(value, comparator, int(timeout*1000))


    def close(self):
        self.channel.destroy()        
        

################################################################################
#Concurrent execution
################################################################################
class Callable(java.util.concurrent.Callable):
    def __init__(self, method, *args):
        self.method = method
        self.args = args
    def call(self):
        #try:
            return self.method(*self.args)
        #except:
        #    traceback.print_exc(file=sys.stderr)


def fork(*methods):
    """Start execution of methods in parallel. 

    Args:
        *methods(method references)

    Returns:
        List of callable objects
    """      
    callables = []
    for m in methods:
        if is_list(m):
            callables.append(Callable(m[0],*m[1]))
        else:
            callables.append(Callable(m))
    return Threading.fork(callables)

def join(callables):
    """Wait parallel execution of methods.

    Args:
        callables(list of Callables) : as returned from fork
    """      
    try:
        return Threading.join(callables)
    except java.util.concurrent.ExecutionException as ex:
        raise ex.getCause()
        
def parallelize(*methods):
    """Equivalent to fork + join

    Args:
        *methods(method references)
    """   
    callables = fork(*methods)
    return join(callables)


################################################################################
#Background tasks control
################################################################################

def startTask(script, delay = 0.0, interval = -1):
    """Start a background task

    Args:        
         script(str): Name of the script implementing the task
         delay(float, optional): time in seconds for the first execution. Default starts immediately.
         interval(float, optional): time in seconds for between execution. If negative (default), single-execution.
    """      
    delay_ms=int(delay*1000)
    interval_ms=int(interval*1000) if (interval>=0) else int(interval)
    controller.taskManager.create(script, delay_ms, interval_ms)
    controller.taskManager.start(script)

def stopTask(script, force = False):
    """Stop a background task

    Args:        
         script(str): Name of the script implementing the task
         force(boolean, optional): interrupt current execution, if running
    """      
    controller.taskManager.remove(script, force)



################################################################################
#Logging
################################################################################

def log(log):        
    """Writes log to data file.

    Args:        
         log(str): Log string

    """  
    controller.dataManager.appendLog(str(log))

################################################################################
#Standard Libraries
################################################################################
ca_channel_path=os.path.join(controller.setup.getStandardLibraryPath(), "epics")
sys.path.append(ca_channel_path)
#This is to destroy previous context of _ca (it is not shared with PShell)
if run_count > 0:
    if sys.modules.has_key("_ca"):
        import _ca
        _ca.initialize()


################################################################################
#Utilities
################################################################################

#Float range -> Useful for scanning is a range
def frange_gen(start, finish, step):
    while ((step >= 0.0) and (start <= finish)) or ((step < 0.0) and (start >= finish)):
      yield start
      start += step

def frange(start, finish, step, enforce_finish = False, inclusive_finish = False):
    """Create a list with a range of float values
    """
    ret = list(frange_gen(start, finish, step))
    if len(ret) > 0:
        if inclusive_finish == False:
            if ret[-1]==finish:
                del ret[-1]
        if enforce_finish and ret[-1]!=finish:
                ret.append(finish)
    return ret

def set_status(status):
    set_preference("status", status)

def set_preference(name, value):
    """Hints to graphical layer:    
    Args:
        name(str): Preference name
            - "scanPlotDisabled": enable/disable scan plot (True/False)
            - "scanTableDisabled": enable/disable scan table (True/False)
            - "enabledPlots": select Readables to be plotted (list of Readable or String (Readable names))
            - "status": set application status
        value(object): preference valye
    """  
    value = to_array(value, 'o') #If list then convert to Object array
    controller.setPreference(CommandSource.script, name, value)

def inject_vars():
    """Re-inject devices and  interpreter variables 
    """  
    controller.injectVars(CommandSource.script)

def run(script_name, globals = globals(), locals = locals()):
    """Run script: can be absolute path, relative, or short name to be search in the path.
    """  
    global _
    script = controller.scriptManager.library.resolveFile(script_name)
    if script is not None and os.path.isfile(script):
        _ = None
        execfile(script, globals, locals)
        return _
    #print >> sys.stderr, "Invalid script: " + str(script_name)
    raise IOError("Invalid script: " + str(script_name))

def set_return(value):
    global _
    #In Jython, the output of last statement is not returned  when running a file so we have to use a return variable
    _=value         #Used when running file
    return value    #Used when parsing file

################################################################################
#Local
################################################################################
try:
    run("local")
except IOError: #If file not present ignore exception
    pass