import matplotlib
matplotlib.use('agg')

import logging
import numpy as np
import os, sys

from allensdk.internal.core.lims_pipeline_module import PipelineModule, run_module
from allensdk.internal.brain_observatory.run_itracker import (run_itracker, 
                                                              compute_bounding_box, 
                                                              DEFAULT_THRESHOLD_FACTOR,
                                                              get_experiment_info)

def debug(experiment_id, num_frames=None, threshold_factor=None, local=False):
    OUTPUT_DIR = "/data/informatics/CAM/eye_tracking/"
    SDK_PATH = "/data/informatics/CAM/eye_tracking/allensdk"
    SCRIPT_PATH = "/data/informatics/CAM/eye_tracking/allensdk/allensdk/internal/pipeline_modules/run_eye_tracking.py" 
    
    experiment_dir = os.path.join(OUTPUT_DIR, str(experiment_id))
    
    info = get_experiment_info(experiment_id)
    info['output_directory'] = experiment_dir

    optional_args = [ ]
    if num_frames is not None:
        optional_args += ['--num_frames',str(num_frames)]

    run_module(SCRIPT_PATH, 
               info, 
               experiment_dir, 
               sdk_path=SDK_PATH,
               pbs=dict(vmem=160,
                        job_name="itrack_%d"% experiment_id,
                        walltime="10:00:00"),
               local=local,
               optional_args=optional_args)
    
def main():
    mod = PipelineModule()
    mod.parser.add_argument("--num_frames", type=int, default=None)
    mod.parser.add_argument("--threshold_factor", type=float, default=DEFAULT_THRESHOLD_FACTOR)

    data = mod.input_data()
    args = dict(
        movie_file=data['movie_file'],
        metadata_file=data['metadata_file'],
        output_directory=data['output_directory'],
        threshold_factor=data.get('threshold_factor', mod.args.threshold_factor),
        num_frames=mod.args.num_frames,
        auto=True,
        cache_input_frames=True,
        input_block_size=None,
        output_annotated_movie_block_size=None
        )

    if data.get('pupil_points', None):
        args['bbox_pupil'] = compute_bounding_box(data['pupil_points'])
    if data.get('corneal_reflection_points', None):
        args['bbox_cr'] = compute_bounding_box(data['corneal_reflection_points'])
    
    tracker = run_itracker(**args)

    logging.debug("finished running itracker")
    
    output_data = dict(
      pupil_file=tracker.pupil_file,
      corneal_reflection_file=tracker.cr_file,
      mean_frame_file=tracker.mean_frame_file
      )
    
    mod.write_output_data(output_data)
  
if __name__=='__main__': main()
