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

from utils.Stomach_Config import Model_Config

class SeparableConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True):
        super(SeparableConv2d, self).__init__()
        self.depthwise = nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, stride=stride,
                                 padding=padding, groups=in_channels, bias=bias)
        self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1,
                                 padding=0, bias=bias)
    
    def forward(self, x):
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x

class EntryFlow(nn.Module):
    def __init__(self, in_channels=3):
        super(EntryFlow, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=2, padding=1, bias=False, dilation = 2)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1, bias=False, dilation = 2)
        self.bn2 = nn.BatchNorm2d(64)
        
        self.conv3_residual = nn.Sequential(
            SeparableConv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=False),  # 修改這裡
            SeparableConv2d(128, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.MaxPool2d(3, stride=2, padding=1)
        )
        self.conv3_shortcut = nn.Conv2d(64, 128, 1, stride=2, bias=False)
        self.bn3 = nn.BatchNorm2d(128)
        
        self.conv4_residual = nn.Sequential(
            nn.ReLU(inplace=False),  # 修改這裡
            SeparableConv2d(128, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=False),  # 修改這裡
            SeparableConv2d(256, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.MaxPool2d(3, stride=2, padding=1)
        )
        self.conv4_shortcut = nn.Conv2d(128, 256, 1, stride=2, bias=False)
        self.bn4 = nn.BatchNorm2d(256)
        
        self.conv5_residual = nn.Sequential(
            nn.ReLU(inplace=False),  # 修改這裡
            SeparableConv2d(256, 728, 3, padding=1),
            nn.BatchNorm2d(728),
            nn.ReLU(inplace=False),  # 修改這裡
            SeparableConv2d(728, 728, 3, padding=1),
            nn.BatchNorm2d(728),
            nn.MaxPool2d(3, stride=2, padding=1)
        )
        self.conv5_shortcut = nn.Conv2d(256, 728, 1, stride=2, bias=False)
        self.bn5 = nn.BatchNorm2d(728)
    
    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        
        residual = self.conv3_residual(x)
        shortcut = self.conv3_shortcut(x)
        x = F.relu(self.bn3(residual + shortcut))
        
        residual = self.conv4_residual(x)
        shortcut = self.conv4_shortcut(x)
        x = F.relu(self.bn4(residual + shortcut))
        
        residual = self.conv5_residual(x)
        shortcut = self.conv5_shortcut(x)
        x = F.relu(self.bn5(residual + shortcut))
        return x

class MiddleFlow(nn.Module):
    def __init__(self):
        super(MiddleFlow, self).__init__()
        self.conv_residual = nn.Sequential(
            nn.ReLU(inplace=False),  # 修改這裡
            SeparableConv2d(728, 728, 3, padding=1),
            nn.BatchNorm2d(728),
            nn.ReLU(inplace=False),  # 修改這裡
            SeparableConv2d(728, 728, 3, padding=1),
            nn.BatchNorm2d(728),
            nn.ReLU(inplace=False),  # 修改這裡
            SeparableConv2d(728, 728, 3, padding=1),
            nn.BatchNorm2d(728)
        )
    
    def forward(self, x):
        return self.conv_residual(x) + x

class ExitFlow(nn.Module):
    def __init__(self):
        super(ExitFlow, self).__init__()
        self.conv1_residual = nn.Sequential(
            nn.ReLU(inplace=False),  # 修改這裡
            SeparableConv2d(728, 1024, 3, padding=1),
            nn.BatchNorm2d(1024),
            nn.ReLU(inplace=False),  # 修改這裡
            SeparableConv2d(1024, 1024, 3, padding=1),
            nn.BatchNorm2d(1024),
            nn.MaxPool2d(3, stride=2, padding=1)
        )
        self.conv1_shortcut = nn.Conv2d(728, 1024, 1, stride=2, bias=False)
        self.bn1 = nn.BatchNorm2d(1024)
        
        self.conv2 = nn.Sequential(
            SeparableConv2d(1024, 1536, 3, padding=1),
            nn.BatchNorm2d(1536),
            nn.ReLU(inplace=False),  # 修改這裡
            SeparableConv2d(1536, 2048, 3, padding=1),
            nn.BatchNorm2d(2048),
            nn.ReLU(inplace=False)   # 修改這裡
        )
        self.avgpool = nn.AdaptiveAvgPool1d(Model_Config["GPA Output Nodes"])
        self.Hidden = nn.Linear(Model_Config["GPA Output Nodes"], Model_Config["Linear Hidden Nodes"])
        self.fc = nn.Linear(Model_Config["Linear Hidden Nodes"], Model_Config["Output Linear Nodes"])
        self.dropout = nn.Dropout(Model_Config["Dropout Rate"])
    
    def forward(self, x):
        residual = self.conv1_residual(x)
        shortcut = self.conv1_shortcut(x)
        x = F.relu(self.bn1(residual + shortcut))
        
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        x = self.avgpool(x)
        x = F.relu(self.Hidden(x))
        x = self.dropout(x)
        x = self.fc(x)
        return x

class Xception(nn.Module):
    def __init__(self):
        super(Xception, self).__init__()
        self.entry_flow = EntryFlow(in_channels=3)  # 默认输入通道为3
        self.middle_flow = nn.Sequential(*[MiddleFlow() for _ in range(8)])
        self.exit_flow = ExitFlow()
    
    def forward(self, x):
        # 正常的前向傳播
        x = self.entry_flow(x)
        x = self.middle_flow(x)
        x = self.exit_flow(x)
        
        return x