import functools
import logging
import os

import attr
from keras.optimizers import Adam
from keras import backend as K
import numpy as np

from lipnet.model import LipNet
from lipnet.core.decoders import Decoder
from lipnet.lipreading.helpers import labels_to_text
from lipnet.utils.spell import Spell
from lipnet.model2 import LipNet

import clairvoyance.core

import pkg_resources

from lipnet.lipreading.videos import Video

class Reader(clairvoyance.core.Reader):
    def __init__(self, config):
        self._log = logging.getLogger(self.__class__.__name__)
        self._image_width = config.image_width
        self._image_height = config.image_height
        self._image_channels = config.image_channels
        self._frames = config.frames
        self._predict_greedy = config.predict_greedy
        self._predict_beam_width = config.predict_beam_width
        if config.predict_dictionary is None:
            self._predict_dictionary = pkg_resources.resource_filename(__name__, os.path.join('..','common','dictionaries','{}.txt'.format(config.predict_dictionary_type)))
        else:
            self._predict_dictionary = config.predict_dictionary
        if config.weight_path is None:
            self._weight_path = pkg_resources.resource_filename(__name__, os.path.join('..','evaluation','models','{}.h5'.format(config.weight_type)))
        else:
            self._weight_path = config.weight_path

    @functools.lru_cache(maxsize=1)
    def _decoder(self):
        return Decoder(greedy=self._predict_greedy, beam_width=self._predict_beam_width,
                       postprocessors=[labels_to_text, Spell(path=self._predict_dictionary).sentence])

    @functools.lru_cache(maxsize=1)
    def _lipnet(self, c, w, h, n, absolute_max_string_len=32, output_size=28):
        lipnet = LipNet(img_c=c, img_w=w, img_h=h, frames_n=n,
                        absolute_max_string_len=absolute_max_string_len, output_size=output_size)

        adam = Adam(lr=0.0001, beta_1=0.9, beta_2=0.999, epsilon=1e-08)

        lipnet.model.compile(loss={'ctc': lambda y_true, y_pred: y_pred}, optimizer=adam)
        lipnet.model.load_weights(self._weight_path)
        return lipnet

    def warmup(self):
        self._lipnet(c=self._image_channels, w=self._image_width, h=self._image_height, n=self._frames)

    def do(self, frames):
        if K.image_data_format() == 'channels_first':
            img_c, frames_n, img_w, img_h = frames.shape
        else:
            frames_n, img_w, img_h, img_c = frames.shape

        assert (self._image_width, self._image_height, self._image_channels) == (img_w, img_h, img_c), '{} != {}'.format((self._image_width, self._image_height, self._image_channels), (img_w, img_h, img_c))

        X_data       = np.array([frames]).astype(np.float32) / 255
        input_length = np.array([len(frames)])

        y_pred         = self._lipnet(c=img_c, w=img_w, h=img_h, n=frames_n).predict(X_data)
        result         = self._decoder().decode(y_pred, input_length)[0]
        return result

    @attr.s
    class Config:
        image_width = attr.ib(default=100)
        image_height = attr.ib(default=50)
        image_channels = attr.ib(default=3)
        frames = attr.ib(default=256)
        predict_greedy = attr.ib(default=False)
        predict_beam_width = attr.ib(default=200)
        predict_dictionary_type = attr.ib(default='big')
        weight_type = attr.ib(default='overlapped-weights368')
        predict_dictionary = attr.ib(default=None)
        weight_path = attr.ib(default=None)


from __future__ import print_function, division
from serial_device2 import find_serial_device_ports
from modular_device import ModularDevices
# from bioshake_device import BioshakeDevice, BioshakeError
from exceptions import Exception
import os
import time
import yaml
import argparse
from numpy.polynomial.polynomial import polyfit,polyadd,Polynomial
from mettler_toledo_device import MettlerToledoDevice
import csv
import copy
import numpy

# try:
#     from pkg_resources import get_distribution, DistributionNotFound
#     _dist = get_distribution('hybridizer')
#     # Normalize case for Windows systems
#     dist_loc = os.path.normcase(_dist.location)
#     here = os.path.normcase(__file__)
#     if not here.startswith(os.path.join(dist_loc, 'hybridizer')):
#         # not installed, but there is another version that *is*
#         raise DistributionNotFound
# except (ImportError,DistributionNotFound):
#     __version__ = None
# else:
#     __version__ = _dist.version

DEBUG = True
BAUDRATE = 9600

class HybridizerError(Exception):
    def __init__(self,value):
        self.value = value
    def __str__(self):
        return repr(self.value)

