from torch import nn
import torch
from sudoku.symetries import mat_sym


class SudokuNet(nn.Module):
    def __init__(self, n_output=2, coef_hidden=4):
        super(SudokuNet, self).__init__()
        self.hidden_neural_number = mat_sym.shape[1]

        self.conv111 = nn.Conv1d(
            self.hidden_neural_number * 2,
            self.hidden_neural_number * 2 * coef_hidden,
            1,
            groups=self.hidden_neural_number * 2,
        )
        self.conv111_last = nn.Conv1d(
            self.hidden_neural_number * 2 * coef_hidden, n_output, 1
        )

        sym_tensor = torch.from_numpy(mat_sym).type(torch.FloatTensor)
        self.sym_tensor = nn.Parameter(sym_tensor, requires_grad=False)

    def forward(self, x):
        x = torch.tensordot(x, self.sym_tensor, dims=([2], [2]))
        x = x.view(-1, 2, 9 * 9 * 9, self.hidden_neural_number)
        x = x.permute(0, 1, 3, 2)
        x = x.contiguous().view(-1, self.hidden_neural_number * 2, 9 * 9 * 9)
        x = torch.relu(self.conv111(x))
        x = self.conv111_last(x)
        return x


class SymPreprocess(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden_neural_number = mat_sym.shape[1]
        sym_tensor = torch.from_numpy(mat_sym).type(torch.FloatTensor)
        self.sym_tensor = nn.Parameter(sym_tensor, requires_grad=False)

    def forward(self, x):
        n_channel = x.shape[1]
        x = torch.tensordot(x, self.sym_tensor, dims=([2], [2]))
        x = x.view(-1, n_channel, 9 * 9 * 9, self.hidden_neural_number)
        x = x.permute(0, 1, 3, 2)
        x = x.contiguous().view(-1, self.hidden_neural_number * n_channel, 9 * 9 * 9)
        return x


class SmallNet(nn.Module):
    def __init__(self, n_output=2, coef_hidden=4, n_input_channel=2):
        super(SmallNet, self).__init__()
        self.hidden_neural_number = mat_sym.shape[1]

        self.conv111 = nn.Conv1d(
            self.hidden_neural_number * n_input_channel,
            self.hidden_neural_number * n_input_channel * coef_hidden,
            1,
            groups=self.hidden_neural_number * n_input_channel,
        )
        self.conv111_last = nn.Conv1d(
            self.hidden_neural_number * n_input_channel * coef_hidden, n_output, 1
        )

    def forward(self, x):
        x = torch.relu(self.conv111(x))
        x = self.conv111_last(x)
        return x


class SmallNetBis(nn.Module):
    def __init__(self, n_output=2, coef_hidden=2, n_input_channel=2):
        super(SmallNetBis, self).__init__()
        self.hidden_neural_number = mat_sym.shape[1] * 2

        self.conv111 = nn.Conv1d(
            self.hidden_neural_number * n_input_channel,
            self.hidden_neural_number * n_input_channel * coef_hidden,
            1,
            groups=self.hidden_neural_number * n_input_channel,
        )
        self.conv111_last = nn.Conv1d(
            self.hidden_neural_number * n_input_channel * coef_hidden, n_output, 1
        )

    def forward(self, x):
        x = torch.cat([x, 1 - x], dim=1)
        x = torch.relu(self.conv111(x))
        x = self.conv111_last(x)
        return x


class SplittedSmallNet(nn.Module):
    def __init__(self, coef_hidden=4, n_input_channel=2):
        super().__init__()
        self.hidden_neural_number = mat_sym.shape[1]

        self.conv111_0 = nn.Conv1d(
            self.hidden_neural_number * n_input_channel,
            self.hidden_neural_number * n_input_channel * coef_hidden,
            1,
            groups=self.hidden_neural_number * n_input_channel,
        )
        self.conv111_1 = nn.Conv1d(
            self.hidden_neural_number * n_input_channel,
            self.hidden_neural_number * n_input_channel * coef_hidden,
            1,
            groups=self.hidden_neural_number * n_input_channel,
        )

        self.conv111_last_0 = nn.Conv1d(
            self.hidden_neural_number * n_input_channel * coef_hidden, 1, 1
        )
        self.conv111_last_1 = nn.Conv1d(
            self.hidden_neural_number * n_input_channel * coef_hidden, 1, 1
        )

    def forward(self, x):
        x_0 = torch.relu(self.conv111_0(x))
        x_0 = self.conv111_last_0(x_0)

        x_1 = torch.relu(self.conv111_1(x))
        x_1 = self.conv111_last_1(x_1)
        return torch.cat([x_0, x_1], dim=1)