import os
import time
import copy
import numpy as np
import pandas as pd
import datetime
from tifffile import *
import matplotlib.pyplot as plt
import json

from xpdacq.config import datapath
from xpdacq.utils import composition_analysis
from xpdacq.control import _get_obj
from xpdacq.analysis import *

pd.set_option('max_colwidth',40)
pd.set_option('colheader_justify','left')


db = _get_obj('db')
get_images = _get_obj('get_images')
get_events = _get_obj('get_events')


##### common functions #####


def table_gen(headers):
    ''' Takes in a header list generated by search functions and return a table
    with metadata information

    Argument:
    headers - list - a list of bluesky header objects

    '''
    plt_list = list()
    feature_list = list()
    uid_list = list()

    if type(list(headers)[0]) == str:
        header_list = []
        header_list.append(headers)
    else:
        header_list = headers

    for header in header_list:
        f_name =_feature_gen(header)
        feature_list.append(f_name)
        try:
            uid_list.append(header.start['uid'][:5])
        except KeyError:
            # jsut in case, it should never happen
            print('Some of your data do not even have a uid, it is very dangerous, please contact beamline scientist immediately')
    plt_list = [feature_list, comment_list] # u_id for ultimate search
    inter_tab = pd.DataFrame(plt_list)
    tab = inter_tab.transpose()
    tab.columns=['Features', 'Comments' ]

    return tab


def time_search(startTime,stopTime=False,exp_day1=False,exp_day2=False):
    '''return list of experiments run in the interval startTime to stopTime

    this function will return a set of events from dataBroker that happened
    between startTime and stopTime on exp_day

    arguments:
    startTime - datetime time object or string or integer - time a the beginning of the
                period that you want to pull data from.  The format could be an integer
                between 0 and 24 to set it at a  whole hour, or a datetime object to do
                it more precisely, e.g., datetime.datetime(13,17,53) for 53 seconds after
                1:17 pm, or a string in the time form, e.g., '13:17:53' in the example above
    stopTime -  datetime object or string or integer - as starTime but the latest time
                that you want to pull data from
    exp_day - str or datetime.date object - the day of the experiment.
    '''
    # date part
    if exp_day1:
        if exp_day2:
            d0 = exp_day1
            d1 = exp_day2
        else:
            d0 = exp_day1
            d1 = d0
    else:
        d0 = datetime.datetime.today().date()
        d1 = d0

    # time part
    if stopTime:

        t0 = datetime.time(startTime)
        t1 = datetime.time(stopTime)

    else:
        now = datetime.datetime.now()
        hr = now.hour
        minu = now.minute
        sec = now.second
        stopTime = datetime.time(hr,minu,sec) # if stopTime is not specified, set current time as stopTime

        t0 = datetime.time(startTime)
        t1 = stopTime

    timeHead = str(d0)+' '+str(t0)
    timeTail = str(d1)+' '+str(t1)

    header_time=db(start_time=timeHead,
                   stop_time=timeTail)

    event_time = get_events(header_time, fill=False)
    event_time = list(event_time)

    print('||You assign a time search in the period:\n'+str(timeHead)+' and '+str(timeTail)+'||' )
    print('||Your search gives out '+str(len(header_time))+' results||')

    return header_time

# FIXME - Refactor search function !!!!!!
#### block of search functions ####
def _list_keys( d, container):
    ''' list out all keys in dictionary, d
    Argument:
        d - dict - input dictionary
        containter - list - auxiliary list
    '''
    if isinstance(d, dict):
        # append keys in the first layer
        container += d.keys()
        # add keys in next layer
        list(map(lambda x: _list_keys(x, container), d.values()))
    elif isinstance(d, list):
        list(map(lambda x: _list_keys(x, container), d))
    


def get_keys(fuzzy_key, d=None, verbose=0):
    ''' fuzzy search on key names contained in a nested dictionary.
    Return all possible key names starting with fuzzy_key:
    Arguments:

    fuzzy_key - str - possible key name, can be fuzzy like 'exp', 'sca' or nearly complete like 'experiment'
    d        -- dictionary you want to search.  Use bluesky metadata store
                when not specified.
    '''
    if d is None:
        #d = _bluesky_metadata_store()
        d = gs.RE.md

    container = []
    _list_keys(d, container)
    
    if verbose:
        # default is not verbose
        print('All keys in target dictionary are: %s' % str(container))
    
    # filter out desired name
    out = list(filter(lambda x: x.startswtih(fuzzy_key), container))
        # just a practice. It is equivalent to out = [ x for x in container if x.startswith(fuzzy_key)]
    return out

