import matplotlib
matplotlib.use('agg')

import os, shutil
import allensdk.core.json_utilities as ju
import shutil
import numpy as np
import argparse
import scipy.misc
from scipy.stats import gaussian_kde

import multiprocessing
import functools
import traceback
import logging

from allensdk.brain_observatory.drifting_gratings import DriftingGratings
from allensdk.brain_observatory.static_gratings import StaticGratings
from allensdk.brain_observatory.locally_sparse_noise import LocallySparseNoise
from allensdk.brain_observatory.natural_scenes import NaturalScenes
from allensdk.brain_observatory.natural_movie import NaturalMovie
from allensdk.brain_observatory import observatory_plots as oplots
from allensdk.core.brain_observatory_nwb_data_set import (BrainObservatoryNwbDataSet,
                                                          MissingStimulusException,
                                                          NoEyeTrackingException)
from allensdk.config.manifest import Manifest
from allensdk.internal.core.lims_pipeline_module import PipelineModule, run_module
import allensdk.internal.core.lims_utilities as lu
import allensdk.brain_observatory.stimulus_info as si
from contextlib import contextmanager

LARGE_HEIGHT = 500
SMALL_HEIGHT = 150
SMALL_FONT = 4
LARGE_FONT = 12
PLOT_CONFIGS = { 'small': dict(height_px=SMALL_HEIGHT, pattern="%s_small.png", font_size=SMALL_FONT),
                 'large': dict(height_px=LARGE_HEIGHT, pattern="%s_large.png", font_size=LARGE_FONT),
                 'svg': dict(height_px=LARGE_HEIGHT, pattern="%s.svg", font_size=SMALL_FONT) }
PLOT_TYPES = ["dg", "sg", "ns", "lsn_on", 
              "lsn_off", "rf",
              "nm1", "nm2", "nm3", "sp", 
              "corr", "eye"]

def get_experiment_analysis_file(experiment_id):
    res = lu.query("""
select * from well_known_files wkf 
join well_known_file_types wkft on wkft.id = wkf.well_known_file_type_id
where attachable_id = %d
and wkft.name = 'OphysExperimentCellRoiMetricsFile'
""" % experiment_id)
    return os.path.join(res[0]['storage_directory'], res[0]['filename'])

def get_experiment_nwb_file(experiment_id):
    res = lu.query("""
select * from well_known_files wkf 
join well_known_file_types wkft on wkft.id = wkf.well_known_file_type_id
where attachable_id = %d
and wkft.name = 'NWBOphys'
""" % experiment_id)
    return os.path.join(res[0]['storage_directory'], res[0]['filename'])

def get_experiment_files(experiment_id):
    nwb_file = get_experiment_nwb_file(experiment_id)
    try:
        analysis_file = get_experiment_analysis_file(experiment_id)
    except:
        analysis_file = None

    if not os.path.exists(nwb_file):
        raise Exception("nwb file does not exist: %s" % nwb_file)

    #if not os.path.exists(analysis_file):
   #     raise Exception("analysis file does not exist: %s" % analysis_file)

    return nwb_file, analysis_file

def get_input_data(experiment_id):
    OUTPUT_DIR = "/data/informatics/CAM/analysis/"

    nwb_file, analysis_file = get_experiment_files(experiment_id)
    output_directory = os.path.join(OUTPUT_DIR, str(experiment_id), "thumbnails")
    
    my_file = "/data/informatics/CAM/analysis/%d/%d_analysis.h5" % (experiment_id, experiment_id)

    if os.path.exists(my_file):
        analysis_file = my_file

    input_data = {
        'nwb_file': nwb_file,
        #'analysis_file': analysis_file,
        'analysis_file': analysis_file,
        'output_directory': output_directory
        }

    return input_data

def debug(experiment_id, plots=None, local=False):
    SDK_PATH = "/data/informatics/CAM/analysis/allensdk/"

    input_data = get_input_data(experiment_id)

    run_module(os.path.abspath(__file__),
               input_data,
               input_data["output_directory"],
               sdk_path=SDK_PATH,
               pbs=dict(vmem=32,
                        job_name="bobthumbs_%d"% experiment_id,
                        walltime="10:00:00"),
               local=local,
               optional_args=['--types='+','.join(plots)] if plots else None)