class Hybridizer(object):
    '''
    This Python package (hybridizer) creates a class named Hybridizer to
    communcate with and control the Janelia Hybridizer. The hybridizer
    uses two hardware control devices, the mixed_signal_controller
    modular_device, and the bioshake_device. The
    mixed_signal_controller both switches the valves and reads the
    analog signals from the cylinder hall effect sensors. The
    bioshake_device controls the heater/shaker.
    Example Usage:

    hyb = Hybridizer('example_config.yaml')
    hyb.run_protocol()
    '''

    # def __init__(self,
    #              config_file_path,
    #              *args,**kwargs):
    def __init__(self,
                 config_file_path,
                 calibration_file_path,
                 *args,**kwargs):
        if 'debug' in kwargs:
            self._debug = kwargs['debug']
        else:
            kwargs.update({'debug': DEBUG})
            self._debug = DEBUG
        with open(config_file_path,'r') as config_stream:
            self._config = yaml.load(config_stream)
        with open(calibration_file_path,'r') as calibration_stream:
            self._calibration = yaml.load(calibration_stream)
        self._valves = self._config['head']
        self._valves.update(self._config['manifold'])
        ports = find_serial_device_ports(debug=self._debug)
        self._debug_print('Found serial devices on ports ' + str(ports))
        self._debug_print('Identifying connected devices (may take some time)...')
        self._balance = MettlerToledoDevice()
        ports.remove(self._balance.get_port())
        # try:
        #     self._bsc = BioshakeDevice()
        # except RuntimeError:
        #     # try one more time
        #     self._bsc = BioshakeDevice()
        # self._debug_print('Found bioshake device on port ' + str(self._bsc.get_port()))
        # ports.remove(self._bsc.get_port())
        # self._SHAKE_SPEED_MIN = self._bsc.get_shake_speed_min()
        # self._SHAKE_SPEED_MAX = self._bsc.get_shake_speed_max()
        # self._SHAKE_DURATION_MIN = 10
        # self._SHAKE_ATTEMPTS = 2
        # self._POST_SHAKE_OFF_DURATION = 5
        modular_devices = ModularDevices(try_ports=ports)

        try:
            msc_dict = modular_devices['mixed_signal_controller']
        except KeyError:
            raise HybridizerError('Could not find mixed_signal_controller. Check connections and permissions.')
        if len(msc_dict) > 1:
            raise HybridizerError('More than one mixed_signal_controller found. Only one should be connected.')
        self._msc = msc_dict[msc_dict.keys()[0]]
        self._debug_print('Found mixed_signal_controller on port ' + str(self._msc.get_port()))

        # self._adc_values_min = None
        # self._adc_values_max = None
        self._adc_sample_count = 21
        self._fill_duration_all_cylinders = 250
        self._fill_duration_one_cylinder = 100
        self._volume_crossover = 6
        self._volume_threshold_initial = 1.0

    def prime_system(self):
        self._setup()
        self._debug_print('priming system...')
        manifold = self._config['manifold']
        chemicals = manifold.keys()
        try:
            chemicals.remove('aspirate')
        except ValueError:
            pass
        try:
            chemicals.remove('separate')
        except ValueError:
            pass
        self._set_valves_on(['separate','aspirate'])
        for chemical in chemicals:
            self._prime_chemical(chemical,self._config['system_prime_count'])
        self._set_all_valves_off()
        self._debug_print('priming finished!')

    def run_protocol(self):
        self._setup()
        self.protocol_start_time = time.time()
        self._debug_print('running protocol...')
        self._set_valves_on(['separate','aspirate'])
        for chemical_info in self._config['protocol']:
            chemical = chemical_info['chemical']
            try:
                prime_count = chemical_info['prime_count']
            except KeyError:
                prime_count = 1
            try:
                dispense_count = chemical_info['dispense_count']
            except KeyError:
                dispense_count = 1
            try:
                shake_speed = chemical_info['shake_speed']
            except KeyError:
                shake_speed = None
            try:
                shake_duration = chemical_info['shake_duration']
            except KeyError:
                shake_duration = None
            try:
                post_shake_duration = chemical_info['post_shake_duration']
            except KeyError:
                post_shake_duration = 0
            try:
                separate = chemical_info['separate']
            except KeyError:
                separate = False
            try:
                aspirate = chemical_info['aspirate']
            except KeyError:
                aspirate = True
            try:
                temperature = chemical_info['temperature']
            except KeyError:
                temperature = None
            try:
                repeat = chemical_info['repeat']
            except KeyError:
                repeat = 0
            self._run_chemical(chemical,
                               prime_count,
                               dispense_count,
                               shake_speed,
                               shake_duration,
                               post_shake_duration,
                               separate,
                               aspirate,
                               temperature,
                               repeat)
        self._set_all_valves_off()
        self.protocol_end_time = time.time()
        protocol_run_time = self.protocol_end_time - self.protocol_start_time
        self._debug_print('protocol finished! it took ' + str(round(protocol_run_time/60)) + ' mins to run.')

    def _setup(self):
        # self._bsc.reset_device()
        self._msc.remove_all_set_fors()
        self._set_all_valves_off()
        self._set_valves_on(['primer','quad1','quad2','quad3','quad4','quad5','quad6'])
        self._debug_print('setting up for ' + str(self._config['setup_duration']) + 's...')
        time.sleep(self._config['setup_duration'])
        self._set_all_valves_off()
        self._debug_print('setup finished!')

    def _prime_chemical(self,chemical,prime_count):
        if prime_count > 0:
            self._set_valve_on(chemical)
        for i in range(prime_count):
            self._set_valves_on(['primer','system'])
            self._debug_print('priming ' + chemical + ' for ' + str(self._config['prime_duration']) + 's ' + str(i+1) + '/' + str(prime_count) + '...')
            time.sleep(self._config['prime_duration'])
            self._set_valves_off(['system'])
            self._debug_print('emptying ' + chemical + ' for ' + str(self._config['prime_aspirate_duration']) + 's ' + str(i+1) + '/' + str(prime_count) + '...')
            time.sleep(self._config['prime_aspirate_duration'])
            self._set_valve_off('primer')
        if prime_count > 0:
            self._set_valve_off(chemical)

    def _run_chemical(self,
                      chemical,
                      prime_count=1,
                      dispense_count=1,
                      shake_speed=None,
                      shake_duration=None,
                      post_shake_duration=0,
                      separate=False,
                      aspirate=True,
                      temp_target=None,
                      repeat=0):
        if (chemical not in self._valves):
            raise HybridizerError(chemical + ' is not listed as part of the manifold in the config file!')
        if repeat < 0:
            repeat = 0
        run_count = repeat + 1
        if temp_target is not None:
            self._debug_print('turning on temperature control for ' + chemical + '...')
            self._bsc.temp_on(temp_target)
            temp_actual = self._bsc.get_temp_actual()
            self._debug_print('actual temperature: ' + str(temp_actual) + ',target temperature: ' + str(temp_target))
            while abs(temp_target - temp_actual) > 0.5:
                time.sleep(1)
                temp_actual = self._bsc.get_temp_actual()
                self._debug_print('actual temperature: ' + str(temp_actual) + ',target temperature: ' + str(temp_target))
            self._debug_print()
        self._prime_chemical(chemical,prime_count)
        for run in range(run_count):
            self._debug_print('running ' + chemical + ' ' + str(run+1) + '/' + str(run_count) + '...')
            self._set_valve_on(chemical)
            self._set_valves_on(['quad1','quad2','quad3','quad4','quad5','quad6','aspirate'])
            for i in range(dispense_count):
                if i > 0:
                    dispense_shake_duration = self._config['inter_dispense_shake_duration']
                    if dispense_shake_duration < self._SHAKE_DURATION_MIN:
                        dispense_shake_duration = self._SHAKE_DURATION_MIN
                    dispense_shake_speed = self._shake_on(self._config['inter_dispense_shake_speed'])
                    self._debug_print('shaking at ' + str(dispense_shake_speed) + 'rpm for ' + str(dispense_shake_duration) + 's...')
                    time.sleep(dispense_shake_duration)
                    self._shake_off(dispense_shake_speed)
                self._set_valve_on('system')
                self._debug_print('loading ' + chemical + ' into syringes for ' + str(self._config['load_duration']) + 's ' + str(i+1) + '/' + str(dispense_count) + '...')
                time.sleep(self._config['load_duration'])
                self._set_valve_off('system')
                self._debug_print('dispensing ' + chemical + ' into microplate for ' + str(self._config['dispense_duration']) + 's ' + str(i+1) + '/' + str(dispense_count) + '...')
                time.sleep(self._config['dispense_duration'])
            self._set_valves_off(['quad1','quad2','quad3','quad4','quad5','quad6'])
            if not ((shake_duration is None) or (shake_duration <= 0)):
                actual_shake_duration = shake_duration
                if shake_duration < self._SHAKE_DURATION_MIN:
                    actual_shake_duration = self._SHAKE_DURATION_MIN
                actual_shake_speed = self._shake_on(shake_speed)
                self._debug_print('shaking at ' + str(actual_shake_speed) + 'rpm for ' + str(actual_shake_duration) + 's...')
                time.sleep(actual_shake_duration)
                self._shake_off(actual_shake_speed)
            if (post_shake_duration > 0):
                self._debug_print('waiting post shake for ' + str(post_shake_duration) + 's...')
                time.sleep(post_shake_duration)
            if separate:
                separate_shake_speed = self._shake_on(self._config['separate_shake_speed'])
                self._set_valve_off('separate')
                self._debug_print('separating ' + chemical + ' for ' + str(self._config['chemical_separate_duration']) + 's...')
                time.sleep(self._config['chemical_separate_duration'])
                self._set_valve_on('separate')
                self._shake_off(separate_shake_speed)
            if aspirate:
                aspirate_shake_speed = self._shake_on(self._config['aspirate_shake_speed'])
                self._set_valve_off('aspirate')
                self._debug_print('aspirating ' + chemical + ' from microplate for ' + str(self._config['chemical_aspirate_duration']) + 's...')
                time.sleep(self._config['chemical_aspirate_duration'])
                self._set_valve_on('aspirate')
                self._shake_off(aspirate_shake_speed)
            self._set_valve_off(chemical)
            self._debug_print(chemical + ' finished!')
            self._debug_print()
        if temp_target is not None:
            self._debug_print('turning off temperature control for ' + chemical + '...')
            try:
                self._bsc.temp_off()
            except BioshakeError:
                pass
            self._debug_print()

    def _shake_on(self,shake_speed):
        if (shake_speed is None) or (shake_speed < self._SHAKE_SPEED_MIN):
            shake_speed = 0
        elif shake_speed > self._SHAKE_SPEED_MAX:
            shake_speed = self._SHAKE_SPEED_MAX
        if shake_speed != 0:
            shook = False
            shake_try = 0
            while (not shook) and (shake_try < self._SHAKE_ATTEMPTS):
                shake_try += 1
                try:
                    self._bsc.shake_on(shake_speed)
                    shook = True
                except BioshakeError:
                    self._debug_print('bioshake_device.get_error_list(): ' + str(self._bsc.get_error_list()))
                    self._debug_print('BioshakeError! Resetting for ' + str(self._config['setup_duration']) + 's and trying again...')
                    self._bsc.reset_device()
                    time.sleep(self._config['setup_duration'])
        return shake_speed

    def _shake_off(self,shake_speed):
        if shake_speed != 0:
            shook = False
            shake_try = 0
            while (not shook) and (shake_try < self._SHAKE_ATTEMPTS):
                shake_try += 1
                try:
                    self._bsc.shake_off()
                    shook = True
                except BioshakeError:
                    self._debug_print('bioshake_device.get_error_list(): ' + str(self._bsc.get_error_list()))
                    self._debug_print('BioshakeError! Resetting for ' + str(self._config['setup_duration']) + 's and trying again...')
                    self._bsc.reset_device()
                    time.sleep(self._config['setup_duration'])
            time.sleep(self._POST_SHAKE_OFF_DURATION)

    def _debug_print(self,*args):
        if self._debug:
            print(*args)

    def _set_valve_on(self,valve_key):
        try:
            valve = self._valves[valve_key]
            channels = [valve['channel']]
            self._msc.set_channels_on(channels)
        except KeyError:
            raise HybridizerError('Unknown valve: ' + str(valve_key) + '. Check yaml config file for errors.')

    def _set_valves_on(self,valve_keys):
        try:
            channels = [self._valves[valve_key]['channel'] for valve_key in valve_keys]
            self._msc.set_channels_on(channels)
        except KeyError:
            raise HybridizerError('Unknown valve: ' + str(valve_key) + '. Check yaml config file for errors.')

    def _set_valve_off(self,valve_key):
        try:
            valve = self._valves[valve_key]
            channels = [valve['channel']]
            self._msc.set_channels_off(channels)
        except KeyError:
            raise HybridizerError('Unknown valve: ' + str(valve_key) + '. Check yaml config file for errors.')

    def _set_valves_off(self,valve_keys):
        try:
            channels = [self._valves[valve_key]['channel'] for valve_key in valve_keys]
            self._msc.set_channels_off(channels)
        except KeyError:
            raise HybridizerError('Unknown valve: ' + str(valve_key) + '. Check yaml config file for errors.')

    def _set_all_valves_off(self):
        valve_keys = self._get_valves()
        self._set_valves_off(valve_keys)

    def _get_valves(self):
        valve_keys = self._valves.keys()
        valve_keys.sort()
        return valve_keys

    def _get_adc_values_filtered(self):
        adc_values = None
        for sample_n in range(self._adc_sample_count):
            sample_values = self._msc.get_analog_inputs_filtered()
            if adc_values is None:
                adc_values = numpy.array([sample_values],int)
            else:
                adc_values = numpy.append(adc_values,[sample_values],axis=0)
            time.sleep(0.1)
        adc_values_filtered = numpy.median(adc_values,axis=0)
        adc_values_filtered = adc_values_filtered.astype(int)
        return adc_values_filtered

    # def _store_adc_values_min(self):
    #     self._adc_values_min = {}
    #     adc_values_filtered = self._get_adc_values_filtered()
    #     head_valves = self._config['head']
    #     for head_valve in head_valves:
    #         try:
    #             ain = self._config['head'][head_valve]['analog_inputs']['low']
    #             adc_value = adc_values_filtered[ain]
    #             self._adc_values_min[head_valve] = adc_value
    #         except KeyError:
    #             continue
    #     self._debug_print(self._adc_values_min)

    # def _set_valves_on_until(self,valve_keys,volume):
    #     for valve_key in valve_keys:
    #         valve = self._valves[valve_key]
    #         channels = [valve['channel']]
    #         adc_value_goal,ain = self._volume_to_adc_and_ain(valve_key,volume)
    #         set_until_index = self._msc.set_channels_on_until(channels,ain,adc_value_goal)
    #     while not self._msc.are_all_set_untils_complete():
    #         self._debug_print('Waiting...')
    #         time.sleep(1)
    #     self._msc.remove_all_set_untils()
    #     for valve_key in valve_keys:
    #         valve = self._valves[valve_key]
    #         adc_value_goal,ain = self._volume_to_adc_and_ain(valve_key,volume)
    #         adc_value = self._msc.get_analog_input(ain)
    #         volume = self._adc_to_volume_low(valve_key,adc_value)

    def _set_valves_on_until(self,valve_keys,volume):
        channels = []
        adc_value_goals = []
        ains = []
        jumps = {}
        valve_keys_copy = copy.copy(valve_keys)
        for valve_key in valve_keys:
            valve = self._valves[valve_key]
            channels.append(valve['channel'])
            adc_value_goal,ain = self._volume_to_adc_and_ain(valve_key,volume)
            adc_value_goals.append(adc_value_goal)
            ains.append(ain)
            jumps[valve_key] = 0

        volume_goal_initial = volume - self._volume_threshold_initial
        fill_duration_initial_max = 0
        fill_durations_initial = []
        if volume_goal_initial >= self._volume_threshold_initial/2:
            for valve_key in valve_keys:
                fill_duration_initial = self._volume_to_fill_duration(valve_key,volume_goal_initial)
                fill_durations_initial.append(fill_duration_initial)
            fill_duration_initial_min = min(fill_durations_initial)
            self._msc.set_channels_on_for(channels,fill_duration_initial_min)
            while not self._msc.are_all_set_fors_complete():
                self._debug_print('Waiting...')
                time.sleep(0.5 + fill_duration_initial_min/1000)
            self._msc.remove_all_set_fors()

        fill_duration_base = self._fill_duration_one_cylinder
        fill_duration_per_cylinder = (self._fill_duration_all_cylinders - self._fill_duration_one_cylinder)//(len(channels)-1)
        while len(channels) > 0:
            fill_duration = fill_duration_base + fill_duration_per_cylinder*(len(channels)-1)
            self._debug_print("Setting {0} valves on for {1}ms".format(valve_keys_copy,fill_duration))
            self._msc.set_channels_on_for(channels,fill_duration)
            while not self._msc.are_all_set_fors_complete():
                self._debug_print('Waiting...')
                time.sleep(fill_duration/1000)
            self._msc.remove_all_set_fors()
            adc_values_filtered = self._get_adc_values_filtered()
            ains_copy = copy.copy(ains)
            for ain in ains_copy:
                index = ains.index(ain)
                valve_key_copy = valve_keys_copy[index]
                jumps[valve_key_copy] += 1
                if adc_values_filtered[ain] >= adc_value_goals[index]:
                    channels.pop(index)
                    adc_value_goals.pop(index)
                    ains.pop(index)
                    valve_keys_copy.pop(index)
        adc_values_filtered = self._get_adc_values_filtered()
        final_adc_values = []
        jumps_list = []
        for valve_key in valve_keys:
            adc_value_goal,ain = self._volume_to_adc_and_ain(valve_key,volume)
            adc_value = adc_values_filtered[ain]
            final_adc_values.append(adc_value)
            jumps_list.append(jumps[valve_key])
            # volume = self._adc_to_volume_low(valve_key,adc_value)
        return final_adc_values,jumps_list

    # def _volume_to_adc_and_ain(self,valve_key,volume):
    #     valve = self._valves[valve_key]
    #     if volume <= self._config['volume_crossover']:
    #         ain = valve['analog_inputs']['low']
    #     else:
    #         ain = valve['analog_inputs']['low']
    #     if volume > self._config['volume_max']:
    #         raise HybridizerError('Asking for volume greater than the max volume of {0}!'.format(self._config['volume_max']))
    #     if volume <= self._config['volume_crossover']:
    #         poly = Polynomial(self._config['poly_coefficients']['volume_to_adc_low'])
    #         adc_value = int(round(poly(volume)))
    #         adc_value += self._adc_values_min[valve_key]
    #         self._debug_print("valve: {0}, adc_value: {1}, ain: {2}".format(valve_key,adc_value,ain))
    #         return adc_value,ain
    #     else:
    #         return 400,ain

    def _volume_to_adc_and_ain(self,valve_key,volume):
        valve = self._valves[valve_key]
        if volume <= self._volume_crossover:
            ain = valve['analog_inputs']['low']
        else:
            ain = valve['analog_inputs']['high']
        if volume > self._config['volume_max']:
            raise HybridizerError('Asking for volume greater than the max volume of {0}!'.format(self._config['volume_max']))
        if volume <= self._volume_crossover:
            poly = Polynomial(self._calibration[valve_key]['volume_to_adc_low'])
            adc_value = int(round(poly(volume)))
            self._debug_print("valve: {0}, adc_value: {1}, ain: {2}".format(valve_key,adc_value,ain))
            return adc_value,ain
        else:
            poly = Polynomial(self._calibration[valve_key]['volume_to_adc_high'])
            adc_value = int(round(poly(volume)))
            self._debug_print("valve: {0}, adc_value: {1}, ain: {2}".format(valve_key,adc_value,ain))
            return adc_value,ain

    def _volume_to_fill_duration(self,valve_key,volume):
        poly = Polynomial(self._calibration[valve_key]['volume_to_fill_duration'])
        fill_duration = int(round(poly(volume)))
        return fill_duration

    # def _adc_to_volume_low(self,valve_key,adc_value):
    #     valve = self._valves[valve_key]
    #     adc_value -= self._adc_values_min[valve_key]
    #     poly = Polynomial(self._config['poly_coefficients']['adc_to_volume_low'])
    #     volume = poly(adc_value)
    #     self._debug_print("valve: {0}, adc_value: {1}, volume: {2}".format(valve_key,adc_value,volume))
    #     return volume

    def run_dispense_tests(self):
        self._debug_print('pre setup sequence...')
        valves = ['quad1','quad2','quad3','quad4','quad5','quad6']
        self._set_valve_on('aspirate')
        self._set_valve_on('system')
        self._set_valves_on(valves)
        time.sleep(10)
        self._set_valve_off('system')
        time.sleep(10)
        self._set_valve_off('aspirate')
        time.sleep(20)
        self._setup()
        self._set_valve_on('aspirate')
        time.sleep(10)
        # self._debug_print('zeroing hall effect sensors...')
        # self._store_adc_values_min()
        self._debug_print('zeroing balance...')
        self._balance.zero()
        self._debug_print('running dispense tests...')
        timestr = time.strftime("%Y%m%d-%H%M%S")
        data_file = open(timestr+'.csv','w')
        data_writer = csv.writer(data_file)
        header = ['dispense_goal','initial_weight']
        valve_adc = [valve+'_adc' for valve in valves]
        header.extend(valve_adc)
        valve_jumps = [valve+'_jumps' for valve in valves]
        header.extend(valve_jumps)
        header.extend(valves)
        data_writer.writerow(header)
        # dispense_goals = [5,4,3,2,1]
        # dispense_goals = [1,0.75,0.5,0.25]
        dispense_goals = [5,1,0.5]
        # dispense_goals = [2,1]
        # run_count = 10
        run_count = 1
        # dispense_goals = [1]
        # run_count = 2
        for dispense_goal in dispense_goals:
            for run in range(run_count):
                self._set_valve_on('aspirate')
                time.sleep(2)
                self._debug_print('dispense_goal: {0}, run: {1} out of {2}'.format(dispense_goal,run+1,run_count))
                row_data = []
                row_data.append(dispense_goal)
                initial_weight = self._balance.get_weight()[0]
                self._debug_print('initial_weight: {0}'.format(initial_weight))
                row_data.append(initial_weight)
                self._set_valve_on('system')
                time.sleep(2)
                final_adc_values,jumps = self._set_valves_on_until(valves,dispense_goal)
                row_data.extend(final_adc_values)
                row_data.extend(jumps)
                self._set_valve_off('system')
                time.sleep(4)
                weight_prev = initial_weight
                for valve in valves:
                    self._debug_print('Dispensing {0}'.format(valve))
                    self._set_valve_on(valve)
                    time.sleep(4)
                    self._set_valve_off(valve)
                    time.sleep(2)
                    weight_total = self._balance.get_weight()[0]
                    weight = weight_total - weight_prev
                    self._debug_print('{0} measured {1}'.format(valve,weight))
                    row_data.append(weight)
                    weight_prev = weight_total
                self._set_valve_off('aspirate')
                self._debug_print('aspirating...')
                time.sleep(20)
                self._set_all_valves_off()
                data_writer.writerow(row_data)
        data_file.close()

    def run_calibration(self):
        self._debug_print('pre setup sequence...')
        valves = ['quad1','quad2','quad3','quad4','quad5','quad6']
        self._set_valve_on('aspirate')
        self._set_valve_on('system')
        self._set_valves_on(valves)
        time.sleep(10)
        self._set_valve_off('system')
        time.sleep(10)
        self._set_valve_off('aspirate')
        time.sleep(20)
        self._setup()
        self._set_valve_on('aspirate')
        time.sleep(10)
        # self._debug_print('zeroing hall effect sensors...')
        # self._store_adc_values_min()
        self._debug_print('zeroing balance...')
        self._balance.zero()
        self._debug_print('running calibration...')
        timestr = time.strftime("%Y%m%d-%H%M%S")
        data_file = open(timestr+'.csv','w')
        data_writer = csv.writer(data_file)
        header = ['fill_duration','initial_weight']
        valve_adc_low = [valve+'_adc_low' for valve in valves]
        header.extend(valve_adc_low)
        valve_adc_high = [valve+'_adc_high' for valve in valves]
        header.extend(valve_adc_high)
        header.extend(valves)
        data_writer.writerow(header)
        duration_inc = 250
        duration_max = 10000
        fill_durations = range(duration_inc,duration_max+duration_inc,duration_inc)
        run_count = 3
        for run in range(run_count):
            for fill_duration in fill_durations:
                self._set_valve_on('aspirate')
                time.sleep(2)
                self._debug_print('fill_duration: {0}, run: {1} out of {2}'.format(fill_duration,run+1,run_count))
                row_data = []
                row_data.append(fill_duration)
                initial_weight = self._balance.get_weight()[0]
                self._debug_print('initial_weight: {0}'.format(initial_weight))
                row_data.append(initial_weight)
                self._set_valve_on('system')
                time.sleep(2)
                channels = []
                adc_low_ain = []
                adc_high_ain = []
                for valve_key in valves:
                    valve = self._valves[valve_key]
                    channels.append(valve['channel'])
                    adc_low_ain.append(valve['analog_inputs']['low'])
                    adc_high_ain.append(valve['analog_inputs']['high'])
                self._msc.set_channels_on_for(channels,fill_duration)
                while not self._msc.are_all_set_fors_complete():
                    self._debug_print('Waiting...')
                    time.sleep(fill_duration/1000)
                self._msc.remove_all_set_fors()
                adc_values_filtered = self._get_adc_values_filtered()
                adc_low_values = [adc_values_filtered[ain] for ain in adc_low_ain]
                adc_high_values = [adc_values_filtered[ain] for ain in adc_high_ain]
                row_data.extend(adc_low_values)
                row_data.extend(adc_high_values)
                self._set_valve_off('system')
                time.sleep(4)
                weight_prev = initial_weight
                for valve in valves:
                    self._debug_print('Dispensing {0}'.format(valve))
                    self._set_valve_on(valve)
                    time.sleep(4)
                    self._set_valve_off(valve)
                    time.sleep(2)
                    weight_total = self._balance.get_weight()[0]
                    weight = weight_total - weight_prev
                    self._debug_print('{0} measured {1}'.format(valve,weight))
                    row_data.append(weight)
                    weight_prev = weight_total
                self._set_valve_off('aspirate')
                self._debug_print('aspirating...')
                time.sleep(20)
                self._set_all_valves_off()
                data_writer.writerow(row_data)
        data_file.close()