def get_keychain(wanted_key, d=None):
    ''' Return keychian(s) of specific key(s) in a nested dictionary

    argumets:
    wanted_key - str - name of key you want to search for
    d        -- dictionary you want to search.  Use bluesky metadata store
                when not specified.
    '''
    if d is None:
        #d = _bluesky_metadata_store()
        d = gs.RE.md
    for k, v in d.items():
        if isinstance(v, dict):
            result = get_keychain(wanted_key, v) # dig in nested element
            if result:
                return [k]+ result
        elif k == wanted_key:
            return [k]

def set_value(key, new_value, d):
    ''' update value of corresponding key in a nested dictionary
    
    argumet:
    key - str - key you want to change
    new_value - str - value of key to-be-updated
    d - dict - target dictionary
    '''
    cur = d
    keychain = get_keychain(key, d)
    key_oper = keychain[-1]
    for path_item in keychain[:-1]:
        try:
            cur = cur[path_item]
        except KeyError:
            cur = cur[path_item] = {}

    old_value = cur[keychain[-1]]
    cur[keychain[-1]] = new_value
    print('Values to key %s has been updated from %s to %s' %(key, old_value, new_value))
    print(d)


def build_keychain_list(key_list, d=None, verbose = 1):
    ''' Return a keychain list that yields all parent keys for every key in key_list
        E.g. d = {'layer1':{'layer2':{'mykey':'value'}}}
            build_keychain_list([layer2, mykey],d) = ['layer1', 'layer1.layer2']
    argumets:
    key_list - str or list - name of key(s) you want to search for
    d        -- dictionary you want to search.  Use bluesky metadata store
                when not specified.
    '''
    if d is None:
        #d = _bluesky_metadata_store()
        d = gs.RE.md
    result = []
    if isinstance(key_list, str):
        key_list_operate = []
        key_list_operate.append(key_list)
    elif isinstance(key_list, list):
        key_list_operate = key_list

    for key in key_list_operate:
        dummy = get_keychain(key)
        if dummy:
            if len(dummy) > 1:
                path = '.'.join(dummy)
                result.append(path)
            elif len(dummy) == 1: # key at first level
                path = dummy[0]
                result.append(path)
        else:  # find an empty dictionary
            path = key
            result.append(path)
        if verbose:
            print('keychain to your desired key %s is "%s"' % (key, path))
        else:
            pass
    return result

def search(desired_value, *args, **kwargs):
    '''Return all possible header(s) that satisfy your searching criteria

    this function operates in two logics:
    1) When desired_value and args are both given. It will search on all headers matches args = desired_value.
        args can be incomplete and in this case, this function yields multiple searches

    example:
    desired_value = 'TiO2'
    search(desired_value, *'sa') will return all headers that has keys starting with 'sa' and its corresponding
    values is 'TiO2' in metadata dictionary. Nanmely, it yields searchs on headers with sample = TiO2, sadness = TiO2 ...

    2) When desired_value is not given. It implies you already knew your searching criteria and are ready to type them explicitly,
        even with additional constrains.

    example:
    desired_value = 'TiO2'
    search (False, **{'sample_name':desired_value, 'additonal_field': 'additional_value ....}) will return
    headers that have exactly key pairs **{'sample_name':desired_value, 'additonal_field': 'additional_value ....}

    General stratege is to use the first logic to figure out what is your desired key.
    Then user the second logic to restrain your search

    arguments:
    desired_value - str - desired value you are looking for
    args - str - key name you want to search for. It can be fuzzy or complete. If it is fuzzy, all possibility will be listed.
    kwargs - dict - an dictionary that contains exact key-value pairs you want to search for

    '''
    if desired_value and args:
        possible_keys = get_keys(args)
        keychain_list = build_keychain_list(possible_keys, verbose =0)
        search_header_list = []
        for i in range(len(keychain_list)):
            dummy_search_dict = {}
            dummy_search_dict[keychain_list[i]] = desired_value
            dummy_search_dict['group'] = 'XPD' # create an anchor as mongoDB and_search needs at least 2 key-value pairs
            search_header = db(**dummy_search_dict)
            search_header_list.append(search_header)
            print('Your %ith search "%s=%s" yields %i headers' % (i,
                keychain_list[i], desired_value, len(search_header)))
            return search_header_list
    elif not desired_value and kwars:
        if len(kwargs)>1:
            search_header = db(**kwargs)
        elif len(kwargs) == 1:
            kwargs['group'] = 'XPD'
            search_header = db(**kwargs)
        else:
            print('You gave empty search criteria. Please try again')
            return
        return search_header
    else:
        print('Sorry, your search is somehow unrecongnizable. Please make sure you are putting values to right fields')


