# -*- coding: utf-8 -*-
"""
Created on Sat Aug  1 18:13:53 2020

@author: g_dic
"""
import numpy as np
import os
import json, codecs
import pandas as pd
from distutils.version import StrictVersion
from qtpy.QtCore import QUrl, QRect, QPointF, Qt
from qtpy.QtGui import QDesktopServices, QIcon, QPainterPath, QPen, QColor
from qtpy.QtWidgets import QHBoxLayout, QGraphicsPathItem, qApp
from qtpy import uic

import flika
from flika import global_vars as g
from flika.window import Window
from flika.process.file_ import save_file_gui, open_file_gui, open_file
from flika import *

flika_version = flika.__version__
if StrictVersion(flika_version) < StrictVersion('0.2.23'):
    from flika.process.BaseProcess import BaseProcess, WindowSelector, SliderLabel, CheckBox
else:
    from flika.utils.BaseProcess import BaseProcess, WindowSelector, SliderLabel, CheckBox


halt_current_computation = False

def launch_docs():
    url='https://github.com/kyleellefsen/pynsight'
    QDesktopServices.openUrl(QUrl(url))



def Export_pts_from_MotilityTracking():
    tracks = g.m.trackPlot.all_tracks
    t_out = []
    x_out = []
    y_out = []
    for i in np.arange(len(tracks)):
        track = tracks[i]
        t_out.extend(track['frames'])
        x_out.extend(track['x_cor'])
        y_out.extend(track['y_cor'])
    p_out = np.array([t_out, x_out, y_out]).T
    filename = r'C:\Users\kyle\Desktop\trial8_pts.txt'
    np.savetxt(filename, p_out)


def getSigma():
    ''' This function isn't complete.  I need to cut out a 20x20 pxl window around large amplitude particles '''
    I = g.m.currentWindow.image
    xorigin = 8
    yorigin = 9
    sigma = 2
    amplitude = 50
    p0 = [xorigin, yorigin, sigma, amplitude]
    p, I_fit, _ = fitGaussian(I, p0)
    xorigin, yorigin, sigma, amplitude = p
    return sigma


def convolve(I, sigma):
    from scipy.signal import convolve2d
    G = generate_gaussian(17, sigma)
    newI = np.zeros_like(I)
    for t in np.arange(len(I)):
        print(t)
        newI[t] = convolve2d(I[t], G, mode='same', boundary='fill', fillvalue=0)
    return newI


def get_points(I):
    import scipy.ndimage
    s = scipy.ndimage.generate_binary_structure(3, 1)
    s[0] = 0
    s[2] = 0
    labeled_array, num_features = scipy.ndimage.measurements.label(I, structure=s)
    objects = scipy.ndimage.measurements.find_objects(labeled_array)

    all_pts = []
    for loc in objects:
        offset = np.array([a.start for a in loc])
        pts = np.argwhere(labeled_array[loc] != 0) + offset
        ts = np.unique(pts[:, 0])
        for t in ts:
            pts_t = pts[pts[:, 0] == t]
            x = np.mean(pts_t[:, 1])
            y = np.mean(pts_t[:, 2])
            all_pts.append([t, x, y])
    all_pts = np.array(all_pts)
    return all_pts


def cutout(pt, Movie, width):
    assert width % 2 == 1  # mx must be odd
    t, x, y = pt
    t=int(t)
    mid = int(np.floor(width / 2))
    x0 = int(x - mid)
    x1 = int(x + mid)
    y0 = int(y - mid)
    y1 = int(y + mid)
    mt, mx, my = Movie.shape
    if y0 < 0: y0 = 0
    if x0 < 0: x0 = 0
    if y1 >= my: y1 = my - 1
    if x1 >= mx: x1 = mx - 1
    corner = [x0, y0]
    I = Movie[t, x0:x1 + 1, y0:y1 + 1]
    return I, corner


def refine_pts(pts, blur_window, sigma, amplitude):
    global halt_current_computation
    if blur_window is None:
        g.alert("Before refining points, you must select a 'blurred window'")
        return None, False
    new_pts = []
    old_frame = -1
    for pt in pts:
        new_frame = int(pt[0])
        if old_frame != new_frame:
            old_frame = new_frame
            blur_window.imageview.setCurrentIndex(old_frame)
            qApp.processEvents()
            if halt_current_computation:
                halt_current_computation = False
                return new_pts, False
        width = 9
        mid = int(np.floor(width / 2))
        I, corner = cutout(pt, blur_window.image, width)
        xorigin = mid
        yorigin = mid
        p0 = [xorigin, yorigin, sigma, amplitude]
        fit_bounds = [(0, 9), (0, 9), (0, 4), (0, 1000)]
        p, I_fit, _ = fitGaussian(I, p0, fit_bounds)
        xfit = p[0] + corner[0]
        yfit = p[1] + corner[1]
        #                t,  old x, old y, new_x, new_y, sigma, amplitude
        new_pts.append([pt[0], pt[1], pt[2], xfit, yfit, p[2], p[3]])
    new_pts = np.array(new_pts)
    return new_pts, True


