import torch
import torch.nn as nn
import torch.nn.functional as F
from models.misc import modules

constrain_path = {
    ('threeD', 'normal'): (True, True, ''),
    ('threeD', 'depth'): (True, True, ''),
    ('normal', 'depth'): (True, True, ''),
    ('depth', 'normal'): (True, True, ''),

}


class UnwarpNet(nn.Module):
    def __init__(self, use_simple=False, combine_num=3, use_constrain=True, constrain_configure=None):
        super(UnwarpNet, self).__init__()
        self.combine_num = combine_num
        self.use_simple = use_simple
        self.use_constrain = use_constrain
        self.constrain_configure = constrain_configure
        
        self.geo_encoder = modules.Encoder(downsample=6, in_channels=3)
        self.threeD_decoder = modules.Decoder(downsample=6, out_channels=3, combine_num=self.combine_num)
        self.normal_decoder = modules.Decoder(downsample=6, out_channels=3, combine_num=self.combine_num)
        self.depth_decoder = modules.Decoder(downsample=6, out_channels=1, combine_num=self.combine_num)
        self.mask_decoder = modules.Decoder(downsample=6, out_channels=1, combine_num=0)
        bottle_neck = sum([2 ** (i + 4) for i in range(self.combine_num)])
        self.second_encoder = modules.Encoder(downsample=6, in_channels=bottle_neck * 3 + 3)
        self.uv_decoder = modules.Decoder(downsample=6, out_channels=2, combine_num=0)
        # self.albedo_decoder = modules.AlbedoDecoder(downsample=6, out_channels=1)
        self.albedo_decoder = modules.Decoder(downsample=6, out_channels=1, combine_num=0)
        self.deform_decoder = modules.Decoder(downsample=6, out_channels=2, combine_num=0)
        self.dep2nor = None
        self.threeD_to_nor2dep = None
        self.nor2dep = None
        
    def forward(self, x):
        gxvals, gx_encode = self.geo_encoder(x)
        threeD_map, threeD_feature = self.threeD_decoder(gxvals, gx_encode)
        threeD_map = nn.functional.tanh(threeD_map)
        dep_map, dep_feature = self.depth_decoder(gxvals, gx_encode)
        dep_map = nn.functional.tanh(dep_map)
        nor_map, nor_feature = self.normal_decoder(gxvals, gx_encode)
        nor_map = nn.functional.tanh(nor_map)
        mask_map, mask_feature = self.mask_decoder(gxvals, gx_encode)
        mask_map = torch.nn.functional.sigmoid(mask_map)
        # geo_feature = torch.cat([threeD_feature, nor_feature, dep_feature], dim=1)
        geo_feature = torch.cat([threeD_feature, nor_feature, dep_feature, x], dim=1)
        b, c, h, w = geo_feature.size()
        geo_feature_mask = geo_feature.mul(mask_map.expand(b, c, h, w))
        secvals, sec_encode = self.second_encoder(geo_feature_mask)
        uv_map, _ = self.uv_decoder(secvals, sec_encode)
        uv_map = nn.functional.tanh(uv_map)
        alb_map, _ = self.albedo_decoder(secvals, sec_encode)
        alb_map = nn.functional.tanh(alb_map)
        deform_map, _ = self.deform_decoder(secvals, sec_encode)
        deform_map = nn.functional.tanh(deform_map)
        return uv_map, threeD_map, nor_map, alb_map, dep_map, mask_map, \
               None, None, None, None, None, deform_map

