import functools
import math
import re
from collections import OrderedDict
from typing import Mapping

from typing_extensions import override

from .arch import RRDBNet
from ...factory import KeyCondition, Architecture
from ...utilities.state_dict import get_seq_len


def _new_to_old_arch(state: Mapping[str, object], state_map: dict, num_blocks: int):
    """Convert a new-arch model state dictionary to an old-arch dictionary."""
    if 'params_ema' in state:
        state = state['params_ema']

    if 'conv_first.weight' not in state:
        # model is already old arch, this is a loose check, but should be sufficient
        return state

    # add nb to state keys
    for kind in ('weight', 'bias'):
        state_map[f'model.1.sub.{num_blocks}.{kind}'] = state_map[f'model.1.sub./NB/.{kind}']
        del state_map[f'model.1.sub./NB/.{kind}']

    old_state = OrderedDict()
    for old_key, new_keys in state_map.items():
        for new_key in new_keys:
            if r'\1' in old_key:
                for k, v in state.items():
                    sub = re.sub(new_key, old_key, k)
                    if sub != k:
                        old_state[sub] = v
            else:
                if new_key in state:
                    old_state[old_key] = state[new_key]

    # upconv layers
    max_upconv = 0
    for key in state.keys():
        match = re.match(r'(upconv|conv_up)(\d)\.(weight|bias)', key)
        if match is not None:
            _, key_num, key_type = match.groups()
            old_state[f'model.{int(key_num) * 3}.{key_type}'] = state[key]
            max_upconv = max(max_upconv, int(key_num) * 3)

    # final layers
    for key in state.keys():
        if key in ('HRconv.weight', 'conv_hr.weight'):
            old_state[f'model.{max_upconv + 2}.weight'] = state[key]
        elif key in ('HRconv.bias', 'conv_hr.bias'):
            old_state[f'model.{max_upconv + 2}.bias'] = state[key]
        elif key in ('conv_last.weight',):
            old_state[f'model.{max_upconv + 4}.weight'] = state[key]
        elif key in ('conv_last.bias',):
            old_state[f'model.{max_upconv + 4}.bias'] = state[key]

    # Sort by first numeric value of each layer
    def compare(item1: str, item2: str):
        parts1 = item1.split('.')
        parts2 = item2.split('.')
        int1 = int(parts1[1])
        int2 = int(parts2[1])
        return int1 - int2

    sorted_keys = sorted(old_state.keys(), key=functools.cmp_to_key(compare))

    # Rebuild the output dict in the right order
    out_dict = OrderedDict((k, old_state[k]) for k in sorted_keys)

    return out_dict


def _get_scale(state: Mapping[str, object]) -> int:
    # model is composed of a few blocks that look like this flattened:
    #
    #   Conv2d
    #   B.ShortcutBlock
    #   [nn.Upsample, Conv2d, nn.LeakyReLU] for i in range(log2(scale))
    #   Conv2d
    #   nn.LeakyReLU (activation)
    #   Conv2d
    seq_len = get_seq_len(state, 'model')
    log2_scale = (seq_len - 5) // 3
    return 2**log2_scale


def _get_num_blocks(state: Mapping[str, object], state_map: dict) -> int:
    nbs = []
    state_keys = state_map[r'model.1.sub.\1.RDB\2.conv\3.0.\4'] + (r'model\.\d+\.sub\.(\d+)\.RDB(\d+)\.conv(\d+)\.0\.(weight|bias)',)
    for state_key in state_keys:
        for k in state:
            m = re.search(state_key, k)
            if m:
                nbs.append(int(m.group(1)))
        if nbs:
            break
    return max(*nbs) + 1


def _to_old_arch(state: Mapping[str, object]) -> Mapping[str, object]:
    state_map = {
        # currently supports old, new, and newer RRDBNet arch models
        # ESRGAN, BSRGAN/RealSR, Real-ESRGAN
        'model.0.weight': ('conv_first.weight',),
        'model.0.bias': ('conv_first.bias',),
        'model.1.sub./NB/.weight': ('trunk_conv.weight', 'conv_body.weight'),
        'model.1.sub./NB/.bias': ('trunk_conv.bias', 'conv_body.bias'),
        r'model.1.sub.\1.RDB\2.conv\3.0.\4': (
            r'RRDB_trunk\.(\d+)\.RDB(\d)\.conv(\d+)\.(weight|bias)',
            r'body\.(\d+)\.rdb(\d)\.conv(\d+)\.(weight|bias)',
        ),
    }

    if 'params_ema' in state:
        state = state['params_ema']

    num_blocks = _get_num_blocks(state, state_map)
    return _new_to_old_arch(state, state_map, num_blocks)


class ESRGANArch(Architecture[RRDBNet]):
    def __init__(self) -> None:
        super().__init__(
            uid='ESRGAN',
            detect=KeyCondition.has_any(
                KeyCondition.has_all(
                    'model.0.weight',
                    'model.1.sub.0.RDB1.conv1.0.weight',
                ),
                KeyCondition.has_all(
                    'conv_first.weight',
                    'body.0.rdb1.conv1.weight',
                    'conv_body.weight',
                    'conv_last.weight',
                ),
                KeyCondition.has_all(
                    # BSRGAN/RealSR
                    'conv_first.weight',
                    'RRDB_trunk.0.RDB1.conv1.weight',
                    'trunk_conv.weight',
                    'conv_last.weight',
                ),
                KeyCondition.has_all(
                    # ESRGAN+
                    'model.0.weight',
                    'model.1.sub.0.RDB1.conv1x1.weight',
                ),
            ),
        )

    @override
    def load(self, state_dict: Mapping[str, object]):
        # default values
        plus: bool = False

        state_dict = _to_old_arch(state_dict)

        model_seq_len = get_seq_len(state_dict, 'model')

        in_nc = state_dict['model.0.weight'].shape[1]
        out_nc = state_dict[f'model.{model_seq_len - 1}.weight'].shape[0]

        scale = _get_scale(state_dict)
        num_blocks = get_seq_len(state_dict, 'model.1.sub') - 1
        num_filters = state_dict['model.0.weight'].shape[0]

        if any('.conv1x1.' in k for k in state_dict.keys()):
            plus = True

        # Detect if pixelunshuffle was used (Real-ESRGAN)
        if in_nc in (out_nc * 4, out_nc * 16) and out_nc in (
            in_nc / 4,
            in_nc / 16,
        ):
            shuffle_factor = int(math.sqrt(in_nc / out_nc))
        else:
            shuffle_factor = None

        model = RRDBNet(
            in_nc=in_nc,
            out_nc=out_nc,
            num_filters=num_filters,
            num_blocks=num_blocks,
            scale=scale,
            plus=plus,
            shuffle_factor=shuffle_factor,
        )
        if shuffle_factor:
            in_nc //= shuffle_factor**2
            scale //= shuffle_factor
        return self._enhance_model(model=model, in_channels=in_nc, out_channels=out_nc, upscale=scale, name='ESRGAN')
