'''Independent-running GSAS-II based auto-integration program with minimal
GUI, no visualization but intended to implement significant levels of
parallelization.
'''
# Autointegration from
# $Id: GSASIIimgGUI.py 3926 2019-04-23 18:11:07Z toby $
# hacked for stand-alone use
#
# idea: select image file type & set filter from that
#
from __future__ import division, print_function
import os
import copy
import glob
import time
import re
import math
import sys
import wx
import wx.lib.mixins.listctrl  as  listmix
import wx.grid as wg
import numpy as np
import GSASIIpath
GSASIIpath.SetBinaryPath(True)
GSASIIpath.SetVersionNumber("$Revision: $")
import GSASIIIO as G2IO
import GSASIIctrlGUI as G2G
import GSASIIobj as G2obj
import GSASIIpy3 as G2py3
import GSASIIimgGUI as G2imG
import GSASIIfiles as G2fil
import GSASIIscriptable as G2sc
import multiprocessing as mp

class AutoIntFrame(wx.Frame):
    '''Creates a wx.Frame window for the Image AutoIntegration.
    The intent is that this will be used as a non-modal dialog window.

    Implements a Start button that morphs into a pause and resume button.
    This button starts a processing loop that is repeated every
    :meth:`PollTime` seconds.

    :param wx.Frame G2frame: main GSAS-II frame
    :param float PollTime: frequency in seconds to repeat calling the
      processing loop. (Default is 30.0 seconds.)
    '''
    def __init__(self,G2frame,PollTime=30.0):
        def OnStart(event):
            '''Called when the start button is pressed. Changes button label
            to Pause. When Pause is pressed the label changes to Resume.
            When either Start or Resume is pressed, the processing loop
            is started. When Pause is pressed, the loop is stopped.
            '''
            self.Pause = False
            # change button label
            if self.btnstart.GetLabel() != 'Pause':
                self.btnstart.SetLabel('Pause')
                self.Status.SetStatusText('Press Pause to delay integration or Reset to prepare to reintegrate all images')
                if self.timer.IsRunning(): self.timer.Stop()
                self.PreventTimerReEntry = False
                if self.StartLoop():
                    G2G.G2MessageBox(self,'Error in setting up integration. See console')
                    return
                # delete pool of slaved interpreters in case input has changed
                if self.MPpool:
                    self.MPpool.terminate()
                    self.MPpool = None
#                if self.ncores >= 1: # debug
                if self.ncores > 1:
