# ===============================================================================
# Copyright 2011 Jake Ross
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===============================================================================

# =============enthought library imports=======================
from
 __future__ import absolute_import
from __future__ import print_function
import six.moves.cPickle as pickle

from traits.api import Bool, Any, Instance, Button, Property, Event, on_trait_change
from traitsui.api import View, Item, Handler, HGroup

# ============= standard library imports
 ========================
# from threading import Thread
from threading import Event as TEvent
from numpy import linspace, argmin, argmax, random, asarray
import time
import os

# ============= local library imports  ==========================
from pychron.core.time_series.time_series import smooth
from pychron.image.cv_wrapper import grayspace, crop, get_focus_measure

# from pychron.image.cvwrapper import grayspace, get_focus_measure, crop, resize
from scipy.ndimage.measurements import variance
from scipy.ndimage.filters import generic_gradient_magnitude, sobel
from scipy.ndimage import sum as ndsum
from pychron.paths import paths
from pychron.managers.manager import Manager
from pychron.image.image import Image

# from pychron.machine_vision.focus_parameters import FocusParameters
# from pychron.image.image_editor import ImageEditor
from pychron.graph.graph import Graph
from pychron.mv.focus.focus_parameters import FocusParameters
from pychron.core.ui.image_editor import ImageEditor
from pychron.core.ui.gui import invoke_in_main_thread
from pychron.core.ui.thread import Thread


class ConfigureHandler(Handler):
    def closed(self, info, isok):
        if isok:
            info.object.dump_parameters()


class AutoFocusManager(Manager):
    """
    currently uses passive focus techniques
    see

    http://en.wikipedia.org/wiki/Autofocus

    """

    video = Any
    laser_manager = Any
    stage_controller = Any
    canvas = Any
    parameters = Instance(FocusParameters)
    configure_button = Button("configure")

    autofocus_button = Event
    autofocus_label = Property(depends_on="autofocusing")
    autofocusing = Bool

    # threading event for cancel signal
    _evt_autofocusing = None

    image = Instance(Image, ())

    graph = None

    def dump_parameters(self):
        p = os.path.join(paths.hidden_dir, "autofocus_configure")
        self.info("dumping parameters to {}".format(p))
        with open(p, "wb") as f:
            pickle.dump(self.parameters, f)

    def load_parameter(self):
        p = os.path.join(paths.hidden_dir, "autofocus_configure")
        if os.path.isfile(p):
            with open(p, "rb") as f:
                try:
                    params = pickle.load(f)
                    self.info("loading parameters from {}".format(p))

                    if not isinstance(params, FocusParameters):
                        self.info("out of date parameters file. using default")
                        params = FocusParameters()
                    return params

                except Exception as e:
                    print("autofocus load parameter", e)
                    return FocusParameters()
        else:
            return FocusParameters()

    def passive_focus(self, block=False, **kw):

        self._evt_autofocusing = TEvent()
        self._evt_autofocusing.clear()
        #        manager = self.laser_manager
        oper = self.parameters.operator
        self.info("passive focus. operator = {}".format(oper))

        g = self.graph
        if not g:
            g = Graph(
                plotcontainer_dict=dict(padding=10),
                window_x=0.70,
                window_y=20,
                window_width=325,
                window_height=325,
                window_title="Autofocus",
            )
            self.graph = g

        g.clear()

        g.new_plot(
            padding=[40, 10, 10, 40],
            xtitle="Z (mm)",
            ytitle="Focus Measure ({})".format(oper),
        )
        g.new_series()
        g.new_series()

        invoke_in_main_thread(self._open_graph)

        target = self._passive_focus
        self._passive_focus_thread = Thread(
            name="autofocus", target=target, args=(self._evt_autofocusing,), kwargs=kw
        )
        self._passive_focus_thread.start()
        if block:
            #             while 1:
            #                 if not self._passive_focus_thread.isRunning():
            #                     break
            #                 time.sleep(0.25)
            self._passive_focus_thread.join()

    def _open_graph(self):
        ui = self.graph.edit_traits()
        self.add_window(ui)

    def stop_focus(self):

        if self.stage_controller:
            self.stage_controller.stop()

        self.info("autofocusing stopped by user")

    def _passive_focus(self, stop_signal, set_zoom=True):
        """
        sweep z looking for max focus measure
        FMgrad= roberts or sobel (sobel removes noise)
        FMvar = intensity variance
        """

        self.autofocusing = True

        manager = self.laser_manager
        fstart = self.parameters.fstart
        fend = self.parameters.fend
        step_scalar = self.parameters.step_scalar
        zoom = self.parameters.zoom
        operator = self.parameters.operator

        steps = step_scalar * (max(fend, fstart) - min(fend, fstart)) + 1

        prev_zoom = None
        if set_zoom and manager is not None and zoom:
            motor = manager.get_motor("zoom")
            if motor:
                prev_zoom = motor.data_position
                self.info("setting zoom: {}".format(zoom))
                manager.set_motor("zoom", zoom, block=True)
                time.sleep(1.5)

        args = self._do_focusing(fstart, fend, steps, operator)

        if manager is not None:
            if prev_zoom is not None:
                self.info("returning to previous zoom: {}".format(prev_zoom))
                manager.set_motor("zoom", prev_zoom, block=True)

        if args:
            mi, fmi, ma, fma = args

            self.info(
                """passive focus results:Operator={}
ImageGradmin={} (z={})
ImageGradmax={}, (z={})""".format(
                    operator, mi, fmi, ma, fma
                )
            )

            focus_pos = fma
            self.graph.add_vertical_rule(focus_pos)
            self.graph.redraw()
            #            self.graph.add_vertical_rule(fma)

            self.info("calculated focus z= {}".format(focus_pos))

            #            if set_z:
            controller = self.stage_controller
            if controller is not None:
                if not stop_signal.isSet():
                    controller.single_axis_move("z", focus_pos, block=True)
                    controller._z_position = focus_pos
                    controller.z_progress = focus_pos

        self.autofocusing = False

    def _cancel_sweep(self, vo):
        if self._evt_autofocusing.isSet():
            # return to original velocity
            self.autofocusing = False
            self._reset_velocity(vo)
            return True

    def _reset_velocity(self, vo):
        if self.stage_controller:
            pdict = dict(velocity=vo, key="z")
            self.stage_controller.set_single_axis_motion_parameters(pdict=pdict)

    def _do_focusing(self, start, end, steps, operator):
        screen_roi = self._get_roi()
        self._add_focus_area_rect(*screen_roi)

        src = self._load_source()
        src = asarray(src)
        h, w, _d = src.shape

       