# Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================


import brainstate
import brainunit as u
import jax
import jax.numpy as jnp
import numpy as np
import pytest
import scipy.sparse as sp

import brainevent
from brainevent._csr.plasticity_binary import (
    update_csr_on_binary_pre,
    update_csr_on_binary_post,
    update_csr_on_binary_pre_p,
    update_csr_on_binary_post_p,
)

PLATFORM = jax.default_backend()
PRE_BACKENDS = tuple(update_csr_on_binary_pre_p.available_backends(PLATFORM))
POST_BACKENDS = tuple(update_csr_on_binary_post_p.available_backends(PLATFORM))

shapes = [
    (20, 100),
    (100, 100),
    (100, 50),
]


class Test_csr_on_pre:
    @pytest.mark.parametrize('backend', PRE_BACKENDS)
    @pytest.mark.parametrize('shape', shapes)
    def test_csr_on_pre_v1(self, backend, shape):
        n_pre, n_post = shape
        mat = brainstate.random.random((n_pre, n_post))
        mask = mat < 0.5
        mat = jnp.where(mask, mat, 0.)

        pre_spike = brainstate.random.random((n_pre,)) < 0.5
        post_trace = brainstate.random.random((n_post,))

        csr = brainevent.CSR.fromdense(mat)
        csr2 = csr.with_data(
            update_csr_on_binary_pre(csr.data, csr.indices, csr.indptr, pre_spike, post_trace,
                                     shape=csr.shape, backend=backend)
        )
        dense2 = jnp.where(mask, mat + jnp.outer(pre_spike.astype(float), post_trace), 0.)

        assert jnp.allclose(csr2.todense(), dense2)

        jax.block_until_ready((mat, pre_spike, post_trace))

    @pytest.mark.parametrize('backend', PRE_BACKENDS)
    @pytest.mark.parametrize('shape', shapes)
    @pytest.mark.parametrize('mat_unit', [u.mV, u.ms])
    @pytest.mark.parametrize('trace_unit', [u.mV, u.ms])
    def test_csr_on_pre_with_unit(self, backend, shape, mat_unit, trace_unit):
        def run():
            n_pre, n_post = shape
            mat = brainstate.random.random((n_pre, n_post))
            mask = mat < 0.5
            mat = jnp.where(mask, mat, 0.) * mat_unit
            pre_spike = brainstate.random.random((n_pre,)) < 0.1
            post_trace = brainstate.random.random((n_post,)) * trace_unit

            csr = brainevent.CSR.fromdense(mat)
            csr = csr.with_data(
                update_csr_on_binary_pre(csr.data, csr.indices, csr.indptr, pre_spike, post_trace,
                                         shape=csr.shape, backend=backend)
            )

            dense = mat + u.math.outer(pre_spike.astype(float), post_trace)
            dense = u.math.where(mask, dense, 0. * mat_unit)

            assert u.math.allclose(csr.todense(), dense)

            jax.block_until_ready((mat, pre_spike, post_trace))

        if mat_unit.has_same_dim(trace_unit):
            run()
        else:
            with pytest.raises(u.UnitMismatchError):
                run()

    @pytest.mark.parametrize('backend', PRE_BACKENDS)
    @pytest.mark.parametrize('shape', shapes)
    @pytest.mark.parametrize('w_in', [None, 0.1])
    @pytest.mark.parametrize('w_max', [None, 0.5])
    def test_csr_on_pre_v2(self, backend, shape, w_in, w_max):
        n_pre, n_post = shape
        mat = brainstate.random.random((n_pre, n_post))
        mask = mat < 0.5
        mat = jnp.where(mask, mat, 0.)
        pre_spike = brainstate.random.random((n_pre,)) < 0.1
        post_trace = brainstate.random.random((n_post,))

        csr = brainevent.CSR.fromdense(mat)
        csr = csr.with_data(
            update_csr_on_binary_pre(csr.data, csr.indices, csr.indptr, pre_spike, post_trace,
                                     w_min=w_in, w_max=w_max, shape=csr.shape, backend=backend)
        )

        mat = mat + jnp.outer(pre_spike.astype(float), post_trace)
        mat = u.math.clip(mat, a_min=w_in, a_max=w_max)

        mat = jnp.where(mask, mat, 0.)
        assert jnp.allclose(csr.todense(), mat)

        jax.block_until_ready((mat, pre_spike, post_trace))


def _csr_to_csc_with_weight_indices(csr_data, csr_indices, csr_indptr, shape):
    """Convert CSR format to CSC format and return weight indices mapping.

    Returns:
        csc_indices: Row indices in CSC format (presynaptic neuron indices)
        csc_indptr: Column pointers in CSC format
        weight_indices: Mapping from CSC positions to original CSR data positions
    """
    # Handle Quantity data by extracting mantissa
    if hasattr(csr_data, 'mantissa'):
        csr_data_np = np.asarray(csr_data.mantissa)
    else:
        csr_data_np = np.asarray(csr_data)

    # Create scipy CSR matrix
    scipy_csr = sp.csr_matrix(
        (csr_data_np, np.asarray(csr_indices), np.asarray(csr_indptr)),
        shape=shape
    )
    # Convert to CSC
    scipy_csc = scipy_csr.tocsc()

    # The weight_indices maps CSC position -> CSR position
    # We can compute this by tracking where each element came from
    # CSC data comes from reordering CSR data based on column sorting

    # Create a CSR matrix with data = position indices
    position_indices = np.arange(len(csr_data_np), dtype=np.int32)
    scipy_csr_indices = sp.csr_matrix(
        (position_indices, np.asarray(csr_indices), np.asarray(csr_indptr)),
        shape=shape
    )
    scipy_csc_indices = scipy_csr_indices.tocsc()
    weight_indices = scipy_csc_indices.data

    return (
        jnp.asarray(scipy_csc.indices, dtype=jnp.int32),
        jnp.asarray(scipy_csc.indptr, dtype=jnp.int32),
        jnp.asarray(weight_indices, dtype=jnp.int32)
    )


