import pathlib
import numpy as np
from .utils import replace_all

class logviewer:
    """ log viewer
        Since lognflow makes lots of files and folders, maybe it is nice
        to have a logviewer that loads those information. In this module we
        provide a set of functions for a logged object that can load variables,
        texts, file lists and etc.. Use it simply by::
 
            from lognflow import logviewer
            logged = logviewer(log_dir = 'dir_contatining_files')
            var = logged.get_single('variable_name')
    """ 
    def __init__(self,
                 log_dir : pathlib.Path,
                 logger = print):
        self.log_dir = pathlib.Path(log_dir)
        assert self.log_dir.is_dir(), \
            f'lognflow.logviewer| No such directory: '+ str(self.log_dir)
        self.logger = logger
        self.load = self.get_single
    
    def get_text(self, log_name='main_log'):
        """ get text log files
            Given the log_name, this function returns the text therein.

            Parameters
            ----------
            :param log_name:
                the log name. If not given then it is the main log.
        """
        flist = list(self.log_dir.glob(f'{log_name}*.txt'))
        flist.sort()
        n_files = len(flist)
        if (n_files>0):
            txt = []
            for fcnt in range(n_files):
                with open(flist[fcnt]) as f_txt:
                    txt.append(f_txt.readlines())
            if(n_files == 1):
                txt = txt[0]
            return txt

    def get_flist(self, var_name, suffix = None):
        """ get list of files
            return the list of files for a saved variable.

            Parameters
            ----------
            :param var_name:
                variable name
            :param suffix:
                If there are different suffixes availble for a variable
                this input needs to be set. npy, npz, mat, and torch are
                supported.
        """
        var_name = var_name.replace('\t', '\\t').replace('\n', '\\n')\
            .replace('\r', '\\r').replace('\b', '\\b')
        
        if suffix is None:
            if len(var_name.split('.')) > 1:
                suffix = var_name.split('.')[-1]
                name_before_suffix = var_name.split('.')[:-1]
                if((len(name_before_suffix) == 1) & 
                   (name_before_suffix[0] == '')):
                    var_name = '*'
                else:
                    var_name = ('.').join(var_name.split('.')[:-1])
            else:
                suffix = '.np*'

        suffix = suffix.strip('.')        

        flist = []            
        if((self.log_dir / var_name).is_file()):
            flist = [self.log_dir / var_name]
        elif((self.log_dir / f'{var_name}.{suffix}').is_file()):
            flist = [self.log_dir / f'{var_name}.{suffix}']
        else:
            _var_name = (self.log_dir / var_name).name
            _var_dir = (self.log_dir / var_name).parent
            search_patt = f'{_var_name}*.{suffix}'
            search_patt = replace_all(search_patt, '**', '*')
            flist = list(_var_dir.glob(search_patt))
            if(len(flist) == 0):
                search_patt = f'{_var_name}*.*'
                search_patt = replace_all(search_patt, '**', '*')
                flist = list(_var_dir.glob(search_patt))
                if(len(flist) > 0):
                    self.logger(
                        'I Can not find the file with the given suffix, '\
                        + 'but found some with a different suffix, '\
                        + f'one file is: {flist[-1]}. This is what I'\
                        + ' will return.' )
                    
        if(len(flist) > 0):
            flist.sort()
        else:
            var_dir = self.log_dir / var_name
            if(var_dir.is_dir()):
                flist = list(var_dir.glob('*.*'))
            if(len(flist) > 0):
                flist.sort()
        return flist
            
    def get_single(self, var_name, file_index = -1, suffix = None):
        """ get a single variable
            return the value of a saved variable.

            Parameters
            ----------
            :param var_name:
                variable name
            :param file_index:
                If there are many snapshots of a variable, this input can
                limit the returned to a set of indices.
            :param suffix:
                If there are different suffixes availble for a variable
                this input needs to be set. npy, npz, mat, and torch are
                supported.
                
            .. note::
                when reading a MATLAB file, the output is a dictionary.
        """
        assert file_index == int(file_index), \
                    f'file_index {file_index} must be an integer'
        flist = self.get_flist(var_name, suffix)
        var_path = None
        if flist:
            var_path = flist[file_index]
    
            if(var_path.is_file()):
                self.logger(f'Loading {var_path}')
                if(var_path.suffix == '.npz'):
                    buf = np.load(var_path)
                    try:
                        time_array = buf['time_array']
                        n_logs = (time_array > 0).sum()
                        time_array = time_array[:n_logs]
                        data_array = buf['data_array']
                        data_array = data_array[:n_logs]
                        return((time_array, data_array))
                    except:
                        return(buf)
                if(var_path.suffix == '.npy'):
                    return(np.load(var_path))
                if(var_path.suffix == '.mat'):
                    from scipy.io import loadmat
                    return(loadmat(var_path))
                if( (var_path.suffix == '.txt')
                   |(var_path.suffix == '.pdb')
                   |(var_path.suffix == '.json')):
                    with open(var_path) as f_txt:
                        return(f_txt.read())
                if((var_path.suffix == '.tif') | (var_path.suffix == '.tiff')):
                    from tifffile import imread
                    return(imread(var_path))
                if(var_path.suffix == '.torch'):      
                    from torch import load as torch_load 
                    return(torch_load(var_path))
                try:
                    from matplotlib.pyplot import imread
                    img = imread(var_path)
                    return(img)
                except:
                    var_path = None
            else:
                var_path = None
                
        if var_path is None:
            self.logger(f'{var_name} does not resemble the name of any file '
                        f'or directory in the log_dir: {self.log_dir}')
    
    def get_stack_of_files(self, 
        var_name = None, flist = [], suffix = None,
        return_data = False, return_flist = True, read_func = None,
        data_makes_a_block = True, verbose = True):
       
        """ Get list or data of all files in a directory
       
            This function gives the list of paths of all files in a directory
            for a single variable.

            Parameters
            ----------
            :param var_name:
                The directory or variable name to look for the files
            :type var_name: str
           
            :param flist:
                list of Paths, if data is returned, this flist input can limit
                the data requested to this list.
            :type flist: list
           
            :param suffix:
                the suffix of files to look for, e.g. 'txt'
            :type siffix: str
           
            :param return_data:
                    with flist you can limit the data that is returned.
                    Otherwise the data for all files in the directory will be
                    returned
            :param return_flist:
                    Maybe you are only intrested in the flist.
                   
            :param data_makes_a_block:
                    if you know that the shape of all numpy arrays in files, or
                    images are the same, set this as true and receive a numpy
                    array. Otherwise returns a list. default: False
           
            Output
            ----------
           
                It returns a tuple, (dataset, flist),
                dataset will be a list of files contains or numpy array in
                case all files produce same shape numpy arrays.
                flist is type pathlib.Path
           
        """
        if not flist:
            flist = self.get_flist(var_name, suffix)
        
        if flist:
            flist.sort()
            n_files = len(flist)
            if((not return_data) & return_flist):
                return(flist)
           
            if(read_func is None):
                try:
                    fdata = np.load(flist[0])
                    read_func = np.load
                except:
                    pass
            if(read_func is None):
                try:
                    from matplotlib.pyplot import imread
                    fdata = imread(flist[0])
                    read_func = imread
                except:
                    pass
            if(read_func is not None):
                assert callable(read_func), \
                    f'given read_func: {read_func} is not callable.'
                fdata = read_func(flist[0])
                if(data_makes_a_block):
                    dataset = np.zeros((n_files, ) + fdata.shape,
                                       dtype=fdata.dtype)
                    if(verbose):
                        self.logger(f'logviewer: Reading dataset from '
                                    f'{flist[0].parent}'
                                    f', the shape would be: {dataset.shape}')
                    for fcnt, fpath in enumerate(flist):
                        try:
                            dataset[fcnt] = read_func(fpath)
                        except e:
                            self.logger(
                                'I cannot concatenate the data I just read on'
                                ' top of last files data. You need to set'
                                ' data_makes_a_block = False in the input'
                                f' of get_stack_of_files({var_name}, ...)'
                                ' This will return a list of data of files')
                            raise e
                else:
                    dataset = [read_func(fpath) for fpath in flist]
                if(return_flist):
                    return(dataset, flist)
                else:
                    return(dataset)
            else:
                if(verbose):
                    self.logger(f'File {flist[0].name} cannot be opened by '\
                          + r'np.load() or plt.imread(), provide read_func.')

    def get_common_files(self, var_name_A, var_name_B):
        """ get common files in two directories
        
            It happens often in ML that there are two directories, A and B,
            and we are interested to get the flist in both that is common 
            between them. returns a tuple of two lists of files.
            
            Parameters
            ----------
            :param var_name_A:
                directory A name
            :param var_name_B:
                directory B name
        """
        flist_A = self.get_stack_of_files(
            var_name_A, return_data = False, return_flist = True)
        flist_B = self.get_stack_of_files(
            var_name_B, return_data = False, return_flist = True)
        
        suffix_A = flist_A[0].suffix
        suffix_B = flist_B[0].suffix 
        parent_A = flist_A[0].parent
        parent_B = flist_B[0].parent
        
        fstems_A = [_fst.stem for _fst in flist_A]
        fstems_B = [_fst.stem for _fst in flist_B]
        
        fstems_A_set = set(fstems_A)
        fstems_B_set = set(fstems_B)
        common_stems = list(fstems_A_set.intersection(fstems_B_set))

        flist_A_new = [parent_A / (common_stem + suffix_A) \
                          for common_stem in common_stems]
        flist_B_new = [parent_B / (common_stem + suffix_B) \
                          for common_stem in common_stems]

        return(flist_A_new, flist_B_new)
    
    def __repr__(self):
        return f'{self.log_dir}'

    def __bool__(self):
        return self.log_dir.is_dir()