def build_plots(prefix, aspect, configs, output_dir, axes=None, transparent=False):
    Manifest.safe_mkdir(output_dir)

    for config in configs:
        h = config['height_px']
        w = int(h * aspect)
        
        file_name = os.path.join(output_dir, config["pattern"] % prefix)

        logging.debug("file: %s", file_name)
        with oplots.figure_in_px(w, h, file_name, transparent=transparent) as fig:
            matplotlib.rcParams.update({'font.size': config['font_size']})
            yield file_name

def build_cell_plots(cell_specimen_ids, prefix, aspect, configs, output_dir, axes=None, transparent=False):
    for i,csid in enumerate(cell_specimen_ids):
        if np.isnan(csid):
            cell_dir = os.path.join(output_dir, str(i))
        else:
            cell_dir = os.path.join(output_dir, str(csid))

        for fn in build_plots(prefix, aspect, configs, cell_dir, transparent=transparent):
            yield fn, csid, i

def build_drifting_gratings(dga, configs, output_dir):
    for fn in build_plots("drifting_gratings_axes_pref_dir", 1.0, [configs['large'], configs['svg']], output_dir):
        dga.plot_preferred_direction(include_labels=True)
        oplots.finalize_no_axes()

    for fn in build_plots("drifting_gratings_pref_dir", 1.0, [configs['small']], output_dir):
        dga.plot_preferred_direction(include_labels=False)
        oplots.finalize_no_axes()

    for fn in build_plots("drifting_gratings_axes_pref_tf", 1.0, [configs['large'], configs['svg']], output_dir):
        dga.plot_preferred_temporal_frequency()
        oplots.finalize_with_axes()

    for fn in build_plots("drifting_gratings_pref_tf", 1.0, [configs['small']], output_dir):
        dga.plot_preferred_temporal_frequency()
        oplots.finalize_no_labels()

    for fn in build_plots("drifting_gratings_axes_dsi", 1.0, [configs['large'], configs['svg']], output_dir):
        dga.plot_direction_selectivity()
        oplots.finalize_with_axes()

    for fn in build_plots("drifting_gratings_dsi", 1.0, [configs['small']], output_dir):
        dga.plot_direction_selectivity()
        oplots.finalize_no_labels()

    for fn in build_plots("drifting_gratings_axes_osi", 1.0, [configs['large'], configs['svg']], output_dir):
        dga.plot_orientation_selectivity()
        oplots.finalize_with_axes()

    for fn in build_plots("drifting_gratings_osi", 1.0, [configs['small']], output_dir):
        dga.plot_orientation_selectivity()
        oplots.finalize_no_labels()

    csids = dga.data_set.get_cell_specimen_ids()
    for fn, csid, i in build_cell_plots(csids, "drifting_gratings", 1.0, configs.values(), output_dir):
        dga.open_star_plot(csid, include_labels=False, cell_index=i)
        oplots.finalize_no_axes()

def build_static_gratings(sga, configs, output_dir):
    for fn in build_plots("static_gratings_axes_time_to_peak", 1.0, [configs['large'], configs['svg']], output_dir):
        sga.plot_time_to_peak()
        oplots.finalize_with_axes()

    for fn in build_plots("static_gratings_time_to_peak", 1.0, [configs['small']], output_dir):
        sga.plot_time_to_peak()
        oplots.finalize_no_labels()

    for fn in build_plots("static_gratings_axes_pref_ori", 1.5, [configs['large'], configs['svg']], output_dir):
        sga.plot_preferred_orientation(include_labels=True)
        oplots.finalize_no_axes()

    for fn in build_plots("static_gratings_pref_ori", 1.5, [configs['small']], output_dir):
        sga.plot_preferred_orientation(include_labels=False)
        oplots.finalize_no_axes()

    for fn in build_plots("static_gratings_axes_osi", 1.0, [configs['large'], configs['svg']], output_dir):
        sga.plot_orientation_selectivity()
        oplots.finalize_with_axes()

    for fn in build_plots("static_gratings_osi", 1.0, [configs['small']], output_dir):
        sga.plot_orientation_selectivity()
        oplots.finalize_no_labels()

    for fn in build_plots("static_gratings_axes_pref_sf", 1.0, [configs['large'], configs['svg']], output_dir):
        sga.plot_preferred_spatial_frequency()
        oplots.finalize_with_axes()

    for fn in build_plots("static_gratings_pref_sf", 1.0, [configs['small']], output_dir):
        sga.plot_preferred_spatial_frequency()
        oplots.finalize_no_labels()

    csids = sga.data_set.get_cell_specimen_ids()
    for file_name, csid, i in build_cell_plots(csids, "static_gratings_all", 2.0, configs.values(), output_dir):
        sga.open_fan_plot(csid, include_labels=False, cell_index=i)
        oplots.finalize_no_axes()