class Points(object):
    def __init__(self, txy_pts):
        self.frames = np.unique(txy_pts[:, 0]).astype(np.int)
        self.txy_pts = txy_pts
        self.window = None
        self.pathitems = []
        self.pts_by_frame = []
        self.pts_remaining = []
        self.pts_idx_by_frame = []  # this array has the same structure as points_by_array but contains the index of the original txy_pts argument
        self.intensities = [] #GD Edit

    def link_pts(self, maxFramesSkipped, maxDistance):
        print('Linking points')
        self.pts_by_frame = []
        self.pts_remaining = []
        self.pts_idx_by_frame = []  # this array has the same structure as points_by_array but contains the index of the original txy_pts argument
        for frame in np.arange(0, np.max(self.frames) + 1):
            indicies = np.where(self.txy_pts[:, 0] == frame)[0]
            pos = self.txy_pts[indicies, 1:]
            self.pts_by_frame.append(pos)
            self.pts_remaining.append(np.ones(pos.shape[0], dtype=np.bool))
            self.pts_idx_by_frame.append(indicies)

        tracks = []
        for frame in self.frames:
            for pt_idx in np.where(self.pts_remaining[frame])[0]:
                self.pts_remaining[frame][pt_idx] = False
                abs_pt_idx = self.pts_idx_by_frame[frame][pt_idx]
                track = [abs_pt_idx]
                track = self.extend_track(track, maxFramesSkipped, maxDistance)
                tracks.append(track)
        self.tracks = tracks

    def extend_track(self, track, maxFramesSkipped, maxDistance):
        pt = self.txy_pts[track[-1]]
        # pt can move less than two pixels in one frame, two frames can be skipped
        for dt in np.arange(1, maxFramesSkipped+2):
            frame = int(pt[0]) + dt
            if frame >= len(self.pts_remaining):
                return track
            candidates = self.pts_remaining[frame]
            nCandidates = np.count_nonzero(candidates)
            if nCandidates == 0:
                continue
            else:
                distances = np.sqrt(np.sum((self.pts_by_frame[frame][candidates] - pt[1:]) ** 2, 1))
            if any(distances < maxDistance):
                next_pt_idx = np.where(candidates)[0][np.argmin(distances)]
                abs_next_pt_idx = self.pts_idx_by_frame[frame][next_pt_idx]
                track.append(abs_next_pt_idx)
                self.pts_remaining[frame][next_pt_idx] = False
                track = self.extend_track(track, maxFramesSkipped, maxDistance)
                return track
        return track

    def get_tracks_by_frame(self):
        tracks_by_frame = [[] for frame in np.arange(np.max(self.frames)+1)]
        for i, track in enumerate(self.tracks):
            frames = self.txy_pts[track][:,0].astype(np.int)
            for frame in frames:
                tracks_by_frame[frame].append(i)
        self.tracks_by_frame = tracks_by_frame

    def clearTracks(self):
        if self.window is not None and not self.window.closed:
            for pathitem in self.pathitems:
                self.window.imageview.view.removeItem(pathitem)
        self.pathitems = []

    def showTracks(self):
        # clear self.pathitems
        self.clearTracks()

        frame = self.window.imageview.currentIndex
        if frame<len(self.tracks_by_frame):
            tracks = self.tracks_by_frame[frame]
            pen = QPen(Qt.green, .4)
            pen.setCosmetic(True)
            pen.setWidth(2)
            for track_idx in tracks:
                pathitem = QGraphicsPathItem(self.window.imageview.view)
                pathitem.setPen(pen)
                self.window.imageview.view.addItem(pathitem)
                self.pathitems.append(pathitem)
                pts = self.txy_pts[self.tracks[track_idx]]
                x = pts[:, 1]+.5; y = pts[:,2]+.5
                path = QPainterPath(QPointF(x[0],y[0]))
                for i in np.arange(1, len(pts)):
                    path.lineTo(QPointF(x[i],y[i]))
                pathitem.setPath(path)

#### GD EDITS ####

