import functools
import os
import tempfile
from typing import List, Tuple, Union

import numpy as np
import sympy as sp
from numpy.typing import NDArray
from sympy.codegen.rewriting import create_expand_pow_optimization
from sympy.printing.lambdarepr import lambdarepr


@functools.lru_cache
def _deco_basic(num_free_symbol: int, parallel: bool, fastmath: bool) -> str:
    """Decorator for basic functions."""
    tail = ''
    # maybe numba bugs, do not compile if set
    # if parallel:
    #     tail += ', target="parallel"'
    if fastmath:
        tail += ', fastmath=True'
    return '@numba.vectorize("float64({})"{})\n'.format(', '.join(['float64'] * num_free_symbol), tail)


@functools.lru_cache
def _deco_F(parallel: bool, fastmath: bool) -> str:
    """Decorator for aggregate function F (value)."""
    tail = ''
    if parallel:
        tail += ', parallel=True'
    if fastmath:
        tail += ', fastmath=True'
    return '@numba.njit("float64[:](float64[:], int32)"{})\n'.format(tail)


@functools.lru_cache
def _deco_GH(parallel: bool, fastmath: bool) -> str:
    """Decorator for aggregate functions G (gradient) & H (hessian)."""
    tail = ''
    if parallel:
        tail += ', parallel=True'
    if fastmath:
        tail += ', fastmath=True'
    return '@numba.njit("float64[:, :](float64[:], int32)"{})\n'.format(tail)