def build_natural_movie(nma, configs, output_dir, name):
    csids = nma.data_set.get_cell_specimen_ids()
    for file_name, csid, i  in build_cell_plots(csids, name, 1.0, configs.values(), output_dir):
        nma.open_track_plot(csid, cell_index=i)
        oplots.finalize_no_axes()

def build_natural_scenes(nsa, configs, output_dir):
    for fn in build_plots("natural_scenes_axes_time_to_peak", 1.0, [configs['large'], configs['svg']], output_dir):
        nsa.plot_time_to_peak()
        oplots.finalize_with_axes()

    for fn in build_plots("natural_scenes_time_to_peak", 1.0, [configs['small']], output_dir):
        nsa.plot_time_to_peak()
        oplots.finalize_no_labels()

    csids = nsa.data_set.get_cell_specimen_ids()
    for file_name, csid, i in build_cell_plots(csids, "natural_scenes", 1.0, configs.values(), output_dir):
        nsa.open_corona_plot(csid, cell_index=i)
        oplots.finalize_no_axes()

def build_locally_sparse_noise(lsna, configs, output_dir, on):
    prefix = "locally_sparse_noise_" + ("on" if on else "off")

    csids = lsna.data_set.get_cell_specimen_ids()
    for file_name, csid, i in build_cell_plots(csids, prefix, 1.754, [configs['large'], configs['small']], output_dir):
        lsna.open_pincushion_plot(on, cell_specimen_id=csid, cell_index=i)
        oplots.finalize_no_axes()

def build_receptive_field(lsna, configs, output_dir):
    lsn_movie, lsn_mask = lsna.data_set.get_locally_sparse_noise_stimulus_template(lsna.stimulus, 
                                                                                   mask_off_screen=False)

    if lsna.cell_index_receptive_field_analysis_data is None:
        logging.warning("receptive field analysis not performed, so no receptive field plots will be made")
        return

    clim = np.nanpercentile(lsna.receptive_field, [1.0,99.0], axis=None)

    for fn in build_plots("population_receptive_field", 1.754, [configs["large"]], output_dir, transparent=True):
        lsna.plot_population_receptive_field(mask=lsn_mask, scalebar=True)
        oplots.finalize_no_axes()

    for fn in build_plots("population_receptive_field", 1.754, [configs["small"]], output_dir, transparent=True):
        lsna.plot_population_receptive_field(mask=lsn_mask, scalebar=False)
        oplots.finalize_no_axes()

    csids = lsna.data_set.get_cell_specimen_ids()
    for file_name, csid, i in build_cell_plots(csids, "receptive_field_on", 1.754, [configs["large"]], output_dir, transparent=True):
        lsna.plot_cell_receptive_field(True, cell_specimen_id=csid, clim=clim, mask=lsn_mask, cell_index=i, scalebar=True)
        oplots.finalize_no_axes()

    for file_name, csid, i in build_cell_plots(csids, "receptive_field_on", 1.754, [configs["small"]], output_dir, transparent=True):
        lsna.plot_cell_receptive_field(True, cell_specimen_id=csid, clim=clim, mask=lsn_mask, cell_index=i, scalebar=False)
        oplots.finalize_no_axes()

    for file_name, csid, i in build_cell_plots(csids, "receptive_field_off", 1.754, [configs["large"]], output_dir, transparent=True):
        lsna.plot_cell_receptive_field(False, cell_specimen_id=csid, clim=clim, mask=lsn_mask, cell_index=i, scalebar=True)
        oplots.finalize_no_axes()

    for file_name, csid, i in build_cell_plots(csids, "receptive_field_off", 1.754, [configs["small"]], output_dir, transparent=True):
        lsna.plot_cell_receptive_field(False, cell_specimen_id=csid, clim=clim, mask=lsn_mask, cell_index=i, scalebar=False)
        oplots.finalize_no_axes()
       