# -----------------------------------------------------------------------------------------
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("config_file_path",help="Path to yaml config file.")
    parser.add_argument("calibration_file_path",help="Path to yaml calibration file.")

    args = parser.parse_args()
    config_file_path = args.config_file_path
    calibration_file_path = args.calibration_file_path
    print("Config File Path: {0}".format(config_file_path))
    print("Calibration File Path: {0}".format(calibration_file_path))

    debug = True
    # hyb = Hybridizer(debug=debug,config_file_path=config_file_path)
    hyb = Hybridizer(debug=debug,
                     config_file_path=config_file_path,
                     calibration_file_path=calibration_file_path)
    # hyb.run_protocol()
    hyb.run_dispense_tests()
    # hyb.run_calibration()


import tensorflow as tf
from .displacement import batch_displacement_warp2d
from .ops import conv2d, upconv2d
from .utils import save_image
import medpy.io.save
import os
import numpy as np

class VectorCNN(object):

    def __init__(self, name, is_train):
        self.name = name
        self.is_train = is_train
        self.reuse = None

    def __call__(self, x):
        non_linearity = tf.nn.elu

        with tf.variable_scope(self.name, reuse=self.reuse):
            conv1_a = conv2d(x, 'conv1_a', 16, 3, 1, 'SAME',
                             True, non_linearity, self.is_train)
            conv1_b = conv2d(conv1_a, 'conv1_b', 16, 3, 1, 'SAME',
                             True, None, self.is_train)
            act1 = non_linearity(conv1_b)
            pool1 = tf.nn.avg_pool(act1, [1, 2, 2, 1], [1, 2, 2, 1], 'SAME')

            conv2_a = conv2d(pool1, 'conv2_a', 32, 3, 1, 'SAME',
                             True, non_linearity, self.is_train)
            conv2_b = conv2d(conv2_a, 'conv2_b', 32, 3, 1, 'SAME',
                             True, None, self.is_train)
            act2 = non_linearity(conv2_b)
            pool2 = tf.nn.avg_pool(act2, [1, 2, 2, 1], [1, 2, 2, 1], 'SAME')

            conv3_a = conv2d(pool2, 'conv3_a', 64, 3, 1, 'SAME',
                             True, non_linearity, self.is_train)
            conv3_b = conv2d(conv3_a, 'conv3_b', 64, 3, 1, 'SAME',
                             True, None, self.is_train)
            act3 = non_linearity(conv3_b)
            pool3 = tf.nn.avg_pool(act3, [1, 2, 2, 1], [1, 2, 2, 1], 'SAME')

            conv4_a = conv2d(pool3, 'conv4_a', 128, 3, 1, 'SAME',
                             True, non_linearity, self.is_train)
            conv4_b = conv2d(conv4_a, 'conv4_b', 128, 3, 1, 'SAME',
                             True, None, self.is_train)
            act4 = non_linearity(conv4_b)
            drop4 = tf.layers.dropout(act4, rate=0.5,
                                      training=self.is_train,
                                      name='drop4')

            deconv4 = upconv2d(drop4, 2, 'deconv4', 64, 3, 1, 'SAME',
                               True, None, self.is_train)
            concat5 = non_linearity(tf.add(deconv4, conv3_b))
            conv5_a = conv2d(concat5, 'conv5_a', 64, 3, 1, 'SAME',
                             True, non_linearity, self.is_train)
            conv5_b = conv2d(conv5_a, 'conv5_b', 64, 3, 1, 'SAME',
                             True, non_linearity, self.is_train)
            drop5 = tf.layers.dropout(conv5_b, rate=0.5,
                                      training=self.is_train,
                                      name='drop5')

            deconv5 = upconv2d(drop5, 2, 'deconv5', 32, 3, 1, 'SAME',
                               True, None, self.is_train)
            concat6 = non_linearity(tf.add(deconv5, conv2_b))
            conv6_a = conv2d(concat6, 'conv6_a', 32, 3, 1, 'SAME',
                             True, non_linearity, self.is_train)
            conv6_b = conv2d(conv6_a, 'conv6_b', 32, 3, 1, 'SAME',
                             True, non_linearity, self.is_train)
            drop6 = tf.layers.dropout(conv6_b, rate=0.5,
                                      training=self.is_train,
                                      name='drop6')

            deconv6 = upconv2d(drop6, 2, 'deconv6', 16, 3, 1, 'SAME',
                               True, None, self.is_train)
            concat7 = non_linearity(tf.add(deconv6, conv1_b))
            conv7_a = conv2d(concat7, 'conv7_a', 16, 3, 1, 'SAME',
                             True, non_linearity, self.is_train)
            conv7_b = conv2d(conv7_a, 'conv7_b', 16, 3, 1, 'SAME',
                             True, non_linearity, self.is_train)
            conv_out = conv2d(conv7_b, 'conv_out', 2, 3, 1, 'SAME',
                              True, None, self.is_train)

        if self.reuse is None:
            self.var_list = tf.get_collection(
                tf.GraphKeys.GLOBAL_VARIABLES, scope=self.name)
            self.saver = tf.train.Saver(
                var_list=self.var_list, max_to_keep=100)
            self.reuse = True

        return conv_out

    def save(self, sess, ckpt_path, step):
        self.saver.save(sess, ckpt_path, global_step=step)

    def restore(self, sess, ckpt_path):
        self.saver.restore(sess, ckpt_path)

