
import copy
import sys

import numpy as num
from wx import AboutDialogInfo

from version import DEBUG, __version__

class ShapeParams:
    def __init__(self,major=0,minor=0,area=0,ecc=0):
        self.major = major
        self.minor = minor
        self.area = area
        self.ecc = ecc

    def copy(self):
        return ShapeParams(self.major,self.minor,self.area,self.ecc)

    def __print__(self):
        return 'major = %.2f, minor = %.2f, area = %.2f, ecc = %.2f'%(self.major,self.minor,self.area,self.ecc)

    def __repr__( self ): return self.__print__()

    def __str__( self ): return self.__print__()

    def __eq__(self,other):
        for i,j in self.__dict__.iteritems():
            if not hasattr(other,i):
                return False
            if not (j == other.__dict__[i]):
                return False
        for i,j in other.__dict__.iteritems():
            if not hasattr(self,i):
                return False
            if not (j == self.__dict__[i]):
                return False
        return True

def averageshape(shape1,shape2):
    shape3 = ShapeParams()
    shape3.major = (shape1.major+shape2.major)/2.
    shape3.minor = (shape1.minor+shape2.minor)/2.
    shape3.area = (shape1.area+shape2.area)/2.
    shape3.ecc = (shape1.ecc+shape2.ecc)/2.
    return shape3

class Grid:
    def __init__(self):
        self.X = None
        self.Y = None
        self.X2 = None
        self.Y2 = None
        self.XY = None
    def setsize(self,sz):
        [self.Y,self.X] = num.mgrid[0:sz[0],0:sz[1]]
        self.X = self.X.astype( num.int64 )
        self.Y = self.Y.astype( num.int64 )
        self.Y2 = self.Y**2
        self.X2 = self.X**2
        self.XY = self.X*self.Y
    def __eq__(self,other):
        return True

