#   nengo_bio -- Extensions to Nengo for more biological plausibility
#   Copyright (C) 2019  Andreas Stöckel
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <https://www.gnu.org/licenses/>.

from nengo_bio.internal.multi_compartment_lif_sim import make_simulator_class
from nengo_bio.internal.multi_compartment_lif_parameters import (
    SomaticParameters,
    DendriticParameters,
)

_tmp_dir = None  # Directory compiled libraries are stored in
_compiled_library_map = {}  # Map storing the Simulator objects
_supports_cpp = True


def _generate_simulator_cpp_code(f, params_som, params_den, dt=1e-3, ss=10):
    """
    Generates the C++ code representing a specific neuron simulator and stores
    it in the given file descriptor f. This function is used internally by
    "compile_simulator_cpp".
    """

    def _fmt(x):
        res = "{:g}".format(x)
        if (not "." in res):
            s = res.split('e')
            s[0] += '.0'  # Ensure number is C++ double, just '.' would work, but is ugly
            return 'e'.join(s)
        return res

    # Code below compares the result of _fmt to the string "0.0"
    assert _fmt(0) == "0.0" and _fmt(0.0) == "0.0"

    def _generate_b_vec_entry(i, B, b_const):
        res = _fmt(b_const[i])
        if res == "0.0":
            res = ""
        for k in range(B.shape[1]):
            if B[i, k] != 0.0:
                if len(res) > 0:
                    res += " + "
                res += "({} * x[{}])".format(_fmt(B[i, k]), k)
        if res == "":
            return "0.0"
        return res

    def _generate_A_mat_entry(i, j, A, a_const, C):
        # Lookup off-diagonal entries in the C matrix
        if i != j:
            return _fmt(C[i, j])

        # Assemble diagonal entries according to the Aconst and A matrix
        return _generate_b_vec_entry(i, A, a_const)

    # Handy aliases for params_som and params_den
    pS, pD = params_som, params_den

    # Compute v_min, v_max
    v_min, v_max = params_den.vEq_extreme(params_som)

    f.write(
        """// This code was automatically generated by multi_compartment_lif_cpp.py

#include "multi_compartment_lif.hpp"

namespace { // Do not export the following symbols
struct Parameters {
    static constexpr double tau_ref = """ + _fmt(pS.tau_ref) + """;
    static constexpr double tau_spike = """ + _fmt(pS.tau_spike) + """;
    static constexpr double v_th = """ + _fmt(pS.v_th) + """;
    static constexpr double v_reset = """ + _fmt(pS.v_reset) + """;
    static constexpr double v_spike = """ + _fmt(pS.v_spike) + """;
    static constexpr size_t n_comp = """ + str(pD.n_comp) + """;
    static constexpr size_t n_inputs = """ + str(pD.n_inputs) + """;
    static constexpr double dt = """ + _fmt(dt) + """;
    static constexpr size_t ss = """ + str(ss) + """;

    using MatA = Matrix<double, n_comp, n_comp>;
    using VecASom = Matrix<double, 1, n_comp>;
    using VecB = Matrix<double, n_comp, 1>;
    using VecX = Matrix<double, n_inputs, 1>;
    using VecV = Matrix<double, n_comp, 1>;

    static const VecV v_min;
    static const VecV v_max;

    static MatA calc_A(const VecX &x) {
        (void)x; // x may be unused, suppress warnings
        return (MatA() <<\n""")
    for i in range(pD.n_comp):
        for j in range(pD.n_comp):
            f.write("            {}".format(
                _generate_A_mat_entry(i, j, pD.A, pD.a_const, pD.C)))
            if i != pD.n_comp - 1 or j != pD.n_comp - 1:
                f.write(",")
            f.write("\n")
    f.write("""        ).finished();
    }

    static VecB calc_b(const VecX &x) {
        (void)x; // x may be unused, suppress warnings
        return (VecB() <<\n""")
    for i in range(pD.n_comp):
        f.write("            {}".format(
            _generate_b_vec_entry(i, pD.B, pD.b_const)))
        if i != pD.n_comp - 1:
            f.write(",")
        f.write("\n")
    f.write("""        ).finished();
    }
};

const Parameters::VecV Parameters::v_min{""" + ", ".join(map(_fmt, v_min)) + """};
const Parameters::VecV Parameters::v_max{""" + ", ".join(map(_fmt, v_max)) + """};
}

extern "C" { // Exported C API
void run_step_from_memory(uint32_t n_neurons, double *state, double *out""")
    for i in range(pD.n_inputs):
        f.write(", double *x{}".format(i))
    f.write(""") {
    const double *xs[""" + str(pD.n_inputs) + """] = {""")
    f.write(", ".join("x{}".format(i) for i in range(pD.n_inputs)))
    f.write("""};
    Simulator<Parameters>::run_step_from_memory(n_neurons, state, out, xs);
}

void run_single_with_constant_input(uint32_t n_samples, double *state, double *out, double *xs) {
    Simulator<Parameters>::run_single_with_constant_input(n_samples, state, out, xs);
}

void run_single_with_poisson_sources(uint32_t n_samples, double *state, double *out, PoissonSource *sources) {
    Simulator<Parameters>::run_single_with_poisson_sources(n_samples, state, out, sources);
}

void run_single_with_gaussian_sources(uint32_t n_samples, double *state, double *out, GaussianSource *sources) {
    Simulator<Parameters>::run_single_with_gaussian_sources(n_samples, state, out, sources);
}

};
""")


