import sys

import pdb
import os.path as osp

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms
import numpy as np

import utils
import model.dpc as dpc
import model.resnet as resnet
import model.hrnet as hrnet
import model.random_feat_generator as random_feat_generator

sys.path.append('/home/junqi/videomontage/DPC/backbone')

class AppearanceModel(nn.Module):
    def __init__(self, args):
        super(AppearanceModel, self).__init__()
        self.args = args
        self.model_type = args.model_type
        self.use_diff_feat = args.use_diff_feat if hasattr(args, 'use_diff_feat') else False
        self.model = make_encoder(args,).to(self.args.device)

    def forward(self, x):
        if self.use_diff_feat:
            #print('x', x.shape)
            z_t, z_d = self.model(x)
            #print('z_t', z_t.shape, 'z_d', z_d.shape)
            return z_t, z_d
        else:
            #print('x', x.shape)
            z = self.model(x)
            #print('z', z.shape)
            return z

def partial_load(pretrained_dict, model, skip_keys=[], log=False):
    model_dict = model.state_dict()

    # 1. filter out unnecessary keys
    filtered_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and not any([sk in k for sk in skip_keys])}
    skipped_keys = [k for k in pretrained_dict if k not in filtered_dict]
    unload_keys = [k for k in model_dict if k not in pretrained_dict]

    # 2. overwrite entries in the existing state dict
    model_dict.update(filtered_dict)

    # 3. load the new state dict
    model.load_state_dict(model_dict)

    if log:
        print('\nSkipped keys: ', skipped_keys)
        print('\nLoading keys: ', filtered_dict.keys())
        print('\nUnLoaded keys: ', unload_keys)

def load_vince_model(path):
    checkpoint = torch.load(path, map_location={'cuda:0': 'cpu'})
    checkpoint = {k.replace('feature_extractor.module.model.', ''): checkpoint[k] for k in checkpoint if 'feature_extractor' in k}
    return checkpoint

def load_uvc_model(ckpt_path):
    net = resnet.resnet18()
    net.avgpool, net.fc = None, None

    ckpt = torch.load(ckpt_path, map_location='cpu')
    state_dict = {k.replace('module.gray_encoder.', ''):v for k,v in ckpt['state_dict'].items() if 'gray_encoder' in k}
    partial_load(state_dict, net)

    return net

def load_tc_model(ckpt_path):
    model_state = torch.load(ckpt_path, map_location='cpu')['state_dict']

    net = resnet.resnet50()
    net_state = net.state_dict()

    for k in [k for k in model_state.keys() if 'encoderVideo' in k]:
        kk = k.replace('module.encoderVideo.', '')
        tmp = model_state[k]
        if net_state[kk].shape != model_state[k].shape and net_state[kk].dim() == 4 and model_state[k].dim() == 5:
            tmp = model_state[k].squeeze(2)
        net_state[kk][:] = tmp[:]

    net.load_state_dict(net_state)

    return net

def load_dpc_model(args):
    net = dpc.DPC(args)
    net.to(torch.device('cuda'))

    if args.resume != 'none':
        net_dict = net.state_dict()

        checkpoint = torch.load(args.resume)
        pretrained_dict = checkpoint['state_dict']

        net = neq_load(net, pretrained_dict)

    return net
    
def neq_load(model, pretrained_dict, encoder_only=False):
    """ Loads pre-trained model when new model has been partially modified """
    model_dict = model.state_dict()

    keywords = ['blocks'] if encoder_only else ['blocks', 'upsampler', 'agg', 'network_pred']

    model_prefix = [None] * len(keywords)
    pretrained_prefix = [None] * len(keywords)

    for i, w in enumerate(keywords):
        for k in model_dict:
            if w in k:
                model_prefix[i] = k[:k.find(w)]
                break # stops once prefix is found

        for k in pretrained_dict:
            if ('_k' in k) | ('_c' in k) | ('upsampler_agg' in k): continue
            if w in k:
                pretrained_prefix[i] = k[:k.find(w)]
                break # stops once prefix is found
    print(model_prefix, pretrained_prefix)
    # replace none with empty strings
    model_prefix = [p if (p is not None) else '' for p in model_prefix]
    pretrained_prefix = [p if (p is not None) else '' for p in pretrained_prefix]

    tmp = {}
    print('\n=======Check Weights Loading======')
    print('Weights not used from pretrained file:')
    for k, v in pretrained_dict.items():
        if ('_k' in k) | ('upsampler_c' in k) | ('upsampler_agg' in k): 
            print(k)
            continue
        # find corresponding model key
        k_m = k
        for i, w in enumerate(keywords):
            if w in k_m:
                k_m = (model_prefix[i] + k_m[k_m.find(w):])
                if w == 'upsampler': 
                    k_m = k_m.replace('conv_', 'conv')
                    k_m = k_m.replace('norm_', 'norm')
                    k_m = k_m.replace('module.', '') 
                break # stops once associated keyword is found

        if k_m in model_dict:
            tmp[k_m] = v
        else:
            print(k)
    print('---------------------------')
    print('Weights not loaded into new model:')
    for k, v in model_dict.items():
        # find corresponding pretrained key
        if k not in tmp:
            print(k)
    print('===================================\n')
    del pretrained_dict
    model_dict.update(tmp)
    del tmp
    model.load_state_dict(model_dict)
    return model


