import torch
from torch import nn

class ArcHeads(nn.Sequential):
    def __init__(self, in_channels, layers):
        d = []
        next_feature = in_channels
        for out_channels in layers:
            d.append(nn.Conv2d(next_feature, out_channels, 3, stride=1, padding=1))
            d.append(nn.ReLU(inplace=True))
            next_feature = out_channels
        super().__init__(*d)
        for m in self.children():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
                nn.init.constant_(m.bias, 0)


class ArcPredictor(nn.Module):
    def __init__(self, in_channels, out_channels=1 ):
        super().__init__()
        input_features = in_channels
        deconv_kernel = 4
        self.kps_score_lowres = nn.ConvTranspose2d(
            input_features,
            out_channels,
            deconv_kernel,
            stride=2,
            padding=deconv_kernel // 2 - 1,
        )
        nn.init.kaiming_normal_(self.kps_score_lowres.weight, mode="fan_out", nonlinearity="relu")
        nn.init.constant_(self.kps_score_lowres.bias, 0)
        self.up_scale = 2
        self.out_channels = out_channels

    def forward(self, x):
        print(f'before kps_score_lowres x:{x.shape}')
        x = self.kps_score_lowres(x)
        print(f'kps_score_lowres x:{x.shape}')
        return x
        # return torch.nn.functional.interpolate(
        #     x, scale_factor=float(self.up_scale), mode="bilinear", align_corners=False, recompute_scale_factor=False
        # )


import torch
import torch.nn as nn
import torch.nn.functional as F


class ArcEquationHead(nn.Module):
    def __init__(self, num_outputs=7):
        super().__init__()

        # --------------------------------------------------
        # Convolution layers - no fixed H,W assumptions
        # Automatically downsamples using stride=2
        # --------------------------------------------------
        self.conv = nn.Sequential(
            nn.Conv2d(1, 32, 3, stride=2, padding=1),
            nn.ReLU(inplace=True),

            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.ReLU(inplace=True),

            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.ReLU(inplace=True),

            nn.Conv2d(128, 256, 3, stride=2, padding=1),
            nn.ReLU(inplace=True),
        )

        # --------------------------------------------------
        # Global pooling ¡ú no H,W dependency
        # --------------------------------------------------
        self.gap = nn.AdaptiveAvgPool2d((1, 1))

        # --------------------------------------------------
        # MLP
        # --------------------------------------------------
        self.mlp = nn.Sequential(
            nn.Linear(256, 256),
            nn.ReLU(inplace=True),
            nn.Linear(256, num_outputs)
        )


    def forward(self, feature_logits):
        """
        Args:
            feature_logits: Tensor [N, 1, H, W]
        """

        # CNN
        x = self.conv(feature_logits)

        # Global pool
        x = self.gap(x).view(x.size(0), -1)

        # Predict params
        arc_params = self.mlp(x)   # -> [N, 7]

        N, _, H, W = feature_logits.shape

        # --------------------------------------------
        # Apply constraints
        # --------------------------------------------
        arc_params[..., 0] = torch.sigmoid(arc_params[..., 0]) * W   # cx
        arc_params[..., 1] = torch.sigmoid(arc_params[..., 1]) * H   # cy

        arc_params[..., 2] = F.relu(arc_params[..., 2]) + 1e-6        # long axis
        arc_params[..., 3] = F.relu(arc_params[..., 3]) + 1e-6        # short axis

        # angles 0~2¦Ð
        arc_params[..., 4:7] = torch.sigmoid(arc_params[..., 4:7]) * (2 * 3.1415926535)

        return arc_params