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'] = mask
                self.Mask().insert1(key, ignore_extra_fields=True)


@schema
class TrackedVideo(dj.Computed):
    definition = """
    -> Eye
    -> TrackingTask
    ---
    tracking_parameters              : longblob  # tracking parameters
    tracking_ts=CURRENT_TIMESTAMP    : timestamp  # automatic
    """

    class Frame(dj.Part):
        definition = """
        -> TrackedVideo
        frame_id                 : int           # frame id with matlab based 1 indexing
        ---
        rotated_rect=NULL        : tinyblob      # rotated rect (center, sidelength, angle) containing the ellipse
        contour=NULL             : longblob      # eye contour relative to ROI
        center=NULL              : tinyblob      # center of the ellipse in (x, y) of image
        major_r=NULL             : float         # major radius of the ellipse
        frame_intensity=NULL     : float         # std of the frame
        """

    key_source = Eye() * TrackingTask() - TrackingTask.Ignore()

    def _make_tuples(self, key):
        print("Populating", key)
        param = DEFAULT_PARAMETERS
        if TrackingTask.ManualParameters() & key:
            param = json.loads((TrackingTask.ManualParameters() & key).fetch1('tracking_parameters'))
            print('Using manual set parameters', param, flush=True)

        roi = (TrackingTask() & key).fetch1('eye_roi')

        avi_path = (Eye() & key).get_video_path()
        print(avi_path)

        if TrackingTask.Mask() & key:
            mask = (TrackingTask.Mask() & key).fetch1('mask')
        else:
            mask = None

        tr = PupilTracker(param, mask=mask)
        traces = tr.track(avi_path, roi - 1, display=config['display.tracking'])  # -1 because of matlab indices

        key['tracking_parameters'] = json.dumps(param)
        self.insert1(key)
        fr = self.Frame()
        for trace in traces:
            trace.update(key)
            fr.insert1(trace, ignore_extra_fields=True)

        self.notify(key)

    @notify.ignore_exceptions
    def notify(self, key):
        msg = 'Pupil tracking for {} has been populated'.format(key)
        (notify.SlackUser() & (experiment.Session() & key)).notify(msg)

    def plot_traces(self, outdir='./', show=False):
        """
        Plot existing traces to output directory.

        :param outdir: destination of plots
        """
        import seaborn as sns
        import matplotlib.pyplot as plt
        plt.switch_backend('GTK3Agg')

        for key in self.fetch('KEY'):
            print('Processing', key)
            with sns.axes_style('ticks'):
                fig, ax = plt.subplots(3, 1, figsize=(10, 6), sharex=True)

            r, center, contrast = (TrackedVideo.Frame() & key).fetch('major_r', 'center',
                                                                     'frame_intensity', order_by='frame_id')
            ax[0].plot(r)
            ax[0].set_title('Major Radius')
            c = np.vstack([cc if cc is not None else np.NaN * np.ones(2) for cc in center])

            ax[1].plot(c[:, 0], label='x')
            ax[1].plot(c[:, 1], label='y')
            ax[1].set_title('Pupil Center Coordinates')
            ax[1].legend()

            ax[2].plot(contrast)
            ax[2].set_title('Contrast (frame std)')
            ax[2].set_xlabel('Frames')
            try:
                os.mkdirs(os.path.expanduser(outdir) + '/{animal_id}/'.format(**key), exist_ok=True)
            except:
                pass

            fig.suptitle(
                'animal id {animal_id} session {session} scan_idx {scan_idx} eye quality {eye_quality}'.format(**key))
            fig.tight_layout()
            sns.despine(fig)
            fig.savefig(outdir + '/{animal_id}/AI{animal_id}SE{session}SI{scan_idx}EQ{eye_quality}.png'.format(**key))
            if show:
                plt.show()
            else:
                plt.close(fig)

    def show_video(self, from_frame, to_frame, framerate=1000):
        """
        Shows the video from from_frame to to_frame (1-based) and the corrsponding tracking results.
        Needs opencv installation.

        :param from_frame: start frame (1 based)
        :param to_frame:  end frame (1 based)
        """
        if not len(self) == 1:
            raise PipelineException("Restrict EyeTracking to one video only.")
        import cv2
        video_info = (experiment.Session() * experiment.Scan.EyeVideo() & self).fetch1()
        videofile = "{path_prefix}/{behavior_path}/{filename}".format(path_prefix=config['path.mounts'], **video_info)
        eye_roi = (Eye() & self).fetch1('eye_roi') - 1

        contours, ellipses = ((TrackedVideo.Frame() & self) \
                              & 'frame_id between {0} and {1}'.format(from_frame, to_frame)
                              ).fetch('contour', 'rotated_rect', order_by='frame_id')
        cap = cv2.VideoCapture(videofile)
        no_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        font = cv2.FONT_HERSHEY_SIMPLEX
        if not from_frame < no_frames:
            raise PipelineException('Starting frame exceeds number of frames')

        cap.set(cv2.CAP_PROP_POS_FRAMES, from_frame - 1)
        fr_count = from_frame - 1

        elem_count = 0
        while cap.isOpened():
            fr_count += 1
            ret, frame = cap.read()
            if fr_count < from_frame:
                continue

            if fr_count >= to_frame or fr_count >= no_frames:
                print("Reached end of videofile ", videofile)
                break
            contour = contours[elem_count]
            ellipse = ellipses[elem_count]
            elem_count += 1

            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

            cv2.putText(gray, str(fr_count), (10, 30), font, 1, (127, 127, 127), 2)

            if contour is not None:
                ellipse = (tuple(eye_roi[::-1, 0] + ellipse[:2]), tuple(ellipse[2:4]), ellipse[4])
                cv2.drawContours(gray, [contour.astype(np.int32)], 0, (255, 0, 0), 1, offset=tuple(eye_roi[::-1, 0]))
                cv2.ellipse(gray, ellipse, (0, 0, 255), 2)
            cv2.imshow('frame', gray)

            if (cv2.waitKey(int(1000 / framerate)) & 0xFF == ord('q')):
                break

        cap.release()
        cv2.destroyAllWindows()