class From3D(nn.Module):
    ''' Use a 2D convnet as a 3D convnet '''
    def __init__(self, resnet):
        super(From3D, self).__init__()
        self.model = resnet

    def forward(self, x):
        N, C, T, h, w = x.shape
        xx = x.permute(0, 2, 1, 3, 4).contiguous().view(-1, C, h, w)
        m = self.model(xx)

        return m.view(N, T, *m.shape[-3:]).permute(0, 2, 1, 3, 4)


def make_encoder(args):
    SSL_MODELS = ['byol', 'deepcluster-v2', 'infomin', 'insdis', 'moco-v1', 'moco-v2',
            'pcl-v1', 'pcl-v2','pirl', 'sela-v2', 'swav', 'simclr-v1', 'simclr-v2',
            'pixpro', 'detco', 'barlowtwins']
    model_type = args.model_type

    if model_type == 'x3d':
        net = load_dpc_model(args)

    elif model_type == 'crw':
        net = resnet.resnet18()
        if osp.isfile(args.resume):
            ckpt = torch.load(args.resume)
            state = {}
            for k, v in ckpt['model'].items():
                if 'conv1.1.weight' in k or 'conv2.1.weight' in k:
                    state[k.replace('.1.weight', '.weight')] = v
                if 'encoder.model' in k:
                    state[k.replace('encoder.model.', '')] = v
                else:
                    state[k] = v
            partial_load(state, net, skip_keys=['head',])
            del ckpt
    elif model_type == 'random18':
        net = resnet.resnet18(pretrained=False)
    elif model_type == 'random50':
        net = resnet.resnet50(pretrained=False)
    elif model_type == 'imagenet18':
        net = resnet.resnet18(pretrained=True)
    elif model_type == 'imagenet50':
        net = resnet.resnet50(pretrained=True)
    elif model_type == 'imagenet101':
        net = resnet.resnet101(pretrained=True)
    elif model_type == 'imagenet_resnext50':
        net = resnet.resnext50_32x4d(pretrained=True)
    elif model_type == 'imagenet_resnext101':
        net = resnet.resnext101_32x8d(pretrained=True)
    elif model_type == 'mocov2':
        net = resnet.resnet50(pretrained=False)
        net_ckpt = torch.load(args.resume)
        net_state = {k.replace('module.encoder_q.', ''):v for k,v in net_ckpt['state_dict'].items() \
                if 'module.encoder_q' in k}
        partial_load(net_state, net)
    elif model_type == 'ssib':
        net = resnet.resnet50(pretrained=False)
        net_ckpt = torch.load(args.resume)
        net_state = {k.replace('module.encoder.', ''):v for k,v in net_ckpt.items() \
                if 'module.encoder' in k}
        partial_load(net_state, net)
    elif model_type == 'uvc':
        net = load_uvc_model(args.resume)
    elif model_type == 'timecycle':
        net = load_tc_model(args.resume)
    elif model_type in SSL_MODELS:
        net = resnet.resnet50(pretrained=False)
        net_ckpt = torch.load(args.resume)
        partial_load(net_ckpt, net)
    elif 'hrnet' in model_type:
        net = hrnet.get_cls_net(model_type, return_stage=args.return_stage, pretrained=args.resume)
    elif model_type == 'random':
        net = random_feat_generator.RandomFeatGenerator(args)
    else:
        raise ValueError('Invalid model_type.')


    if hasattr(net, 'modify'):
        net.modify(remove_layers=args.remove_layers)

    if 'Conv2d' in str(net) and not args.infer2D:
        net = From3D(net)

    return net