#                    print('Creating pool of ',self.ncores,'cores') # debug
                    self.MPpool = mp.Pool(self.ncores)
                self.OnTimerLoop(None) # run once immediately
                if not self.Pause:
                    # no pause, so start timer to check for new files
                    self.timer.Start(int(1000*PollTime),oneShot=False)
                    return
            # we will get to this point if Paused
            self.OnPause()

        def OnReset(event):
            '''Called when Reset button is pressed. This stops the
            processing loop and resets the list of integrated files so
            all images can be reintegrated.
            '''
            self.btnstart.SetLabel('Restart')
            self.Status.SetStatusText('Press Restart to reload and re-integrate images matching filter')
            if self.timer.IsRunning(): self.timer.Stop()
            self.Reset = True
            self.Pause = True
            self.ProcessedList = []
            self.ShowMatchingFiles(None)

        def OnQuit(event):
            '''Stop the processing loop and close the Frame
            '''
            if self.timer.IsRunning(): self.timer.Stop() # make sure we stop first
            wx.CallAfter(self.Destroy)

        def OnBrowse(event):
            '''Responds when the Browse button is pressed to load a file.
            The routine determines which button was pressed and gets the
            appropriate file type and loads it into the appropriate place
            in the dict.
            '''
            if btn3 == event.GetEventObject():
                dlg = wx.DirDialog(
                    self, 'Select directory for output files',
                    self.params['outdir'],wx.DD_DEFAULT_STYLE)
                dlg.CenterOnParent()
                try:
                    if dlg.ShowModal() == wx.ID_OK:
                        self.params['outdir'] = dlg.GetPath()
                        fInp3.SetValue(self.params['outdir'])
                finally:
                    dlg.Destroy()
                return

        def showPDFctrls(event):
            '''Called to show or hide AutoPDF widgets. Note that TextCtrl's
            must be included in the sizer layout with .Show(True) before
            .Show(False) will work properly.
            '''
            TestInput()
            if len(self.pdfList) == 0 or self.params['TableMode']:
                lbl4b.SetValue(False)
                self.params['ComputePDF'] = False
                lbl4b.Enable(False)
            else:
                lbl4b.Enable(True)

            for w in [self.pbkg[i][4] for i in (0,1,2)]:
                w.Enable(self.params['ComputePDF'])
                w.Show(self.params['ComputePDF'])
            #for o in pdfwidgets+[self.pbkg[i][2] for i in (0,1,2)]:
            for o in pdfwidgets+[self.pdfSel]+[self.pbkg[i][1] for i in range(3)]:
                o.Enable(self.params['ComputePDF'])
            if self.params['ComputePDF']:
                c = "black"
            else:
                c = "gray"
            for l in ([lbl4,lbl4a,lbl5,lbl5a,lbl5b]
                          +[self.pbkg[i][0] for i in (0,1,2)]
                          +[self.pbkg[i][5] for i in (0,1,2)]):
                l.SetForegroundColour(c)
            checkPDFselection()
            self.SendSizeEvent()

        def GetGPXInputFile(event):
            'Get and read input from input GPX file'
            pth = self.params['readdir']
            lbl = 'Select a project input file'
            filtyp = 'project file (.gpx)|*.gpx'
            if event is not None:
                dlg = wx.FileDialog(self,lbl, pth,
                        style=wx.FD_OPEN,wildcard=filtyp)
                dlg.CenterOnParent()
                if dlg.ShowModal() == wx.ID_OK:
                    self.gpxin[3] = dlg.GetPath()
                dlg.Destroy()
            if not os.path.exists(self.gpxin[3]):
                G2G.G2MessageBox(self,'Error: file {} not found.'.format(self.gpxin[3]))
                return
            SetGPXInputFile()
        def SetGPXInputFile():
            gpx = G2sc.G2Project(self.gpxin[3])
            self.imgList = gpx.images()
            self.histList = gpx.histograms()
            self.pdfList = gpx.pdfs()
            if not self.imgList:
                G2G.G2MessageBox(self,'Error: no images in {}.'.format(self.gpxin[3]))
                return
            self.gpxin[1].SetValue(self.gpxin[3])
            self.imprm[1].Clear()
            self.imprm[1].AppendItems([i.name for i in self.imgList])
            if len(self.imgList) == 1:
                self.imprm[1].SetSelection(0)
            self.maskfl[1].Clear()
            self.maskfl[1].AppendItems(['']+[i.name for i in self.imgList])
            for i in range(3):
                self.pbkg[i][1].Clear()
                self.pbkg[i][1].AppendItems(['']+[i.name for i in self.histList])
            self.pdfSel.Clear()
            self.pdfSel.AppendItems([i.name for i in self.pdfList])
            showPDFctrls(None)

        def TestInput(*args,**kwargs):
            'Determine if the start button should be enabled'
            writingSomething = False
            writingPDF = False
            # is at least one output file selected?
            for dfmt in self.fmtlist:
                fmt = dfmt[1:]
                if self.params['outsel'][fmt]:
                    writingSomething = True
                    break
            if self.params['ComputePDF']:
                for fmt in self.PDFformats:
                    if self.params['outsel'][fmt]:
                        writingPDF = writingSomething = True
                        break
            if not writingSomething:
                self.EnableIntButtons(False)
                return
            # do we have integration input?
            if not self.params['TableMode']:
                if not self.gpxin[3]:
                    self.EnableIntButtons(False)
                    return
                elif not os.path.exists(self.gpxin[3]):
                    self.EnableIntButtons(False)
                    return
                elif len(self.imgList) == 0:
                    self.EnableIntButtons(False)
                    return
                else:
                    self.EnableIntButtons(True)
            else:
                if self.ImgTblParms:
                    self.EnableIntButtons(True)
                else:
                    self.EnableIntButtons(False)
                    return
            # do we have PDF input, if requested
            if self.params['ComputePDF']:
                if len(self.pdfList) == 0 or not writingPDF:
                    self.EnableIntButtons(False)
                elif 'Error' in self.formula:
                    self.EnableIntButtons(False)

        def checkPDFselection():
            'Read PDF entry from input GPX file & show in GUI'
            pdfEntry = self.pdfSel.GetStringSelection()
            if not self.pdfList:
                self.params['ComputePDF'] = False
                lbl4b.Enable(False)
                return
            else:
                lbl4b.Enable(True)
            if pdfEntry not in [i.name for i in self.pdfList]:
                # strange something -- not selected
                self.pdfSel.SetSelection(0)
                pdfEntry = self.pdfSel.GetStringSelection()
            if self.gpxInp is None or self.gpxInp.filename != self.gpxin[3]:
                self.gpxInp = G2sc.G2Project(self.gpxin[3])
            try:
                PDFobj = self.gpxInp.pdf(pdfEntry)
            except KeyError:
                print("PDF entry not found: {}".format(pdfEntry))
                return
            histNames = [i.name for i in self.histList]
            for i,lbl in enumerate(('Sample Bkg.','Container',
                                    'Container Bkg.')):
                self.pbkg[i][4].SetValue(str(PDFobj.data['PDF Controls'][lbl]['Mult']))
                self.pbkg[i][6] = PDFobj.data['PDF Controls'][lbl]['Mult']
                try:
                    i = 1 + histNames.index(PDFobj.data['PDF Controls'][lbl]['Name'])
                    self.pbkg[i][1].SetSelection(i)
                except ValueError:
                    i = 0
                    self.pbkg[i][1].SetSelection(0)
                    if PDFobj.data['PDF Controls'][lbl]['Name']:
                        print('PDF {} hist entry {} not found'.format(
                            lbl,PDFobj.data['PDF Controls'][lbl]['Name']))
                        PDFobj.data['PDF Controls'][lbl]['Name'] = ''
            self.formula = ''
            for el in PDFobj.data['PDF Controls']['ElList']:
                i = PDFobj.data['PDF Controls']['ElList'][el]['FormulaNo']
                if i <= 0:
                    continue
                elif i == 1:
                    if self.formula: self.formula += ' '
                    self.formula += '{}'.format(el)
                else:
                    if self.formula: self.formula += ' '
                    self.formula += '{}({:.1f})'.format(el,i)
            if not self.formula:
                self.formula = 'Error: no chemical formula'
            lbl5b.SetLabel(self.formula)
            TestInput()

        def ShowbyMode():
            'create table or non-table integration section of GUI'
            intPrmSizer.Clear(True)
            if not self.params['TableMode']:
                sizer = wx.BoxSizer(wx.HORIZONTAL)
                self.gpxin[0] = wx.StaticText(mnpnl, wx.ID_ANY,'Project (gpx) file:')
                sizer.Add(self.gpxin[0])
                self.gpxin[1] = G2G.ValidatedTxtCtrl(mnpnl,self.gpxin,3,
                                OKcontrol=TestInput,OnLeave=TestInput)
                sizer.Add(self.gpxin[1],1,wx.EXPAND,1)
                self.gpxin[2] = wx.Button(mnpnl, wx.ID_ANY, "Browse")
                self.gpxin[2].Bind(wx.EVT_BUTTON, GetGPXInputFile)
                sizer.Add(self.gpxin[2],0,wx.ALIGN_RIGHT|wx.ALIGN_CENTER_VERTICAL)
                intPrmSizer.Add(sizer,0,wx.EXPAND,0)

                sizer = wx.BoxSizer(wx.HORIZONTAL)
                self.imprm[0] = wx.StaticText(mnpnl, wx.ID_ANY,'Image parms from:')
                sizer.Add(self.imprm[0])
                self.imprm[3] = 0
                self.imprm[1] = G2G.G2ChoiceButton(mnpnl,[''],self.imprm,3,onChoice=TestInput)
                sizer.Add(self.imprm[1],1,wx.EXPAND,1)
                intPrmSizer.Add(sizer,0,wx.EXPAND,0)

                sizer = wx.BoxSizer(wx.HORIZONTAL)
                self.maskfl[0] = wx.StaticText(mnpnl, wx.ID_ANY,'Mask parms from:')
                sizer.Add(self.maskfl[0])
                self.maskfl[3] = 0
                self.maskfl[1] = G2G.G2ChoiceButton(mnpnl,[''],self.maskfl,3,onChoice=TestInput)
                sizer.Add(self.maskfl[1],1,wx.EXPAND,1)
                intPrmSizer.Add(sizer,0,wx.EXPAND,0)
            else:
                sizer = wx.BoxSizer(wx.HORIZONTAL)
                self.table = [None,None,None,None]
                self.table[0] = wx.Button(mnpnl,  wx.ID_ANY, "Create table")
                sizer.Add(self.table[0],0,wx.ALIGN_LEFT|wx.ALL,5)
                self.table[0].Bind(wx.EVT_BUTTON, OnTableButton)
                self.table[1] = wx.Button(mnpnl,  wx.ID_ANY, "Read table")
                sizer.Add(self.table[1],0,wx.ALIGN_LEFT|wx.ALL,5)
                self.table[1].Bind(wx.EVT_BUTTON, OnTableButton)
                self.table[2] = wx.Button(mnpnl,  wx.ID_ANY, "Edit table")
                sizer.Add(self.table[2],0,wx.ALIGN_LEFT|wx.ALL,5)
                self.table[2].Bind(wx.EVT_BUTTON, OnTableButton)
                #self.table[3] = wx.Button(mnpnl,  wx.ID_ANY, "Save table")
                #sizer.Add(self.table[3],0,wx.ALIGN_LEFT|wx.ALL,5)
                #self.table[3].Bind(wx.EVT_BUTTON, OnTableButton)
                intPrmSizer.Add(sizer,0,wx.EXPAND,0)
            # enable/disable based on status of files/table
            TestInput()
            mnsizer.Fit(self)

        def OnTableButton(event):
            '''Called to edit/create the distance-dependent parameter look-up table.
            '''
            pth = self.params['readdir']
            readFileList = []
            parms,fileList = [], []
            if event.GetEventObject() == self.table[0]:
                dlg = wx.FileDialog(self, 'Build new table by selecting image control files', pth,
                    style=wx.FD_OPEN| wx.FD_MULTIPLE,
                    wildcard='Image control files (.imctrl)|*.imctrl')
                dlg.CenterOnParent()
                if dlg.ShowModal() == wx.ID_OK:
                    readFileList = dlg.GetPaths()
                dlg.Destroy()
                if len(readFileList) <= 0: return
            elif event.GetEventObject() == self.table[1]:
                dlg = wx.FileDialog(self, 'Reload table by selecting saved file', pth,
                    style=wx.FD_OPEN,
                    wildcard='Integration table (*.imtbl)|*.imtbl')
                dlg.CenterOnParent()
                if dlg.ShowModal() == wx.ID_OK:
                    readFileList = [dlg.GetPath()]
                dlg.Destroy()
                if len(readFileList) <= 0: return
            elif event.GetEventObject() == self.table[2]:
                parms = copy.deepcopy(self.ImgTblParms)
                fileList = copy.copy(self.IMfileList)
                if not parms:
                    G2G.G2MessageBox(self,'Create or Read table first')
                    return
            dlg = None
            try:
                dlg = G2imG.IntegParmTable(self,parms,fileList,readFileList)
                dlg.CenterOnParent()
                if dlg.ShowModal() == wx.ID_OK:
                    self.params['InterVals'] = SetupInterpolation(dlg)
                    self.ImgTblParms = dlg.ReadImageParmTable()
                    self.IMfileList = dlg.IMfileList
                    self.params['TableMode'] = True
                    self.params['ControlsTable'] = {}
                    self.params['MaskTable'] = {}
                    for f,m in zip(self.IMfileList,self.ImgTblParms[-1]):
                        n = os.path.split(f)[1]
                        if n in self.params['ControlsTable']:
                            print('Warning overwriting entry {}'.format(n))
                        self.params['ControlsTable'][n] = G2imG.ReadControls(f)
                        if m and os.path.exists(m):
                            self.params['MaskTable'][n] = G2imG.ReadMask(m)
                        elif m != "(none)":
                            print("Error: Mask file {} not found".format(m))
                else:
                    self.params['TableMode'] = False
                    self.params['ControlsTable'] = {}
                    self.params['MaskTable'] = {}
                    self.imageBase = G2frame.Image
            finally:
                if dlg: dlg.Destroy()
            TestInput()

        ##################################################
        # beginning of __init__ processing
        ##################################################
        self.G2frame = G2frame
        self.ImgTblParms = None
        self.IMfileList = None
        self.params = {}
        self.Reset = False
        self.Pause = False
        self.PreventReEntryShowMatch = False
        self.PreventTimerReEntry = False
        self.params['ControlsTable'] = {}
        self.params['MaskTable'] = {}
        G2sc.LoadG2fil()
        self.fmtlist = G2sc.exportersByExtension.get('powder',{})
        self.fmtlist['.gpx'] = None
        self.timer = wx.Timer()
        self.timer.Bind(wx.EVT_TIMER,self.OnTimerLoop)
        self.imageBase = G2frame.Image
        self.params['ComputePDF'] = False
        self.params['pdfDmax'] = 0.0
        self.params['pdfprm'] = ''
        self.params['optPDF'] = True
        self.params['TableMode'] = False
        self.params['outsel'] = {}
        self.formula = 'Error'
        self.pdfControls = {}
        self.imgList = []
        self.histList = []
        self.pdfList = []
        self.Scale = [1.0,]
        self.ProcessedList = [] # files that have been integrated
        self.currImageList = [] # files that will be integrated
        self.gpxInp = None
        self.gpxin = [None,None,None,'']
        self.imprm = [None,None,None,'']
        self.maskfl = [None,None,None,'']
        self.ncores = mp.cpu_count()//2
        self.MPpool = None
        self.params['readdir'] = os.getcwd()
        self.params['filter'] = '*.tif'
        self.params['outdir'] = os.getcwd()
        self.PDFformats = ('I(Q)', 'S(Q)', 'F(Q)', 'G(r)', 'PDFgui')
        #GSASIIpath.IPyBreak_base()

        wx.Frame.__init__(self, None, title='Automatic Integration',
                          style=wx.DEFAULT_FRAME_STYLE)
        self.Status = self.CreateStatusBar()
        self.Status.SetStatusText('Press Start to load and integrate images matching filter')
        mnpnl = wx.Panel(self)
        mnsizer = wx.BoxSizer(wx.VERTICAL)
        # box for integration controls & masks input
        intSizer = wx.BoxSizer(wx.VERTICAL)
        sizer = wx.BoxSizer(wx.HORIZONTAL)
        sizer.Add((-1,-1),1,wx.EXPAND,1)
        sizer.Add(wx.StaticText(mnpnl, wx.ID_ANY,'Calibration/Integration parameters'))
        sizer.Add((-1,-1),1,wx.EXPAND,1)
        intSizer.Add(sizer,1,wx.EXPAND,1)
        def ontblModeBtn(event):
            if tblModeBtn.GetValue():
                self.params['TableMode'] = True
            else:
                self.params['TableMode'] = False
            ShowbyMode()
        tblModeBtn = wx.CheckBox(mnpnl,label='Use distance lookup table')
        tblModeBtn.SetValue(False)
        tblModeBtn.Bind(wx.EVT_CHECKBOX, ontblModeBtn)
        intSizer.Add(tblModeBtn)
        mnsizer.Add(intSizer,0,wx.EXPAND,0)
        intPrmSizer = wx.BoxSizer(wx.VERTICAL)

        mnsizer.Add(intPrmSizer,0,wx.EXPAND,0)
        # file filter stuff
        mnsizer.Add((-1,15))
        sizer = wx.BoxSizer(wx.HORIZONTAL)
        sizer.Add(wx.StaticText(mnpnl, wx.ID_ANY,'Read images from '))
        self.readDir = G2G.ValidatedTxtCtrl(mnpnl,self.params,'readdir',
                            OnLeave=self.ShowMatchingFiles,size=(200,-1))
        sizer.Add(self.readDir,1,wx.EXPAND,1)
        btn3 = wx.Button(mnpnl, wx.ID_ANY, "Browse")
        btn3.Bind(wx.EVT_BUTTON, self.SetSourceDir)
        sizer.Add(btn3,0,wx.ALIGN_RIGHT|wx.ALIGN_CENTER_VERTICAL)
        mnsizer.Add(sizer,0,wx.EXPAND,0)
        sizer = wx.BoxSizer(wx.HORIZONTAL)
        sizer.Add((-1,-1),1,wx.EXPAND,1)
        sizer.Add(wx.StaticText(mnpnl, wx.ID_ANY,'  Image filter'))
        flterInp = G2G.ValidatedTxtCtrl(mnpnl,self.params,'filter',
                                        OnLeave=self.ShowMatchingFiles)
        sizer.Add(flterInp)
        mnsizer.Add(sizer,0,wx.EXPAND,0)

        self.ListBox = wx.ListBox(mnpnl,size=(-1,100))
        mnsizer.Add(self.ListBox,1,wx.EXPAND,1)
        self.ShowMatchingFiles(None)

        # box for output selections
        mnsizer.Add((-1,15))
        lbl = wx.StaticBox(mnpnl, wx.ID_ANY, "Integration output")
        lblsizr = wx.StaticBoxSizer(lbl, wx.VERTICAL)
        sizer = wx.BoxSizer(wx.HORIZONTAL)
        sizer.Add(wx.StaticText(mnpnl, wx.ID_ANY,'Write to: '),0,wx.ALIGN_CENTER_VERTICAL)
        fInp3 = G2G.ValidatedTxtCtrl(mnpnl,self.params,'outdir',notBlank=False,size=(300,-1))
        sizer.Add(fInp3,1,wx.ALIGN_CENTER_VERTICAL|wx.EXPAND)
        btn3 = wx.Button(mnpnl,  wx.ID_ANY, "Browse")
        btn3.Bind(wx.EVT_BUTTON, OnBrowse)
        sizer.Add(btn3,0,wx.ALIGN_RIGHT|wx.ALIGN_CENTER_VERTICAL)
        lblsizr.Add(sizer,0,wx.EXPAND)
        sizer = wx.BoxSizer(wx.HORIZONTAL)
        sizer.Add(wx.StaticText(mnpnl, wx.ID_ANY,'Select format(s):'))
        for dfmt in self.fmtlist:
            sizer.Add((6,2)) # add a bit of extra space
            fmt = dfmt[1:]
            if fmt not in self.params['outsel']: self.params['outsel'][fmt] = False
            btn = G2G.G2CheckBox(mnpnl,dfmt,self.params['outsel'],fmt,
                                     OnChange=TestInput)
            sizer.Add(btn)
        lblsizr.Add(sizer)
        sizer = wx.BoxSizer(wx.HORIZONTAL)
        sizer.Add(wx.StaticText(mnpnl, wx.ID_ANY,'Separate dir for each format: '))
        self.params['SeparateDir'] = False
        sizer.Add(G2G.G2CheckBox(mnpnl,'',self.params,'SeparateDir'))
        lblsizr.Add(sizer)
        mnsizer.Add(lblsizr,0,wx.ALIGN_CENTER|wx.EXPAND,1)

        mnsizer.Add((-1,15))
        lbl = wx.StaticBox(mnpnl, wx.ID_ANY, "PDF settings")
        pdfwidgets = []
        lblsizr = wx.StaticBoxSizer(lbl, wx.VERTICAL)
        sizer = wx.BoxSizer(wx.HORIZONTAL)
        sizer.Add(wx.StaticText(mnpnl, wx.ID_ANY,'Autocompute PDF:'),0,wx.ALIGN_CENTER_VERTICAL)
        lbl4b = G2G.G2CheckBox(mnpnl,'',self.params,'ComputePDF',
                                     OnChange=showPDFctrls)
        sizer.Add(lbl4b)
        lbl4a = wx.StaticText(mnpnl, wx.ID_ANY,'Max detector distance: ')
        sizer.Add(lbl4a,0,wx.ALIGN_CENTER_VERTICAL)
        fInp4a = G2G.ValidatedTxtCtrl(mnpnl,self.params,'pdfDmax',min=0.0)
        pdfwidgets.append(fInp4a)
        sizer.Add(fInp4a,0,wx.ALIGN_CENTER_VERTICAL)
        cOpt = G2G.G2CheckBox(mnpnl,'Optimize',self.params,'optPDF')
        pdfwidgets.append(cOpt)
        sizer.Add(cOpt)
        lblsizr.Add(sizer,0)

        sizer = wx.BoxSizer(wx.HORIZONTAL)
        lbl4 = wx.StaticText(mnpnl, wx.ID_ANY,'PDF control: ')
        sizer.Add(lbl4,0,wx.ALIGN_CENTER_VERTICAL)
        self.pdfSel = G2G.G2ChoiceButton(mnpnl,[''],self.params,'pdfprm',
                                       onChoice=checkPDFselection)
        sizer.Add(self.pdfSel,1,wx.ALIGN_CENTER_VERTICAL|wx.EXPAND,1)
        lblsizr.Add(sizer,0,wx.EXPAND)
        sizer = wx.BoxSizer(wx.HORIZONTAL)
        lbl5a = wx.StaticText(mnpnl, wx.ID_ANY,'Chemical formula: ')
        sizer.Add(lbl5a)
        lbl5b = wx.StaticText(mnpnl, wx.ID_ANY,'(formula)')
        sizer.Add(lbl5b)
        lblsizr.Add(sizer,0,wx.EXPAND)

        sizer = wx.BoxSizer(wx.HORIZONTAL)
        lbl5 = wx.StaticText(mnpnl, wx.ID_ANY,'Select format(s):')
        sizer.Add(lbl5)
        for fmt in self.PDFformats:
            sizer.Add((6,2)) # add a bit of extra space
            if fmt not in self.params['outsel']: self.params['outsel'][fmt] = False
            btn = G2G.G2CheckBox(mnpnl,fmt,self.params['outsel'],fmt,
                                     OnChange=TestInput)
            sizer.Add(btn)
            pdfwidgets.append(btn)
        lblsizr.Add(sizer,0,wx.EXPAND)

        self.pbkg = 3*[None]
        for i,lbl in enumerate((' Sample bkg:',' Container:',
                                   'Container bkg:')):
            self.pbkg[i] = [None,None,None,'',None,None,-1.0]
            sizer = wx.BoxSizer(wx.HORIZONTAL)
            self.pbkg[i][0] = wx.StaticText(mnpnl, wx.ID_ANY,lbl)
            sizer.Add(self.pbkg[i][0])
            self.pbkg[i][1] = G2G.G2ChoiceButton(mnpnl,[''],self.pbkg[i],3)
            sizer.Add(self.pbkg[i][1],1,wx.EXPAND,1)
            self.pbkg[i][5] = wx.StaticText(mnpnl, wx.ID_ANY,' mult:')
            sizer.Add(self.pbkg[i][5])
            self.pbkg[i][4] = G2G.ValidatedTxtCtrl(mnpnl,self.pbkg[i],6,
                            (6,3),typeHint=float,size=(50,-1),
                            OnLeave=TestInput,notBlank=False)
            sizer.Add(self.pbkg[i][4])
            lblsizr.Add(sizer,0,wx.EXPAND,0)

        mnsizer.Add(lblsizr,0,wx.ALIGN_CENTER|wx.EXPAND,1)

        # buttons on bottom
        sizer = wx.BoxSizer(wx.HORIZONTAL)
        sizer.Add((20,-1))
        self.btnstart = wx.Button(mnpnl,  wx.ID_ANY, "Start")
        self.btnstart.Bind(wx.EVT_BUTTON, OnStart)
        sizer.Add(self.btnstart)
        self.btnreset = wx.Button(mnpnl,  wx.ID_ANY, "Reset")
        self.btnreset.Bind(wx.EVT_BUTTON, OnReset)
        sizer.Add(self.btnreset)
        sizer.Add((20,-1),wx.EXPAND,1)

        sizer.Add(wx.StaticText(mnpnl, wx.ID_ANY,' cores:'))
        spin = wx.SpinCtrl(mnpnl, wx.ID_ANY,size=(50,-1))
        spin.SetRange(1, mp.cpu_count())