class ACRegNet(object):

    def __init__(self, sess, name, im_size):
        self.sess = sess
        self.name = name

        im_shape = [1] + im_size + [1]

        self.x = tf.placeholder(tf.float32, im_shape)
        self.y = tf.placeholder(tf.float32, im_shape)
        self.xy = tf.concat([self.x, self.y], axis=3)

        if im_shape[1:-1] != [64, 64]:
            dsfac = im_shape[1] / 64.
            x_reshaped = tf.image.resize_images(self.x, size=[64, 64])
            y_reshaped = tf.image.resize_images(self.y, size=[64, 64])

            self.xy = tf.concat([x_reshaped, y_reshaped], axis=3)

        self.VectorCNN = VectorCNN('VectorCNN', is_train=False)
        self.v = self.VectorCNN(self.xy)

        if list(im_shape[1:-1]) != [64, 64]:
            self.v = tf.image.resize_images(self.v, size=im_shape[1:-1])
            self.v = self.v * dsfac

        self.z = batch_displacement_warp2d(
            self.x, self.v, vector_fields_in_pixel_space=True)

        self.sess.run(tf.global_variables_initializer())

    def deploy(self, dir_path, x, y, save_def_field_info=False):
        z, df = self.sess.run(
            [self.z, self.v],
            {
                self.x: x,
                self.y: y})

        save_image(os.path.join(dir_path, 'result_image.png'), z[0, :, :, 0])
        if save_def_field_info:
            medpy.io.save(
                np.squeeze(np.transpose(df[0, :, :, 0])),
                os.path.join(
                    dir_path, 'df_u.nii.gz'),
                hdr=False, force=True)
            medpy.io.save(
                np.squeeze(np.transpose(df[0, :, :, 1])),
                os.path.join(
                    dir_path, 'df_v.nii.gz'),
                hdr=False, force=True)
            grad_magn = np.sqrt(np.multiply(df[0, :, :, 0],
                                            df[0, :, :, 0]) +
                                np.multiply(df[0, :, :, 1],
                                            df[0, :, :, 1]))
            grad_magn = np.transpose(grad_magn)
            medpy.io.save(
                np.squeeze(grad_magn),
                os.path.join(
                    dir_path, 'df_grad.nii.gz'),
                hdr=False, force=True)

        return z, df

    def save(self, ckpt_path, step=None):
        self.VectorCNN.save(
            self.sess, os.path.join(ckpt_path, 'model.ckpt'), step)

    def restore(self, ckpt_path, step=None):
        ckpt_file = 'model.ckpt'
        if step is not None:
            ckpt_file = ckpt_file + '-' + str(step)
        self.VectorCNN.restore(
            self.sess, os.path.join(ckpt_path, ckpt_file))


