#!/usr/bin/env python3
'''
Analyse Gonio Imsoft data and output the results.
'''
import sys
import os
import datetime
import argparse

import numpy as np
import matplotlib.pyplot as plt

from gonioanalysis.drosom import analyser_commands
from gonioanalysis.drosom.analyser_commands import (
        ANALYSER_CMDS,
        DUALANALYSER_CMDS,
        MULTIANALYSER_CMDS,
        )
from gonioanalysis.directories import ANALYSES_SAVEDIR, PROCESSING_TEMPDIR_BIGFILES
from gonioanalysis.droso import DrosoSelect
from gonioanalysis.antenna_level import AntennaLevelFinder
from gonioanalysis.drosom.analysing import MAnalyser, MAverager
from gonioanalysis.drosom.orientation_analysis import OAnalyser
from gonioanalysis.drosom.optic_flow import FAnalyser
from gonioanalysis.drosom.transmittance_analysis import TAnalyser
from gonioanalysis.drosom import plotting
from gonioanalysis.drosom.plotting.plotter import MPlotter
from gonioanalysis.drosom.plotting import complete_flow_analysis, error_at_flight
from gonioanalysis.drosom.special.norpa_rescues import norpa_rescue_manyrepeats
from gonioanalysis.drosom.special.paired import cli_group_and_compare
import gonioanalysis.drosom.reports as reports


if '--tk_waiting_window' in sys.argv:
    from gonioanalysis.tkgui.widgets import WaitingWindow



Analysers = {'orientation': OAnalyser, 'motion': MAnalyser, 'flow': FAnalyser,
        'transmittance': TAnalyser}

analyses = {**ANALYSER_CMDS, **DUALANALYSER_CMDS, **MULTIANALYSER_CMDS}


def roimovement_video(analyser):
    '''
    Create a video where the imaging data is played and the analysed ROI is moved on the
    image, tracking the moving feature.

    Good for confirming visually that the movment analysis works.
    '''

    print(analyser.getFolderName())
    images, ROIs, angles = analyser.get_time_ordered()
    
    workdir = os.path.join(PROCESSING_TEMPDIR_BIGFILES, 'movie_{}'.format(str(datetime.datetime.now())))
    os.makedirs(workdir, exist_ok=True)

    newnames = [os.path.join(workdir, '{:>0}.jpg'.format(i)) for i in range(len(images))]
    

    adj = ROIAdjuster()
    newnames = adj.writeAdjusted(images, ROIs, newnames, extend_factor=3, binning=1)
    
    enc = Encoder()
    fps = 25
    enc.encode(newnames, os.path.join(ANALYSES_SAVEDIR, 'movies','{}_{}fps.mp4'.format(analyser.getFolderName(), fps)), fps)
 
    for image in newnames:
        os.remove(image)
    try:
        os.rmdir(workdir)
    except OSError:
        print("Temporal directory {} left behind because it's not empty".format(workdir))
       

   
def export_optic_flow():
    '''
    Exports the optic flow vectors.
    '''
    import json
    from gonioanalysis.coordinates import optimal_sampling
    from gonioanalysis.drosom.optic_flow import flow_vectors
    points = optimal_sampling(np.arange(-90, 90, 5), np.arange(-180, 180, 5))
    vectors = flow_vectors(points)
    
    with open('optic_flow_vectors.json', 'w') as fp:
        json.dump({'points': np.array(points).tolist(), 'vectors': np.array(vectors).tolist()}, fp)
            


         