class Tracks(object):
    def __init__(self):
        self.txy_pts = []
        self.points = None
        self.window = None
        self.intensities = []
        self.dataArray = None

    def loadtracksjson(self, fileString=''):
        if fileString != '':
            filename = fileString
        else:
            filename = open_file_gui("Open tracks from json", filetypes='*.json')
            if filename is None:
                return None
        obj_text = codecs.open(filename, 'r', encoding='utf-8').read()
        pts = json.loads(obj_text)
        txy_pts = np.array(pts['txy_pts'])
        self.txy_pts = txy_pts
        points = Points(txy_pts)
        points.tracks = pts['tracks']
        points.get_tracks_by_frame()
        self.points = points
        nTracks = len(points.tracks)

    def savetracks(self):
        tracks = self.points.tracks
        if isinstance(tracks[0][0], np.int64):
            tracks = [[np.asscalar(a) for a in b] for b in tracks]
        txy_pts = self.points.txy_pts.tolist()
        txy_intensities = self.points.intensities
        filename = save_file_gui("Save all tracks as json  (also saves linked tracks as csv)", filetypes='*.json')
        base = os.path.splitext(filename)[0]
        filename_csv = base + '.csv'
        pts = {'tracks': tracks, 'txy_pts': txy_pts, 'txy_intensities': txy_intensities}
        json.dump(pts, codecs.open(filename, 'w', encoding='utf-8'), separators=(',', ':'), sort_keys=True, indent=4)  ### this saves the array in .json format
        self.savetrackscsv(filename_csv)

    def savetrackscsv(self,filename):
        tracks = self.points.tracks
        if isinstance(tracks[0][0], np.int64):
            tracks = [[np.asscalar(a) for a in b] for b in tracks]
        txy_pts = self.points.txy_pts.tolist()
        txy_intensities = self.points.intensities
        #filter tracks list to only include linked tracks
        linkedTracks = [item for item in tracks if len(item) > 1]
        #get xy coordinates and intensities for linked tracks
        trackNumber = []
        txy_ptsByTrack = []
        txy_intensitiesByTrack = []
        for i, indices in enumerate(linkedTracks):
            trackNumber.append( i )
            txy_ptsByTrack.append(list(txy_pts[j] for j in indices))
            txy_intensitiesByTrack.append(list(txy_intensities[k] for k in indices))

        #make dataframe of tracks, xy coordianates and intensities for linked tracks
        dict = {'track_number': trackNumber, 'pts': txy_ptsByTrack, 'intensities': txy_intensitiesByTrack}
        self.linkedtrack_DF = pd.DataFrame(dict)
        #save df as csv
        self.linkedtrack_DF.to_csv(filename)
        
        return


    def setWindow(self, window, display=True):
        self.window = window
        self.points.window = window
        self.dataArray = self.window.imageArray()
        #display tracks
        if display:
            self.points.window.sigTimeChanged.connect(self.points.showTracks)

    def getIntensities(self):
        #intensities retrieved from image stack using point data (converted from floats to ints)
        
        n, w, h = self.dataArray.shape
        
        #clear intensity list
        self.intensities = []
        
        for point in self.txy_pts:
            frame = int(round(point[0]))
            x = int(round(point[1]))
            y = int(round(point[2]))
            
            #set x,y bounds for 3x3 pixel square
            xMin = x - 1
            xMax = x + 2
            
            yMin = y - 1
            yMax = y + 2
            
            #deal with edge cases
            if xMin < 0:
                xMin = 0
            if xMax > w:
                xMax = w
                
            if yMin <0:
                yMin = 0
            if yMax > h:
                yMax = h
            
            #get mean pixels values for 3x3 square
            self.intensities.append((np.mean(self.dataArray[frame][xMin:xMax,yMin:yMax])))
        
        #update points instance
        self.points.intensities = self.intensities
        

        
tracks = Tracks()


if __name__ == '__main__':
    start_flika() 
    #load particle image stack that you want to get intensities from
    imgStack_filename = open_file_gui()
    #data_window = open_file(r"C:\Users\g_dic\OneDrive\Desktop\testing\particles.tif")
    data_window = open_file(imgStack_filename)
    #load tracks from JSON file
    tracks.loadtracksjson()
    #set window to get intensities from
    tracks.setWindow(data_window)
    #get intensities from point positions
    tracks.getIntensities()
    #save new json file including point intensities + csv of linked track values
    tracks.savetracks()
    
    #load json
    #tracks.loadtracksjson(r"C:\Users\g_dic\OneDrive\Desktop\AL_55_2020-06-29-TIRFM_Diff_tdt-MEFs_B_4_trackintensities.json")   
    #save csv file  
    #tracks.savetrackscsv(r"C:\Users\g_dic\OneDrive\Desktop\AL_55_2020-06-29-TIRFM_Diff_tdt-MEFs_B_4_trackintensities.csv")    
    
