# -*- coding: utf-8 -*-
"""
Created on Wed Jun 26 00:08:07 2019

@author: zehaojin
"""
import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
import os
import matplotlib.pyplot as plt

import torch.nn as nn
import torch.utils.model_zoo as model_zoo


__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
           'resnet152']


model_urls = {
    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}


def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out

class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)

class ResNet(nn.Module):

    def __init__(self, block, layers, num_classes=1000, num_classes_=[1,56,56]):
        self.inplanes = 64
        super(ResNet, self).__init__()
        self.num_classes = num_classes
        self.num_classes_ = num_classes_

        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=1)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=1)

        number_features = 256
        kernel_size = 3
        final_conv = nn.Conv2d(number_features, self.num_classes, kernel_size=kernel_size)

        self.classifier = nn.Sequential(
            nn.Dropout(p=0.5),
            final_conv,
            nn.AvgPool2d(kernel_size=self.num_classes_[1]-(kernel_size-1),stride=1)
        )



        final_conv_ = nn.Conv2d(number_features, self.num_classes_[0], kernel_size=1)
        self.classifier_ = nn.Sequential(
            nn.Dropout(p=0.5),
            final_conv_
        )

        final_conv_recon = nn.Conv2d(number_features, self.num_classes_[0], kernel_size=1)
        self.classifier_recon = nn.Sequential(
            nn.Dropout(p=0.5),
            final_conv_recon
        )

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)

        x_macro = self.classifier(x)
        x_micro = self.classifier_(x)
        x_recon = self.classifier_recon(x)

        # DEBUG
        #print(x_macro.size())
        #print(x_micro.size())

        x_macro = x_macro.squeeze()
        x_micro = x_micro.squeeze()
        x_recon = x_recon.squeeze()


        # DEBUG
        #print(x_macro.size())
        #print(x_micro.size())
        return x_macro, x_micro, x_recon


def resnetDMS(pretrained=False, **kwargs):
    """
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
    return model

def resnetDMSV2(pretrained=False, **kwargs):
    """
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(BasicBlock, [4, 4, 4, 4], **kwargs)
    return model


def resnet18(pretrained=False, **kwargs):
    """Constructs a ResNet-18 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
    return model


def resnet34(pretrained=False, **kwargs):
    """Constructs a ResNet-34 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
    return model


def resnet50(pretrained=False, **kwargs):
    """Constructs a ResNet-50 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
    return model


def resnet101(pretrained=False, **kwargs):
    """Constructs a ResNet-101 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
    return model


def resnet152(pretrained=False, **kwargs):
    """Constructs a ResNet-152 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
    return model


def sigmoid(x):
    return 1/(1+np.exp(-x))
def Visualization(net,model_name, xdata, ydata, num_test):

    for index in range(num_test//10):
        i=index*10
        image=xdata[i]
        x=np.zeros((3, 224, 224))
        for j in range(3):
            x[j,:,:]+=image
        x=np.array([x])
        x=torch.from_numpy(x).float()
        x = Variable(x).cuda()
        output_macro,output,output_source  = net(x)

        y=ydata[i]
        print(np.sum(y))
        y=scipy.ndimage.zoom(y, 56/224, order=1)

        x=x.detach().cpu().numpy()
        x=x[0][0]
        output=output.detach().cpu().numpy()
        output = sigmoid(output)


        #print(x.shape)
        #print(output)
        #print(y.shape)
        #print(output.shape)

        #x = Image.fromarray(x).convert("RGB")




        '''
        orig_img = np.asarray(orig_img/ np.max(orig_img) * 255., dtype=np.int32)
        lens_image = np.asarray(lens_image/ np.max(lens_image) * 255., dtype=np.int32)
        smooth_model_image = np.asarray(smooth_model_image/ np.max(smooth_model_image) * 255., dtype=np.int32)
        lens_source_image = np.asarray(lens_source_image/ np.max(lens_source_image) * 255., dtype=np.int32)
        '''
        plt.figure(num=index, figsize=(12, 6), dpi=80, facecolor='w', edgecolor='k')
        plt.subplots_adjust(wspace=0.3)

        plt.subplot(131)
        plt.imshow(x)
        plt.title("galaxy map")
        plt.colorbar(fraction=0.046, pad=0.04)
        #plt.colorbar()


        plt.subplot(132)
        plt.imshow(y)
        plt.title("void map(ground truth)")
        plt.colorbar(fraction=0.046, pad=0.04)
        #plt.colorbar()


        plt.subplot(133)
        plt.imshow(output)
        plt.title("void map(prediction)")
        plt.colorbar(fraction=0.046, pad=0.04)
        #plt.colorbar()
        plt.show()

        plt.savefig('void_project/result/demofigure_'+str(index)+'.png')
        plt.close('all')


        print('#-------------------------------------------------------------------------')

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


import torch
from  torch.autograd import Variable
import matplotlib.pyplot as plt
import numpy as np
import scipy.ndimage
from PIL import Image


loaded_model_path = 'void_project/saved_model'

model_names = ['densenet_7','squeezenet_7','resnet_7']

mod = model_names[2]

if os.path.exists(loaded_model_path+mod):
    net = torch.load(loaded_model_path+mod)
    net.cuda()
    print('loaded {}!'.format(mod))

net.eval()
num_test = 2000


print("Total number of trainable parameters: ", count_parameters(net))
xdata=np.load('void_project/data/xTest.npy')
ydata=np.load('void_project/data/yTest.npy')
Visualization(net, mod, xdata, ydata, num_test)
del net