import os
import glob
import json
import datetime
from collections import defaultdict
import cv2
import numpy as np
import pandas as pd
from sklearn.neighbors import KDTree

from PySide6.QtCore import QObject, Signal

from ..utils.common import get_immediate_subdirectories, to_celsius
from ..utils.geojson import load_geojson, save_geojson, coords_wgs84_to_ltp

def load_modules(file):
    df = load_geojson(open(file, "r"))
    df = df.set_index("track_id")
    df_ltp = coords_wgs84_to_ltp(df)
    df_corners = df_ltp.loc[df_ltp["geometry_type"] == "Polygon"]
    df_centers = df_ltp.loc[df_ltp["geometry_type"] == "Point"]
    return df, df_corners, df_centers

def truncate_patch(patch, margin=0.05):
    """Truncates module edges by margin (percent of width) to remove module frame."""
    width = patch.shape[1]
    margin_px = int(margin*width)
    patch = patch[margin_px:-margin_px, margin_px:-margin_px]
    return patch

def remove_patches_with_sun_reflection(patch_files, sun_reflections):
    """Returns a copy of the patch_files list withput patches that contain a sun reflections as per 'sun_reflections'."""
    patch_files_filtered = []
    for patch_file in patch_files:
        patch_name = os.path.splitext(os.path.basename(patch_file))[0]
        if patch_name not in sun_reflections:
            patch_files_filtered.append(patch_file)
    return patch_files_filtered

def get_patch_temps(patch_files, margin, to_celsius_gain, to_celsius_offset):
    """Returns min, max, mean and median temperatures for each patch of a module."""
    temps = defaultdict(list)
    for patch_file in patch_files:
        patch = cv2.imread(patch_file, cv2.IMREAD_ANYDEPTH)
        if patch is not None:
            patch = truncate_patch(patch, margin)
            temps["min"].append(to_celsius(np.min(patch), to_celsius_gain, to_celsius_offset))
            temps["max"].append(to_celsius(np.max(patch), to_celsius_gain, to_celsius_offset))
            temps["mean"].append(to_celsius(np.mean(patch), to_celsius_gain, to_celsius_offset))
            temps["median"].append(to_celsius(np.median(patch), to_celsius_gain, to_celsius_offset))
    return temps

def mean_over_patches(dataframe, temps):
    """Compute the mean of the module temperatures over all patches of a module."""
    for patch_area_agg in ["min", "max", "mean", "median"]:
        dataframe["{}_temp".format(patch_area_agg)] = pd.Series({track_id: np.mean(t[patch_area_agg]) for track_id, t in temps.items()})