#        spin.SetRange(0, mp.cpu_count()) # debug
        spin.SetValue(self.ncores)
        def _onSpin(event):
            if event: event.Skip()
            self.ncores = spin.GetValue()
        spin.Bind(wx.EVT_SPINCTRL, _onSpin)
        spin.Bind(wx.EVT_KILL_FOCUS, _onSpin)
        sizer.Add(spin)

        sizer.Add((20,-1),wx.EXPAND,1)
        self.btnclose = wx.Button(mnpnl,  wx.ID_ANY, "Exit")
        self.btnclose.Bind(wx.EVT_BUTTON, OnQuit)
        self.EnableIntButtons(False)
        sizer.Add(self.btnclose)
        sizer.Add((20,-1))
        mnsizer.Add(sizer,0,wx.EXPAND|wx.BOTTOM|wx.TOP,5)
        # finish up window
        mnpnl.SetSizer(mnsizer)
        #OnRadioSelect(None) # disable widgets
        mnsizer.Fit(self)
        ShowbyMode()
        if len(sys.argv) > 1:
            fil = os.path.splitext(sys.argv[1])[0]+'.gpx'
            if os.path.exists(fil):
                self.gpxin[3] = fil
                SetGPXInputFile()
        showPDFctrls(None)

    def SetSourceDir(self,event):
        '''Use a dialog to get a directory for image files
        '''
        dlg = wx.DirDialog(self, 'Select directory for image files',
                        self.params['readdir'],wx.DD_DEFAULT_STYLE)
        dlg.CenterOnParent()
        try:
            if dlg.ShowModal() == wx.ID_OK:
                self.params['readdir'] = dlg.GetPath()
            self.readDir.SetValue(self.params['readdir'])
            self.ShowMatchingFiles(None)
        finally:
            dlg.Destroy()
        return

    def ShowMatchingFiles(self,value,invalid=False,**kwargs):
        '''Find and image files matching the image
        file directory (self.params['readdir']) and the image file filter
        (self.params['filter']) and add this information to the GUI list box
        '''
        if invalid: return
        if self.PreventReEntryShowMatch: return
        self.PreventReEntryShowMatch = True
        filmsg = ""
        self.currImageList = []
        if os.path.exists(self.params['readdir']):
            imageList = sorted(
                glob.glob(os.path.join(self.params['readdir'],self.params['filter'])))
            if not imageList:
                msg = 'Warning: No files match search string '+os.path.join(
                        self.params['readdir'],self.params['filter'])
            else:
                for fil in imageList:
                    if fil not in self.ProcessedList:
                        filmsg += '\n  '+fil
                        self.currImageList.append(fil)
                if filmsg:
                    msg = 'Files to integrate from '+os.path.join(
                            self.params['readdir'],self.params['filter'])+filmsg
                else:
                    msg = 'No files found to process in '+self.params['readdir']
        else:
            msg = 'Warning, does not exist: '+self.params['readdir']
        if self.ProcessedList:
            msg += '\nIntegrated files:'
            for fil in self.ProcessedList:
                msg += '\n  '+fil
        self.ListBox.Clear()
        self.ListBox.AppendItems(msg.split('\n'))
        self.PreventReEntryShowMatch = False
        return

    def OnPause(self):
        '''Respond to Pause, changes text on button/Status line, if needed
        Stops timer
        self.Pause should already be True
        '''
        if self.timer.IsRunning(): self.timer.Stop()
        if self.btnstart.GetLabel() == 'Restart':
            return
        if self.btnstart.GetLabel() != 'Resume':
            print('\nPausing autointegration\n')
            self.btnstart.SetLabel('Resume')
            self.Status.SetStatusText(
                    'Press Resume to continue integration or Reset to prepare to reintegrate all images')
        self.Pause = True

    def EnableIntButtons(self,flag):
        for item in (self.btnstart,self.btnreset): item.Enable(flag)

    def StartLoop(self):
        '''Prepare to start autointegration timer loop.
        Save current Image params for use in future integrations
        also label the window so users understand what is being used
        '''
        print('\nStarting new autointegration\n')
        # make sure all output directories exist
        if self.params['SeparateDir']:
            for dfmt in self.fmtlist:
                if not self.params['outsel'][dfmt[1:]]: continue
                dir = os.path.join(self.params['outdir'],dfmt[1:])
                if not os.path.exists(dir): os.makedirs(dir)
        else:
            if not os.path.exists(self.params['outdir']):
                os.makedirs(self.params['outdir'])
        if self.Reset: # special things to do after Reset has been pressed
            self.G2frame.IntegratedList = []
            wx.Yield()
            self.Reset = False
        if self.params['ComputePDF'] and self.params['SeparateDir']:
            for fmt in self.PDFformats:
                if not self.params['outsel'][fmt]: continue
                dir = os.path.join(self.params['outdir'],
                                   fmt.replace("(","_").replace(")",""))
                if not os.path.exists(dir): os.makedirs(dir)
        return False

    def ArgGen(self,PDFobj,imgprms,mskprms,xydata):
        '''generator for arguments for integration/PDF calc
        '''
        for newImage in self.currImageList:
            self.Pause |= self.G2frame.PauseIntegration
            if self.Pause:
                self.OnPause()
                self.PreventTimerReEntry = False
                self.Raise()
                return
            TableMode = self.params['TableMode']
            ComputePDF = self.params['ComputePDF']
            SeparateDir = self.params['SeparateDir']
            optPDF = self.params['optPDF']
            outdir = self.params['outdir']
            calcModes = (TableMode,ComputePDF,SeparateDir,optPDF)
            InterpVals = self.params.get('InterVals')
            outputSelect = self.params['outsel']
            PDFformats = self.PDFformats
            outputModes = (outputSelect,PDFformats,self.fmtlist,outdir)
            if PDFobj:
                PDFdict = PDFobj.data
            else:
                PDFdict = None
            yield (newImage,imgprms,mskprms,xydata,PDFdict,InterpVals,calcModes,outputModes)
    def OnTimerLoop(self,event):
        '''A method that is called every :meth:`PollTime` seconds that is
        used to check for new files and process them. Integrates new images.
        Also optionally sets up and computes PDF.
        This is called only after the "Start" button is pressed (then its label reads "Pause").
        '''

        if GSASIIpath.GetConfigValue('debug'):
            import datetime
            print ("DBG_Timer tick at {:%d %b %Y %H:%M:%S}\n".format(datetime.datetime.now()))
        if self.PreventTimerReEntry: return
        self.PreventTimerReEntry = True
        self.ShowMatchingFiles(None)
        if not self.currImageList:
            self.PreventTimerReEntry = False
            return
        updateList = False

        # get input for integration
        imgprms = mskprms = None
        if not self.params['TableMode']:
            # read in image controls/masks, used below in loop. In Table mode
            # we will get this image-by image.
            gpxinp = G2sc.G2Project(self.gpxin[3])
            print('reading template project',gpxinp.filename)
            img = gpxinp.image(self.imprm[1].GetStringSelection())
            imgprms = img.getControls(True)
            if self.maskfl[1].GetStringSelection().strip():
                img = gpxinp.image(self.maskfl[1].GetStringSelection())
                mskprms = img.getMasks()
        # setup shared input for PDF computation (for now will not be table mode)
        xydata = {}
        if self.params['ComputePDF']:
            pdfEntry = self.pdfSel.GetStringSelection()
            try:
                PDFobj = gpxinp.pdf(pdfEntry)
            except KeyError:
                print("PDF entry not found: {}".format(pdfEntry))
            # update with GUI input
            for i,lbl in enumerate(('Sample Bkg.','Container',
                                    'Container Bkg.')):
                name = self.pbkg[i][1].GetStringSelection()
                try:
                    xydata[lbl] = gpxinp.histogram(name).data['data']
                except AttributeError:
                    pass
                PDFobj.data['PDF Controls'][lbl]['Mult'] = self.pbkg[i][6]
                PDFobj.data['PDF Controls'][lbl]['Name'] = name
        else:
            PDFobj = None
        if self.MPpool:
            self.MPpool.imap_unordered(ProcessImageMP,
                            self.ArgGen(PDFobj,imgprms,mskprms,xydata))
        else:
            for intArgs in self.ArgGen(PDFobj,imgprms,mskprms,xydata):
                newImage = intArgs[0]
                print('processing ',newImage)
                ProcessImage(*intArgs)
                updateList = True
        for newImage in self.currImageList:
            self.ProcessedList.append(newImage)
        if updateList: self.ShowMatchingFiles(None)
        self.PreventTimerReEntry = False
        self.Raise()

