from typing import Optional

from tha3.nn.conv import create_conv7_block_from_block_args, create_conv3_block_from_block_args, \
    create_downsample_block_from_block_args, create_conv3
from tha3.nn.resnet_block import ResnetBlock
from tha3.nn.resnet_block_seperable import ResnetBlockSeparable
from tha3.nn.separable_conv import create_separable_conv7_block, create_separable_conv3_block, \
    create_separable_downsample_block, create_separable_conv3
from tha3.nn.util import BlockArgs


class ConvBlockFactory:
    def __init__(self,
                 block_args: BlockArgs,
                 use_separable_convolution: bool = False):
        self.use_separable_convolution = use_separable_convolution
        self.block_args = block_args

    def create_conv3(self,
                     in_channels: int,
                     out_channels: int,
                     bias: bool,
                     initialization_method: Optional[str] = None):
        if initialization_method is None:
            initialization_method = self.block_args.initialization_method
        if self.use_separable_convolution:
            return create_separable_conv3(
                in_channels, out_channels, bias, initialization_method, self.block_args.use_spectral_norm)
        else:
            return create_conv3(
                in_channels, out_channels, bias, initialization_method, self.block_args.use_spectral_norm)

    def create_conv7_block(self, in_channels: int, out_channels: int):
        if self.use_separable_convolution:
            return create_separable_conv7_block(in_channels, out_channels, self.block_args)
        else:
            return create_conv7_block_from_block_args(in_channels, out_channels, self.block_args)

    def create_conv3_block(self, in_channels: int, out_channels: int):
        if self.use_separable_convolution:
            return create_separable_conv3_block(in_channels, out_channels, self.block_args)
        else:
            return create_conv3_block_from_block_args(in_channels, out_channels, self.block_args)

    def create_downsample_block(self, in_channels: int, out_channels: int, is_output_1x1: bool):
        if self.use_separable_convolution:
            return create_separable_downsample_block(in_channels, out_channels, is_output_1x1, self.block_args)
        else:
            return create_downsample_block_from_block_args(in_channels, out_channels, is_output_1x1)

    def create_resnet_block(self, num_channels: int, is_1x1: bool):
        if self.use_separable_convolution:
            return ResnetBlockSeparable.create(num_channels, is_1x1, block_args=self.block_args)
        else:
            return ResnetBlock.create(num_channels, is_1x1, block_args=self.block_args)