###############################################################################
#                                                                              #
#   sa2d_decomp_value.py copyright(c) Qiqi Wang 2015 (qiqi.wang@gmail.com)     #
#                                                                              #
################################################################################

import os
import sys
import time
import collections
import copy as copymodule
from subprocess import Popen, PIPE
from io import BytesIO

import numpy as np

def _is_like_sa_value(a):
    '''
    Check attributes of symbolic array value
    '''
    if hasattr(a, 'owner'):
        return a.owner is None or hasattr(a.owner, 'inputs')
    else:
        return False

# ============================================================================ #
#                             symbolic array value                             #
# ============================================================================ #

class symbolic_array_value(object):
    def __init__(self, shape=(), owner=None):
        self.shape = np.empty(shape).shape
        self.owner = owner

    def __repr__(self):
        name = 'value of shape {0}'.format(self.shape)
        if self.owner:
            return 'Dependent {0} generated by {1}'.format(name, self.owner)
        else:
            return 'Independent {0}'.format(name)

    # --------------------------- properties ------------------------------ #

    @property
    def ndim(self):
        return len(self.shape)

    @property
    def size(self):
        return int(np.prod(self.shape))

    def __len__(self):
        return 1 if not self.shape else self.shape[0]


class builtin:
    ZERO = symbolic_array_value()
    I = symbolic_array_value()
    J = symbolic_array_value()
    K = symbolic_array_value()


# ============================================================================ #
#                             computational graph                              #
# ============================================================================ #

def discover_operations_and_inputs(input_values, output_values):
    '''
    Discover all operations required to compute output_values from input_values,
    plus any additional input values
    '''
    discovered_additional_input_values = []
    discovered_operations = []
    def discover_from(v):
        if not hasattr(v, 'owner'):
            return
        if v in input_values:
            return
        if v.owner is None:
            if v not in discovered_additional_input_values:
                discovered_additional_input_values.append(v)
        elif v.owner not in discovered_operations:
            discovered_operations.append(v.owner)
            for v_inp in v.owner.inputs:
                discover_from(v_inp)
    for v in output_values:
        discover_from(v)
    return (sort_operations(discovered_operations),
            discovered_additional_input_values)

def sort_operations(unsorted_operations):
    'Sort operations so that they can be performed in order'
    sorted_operations = []
    def is_computable(v):
        return (not _is_like_sa_value(v) or
                v.owner is None or
                v.owner not in unsorted_operations)
    while len(unsorted_operations):
        removed_any = False
        for op in copymodule.copy(unsorted_operations):
            if all([is_computable(inp) for inp in op.inputs]):
                unsorted_operations.remove(op)
                sorted_operations.append(op)
                removed_any = True
        assert removed_any
    return sorted_operations

def find_dependent(sorted_operations, inputs):
    'Find in a list of sorted operations those that depends on inputs'
    dependent_values = set(inputs)
    dependent_operations = []
    for op in sorted_operations:
        for inp in op.inputs:
            if _is_like_sa_value(inp) and inp in dependent_values:
                dependent_values.add(op.output)
                dependent_operations.append(op)
                break
    return dependent_operations

class AtomicStage(object):
    '''
    Immutable compact stage
    '''
    def __init__(self, input_values, output_values):
        self.input_values = copymodule.copy(input_values)
        self.output_values = copymodule.copy(output_values)
        self.sorted_operations, self.additional_input_values \
            = discover_operations_and_inputs(input_values, output_values)
        for v in output_values:
            assert v.owner is None or v in input_values or \
                   v.owner in self.sorted_operations

    def __call__(self, input_map):
        input_values = self.input_values + self.additional_input_values
        # if hasattr(input_map, '__call__'):
        #     actual_inputs = [input_map(v) for v in input_values]
        # elif hasattr(input_map, '__getitem__'):
        actual_inputs = [input_map[v] for v in input_values]
        # _act attributes are assigned to inputs
        values_assigned_act = set()
        input_values = self.input_values + self.additional_input_values
        for v, v_act in zip(input_values, actual_inputs):
            assert not hasattr(v, '_act')
            v._act = v_act
            values_assigned_act.add(v)
        # _act attributes are computed to each value
        _act = lambda v : v._act if _is_like_sa_value(v) else v
        for op in self.sorted_operations:
            assert not hasattr(op.output, '_act')
            inputs_act = [_act(inp) for inp in op.inputs]
            output_act = op.perform(inputs_act)
            op.output._act = output_act
            values_assigned_act.add(op.output)
        # _act attributes are extracted from outputs then deleted from all
        actual_outputs = tuple(v._act for v in self.output_values)
        for v in values_assigned_act:
            del v._act
        return actual_outputs

    def __hash__(self):
        return id(self)