@schema
@gitlog
class ManuallyTrackedContours(dj.Manual, AutoPopulate):
    definition = """
    -> Eye
    ---
    tracking_ts=CURRENT_TIMESTAMP    : timestamp  # automatic
    min_lambda=null                  : float      # minimum mixing weight for current frame in running average computation (1 means no running avg was used)
    """

    class Frame(dj.Part):
        definition = """
        -> master
        frame_id                 : int           # frame id with matlab based 1 indexing
        ---
        contour=NULL             : longblob      # eye contour relative to ROI
        """

    class Parameter(dj.Part):
        definition = """
        -> master.Frame
        ---
        roi=NULL                : longblob  # roi of eye
        gauss_blur=NULL         : float     # bluring of ROI
        exponent=NULL           : tinyint   # exponent for contrast enhancement
        dilation_iter=NULL      : tinyint   # number of dilation and erosion operations
        min_contour_len=NULL    : tinyint   # minimal contour length
        running_avg_mix=NULL    : float     # weight a in a * current_frame + (1-a) * running_avg 
        """

    def make(self, key, backup_file=None):
        print("Populating", key)


        if backup_file is None:
            avi_path = (Eye() & key).get_video_path()
            tracker = ManualTracker(avi_path)
            tracker.backup_file = '/tmp/tracker_state{animal_id}-{session}-{scan_idx}.pkl'.format(**key)
        else:
            tracker = ManualTracker.from_backup(backup_file)

        try:
            tracker.run()
        except:
            tracker.backup()
            raise

        logtrace = tracker.mixing_constant.logtrace.astype(float)
        self.insert1(dict(key, min_lambda=logtrace[logtrace > 0].min()))
        self.log_git(key)
        frame = self.Frame()
        parameters = self.Parameter()
        for frame_id, ok, contour, params in tqdm(zip(count(), tracker.contours_detected, tracker.contours,
                                              tracker.parameter_iter()),
                                          total=len(tracker.contours)):
            assert frame_id == params['frame_id']
            if ok:
                frame.insert1(dict(key, frame_id=frame_id, contour=contour))
            else:
                frame.insert1(dict(key, frame_id=frame_id))
            parameters.insert1(dict(key, **params), ignore_extra_fields=True)


    def warm_start(self, key, backup_file):
        assert not key in self, '{} should not be in the table already!'
        with self.connection.transaction:
            self.make(key, backup_file)


    def update(self, key):
        print("Populating", key)

        avi_path = (Eye() & key).get_video_path()

        tracker = ManualTracker(avi_path)
        contours = (self.Frame() & key).fetch('contour', order_by='frame_id')
        tracker.contours = np.array(contours)
        tracker.contours_detected = np.array([e is not None for e in contours])
        tracker.backup_file = '/tmp/tracker_update_state{animal_id}-{session}-{scan_idx}.pkl'.format(**key)

        try:
            tracker.run()
        except Exception as e:
            print(str(e))
            answer = input('Tracker crashed. Do you want to save the content anyway [y/n]?').lower()
            while answer not in ['y', 'n']:
                answer = input('Tracker crashed. Do you want to save the content anyway [y/n]?').lower()
            if answer == 'n':
                raise
        if input('Do you want to delete and replace the existing entries? Type "YES" for acknowledgement.') == "YES":
            with dj.config(safemode=False):
                with self.connection.transaction:
                    (self & key).delete()

                    logtrace = tracker.mixing_constant.logtrace.astype(float)
                    self.insert1(dict(key, min_lambda=logtrace[logtrace > 0].min()))
                    self.log_key(key)

                    frame = self.Frame()
                    parameters = self.Parameter()
                    for frame_id, ok, contour, params in tqdm(zip(count(), tracker.contours_detected, tracker.contours,
                                                                  tracker.parameter_iter()),
                                                              total=len(tracker.contours)):
                        assert frame_id == params['frame_id']
                        if ok:
                            frame.insert1(dict(key, frame_id=frame_id, contour=contour))
                        else:
                            frame.insert1(dict(key, frame_id=frame_id))
                        parameters.insert1(dict(key, **params), ignore_extra_fields=True)
