from __future__ import absolute_import
from copy import copy

import numpy as nm

from sfepy.base.testing import TestCommon
from sfepy.base.base import ordered_iteritems
from sfepy import data_dir

filename_meshes = [data_dir + '/meshes/elements/%s_2.mesh' % geom
                   for geom in ['1_2', '2_3', '2_4', '3_4', '3_8', '3_2_4']]

def make_term_args(arg_shapes, arg_kinds, arg_types, ats_mode, domain,
                   material_value=None, poly_space_base=None):
    from sfepy.base.base import basestr
    from sfepy.discrete import FieldVariable, Material, Variables, Materials
    from sfepy.discrete.fem import Field
    from sfepy.solvers.ts import TimeStepper
    from sfepy.mechanics.tensors import dim2sym

    omega = domain.regions['Omega']
    dim = domain.shape.dim
    sym = dim2sym(dim)

    def _parse_scalar_shape(sh):
        if isinstance(sh, basestr):
            if sh == 'D':
                return dim

            elif sh == 'D2':
                return dim**2

            elif sh == 'S':
                return sym

            elif sh == 'N': # General number ;)
                return 1

            else:
                return int(sh)

        else:
            return sh

    def _parse_tuple_shape(sh):
        if isinstance(sh, basestr):
            return [_parse_scalar_shape(ii.strip()) for ii in sh.split(',')]

        else:
            return (int(sh),)

    args = {}
    str_args = []
    materials = []
    variables = []
    for ii, arg_kind in enumerate(arg_kinds):
        if arg_kind != 'ts':
            if ats_mode is not None:
                extended_ats = arg_types[ii] + ('/%s' % ats_mode)

            else:
                extended_ats = arg_types[ii]

            try:
                sh = arg_shapes[arg_types[ii]]

            except KeyError:
                sh = arg_shapes[extended_ats]

        if arg_kind.endswith('variable'):
            shape = _parse_scalar_shape(sh[0] if isinstance(sh, tuple) else sh)
            field = Field.from_args('f%d' % ii, nm.float64, shape, omega,
                                    approx_order=1,
                                    poly_space_base=poly_space_base)

            if arg_kind == 'virtual_variable':
                if sh[1] is not None:
                    istate = arg_types.index(sh[1])

                else:
                    # Only virtual variable in arguments.
                    istate = -1
                    # -> Make fake variable.
                    var = FieldVariable('u-1', 'unknown', field)
                    var.set_constant(0.0)
                    variables.append(var)

                var = FieldVariable('v', 'test', field,
                                    primary_var_name='u%d' % istate)

            elif arg_kind == 'state_variable':
                var = FieldVariable('u%d' % ii, 'unknown', field)
                var.set_constant(0.0)

            elif arg_kind == 'parameter_variable':
                var = FieldVariable('p%d' % ii, 'parameter', field,
                                    primary_var_name='(set-to-None)')
                var.set_constant(0.0)

            variables.append(var)
            str_args.append(var.name)
            args[var.name] = var

        elif arg_kind.endswith('material'):
            if sh is None: # Switched-off opt_material.
                continue

            prefix = ''
            if isinstance(sh, basestr):
                aux = sh.split(':')
                if len(aux) == 2:
                    prefix, sh = aux

            if material_value is None:
                material_value = 1.0

            shape = _parse_tuple_shape(sh)
            if (len(shape) > 1) or (shape[0] > 1):
                if ((len(shape) == 2) and (shape[0] ==  shape[1])
                    and (material_value != 0.0)):
                    # Identity matrix.
                    val = nm.eye(shape[0], dtype=nm.float64)

                else:
                    # Array.
                    val = nm.empty(shape, dtype=nm.float64)
                    val.fill(material_value)

                values = {'%sc%d' % (prefix, ii)
                          : val}

            elif (len(shape) == 1) and (shape[0] == 1):
                # Single scalar as a special value.
                values = {'.c%d' % ii : material_value}

            else:
                raise ValueError('wrong material shape! (%s)' % shape)

            mat = Material('m%d' % ii, values=values)

            materials.append(mat)
            str_args.append(mat.name + '.' + 'c%d' % ii)
            args[mat.name] = mat

        elif arg_kind == 'ts':
            ts = TimeStepper(0.0, 1.0, 1.0, 5)
            str_args.append('ts')
            args['ts'] = ts

        else:
            str_args.append('user%d' % ii)
            args[str_args[-1]] = None

    materials = Materials(materials)
    variables = Variables(variables)

    return args, str_args, materials, variables

class Test(TestCommon):

    @staticmethod
    def from_conf(conf, options):
        from sfepy.discrete import Integral
        from sfepy.discrete.fem import Mesh, FEDomain

        domains = []
        for filename in filename_meshes:
            mesh = Mesh.from_file(filename)
            domain = FEDomain('domain_%s' % mesh.name.replace(data_dir, ''),
                              mesh)
            domain.create_region('Omega', 'all')
            domain.create_region('Gamma', 'vertices of surface', 'facet')

            domains.append(domain)

        integral = 
Integral('i', order=3)