def _compile_cpp_library(code, tar, debug=False):
    import os
    import subprocess
    import posix_ipc

    # Create a semaphore for the given target file. Only one process should try
    # to compile code at a time
    with posix_ipc.Semaphore(
            os.path.basename(tar), posix_ipc.O_CREAT, initial_value=1) as lock:
        # Someone else compiled the code for us, return!
        if os.path.exists(tar):
            return

        # Fetch the module directory which is used as include directory
        curdir = os.path.dirname(os.path.abspath(__file__))

        # Compile the code
        args = list(
            filter(
                lambda x: bool(x),
                [
                    'g++',
                    '-std=c++11',
                    '-march=native',  # Compile for this particular computer
                    '-g' if debug else '-DNDEBUG',  # Toggle debug code
                    '-O1' if debug else '-O3',  # Toggle optimisation
                    '-Wall',
                    '-Wextra',
                    '-Wno-deprecated-copy',
                    '-fno-exceptions',  # This code does not use exceptions
                    '-fno-rtti',  # We don't need no runtime-type information
                    '-ffast-math',  # Be sloppy, we don't have non-finite math
                    '-I' + os.path.expanduser('~/.local/include/eigen3'),
                    '-I/usr/include/eigen3',
                    '-I' + curdir,
                    '-x',  # Specify the language since we're compiling from stdin
                    'c++',
                    '-fPIC',
                    '-shared',
                    None if debug else '-s',
                    '-Wl,--as-needed,-soname,' + tar,
                    '-lm',
                    '-o',
                    tar
                ]))

        # Write the source code to a temporary file when debugging
        if debug:
            src = tar + '.cpp'
            with open(src, 'wb') as f:
                f.write(code)
            args.append(src)
        else:
            args.append('-')

        # Call the compiler
        process = subprocess.Popen(
            args,
            stdin=None if debug else subprocess.PIPE,
            stderr=subprocess.PIPE)
        _, stderr = process.communicate(None if debug else code)
        stderr = str(stderr, 'utf-8')
        if process.returncode != 0:
            raise Exception("Error while compiling the C++ code:\n" + stderr)
        elif len(stderr) > 0:
            print(stderr)


def _create_tmp_dir():
    import os
    import tempfile

    global _tmp_dir, _compiled_library_map

    # Create a user-accessible temporary directory
    if _tmp_dir is None or not os.path.exists(_tmp_dir):
        _compiled_library_map = {}
        _tmp_dir = os.path.expanduser('~/.cache/multi_compartment_lif/')
        os.makedirs(_tmp_dir, exist_ok=True, mode=0O755)

    return _tmp_dir


def _file_hash(filename, hasher):
    """
    Feeds the content of the file with the specified file into the hasher.
    """

    with open(filename, 'rb') as f:
        while True:
            data = f.read(65536)
            if not data:
                break
            hasher.update(data)