def ProcessImageMP(onearg):
    newImage = onearg[0]
    print('processing ',newImage)
    t0 = time.time() # debug
    ProcessImage(*onearg)
    print('ProcessImageMP time',time.time()-t0)

MapCache = {'maskMap':{}, 'ThetaAzimMap':{}, 'distanceList':[]}
'caches for TA and Mask maps'

def ProcessImage(newImage,imgprms,mskprms,xydata,PDFdict,InterpVals,calcModes,outputModes):
    '''Process one image that is read from file newImage and is integrated into
    one or more diffraction patterns and optionally each diffraction pattern can
    be transformed into a pair distribution function.

    :param str newImage: file name (full path) for input image
    :param dict imgprms: dict with some nested lists & dicts describing the image
      settings and integration parameters
    :param dict mskprms: dict with areas of image to be masked
    :param dict xydata: contains histogram information with about background
      contributions, used for PDF computation (used if ComputePDF is True)
    :param PDFdict: contains PDF parameters (used if ComputePDF is True)
    :param InterpVals: contains interpolation table (used if TableMode is True)
    :param tuple calcModes: set of values for which computations are
      performed and how
    :param tuple outputModes: determines which files are written and where
    '''
    (TableMode,ComputePDF,SeparateDir,optPDF) = calcModes
    (outputSelect,PDFformats,fmtlist,outdir) = outputModes
    if SeparateDir:
        savedir = os.path.join(outdir,'gpx')
        if not os.path.exists(savedir): os.makedirs(savedir)
    else:
        savedir = outdir
    outgpx = os.path.join(savedir,os.path.split(os.path.splitext(newImage)[0]+'.gpx')[1])
    gpxout = G2sc.G2Project(filename=outgpx)
    print('creating',gpxout.filename)
    # looped because a file can contain multiple images
    if TableMode: # look up parameter values from table
        imgprms,mskprms = LookupFromTable(im.data['Image Controls'].get('setdist'),
                                                  InterpVals)
    for im in gpxout.add_image(newImage):
        # apply image parameters
        im.setControls(imgprms)
        setdist = '{:.2f}'.format(im.getControls()['setdist']) # ignore differences in position less than 0.01 mm
        if setdist not in MapCache['distanceList']:
            if mskprms:
                im.setMasks(mskprms)
            else:
                im.initMasks()
            MapCache['distanceList'].append(setdist)
            MapCache['maskMap'][setdist] = G2sc.calcMaskMap(im.getControls(),
                                                            im.getMasks())
            MapCache['ThetaAzimMap'][setdist] = G2sc.calcThetaAzimMap(im.getControls())
#        else: # debug
#            print('*** reusing',setdist)
        #if mskprms:
        #    im.setMasks(mskprms)
        #else:
        #    im.initMasks()
        hists = im.Integrate(MaskMap=MapCache['maskMap'][setdist],
                             ThetaAzimMap=MapCache['ThetaAzimMap'][setdist])
        # write requested files
        for dfmt in fmtlist:
            fmt = dfmt[1:]
            if not outputSelect[fmt]: continue
            if fmtlist[dfmt] is None: continue
            if SeparateDir:
                savedir = os.path.join(outdir,fmt)
            else:
                savedir = outdir
            if not os.path.exists(savedir): os.makedirs(savedir)
            # loop over created histgrams (multiple if caked), writing them as requested
            for i,h in enumerate(hists):
                fname = h.name[5:].replace(' ','_')
                try:
                    fil = os.path.join(savedir,fname)
                    print('Wrote',h.Export(fil,dfmt))
                except Exception as msg:
                    print('Failed to write {} as {}. Error msg\n{}'
                              .format(fname,dfmt,msg))
        if ComputePDF:  # compute PDF
            for h in hists:
                pdf = gpxout.copy_PDF(PDFdict,h)
                pdf.data['PDF Controls']['Sample']['Name'] = h.name
                xydata['Sample'] = h.data['data']
                fname = h.name[5:].replace(' ','_')
                limits = h.data['Limits'][1]
                inst = h.data['Instrument Parameters'][0]
                pdf.calculate(copy.deepcopy(xydata),limits,inst)
                if optPDF:
                    for i in range(5):
                        if pdf.optimize(True,5,copy.deepcopy(xydata),limits,inst):
                            break
                    pdf.calculate(copy.deepcopy(xydata),limits,inst)
                for fmt in PDFformats:
                    if not outputSelect[fmt]: continue
                    if SeparateDir:
                        savedir = os.path.join(outdir,fmt.replace("(","_").replace(")",""))
                    else:
                        savedir = outdir
                    pdf.export(os.path.join(savedir,fname),fmt)
    if outputSelect.get('gpx'):
        gpxout.save()
    else:
        del gpxout
