# -*- coding: utf-8 -*-

from abc import ABCMeta, abstractmethod
import re

import numpy as np


class Arg(object):
    def __init__(self, name, spec, body):
        self.name = name

        specptn = (r'(?:(in|inout|out)\s+)?'                # Intent
                   r'(?:(broadcast|mpi|scalar|view)\s+)?'   # Attrs
                   r'([A-Za-z_]\w*)'                        # Data type
                   r'((?:\[\d+\]){0,2})$')                  # Dimensions
        dimsptn = r'(?<=\[)\d+(?=\])'
        usedptn = r'(?:[^A-Za-z]|^){0}[^A-Za-z0-9]'.format(name)

        # Parse our specification
        m = re.match(specptn, spec)
        if not m:
            raise ValueError('Invalid argument specification')

        g = m.groups()

        # Properties
        self.intent = g[0] or 'in'
        self.attrs = g[1] or ''
        self.dtype = g[2]

        # Dimensions
        self.cdims = [int(d) for d in re.findall(dimsptn, g[3])]
        self.ncdim = len(self.cdims)

        # Attributes
        self.isbroadcast = 'broadcast' in self.attrs
        self.ismpi = 'mpi' in self.attrs
        self.isused = bool(re.search(usedptn, body))
        self.isview = 'view' in self.attrs
        self.isscalar = 'scalar' in self.attrs
        self.isvector = 'scalar' not in self.attrs

        # Validation
        if self.isbroadcast and self.intent != 'in':
            raise ValueError('Broadcast arguments must be of intent in')
        if self.isbroadcast and self.ncdim != 0:
            raise ValueError('Broadcast arguments can not have dimensions')
        if self.isscalar and self.dtype != 'fpdtype_t':
            raise ValueError('Scalar arguments must be of type fpdtype_t')