# ============================================================================ #
#                                decomposition                                 #
# ============================================================================ #

def build_graph(all_values):
    weights = []
    for i, v in enumerate(all_values):
        v._value_id = i
        weights.append(v.size)
    weights.append(1)
    edges = []
    for i, v in enumerate(all_values):
        if not v.owner: continue
        for v_inp in v.owner.inputs:
            if hasattr(v_inp, '_value_id'):
                e = (v_inp._value_id, v._value_id, v.owner.access_neighbor)
                edges.append(e)
    for v in all_values:
        del v._value_id
    return np.array(weights, int), np.array(edges, int)

def decompose_graph(weights, edges, comp_graph_output_file=None):
    my_path = os.path.dirname(os.path.abspath(__file__))
    bin_path = os.path.abspath(os.path.join(my_path, '..', 'bin'))
    quarkflow_bin = os.path.join(bin_path, 'quarkflow')
    p = Popen(quarkflow_bin, stdin=PIPE, stdout=PIPE, stderr=PIPE)
    first_line = '{0} {1}'.format(len(weights) - 1, len(edges))
    weights = ['{0}'.format(w) for w in weights]
    edges = ['{0} {1} {2}'.format(i, j, s) for i, j, s in edges]
    inp = '\n'.join([first_line] + weights + edges)
    if comp_graph_output_file:
        open(comp_graph_output_file, 'w').write(inp)
    out, err = p.communicate(inp.encode())
    assert len(err.strip()) == 0
    return np.loadtxt(BytesIO(out), int).T

def init_operation_input_and_output(sorted_operations, decomp_operations):
    '''
    Find the initialization stage, whose output contains
    values that depends on I,J,K and are required by decomp_operations
    '''
    decomp_inputs = set()
    for op in decomp_operations:
        decomp_inputs.update(filter(_is_like_sa_value, op.inputs))
    init_input = [builtin.I, builtin.J, builtin.K]
    IJK_operations = find_dependent(sorted_operations, init_input)
    init_operations = set(IJK_operations).difference(decomp_operations)
    init_operations = sort_operations(init_operations)
    init_output = set([op.output for op in init_operations]).intersection(
                  decomp_inputs)
    return init_operations, init_input, list(init_output)

def decompose_operations(operations, input_values, output_values,
                         additional_input_values, comp_graph_output_file):
    decomp_values = list(input_values) + [op.output for op in operations]
    weights, edges = build_graph(decomp_values)
    c, d, e = decompose_graph(weights, edges, comp_graph_output_file)
    num_stages = d.max()
    for i, v in enumerate(decomp_values):
        v.create_stage = c[i]
        v.discard_stage = d[i]
    stages = []
    stage_input_values = list(input_values)
    for k in range(1, num_stages):
        next_stage_input_values = list(filter(
                lambda v : v.create_stage <= k and v.discard_stage > k,
                decomp_values))
        stages.append(AtomicStage(stage_input_values + additional_input_values,
                                  next_stage_input_values))
        stage_input_values = next_stage_input_values
    stages.append(AtomicStage(stage_input_values + additional_input_values,
                              list(output_values)))
    return stages

def init_stages_and_output(sorted_operations, decomp_operations,
                           init_graph_output_file):
    init_operations, init_input, init_output = \
            init_operation_input_and_output(sorted_operations,
                                            decomp_operations)
    if len(init_output) == 0:
        return (), []
    init_stages = decompose_operations(init_operations,
                                       init_input, init_output,
                                       [], init_graph_output_file)
    for s in init_stages:
        s.additional_input_values += init_input
    init_stages[0].input_values = [symbolic_array_value(shape=(0,))]
    return init_stages, init_output

def decompose(input_values, output_values, comp_graph_output_file=None,
                                           init_graph_output_file=None):
    sorted_operations, additional_input_values = \
            discover_operations_and_inputs(input_values, output_values)
    decomp_operations = find_dependent(sorted_operations, input_values)
    init_stages, init_output = init_stages_and_output(
            sorted_operations, decomp_operations, init_graph_output_file)
    decomp_stages = decompose_operations(decomp_operations,
                                         input_values, output_values,
                                         init_output, comp_graph_output_file)
    return decomp_stages, init_stages

################################################################################
################################################################################
################################################################################
