from __future__ import absolute_import
import re
from copy import copy

import numpy as nm

from sfepy.base.base import (as_float_or_complex, get_default, assert_,
                             Container, Struct, basestr, goptions)
from sfepy.base.compat import in1d

# Used for imports in term files.
from sfepy.terms.extmods import terms
import six
from six.moves import range
from functools import reduce

_match_args = re.compile('^([^\(\}]*)\((.*)\)$').match
_match_virtual = re.compile('^virtual$').match
_match_state = re.compile('^state(_[_a-zA-Z0-9]+)?$').match
_match_parameter = re.compile('^parameter(_[_a-zA-Z0-9]+)?$').match
_match_material = re.compile('^material(_[_a-zA-Z0-9]+)?$').match
_match_material_opt = re.compile('^opt_material(_[_a-zA-Z0-9]+)?$').match
_match_material_root = re.compile('(.+)\.(.*)').match
_match_ts = re.compile('^ts$').match

def get_arg_kinds(arg_types):
    """
    Translate `arg_types` of a Term to a canonical form.

    Parameters
    ----------
    arg_types : tuple of strings
        The term argument types, as given in the `arg_types` attribute.

    Returns
    -------
    arg_kinds : list of strings
        The argument kinds - one of 'virtual_variable', 'state_variable',
        'parameter_variable', 'opt_material', 'ts', 'user'.
    """
    arg_kinds = []
    for ii, arg_type in enumerate(arg_types):
        if _match_virtual(arg_type):
            arg_kinds.append('virtual_variable')

        elif _match_state(arg_type):
            arg_kinds.append('state_variable')

        elif _match_parameter(arg_type):
            arg_kinds.append('parameter_variable')

        elif _match_material(arg_type):
            arg_kinds.append('material')

        elif _match_material_opt(arg_type):
            arg_kinds.append('opt_material')
            if ii > 0:
                msg = 'opt_material at position %d, must be at 0!' % ii
                raise ValueError(msg)

        elif _match_ts(arg_type):
            arg_kinds.append('ts')

        else:
            arg_kinds.append('user')

    return arg_kinds

def get_shape_kind(integration):
    """
    Get data shape kind for given integration type.
    """
    if integration == 'surface':
        shape_kind = 'surface'

    elif integration in ('volume', 'plate', 'surface_extra'):
        shape_kind = 'volume'

    elif integration == 'point':
        shape_kind = 'point'

    else:
        raise NotImplementedError('unsupported term integration! (%s)'
                                  % integration)

    return shape_kind

def split_complex_args(args):
    """
    Split complex arguments to real and imaginary parts.

    Returns
    -------
    newargs : dictionary
        Dictionary with lists corresponding to `args` such that each
        argument of numpy.complex128 data type is split to its real and
        imaginary part. The output depends on the number of complex
        arguments in 'args':

          - 0: list (key 'r') identical to input one

          - 1: two lists with keys 'r', 'i' corresponding to real
            and imaginary parts

          - 2: output dictionary contains four lists:

            - 'r' - real(arg1), real(arg2)
            - 'i' - imag(arg1), imag(arg2)
            - 'ri' - real(arg1), imag(arg2)
            - 'ir' - imag(arg1), real(arg2)
    """
    newargs = {}
    cai = []

    for ii, arg in enumerate(args):
        if isinstance(arg, nm.ndarray) and (arg.dtype == nm.complex128):
            cai.append(ii)

    if len(cai) > 0:
        newargs['r'] = list(args[:])
        newargs['i'] = list(args[:])

        arg1 = cai[0]
        newargs['r'][arg1] = args[arg1].real.copy()
        newargs['i'][arg1] = args[arg1].imag.copy()

        if len(cai) == 2:
            arg2 = cai[1]
            newargs['r'][arg2] = args[arg2].real.copy()
            newargs['i'][arg2] = args[arg2].imag.copy()

            newargs['ri'] = list(args[:])
            newargs['ir'] = list(args[:])
            newargs['ri'][arg1] = newargs['r'][arg1]
            newargs['ri'][arg2] = newargs['i'][arg2]
            newargs['ir'][arg1] = newargs['i'][arg1]
            newargs['ir'][arg2] = newargs['r'][arg2]

        elif len(cai) > 2:
            raise NotImplementedError('more than 2 complex arguments! (%d)'
                                      % len(cai))

    else:
        newargs['r'] = args[:]

    return newargs

