import argparse
import allensdk.internal.core.lims_utilities as lu
import glob
import time
import shutil
import logging
from allensdk.config.manifest import Manifest
from allensdk.internal.brain_observatory.itracker import iTracker
from allensdk.internal.brain_observatory.frame_stream import FfmpegInputStream, FfmpegOutputStream
import h5py
import ast
import sys
import numpy as np

DEFAULT_THRESHOLD_FACTOR = 1.6

if sys.platform=='linux2':
    FFMPEG_BIN = "/shared/utils.x86_64/ffmpeg/bin/ffmpeg"
elif sys.platform=='darwin':
    FFMPEG_BIN = "/usr/local/bin/ffmpeg"

def compute_bounding_box(points):
    if not points:
        return None
    points = np.array(points)
    return [ points[:,0].min(), points[:,0].max(),
             points[:,1].min(), points[:,1].max() ]

def get_polygon(experiment_id, group_name):
    query = """
select ago.* from avg_graphic_objects ago
join avg_graphic_objects pago on pago.id = ago.parent_id
join avg_group_labels agl on pago.group_label_id = agl.id
join sub_images si on si.id = pago.sub_image_id
join specimens sp on sp.id = si.specimen_id
join ophys_sessions os on os.specimen_id = sp.id
where agl.name = '%s'
and os.id = %d
""" % (group_name, experiment_id)

    try:
        path = np.array([ int(v) for v in lu.query(query)[0]['path'].split(',') ])
    except KeyError as e:
        return []
    except IndexError as e:
        return []

    points = path.reshape((len(path)/2, 2))

    return points

def get_experiment_info(experiment_id):
    logging.info("Downloading paths/metadata for experiment ID: %d", experiment_id)
    query = "select storage_directory, id from ophys_sessions where id = "+str(experiment_id)

    storage_directory = lu.query(query)[0]['storage_directory']
    logging.info("\tStorage directory: %s", storage_directory)

    movie_file = glob.glob(storage_directory+'*video-1.avi')[0]
    metadata_file = glob.glob(storage_directory+'*video-1.h5')[0]

    cr_points = get_polygon(experiment_id, 'Corneal Reflection Bounding Box')
    pupil_points = get_polygon(experiment_id, 'Pupil Bounding Box')

    logging.info("\tmovie file: %s", movie_file)
    logging.info("\tmetadata file: %s", metadata_file)

    return dict(movie_file=movie_file,
                metadata_file=metadata_file,
                corneal_reflection_points=cr_points,
                pupil_points=pupil_points)

def get_movie_shape_from_metadata(metadata_file):
    with h5py.File(metadata_file, "r") as f:
        metadata_str = f["video_metadata"].value
        metadata = ast.literal_eval(metadata_str)

    # assuming 3 channels
    # movie_shape = (metadata['frames'], metadata['height'], metadata['width'], 3)
    # in the metadata file from lims, the 'width' and 'height' variables are swapped,
    # hopefully this is the same for every single experiment.
    movie_shape = (metadata['frames'], metadata['width'], metadata['height'], 3)
    logging.info("movie_shape from metadata_file = %s", str(movie_shape))

    return movie_shape

def run_itracker(movie_file, output_directory,
                 output_frames=False,
                 output_annotation_frames=False,
                 output_annotated_movie=True,
                 output_annotated_movie_block_size=1,
                 estimate_bbox=False,
                 num_frames=None,
                 output_QC=True,
                 image_type='png',
                 cache_input_frames=False,
                 input_block_size=1,
                 metadata_file=None,
                 movie_shape=None,
                 **kwargs):

    if output_directory is not None:
        Manifest.safe_mkdir(output_directory)

    assert(metadata_file is not None and movie_shape is not None, "Must provide either metadata_file or movie_shape")

    if metadata_file:
        movie_shape = get_movie_shape_from_metadata(metadata_file)
       
    frame_shape = movie_shape[1:]

    if num_frames is None:
        num_frames = movie_shape[0]


    input_stream = FfmpegInputStream(movie_file, frame_shape,
                                     ffmpeg_bin=FFMPEG_BIN,
                                     num_frames=num_frames,
                                     cache_frames=cache_input_frames,
                                     block_size=input_block_size)

    movie_output_stream = FfmpegOutputStream(frame_shape, 
                                             block_size=output_annotated_movie_block_size,
                                             ffmpeg_bin=FFMPEG_BIN) if output_annotated_movie else None

    itracker = iTracker(output_directory, input_stream=input_stream,
                        im_shape=(movie_shape[1], movie_shape[2]),
                        num_frames=num_frames,
                        **kwargs)

    # open this early to avoid duplicating massive memory
    movie_output_stream.open(itracker.annotated_movie_file)

    itracker.set_movie(movie_file)

    itracker.mean_frame = itracker.compute_mean_frame()

    if estimate_bbox:
        bbox_pupil, bbox_cr = itracker.estimate_bbox_from_mean_frame()


    itracker.process_movie(movie_output_stream=movie_output_stream,
                           output_frames=output_frames,
                           output_annotation_frames=output_annotation_frames)

    if output_QC:
        itracker.output_QC(image_type=image_type)

    return itracker


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--experiment_id', default=None, type=int)
    parser.add_argument('--movie_file', default=None)
    parser.add_argument('--metadata_file', default=None)
    parser.add_argument('--output_directory', default='.')
    parser.add_argument('--estimate_bbox', action='store_true')
    parser.add_argument('--num_frames', default=None, type=int)
    parser.add_argument('--threshold_factor', default=DEFAULT_THRESHOLD_FACTOR)
    parser.add_argument('--log_level', default=logging.DEBUG)
    args = parser.parse_args()

    logging.getLogger().setLevel(args.log_level)

    data = dict(
        threshold_factor=args.threshold_factor,
        output_directory=args.output_directory,
        num_frames=args.num_frames,
        estimate_bbox=args.estimate_bbox
        )

    if args.experiment_id:
        info = get_experiment_info(args.experiment_id)

        data['movie_file'] = info['movie_file']
        data['metadata_file'] = info['metadata_file']

        if info.get('pupil_points', None):
            data['bbox_pupil'] = compute_bounding_box(info['pupil_points'])
        if info.get('corneal_reflection_points', None):
            data['bbox_cr'] = compute_bounding_box(info['corneal_reflection_points'])
    else:
        data['movie_file'] = args.movie_file
        data['metadata_file'] = args.metdata_file

    run_itracker(**data)

if __name__ == "__main__": main()