@schema
class FittedContour(dj.Computed):
    definition = """
    -> ManuallyTrackedContours
    ---
    tracking_ts=CURRENT_TIMESTAMP    : timestamp  # automatic
    """

    class Ellipse(dj.Part):
        definition = """
        -> master
        frame_id                 : int           # frame id with matlab based 1 indexing
        ---
        center=NULL              : tinyblob      # center of the ellipse in (x, y) of image
        major_r=NULL             : float         # major radius of the ellipse
        """

    def display_frame_number(self, img, frame_number, n_frames):
        font = cv2.FONT_HERSHEY_SIMPLEX
        fs = .7
        cv2.putText(img, "[{fr_count:05d}/{frames:05d}]".format(
            fr_count=frame_number, frames=n_frames),
                    (10, 30), font, fs, (255, 144, 30), 2)

    def make(self, key):
        print("Populating", key)

        avi_path = (Eye() & key).get_video_path()

        contours = (ManuallyTrackedContours.Frame() & key).fetch(order_by='frame_id ASC', as_dict=True)
        self._cap = cap = cv2.VideoCapture(avi_path)

        frame_number = 0
        n_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        while cap.isOpened():
            if frame_number >= n_frames - 1:
                print("Reached end of videofile ", avi_path)
                break

            ret, frame = self._cap.read()
            ckey = contours[frame_number]
            if ret and frame is not None and ckey['contour'] is not None:
                if ckey['contour'] is not None and len(ckey['contour']) >= 5:
                    contour = ckey['contour']
                    center = contour.mean(axis=0)
                    cv2.drawContours(frame, [contour], -1, (0, 255, 0), 1)
                    cv2.circle(frame, tuple(center.squeeze().astype(int)), 4, (0, 165, 255), -1)
                    ellipse = cv2.fitEllipse(contour)
                    cv2.ellipse(frame, ellipse, (255, 0, 255), 2)
                    ecenter = ellipse[0]
                    cv2.circle(frame, tuple(map(int, ecenter)), 5, (255, 165, 0), -1)
                    ckey['center'] = np.array(ecenter, dtype=np.float32)
                    ckey['major_r'] = max(ellipse[1])
                self.display_frame_number(frame, frame_number, n_frames)
                cv2.imshow('Sauron', frame)
                if (cv2.waitKey(5) & 0xFF) == ord('q'):
                    break
            frame_number += 1
        cap.release()
        cv2.destroyAllWindows()

        self.insert1(key)
        for ckey in tqdm(contours):
            self.Ellipse().insert1(ckey, ignore_extra_fields=True)