class AnalysisModuleTemperaturesWorker(QObject):
    finished = Signal()
    progress = Signal(float, bool, str)

    def __init__(self, dataset_dir, dataset_version, name, to_celsius_gain, to_celsius_offset,
            border_margin, neighbour_radius, ignore_sun_reflections, sun_reflections):
        super().__init__()
        self.is_cancelled = False
        self.timestamp = datetime.datetime.utcnow().isoformat()
        self.dataset_dir = dataset_dir
        self.dataset_version = dataset_version
        self.name = name
        self.to_celsius_gain = to_celsius_gain
        self.to_celsius_offset = to_celsius_offset
        self.border_margin = 0.01 * border_margin
        self.neighbour_radius = neighbour_radius
        self.ignore_sun_reflections = ignore_sun_reflections
        self.sun_reflections = sun_reflections
        self.progress_last_step = 0.0

    def get_neighbours_median_temp(self, df_centers, neighbour_radius=7, column="mean_of_max_temps"):
        """Returns a list of mean temperatures of the neighbours of each module in `df_centers`.
        The `neighbour_radius` defines the circle radius in which to look for neighbouring modules.
        The `column` specifies which temperature column to use."""
        centers = np.array([[d["coordinates"][0], d["coordinates"][1]] for d in df_centers["geometry"]])
        tree = KDTree(centers)
        neighbor_idxs = tree.query_radius(centers, r=neighbour_radius)

        # get mean temperature of neighbors
        neighbour_mean_temps = []
        for row_idx, neighbor_idx in enumerate(neighbor_idxs):
            progress = self.progress_last_step + (row_idx / len(neighbor_idxs)) / 5
            if self.is_cancelled:
                self.progress.emit(progress, True, "Cancelled")
                self.finished.emit()
                return

            neighbor_idx = np.delete(neighbor_idx, np.nonzero(neighbor_idx == row_idx))  # remove the current module from list of neighbors
            mean_temp = df_centers.iloc[neighbor_idx][column].median()
            neighbour_mean_temps.append(mean_temp)

            self.progress.emit(progress, False, "Computing corrected {}...".format(" ".join(column.split("_"))))
        self.progress_last_step = progress
        return neighbour_mean_temps

    def run(self):
        if self.dataset_version == "v1":
            patches_dir = os.path.join(self.dataset_dir, "patches_final", "radiometric")
        elif self.dataset_version == "v2":
            patches_dir = os.path.join(self.dataset_dir, "patches", "radiometric")

        file = os.path.join(self.dataset_dir, "mapping", "module_geolocations_refined.geojson")
        df, df_corners, df_centers = load_modules(file)

        temps = {}
        track_ids = sorted(get_immediate_subdirectories(patches_dir))
        for i, track_id in enumerate(track_ids):
            progress = (i / len(track_ids)) / 5
            if self.is_cancelled:
                self.progress.emit(progress, True, "Cancelled")
                self.finished.emit()
                return

            patch_files = sorted(glob.glob(os.path.join(patches_dir, track_id, "*")))
            if self.ignore_sun_reflections and self.sun_reflections is not None:
                patch_files = remove_patches_with_sun_reflection(patch_files, self.sun_reflections[track_id])
            temps[track_id] = get_patch_temps(patch_files, self.border_margin, self.to_celsius_gain, self.to_celsius_offset)

            self.progress.emit(progress, False, "Computing temperature distribution...")
        self.progress_last_step = progress

        mean_over_patches(df_corners, temps)
        mean_over_patches(df_centers, temps)

        for patch_area_agg in ["min", "max", "mean", "median"]:
            column = "{}_temp".format(patch_area_agg)
            neighbour_mean_temps = self.get_neighbours_median_temp(df_centers, neighbour_radius=self.neighbour_radius, column=column)
            if neighbour_mean_temps is None: # cancelled
                return

            df_corners["{}_corrected".format(column)] = df_corners.loc[:, column] - neighbour_mean_temps
            df_centers["{}_corrected".format(column)] = df_centers.loc[:, column] - neighbour_mean_temps

        # merge back into single geodataframe
        df_merged = df_corners.append(df_centers)

        # reuse WGS84 coordinates of original dataframe
        df_merged = df_merged.reset_index()
        df_merged = df_merged.set_index(["track_id", "geometry_type"])
        df = df.reset_index()
        df = df.set_index(["track_id", "geometry_type"])
        df_merged.update(df.loc[:, "geometry"])

        # write results to disk
        self.progress.emit(1, False, "Saving analysis results...")
        save_path = os.path.join(self.dataset_dir, "analyses", self.name)
        save_file = os.path.join(save_path, "results.geojson")
        print("Saving module temperature results in {}".format(save_file))
        os.makedirs(save_path, exist_ok=True)
        save_geojson(df_merged, open(save_file, "w"))

        print("Saving meta json in {}".format(os.path.join(save_path, "meta.json")))
        meta = {
            "type": "module_temperatures",
            "timestamp": self.timestamp,
            "dataset_dir": self.dataset_dir,
            "hyperparameters": {
                "border_margin": self.border_margin,
                "neighbour_radius": self.neighbour_radius,
                "ignore_sun_reflections": self.ignore_sun_reflections
            }
        }
        if self.ignore_sun_reflections and self.sun_reflections is not None:
            meta["hyperparameters"]["sun_reflections"] = self.sun_reflections
        json.dump(meta, open(os.path.join(save_path, "meta.json"), "w"))

        self.progress.emit(1, False, "Done")
        self.finished.emit()


# Written by Dr Daniel Buscombe, Marda Science LLC
# for the USGS Coastal Change Hazards Program
#
# MIT License
#
# Copyright (c) 2020, Marda Science LLC
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# =========================================================
import rasterio
from glob import glob
import numpy as np
import os, time, sys, getopt
from scipy.signal import convolve2d
from tkinter import filedialog
from tkinter import *
# from skimage.io import imsave
from tqdm import tqdm
from tempfile import TemporaryFile
from datetime import datetime
from skimage.filters import threshold_otsu
from skimage.filters.rank import median
from skimage.morphology import disk

# =========================================================
USE_GPU = True
CALC_CONF = True
# =======================================================

if USE_GPU == True:
   ##use the first available GPU
   os.environ['CUDA_VISIBLE_DEVICES'] = '0' #'1'
else:
   ## to use the CPU (not recommended):
   os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

import tensorflow as tf #numerical operations on gpu
import tensorflow.keras.backend as K
from tensorflow.keras.models import model_from_json

#suppress tensorflow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

SEED=42
np.random.seed(SEED)
AUTO = tf.data.experimental.AUTOTUNE # used in tf.data.Dataset API

tf.random.set_seed(SEED)

print("Version: ", tf.__version__)
print("Eager mode: ", tf.executing_eagerly())
print('GPU name: ', tf.config.experimental.list_physical_devices('GPU'))
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

# =========================================================
# =========================================================
# def rescale(dat,mn,mx):
#    """
#    rescales an input dat between mn and mx
#    """
#    m = np.min(dat.flatten())
#    M = np.max(dat.flatten())
#    return (mx-mn)*(dat-m)/(M-m)+mn

#-----------------------------------
def std_convoluted(image, N):
    """
    fast windowed stdev based on kernel convolution
    """
    im = np.array(image, dtype=float)
    im2 = im**2
    ones = np.ones(im.shape)

    kernel = np.ones((2*N+1, 2*N+1))
    s = convolve2d(im, kernel, mode="same")
    s2 = convolve2d(im2, kernel, mode="same")
    ns = convolve2d(ones, kernel, mode="same")

    return np.sqrt((s2 - s**2 / ns) / ns)

#-----------------------------------
def z_exp(z,n=0.5):
    """
    empirical exponential scaler
    """
    zmin = np.min(z)
    zmax = np.max(z)
    zrange = zmax-zmin
    ivals = 255*((z-zmin)/(zrange))**n
    return ivals

##==========================================================
def load_ODM_json(weights_file):
    """
    load keras Res-UNet model implementation from json file
    mode=1, dem only
    """
    weights_path = weights_file
    json_file = open(weights_path.replace('.h5','.json'), 'r')
    loaded_model_json = json_file.read()
    json_file.close()
    model = model_from_json(loaded_model_json)

    model.load_weights(weights_path)
    model.compile(optimizer = 'adam', loss = 'categorical_crossentropy')#, metrics = [mean_iou, dice_coef])
    return model

#-----------------------------------
def medianfiltmask(est_label, radius=7):
    """
    mask a 1 or 3 band image, band by band for memory saving
    """
    out_label = np.ones_like(est_label)
    for k in range(est_label.shape[-1]):
        out_label[:,:,k] = median(est_label[:,:,k],disk(radius))
    return out_label

#-----------------------------------
def mask(bigtiff,profile,out_mask):
    """
    mask a 1 or 3 band image, band by band for memory saving
    """
    if profile['count']==4:
        bigtiff1 = bigtiff[0,:,:]
        bigtiff1[out_mask==1] = profile['nodata']
        bigtiff[0,:,:] = bigtiff1
        del bigtiff1

        bigtiff2 = bigtiff[1,:,:]
        bigtiff2[out_mask==1] = profile['nodata']
        bigtiff[1,:,:] = bigtiff2
        del bigtiff2

        bigtiff3 = bigtiff[2,:,:]
        bigtiff3[out_mask==1] = profile['nodata']
        bigtiff[2,:,:] = bigtiff3
        del bigtiff3

        bigtiff4 = bigtiff[3,:,:]
        bigtiff4[out_mask==1] = profile['nodata']
        bigtiff[3,:,:] = bigtiff4
        del bigtiff4
    else:
        bigtiff1 = bigtiff[0,:,:]
        bigtiff1[out_mask==1] = profile['nodata']
        bigtiff[0,:,:] = bigtiff1
        del bigtiff1

    return bigtiff