class BaseKernelGenerator(object, metaclass=ABCMeta):
    def __init__(self, name, ndim, args, body, fpdtype):
        self.name = name
        self.ndim = ndim
        self.fpdtype = fpdtype

        # Parse and sort our argument list
        sargs = sorted((k, Arg(k, v, body)) for k, v in args.items())

        # Eliminate unused arguments
        sargs = [v for k, v in sargs if v.isused]

        # Break arguments into point-scalars and point-vectors
        self.scalargs = [v for v in sargs if v.isscalar]
        self.vectargs = [v for v in sargs if v.isvector]

        # If we are 1D ensure that none of our arguments are broadcasts
        if ndim == 1 and any(v.isbroadcast for v in self.vectargs):
            raise ValueError('Broadcast arguments are not supported in 1D '
                             'kernels')

        # If we are 2D ensure none of our arguments are views
        if ndim == 2 and any(v.isview for v in self.vectargs):
            raise ValueError('View arguments are not supported for 2D '
                             'kernels')

        # Similarly, check for MPI matrices
        if ndim == 2 and any(v.ismpi for v in self.vectargs):
            raise ValueError('MPI matrices are not supported for 2D kernels')

        # Render the main body of our kernel
        self.body = self._render_body(body)

        # Determine the dimensions to be iterated over
        self._dims = ['_nx'] if ndim == 1 else ['_ny', '_nx']

    def argspec(self):
        # Argument names and types
        argn, argt = [], []

        # Dimensions
        argn += self._dims
        argt += [[np.int32]]*self.ndim

        # Scalar args (always of type fpdtype)
        argn += [sa.name for sa in self.scalargs]
        argt += [[self.fpdtype]]*len(self.scalargs)

        # Vector args
        for va in self.vectargs:
            argn.append(va.name)

            # View
            if va.isview:
                argt.append([np.intp]*(2 + (va.ncdim == 2)))
            # Non-stacked vector or MPI type
            elif self.ndim == 1 and (va.ncdim == 0 or va.ismpi):
                argt.append([np.intp])
            # Broadcast vector
            elif va.isbroadcast:
                argt.append([np.intp])
            # Stacked vector/matrix/stacked matrix
            else:
                argt.append([np.intp, np.int32])

        # Return
        return self.ndim, argn, argt

    def needs_ldim(self, arg):
        return ((self.ndim == 2 and not arg.isbroadcast) or
                (arg.ncdim > 0 and not arg.ismpi))

    @abstractmethod
    def render(self):
        pass

    def _deref_arg_view(self, arg):
        ptns = ['{0}_v[{0}_vix[X_IDX]]',
                r'{0}_v[{0}_vix[X_IDX] + SOA_SZ*\1]',
                r'{0}_v[{0}_vix[X_IDX] + {0}_vrstri[X_IDX]*\1 + SOA_SZ*\2]']

        return ptns[arg.ncdim].format(arg.name)

    def _deref_arg_array_1d(self, arg):
        # Leading dimension
        ldim = 'ld' + arg.name if not arg.ismpi else '_nx'

        # Vector:
        #   name => name_v[X_IDX]
        if arg.ncdim == 0:
            ix = 'X_IDX'
        # Stacked vector:
        #   name[\1] => name_v[ldim*\1 + X_IDX]
        elif arg.ncdim == 1:
            ix = r'{0}*\1 + X_IDX'.format(ldim)
        # Doubly stacked MPI vector:
        #   name[\1][\2] => name_v[(nv*\1 + \2)*ldim + X_IDX]
        elif arg.ismpi:
            ix = r'({0}*\1 + \2)*{1} + X_IDX'.format(arg.cdims[1], ldim)
        # Doubly stacked vector:
        #   name[\1][\2] => name_v[ldim*\1 + X_IDX_AOSOA(\2, nv)]
        else:
            ix = (r'ld{0}*\1 + X_IDX_AOSOA(\2, {1})'
                   .format(arg.name, arg.cdims[1]))

        return '{0}_v[{1}]'.format(arg.name, ix)

    def _deref_arg_array_2d(self, arg):
        # Broadcast vector:
        #   name => name_v[X_IDX]
        if arg.isbroadcast:
            ix = 'X_IDX'
        # Matrix:
        #   name => name_v[ldim*_y + X_IDX]
        elif arg.ncdim == 0:
            ix = 'ld{0}*_y + X_IDX'.format(arg.name)
        # Stacked matrix:
        #   name[\1] => name_v[ldim*_y + X_IDX_AOSOA(\1, nv)]
        elif arg.ncdim == 1:
            ix = (r'ld{0}*_y + X_IDX_AOSOA(\1, {1})'
                   .format(arg.name, arg.cdims[0]))
        # Doubly stacked matrix:
        #   name[\1][\2] => name_v[(\1*ny + _y)*ldim + X_IDX_AOSOA(\2, nv)]
        else:
            ix = (r'(\1*_ny + _y)*ld{0} + X_IDX_AOSOA(\2, {1})'
                   .format(arg.name, arg.cdims[1]))

        return '{0}_v[{1}]'.format(arg.name, ix)

    def _render_body(self, body):
        ptns = [r'\b{0}\b', r'\b{0}\[(\d+)\]', r'\b{0}\[(\d+)\]\[(\d+)\]']

        # At single precision suffix all floating point constants by 'f'
        if self.fpdtype == np.float32:
            body = re.sub(r'(?=\d*[.eE])(?=\.?\d)\d*\.?\d*(?:[eE][+-]?\d+)?',
                          r'\g<0>f', body)

        # Dereference vector arguments
        for va in self.vectargs:
            if va.isview:
                darg = self._deref_arg_view(va)
            else:
                if self.ndim == 1:
                    darg = self._deref_arg_array_1d(va)
                else:
                    darg = self._deref_arg_array_2d(va)

            # Substitute
            body = re.sub(ptns[va.ncdim].format(va.name), darg, body)

        return body
