from typing import Iterable, Callable
from types import NoneType

import numpy as np
from numba import njit, prange
from numpy import ndarray

from sigmaepsilon.math.function import Function

from .config import Config


@njit(nogil=True, parallel=True, cache=True)
def _monte_carlo_1d_njit(
    size: float,
    shape: int,
    points: ndarray,
    fvalues: ndarray,
    dsize: float,
    num_MC_tot: int,
    out: ndarray,
) -> ndarray:
    N = shape
    L = size
    multi = (2 / L) * dsize
    factor = np.pi / L

    for n in prange(1, N + 1):
        iN = n - 1
        out[iN, 0] += (
            multi * np.sum(np.sin(points * factor * n) * fvalues[:, 0]) / num_MC_tot
        )
        out[iN, 1] += (
            multi * np.sum(np.cos(points * factor * n) * fvalues[:, 1]) / num_MC_tot
        )

    return out


def _monte_carlo_1d(
    size: float,
    shape: int,
    values: Iterable,
    dsize: float,
    *,
    rpg: Callable[[int], ndarray],
    out: ndarray | NoneType,
    n_MC: int | NoneType = None,
) -> ndarray:
    functions = []
    for i in range(2):
        if isinstance(values[i], str):
            function = Function(values[i], variables=["x"])
        elif isinstance(values[i], (float, int)):

            def function(points) -> ndarray:
                return np.full((len(points),), values[i], dtype=float)

        else:  # pragma: no cover
            raise ValueError(f"Invalid value {values[i]}")
        functions.append(function)

    num_MC_samples = n_MC or Config.NUM_MC_SAMPLES_BEAM
    MC_batch_size = Config.MC_BATCH_SIZE_BEAM
    num_MC_bacthes = int(num_MC_samples // MC_batch_size)
    remaining_MC_points = int(num_MC_samples % MC_batch_size)
    f_vals = np.zeros((MC_batch_size, 2), dtype=float)
    points = np.zeros((MC_batch_size), dtype=float)

    for _ in range(num_MC_bacthes):
        rpg(MC_batch_size, out=points)
        for i in range(2):
            f_vals[:, i] = functions[i]([points])
        _monte_carlo_1d_njit(size, shape, points, f_vals, dsize, num_MC_samples, out)

    MC_batch_size = remaining_MC_points
    if remaining_MC_points > 1:
        points = points[:MC_batch_size]
        f_vals = f_vals[:MC_batch_size]
        rpg(MC_batch_size, out=points)
        for i in range(2):
            f_vals[:, i] = functions[i]([points])
        _monte_carlo_1d_njit(size, shape, points, f_vals, dsize, num_MC_samples, out)

    return out


@njit(nogil=True, parallel=True, cache=True)
def _monte_carlo_2d_njit(
    size: tuple,
    shape: tuple,
    points: ndarray,
    fvalues: ndarray,
    dsize: float,
    num_MC_tot: int,
    out: ndarray,
) -> ndarray:
    M, N = shape
    Lx, Ly = size
    multi = (4 / Lx / Ly) * dsize
    px = points[:, 0]
    py = points[:, 1]
    factor_x = np.pi / Lx
    factor_y = np.pi / Ly

    for m in prange(1, M + 1):
        for n in prange(1, N + 1):
            mn = (m - 1) * N + n - 1
            out[mn, 0] += (
                multi
                * np.sum(
                    np.sin(px * factor_x * m)
                    * np.sin(py * factor_y * n)
                    * fvalues[:, 0]
                )
                / num_MC_tot
            )
            out[mn, 1] += (
                multi
                * np.sum(
                    np.sin(px * factor_x * m)
                    * np.cos(py * factor_y * n)
                    * fvalues[:, 1]
                )
                / num_MC_tot
            )
            out[mn, 2] += (
                multi
                * np.sum(
                    np.cos(px * factor_x * m)
                    * np.sin(py * factor_y * n)
                    * fvalues[:, 2]
                )
                / num_MC_tot
            )

    return out


def _monte_carlo_2d(
    size: tuple,
    shape: tuple,
    values: Iterable,
    dsize: float,
    *,
    rpg: Callable[[int], ndarray],
    out: ndarray | NoneType,
    n_MC: int | NoneType = None,
) -> ndarray:
    functions = []
    for i in range(3):
        if isinstance(values[i], str):
            function = Function(values[i], variables=["x", "y"])
        elif isinstance(values[i], (float, int)):

            def function(points) -> ndarray:
                return np.full((points.shape[1],), values[i], dtype=float)

        else:  # pragma: no cover
            raise ValueError(f"Invalid value {values[i]}")
        functions.append(function)

    num_MC_samples = n_MC or Config.NUM_MC_SAMPLES_PLATE
    MC_batch_size = Config.MC_BATCH_SIZE_PLATE
    num_MC_bacthes = int(num_MC_samples // MC_batch_size)
    remaining_MC_points = int(num_MC_samples % MC_batch_size)
    f_vals = np.zeros((MC_batch_size, 3), dtype=float)
    points = np.zeros((MC_batch_size, 2), dtype=float)

    for _ in range(num_MC_bacthes):
        rpg(MC_batch_size, out=points)
        for i in range(3):
            f_vals[:, i] = functions[i](points.T)
        _monte_carlo_2d_njit(size, shape, points, f_vals, dsize, num_MC_samples, out)

    MC_batch_size = remaining_MC_points
    if remaining_MC_points > 1:
        points = points[:MC_batch_size]
        f_vals = f_vals[:MC_batch_size]
        rpg(MC_batch_size, out=points)
        for i in range(3):
            f_vals[:, i] = functions[i](points.T)
        _monte_carlo_2d_njit(size, shape, points, f_vals, dsize, num_MC_samples, out)

    return out