# Autointegration end

def SetupInterpolation(dlg):
    '''Creates an object for interpolating image parameters at a given distance value
    '''
    parms = dlg.ReadImageParmTable()
    IMfileList = dlg.IMfileList
    cols = dlg.list.GetColumnCount()
    ParmList = dlg.ParmList
    nonInterpVars = dlg.nonInterpVars
    ControlsTable = {}
    MaskTable = {}
    for f,m in zip(IMfileList,parms[-1]):
        n = os.path.split(f)[1]
        if n in ControlsTable:
            print('Warning overwriting entry {}'.format(n))
        ControlsTable[n] = G2imG.ReadControls(f)
        if m and os.path.exists(m):
            MaskTable[n] = G2imG.ReadMask(m)
        elif m != "(none)":
            print("Error: Mask file {} not found".format(m))
    return copy.deepcopy([cols, parms, IMfileList, ParmList, nonInterpVars,ControlsTable,MaskTable])

def LookupFromTable(dist,parmList):
    '''Interpolate image parameters for a supplied distance value

    :param float dist: distance to use for interpolation
    :returns: a list with 2 items:
      * a dict with interpolated parameter values,
      * the closest imctrl
    '''
    cols, parms, IMfileList, ParmList, nonInterpVars,ControlsTable,MaskTable = parmList
    x = np.array([float(i) for i in parms[0]])
    closest = abs(x-dist).argmin()
    D = {'setdist':dist}
    imctfile = IMfileList[closest]
    for c in range(1,cols-1):
        lbl = ParmList[c]
        if lbl in nonInterpVars:
            if lbl in ['outChannels',]:
                D[lbl] = int(float(parms[c][closest]))
            else:
                D[lbl] = float(parms[c][closest])
        else:
            y = np.array([float(i) for i in parms[c]])
            D[lbl] = np.interp(dist,x,y)
    # full integration when angular range is 0
    D['fullIntegrate'] = (D['LRazimuth_min'] == D['LRazimuth_max'])
    # conversion for paired values
    for a,b in ('center_x','center_y'),('LRazimuth_min','LRazimuth_max'),('IOtth_min','IOtth_max'):
        r = a.split('_')[0]
        D[r] = [D[a],D[b]]
        if r in ['LRazimuth',]:
            D[r] = [int(D[a]),int(D[b])]
        del D[a]
        del D[b]
    interpDict,imgctrl = D,imctfile
    if GSASIIpath.GetConfigValue('debug'):
        print ('DBG_interpolated values: ',interpDict)
    f = os.path.split(imgctrl)[1]
    ImageControls = ControlsTable[f]
    ImageControls.update(interpDict)
    ImageControls['showLines'] = True
    ImageControls['ring'] = []
    ImageControls['rings'] = []
    ImageControls['ellipses'] = []
    ImageControls['setDefault'] = False
    for i in 'range','size','GonioAngles':
        if i in ImageControls: del ImageControls[i]
    ImageMasks = MaskTable.get(f)
    return ImageControls,ImageMasks

###########################################################################
if __name__ == "__main__":
    GSASIIpath.InvokeDebugOpts()
    App = wx.App()
    class dummyClass(object):
        '''An empty class where a few values needed from parent are placed
        '''
        def __init__(self):
            self.Image = None
            self.PauseIntegration = False
            self.TutorialImportDir = None
            self.GSASprojectfile = ''
            self.LastExportDir = ''
            self.LastGPXdir = ''

    G2frame = dummyClass()
    frm = AutoIntFrame(G2frame,5)
    App.GetTopWindow().Show(True)
    App.MainLoop()


"""
Rapidly Developing Thunderstorms (RDT)
--------------------------------------
"""
import os
import glob
import re
import datetime as dt
import bokeh
import json
import numpy as np
from forest import (
        geo,
        locate)
from forest.util import timeout_cache
from forest.exceptions import FileNotFound
from bokeh.palettes import GnBu3, OrRd3
import itertools
import math

class RenderGroup(object):
    """Collection of renderers that act as one"""
    def __init__(self, renderers, visible=False):
        self.renderers = renderers
        self.visible = visible

    @property
    def visible(self):
        return any(r.visible for r in self.renderers)

    @visible.setter
    def visible(self, value):
        for r in self.renderers:
            r.visible = value

class View(object):
    """Rapidly Developing Thunderstorms (RDT) visualisation"""
    def __init__(self, loader):
        self.loader = loader
        empty = {
            "type": "FeatureCollection",
            "features": [
                {
                    "type": "Feature",
                    "geometry": {
                        "type": "Polygon",
                        "coordinates": [[[0, 0]]]
                    },
                    "properties": {
                        'CType': 0,
                        'CRainRate': 0,
                        'ConvTypeMethod': 0,
                        'ConvType': 0,
                        'ConvTypeQuality': 0,
                        'SeverityIntensity': 0,
                        'MvtSpeed': 0,
                        'MvtDirection': 0,
                        'NumIdCell': 0,
                        'CTPressure': 0,
                        'CTPhase': '',
                        'CTReff': '',
                        'ExpansionRate': '-',
                        'BTmin': 0,
                        'BTmoy': 0,
                        'CTCot': '-',
                        'CTCwp': '-',
                        'NbPosLightning': 0,
                        'SeverityType': '',
                        'Surface': '',
                        'Duration': 0,
                        'CoolingRate': 0,
                        'PhaseLife': "Triggering"
                    }
                }
            ]
        }
        self.empty_geojson = json.dumps(empty)
        self.empty_tail_line = dict(
                xs=[], ys=[],
                LonTrajCellCG=[],
                LatTrajCellCG=[],
                NumIdCell=[],
                NumIdBirth=[],
                DTimeTraj=[],
                BTempTraj=[],
                BTminTraj=[],
                BaseAreaTraj=[],
                TopAreaTraj=[],
                CoolingRateTraj=[],
                ExpanRateTraj=[],
                SpeedTraj=[],
                DirTraj=[])
        self.empty_tail_point = dict(
                x=[], y=[],
                LonTrajCellCG=[],
                LatTrajCellCG=[],
                NumIdCell=[],
                NumIdBirth=[],
                DTimeTraj=[],
                BTempTraj=[],
                BTminTraj=[],
                BaseAreaTraj=[],
                TopAreaTraj=[],
                CoolingRateTraj=[],
                ExpanRateTraj=[],
                SpeedTraj=[],
                DirTraj=[])
        self.empty_centre_point = dict(
                x1=[], y1=[], x2=[], y2=[], xs=[], ys=[],
                Arrowxs=[],
                Arrowys=[],
                LonG=[],
                LatG=[],
                NumIdCell=[],
                NumIdBirth=[],
                MvtSpeed=[],
                MvtDirection=[])
        self.color_mapper = bokeh.models.CategoricalColorMapper(
                palette=['#fee8c8', '#fdbb84', '#e34a33', '#43a2ca', '#a8ddb5'],
                factors=["Triggering", "Triggering from split", "Growing", "Mature", "Decaying"])
        self.source = bokeh.models.GeoJSONDataSource(geojson=self.empty_geojson)
        self.tail_line_source = bokeh.models.ColumnDataSource(self.empty_tail_line)
        self.tail_point_source = bokeh.models.ColumnDataSource(self.empty_tail_point)
        self.centre_point_source = bokeh.models.ColumnDataSource(self.empty_centre_point)

    def render(self, state):
        """Gets called when a menu button is clicked (or when application state changes)"""
        if state.valid_time is not None:
            date = dt.datetime.strptime(state.valid_time, '%Y-%m-%d %H:%M:%S')
            try:
                (self.source.geojson,
                 self.tail_line_source.data,
                 self.tail_point_source.data,
                 self.centre_point_source.data) = self.loader.load_date(date)
            except FileNotFound:
                print("rdt.View.render caught FileNotFound", date)
                self.source.geojson = self.empty_geojson
                self.tail_line_source.data = self.empty_tail_line
                self.tail_point_source.data = self.empty_tail_point
                self.centre_point_source.data = self.empty_centre_point

    def add_figure(self, figure):
        """This is where all the plotting happens (e.g. when the applciation is loaded)"""
        circles = figure.circle(x="x", y="y", size=3, source=self.tail_point_source)
        cntr_circles = figure.circle_cross(x="x1", y="y1", size=10, line_color='black', fill_color=None, source=self.centre_point_source)
        future_lines = figure.multi_line(xs="xs", ys="ys", line_color='black', source=self.centre_point_source)
        lines = figure.multi_line(xs="xs", ys="ys", line_dash='dashed', source=self.tail_line_source)
        arrows = figure.patches(xs='Arrowxs', ys='Arrowys', fill_color='black', line_color='black', source=self.centre_point_source)
        renderer = figure.patches(
            xs="xs",
            ys="ys",
            fill_alpha=0,
            line_width=2,
            line_color={
                 'field': 'PhaseLife',
                 'transform': self.color_mapper},
            source=self.source)

        tool = bokeh.models.HoverTool(
                tooltips=[
                    ('NumIdCell', '@NumIdCell'),
                    ('Duration (Since Birth)', '@Duration{00:00:00}'),
                    ('Phase life', '@PhaseLife'), # Categorical
                    ('Cloud Type', '@CType'), # Categorical
                    ('Convective Rainfall Rate', '@CRainRate{0.0}' + ' mm/hr'),
                    ('Cloud System', '@ConvType'), # Categorical
                    ('Severity Type', '@SeverityType'), # Categorical
                    ('Severity Intensity', '@SeverityIntensity'), # Categorical
                    ('Cloud Top Phase', '@CTPhase'), # Categorical
                    ('Min. Cloud Top Pressure', '@CTPressure' + ' hPa'),
                    # ('Max. Cloud Top Effective Radius', '@CTReff' + ' metres'),
                    ('Expansion Rate (Past)', '@ExpansionRate{+0.0}' + ' m-2/sec'),
                    ('Rate of Temp. Change', '@CoolingRate{+0.0}' + ' K/15mins'),
                    ('Min. Brightness Temp', '@BTmin{0.0}' + ' K'),
                    ('Average Brightness Temp', '@BTmoy{0.0}' + ' K'),
                    ('Max. Cloud Optical Thickness', '@CTCot'),
                    ('No. of Cloud-Ground Positive Lightning Strokes', '@NbPosLightning')
                ],
                renderers=[renderer])
        figure.add_tools(tool)
        return RenderGroup([renderer, lines, circles, cntr_circles, future_lines, arrows])

