import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from blazeblock import BlazeBlock


class BlazePose(nn.Module):
    """The BlazePose pose detection model from MediaPipe.

    Because we won't be training this model, it doesn't need to have
    batchnorm layers. These have already been "folded" into the conv 
    weights by TFLite.

    The conversion to PyTorch is fairly straightforward, but there are 
    some small differences between TFLite and PyTorch in how they handle
    padding on conv layers with stride 2.

    This version works on batches, while the MediaPipe version can only
    handle a single image at a time.

    Based on code from https://github.com/tkat0/PyTorch_BlazeFace/ and
    https://github.com/google/mediapipe/
    """
    def __init__(self):
        super(BlazePose, self).__init__()

        self._define_layers()

    def load_weights(self, path):
        self.load_state_dict(torch.load(path))
        self.eval() 

    def _define_layers(self):
        self.backbone1 = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=48, kernel_size=5, stride=1, padding=2, bias=True),
            nn.ReLU(inplace=True),

            BlazeBlock(48, 48, 5),
            BlazeBlock(48, 48, 5),
            BlazeBlock(48, 48, 5),
            BlazeBlock(48, 64, 5, 2, skip_proj=True),

            BlazeBlock(64, 64, 5),
            BlazeBlock(64, 64, 5),
            BlazeBlock(64, 64, 5),
            BlazeBlock(64, 96, 5, 2, skip_proj=True),

            BlazeBlock(96, 96, 5),
            BlazeBlock(96, 96, 5),
            BlazeBlock(96, 96, 5),
            BlazeBlock(96, 96, 5),
            BlazeBlock(96, 96, 5),
            BlazeBlock(96, 96, 5),
            BlazeBlock(96, 128, 5, 2, skip_proj=True),

            BlazeBlock(128, 128, 5),
            BlazeBlock(128, 128, 5),
            BlazeBlock(128, 128, 5),
            BlazeBlock(128, 128, 5),
            BlazeBlock(128, 128, 5),
            BlazeBlock(128, 128, 5),
            BlazeBlock(128, 128, 5),

        )

        self.backbone2 = nn.Sequential(
            BlazeBlock(128, 256, 5, 2, skip_proj=True),
            BlazeBlock(256, 256, 5),
            BlazeBlock(256, 256, 5),
            BlazeBlock(256, 256, 5),
            BlazeBlock(256, 256, 5),
            BlazeBlock(256, 256, 5),
            BlazeBlock(256, 256, 5),

        )

        self.classifier_8 = nn.Conv2d(128, 2, 1, bias=True)
        self.classifier_16 = nn.Conv2d(256, 6, 1, bias=True)

        self.regressor_8 = nn.Conv2d(128, 24, 1, bias=True)
        self.regressor_16 = nn.Conv2d(256, 72, 1, bias=True)

    def forward(self, x):
        # TFLite uses slightly different padding on the first conv layer
        # than PyTorch, so do it manually.
        # x = F.pad(x, (1, 2, 1, 2), "constant", 0)
        
        b = x.shape[0]      # batch size, needed for reshaping later

        x = self.backbone1(x)         
        h = self.backbone2(x)         
        
        # Note: Because PyTorch is NCHW but TFLite is NHWC, we need to
        # permute the output from the conv layers before reshaping it.
        
        c1 = self.classifier_8(x)
        # print(c1)
        # print(c1.shape)

        c1 = c1.permute(0, 2, 3, 1) 
        c1 = c1.reshape(b, -1, 1)   

        c2 = self.classifier_16(h)   
        c2 = c2.permute(0, 2, 3, 1)  
        c2 = c2.reshape(b, -1, 1)    

        c = torch.cat((c1, c2), dim=1) 

        r1 = self.regressor_8(x)       
        r1 = r1.permute(0, 2, 3, 1)    
        r1 = r1.reshape(b, -1, 12)     

        r2 = self.regressor_16(h)      
        r2 = r2.permute(0, 2, 3, 1)    
        r2 = r2.reshape(b, -1, 12)     

        r = torch.cat((r1, r2), dim=1)  
        return [r, c]