def create_arg_parser():
    from pyparsing import Literal, Word, delimitedList, Group, \
         StringStart, StringEnd, Optional, nums, alphas, alphanums

    inumber = Word("+-" + nums, nums)

    history = Optional(Literal('[').suppress() + inumber
                       + Literal(']').suppress(), default=0)("history")
    history.setParseAction(lambda str, loc, toks: int(toks[0]))

    variable = Group(Word(alphas, alphanums + '._') + history)

    derivative = Group(Literal('d') + variable\
                       + Literal('/').suppress() + Literal('dt'))

    trace = Group(Literal('tr') + Literal('(').suppress() + variable \
                  + Literal(')').suppress())

    generalized_var = derivative | trace | variable

    args = StringStart() + delimitedList(generalized_var) + StringEnd()

    return args

class ConnInfo(Struct):

    def get_region(self, can_trace=True):
        if self.is_trace and can_trace:
            return self.region.get_mirror_region()
        else:
            return self.region

    def get_region_name(self, can_trace=True):
        if self.is_trace and can_trace:
            reg = self.region.get_mirror_region()
        else:
            reg = self.region

        if reg is not None:
            return reg.name
        else:
            return None

class Terms(Container):

    @staticmethod
    def from_desc(term_descs, regions, integrals=None):
        """
        Create terms, assign each term its region.
        """
        from sfepy.terms import term_table

        terms = Terms()
        for td in term_descs:
            try:
                constructor = term_table[td.name]
            except:
                msg = "term '%s' is not in %s" % (td.name,
                                                  sorted(term_table.keys()))
                raise ValueError(msg)

            try:
                region = regions[td.region]
            except IndexError:
                raise KeyError('region "%s" does not exist!' % td.region)

            term = Term.from_desc(constructor, td, region, integrals=integrals)
            terms.append(term)

        return terms

    def __init__(self, objs=None):
        Container.__init__(self, objs=objs)

        self.update_expression()

    def insert(self, ii, obj):
        Container.insert(self, ii, obj)
        self.update_expression()

    def append(self, obj):
        Container.append(self, obj)
        self.update_expression()

    def update_expression(self):
        self.expression = []
        for term in self:
            aux = [term.sign, term.name, term.arg_str,
                   term.integral_name, term.region.name]
            self.expression.append(aux)

    def __mul__(self, other):
        out = Terms()
        for name, term in self.iteritems():
            out.append(term * other)

        return out

    def __rmul__(self, other):
        return self * other

    def __add__(self, other):
        if isinstance(other, Term):
            out = self.copy()
            out.append(other)

        elif isinstance(other, Terms):
            out = Terms(self._objs + other._objs)

        else:
            raise ValueError('cannot add Terms with %s!' % other)

        return out

    def __radd__(self, other):
        return self + other

    def __sub__(self, other):
        if isinstance(other, Term):
            out = self + (-other)

        elif isinstance(other, Terms):
            out = self + (-other)

        else:
            raise ValueError('cannot subtract Terms with %s!' % other)

        return out

    def __rsub__(self, other):
        return -self + other

    def __pos__(self):
        return self

    def __neg__(self):
        return -1.0 * self

    def setup(self):
        for term in self:
            term.setup()

    def assign_args(self, variables, materials, user=None):
        """
        Assign all term arguments.
        """
        for term in self:
            term.assign_args(variables, materials, user)

    def get_variable_names(self):
        out = []
        for term in self:
            out.extend(term.get_variable_names())
        return list(set(out))

    def get_material_names(self):
        out = []
        for term in self:
            out.extend(term.get_material_names())
        return list(set(out))

    def get_user_names(self):
        out = []
        for term in self:
            out.extend(term.get_user_names())
        return list(set(out))

class Term(Struct):
    name = ''
    arg_types = ()
    arg_shapes = {}
    integration = 'volume'
    geometries = ['1_2', '2_3', '2_4', '3_4', '3_8']

    @staticmethod
    def new(name, integral, region, **kwargs):
        from sfepy.terms import term_table

        arg_str = _match_args(name)
        if arg_str is not None:
            name, arg_str = arg_str.groups()

        else:
            raise ValueError('bad term syntax! (%s)' % name)

        if name in term_table:
            constructor = term_table[name]

        else:
            msg = "term '%s' is not in %s" % (name, sorted(
term_table.keys()