# -*- coding: utf-8 -*-
"""
Created on Mon Mar 15 15:08:46 2021

@author: Nick Sauerwein
"""

import lyse
import numpy as np

import os.path
import time

import h5py


def get_mtime(filename):
    return time.ctime(os.path.getmtime(filename))


class DataExtractorManager:

    def __init__(self, data_path=None):

        self.data_extractors = {}
        self.h5_file = None
        self.data_path = data_path
        self.local_data_changed = False

    def update_local_data(self, h5_path):
        with h5py.File(h5_path, 'r') as h5_file:
            self.local_data_changed = False
            for key in self.data_extractors:
                self.data_extractors[key].update_local_data(h5_path, h5_file=h5_file)
                if self.data_extractors[key].local_data_changed:
                    self.local_data_changed = True

    def clean_memory(self, h5_paths):
        for key in self.data_extractors:
            self.data_extractors[key].clean_memory(h5_paths)

    def __getitem__(self, key):
        return self.data_extractors[key]

    def __setitem__(self, key, de):
        self.data_extractors[key] = de

    def add_data_extractor(self, label, type, data_path):

        result = DataExtractor(label, type, data_path)
        self.data_extractors[label] = result

        return result

    def get_data(self, h5_path, h5_file=None):
        result = {}
        for key in self.data_extractors:
            result[key] = self.data_extractors[key].get_data(h5_path, h5_file=h5_file)

        return result


class DataExtractor:
    def __init__(self, label, type, data_path, load_to_ram=True):

        self.load_to_ram = load_to_ram

        self.local_datas = {}
        self.local_mtimes = {}

        if self.load_to_ram:
            self.local_data_changed = True

        # The path of data in hdf5 file
        self.data_path = data_path
        self.data_type = type
        self.data_label = label

    def extract_data(self, h5_path, h5_file=None):

        result = None

        data_handle = lyse.Run(h5_path, no_write=True)

        if self.data_type == 'array':
            try:
                result = data_handle.get_result_array(self.data_path[0], self.data_path[1],
                                                      h5_file=h5_file)
                result = np.array(result)
            except:
                result = None

        elif self.data_type == 'single':
            try:
                    result = data_handle.get_result(self.data_path[0], self.data_path[1],
                                                    h5_file=h5_file)
            except:
                result = None

        elif self.data_type == 'global':
            try:
                    result = data_handle.get_globals(self.data_path[0], h5_file=h5_file)
                    result = result[self.data_path[1]]
            except:
                result = None

        elif self.data_type == 'image':
            try:
                result = data_handle.get_image(self.data_path[0], self.data_path[1],
                                               self.data_path[2], h5_file=h5_file)
                result = np.array(result)
            except:
                result = None

        return result

    def update_local_data(self, h5_path, h5_file=None):
        if h5_path in self.local_datas and self.local_mtimes[h5_path] == get_mtime(h5_path):
            self.local_data_changed = False
        elif self.load_to_ram:
            self.local_datas[h5_path] = self.extract_data(h5_path, h5_file=h5_file)
            self.local_mtimes[h5_path] = get_mtime(h5_path)
            self.local_data_changed = True

    def update_local_datas(self):
        for key in self.local_datas:
            self.update_local_data(key)

    def get_data(self, h5_path, h5_file=None):

        if self.load_to_ram:

            self.update_local_data(h5_path)

            return self.local_datas[h5_path]
        else:
            return self.extract_data(h5_path, h5_file=h5_file)

    def clean_memory(self, h5_paths):
        for key in list(self.local_datas):
            if key not in h5_paths.to_list():
                del self.local_datas[key]
                del self.local_mtimes[key]

                self.local_data_changed = True
        self.update_local_datas()