class Parameters:
    def __init__(self):

        self.DOBREAK = False

        # set constants

        # for timing various parts of the code
        self.last_time = 0

        # default frame rate
        self.DEFAULT_FRAME_RATE = 25.

        # for fast computation of weighted region props
        self.GRID = Grid()

        # max frac above-threshold points in a frame
        self.max_n_points_ratio = 1./250.
        # for displaying zoom windows
        self.id_spinner_width = 50
        # palettes
        self.normal_palette = [[255,0,0],     # red
                               [0,255,0],     # green
                               [0,0,255],     # blue
                               [255,0,255],   # magenta
                               [255,255,0],   # yellow
                               [0,255,255],   # cyan
                               [255,127,127], # light red
                               [127,255,127], # light green
                               [127,127,255], # light blue
                               [255,127,255], # light magenta
                               [255,255,127], # light yellow
                               [127,255,255]] # light cyan
        # colorblind-friendly palette from
        # http://jfly.iam.u-tokyo.ac.jp/color/index.html
        self.colorblind_palette = [[230,159,0],   # orange
                                   [86,180,233],  # sky blue
                                   [0,158,115],   # blue-green
                                   [240,228,66],  # yellow
                                   [0,114,178],   # blue
                                   [213,94,0],    # vermillion
                                   [204,121,167]] # red-purple

        self.zoom_drag_rectangle_color = (1, 0.6, 0, 1)

        # background constants
        self.SHOW_RAW_FRAME = -1
        self.SHOW_BACKGROUND = 0
        self.SHOW_DISTANCE = 1
        self.SHOW_THRESH = 2
        self.SHOW_NONFORE = 3
        self.SHOW_DEV = 4
        self.SHOW_CC = 5
        self.SHOW_ELLIPSES = 6
        self.SHOW_EXPBGFGMODEL_LLR = 7
        self.SHOW_EXPBGFGMODEL_ISBACK = 8
        self.SHOW_EXPBGFGMODEL_BGMU = 9
        self.SHOW_EXPBGFGMODEL_FGMU = 10
        self.SHOW_EXPBGFGMODEL_BGSIGMA = 11
        self.SHOW_EXPBGFGMODEL_FGSIGMA = 12
        self.SHOW_EXPBGFGMODEL_FRACFRAMESISBACK = 13
        self.SHOW_EXPBGFGMODEL_MISSINGDATA = 14
        
        self.BG_SHOW_STRINGS = ['Background Image',
                                'Distance from Background',
                                'Foreground/Background Classification',
                                'Background-Only Areas',
                                'Normalization Image',
                                'Connected Components',
                                'Ellipse Fits']
        self.EXPBGFGMODEL_SHOW_STRINGS = ['Prior Log-Likelihood Ratio',
                                          'Use in Background Model',
                                          'Prior Background Mean Px Intensity',
                                          'Prior Foreground Mean Px Intensity',
                                          'Prior Background Std Px Intensity',
                                          'Prior Foreground Std Px Intensity',
                                          'Frac Frames in Bg Model',
                                          'Bg Missing Data']
        self.SHOW_EXPBGFGMODEL = [self.SHOW_EXPBGFGMODEL_LLR,self.SHOW_EXPBGFGMODEL_ISBACK,
                                  self.SHOW_EXPBGFGMODEL_BGMU,self.SHOW_EXPBGFGMODEL_FGMU,
                                  self.SHOW_EXPBGFGMODEL_BGSIGMA,self.SHOW_EXPBGFGMODEL_FGSIGMA,
                                  self.SHOW_EXPBGFGMODEL_FRACFRAMESISBACK,self.SHOW_EXPBGFGMODEL_MISSINGDATA]
        
        self.ALGORITHM_MEDIAN = 0
        self.ALGORITHM_MEAN = 1

        # display stuff
        self.status_green = "#66FF66"
        self.status_blue = "#AAAAFF"
        self.status_red = "#FF6666"
        self.status_yellow = "#FFFF66"
        self.wxvt_bg = '#DDFFDD'

        self.MAXDSHOWINFO = 10 # maximum allowed click distance from selection
        self.DRAW_MOTION_SCALE = 10.

        ## User interaction mode ###
        self.interactive = True # whether Ctrax was started in interactive mode or not
        self.noninteractive_resume_tracking = False # in non-interactive mode, if we read in trajectories, should we resume tracking
        self.feedback_enabled = True # whether user feedback is currently enabled
        self.app_instance = None # pointer to the global Application

        # set parameters to default values
        self.start_frame = 0

        # Number of identities assigned so far
        self.nids = long( 0 )

        self.version = 0

        ## Movie Parameters ###

        # number of frames in the movie
        self.n_frames = 0
        # size of a frame
        self.movie_size = (0,0)
        # number of pixels in a frame
        self.npixels = 0
        # movie object
        self.movie = None
        self.movie_name = ''
        self.annotation_movie_name = ''

        self.movie_index_transpose = True

        ## Background Estimation Parameters ###

        # maximum number of pixels (should = bytes?) to allocate for temporary storage while computing the bg median
        # equivalent to 100 x (480x640) images
        self.bg_median_maxbytesallocate = 30720000
        # Background Subtraction Parameters

        # homomorphic filtering constants
        # these defaults are based on a painstaking and probably
        # inaccurate parameterization, but they're a good starting point
        self.hm_cutoff = 0.35
        self.hm_boost = 2
        self.hm_order = 2
        # number of standard deviations to threshold background at
        self.n_bg_std_thresh = 20.
        # do hysteresis
        self.n_bg_std_thresh_low = 10.
        # minimum number of standard deviations for background
        self.bg_std_min = 1.
        self.bg_std_max = 10.
        # if background intensity is greater than min_nonarena, then it
        # is not an area flies can be in, so don't allow foreground in
        # these areas
        self.min_nonarena = 256.
        # if background intensity is less than max_nonarena, then it
        # is not an area flies can be in, so don't allow foreground in
        # these areas
        self.max_nonarena = -1.
        
        # regions of interest
        self.roipolygons = []
        
        # location of arena
        self.arena_center_x = None
        self.arena_center_y = None
        self.arena_radius = None
        self.arena_edgethresh = None
        self.do_set_circular_arena = True

        # search space for arena circle
        self.min_arena_center_x = .4
        self.max_arena_center_x = .6
        self.min_arena_center_y = .4
        self.max_arena_center_y = .6
        self.min_arena_radius = .25
        self.max_arena_radius = .5

        # batch processing & auto-detecting arena
        self.batch_autodetect_arena = True
        self.batch_autodetect_shape = True
        self.batch_autodetect_bg_model = True

        # morphology
        self.do_use_morphology = False
        self.opening_radius = 0
        self.closing_radius = 0

        ## Shape Parameters ###
        
        # upper bounds on shape
        self.maxshape = ShapeParams(9999.,9999.,9999.,1.)
        # lower bounds on shape
        self.minshape = ShapeParams(1.,1.,1.,0.)
        self.meanshape = ShapeParams(2.64,3.56,40.25,0.98)
        # foreground background thresh set to minbackthresh when trying to increase target area
        self.minbackthresh = 1.
        # maximum number of clusters to split a foreground connected component into during the forward pass
        self.maxclustersperblob = 5
        # maximum penalty for merging together two ccs
        self.maxpenaltymerge = 40
        # maximum area of deleted target
        self.maxareadelete = 5
        # minimum area of ignored connected components
        self.minareaignore = 2500
        # max n objects per frame to return
        self.max_n_clusters = 100 
        # number of frames used to compute shape bounds
        self.n_frames_size = 50
        # number of standard deviations from mean for upper and lower shape bounds
        self.n_std_thresh = 4.

        ## Motion model parameters ###
        
        # weight of angle in distance measure
        self.ang_dist_wt = 100.
        # maximum distance a fly can move between frames
        self.max_jump = 100.
        # KB 20120109: maximum distance a fly can move between frames and end up within a split connected component
        self.max_jump_split = -1.
        # minimum distance to be called a jump
        self.min_jump = 50.
        # dampening constant
        self.dampen = 0. # weighting term used in cvpred()
        self.angle_dampen = 0.5 

        # weight of velocity angle in choosing orientation mod 2pi
        self.velocity_angle_weight = .05
        self.max_velocity_angle_weight = .25

        ## Fix errors parameters ###
        
        self.do_fix_split = True
        self.do_fix_merged = True
        self.do_fix_spurious = True
        self.do_fix_lost = True
        self.lostdetection_length = 50
        self.lostdetection_distance = 100.
        self.spuriousdetection_length = 50
        self.mergeddetection_distance = 20.
        self.mergeddetection_length = 50
        self.splitdetection_cost = 40.
        self.splitdetection_length = 50
        self.maxdcentersextra = 3.
        self.bigboundingboxextra = 2.
        # maximum number of frames to buffer for faster hindsight
        self.maxnframesbuffer = 100

        ## Drawing/Display parameters ###
        # thickness of lines of drawn ellipses
        if 'darwin' in sys.platform:
            self.ellipse_thickness = 1
        else:
            self.ellipse_thickness = 2
        # palette of colors to use
        self.use_colorblind_palette = False
        # colors
        self.colors = self.normal_palette

        self.tail_length = 10

        self.status_box = 0
        self.file_box = 1
        self.file_box_max_width = 40

        # computing prior background/foreground models

        # number of frames to sample per video
        self.prior_nframessample = 100

        self.expbgfgmodel_filename = None
        self.use_expbgfgmodel = False
        
        # threshold for log-likelihood ratio of foreground over background
        # if llr is >= thresh_low and near a pixel > thresh then don't 
        # include in the background model
        self.expbgfgmodel_llr_thresh = 0
        self.expbgfgmodel_llr_thresh_low = 0
        
        # fraction of sampled frames that we require to be classified as background according
        # to the prior expbgfgmodel to trust our estimates. otherwise, we will fill missing sections
        # either with the prior mean and standard deviation or by interpolation 
        self.min_frac_frames_isback = .1
        self.expbgfgmodel_fill = 'Interpolation'
        self.EXPBGFGMODEL_FILL_STRINGS = ['Prior BG Model','Interpolation']

        # number of foreground pixels to sample per location
        self.prior_fg_nsamples_pool = 25
        # number of background pixels to sample per location
        self.prior_bg_nsamples_pool = 25

        # how much to increase the sample radius per iteration
        self.prior_fg_pool_radius_factor = 1.25
        # how much to increase the sample radius per iteration
        self.prior_bg_pool_radius_factor = 1.25

        # whether to use our uncompressed AVI reader
        self.use_uncompressed_avi = True

    def enable_feedback( self, now_enabled ):
        """Change enablement of GUI/user feedback, dependent on interactivity mode.""" 
        self.feedback_enabled = now_enabled and self.interactive

    def __print__(self):
        s = ""
        for i,j in self.__dict__.iteritems():
            if j is None:
                s += i + ": None\n"
            else:
                s += i + ": " + str(j) + "\n"
        return s

    def __repr__( self ): return self.__print__()

    def __str__( self ): return self.__print__()

    def copy(self):
        v = Parameters()
        for i,j in self.__dict__.iteritems():
            if i == 'movie':
                continue
            try:
                v.__dict__[i] = copy.deepcopy(j)
            except:
                v.__dict__[i] = j
        return v

    def __eq__( self, other ):
        for i,j in self.__dict__.iteritems():
            if not hasattr(other,i):
                return False
            if not (j == other.__dict__[i]):
                return False
        for i,j in other.__dict__.iteritems():
            if not hasattr(self,i):
                return False
            if not (j == self.__dict__[i]):
                return False
        return True
        