def build_speed_tuning(analysis, configs, output_dir):
    csids = analysis.data_set.get_cell_specimen_ids()

    for fn in build_plots("running_speed", 1.0, [configs['large'], configs['svg']], output_dir):
        analysis.plot_running_speed_histogram()
        oplots.finalize_with_axes()

    for fn in build_plots("running_speed", 1.0, [configs['small']], output_dir):
        analysis.plot_running_speed_histogram()
        oplots.finalize_no_labels()

    for fn, csid, i in build_cell_plots(csids, "speed_tuning", 1.0, [configs['large'], configs['svg']], output_dir):
        analysis.plot_speed_tuning(csid, cell_index=i)
        oplots.finalize_with_axes()

    for fn, csid, i in build_cell_plots(csids, "speed_tuning", 1.0, [configs['small']], output_dir):
        analysis.plot_speed_tuning(csid, cell_index=i)
        oplots.finalize_no_axes()

def build_correlation_plots(data_set, analysis_file, configs, output_dir):
    sig_corrs = []
    noise_corrs = []

    avail_stims = si.stimuli_in_session(data_set.get_session_type())
    ans = []
    labels = []
    colors = []
    if si.DRIFTING_GRATINGS in avail_stims:
        dg = DriftingGratings.from_analysis_file(data_set, analysis_file)

        if hasattr(dg, 'representational_similarity'):
            ans.append(dg)
            labels.append(si.DRIFTING_GRATINGS_SHORT)
            colors.append(si.DRIFTING_GRATINGS_COLOR)
            setups = [ ( [configs['large']], True ), ( [configs['small']], False )]
            for cfgs, show_labels in setups:
                for fn in build_plots("drifting_gratings_representational_similarity", 1.0, cfgs, output_dir):
                    oplots.plot_representational_similarity(dg.representational_similarity,
                                                            dims=[dg.orivals, dg.tfvals[1:]],
                                                            dim_labels=["dir", "tf"],
                                                            dim_order=[1,0],
                                                            colors=['r','b'],
                                                            labels=show_labels)

    if si.STATIC_GRATINGS in avail_stims:
        sg = StaticGratings.from_analysis_file(data_set, analysis_file)
        if hasattr(sg, 'representational_similarity'):
            ans.append(sg)
            labels.append(si.STATIC_GRATINGS_SHORT)
            colors.append(si.STATIC_GRATINGS_COLOR)
            setups = [ ( [configs['large']], True ), ( [configs['small']], False )]
            for cfgs, show_labels in setups:
                for fn in build_plots("static_gratings_representational_similarity", 1.0, cfgs, output_dir):
                    oplots.plot_representational_similarity(sg.representational_similarity,
                                                            dims=[sg.orivals, sg.sfvals[1:], sg.phasevals],
                                                            dim_labels=["ori", "sf", "ph"],
                                                            dim_order=[1,0,2],
                                                            colors=['r','g','b'], 
                                                            labels=show_labels)

    if si.NATURAL_SCENES in avail_stims:
        ns = NaturalScenes.from_analysis_file(data_set, analysis_file)
        if hasattr(ns, 'representational_similarity'):
            ans.append(ns)
            labels.append(si.NATURAL_SCENES_SHORT)
            colors.append(si.NATURAL_SCENES_COLOR)
            setups = [ ( [configs['large']], True ), ( [configs['small']], False )]
            for cfgs, show_labels in setups:
                for fn in build_plots("natural_scenes_representational_similarity", 1.0, cfgs, output_dir):
                    oplots.plot_representational_similarity(ns.representational_similarity, labels=show_labels)

    if len(ans):
        for an in ans:
            sig_corrs.append(an.signal_correlation)
            extra_dims = range(2,len(an.noise_correlation.shape))
            noise_corrs.append(an.noise_correlation.mean(axis=tuple(extra_dims)))

        for fn in build_plots("correlation", 1.0, [configs['large'], configs['svg']], output_dir):
            oplots.population_correlation_scatter(sig_corrs, noise_corrs, labels, colors, scale=16.0)
            oplots.finalize_with_axes()

        for fn in build_plots("correlation", 1.0, [configs['small']], output_dir):
            oplots.population_correlation_scatter(sig_corrs, noise_corrs, labels, colors, scale=4.0)
            oplots.finalize_no_labels()

        csids = ans[0].data_set.get_cell_specimen_ids()
        for fn, csid, i in build_cell_plots(csids, "signal_correlation", 1.0, [configs['large']], output_dir):
            row = ans[0].row_from_cell_id(csid, i)
            oplots.plot_cell_correlation([ np.delete(sig_corr[row],i) for sig_corr in sig_corrs ], 
                                         labels, colors)
            oplots.finalize_with_axes()

        for fn, csid, i in build_cell_plots(csids, "signal_correlation", 1.0, [configs['small']], output_dir):
            row = ans[0].row_from_cell_id(csid, i)
            oplots.plot_cell_correlation([ np.delete(sig_corr[row],i) for sig_corr in sig_corrs ], 
                                         labels, colors)
            oplots.finalize_no_labels()

        