def _compute_code_hash(code):
    import hashlib
    import os

    # Find the header file. Changing this file should also result in a new
    # library being compiled
    header_file = os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        "multi_compartment_lif.hpp")

    hasher = hashlib.sha1()
    _file_hash(header_file, hasher)
    _file_hash(__file__, hasher)
    hasher.update(str(os.getuid()).encode('utf-8'))
    hasher.update(code)
    return hasher.hexdigest()


def compile_simulator_cpp(params_som, params_den, dt=1e-3, ss=10):
    """
    Generates and compiles the C++ simulator library. Returns a callable that
    can be used to  To compiled dynamic library will be stored in a temporary
    directory accessible only to the user running this Python process.
    """
    from ctypes import cdll, POINTER, cast, c_double, c_uint32, c_void_p
    import io
    import numpy as np
    import os

    # Generate the C++ code
    f = io.StringIO()
    _generate_simulator_cpp_code(f, params_som, params_den, dt, ss)

    # Compute the hash, generate the lookup table key -- if this neuron model
    # has already been compiled, return the corresponding compiled code
    code = f.getvalue().encode('utf-8')
    sha_hash = _compute_code_hash(code)
    key = "multi_comp_lif_" + sha_hash[:8]
    if not key in _compiled_library_map:
        # Compile the C++ code to a library
        libpath = os.path.join(_create_tmp_dir(), key + '.so')
        if not os.path.exists(libpath):
            _compile_cpp_library(code, libpath)

        # Load the C library
        lib = cdll.LoadLibrary(libpath)
        c_run_step_from_memory = lib.run_step_from_memory
        c_run_single_with_constant_input = lib.run_single_with_constant_input
        c_run_single_with_poisson_sources = lib.run_single_with_poisson_sources
        c_run_single_with_gaussian_source = lib.run_single_with_gaussian_sources

        c_double_p = POINTER(c_double)

        class CppImpl:
            @staticmethod
            def run_step_from_memory(self, out, *xs):
                pstate = self.state.ctypes.data_as(c_double_p)
                pout = out.ctypes.data_as(c_double_p)
                pxs = [x.ctypes.data_as(c_double_p) for x in xs]
                c_run_step_from_memory(
                    c_uint32(self.n_neurons), pstate, pout, *pxs)

            @staticmethod
            def run_single_with_constant_input(self, out, xs):
                pstate = self.state.ctypes.data_as(c_double_p)
                pout = out.ctypes.data_as(c_double_p)
                pxs = xs.ctypes.data_as(c_double_p)
                c_run_single_with_constant_input(
                    c_uint32(out.shape[0]), pstate, pout, pxs)

            @staticmethod
            def run_single_with_poisson_sources(self, out, sources):
                pstate = self.state.ctypes.data_as(c_double_p)
                pout = out.ctypes.data_as(c_double_p)
                psources = cast(sources, c_void_p)
                c_run_single_with_poisson_sources(out.size, pstate, pout,
                                                  psources)

            @staticmethod
            def run_single_with_gaussian_sources(self, out, sources):
                pstate = self.state.ctypes.data_as(c_double_p)
                pout = out.ctypes.data_as(c_double_p)
                psources = cast(sources, c_void_p)
                c_run_single_with_gaussian_source(out.size, pstate, pout,
                                                  psources)

        # Register the above function for the specific neuron type and return it
        _compiled_library_map[key] = make_simulator_class(
            CppImpl, params_som, params_den, dt, ss)
    return _compiled_library_map[key]


def supports_cpp():
    global _supports_cpp
    if _supports_cpp is None:
        import os
        import warnings

        # Create a user-accessible temporary directory
        tar = os.path.join(_create_tmp_dir(), 'test.so')

        # Try to compile a minimal library which just includes eigen
        try:
            compile_cpp_library(
                b"\"multi_compartment_lif.hpp\"\n#include <Eigen/Dense>\n#include <Eigen/Eigenvalues>",
                tar)
            _supports_cpp = True
        except:
            _supports_cpp = False
            warnings.warn(
                "Cannot activate C++ support for multi-compartment LIF "
                "neurons. The simulator will fallback to a Python "
                "implementation. This will be slow. Please make sure you have "
                "a recent version of g++ and Eigen3 installed.",
                RuntimeWarning)

        # Remove the target file
        if os.path.exists(tar):
            os.remove(tar)

    return _supports_cpp

