# coding: utf-8

import sys
import os
import glob as glob
import copy

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 = copy.deepcopy(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, unit=None, unit_old=None):
        self.open()
        try:
            node = self.save.get_node(node_name)

            if "unit" in node._v_attrs:
                unit_old = node._v_attrs.unit

            if node._v_attrs.CLASS == "GROUP":
                value = {}
                for child_name in node._v_children:
                    value[child_name] = self.get_value(
                        node_name + "/" + child_name, unit, unit_old
                    )
            else:
                value = node.read()
                if not (unit is None or unit_old is None or unit_old == cst.none):
                    value = value * unit_old.express(unit)
        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:
            for key in attrs:
                self.save.get_node(name_full)._v_attrs[key] = attrs[key]

    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 set_attribute(self, node_name, attr_name, attr_value):
        self.open()
        try:
            node = self.save.get_node(node_name)
            node._v_attrs[attr_name] = attr_value
        finally:
            self.close()

    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

    def apply(
        self,
        fn,
        name_array_in,
        name_array_out,
        unit_out=cst.none,
        description="",
        recursive=True,
    ):
        array_in = self.get_value(name_array_in)
        if recursive and isinstance(array_in, dict):
            for key in array_in:
                self.apply(
                    fn,
                    name_array_in + "/" + key,
                    name_array_out + "/" + key,
                    unit_out,
                    description,
                )
                self.set_attribute(name_array_out, "unit", unit_out)
        else:
            try:
                array_out = fn(array_in)
            except TypeError:
                array_out = array_in

            self.set_value(name_array_out, array_out, description, unit_out)


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