import matplotlib
matplotlib.use('agg')
# matplotlib.use('module://kivy.garden.matplotlib.backend_kivy')

from point_separation import  build_multi_problem, multi_problem_iteration,load_data

from kivy.app import App
from kivy.uix.widget import Widget
from kivy.uix.button import Button
from kivy.uix.label import Label
from kivy.uix.image import Image
from kivy.properties import ObjectProperty, StringProperty, NumericProperty
from kivy.uix.boxlayout import BoxLayout
from kivy.clock import Clock, mainthread
from os.path import sep, expanduser, isdir, dirname, join
from kivy.garden.filebrowser import FileBrowser
from kivy.utils import platform

from kivy.uix.textinput import TextInput
from kivy.uix.screenmanager import ScreenManager, Screen, NoTransition
import numpy as np
import re
import threading
from matplotlib import pyplot as plt
import nifty2go as ift
import matplotlib.pyplot as plt
import time
from kivy.uix.progressbar import ProgressBar

class FloatInput(TextInput):

    pat = re.compile('[^0-9]')
    def insert_text(self, substring, from_undo=False):
        print substring
        pat = self.pat
        if '.' in self.text:
            s = re.sub(pat, '', substring)
        else:
            s = '.'.join([re.sub(pat, '', s) for s in substring.split('.', 1)])
        return super(FloatInput, self).insert_text(s, from_undo=from_undo)
class IntInput(TextInput):

    pat = re.compile('[^0-9]')
    def insert_text(self, substring, from_undo=False):
        pat = self.pat
        s = re.sub(pat, '', substring)

        return super(IntInput, self).insert_text(s, from_undo=from_undo)



class MyImage(BoxLayout):
    text = StringProperty('')
    source = StringProperty('')

    def reload(self):
        self.img.reload()

class MyAlphaWidget(BoxLayout):
    alpha = NumericProperty(None)
    pass
class IterationWidget(BoxLayout):
    iteration = NumericProperty(None)
    pass
class ResultsPathWidget(BoxLayout):

    pass

class DataPathWidget(BoxLayout):
    pass

class DisplayWidget(ScreenManager):
    def reload(self):
        for child in self.children:
            child.reload()
class ImageWidget(BoxLayout):
    def reload(self):
        self.image_widget.reload()

class MenuWidget(BoxLayout):
    pass
class SingleImageScreen(Screen):
    source = StringProperty('')
    def reload(self):
        self.img.reload()
    pass
class AllImageScreen(Screen):
    def reload(self):
        self.all.reload()
    pass
class AllImageWidget(BoxLayout):
    def reload(self):
        self.data.reload()
        self.points.reload()
        self.diffuse.reload()
        self.power.reload()
    pass

class ActionWidget(BoxLayout):
    pass
class DisplayChoiceWidget(BoxLayout):
    pass

class GlobalScreenManager(ScreenManager):
    pass

class MainScreen(Screen):
    pass

class FileScreen(Screen):
    pass
class PathScreen(Screen):
    def is_dir(self, directory, filename):
        return isdir(join(directory, filename))


class MyWidget(BoxLayout):
    image_widget = ObjectProperty(None)
    menu_widget = ObjectProperty(None)

    data_path = StringProperty(None)
    result_path = StringProperty(None)
    alpha = NumericProperty(None)


class MyPathBrowser(FileBrowser):
    pass

class MyFileBrowser(FileBrowser):
    filters = ['*.fits', '*.png', '*.jpg']
    pass


