# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02_models.ipynb (unless otherwise specified).

__all__ = ['expand_filter', 'ChLin', 'FireHR', 'download_model_weights', 'load_pretrained_model']

# Cell
import os, sys
import requests
from fastai.vision.all import *
import FireHR

# Cell
def expand_filter(x, ks=3):
    with torch.no_grad():
        k5 = nn.Conv2d(1, 1, kernel_size=ks, padding=ks//2, padding_mode='reflect', bias=False)
        k5.weight.data = torch.ones(1, 1, ks, ks)/(ks*ks)
        xbuffer = k5(x[:,-1].unsqueeze(1))
        x = torch.cat([x[:,:-1], xbuffer], dim=1)
        return x

class ChLin(Module):
    def __init__(self, ni, nf):
        self.chlin = nn.Sequential(
            nn.Linear(ni, nf, bias=False), nn.BatchNorm1d(nf), nn.ReLU(inplace=True))
    def forward(self, x):
        sh = x.shape
        x = x.permute(0,2,3,1).contiguous().view(sh[0]*sh[2]*sh[3], sh[1])
        x = self.chlin(x).view(sh[0],sh[2],sh[3], -1).permute(0,3,1,2).contiguous()
        return x

class FireHR(Module):
    def __init__(self, ni, nc):
        self.conv = ConvLayer(1, 8)
        self.chlin = nn.Sequential(ChLin(ni+8, 128), ChLin(128, 64))
        self.middleconv = nn.Sequential(ConvLayer(64, 128), ConvLayer(128, 64))
        self.finalconv = nn.Conv2d(64, nc, kernel_size=1, bias=True)

    def forward(self, x):
        x = torch.cat([x[:,:-1], self.conv(x[:,-1].unsqueeze(1))], dim=1)
        x = self.chlin(x)
        x = self.middleconv(x)
        return self.finalconv(x)

# Cell
def download_model_weights(weight_file='model512.pth'):
    """Download model weights if they don't exist yet on ~/.firehr."""
    path_save = Path(os.path.expandvars('$HOME'))/'.firehr'
    path_save.mkdir(exist_ok=True)
    file_save = path_save/weight_file
    if not file_save.is_file():
        print(f'Downloading model weights {weight_file}.')
        url = 'https://github.com/mnpinto/FireHR_weights/raw/main/model512.pth'
        file = requests.get(url)
        open(str(file_save), 'wb').write(file.content)
    else:
        print(f'Using local model weights {file_save}')

# Cell
_WEIGHTS = Path(os.path.expandvars('$HOME'))/'.firehr/model512.pth'
def load_pretrained_model(weights=_WEIGHTS, ni=6, nc=1, half_precision=True, gpu=True):
    download_model_weights()
    model = FireHR(ni,nc)
    st = torch.load(weights, map_location=torch.device('cpu'))
    model.load_state_dict(st['model'])
    if gpu:
        if half_precision: model = model.half()
        if torch.cuda.is_available():
            model = model.cuda()
        else:
            warnings.warn('GPU is not available. torch.cuda.is_available() returned False.')
    return model