import torch
from .resnet import ResNet, Bottleneck

__all__ = ['resnest50', 'resnest101', 'resnest200', 'resnest269']

_url_format = 'https://s3.us-west-1.wasabisys.com/resnest/torch/{}-{}.pth'

_model_sha256 = {
    name: checksum
    for checksum, name in [
        ('528c19ca', 'resnest50'),
        ('22405ba7', 'resnest101'),
        ('75117900', 'resnest200'),
        ('0cc87c48', 'resnest269'),
    ]
}


def short_hash(name):
    if name not in _model_sha256:
        raise ValueError(
            'Pretrained model for {name} is not available.'.format(name=name))
    return _model_sha256[name][:8]


resnest_model_urls = {
    name: _url_format.format(name, short_hash(name))
    for name in _model_sha256.keys()
}


def resnest50(pretrained=False, root='~/.encoding/models', **kwargs):
    model = ResNet(Bottleneck, [3, 4, 6, 3],
                   radix=2,
                   groups=1,
                   bottleneck_width=64,
                   deep_stem=True,
                   stem_width=32,
                   avg_down=True,
                   avd=True,
                   avd_first=False,
                   **kwargs)
    if pretrained:
        model.load_state_dict(
            torch.hub.load_state_dict_from_url(resnest_model_urls['resnest50'],
                                               progress=True,
                                               check_hash=True))
    return model


def resnest101(pretrained=False, root='~/.encoding/models', **kwargs):
    model = ResNet(Bottleneck, [3, 4, 23, 3],
                   radix=2,
                   groups=1,
                   bottleneck_width=64,
                   deep_stem=True,
                   stem_width=64,
                   avg_down=True,
                   avd=True,
                   avd_first=False,
                   **kwargs)
    if pretrained:
        model.load_state_dict(
            torch.hub.load_state_dict_from_url(
                resnest_model_urls['resnest101'],
                progress=True,
                check_hash=True))
    return model


def resnest200(pretrained=False, root='~/.encoding/models', **kwargs):
    model = ResNet(Bottleneck, [3, 24, 36, 3],
                   radix=2,
                   groups=1,
                   bottleneck_width=64,
                   deep_stem=True,
                   stem_width=64,
                   avg_down=True,
                   avd=True,
                   avd_first=False,
                   **kwargs)
    if pretrained:
        model.load_state_dict(
            torch.hub.load_state_dict_from_url(
                resnest_model_urls['resnest200'],
                progress=True,
                check_hash=True))
    return model


def resnest269(pretrained=False, root='~/.encoding/models', **kwargs):
    model = ResNet(Bottleneck, [3, 30, 48, 8],
                   radix=2,
                   groups=1,
                   bottleneck_width=64,
                   deep_stem=True,
                   stem_width=64,
                   avg_down=True,
                   avd=True,
                   avd_first=False,
                   **kwargs)
    if pretrained:
        model.load_state_dict(
            torch.hub.load_state_dict_from_url(
                resnest_model_urls['resnest269'],
                progress=True,
                check_hash=True))
    return model