class FastFunc:
    """FastFunc take a sympy expression and a list of sympy symbols as function arguments,
    and generate JITed, vectorized functions for value, gradient and hessian.
    Gradient and hessian are generated by automatic differentiation and in sparse format,
    i.e., only non-zero elements are computed. The index of non-zero elements are stored in
    G_index and H_index, respectively.

    Both F, G, and H take a second argument for length of a single variable, the allows the
    user change the discretization scheme without recompiling the functions.
    """

    def __init__(self, function: Union[int, float, sp.Expr], args: List[sp.Symbol], /,
                 simplify: bool = False, parallel: bool = False, fastmath: bool = False) -> None:
        self._function = sp.sympify(function)
        self._args = args

        # random valid python name to avoid conflict
        self._valid_args = [sp.Symbol('xbN4dRhrnN_{}'.format(i)) for i in range(len(args))]
        for i in range(len(args)):
            self._function = self._function.subs(args[i], self._valid_args[i])
        self._args = self._valid_args

        self._simplify = simplify
        self._parallel = parallel
        self._fastmath = fastmath

        self._value_index = []
        self._grad_index = []
        self._hess_index = []

        self._value_consts = None
        self._grad_consts = {}
        self._hess_consts = {}

        self.__expand_opt = create_expand_pow_optimization(3)

        basic_str = self._gen_basic_funcs()
        aggregate_str = self._gen_aggregate_funcs()

        self._compile(basic_str + aggregate_str)
        self.G_index = self._gen_G_index()
        self.H_index = self._gen_H_index()

    def _compile(self, code: str) -> None:
        tmp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False)

        tmp_file.write('from numpy import *\n')
        tmp_file.write('import numba\n')

        tmp_file.write(code)
        tmp_file.close()

        import importlib.util
        import sys

        spec = importlib.util.spec_from_file_location('tmp_module', tmp_file.name)
        tmp_module = importlib.util.module_from_spec(spec)
        sys.modules[spec.name] = tmp_module
        spec.loader.exec_module(tmp_module)

        self.F = tmp_module.F
        self.G = tmp_module.G
        self.H = tmp_module.H

        os.unlink(tmp_file.name)

    def _free_symbols(self, expr: sp.Expr) -> List[int]:
        return sorted([self._args.index(sym) for sym in expr.free_symbols])

    def _gen_basic_funcs(self) -> str:
        result_str = ''
        func = self._function

        # value
        if self._simplify:
            func = sp.simplify(func)
        if func.is_constant():
            func = sp.sympify(func)
            self._value_consts = float(func)
            return ''

        free_symbols_value = self._free_symbols(func)
        self._value_index = free_symbols_value
        deco_str = _deco_basic(len(free_symbols_value), self._parallel, self._fastmath)
        func_str = lambdarepr(self.__expand_opt(func)).replace('math.', '')
        param_str = ', '.join([self._args[i].name for i in free_symbols_value])
        result_str += deco_str + 'def f({}):\n'.format(param_str) + '    return ' + func_str + '\n\n'

        for j in free_symbols_value:  # row in hess
            # gradient
            grad = sp.diff(func, self._args[j])
            if self._simplify:
                grad = sp.simplify(grad)
            if grad == 0:
                continue
            if grad.is_constant():
                grad = sp.simplify(grad)
                self._grad_consts[j] = float(grad)
                self._grad_index.append((j, []))
                continue

            free_symbols_grad = self._free_symbols(grad)
            self._grad_index.append((j, free_symbols_grad))
            deco_str = _deco_basic(len(free_symbols_grad), self._parallel, self._fastmath)
            grad_str = lambdarepr(self.__expand_opt(grad)).replace('math.', '')
            param_str = ', '.join([self._args[k].name for k in free_symbols_grad])
            result_str += deco_str + 'def df_d{}({}):\n'.format(j, param_str) \
                          + '    return ' + grad_str + '\n\n'

            # hessian
            for k in free_symbols_value:  # col
                if j < k:
                    break
                hess = sp.diff(grad, self._args[k])
                if self._simplify:
                    hess = sp.simplify(hess)
                if hess == 0:
                    continue
                if hess.is_constant():
                    hess = sp.simplify(hess)
                    self._hess_consts[(j, k)] = float(hess)
                    self._hess_index.append((j, k, []))
                    continue

                free_symbols_hess = self._free_symbols(hess)
                self._hess_index.append((j, k, free_symbols_hess))
                deco_str = _deco_basic(len(free_symbols_hess), self._parallel, self._fastmath)
                hess_str = lambdarepr(self.__expand_opt(hess)).replace('math.', '')
                param_str = ', '.join([self._args[l].name for l in free_symbols_hess])
                result_str += deco_str + 'def d2f_d{}d{}({}):\n'.format(j, k, param_str) \
                              + '    return ' + hess_str + '\n\n'
        return result_str

    def _gen_aggregate_funcs(self) -> str:
        result_str = ''
        # value
        result_str += _deco_F(self._parallel, self._fastmath)
        result_str += 'def F(x, n):\n'
        result_str += '    r = zeros(n, dtype=float64)\n'
        if self._value_consts is not None:
            result_str += f'    r[:] = {self._value_consts}\n'
        else:
            param_str = ', '.join([f'x[{j} * n:{j + 1} * n]' for j in self._value_index])
            result_str += '    r = f({})\n'.format(param_str)
        result_str += '    return r\n\n'

        # gradient
        result_str += _deco_GH(self._parallel, self._fastmath)
        result_str += 'def G(x, n):\n'
        result_str += f'    r = zeros(({len(self._grad_index)}, n), dtype=float64)\n'
        for row, (x, fs) in enumerate(self._grad_index):  # fs for free symbols
            if x in self._grad_consts:
                result_str += f'    r[{row}] = {self._grad_consts[x]}\n'
            else:
                param_str = ', '.join([f'x[{j} * n:{j + 1} * n]' for j in fs])
                result_str += '    r[{}] = df_d{}({})\n'.format(row, x, param_str)
        result_str += '    return r\n\n'

        # hessian
        result_str += _deco_GH(self._parallel, self._fastmath)
        result_str += 'def H(x, n):\n'
        result_str += f'    r = zeros(({len(self._hess_index)}, n), dtype=float64)\n'
        for row, (x, y, fs) in enumerate(self._hess_index):
            if (x, y) in self._hess_consts:
                result_str += f'    r[{row}] = {self._hess_consts[(x, y)]}\n'
            else:
                param_str = ', '.join([f'x[{j} * n:{j + 1} * n]' for j in fs])
                result_str += '    r[{}] = d2f_d{}d{}({})\n'.format(row, x, y, param_str)
        result_str += '    return r\n\n'

        return result_str

    def _gen_G_index(self) -> NDArray[np.int32]:
        return np.array([x for x, _ in self._grad_index], dtype=np.int32)

    def _gen_H_index(self) -> Tuple[NDArray[np.int32], NDArray[np.int32]]:
        return np.array([x for x, _, _ in self._hess_index], dtype=np.int32), \
               np.array([y for _, y, _ in self._hess_index], dtype=np.int32)