class Loader(object):
    """High-level RDT loader"""
    def __init__(self, pattern):
        self.locator = Locator(pattern)

    def load_date(self, date):
        file_name = self.locator.find_file(date)
        return (
                self.load_polygon(file_name),
                self.load_tail_lines(file_name),
                self.load_tail_points(file_name),
                self.load_centre_points(file_name))

    @staticmethod
    def load_polygon(path):
        """Load GeoJSON string representation of Polygons from file

        :returns: GeoJSON str
        """
        with open(path) as stream:
            rdt = json.load(stream)

        # Convert units from the netcdf / geojson data file units into something more readable (e.g. could be Kelvin to degrees C or Pa to hPa)
        unitsToRescale = {'Pa' : {'scale':100, 'offset':0, 'Units': 'hPa'} }
        # Get text labels instead of numbers for certain fields
        fieldsToLookup = ['PhaseLife', 'SeverityType', 'SeverityIntensity', 'ConvType', 'CType']

        copy = dict(rdt)
        for i, feature in enumerate(rdt["features"]):
            coordinates = feature['geometry']['coordinates'][0]
            lons, lats = np.asarray(coordinates).T
            x, y = geo.web_mercator(lons, lats)
            c = np.array([x, y]).T.tolist()
            copy["features"][i]['geometry']['coordinates'][0] = c

            for k in feature['properties'].keys():

                # Might be easier to have a units lookup instead of this ...
                deldata, myunits = descale_rdt(k, feature['properties'][k])
                if myunits in unitsToRescale.keys():
                    try:
                        mydict = unitsToRescale.get(myunits)
                        scale, offset, units = mydict.values()
                        conv_data = (feature['properties'][k] / scale) + offset
                        copy['features'][i]['properties'][k] = conv_data
                    except:
                        continue

                if k in fieldsToLookup:
                    try:
                        copy['features'][i]['properties'][k] = fieldValueLUT(k, feature['properties'][k])
                    except:
                        continue
        return json.dumps(copy)

    @staticmethod
    def load_tail_lines(path):
        """Load tail line data from file

        :returns: dict representation suitable for ColumnDataSource
        """
        with open(path) as stream:
            rdt = json.load(stream)

        # Create an empty dictionary
        mydict = dict(
                xs=[], ys=[],
                LonTrajCellCG=[],
                LatTrajCellCG=[],
                NumIdCell=[],
                NumIdBirth=[],
                DTimeTraj=[],
                BTempTraj=[],
                BTminTraj=[],
                BaseAreaTraj=[],
                TopAreaTraj=[],
                CoolingRateTraj=[],
                ExpanRateTraj=[],
                SpeedTraj=[],
                DirTraj=[])

        # Loop through features
        for i, feature in enumerate(rdt["features"]):
            # Append data from the feature properties to the dictionary
            for k in mydict.keys():
                try:
                    thisdata, units = descale_rdt(k, feature['properties'][k])
                    mydict[k].append(thisdata)
                except:
                    # Do nothing at the moment with the xs and ys
                    if not k in ['xs', 'ys']:
                        mydict[k].append(None)
                    else:
                        continue
            # Takes the trajectory lat/lons, reprojects and
            # puts them in a list within a list
            lons = feature['properties']['LonTrajCellCG']
            lats = feature['properties']['LatTrajCellCG']
            xs, ys = geo.web_mercator(lons, lats)
            mydict['xs'].append(xs)
            mydict['ys'].append(ys)
        return mydict

    @staticmethod
    def load_tail_points(path):
        with open(path) as stream:
            rdt = json.load(stream)

        # Create an empty dictionary
        mydict = dict(
                x=[], y=[],
                LonTrajCellCG=[],
                LatTrajCellCG=[],
                NumIdCell=[],
                NumIdBirth=[],
                DTimeTraj=[],
                BTempTraj=[],
                BTminTraj=[],
                BaseAreaTraj=[],
                TopAreaTraj=[],
                CoolingRateTraj=[],
                ExpanRateTraj=[],
                SpeedTraj=[],
                DirTraj=[])

        # Loop through features
        for i, feature in enumerate(rdt["features"]):
            # Append data from the feature properties to the dictionary
            # First though, how many points do we have in the trajectory tail?
            npts = len(feature['properties']['LonTrajCellCG'])
            mykeys = [k for k in mydict.keys() if k not in ['x', 'y']]
            for k in mykeys:
                # print(k, type(feature['properties'][k]), sep=':')
                try:
                    if not isinstance(feature['properties'][k], list):
                        datalist = [i for i in itertools.repeat(feature['properties'][k],npts)]
                        thisdata, units = descale_rdt(k, datalist)
                        mydict[k].extend(datalist)
                    else:
                        thisdata, units = descale_rdt(k, feature['properties'][k])
                        mydict[k].extend(thisdata)
                except:
                    datalist = [i for i in itertools.repeat(None, npts)]
                    mydict[k].extend(datalist)
            # Takes the trajectory lat/lons, reprojects and puts them in a list within a list
            lons = feature['properties']['LonTrajCellCG']
            lats = feature['properties']['LatTrajCellCG']
            x, y = geo.web_mercator(lons, lats)
            mydict['x'].extend(x)
            mydict['y'].extend(y)
        return mydict

    @staticmethod
    def load_centre_points(path):
        """Holds a centre point, future point and future movement line"""
        with open(path) as stream:
            rdt = json.load(stream)

        # Create an empty dictionary
        mydict = dict(
                x1=[], y1=[], x2=[], y2=[], xs=[], ys=[],
                Arrowxs=[],
                Arrowys=[],
                LonG=[],
                LatG=[],
                NumIdCell=[],
                NumIdBirth=[],
                MvtSpeed=[],
                MvtDirection=[])

        # Loop through features
        for i, feature in enumerate(rdt["features"]):
            # Append data from the feature properties to the dictionary
            mykeys = [k for k in mydict.keys() if not (('x' in k) or ('y' in k))]
            for k in mykeys:
                try:
                    thisdata, units = descale_rdt(k, feature['properties'][k])
                    mydict[k].append(thisdata)
                except:
                    mydict[k].append(None)
            # Takes the trajectory lat/lons, reprojects and puts them in a list within a list
            lon = feature['properties']['LonG']
            lat = feature['properties']['LatG']
            x1, y1 = geo.web_mercator(lon, lat)
            mydict['x1'].extend(x1)
            mydict['y1'].extend(y1)

            # Now calculate future point and line
            try:
                speed = float(feature['properties']['MvtSpeed'])
            except ValueError:
                speed = 0
            try:
                direction = float(feature['properties']['MvtDirection'])
            except ValueError:
                direction = 0
            lon2, lat2 = calc_dst_point(lon, lat, speed, direction)
            x2, y2 = geo.web_mercator(lon2, lat2)
            mydict['x2'].extend(x2)
            mydict['y2'].extend(y2)
            mydict['xs'].append([x1, x2])
            mydict['ys'].append([y1, y2])

            # Now calculate arrow polygon
            x3d, y3d, x4d, y4d = get_arrow_poly(lon2, lat2, speed, direction)
            [x3, x4], [y3, y4] = geo.web_mercator([x3d, x4d], [y3d, y4d])

            mydict['Arrowxs'].append([x2[0], x3, x4])
            mydict['Arrowys'].append([y2[0], y3, y4])
        return mydict

def calc_dst_point(x1d, y1d, speed, angle):
    """Calculate destination point

    Estimates positions in longitude/latitude space from speed and
    angle on the surface of a sphere, in this case Earth.

    :param x1d: longitude
    :param y1d: latitude
    """
    # NB: X and Y need to be lat/lons

    # Distance travelled (m) = speed (m/s) * 60 seconds * 60 minutes
    # NB: 60 mins may change depending on the time frequency of the display (currently 1 hour)
    d = (speed * 60 * 60)

    # Radius of the earth (m)
    R = 6378137

    x1 = math.radians(x1d)
    y1 = math.radians(y1d)

    # Convert degrees to radians
    direction = math.radians(angle)

    y2 = math.asin(math.sin(y1) * math.cos(d / R) +
                   math.cos(y1) * math.sin(d / R) * math.cos(direction))

    x2 = x1 + math.atan2(math.sin(direction) * math.sin(d / R) * math.cos(y1),
                         math.cos(d / R) - math.sin(y1) * math.sin(y2))

    x2d = math.degrees(x2)
    y2d = math.degrees(y2)
    return x2d, y2d

def get_arrow_poly(x2,y2, speed, direction):
    """Draw a polygon representing an arrow in geographical coordinates

    .. note:: The arrows are scaled in longitude/latitude space not
              in screen coordinates

    :param x2: longitude of point
    :param y2: latitude of point
    :param speed: scalar velocity in m/s
    :param direction: angle in degrees relative to north
    """
    timestep = 60 # See above function re: 60 mins
    mvt_line_len = speed * 60 * timestep
    mvt_line_dir = direction
    arrow_angl = 20
    arrow_linefrac = 1./5

    # First point
    pt1_dir = (mvt_line_dir - 180) % 360 - arrow_angl
    pt1_len = math.sqrt(3. * math.pow( mvt_line_len * arrow_linefrac, 2 ) / 2) # Metres
    # Convert len back to speed for the function
    pt1_speed = pt1_len / (timestep * 60)
    # Calculate x3, y3
    x3, y3 = calc_dst_point(x2, y2, pt1_speed, pt1_dir)

    # Second point
    pt2_dir = (mvt_line_dir - 180) % 360 + arrow_angl
    pt2_len = math.sqrt(3.* math.pow( mvt_line_len * arrow_linefrac, 2 ) / 2) # Metres
    # Convert len back to speed for the function
    pt2_speed = pt2_len / (timestep * 60)
    # Calculate x3, y3
    x4, y4 = calc_dst_point(x2, y2, pt2_speed, pt2_dir)
    return x3, y3, x4, y4