#-----------------------------------
def main(dem, rgb, model_mode, overlap, medfilt_radius, use_otsu):
    """
    main program
	-o is percentage overlap between tiles (0 for no overlap, 25 for 25% overlap, etc)
	-f is an integer >0 for median filtering. Refers to radius of smoothing kernel. Default = no median filter (0)
	-t is 0 for "no, dont use Otsu - use a threshold of 0.5' and 1 for "yes please, use Otsu"

	Defaults:
	-m  = 1 (dem only - the other model modes are not implemented)
	-o = 25 (percentage overlap)
	-f = 0 (median filter radius)
	-t = 0 (dont use otsu, use threshold=0.5)
    """
    print('.....................................')
    print('.....................................')
    print(datetime.now())
    print("Working on %s" % (dem))
    tilesize = 1024 #
    verbose = True
    # thres = .25  #channel 2 (bad) threshold - deliberately <=.5

    if model_mode==1: #dem only
        print('DEM-only mode')
        weights_file = 'model'+os.sep+'orthoclip_demonly_3class_batch_6.h5'
    elif model_mode==2: #stdev only
        print('STDEV-only mode')
        weights_file = 'model'+os.sep+'orthoclip_stdevonly_3class_batch_6.h5'
    elif model_mode==3: #rgb only
        print('Ortho-only mode')
        weights_file = 'model'+os.sep+'orthoclip_orthoonly_3class_batch_6.h5'
    elif model_mode==4: #rgb+dem
        print('Ortho+DEM mode')
        weights_file = 'model'+os.sep+'orthoclip_RGBdem_3class_batch_6.h5'
    else: #rgb+stdev
        print('Ortho+STDEV mode')
        weights_file = 'model'+os.sep+'orthoclip_RGBstdev_3class_batch_6.h5'

    try:
        os.remove(dem.replace('.tif','_automask.tif'))
    except:
        pass

    try:
        os.remove(dem.replace('.tif','_automasked_dem.tif'))
    except:
        pass

    do_plot=False #True
    if do_plot:
        import matplotlib.pyplot as plt

    if verbose:
        # # start timer
        start = time.time()

    #====================================================
    with rasterio.open(dem) as src:
        profile = src.profile

    width = profile['width']
    height = profile['height']

    padwidth = width + (tilesize - width % tilesize)
    padheight = height + (tilesize - height % tilesize)

    try:
        #=======================================================
        # step 2: make an instance of the model and load the weights
        model = load_ODM_json(weights_file)
    except:
        print("Model loading failed - check inputs. Exiting ...")
        sys.exit(2)

    #=======================================================
    ##step 3: pre-allocate arrays to be written as big geotiff later, and start filling it
    # by applying the model to windows of data
    print('.....................................')
    print('Using model for prediction on images ...')
    print("Started at %s" % datetime.now())

    try:
        #pre-allocate large arrays to be eventually written to geotiff format files,
        out_mask = np.zeros((padwidth, padheight), dtype=np.uint8)
        accumulator = np.zeros((padwidth, padheight), dtype=np.uint8)

        # if CALC_CONF:
            # out_conf = np.ones((padwidth, padheight), dtype=np.float32)

        # counter=0
        # prc_overlap = 25 #100
        prc_overlap = overlap/100
        increment = int((1-prc_overlap)*tilesize)
        # if prc_overlap<.25:
            # increment = tilesize
        # print("Using overlap of %i pixels" % (increment))

        tilesize=1024
        # window1d = np.abs(np.hanning(tilesize/2))
        # window2d = np.sqrt(np.outer(window1d,window1d))

        # window2d_up = np.roll(window2d,int(tilesize/4),axis=1)
        # window2d_up[:,:int(tilesize/4)] = 0

        for i in tqdm(range(0, padwidth, increment)):
            for j in range(0, padheight, increment):
                window = rasterio.windows.Window(i,j,tilesize, tilesize)

                with rasterio.open(dem) as src:
                    profile = src.profile
                    orig = src.read(window=window).squeeze() #/255.

                if len(np.unique(orig))>1:
                    #print(i,j)

                    orig[orig==profile['nodata']] = 0

                    # if do_plot:
                    #     plt.subplot(233); plt.imshow(image, cmap='gray'); plt.axis('off'); plt.colorbar(shrink=.5)
                    if model_mode<3:
                        image = np.squeeze(orig).T
                    elif model_mode==3: #rgb only
                        image = np.squeeze(orig[:3,:,:]).T

                    if model_mode==2: #stdev only
                        ##compute stdev raster using convoluttion with a 3x3 kernel
                        image = std_convoluted(orig, 3)
                        del orig
                        ##scale image exponentially using n=.5
                        image = z_exp(image,n=0.5)

                    if model_mode==4: #rgb + dem
                        image = np.squeeze(orig).T

                        if rgb is not None: #when adding rgb to dem or stdev
                            with rasterio.open(rgb) as src:
                                profile = src.profile
                                orig = src.read(window=window).squeeze() #/255.
                            orig = np.squeeze(orig).T

                    if do_plot:
                        plt.subplot(221); plt.imshow(image, cmap='gray'); plt.axis('off'); plt.colorbar(shrink=.5)

                    do_pad = 0
                    if image.shape[0] < tilesize:

                        if model_mode<3:
                            subset_pad = np.zeros((tilesize, tilesize), dtype=profile['dtype'])
                            x,y=image.shape
                            subset_pad[:x,:y] = image
                        elif model_mode==3:
                            x,y,z=image.shape
                            subset_pad = np.zeros((tilesize, tilesize,z), dtype=profile['dtype'])
                            subset_pad[:x,:y,:z] = image
                        elif model_mode>3:
                            x,y=image.shape
                            subset_pad = np.zeros((tilesize, tilesize), dtype=profile['dtype'])
                            subset_pad[:x,:y] = image

                            x,y,z=orig.shape
                            subset_pad2 = np.zeros((tilesize, tilesize,z), dtype=profile['dtype'])
                            subset_pad2[:x,:y,:z] = orig

                        del image
                        image = subset_pad.copy()
                        del subset_pad
                        do_pad = 1

                        if model_mode>3:
                            if 'subset_pad2' in locals():
                                #del orig
                                orig = subset_pad2.copy()#/255
                                del subset_pad2

                    elif image.shape[1] < tilesize:
                        if model_mode<3:
                            subset_pad = np.zeros((tilesize, tilesize), dtype=profile['dtype'])
                            x,y=image.shape
                            subset_pad[:x,:y] = image
                        elif model_mode==3:
                            x,y,z=image.shape
                            subset_pad = np.zeros((tilesize, tilesize,z), dtype=profile['dtype'])
                            subset_pad[:x,:y,:z] = image
                        elif model_mode>3:
                            x,y=image.shape
                            subset_pad = np.zeros((tilesize, tilesize), dtype=profile['dtype'])
                            subset_pad[:x,:y] = image

                            x,y,z=orig.shape
                            subset_pad2 = np.zeros((tilesize, tilesize,z), dtype=profile['dtype'])
                            subset_pad2[:x,:y,:z] = orig
                        do_pad = 1

                        del image
                        image = subset_pad.copy()#/255
                        del subset_pad

                        if model_mode>3:
                            if 'subset_pad2' in locals():
                                #del orig
                                orig = subset_pad2.copy()#/255
                                del subset_pad2

                    if model_mode>3:
                        image = np.dstack((orig[:,:,0], orig[:,:,1], orig[:,:,2], image))
                        #del orig

                    image = tf.expand_dims(image, 0)
                    #print(image.shape)
                    # use model in prediction mode
                    est_label = model.predict(image , batch_size=1).squeeze()
                    image = image.numpy().squeeze()

                    conf = np.std(est_label, -1)

                    # est_label = np.argmax(est_label, -1) #2=bad, 1=good, 0=bad
                    # est_label = (est_label==1).astype('uint8') #only good

                    est_label_orig = est_label.copy()
                    est_label = np.zeros((est_label.shape[0], est_label.shape[1])) #np.argmax(est_label, -1).astype(np.uint8).T

                    if use_otsu:
                       otsu_nodata = threshold_otsu(est_label_orig[:,:,0])
                       otsu_baddata = threshold_otsu(est_label_orig[:,:,1])
                       otsu_gooddata = threshold_otsu(est_label_orig[:,:,2])

                       est_label[est_label_orig[:,:,0]>otsu_nodata] = 0  #nodata
                       est_label[est_label_orig[:,:,1]>otsu_baddata] = 0  #bad

                       est_label[est_label_orig[:,:,2]>otsu_gooddata] = 1  #good
                       est_label[est_label>1] = 1

                    else:
                       # if np.any(est_label_orig[:,:,2]>.05):
                       est_label[est_label_orig[:,:,0]>0.5] = 0  #nodata
                       est_label[est_label_orig[:,:,1]>0.5] = 0  #bad

                       est_label[est_label_orig[:,:,2]>0.5] = 1  #good
                       est_label[est_label>1] = 1

                    if do_plot:
                        plt.subplot(222); plt.imshow(est_label, cmap='bwr'); plt.axis('off'); plt.colorbar(shrink=.5)
                        plt.subplot(223); plt.imshow(image, cmap='gray'); plt.imshow(est_label, cmap='bwr', alpha=0.3); plt.axis('off'); plt.colorbar(shrink=.5)

                    if do_plot:
                        plt.subplot(224); plt.imshow(conf, cmap='Blues'); plt.axis('off'); plt.colorbar(shrink=.5)
                        plt.savefig('aom_ex'+str(i)+'_'+str(j)+'.png', dpi=300); plt.close()

                    #fill out big rasters
                    if do_pad==1:
                        tilesize_j, tilesize_i = orig.shape
                        out_mask[i:i+tilesize_i,j:j+tilesize_j] += est_label[:tilesize_i,:tilesize_j].astype('uint8') #fill in that portion of big mask
                    else:
                        out_mask[i:i+tilesize,j:j+tilesize] += est_label.astype('uint8') #fill in that portion of big mask

                    #fill out big rasters
                    if do_pad==1:
                        tilesize_j, tilesize_i = orig.shape
                        accumulator[i:i+tilesize_i,j:j+tilesize_j] += np.ones((tilesize_i,tilesize_j)).astype('uint8') #fill in that portion of accumulator
                    else:
                        accumulator[i:i+tilesize,j:j+tilesize] += np.ones((tilesize,tilesize)).astype('uint8') #fill in that portion of accumulator

                    del est_label, orig

        #=======================================================
        ##step 4: crop, rotate arrays, divide by N, and write to memory mapped temporary file
        #(i know, classy)

        out_mask = out_mask[:width,:height]
        accumulator = accumulator[:width,:height]

        accumulator[accumulator<1] = 1
        out_mask2 = out_mask.copy()
        out_mask2 = out_mask2/accumulator
        out_mask2[np.isnan(out_mask2)] = 0
        out_mask2[np.isinf(out_mask2)] = 0
        out_mask2[out_mask2>0] = 1
        out_mask2 = out_mask2.astype('uint8')
        out_mask = out_mask2.copy()

        out_mask = np.rot90(np.fliplr(out_mask))
        out_shape = out_mask.shape

        # out_conf = out_conf[:width,:height]
        # out_conf = np.rot90(np.fliplr(out_conf))

        out_mask = out_mask.astype('uint8')
        out_mask = (out_mask!=1).astype('uint8')

        if medfilt_radius>0:
           out_mask = median(out_mask,disk(medfilt_radius))

        # if CALC_CONF:
            # # out_conf = np.divide(out_conf, np.maximum(0,n.astype('float')), out=out_conf, casting='unsafe')
            # out_conf[out_conf==np.nan]=0
            # out_conf[out_conf==np.inf]=0
            # out_conf[out_conf<=0] = 0
            # out_conf = out_conf-1
            # #out_conf[out_conf>.5] = 0
            # # out_conf[out_conf==-1] = 0
            # #out_mask = np.memmap(outfile, dtype='uint8', mode='r', shape=out_shape)
            # out_conf[out_mask==0] = 0

        # # out_mask[out_conf<.25] = 0

        outfile = TemporaryFile()
        fp = np.memmap(outfile, dtype='uint8', mode='w+', shape=out_mask.shape)
        fp[:] = out_mask[:]
        fp.flush()
        del out_mask
        del fp

        # #write out to temporary memory mapped file
        # outfile2 = TemporaryFile()
        # fp = np.memmap(outfile2, dtype='float32', mode='w+', shape=out_conf.shape)
        # fp[:] = out_conf[:]
        # fp.flush()
        # del out_conf
        # del fp

        # #read back in again without using any memory
        # out_conf = np.memmap(outfile2, dtype='float32', mode='r', shape=out_shape)

        if 'out_mask' not in locals():
            out_mask = np.memmap(outfile, dtype='uint8', mode='r', shape=out_shape)

        #=======================================================
        ##step 5: write out masked dem and mask to file
    except:
        print("Error occurred with creating mask for file %s" % (dem))

    # try:
        # # get the profile of the geotif which contains info we need to write out rasters with the same CRS, dtype, etc
        # with rasterio.open(dem) as src: #rgb
            # profile = src.profile

        # #=============================
        # ##conf
        # if CALC_CONF:
            # if verbose:
                # print('.....................................')
                # print('Writing out mask confidence raster ...')
                # print(datetime.now())

            # profile['dtype'] = 'float32'
            # profile['nodata'] = 0.0
            # profile['count'] = 1

            # with rasterio.Env():
                # with rasterio.open(dem.replace('.tif','_automask_conf.tif'), 'w', **profile) as dst:
                    # dst.write(np.expand_dims(out_conf.astype('float32'),0))
            # del out_conf

    # except:
        # print("Error occurred with creating masked confidence file for %s" % (dem))

    try:
        #=============================
        ## dem
        if verbose:
            print('.....................................')
            print('Writing out masked DEM ...')
            print(datetime.now())

        with rasterio.open(dem) as src:
            profile = src.profile
            bigtiff = src.read()

        bigtiff = mask(bigtiff,profile,out_mask)

        with rasterio.Env():
            with rasterio.open(dem.replace('.tif','_automasked_dem.tif'), 'w', **profile) as dst:
                dst.write(bigtiff)
        del bigtiff

    except:
        print("Error occurred with creating masked dem file for %s" % (dem))

    try:
        #=============================
        ##mask
        if verbose:
            print('.....................................')
            print('Writing out mask raster ...')
            print(datetime.now())

        profile['nodata'] = 0.0
        profile['dtype'] = 'uint8'
        profile['count'] = 1

        with rasterio.Env():
            with rasterio.open(dem.replace('.tif','_automask.tif'), 'w', **profile) as dst:
                dst.write(np.expand_dims(out_mask,0))

    except:
        print("Error occurred with creating mask file for %s" % (dem))

    try:
        #=======================================================
        ##step 6: clean up temp files
        if 'outfile' in locals():
            outfile.close()
        if 'outfile2' in locals():
            outfile2.close()
        del out_mask

        if verbose:
            print("Program finished at %s" % datetime.now())
            # if os.name=='posix': # true if linux/mac
            elapsed = (time.time() - start)/60
            # else: # windows
            #    elapsed = (time.clock() - start)/60
            print("Processing took "+ str(elapsed) + " minutes")

    except:
        print("Error")

