from typing import List

import torch
import torch.nn as nn
from torch.optim.lr_scheduler import MultiStepLR

from isr.models.lightning_model import LightningIsr


class _ResidualBlock(nn.Module):
    def __init__(self, n_channels: int = 64):
        super(_ResidualBlock, self).__init__()
        self.conv1 = nn.Conv2d(n_channels, n_channels, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(n_channels)
        self.prelu = nn.PReLU(n_channels)
        self.conv2 = nn.Conv2d(n_channels, n_channels, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(n_channels)

    def forward(self, x):
        hid = self.conv1(x)
        hid = self.bn1(hid)
        hid = self.prelu(hid)
        hid = self.conv2(hid)
        out = self.bn2(hid)

        return torch.add(x, out)


class _SubPixelBlock(nn.Module):
    def __init__(self, in_channels: int = 64, out_channels: int = 64, scale_factor: int = 2):
        super(_SubPixelBlock, self).__init__()
        n_out = out_channels * scale_factor ** 2
        self.conv = nn.Conv2d(in_channels, n_out, kernel_size=3, stride=1, padding=1)
        self.shuffle = nn.PixelShuffle(scale_factor)
        self.prelu = nn.PReLU(out_channels)

    def forward(self, x):
        hid = self.conv(x)
        hid = self.shuffle(hid)
        out = self.prelu(hid)

        return out


class _ResNetModule(nn.Module):
    def __init__(self, in_channels: int = 64, out_channels: int = 64, n_blocks: int = 5):
        super(_ResNetModule, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=9, stride=1, padding=4)
        self.prelu = nn.PReLU(out_channels)

        self.resnet = nn.Sequential(*[_ResidualBlock(out_channels) for _ in range(n_blocks)])

        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
        self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        inp = self.prelu(self.conv1(x))
        hid = self.resnet(inp)
        out = self.bn(self.conv2(hid))

        return torch.add(inp, out)


class _ScaleModule(nn.Module):
    def __init__(self, in_channels: int = 64, out_channels: int = 64, scales: List[int] = None):
        super(_ScaleModule, self).__init__()

        scales = scales or [2, 2]
        channels = in_channels

        layers = []
        for scale in scales:
            layers.append(_SubPixelBlock(channels, out_channels, scale))
            channels = out_channels

        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)


class SrResNet(LightningIsr):
    """
    https://arxiv.org/pdf/1609.04802.pdf
    """
    def __init__(self, hparams):

        super(SrResNet, self).__init__(hparams)
        self.resnet = _ResNetModule(self.in_channels, hparams.hid_channels, hparams.n_blocks)

        # attempt to scale up in multiple steps
        scales = self._compute_scale_steps(hparams.scale_factor)

        self.scaler = _ScaleModule(hparams.hid_channels, hparams.hid_channels, scales)
        self.conv = nn.Conv2d(
            hparams.hid_channels, self.in_channels, kernel_size=9, stride=1, padding=4
        )

    @staticmethod
    def _compute_scale_steps(scale_factor: int):
        scales = []
        scale_factor = scale_factor
        while scale_factor % 2 == 0 and scale_factor > 2:
            scales.append(2)
            scale_factor = scale_factor // 2
        scales.append(scale_factor)
        return scales

    def forward(self, x):
        hid = self.resnet(x)
        upscale = self.scaler(hid)
        out = self.conv(upscale)
        return torch.clamp(out, 0., 1.)

    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = LightningIsr.add_model_specific_args(parent_parser)
        parser.add_argument('--learning_rate', type=float, default=0.002, help='base learning rate')
        parser.add_argument('--weight_decay', type=float, default=0.,
                            help='weight decay penalty (default=0)')
        parser.add_argument('--hid_channels', type=int, default=64)
        parser.add_argument('--n_blocks', type=int, default=5)
        parser.add_argument('--b1', type=float, default=0.9)
        parser.add_argument('--b2', type=float, default=0.999)
        return parser

    def configure_optimizers(self):
        optim = torch.optim.Adam(
            self.parameters(),
            lr=self.hparams.learning_rate,
            weight_decay=self.hparams.weight_decay
        )
        scheduler = MultiStepLR(
            optim,
            milestones=[self.hparams.lr_epochs // 2, 3 * (self.hparams.lr_epochs // 4)]
        )
        return [optim], [scheduler]

from typing import List

import torch
import torch.nn as nn
from torch.optim.lr_scheduler import MultiStepLR

from isr.models.lightning_model import LightningIsr


class _ResidualBlock(nn.Module):
    def __init__(self, n_channels: int = 64):
        super(_ResidualBlock, self).__init__()
        self.conv1 = nn.Conv2d(n_channels, n_channels, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(n_channels)
        self.prelu = nn.PReLU(n_channels)
        self.conv2 = nn.Conv2d(n_channels, n_channels, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(n_channels)

    def forward(self, x):
        hid = self.conv1(x)
        hid = self.bn1(hid)
        hid = self.prelu(hid)
        hid = self.conv2(hid)
        out = self.bn2(hid)

        return torch.add(x, out)


class _SubPixelBlock(nn.Module):
    def __init__(self, in_channels: int = 64, out_channels: int = 64, scale_factor: int = 2):
        super(_SubPixelBlock, self).__init__()
        n_out = out_channels * scale_factor ** 2
        self.conv = nn.Conv2d(in_channels, n_out, kernel_size=3, stride=1, padding=1)
        self.shuffle = nn.PixelShuffle(scale_factor)
        self.prelu = nn.PReLU(out_channels)

    def forward(self, x):
        hid = self.conv(x)
        hid = self.shuffle(hid)
        out = self.prelu(hid)

        return out


class _ResNetModule(nn.Module):
    def __init__(self, in_channels: int = 64, out_channels: int = 64, n_blocks: int = 5):
        super(_ResNetModule, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=9, stride=1, padding=4)
        self.prelu = nn.PReLU(out_channels)

        self.resnet = nn.Sequential(*[_ResidualBlock(out_channels) for _ in range(n_blocks)])

        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
        self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        inp = self.prelu(self.conv1(x))
        hid = self.resnet(inp)
        out = self.bn(self.conv2(hid))

        return torch.add(inp, out)


class _ScaleModule(nn.Module):
    def __init__(self, in_channels: int = 64, out_channels: int = 64, scales: List[int] = None):
        super(_ScaleModule, self).__init__()

        scales = scales or [2, 2]
        channels = in_channels

        layers = []
        for scale in scales:
            layers.append(_SubPixelBlock(channels, out_channels, scale))
            channels = out_channels

        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)


class SrResNet(LightningIsr):
    """
    https://arxiv.org/pdf/1609.04802.pdf
    """
    def __init__(self, hparams):

        super(SrResNet, self).__init__(hparams)
        self.resnet = _ResNetModule(self.in_channels, hparams.hid_channels, hparams.n_blocks)

        # attempt to scale up in multiple steps
        scales = self._compute_scale_steps(hparams.scale_factor)

        self.scaler = _ScaleModule(hparams.hid_channels, hparams.hid_channels, scales)
        self.conv = nn.Conv2d(
            hparams.hid_channels, self.in_channels, kernel_size=9, stride=1, padding=4
        )

    @staticmethod
    def _compute_scale_steps(scale_factor: int):
        scales = []
        scale_factor = scale_factor
        while scale_factor % 2 == 0 and scale_factor > 2:
            scales.append(2)
            scale_factor = scale_factor // 2
        scales.append(scale_factor)
        return scales

    def forward(self, x):
        hid = self.resnet(x)
        upscale = self.scaler(hid)
        out = self.conv(upscale)
        return torch.clamp(out, 0., 1.)

    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = LightningIsr.add_model_specific_args(parent_parser)
        parser.add_argument('--learning_rate', type=float, default=0.002, help='base learning rate')
        parser.add_argument('--weight_decay', type=float, default=0.,
                            help='weight decay penalty (default=0)')
        parser.add_argument('--hid_channels', type=int, default=64)
        parser.add_argument('--n_blocks', type=int, default=5)
        parser.add_argument('--b1', type=float, default=0.9)
        parser.add_argument('--b2', type=float, default=0.999)
        return parser

    def configure_optimizers(self):
        optim = torch.optim.Adam(
            self.parameters(),
            lr=self.hparams.learning_rate,
            weight_decay=self.hparams.weight_decay
        )
        scheduler = MultiStepLR(
            optim,
            milestones=[self.hparams.lr_epochs // 2, 3 * (self.hparams.lr_epochs // 4)]
        )
        return [optim], [scheduler]

https://arxiv.org/pdf/1609.04802.pdf # attempt to scale up in multiple steps