from typing import Optional

from torch.nn import Sigmoid, Sequential, Tanh

from tha3.nn.conv import create_conv3, create_conv3_from_block_args
from tha3.nn.nonlinearity_factory import ReLUFactory
from tha3.nn.normalization import InstanceNorm2dFactory
from tha3.nn.util import BlockArgs


class PoserArgs00:
    def __init__(self,
                 image_size: int,
                 input_image_channels: int,
                 output_image_channels: int,
                 start_channels: int,
                 num_pose_params: int,
                 block_args: Optional[BlockArgs] = None):
        self.num_pose_params = num_pose_params
        self.start_channels = start_channels
        self.output_image_channels = output_image_channels
        self.input_image_channels = input_image_channels
        self.image_size = image_size
        if block_args is None:
            self.block_args = BlockArgs(
                normalization_layer_factory=InstanceNorm2dFactory(),
                nonlinearity_factory=ReLUFactory(inplace=True))
        else:
            self.block_args = block_args

    def create_alpha_block(self):
        from torch.nn import Sequential
        return Sequential(
            create_conv3(
                in_channels=self.start_channels,
                out_channels=1,
                bias=True,
                initialization_method=self.block_args.initialization_method,
                use_spectral_norm=False),
            Sigmoid())

    def create_all_channel_alpha_block(self):
        from torch.nn import Sequential
        return Sequential(
            create_conv3(
                in_channels=self.start_channels,
                out_channels=self.output_image_channels,
                bias=True,
                initialization_method=self.block_args.initialization_method,
                use_spectral_norm=False),
            Sigmoid())

    def create_color_change_block(self):
        return Sequential(
            create_conv3_from_block_args(
                in_channels=self.start_channels,
                out_channels=self.output_image_channels,
                bias=True,
                block_args=self.block_args),
            Tanh())

    def create_grid_change_block(self):
        return create_conv3(
            in_channels=self.start_channels,
            out_channels=2,
            bias=False,
            initialization_method='zero',
            use_spectral_norm=False)