def descale_rdt(fn, data):
    # Converts units according to netcdf files definition
    rdtUnitsLUT = {
        'DecTime': {'scale': 1, 'offset': 0, 'Units': 's'},
        'LeadTime': {'scale': 1, 'offset': 0, 'Units': 's'},
        'Duration': {'scale': 1, 'offset': 0, 'Units': 's'},
        'MvtSpeed': {'scale': 0.001, 'offset': 0, 'Units': 'm s-1'},
        'MvtDirection': {'scale': 1, 'offset': 0, 'Units': 'degree'},
        'DtTimeRate': {'scale': 1, 'offset': 0, 'Units': 's'},
        'ExpansionRate': {'scale': 2e-07, 'offset': -0.005, 'Units': 's-1'},
        'CoolingRate': {'scale': 2e-06, 'offset': -0.05, 'Units': 'K s-1'},
        'LightningRate': {'scale': 1e-04, 'offset': -2.0, 'Units': 's-1'},
        'CTPressRate': {'scale': 0.001, 'offset': -25.0, 'Units': 'Pa s-1'},
        'BTemp': {'scale': 0.01, 'offset': 130.0, 'Units': 'K'},
        'BTmoy': {'scale': 0.01, 'offset': 130.0, 'Units': 'K'},
        'BTmin': {'scale': 0.01, 'offset': 130.0, 'Units': 'K'},
        'Surface': {'scale': 5000000.0, 'offset': 0, 'Units': 'm2'},
        'EllipseGaxe': {'scale': 20.0, 'offset': 0, 'Units': 'm'},
        'EllipsePaxe': {'scale': 20.0, 'offset': 0, 'Units': 'm'},
        'EllipseAngle': {'scale': 1, 'offset': 0, 'Units': 'degrees_north'},
        'DtLightning': {'scale': 1, 'offset': 0, 'Units': 's'},
        'CTPressure': {'scale': 10.0, 'offset': 0, 'Units': 'Pa'},
        'CTCot': {'scale': 0.01, 'offset': 0, 'Units': '1'},
        'CTReff': {'scale': 1e-08, 'offset': 0, 'Units': 'm'},
        'CTCwp': {'scale': 0.001, 'offset': 0, 'Units': 'kg m-2'},
        'CRainRate': {'scale': 0.1, 'offset': 0, 'Units': 'mm/h'},
        'BTempSlice': {'scale': 0.01, 'offset': 130.0, 'Units': 'K'},
        'SurfaceSlice': {'scale': 5000000.0, 'offset': 0, 'Units': 'm2'},
        'DTimeTraj': {'scale': 1, 'offset': 0, 'Units': 's'},
        'BTempTraj': {'scale': 0.01, 'offset': 130.0, 'Units': 'K'},
        'BTminTraj': {'scale': 0.01, 'offset': 130.0, 'Units': 'K'},
        'BaseAreaTraj': {'scale': 5000000.0, 'offset': 0, 'Units': 'm2'},
        'TopAreaTraj': {'scale': 5000000.0, 'offset': 0, 'Units': 'm2'},
        'CoolingRateTraj': {'scale': 2e-06, 'offset': -0.05, 'Units': 'K s-1'},
        'ExpanRateTraj': {'scale': 2e-07, 'offset': -0.005, 'Units': 's-1'},
        'SpeedTraj': {'scale': 0.001, 'offset': 0, 'Units': 'm s-1'},
        'DirTraj': {'scale': 1, 'offset': 0, 'Units': 'degree'}
    }

    try:
        dict = rdtUnitsLUT.get(fn, {'scale': 1, 'offset': 0, 'units': '-'})
        scale, offset, units = dict.values()
        conv_data = ( data / scale ) + offset
        return(conv_data, units)
    except:
        return(data, '-')

def fieldNameLUT(fn):
    short2long = {
                    'NbSigCell': 'Number of encoded significant RDT-CW cloud cells',
                    'NbConvCell': 'Number of convective cloud cells',
                    'NbCloudCell': 'Number of analyzed and tracked cloud cells',
                    'NbElecCell': 'Number of electric cloud cells',
                    'NbHrrCell': 'Number of cloud cells with High Rain Rate values ',
                    'MapCellCatType': 'NWC GEO RDT-CW Type and phase of significant cells',
                    'NumIdCell': 'Identification Number of cloud cell',
                    'NumIdBirth': 'Identification Number of cloud cell at birth',
                    'DecTime': 'time gap between radiometer time and slot time',
                    'LeadTime': 'lead time from slot for forecast cloud cell',
                    'Duration': 'Duration of cloud system since birth',
                    'ConvType': 'Type (conv or not) of cloud system',
                    'ConvTypeMethod': 'Method used for convective diagnosis',
                    'ConvTypeQuality': 'Quality of convective diagnosis ',
                    'PhaseLife': 'Phase Life of Cloud system',
                    'MvtSpeed': 'Motion speed of cloud cell',
                    'MvtDirection': 'Direction of Motion of cloud cell',
                    'MvtQuality': 'Quality of motion estimation of cloud cell',
                    'DtTimeRate': 'gap time to compute rates of cloud system',
                    'ExpansionRate': 'Expansion rate of cloud system',
                    'CoolingRate': 'Temperature change rate of cloud system',
                    'LightningRate': 'Lightning trend of cloud system',
                    'CTPressRate': 'Top Pressure trend of cloud system',
                    'SeverityType': 'Type of severity of cloud cell',
                    'SeverityIntensity': 'severity intensity of cloud cell ',
                    'LatContour': 'latitude of contour point of cloud cell',
                    'LonContour': 'longitude of contour point of cloud cell',
                    'LatG': 'latitude of Gravity Centre of cloud cell',
                    'LonG': 'longitude of Gravity Centre of cloud cell',
                    'BTemp': 'Brightness Temperature threshold defining a Cloud cell',
                    'BTmoy': 'Average Brightness Temperature over a Cloud cell',
                    'BTmin': 'Minimum Brightness Temperature of a Cloud cell',
                    'Surface': 'Surface of a Cloud cell',
                    'EllipseGaxe': 'Large axis of Ellipse approaching Cloud cell',
                    'EllipsePaxe': 'Small axis of Ellipse approaching Cloud cell',
                    'EllipseAngle': 'Angle of Ellipse approaching Cloud cell',
                    'NbPosLightning': 'Number of CG positive lightning strokes paired with cloud cell',
                    'NbNegLightning': 'Number of CG negative lightning strokes paired with cloud cell',
                    'NbIntraLightning': 'Number of IntraCloud lightning strokes paired with cloud cell',
                    'DtLightning': 'time interval to pair lighting data with cloud cells',
                    'CType': 'Most frequent Cloud Type over cloud cell extension',
                    'CTPhase': 'Most frequent Cloud Top Phase over cloud cell extension',
                    'CTPressure': 'Minimum Cloud Top Pressure over cloud cell extension',
                    'CTCot': 'maximum cloud_optical_thickness over cloud cell extension',
                    'CTReff': 'maximum radius_effective over cloud cell extension',
                    'CTCwp': 'maximum cloud condensed water_path over cloud cell extension',
                    'CTHicgHzd': 'High altitude Icing Hazard index',
                    'CRainRate': 'maximum convective_rain_rate over cloud cell extension',
                    'BTempSlice': 'Brightness Temperature threshold defining a Cloud cell',
                    'SurfaceSlice': 'Surface of Cloud cell at Temperature threshold',
                    'DTimeTraj': 'time gap between current and past Cloud cell',
                    'LatTrajCellCG': 'latitude of Gravity Centre of past cloud cell',
                    'LonTrajCellCG': 'longitude of Gravity Centre of past cloud cell',
                    'BTempTraj': 'Brightness Temperature threshold defining past Cloud cell',
                    'BTminTraj': 'Minimum Brightness Temperature of past Cloud cell',
                    'BaseAreaTraj': 'Surface of base of past Cloud cell',
                    'TopAreaTraj': 'Surface of top of past Cloud cell',
                    'CoolingRateTraj': 'Temperature change rate of past cloud system',
                    'ExpanRateTraj': 'Expansion rate of past cloud system',
                    'SpeedTraj': 'Motion speed of past cloud cell',
                    'DirTraj': 'Direction of Motion of past cloud cell',
                    'lat': 'Latitude at the centre of each pixel',
                    'lon': 'Longitude at the centre of each pixel',
                    'ny': 'Y Georeferenced Coordinate for each pixel count',
                    'nx': 'X Georeferenced Coordinate for each pixel count',
                    'MapCellCatType_pal': 'RGB palette for MapCellCatType'
    }
    try:
        return short2long.get(fn)
    except:
        return fn

def fieldValueLUT(fn, uid):
    rdtLUT = {
        "MapCellCatType": {
            0: "Non convective",
            1: "Convective triggering",
            2: "Convective triggering from split",
            3: "Convective growing",
            4: "Convective mature",
            5: "OvershootingTop mature",
            6: "Convective decaying",
            7: "Electric triggering",
            8: "Electric triggering from split",
            9: "Electric growing",
            10: "Electric mature",
            11: "Electric decaying",
            12: "HighRainRate triggering",
            13: "HighRainRate triggering from split",
            14: "HighRainRate growing",
            15: "HighRainRate mature",
            16: "HighRainRate decaying",
            17: "HighSeverity triggering",
            18: "HighSeverity triggering from split",
            19: "HighSeverity growing",
            20: "HighSeverity mature",
            21: "HighSeverity decaying"
        },
        'ConvType': {
            0: "Non convective",
            1: "Convective",
            2: "Convective inherited",
            3: "Convective forced overshoot",
            4: "Convective forced lightning",
            5: "Convective forced convrainrate",
            6: "Convective forced coldtropical",
            7: "Convective forced inherited",
            8: "Declassified convective",
            9: "Not defined"
        },
        'ConvTypeMethod': {
            1: "Discrimination statistical scheme",
            2: "Electric",
            3: "Overshoot",
            4: "Convective rain rate",
            5: "Tropical"
        },
        'ConvTypeQuality': {
            1: "High quality",
            2: "Moderate quality",
            3: "Low quality",
            4: "Very low quality"
        },
        'PhaseLife': {
            0: "Triggering",
            1: "Triggering from split",
            2: "Growing",
            3: "Mature",
            4: "Decaying"
        },
        'MvtQuality': {
            1: "High quality",
            2: "Moderate quality",
            3: "Low quality",
            4: "Very low quality"
        },
        'SeverityType' : {
            0: "No activity",
            1: "Turbulence",
            2: "Lightning",
            3: "Icing",
            4: "High altitude icing",
            5: "Hail",
            6: "Heavy rainfall",
            7: "not defined"
        },
        'SeverityIntensity' : {
            0: "not defined",
            1: "Low",
            2: "Moderate",
            3: "High",
            4: "Very high"
        },
        'CType' : {
            1: "Cloud-free land",
            2: "Cloud-free sea",
            3: "Snow over land",
            4: "Sea ice",
            5: "Very low clouds",
            6: "Low clouds",
            7: "Mid-level clouds",
            8: "High opaque clouds",
            9: "Very high opaque clouds",
            10: "Fractional clouds",
            11: "High semitransparent thin clouds",
            12: "High semitransparent meanly thick clouds",
            13: "High semitransparent thick clouds",
            14: "High semitransparent above low or medium clouds",
            15: "High semitransparent above snow ice"
        },
        'CTPhase' : {
            1: "Liquid",
            2: "Ice",
            3: "Mixed",
            4: "Cloud-free",
            5: "Undefined separability problems"
        },
        'CTHicgHzd' : {
            0: "not defined",
            1: "Low risk",
            2: "Moderate risk",
            3: "High risk-free",
            4: "Very high risk"
        }
    }
    try:
        return rdtLUT.get(fn)[uid]
    except:
        return "-"