###==================================================================
#===============================================================
if __name__ == '__main__':

    argv = sys.argv[1:]
    try:
        opts, args = getopt.getopt(argv,"h:m:o:f:t:")
    except getopt.GetoptError:
        print('python mask_dem.py -m model mode -o overlap percentage -f median filter radius -t otsu threshold \n')
        print('Defaults:\n')
        print('-m  = 1 (dem only - the other model modes are not implemented)\n')
        print('-o = 25 (percentage overlap between 5 and 95)\n')
        print('-f = 0 (median filter radius)\n')
        print('-t = 0 (Not otsu)\n')
        sys.exit(2)
    for opt, arg in opts:
        if opt == '-h':
            print('Example usage python mask_dem.py')
            print('Example usage python mask_dem.py -m 1 -o 25 -f 5 -t 1')
            print('Example usage: python mask_dem.py -m 2 -o 10 -f 10 -t 0')
            sys.exit()

        elif opt in ("-m"):
            model_mode = arg
            model_mode = int(model_mode)
        elif opt in ("-o"):
            overlap = arg
            overlap = int(overlap)
        elif opt in ("-f"):
            medfilt_radius = arg
            medfilt_radius = int(medfilt_radius)
        elif opt in ("-t"):
            use_otsu = arg
            use_otsu = int(use_otsu)

    if 'model_mode' not in locals():
        model_mode = 1
    print(model_mode)

    if 'overlap' not in locals():
        overlap = 25 #25% overlap
    elif overlap<5:
        overlap = 5
        print("Overlap minimum of 5% imposed")
    elif overlap>95:
        overlap = 95
        print("Overlap maximum of 95% imposed")

    print("Overlap: %i"% (overlap)	)

    if 'medfilt_radius' not in locals():
        medfilt_radius = 0

    print("Median filter radius: %i" % (medfilt_radius))

    if 'use_otsu' not in locals():
        use_otsu = False
    elif type(use_otsu) is int:
        use_otsu = bool(use_otsu)
    print(use_otsu)

    #====================================================
    print('.....................................')
    print('.....................................')
    print('.....................................')
    print('.....................................')
    print('.....................................')
    print('................A.D.M................')
    print('.....................................')
    print('.....................................')
    print('.....................................')
    print('.....................................')
    print('.....................................')

    # model_mode = 1
    #model_mode = 2
    #model_mode = 3
    #1=dem, 2=stdev, 3=rgb, 4=rgb+dem, 5=rgb+stdev

    if model_mode<3:
        root = Tk()
        dems = filedialog.askopenfilenames(initialdir = "./",title = "Select DEM file",filetypes = (("DSM/DEM geotiff file","*.tif"),("all files","*.*")))
        root.withdraw()
        print("%s files selected" % (len(dems)))
        rgbs = [None for k in dems]
    elif model_mode==3:
        root = Tk()
        dems = filedialog.askopenfilenames(initialdir = "./",title = "Select Orthomosaic file",filetypes = (("RGB Ortho geotiff file","*.tif"),("all files","*.*")))
        root.withdraw()
        print("%s files selected" % (len(dems)))
        rgbs = [None for k in dems]
    elif model_mode>3:
        root = Tk()
        rgbs = filedialog.askopenfilenames(initialdir = "./",title = "Select Orthomosaic file",filetypes = (("RGB Ortho geotiff file","*.tif"),("all files","*.*")))
        root.withdraw()
        print("%s files selected" % (len(rgbs)))
        rgbs = sorted(rgbs)

        root = Tk()
        dems = filedialog.askopenfilenames(initialdir = "./",title = "Select corresponding DEM file",filetypes = (("DSM/DEM geotiff file","*.tif"),("all files","*.*")))
        root.withdraw()
        print("%s files selected" % (len(dems)))
        dems = sorted(dems)

    try:
        for dem,rgb in zip(dems,rgbs):
            main(dem, rgb, model_mode, overlap, medfilt_radius, use_otsu)
    except:
        print("Unspecified error. Check inputs. Exiting ...")
        sys.exit(2)

	# dem = r'F:/dbuscombe_github/dems/noisy/clip_20181006_Ophelia_Inlet_to_Beaufort_Inlet_1m_UTM18N_NAVD88_cog.tif'
	# dem = os.path.normpath(dem)