[((680, 724), 'models.misc.modules.Encoder', 'modules.Encoder', ([], {'downsample': '(6)', 'in_channels': '(3)'}), '(downsample=6, in_channels=3)\n', (695, 724), False, 'from models.misc import modules\n'), ((755, 830), 'models.misc.modules.Decoder', 'modules.Decoder', ([], {'downsample': '(6)', 'out_channels': '(3)', 'combine_num': 'self.combine_num'}), '(downsample=6, out_channels=3, combine_num=self.combine_num)\n', (770, 830), False, 'from models.misc import modules\n'), ((861, 936), 'models.misc.modules.Decoder', 'modules.Decoder', ([], {'downsample': '(6)', 'out_channels': '(3)', 'combine_num': 'self.combine_num'}), '(downsample=6, out_channels=3, combine_num=self.combine_num)\n', (876, 936), False, 'from models.misc import modules\n'), ((966, 1041), 'models.misc.modules.Decoder', 'modules.Decoder', ([], {'downsample': '(6)', 'out_channels': '(1)', 'combine_num': 'self.combine_num'}), '(downsample=6, out_channels=1, combine_num=self.combine_num)\n', (981, 1041), False, 'from models.misc import modules\n'), ((1070, 1130), 'models.misc.modules.Decoder', 'modules.Decoder', ([], {'downsample': '(6)', 'out_channels': '(1)', 'combine_num': '(0)'}), '(downsample=6, out_channels=1, combine_num=0)\n', (1085, 1130), False, 'from models.misc import modules\n'), ((1236, 1298), 'models.misc.modules.Encoder', 'modules.Encoder', ([], {'downsample': '(6)', 'in_channels': '(bottle_neck * 3 + 3)'}), '(downsample=6, in_channels=bottle_neck * 3 + 3)\n', (1251, 1298), False, 'from models.misc import modules\n'), ((1325, 1385), 'models.misc.modules.Decoder', 'modules.Decoder', ([], {'downsample': '(6)', 'out_channels': '(2)', 'combine_num': '(0)'}), '(downsample=6, out_channels=2, combine_num=0)\n', (1340, 1385), False, 'from models.misc import modules\n'), ((1500, 1560), 'models.misc.modules.Decoder', 'modules.Decoder', ([], {'downsample': '(6)', 'out_channels': '(1)', 'combine_num': '(0)'}), '(downsample=6, out_channels=1, combine_num=0)\n', (1515, 1560), False, 'from models.misc import modules\n'), ((1591, 1651), 'models.misc.modules.Decoder', 'modules.Decoder', ([], {'downsample': '(6)', 'out_channels': '(2)', 'combine_num': '(0)'}), '(downsample=6, out_channels=2, combine_num=0)\n', (1606, 1651), False, 'from models.misc import modules\n'), ((1926, 1956), 'torch.nn.functional.tanh', 'nn.functional.tanh', (['threeD_map'], {}), '(threeD_map)\n', (1944, 1956), True, 'import torch.nn as nn\n'), ((2044, 2071), 'torch.nn.functional.tanh', 'nn.functional.tanh', (['dep_map'], {}), '(dep_map)\n', (2062, 2071), True, 'import torch.nn as nn\n'), ((2160, 2187), 'torch.nn.functional.tanh', 'nn.functional.tanh', (['nor_map'], {}), '(nor_map)\n', (2178, 2187), True, 'import torch.nn as nn\n'), ((2277, 2314), 'torch.nn.functional.sigmoid', 'torch.nn.functional.sigmoid', (['mask_map'], {}), '(mask_map)\n', (2304, 2314), False, 'import torch\n'), ((2422, 2485), 'torch.cat', 'torch.cat', (['[threeD_feature, nor_feature, dep_feature, x]'], {'dim': '(1)'}), '([threeD_feature, nor_feature, dep_feature, x], dim=1)\n', (2431, 2485), False, 'import torch\n'), ((2740, 2766), 'torch.nn.functional.tanh', 'nn.functional.tanh', (['uv_map'], {}), '(uv_map)\n', (2758, 2766), True, 'import torch.nn as nn\n'), ((2847, 2874), 'torch.nn.functional.tanh', 'nn.functional.tanh', (['alb_map'], {}), '(alb_map)\n', (2865, 2874), True, 'import torch.nn as nn\n'), ((2961, 2991), 'torch.nn.functional.tanh', 'nn.functional.tanh', (['deform_map'], {}), '(deform_map)\n', (2979, 2991), True, 'import torch.nn as nn\n')]