from itertools import count

from .utils.decorators import gitlog
from scipy.misc import imresize
import datajoint as dj
from datajoint.jobs import key_hash
from tqdm import tqdm
import cv2
import numpy as np
import json
import os
from commons import lab
from datajoint.autopopulate import AutoPopulate

from .utils.eye_tracking import ROIGrabber, PupilTracker, CVROIGrabber, ManualTracker
from . import config
from .utils import h5
from . import experiment, notify
from .exceptions import PipelineException


schema = dj.schema('pipeline_eye', locals())


DEFAULT_PARAMETERS = {'relative_area_threshold': 0.002,
                      'ratio_threshold': 1.5,
                      'error_threshold': 0.1,
                      'min_contour_len': 5,
                      'margin': 0.02,
                      'contrast_threshold': 5,
                      'speed_threshold': 0.1,
                      'dr_threshold': 0.1,
                      'gaussian_blur': 5,
                      'extreme_meso': 0,
                      'running_avg': 0.4,
                      'exponent': 9
                      }


@schema
class Eye(dj.Imported):
    definition = """  # eye movie timestamps synchronized to behavior clock

    -> experiment.Scan
    ---
    eye_time                    : longblob  # times of each movie frame in behavior clock
    total_frames                : int       # number of frames in movie.
    preview_frames              : longblob  # 16 preview frames
    eye_ts=CURRENT_TIMESTAMP    : timestamp
    """

    @property
    def key_source(self):
        return experiment.Scan()
 & experiment.Scan.EyeVideo().proj()

    def _make_tuples(self, key):
        # Get behavior filename
        behavior_path = (experiment.Session() & key).fetch1('behavior_path')
        local_path = lab.Paths().get_local_path(behavior_path)
        filename = (experiment.Scan.BehaviorFile() & key).fetch1('filename')
        full_filename = os.path.join(local_path, filename)

        # Read file
        data = h5.read_behavior_file(full_filename)

        # Read counter 
timestamps and convert to seconds
        if data['version'] == '1.0': # older h5 format
            rig = (experiment.Session() & key).fetch('rig')
            timestamps_in_secs = h5.ts2sec(data['cam1_ts' if rig == '2P3' else 'cam2_ts'])
        else:
            timestamps_in_secs = h5.ts2sec(data['eyecam_ts'][0])
        ts = h5.ts2sec(data['ts'], is_packeted=True)
        # edge case when ts and eye ts start in different sides of the master clock max value 2 **32
        if abs(ts[0] - timestamps_in_secs[0]) > 2 ** 31:
            timestamps_in_secs += (2 ** 32 if ts[0] > timestamps_in_secs[0] else -2 ** 32)

        # Fill with NaNs for out-of-range data or mistimed packets (NaNs in ts)
        timestamps_in_secs[timestamps_in_secs < ts[0]] = float('nan')
        timestamps_in_secs[timestamps_in_secs > ts[-1]] = float('nan')
        nan_limits = np.where(np.diff([0, *np.isnan(ts), 0]))[0]
        for start, stop in zip(nan_limits[::2], nan_limits[1::2]):
            lower_ts = float('-inf') if start == 0 else ts[start - 1]
            upper_ts = float('inf') if stop == len(ts) else ts[stop]
            timestamps_in_secs[np.logical_and(timestamps_in_secs > lower_ts,
                                              timestamps_in_secs < upper_ts)] = float('nan')

        # Read video
        filename = (experiment.Scan.EyeVideo() & key).fetch1('filename')
        full_filename = os.path.join(local_path, filename)
        video = cv2.VideoCapture(full_filename) # note: prints many 'Unexpected list ...'

        # Fix inconsistent num_video_frames vs num_timestamps
        num_video_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
        num_timestamps = len(timestamps_in_secs)
        if num_timestamps != num_video_frames:
            if abs(num_timestamps - num_video_frames) > 1:
                msg = ('Number of movie frames and timestamps differ: {} frames vs {} '
                       'timestamps'). format(num_video_frames, num_timestamps)
                raise PipelineException(msg)
            elif num_timestamps > num_video_frames: # cut timestamps to match video frames
                timestamps_in_secs = timestamps_in_secs[:-1]
            else: # fill with NaNs
                timestamps_in_secs = np.array([*timestamps_in_secs, float('nan')])

        # Get 16 sample frames
        frames = []
        for frame_idx in np.round(np.linspace(0, num_video_frames - 1, 16)).astype(int):
            video.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
            _, frame = video.read()
            frames.append(np.asarray(frame, dtype=float)[..., 0])
        frames = np.stack(frames, axis=-1)

        # Insert
        self.insert1({**key, 'eye_time': timestamps_in_secs,
                      'total_frames': len(timestamps_in_secs), 'preview_frames': frames})
        self.notify(key, frames)

    @notify.ignore_exceptions
    def notify(self, key, frames):
        import imageio

        video_filename = '/tmp/' + key_hash(key) + '.gif'
        frames = [imresize(img, 0.25) for img in frames.transpose([2, 0, 1])]
        imageio.mimsave(video_filename, frames, duration=0.5)

        msg = 'eye frames for {animal_id}-{session}-{scan_idx}'.format(**key)
        slack_user = notify.SlackUser() & (experiment.Session() & key)
        slack_user.notify(file=video_filename, file_title=msg, channel='#pipeline_quality')

    def get_video_path(self):
        video_info = (experiment.Session() * experiment.Scan.EyeVideo() & self).fetch1()
        video_path = lab.Paths().get_local_path("{behavior_path}/{filename}".format(**video_info))
        return video_path