def lsna_check_hvas(data_set, data_file):
    avail_stims = si.stimuli_in_session(data_set.get_session_type())    
    targeted_structure = data_set.get_metadata()['targeted_structure']

    stim = None

    if targeted_structure == "VISp":
        if si.LOCALLY_SPARSE_NOISE_4DEG in avail_stims:
            stim = si.LOCALLY_SPARSE_NOISE_4DEG
        elif si.LOCALLY_SPARSE_NOISE in avail_stims:
            stim = si.LOCALLY_SPARSE_NOISE
    else:
        if si.LOCALLY_SPARSE_NOISE_8DEG in avail_stims:
            stim = si.LOCALLY_SPARSE_NOISE_8DEG
        elif si.LOCALLY_SPARSE_NOISE in avail_stims:
            stim = si.LOCALLY_SPARSE_NOISE

    if stim is None:
        raise MissingStimulusException("Could not find appropriate LSN stimulus for session %s", 
                                       data_set.get_session_type())
    else:
        logging.debug("in structure %s, using %s stimulus for plots", targeted_structure, stim)
            
    
    return LocallySparseNoise.from_analysis_file(data_set, data_file, stim)


def build_eye_tracking_plots(data_set, configs, output_dir):
    try:
        pupil_times, xy_deg = data_set.get_pupil_location()
        xy_deg = xy_deg[~np.isnan(xy_deg).any(axis=1)]
        c = gaussian_kde(xy_deg.T)(xy_deg.T)

        for fn in build_plots("eye_tracking_gaze_axes", 1.0,
                              [configs['large'], configs['svg']],
                              output_dir):
            oplots.plot_pupil_location(xy_deg, c=c, include_labels=True)
            oplots.finalize_with_axes()

        for fn in build_plots("eye_tracking_gaze", 1.0, [configs['small']],
                              output_dir):
            oplots.plot_pupil_location(xy_deg, c=c, include_labels=False)
            oplots.finalize_no_axes()
    except NoEyeTrackingException:
        logging.debug("No eye tracking found.")
        

