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


class ElectraDTA(nn.Module):
    def __init__(self, smilen, seq_len, hidden_dim):
        super().__init__()
        self.drug_input = nn.Linear(smilen, hidden_dim)
        self.prot_input = nn.Linear(seq_len, hidden_dim)
        self.num_filters = hidden_dim
        self.filter_length = 3
        self.n_layers = 4
        self.seblock = True
        self.encode_smiles = ConvBlock(hidden_dim, self.seblock, self.num_filters, self.filter_length)
        self.encode_prot = ConvBlock(hidden_dim, self.seblock, self.num_filters, self.filter_length)
        self.global_pool = nn.AdaptiveMaxPool1d(1)
        self.concat = nn.Linear(hidden_dim * 4, hidden_dim * 4)
        self.predictions = FCNet(hidden_dim * 4)  # you need to define this function
        self.interaction_model = nn.Sequential(self.drug_input, self.prot_input, self.encode_smiles, self.encode_prot,
                                               self.global_pool, self.concat, self.predictions)

    def forward(self, x):
        x = self.interaction_model(x)
        return x


class SEBlock(nn.Module):
    def __init__(self, channels, r=8):
        super().__init__()
        self.squeeze = nn.AdaptiveAvgPool1d(1)
        self.excitation = nn.Sequential(
            nn.Linear(channels, channels // r),
            nn.ReLU(),
            nn.Linear(channels // r, channels),
            nn.Sigmoid()
        )

    def forward(self, x):
        out = self.squeeze(x)
        out = self.excitation(out)
        return x * out


class ConvBlock(nn.Module):
    def __init__(self, input_channels, seblock, num_filters, filter_length):
        super().__init__()
        self.conv1 = nn.Conv1d(input_channels, num_filters, filter_length, padding='valid', stride=1)
        self.conv2 = nn.Conv1d(num_filters, num_filters * 2, filter_length, padding='valid', stride=1)
        self.seblock = seblock
        if seblock:
            self.se1 = SEBlock(num_filters)
            self.se2 = SEBlock(num_filters * 2)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        if self.seblock:
            x = self.se1(x)
        x = F.relu(self.conv2(x))
        if self.seblock:
            x = self.se2(x)
        return x


class Highway(nn.Module):
    def __init__(self, dim, n_layers, activation=nn.Tanh(), gate_bias=0):
        super(Highway, self).__init__()
        self.n_layers = n_layers
        self.activation = activation
        self.T_gates = nn.ModuleList([nn.Linear(dim, dim) for _ in range(n_layers)])
        self.transforms = nn.ModuleList([nn.Linear(dim, dim) for _ in range(n_layers)])
        self.sigmoid = nn.Sigmoid()
        nn.init.constant_(self.linear.bias, gate_bias)

    def forward(self, x):
        for i in range(self.n_layers):
            T = self.sigmoid(self.gate(x))
            C = 1 - T
            H = self.activation(self.linear(x))
            x = T * H + C * x
        return x


class FCNet(nn.Module):
    def __init__(self, input_dim):
        super().__init__()
        self.n_layers = 4
        self.highway = Highway(input_dim, self.n_layers, gate_bias=-2)
        self.fc1 = nn.Linear(input_dim, 1024)
        self.fc2 = nn.Linear(1024, 1024)
        self.fc3 = nn.Linear(1024, 512)
        # self.fc4 = nn.Linear(512, 1)
        self.dropout = nn.Dropout(0.4)

    def forward(self, x):
        x = self.highway(x)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        x = F.relu(self.fc3(x))
        # x = self.fc4(x)
        return x