import numpy as np
from dataclasses import dataclass
from functools import cached_property
import functions as fn

Array = np.ndarray


class InvalidExpansion(Exception):
    pass


@dataclass
class Expansion:
    """Generic class for storing surface charge expansion coefficients."""

    l_array: Array
    coeffs: Array

    def __post_init__(self):
        """Validation of the given expansion."""
        if len(self.coeffs) != np.sum(2 * self.l_array + 1):
            raise InvalidExpansion('Number of expansion coefficients does not match the provided l_array.')
        if np.all(np.sort(self.l_array) != self.l_array) or np.all(np.unique(self.l_array) != self.l_array):
            raise InvalidExpansion('Array of l values should be unique and sorted.')

    @property
    def max_l(self) -> int:
        return max(self.l_array)

    @cached_property
    def lm_arrays(self) -> (Array, Array):
        """Return l and m arrays containing all (l, m) pairs."""
        all_m_list = []
        for l in self.l_array:
            for i in range(2 * l + 1):
                all_m_list.append(-l + i)
        return np.repeat(self.l_array, 2 * self.l_array + 1), np.array(all_m_list)

    def repeat_over_m(self, arr: Array, axis=0):
        if not arr.shape[axis] == len(self.l_array):
            raise ValueError('Array length should be equal to the number of l in the expansion.')
        return np.repeat(arr, 2 * self.l_array + 1, axis=axis)


def rot_sym_expansion(l_array: Array, coeffs: Array) -> Array:
    """Create full expansion array for rotationally symmetric distributions with only m=0 terms different form 0."""
    full_coeffs = np.zeros(np.sum(2 * l_array + 1))
    full_coeffs[np.cumsum(2 * l_array + 1) - l_array - 1] = coeffs
    return full_coeffs


class Expansion24(Expansion):

    def __init__(self, sigma2: float, sigma4: float, sigma0: float = 0.):
        l_array = np.array([0, 2, 4])
        coeffs = rot_sym_expansion(l_array, np.array([sigma0, sigma2, sigma4]))
        super().__init__(l_array, coeffs)


class MappedExpansion(Expansion):

    def __init__(self, a_bar: float, kappaR: float, sigma_m: float, max_l: int = 20, sigma0: float = 0):
        l_array = np.array([l for l in range(max_l + 1) if l % 2 == 0])
        coeffs = (2 * sigma_m * fn.coeff_C_diff(l_array, kappaR)
                  * np.sqrt(4 * np.pi * (2 * l_array + 1)) * np.power(a_bar, l_array))
        coeffs[0] = sigma0
        coeffs = rot_sym_expansion(l_array, coeffs)
        super().__init__(l_array, coeffs)


if __name__ == '__main__':

    ex = MappedExpansion(0.44, 3, 1, 10)
    print(ex.coeffs)