@schema
class TrackingTask(dj.Manual):
    definition = """
    # ROI and parameters for tracking the eye
    -> Eye
    ---
    eye_roi                     : tinyblob  # manual roi containing eye in full-size movie
    """

    class ManualParameters(dj.Part):
        definition = """
        # manual tracking parameters overwriting the default settings
        -> master
        ---
        tracking_parameters  : varchar(512)  # tracking parameters
        """

    class Ignore(dj.Part):
        definition = """
        # eyes that are too bad to be tracked
        -> master
        ---
        """

    class Mask(dj.Part):
        definition = """
        # mask for tracking
        -> master
        ---
        mask        : longblob
        """

    @staticmethod
    def _get_modified_parameters():
        new_param = dict(DEFAULT_PARAMETERS)
        for k, v in new_param.items():
            nv = input("{} (default: {}): ".format(k, v))
            new_param[k] = float(nv) if nv else v
        return json.dumps(new_param)

    def enter_roi(self, key, **kwargs):
        key = (Eye() & key).fetch1(dj.key)  # complete key
        frames = (Eye() & key).fetch1('preview_frames')
        try:
            print('Drag window and print q when done')
            rg = CVROIGrabber(frames.mean(axis=2))
            rg.grab()
        except ImportError:
            rg = ROIGrabber(frames.mean(axis=2))

        key['eye_roi'] = rg.roi
        mask = np.asarray(rg.mask, dtype=np.uint8)
        with self.connection.transaction:
            self.insert1(key)
            trackable = input('Is the quality good enough to be tracked? [Y/n]')
            if trackable.lower() == 'n':
                self.insert1(key)
                self.Ignore().insert1(key, ignore_extra_field=True)
            else:
                new_param = dict(DEFAULT_PARAMETERS, **kwargs)
                print('Those are the tracking parameters')
                print(new_param)
                new_param = json.dumps(new_param)
                extra_parameters = input('Do you want to change them? [N/y]')
                if extra_parameters.lower() == 'y':
                    new_param = self._get_modified_parameters()
                self.ManualParameters().insert1(dict(key, tracking_parameters=new_param),
                                                ignore_extra_fields=True)
            if np.any(mask == 0):
                print('Inserting mask')
                key['mask'] = m