def main(custom_args=None):
    
    if custom_args is None:
        custom_args = sys.argv[1:]
    
    parser = argparse.ArgumentParser(description=__doc__)
    
    
    # DATA ARGUMENTS
    parser.add_argument('-D', '--data_directory', nargs='+',
            help='Data directory')

    parser.add_argument('-S', '--specimens', nargs='+',
            help=('Comma separeted list of specimen names.'
                ' Separate groups by space when averaging is on.'
                ' If needed, wanted image folders can be specified with semicolons'
                ' for example, specimen1;imf1;imf2:specimen2;imf1'
                ' (does not work with groups). If image folders contain commas, use'
                ' semicolons, use colons (:) to separate specimens'))
    

    # Analyser settings
    parser.add_argument('-a', '--average', action='store_true',
            help='Average and interpolate the results over the specimens')
    
    parser.add_argument('--short-name', nargs=1,
            help='Short name to set if --average is set')

    parser.add_argument('-t', '--type', nargs='+',
            help='Analyser type, either "motion" or "orientation". Space separate gor groups')
    
    parser.add_argument('-r', '--reselect-rois', action='store_true',
            help='Reselect ROIs')
 
    parser.add_argument('-R', '--recalculate-movements', action='store_true',
            help='Recalculate with Movemeter')

    parser.add_argument('--reverse-directions', action='store_true',
            help='Reverse movement directions')

    parser.add_argument('--active-analysis', nargs='?', default='',
            help='Name of the analysis')

    # Other settings
    parser.add_argument('--tk_waiting_window', action='store_true',
            help='(internal) Launches a tkinter waiting window')
    parser.add_argument('--dont-show', action='store_true',
            help='Skips showing the plots')
    parser.add_argument('--worker-info', nargs=2,
            help='Worker id and total number of parallel workers. Only 3D video plotting now')

    # Different analyses for separate specimens

    parser.add_argument('-A', '--analysis', nargs=1,
            choices=analyses.keys(),
            help='Analysis method or action. Allowed analyses are '+', '.join(analyses.keys()))

    args = parser.parse_args(custom_args)
    



    if args.tk_waiting_window:
        waiting_window = WaitingWindow('terminal.py', 'When ready, this window closes.')

    if args.worker_info:
        analyser_commands.I_WORKER = int(args.worker_info[0])
        analyser_commands.N_WORKERS = int(args.worker_info[1])

    # Getting the data directory
    # --------------------------
    if args.data_directory:
        print('Using data directory {}'.format(args.data_directory[0]))
        
        data_directories = args.data_directory
    else:
        data_directories = input('Data directory >> ')
    
    # Check that the data directory exists
    for directory in data_directories:
        if not os.path.isdir(directory):
            raise ValueError("{} is not a directory".format(directory))

    # {specimen name : [image_folder_1, ...]}
    wanted_imagefolders = {}

    # Getting the specimens
    # ---------------------
    directory_groups = []
    if args.specimens:
        
        for group in args.specimens:
            print('Using specimens {}'.format(group))
            
            if group == 'none':
                directory_groups.append(None)
                continue

            if ':' in group:
                specimen_separator = ':'
            else:
                specimen_separator = ','

            if ';' in group:
                
                for specimen in group.split(specimen_separator):
                    if ';' in specimen:
                        splitted = specimen.split(';')
                        wanted_imagefolders[splitted[0]] = splitted[1:]
                
                # Remove specified imagefolders
                group = ','.join([z.split(';')[0] for z in group.split(specimen_separator)])
            
            
            # dont commit me
            group = group.replace(':', ',')
            
            
            directories = []
            for directory in data_directories:
                selector = DrosoSelect(datadir=directory)
                directories.extend( selector.parse_specimens(group) )
            directory_groups.append(directories)
    else:
        selector = DrosoSelect(datadir=data_directories[0])
        directories = selector.ask_user()
     
            
    # Setting up analysers
    # ---------------------
    
    if not args.type:
        args.type = ['motion' for i in directory_groups]

    analyser_groups = []
    
    for i_group, directories in enumerate(directory_groups):

        analysers = []
        Analyser = Analysers[args.type[i_group]]
        
        print('Using {}'.format(Analyser.__name__))
        
        if directories is None:
            analysers.append(Analyser(None, None))
        else:

            for directory in directories: 
                
                path, folder_name = os.path.split(directory)
                analyser = Analyser(path, folder_name) 
                
                if args.active_analysis:
                    analyser.active_analysis = args.active_analysis
                
                analysers.append(analyser)
 
        # Ask ROIs if not selected
        for analyser in analysers:
            if analyser.are_rois_selected() == False or args.reselect_rois:
                analyser.select_ROIs()

        # Analyse movements if not analysed, othewise load these
        for analyser in analysers:
            if analyser.is_measured() == False or args.recalculate_movements:
                analyser.measure_movement(eye='left')
                analyser.measure_movement(eye='right')
            analyser.load_analysed_movements()
        
        
        if args.reverse_directions:
            for analyser in analysers:
                analyser.receptive_fields = True
        
       

        if args.average:
            
            if len(analysers) >= 2:

                avg_analyser = MAverager(analysers)
                avg_analyser.setInterpolationSteps(5,5)
                
                if args.short_name:
                    avg_analyser.set_short_name(args.short_name[0])
                           
                analysers = avg_analyser
            else:
                analysers = analysers[0]
        else:
            if len(analysers) == 1:
                analysers = analysers[0]

        analyser_groups.append(analysers)
    

    function = analyses[args.analysis[0]]
    
    print(analyser_groups)
    
    kwargs = {}
    if wanted_imagefolders:
        kwargs['wanted_imagefolders'] = wanted_imagefolders
    
    if function in MULTIANALYSER_CMDS.values():
        for analysers in analyser_groups:
            function(analysers, **kwargs)
    elif args.average or function in DUALANALYSER_CMDS.values():
        function(*analyser_groups, **kwargs)
    else:
        for analysers in analyser_groups:
            if not isinstance(analysers, list):
                analysers = [analysers]
            for analyser in analysers:
                function(analyser, **kwargs)

    if args.tk_waiting_window:
        waiting_window.close()

    if not args.dont_show:
        plt.show()


if __name__ == "__main__":
    main()