def build_type(nwb_file, data_file, configs, output_dir, type_name):
    data_set = BrainObservatoryNwbDataSet(nwb_file)
    try:
        if type_name == "dg":
            dga = DriftingGratings.from_analysis_file(data_set, data_file)
            build_drifting_gratings(dga, configs, output_dir)
        elif type_name == "sg":
            sga = StaticGratings.from_analysis_file(data_set, data_file)
            build_static_gratings(sga, configs, output_dir)
        elif type_name == "nm1":
            nma = NaturalMovie.from_analysis_file(data_set, data_file, si.NATURAL_MOVIE_ONE)
            build_natural_movie(nma, configs, output_dir, si.NATURAL_MOVIE_ONE)
        elif type_name == "nm2":
            nma = NaturalMovie.from_analysis_file(data_set, data_file, si.NATURAL_MOVIE_TWO)
            build_natural_movie(nma, configs, output_dir, si.NATURAL_MOVIE_TWO)
        elif type_name == "nm3":
            nma = NaturalMovie.from_analysis_file(data_set, data_file, si.NATURAL_MOVIE_THREE)
            build_natural_movie(nma, configs, output_dir, si.NATURAL_MOVIE_THREE)
        elif type_name == "ns":
            nsa = NaturalScenes.from_analysis_file(data_set, data_file)
            build_natural_scenes(nsa, configs, output_dir)
        elif type_name == "sp":
            nma = NaturalMovie.from_analysis_file(data_set, data_file, si.NATURAL_MOVIE_ONE)
            build_speed_tuning(nma, configs, output_dir)
        elif type_name == "lsn_on":
            lsna = lsna_check_hvas(data_set, data_file)
            build_locally_sparse_noise(lsna, configs, output_dir, True)
        elif type_name == "lsn_off":
            lsna = lsna_check_hvas(data_set, data_file)
            build_locally_sparse_noise(lsna, configs, output_dir, False)
        elif type_name == "rf":
            lsna = lsna_check_hvas(data_set, data_file)
            build_receptive_field(lsna, configs, output_dir)
        elif type_name == "corr":
            build_correlation_plots(data_set, data_file, configs, output_dir)
        elif type_name == "eye":
            build_eye_tracking_plots(data_set, configs, output_dir)

    except MissingStimulusException as e:
        logging.warning("could not load stimulus (%s)", type_name)
    except Exception as e:
        traceback.print_exc()
        logging.critical("error running stimulus (%s)", type_name)
        raise e

def parse_input(data):
    nwb_file = data.get("nwb_file", None)

    if nwb_file is None:
        raise IOError("input JSON missing required field 'nwb_file'")
    if not os.path.exists(nwb_file):
        raise IOError("nwb file does not exists: %s" % nwb_file)

    analysis_file = data.get("analysis_file", None)

    if analysis_file is None:
        raise IOError("input JSON missing required field 'analysis_file'")
    if not os.path.exists(analysis_file):
        raise IOError("analysis file does not exists: %s" % analysis_file)


    output_directory = data.get("output_directory", None)

    if output_directory is None:
        raise IOError("input JSON missing required field 'output_directory'")

    Manifest.safe_mkdir(output_directory)
    
    return nwb_file, analysis_file, output_directory
    
def build_experiment_thumbnails(nwb_file, analysis_file, output_directory, 
                                types=None, threads=4):
    if types is None:
        types = PLOT_TYPES

    logging.info("nwb file: %s", nwb_file)
    logging.info("analysis file: %s", analysis_file)
    logging.info("output directory: %s", output_directory)
    logging.info("types: %s", str(types))
    Manifest.safe_mkdir(output_directory)

    if len(types) == 1:
        build_type(nwb_file, analysis_file, PLOT_CONFIGS, output_directory, types[0])
    elif threads == 1:
        for type_name in types:
            build_type(nwb_file, analysis_file, PLOT_CONFIGS, output_directory, type_name)
    else:
        p = multiprocessing.Pool(threads)

        func = functools.partial(build_type, nwb_file, analysis_file, PLOT_CONFIGS, output_directory)
        results = p.map(func, types)
        p.close()
        p.join()

    

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-t", "--threads", type=int, default=4)
    parser.add_argument("--log-level", default=logging.DEBUG)
    parser.add_argument("--types", default=','.join(PLOT_TYPES))
    parser.add_argument("input_json")
    parser.add_argument("output_json")
    args = parser.parse_args()

    args.types = args.types.split(',')

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

    input_data = ju.read(args.input_json)

    nwb_file, analysis_file, output_directory = parse_input(input_data)

    build_experiment_thumbnails(nwb_file, analysis_file, output_directory, 
                                args.types, args.threads)

    ju.write(args.output_json, {})

if __name__ == "__main__": main()