params = Parameters()

diagnostics = dict()
diagnostics['nbirths_nohindsight'] = 0
diagnostics['ndeaths_nohindsight'] = 0
diagnostics['ndeaths_notfixed'] = 0
diagnostics['nbirths_notfixed'] = 0
diagnostics['nsplits_fixed'] = 0
diagnostics['nspurious_fixed'] = 0
diagnostics['nmerged_fixed'] = 0
diagnostics['nlost_fixed'] = 0
diagnostics['nhindsight_fixed'] = 0
diagnostics['nlarge_notfixed'] = 0
diagnostics['nsmall_notfixed'] = 0
diagnostics['nlarge_split'] = 0
diagnostics['nsmall_merged'] = 0
diagnostics['nsmall_lowerthresh'] = 0
diagnostics['nsmall_deleted'] = 0
diagnostics['max_nsplit'] = 0
diagnostics['sum_nsplit'] = 0
diagnostics['nlarge_ignored'] = 0
diagnostics['nframes_analyzed'] = 0

class GUIConstants:
    def __init__( self ):
        self.info = AboutDialogInfo()
        self.info.SetName( "Ctrax" )
        self.info.SetVersion( __version__ )
        self.info.SetCopyright( "2007-2013, Caltech ethomics project" )
        self.info.SetDescription( """The Caltech Multiple Fly Tracker
Kristin Branson et al.

http://ctrax.sourceforge.net/
http://dx.doi.org/10.1038/nmeth.1328

Distributed under the GNU General Public License
(http://www.gnu.org/licenses/gpl.html) with
ABSOLUTELY NO WARRANTY.

This project has been supported by the NIH and
the HHMI.""" )
        
        self.TRACK_START = "Start Tracking"
        self.TRACK_STOP = "Stop Tracking"
        self.TRACK_PLAY = "Start Playback"
        self.TRACK_PAUSE = "Stop Playback"

        params.version = __version__
# rather than global variables for each of these...

const = GUIConstants()