class Locator(object):
    def __init__(self, pattern):
        self.pattern = pattern

    def find_file(self, valid_date):
        paths = np.array(self.paths)  # Note: timeout cache in use
        bounds = locate.bounds(
                self.dates(paths),
                dt.timedelta(minutes=15))
        pts = locate.in_bounds(bounds, valid_date)
        found = paths[pts]
        if len(found) > 0:
            return found[0]
        else:
            raise FileNotFound("RDT: '{}' not found in {}".format(valid_date, bounds))

    @property
    def paths(self):
        return self.find(self.pattern)

    @staticmethod
    @timeout_cache(dt.timedelta(minutes=10))
    def find(pattern):
        return sorted(glob.glob(pattern))

    def dates(self, paths):
        return np.array([
            self.parse_date(p) for p in paths],
            dtype='datetime64[s]')

    @staticmethod
    def parse_date(path):
        groups = re.search(r"[0-9]{12}", os.path.basename(path))
        if groups is not None:
            return dt.datetime.strptime(groups[0], "%Y%m%d%H%M")

class Coordinates(object):
    """Menu system interface"""
    def initial_time(self, path):
        times = self.valid_times(path, None)
        if len(times) > 0:
            return times[0]
        return None

    def variables(self, path):
        return ["RDT"]

    def valid_times(self, path, variable):
        date = Locator.parse_date(path)
        if date is None:
            return []
        return [str(date)]

    def pressures(self, path, variable):
        return None


#!/usr/bin/python2

# Identify the language variant being used by the author of a number of messages.
# Messages are provided as successive "message" fields, and accuracy improves with the number of messages provided.
# At the moment, it supports only the two main variants of Portuguese.

import numpy
import sdb
import sklearn.svm
import sklearn.cross_validation as kfold
import sklearn.metrics
import normalize
import word
import pickle
import json

premade_model = True  # Set to true if you want to load the model from the disk. False to read data from DB.
overwrite_model = False  # If set to true, overwrite the model on disk with the data read from the DB.
sample_size = 5  # Has no effect in the model data if we use premade model data.
model_file = "mongodata.pickle"  # The name of the file to store the model, pickled.
class_names = ("pt", "br")  # The names to return on classification.

# This is the premade collection of users.

pt = [
2264444869,
53161052,
1011077418,
2153646048,
236138844,
1339859353,
212293323,
472169787,
1339859353,
1460491975,
212293323,
2190608648,
717874436,
549264462,
997860512,
1460491975
]

br = [
87843887,
443048958,
87843887,
87843887,
1069863936,
443048958,
443048958,
199901650,
393401881,
199901650
]

# db.twitter.find({_id:{$in:[436229899020689409,436249107183964161,436259186364915713,436276712255799296,436310382194741248,436581291228950528,436637473570312192,436828457671348224,436935368756432896,437000025143783424,437447659809300480,437615529382002688,437732762444132352,437735194985246721,437760998695796736,438004134244286464,438017630730194944,438017712129048576,438017876201844736,438028190616154112,438079892304183296,440590461724794880,440877790733557761,440903906454171648,440914662084067328,441520873422745600,441532032376930304,441544644821016576,441552085289811968,441555983538864129,441556333515796480,441571124175175680,441667247053238274,441684872110022656,441685606868221953,441686981459058688,441694961298591744,441705391073013760,441710362287931392,441719836008513536,43628406292876492,43628248059401830,43643930947092480,43647142375511244,436533966301298689,43661676535192371,43661903720510259,436623647328174081,436647586057818112,43665534046650368,43686432079924428,436897643110096896,43734180011876761,43757654689185792,43760635080239923,43760651855778611,437608828629032960,43800215712447692,43808724673102643,43809079219298304,438090946866331648,438091003149701122,438108751066071040,438119070412836864,438183703412355072,44000107475451084,440922749469335552,448546420216590337,44174745596555264,44121748764101017,44128025600014336,441288407747407872,442114639761928192,448496517641674752,442334679761698816,442675007383433217,442837696516747264,44336087002306150,443380147656617984,44341479194912358,436534495043268608,436534497224298498,43680106082993356,43732087378857164,43739483702612787,448540356049920000,44775845986673868,44753452293949030,44717209162091724,44756449629152870,44675756054966681,436906361868541952,44428402054570803]}}).count()

# db.twitter.find({_id:{$in:[436279144717230080,436591409936994304,436616567766671360,437244894697230336,437315719991001088,43769251987941785 ,437694655480987648,437697395007033344,43769741752245452 ,437698445277921280,43769942177900953 ,437710307490668544,437710611913654273,437728411868741632,437882465407885312,437979993000660992,437980084163837952,448574019680632832,438271917330792448,44857315812864000 ,44855345911814144 ,44855236922824704 ,44855124969300377 ,44855119099372339 ,44855105471697305 ,44854955110094848 ,44854861952084787 ,44854451310008729 ,44854370674515968 ,44854326877172940 ,44853918844794880 ,44853889084530278 ,44860332177583308 ,44861206420733952 ,44861208704950681 ,448533838353149954,44853104610430566 ,44863699223106355 ,448527385257058304,44852407221827993 ,448575035255816192,436227487811796992,436259855310262272,436459556148105216,436580772271882240,436590972928294913,436591779841081344,436591774392647681,436594262264664064,436595302272090113,436598691345203200,436599510287282176,436604066228371456,436606178060021760,436607059183038464,436610878960852992,436614308748075008,436621335108673536,436624252209758209,436636366651551744,436662520846512128,436864587854794752,436889789783347200,436947235759083520,436965276517924864,436985508749930497,437303413513408513,437339304315338752,437366411410022402,437352154798768128,437368320376209409,437374055025762304,437402882686877696,437402971421556737,437403036684935168,437541652341600256,437586175122219008,437654232364113920,437654857692889088]}},{"user.id":1,"_id":false})

def text_sample(user_id, sample_size = sample_size):
  "Read a sample of the messages of the user referred. You can specify the number of messages to get."
  return normalize.Normalizer.normalize(" ".join(sdb.get_text_sample(user_id, sample_size)))

def merge_matrices(pt, br):
  "Given the matrices for each language, join them together into X and y (classificatoin features and class)."
  X = numpy.vstack((pt, br))
  y = numpy.hstack((numpy.zeros(pt.shape[0]), numpy.ones(br.shape[0])))
  return (X, y)

def create_model(X, y):
  "Creates a trained model given the features matrix and their class."
  classifier = sklearn.svm.SVC(kernel='linear')
  classifier.fit(X, y)
  return classifier

def load_db_data(pt = pt, br = br):
  "Load the user data from the database."
  # First we get the text from the userlist above.
  # Then we extract the features, that are matrices.
  # Finally we join the matrices together to form the language features matrix.
  pt_samples = list( text_sample(u, sample_size) for u in pt )
  print pt_samples
  pt_features = ( word.features_of_text(sample, sample_size) for sample in pt_samples if sample )

  pt_matrix = numpy.vstack(tuple(pt_features))
  br_samples = ( text_sample(u, sample_size) for u in br )
  br_features = ( word.features_of_text(sample, sample_size) for sample in br_samples if sample )
  br_matrix = numpy.vstack(tuple(br_features))
  X, y = merge_matrices(pt_matrix, br_matrix)
  return X, y

def k_fold_test():
  k = 10
  fold = kfold.StratifiedKFold(y, k)

  sum_pred = 0.0
  for train_index, test_index in fold:
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]

    classifier = create_model(X_train, y_train)
    y_prediction = classifier.predict(X_test)
    print(y_prediction, y_test)
    r = sklearn.metrics.precision_score(y_test, y_prediction)
    print(r)
    sum_pred += r

  print(sum_pred / k)

def save_model(model, filename):
  with open(filename, "wb") as f:
    pickle.dump(model, f)

def load_model(filename):
  with open(filename, "rb") as f:
    return pickle.load(f)

def initialize_model():
    # Load the data from either a local file or get it from the DB if absent (and store the new model locally).
    try:
      model = load_model(model_file)
      return model

    except IOError:
      X, y = load_db_data(pt, br)
      model = create_model(X, y)
      save_model(model, model_file)

      return model

# def classify(messages, model = model):
def classify(messages, model):
  "Classify language variant present in all these messages."
  norm_text = normalize.Normalizer.normalize(" ".join(messages))
  features = word.features_of_text(norm_text, len(messages))
  classification = model.predict(features)
  return class_names[int(classification[0])]

# print classify(["lindo!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! porto porto porto http://t.co/HZlswsZuf3"], model)
# print classify(["E VC AINDA ACREDITA QUE ESTE GOVFDRL PT E SOCIAL??  #VemPraRua #VemPraURNA  RT @JoaoBolsoni:  povo? Esse? http://t.co/KGPAI0Ao59"], model)

# if __name__ == "__main__":
#   import cgi
#   form = cgi.FieldStorage()
#   print "Content-Type: application/json"
#   print                               # blank line, end of headers
#   try:
#     messages = form.getlist('message')
#     if len(messages) == 0:
#       print json.dumps({"errno": -1, "error": "No message found. Use parameter \"message\" to supply messages."})
#     else:
#       variant = classify(messages, model)
#       print json.dumps({"errno": 0, "variant": variant})
#   except ValueError:
#     print json.dumps({"errno": -1, "error": "No message found. Use parameter \"message\" to supply messages."})