class SeparatorApp(App):
    stop = threading.Event()

    data_path = StringProperty(None)
    result_path = StringProperty(None)
    alpha = NumericProperty(None)
    data_image = StringProperty(None)
    diffuse_image = StringProperty(None)
    points_image = StringProperty(None)
    power_image = StringProperty(None)
    vmin = None
    vmax = None
    myEnergy = None
    iterations = 3
    user_path = ''
    reconstructing = False
    data_loaded = False

    def build(self):
        self.set_default()
        self.trigger = 0
        self.root = GlobalScreenManager()
        self.image_widget = self.root.main.image_widget.image_widget.image_widget
        self.root.transition = NoTransition()
        self.image_widget.transition = NoTransition()
        return self.root

    def set_default(self):
        self.data_path = ''
        self.result_path = ''
        self.alpha = 1.5
        self.path = ''
        self.data_image = self.path + 'placeholder.png'
        self.diffuse_image = self.path + 'placeholder.png'
        self.points_image = self.path + 'placeholder.png'
        self.power_image = self.path + 'placeholder.png'
        if platform == 'win':
            self.user_path = dirname(expanduser('~')) + sep + 'Documents'
        else:
            self.user_path = expanduser('~') + sep + 'Documents'

    def load_data(self, selection):
        print selection
        self.data_path = selection[0]
        threading.Thread(target=self.load_data_thread).start()
        self.root.current = 'main'

    def load_data_thread(self):
        self.data = load_data(self.data_path)
        self.vmin = np.log(self.data.min())
        self.max = np.log(self.data.max())
        self.plot_data()
        self.set_data_loaded(True)
        self.set_data_image()
        self.update_plots()

    def run_separation(self):
        if not self.reconstructing and self.data_loaded:
            threading.Thread(target= self.run_separation_thread).start()

    def run_separation_thread(self):
        self.set_reconstructing(True)
        self.myEnergy = build_multi_problem(self.data, self.alpha)
        self.plot_components(self.path)
        self.set_image_paths()
        self.update_plots()

        for i in range(self.iterations):
            self.myEnergy = multi_problem_iteration(self.myEnergy)
            self.plot_components(self.path)
            self.update_plots()
        self.set_reconstructing(False)

    def save_results(self):
        if self.myEnergy is not None:
            self.root.current = 'path'

    def select_path(self, path):
        self.result_path = path[0]
        threading.Thread(target=self.save_data_thread).start()
        self.root.current = 'main'

    def save_data_thread(self):
        np.savetxt(self.result_path + '/points.csv', np.exp(self.myEnergy.u.val))
        np.savetxt(self.result_path + '/diffus.csv', np.exp(self.myEnergy.s.val))
        self.plot_components(self.result_path)
    def set_result_path(self, path):
        self.result_path = path

    @mainthread
    def update_plots(self):
        self.image_widget.reload()
    @mainthread
    def set_reconstructing(self, reconstructing):
        self.reconstructing = reconstructing
    @mainthread
    def set_data_loaded(self, loaded):
        self.data_loaded = loaded

    def set_data_path(self, path):
        self.data_path = path
    @mainthread
    def set_data_image(self):

        self.data_image = self.path + 'data.png'
    @mainthread
    def set_image_paths(self):
        self.points_image = self.path + 'points.png'
        self.diffuse_image = self.path + 'diffuse.png'

    def plot_data(self):
        if self.data.shape[0] == 1:
            plt.imsave(self.path+'data.png', self.data[0], vmin=self.vmin, vmax=self.vmax)
        else:
            plt.imsave(self.path+ 'data.png', self.data/255.)

    def plot_components(self, path):
        diffuse = np.empty_like(self.data)
        points = np.empty_like(self.data)
        for i  in range(len(self.myEnergy)):
            diffuse[...,i] = np.exp(self.myEnergy[i].s.val)
            points[...,i] = np.exp(self.myEnergy[i].u.val)
        if len(self.myEnergy) == 1:
            plt.imsave(path+'diffuse.png', diffuse[...,0], vmin=self.vmin, vmax=self.vmax)
            plt.imsave(path+'points.png', points[...,0], vmin=self.vmin, vmax=self.vmax)
        else:
            plt.imsave(self.path+ 'diffuse.png', diffuse/255.)
            plt.imsave(self.path+ 'points.png', points/255.)


    def set_alpha(self, alpha):
        if alpha == '':
            pass
        else:
            self.alpha = alpha
    def set_iterations(self,iterations):
        if iterations == '':
            pass
        else:
            self.iterations = int(iterations)

    def on_stop(self):
        self.stop.set()


if __name__ == '__main__':
    plt.gray()
    SeparatorApp().run()