from dataclasses import field
from typing import Optional, Dict, cast

import flax.linen as nn
import jax
import jax.numpy as jnp

from . import RearrangeMixin, ReduceMixin
from ._einmix import _EinmixMixin

__author__ = "Alex Rogozhnikov"


class Reduce(nn.Module):
    pattern: str
    reduction: str
    sizes: dict = field(default_factory=lambda: {})

    def setup(self):
        self.reducer = ReduceMixin(self.pattern, self.reduction, **self.sizes)

    def __call__(self, input):
        return self.reducer._apply_recipe(input)


class Rearrange(nn.Module):
    pattern: str
    sizes: dict = field(default_factory=lambda: {})

    def setup(self):
        self.rearranger = RearrangeMixin(self.pattern, **self.sizes)

    def __call__(self, input):
        return self.rearranger._apply_recipe(input)


class EinMix(nn.Module, _EinmixMixin):
    pattern: str
    weight_shape: str
    bias_shape: Optional[str] = None
    sizes: dict = field(default_factory=lambda: {})

    def setup(self):
        self.initialize_einmix(
            pattern=self.pattern,
            weight_shape=self.weight_shape,
            bias_shape=self.bias_shape,
            axes_lengths=self.sizes,
        )

    def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound):
        self.weight = self.param("weight", jax.nn.initializers.uniform(weight_bound), weight_shape)

        if bias_shape is not None:
            self.bias = self.param("bias", jax.nn.initializers.uniform(bias_bound), bias_shape)
        else:
            self.bias = None

    def _create_rearrange_layers(
        self,
        pre_reshape_pattern: Optional[str],
        pre_reshape_lengths: Optional[Dict],
        post_reshape_pattern: Optional[str],
        post_reshape_lengths: Optional[Dict],
    ):
        self.pre_rearrange = None
        if pre_reshape_pattern is not None:
            self.pre_rearrange = Rearrange(pre_reshape_pattern, sizes=cast(dict, pre_reshape_lengths))

        self.post_rearrange = None
        if post_reshape_pattern is not None:
            self.post_rearrange = Rearrange(post_reshape_pattern, sizes=cast(dict, post_reshape_lengths))

    def __call__(self, input):
        if self.pre_rearrange is not None:
            input = self.pre_rearrange(input)
        result = jnp.einsum(self.einsum_pattern, input, self.weight)
        if self.bias is not None:
            result += self.bias
        if self.post_rearrange is not None:
            result = self.post_rearrange(result)
        return result