class Test_on_post:
    @pytest.mark.parametrize('backend', POST_BACKENDS)
    @pytest.mark.parametrize('shape', shapes)
    def test_csr_on_post_v1(self, backend, shape):
        n_pre, n_post = shape
        mat = brainstate.random.random((n_pre, n_post))
        mask = mat < 0.5
        mat = jnp.where(mask, mat, 0.)

        post_spike = brainstate.random.random((n_post,)) < 0.5
        pre_trace = brainstate.random.random((n_pre,))

        csr = brainevent.CSR.fromdense(mat)
        # Convert CSR to CSC format for csr2csc_on_post
        csc_indices, csc_indptr, weight_indices = _csr_to_csc_with_weight_indices(
            csr.data, csr.indices, csr.indptr, csr.shape
        )

        new_weights = update_csr_on_binary_post(
            csr.data, csc_indices, csc_indptr, weight_indices,
            pre_trace, post_spike, shape=csr.shape, backend=backend
        )
        csr2 = csr.with_data(new_weights)
        dense2 = jnp.where(mask, mat + jnp.outer(pre_trace, post_spike.astype(float)), 0.)

        assert jnp.allclose(csr2.todense(), dense2)

        jax.block_until_ready((mat, post_spike, pre_trace))

    @pytest.mark.parametrize('backend', POST_BACKENDS)
    @pytest.mark.parametrize('shape', shapes)
    @pytest.mark.parametrize('mat_unit', [u.mV, u.ms])
    @pytest.mark.parametrize('trace_unit', [u.mV, u.ms])
    def test_csr_on_post_with_unit(self, backend, shape, mat_unit, trace_unit):
        def run():
            n_pre, n_post = shape
            mat = brainstate.random.random((n_pre, n_post))
            mask = mat < 0.5
            mat = jnp.where(mask, mat, 0.) * mat_unit
            post_spike = brainstate.random.random((n_post,)) < 0.1
            pre_trace = brainstate.random.random((n_pre,)) * trace_unit

            csr = brainevent.CSR.fromdense(mat)
            # Convert CSR to CSC format for csr2csc_on_post
            csc_indices, csc_indptr, weight_indices = _csr_to_csc_with_weight_indices(
                csr.data, csr.indices, csr.indptr, csr.shape
            )

            new_weights = update_csr_on_binary_post(
                csr.data, csc_indices, csc_indptr, weight_indices,
                pre_trace, post_spike, shape=csr.shape, backend=backend
            )
            csr = csr.with_data(new_weights)

            dense = mat + u.math.outer(pre_trace, post_spike.astype(float))
            dense = u.math.where(mask, dense, 0. * mat_unit)

            assert u.math.allclose(csr.todense(), dense)

            jax.block_until_ready((mat, post_spike, pre_trace))

        if mat_unit.has_same_dim(trace_unit):
            run()
        else:
            with pytest.raises(u.UnitMismatchError):
                run()

    @pytest.mark.parametrize('backend', POST_BACKENDS)
    @pytest.mark.parametrize('shape', shapes)
    @pytest.mark.parametrize('w_in', [None, 0.1])
    @pytest.mark.parametrize('w_max', [None, 0.5])
    def test_csr_on_post_v2(self, backend, shape, w_in, w_max):
        n_pre, n_post = shape
        mat = brainstate.random.random((n_pre, n_post))
        mask = mat < 0.5
        mat = jnp.where(mask, mat, 0.)
        post_spike = brainstate.random.random((n_post,)) < 0.1
        pre_trace = brainstate.random.random((n_pre,))

        csr = brainevent.CSR.fromdense(mat)
        # Convert CSR to CSC format for csr2csc_on_post
        csc_indices, csc_indptr, weight_indices = _csr_to_csc_with_weight_indices(
            csr.data, csr.indices, csr.indptr, csr.shape
        )

        new_weights = update_csr_on_binary_post(
            csr.data, csc_indices, csc_indptr, weight_indices,
            pre_trace, post_spike, w_min=w_in, w_max=w_max, shape=csr.shape, backend=backend
        )
        csr = csr.with_data(new_weights)

        mat = mat + jnp.outer(pre_trace, post_spike.astype(float))
        mat = u.math.clip(mat, a_min=w_in, a_max=w_max)

        mat = jnp.where(mask, mat, 0.)
        assert jnp.allclose(csr.todense(), mat)

        jax.block_until_ready((mat, post_spike, pre_trace))
