f  # coding: utf-8

import sys
import os
import glob as glob


import tables
import pymses
import numpy as np
from numpy.polynomial.polynomial import polyfit
from scipy.stats import linregress

from pymses.sources.ramses import output
from pymses.sources.hop.file_formats import *
from pymses.analysis import Camera, raytracing, slicing, splatting
from pymses.filters import CellsToPoints
from pymses.analysis import ScalarOperator, FractionOperator, MaxLevelOperator
import subprocess

import module_extract as me

from mypool import MyPool as Pool
from functools import partial
from abc import ABCMeta, abstractmethod
import bunch

from run_selector import *
from units import *


class Rule:
    def __init__(
        self,
        postproc,
        process,
        description="",
        group="",
        dependencies=[],
        is_valid=lambda arg: True,
        kind="classic",
        unit=cst.none,
    ):
        self.postproc = postproc
        self.process_fn = process
        self.dependencies = dependencies
        self.is_valid_add = is_valid
        self.group = group
        self.description = description
        self.unit = unit
        self.kind = kind

    def process(self, arg, **kwargs):
        if not arg is None:
            return self.process_fn(arg, **kwargs)
        else:
            return self.process_fn(**kwargs)

    def is_valid(self, arg):
        # save = self.postproc.save
        # valid = True
        # for dep in self.dependencies:
        #     if dep in self.postproc.rules:
        #         rule_dep = self.postproc.rules[dep]
        #         if not arg is None:
        #             valid = valid and rule_dep.group + '/' + dep + '_' + str(arg) in save
        #         else:
        #             valid = valid and rule_dep.group + '/' + dep in save
        # return valid and self.is_valid_add(arg)
        return self.is_valid_add(arg)


class BaseProcessor:
    """
    Base class for processors, should not be instanciated
    """

    __metaclass__ = ABCMeta

    log_id = ""
    rules = {}
    solve_self_dep = True

    def __init__(self, path, path_out=None, pp_params=None, tag=None):
        if pp_params is None:
            self.pp_params = default_params()
        elif type(pp_params) == str:
            self.pp_params = load_params(pp_params)
        else:
            self.pp_params = pp_params

        if tag is not None:
            self.pp_params.out.tag = tag

        # Determining output directory
        if path_out is None:
            self.path_out = path
        else:
            self.path_out = path_out

    def _log(self, string, status=""):
        if self.pp_params.process.verbose:
            if len(status) > 0:
                print(status + ": " + self.log_id + string)
            else:
                print(self.log_id + string)

    def process(
        self,
        to_process_list,
        args=[None],
        overwrite=False,
        overwrite_dep=False,
        **kwargs
    ):
        """
        Render the data in to_process_list and save them
        """

        if type(to_process_list) == str:
            to_process_list = [to_process_list]

        if type(args) == str or args is None:
            args = [args]

        self.overwrite_dep = overwrite_dep
        self.just_done = []  # Computations done within this call

        for name in to_process_list:
            if name in self.rules:
                rule = self.rules[name]
                for arg in args:
                    self._solve_and_process_rule(name, rule, arg, overwrite, **kwargs)
            else:
                self._log(
                    "{} is unknown, allowed rules are {}".format(
                        name, self.rules.keys()
                    ),
                    "ERROR",
                )

        return self.just_done

    def _solve_and_process_rule(self, name, rule, arg, overwrite=False, **kwargs):
        updated = self._solve_dependencies(name, rule, arg, overwrite, **kwargs)
        overwrite_rule = overwrite or updated
        self._process_rule(name, rule, arg, overwrite_rule, **kwargs)

    def _solve_dependencies(self, name, rule, arg, overwrite=False, **kwargs):

        self.done_before_dep = len(self.just_done)

        # Solve dependencies
        for dep in rule.dependencies:
            try:
                dep_arg = rule.dependencies[dep]
            except:
                dep_arg = arg

            if dep_arg == "__parent__":
                dep_arg = arg

            if self.solve_self_dep and dep in self.rules:
                rule_dep = self.rules[dep]
                self._solve_and_process_rule(dep, rule_dep, dep_arg, self.overwrite_dep)
            else:
                self._not_self_dep(name, dep, dep_arg, self.overwrite_dep, **kwargs)

        # Whether dependencies where updated
        return len(self.just_done) > self.done_before_dep

    def _not_self_dep(self, name, dep, dep_arg, overwrite, **kwargs):
        self._log("Dependency {} for {} is unknown".format(dep, name), "ERROR")

    def _needs_computation(self, overwrite, name_full):
        return overwrite

    def _process_rule(self, name, rule, arg, overwrite=False, **kwargs):
        if not arg is None:
            name_full = rule.group + "/" + name + "_" + str(arg)
        else:
            name_full = rule.group + "/" + name

        if rule.is_valid(arg):
            if not name_full in self.just_done:
                if self._needs_computation(overwrite, name_full):
                    self._log("Processing {}".format(name_full))
                    data = rule.process(arg, **kwargs)
                    self._save_data(name_full, data, rule.description, rule.unit)
                    self._log("Data for {} computed".format(name_full), "SUCCESS")
                    self.just_done.append(name_full)
                else:
                    self._log(
                        "Data for {} is already computed, skipping...".format(name_full)
                    )
        else:
            self._log("{} is not valid in this context".format(name_full), "ERROR")

    def def_rules(self):
        for rule in self.rules:
            setattr(self, rule, partial(self.process, rule))


class HDF5Container(BaseProcessor):

    filename = ""
    save = None
    opened = False

    def open(self):
        if not self.opened:
            self.save = tables.open_file(self.filename, mode="a")
            self.opened = True

    def close(self):
        if self.opened:
            self.save.close()
            self.opened = False

    def _needs_computation(self, overwrite, name_full):
        return overwrite or not (name_full in self.save)

    def _process_rule(self, name, rule, arg, overwrite, **kwargs):
        self.open()
        try:
            super(HDF5Container, self)._process_rule(
                name, rule, arg, overwrite, **kwargs
            )
        finally:
            self.close()

    def get_value(self, node_name):
        self.open()
        try:
            node = self.save.get_node(node_name)
            if node._v_attrs.CLASS == "GROUP":
                value = {}
                for child_name in node._v_children:
                    value[child_name] = self.get_value(node_name + "/" + child_name)
            else:
                value = node.read()
        finally:
            self.close()
        return value

    def _save_data(self, name_full, data, description, unit):
        """
        Save data in the HDF5 structure, overwrite if necessary
        """
        if name_full in self.save:
            self.save.remove_node(name_full, recursive=True)

        attrs = None
        if isinstance(data, tuple):
            attrs = data[1]
            data = data[0]

        if isinstance(data, dict):
            if type(description) == str:
                self.save.create_group(
                    os.path.dirname(name_full),
                    os.path.basename(name_full),
                    description,
                    createparents=True,
                )
            else:
                self.save.create_group(
                    os.path.dirname(name_full),
                    os.path.basename(name_full),
                    "",
                    createparents=True,
                )

            if not isinstance(unit, dict):
                self.save.get_node(name_full)._v_attrs.unit = unit

            for key in data:
                if isinstance(description, dict):
                    if isinstance(unit, dict):
                        self._save_data(
                            name_full + "/" + key,
                            data[key],
                            description[key],
                            unit[key],
                        )
                    else:
                        self._save_data(
                            name_full + "/" + key, data[key], description[key], unit
                        )
                else:
                    if isinstance(unit, dict):
                        self._save_data(name_full + "/" + key, data[key], "", unit[key])
                    else:
                        self._save_data(name_full + "/" + key, data[key], "", unit)
        else:
            self.save.create_array(
                os.path.dirname(name_full),
                os.path.basename(name_full),
                data,
                description,
                createparents=True,
            )
            self.save.get_node(name_full).attrs.unit = unit

        if not attrs is None:
            self.save.get_node(name_full).attrs.update(attrs)

    def set_value(self, node_name, data, description, unit):
        self.open()
        try:
            self._save_data(node_name, data, description, unit)
        finally:
            self.close()

    def get_attribute(self, node_name, attr_name):
        self.open()
        try:
            node = self.save.get_node(node_name)
            attr = node._v_attrs[attr_name]
        finally:
            self.close()
        return attr

    def _transform(self, name, transform_fn, group="/maps", **kwargs):
        src = self.save.get_node(group + "/" + name).read()
        return transform_fn(src, **kwargs)

    def _gen_rule_transform(
        self,
        rule_src_name,
        transform_fn,
        transform_name,
        subarray_name=None,
        group=None,
    ):

        rule_src = self.rules[rule_src_name]

        if subarray_name is None:
            src_name = rule_src_name
            group_src = rule_src.group
            unit = rule_src.unit
            description = rule_src.description
        else:
            src_name = subarray_name
            group_src = rule_src.group + "/" + rule_src_name
            unit = rule_src.unit[subarray_name]
            description = rule_src.description[subarray_name]

        def fn(arg=None, **kwargs):
            if arg is None:
                return self._transform(
                    src_name, transform_fn, group=group_src, **kwargs
                )
            else:
                return self._transform(
                    src_name + "_" + str(arg), transform_fn, group=group_src, **kwargs
                )

        if group is None:
            group = group_src

        name = transform_name + "_" + rule_src_name

        self.rules[name] = Rule(
            self,
            fn,
            group=group,
            unit=unit,
            description=description,
            dependencies=[rule_src_name],
        )


def simple_getter(name, dset):
    return dset[name]