#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2020-03-15 19:00 qiang.zhou <theodoruszq@gmail.com>
#
# Distributed under terms of the CC-NC license.

# from dataset import load_processed_ct_images
from dataset import combine_ct_images_and_masks
import cv2
import sys, os
from importlib import import_module
import torch.nn.functional as F
import torch
import time
from scipy.ndimage import zoom
import numpy as np
import logging
from datetime import datetime
import time

SEG_MODEL_UID = "unet"
CLS_MODEL_UID = "covid19_3dcls"

SEG_PRETRAINED_MODEL_PATH = "./pretrained_models/unet.pth"
CLS_PRETRAINED_MODEL_PATH = "./pretrained_models/ncov-Epoch_00140-auc95p9.pth"

SEG_BATCH_NUM = 1           # Can be larger with powerful CPU
SEG_CLIP_RANGE = (0.0, 1.0)
CROP_CLIP_RANGE = (0.0, 1.0)
CROP_OUTPUT_SIZE = (224, 336)

LOG_DIRPATH = "logs"
os.makedirs(LOG_DIRPATH, exist_ok=True)

logger = logging.getLogger()

def load_processed_ct_images(npy_filepath, patient, clip_range):
    files = os.listdir(npy_filepath+patient)
    files.sort()
    images = []
    for file in files:
        images.append(cv2.imread(npy_filepath+patient+"/"+file,cv2.COLOR_BGR2GRAY))
    images = np.array(images)
    num_frames = len(images)
    global left
    global origin_images
    origin_images = images.copy()

    left, right = int(num_frames*clip_range[0]), int(num_frames*clip_range[1])
    images = images[left:right]

    num_frames = len(images)
    shape = images.shape

    if False:
        from zqlib import imgs2vid
        imgs2vid(images, "test_image.avi")
    
    images = np.asarray(images, dtype=np.float32)
    images = images / 255.
    # add channel dimension
    images = np.expand_dims(images, axis=1)
    img_info = {"name": npy_filepath, 
                "num_frames": num_frames,
                "clip_range": clip_range,
                "shape": shape}
    th_images = torch.from_numpy(images.copy()).float()
    return th_images, img_info

### Compute the lung mask as the intermediate results ###
def get_ct_images_and_masks(npy_filepath, patient, seg_clip_range, 
                            batch_num=1, socket_logger=None):
    ct_images, ct_info = load_processed_ct_images(
                            npy_filepath, patient, seg_clip_range)
    # ct_masks = []
    # for i in range(0, len(ct_images), batch_num):
    #     if socket_logger: 
    #         socket_logger("Segmenting {}-{}".format(i, i+batch_num))
    #     ct_image = ct_images[i:i+batch_num]
    #     ct_mask = model(ct_image)
    #     ct_mask = torch.argmax(F.softmax(ct_mask, dim=1), dim=1)
    #     ct_masks.append(ct_mask)
    # ct_np_masks = torch.cat(ct_masks, dim=0).cpu().numpy().astype('uint8')

    # load the pregenerated mask npy files
    ct_np_masks = np.load("./mask_npy/"+patient+".npy")
    masks_sum = ct_np_masks.reshape(ct_np_masks.shape[0],-1).sum(axis=1)
    masks_thres = masks_sum.max()*0.4
    ignore_slices = int(0.2*len(masks_sum))
    max_slide = masks_sum[ignore_slices:-ignore_slices].argmax()+ignore_slices
    global mask_left,mask_right
    mask_left = max_slide - np.where((masks_sum[0:max_slide][::-1]>masks_thres)==0)[0].min()
    if (masks_sum[max_slide:]>masks_thres).all():
        mask_right = 1
    else:
        mask_right = ct_np_masks.shape[0]-np.where((masks_sum[max_slide:]>masks_thres)==0)[0].min()-max_slide
    ct_np_images = np.uint8(ct_images[:, 0].cpu().numpy().astype('float32') * 255)
    return ct_np_images, ct_np_masks

crop_box = None
ori_images = None

def crop_and_resize(ct_images, ct_masks, patient, output_size, crop_clip_range):
    num_frame = len(ct_images)
    left, right = int(num_frame*crop_clip_range[0]), int(num_frame*crop_clip_range[1])
    ct_images, ct_masks = ct_images[left:right], ct_masks[left:right]

    ori_images = ct_images.copy()
    global zz,yy,xx
    zz, yy, xx = np.where(ct_masks > 0)
    cropbox = np.array([[np.min(zz), np.max(zz)], 
                        [np.min(yy), np.max(yy)], 
                        [np.min(xx), np.max(xx)]])
    global supplement_info,ct_zoom_images,ct_zoom_masks
    # load the 3D connected components generated by 3DCC
    supplement_info = np.load("./binary-seg/"+patient+".npy")
    supplement_info = supplement_info[cropbox[0, 0]:cropbox[0, 1],
                               cropbox[1, 0]:cropbox[1, 1],
                               cropbox[2, 0]:cropbox[2, 1]]
    ct_crop_images = ct_images[cropbox[0, 0]:cropbox[0, 1],
                               cropbox[1, 0]:cropbox[1, 1],
                               cropbox[2, 0]:cropbox[2, 1]]
    ct_crop_masks = ct_masks[cropbox[0, 0]:cropbox[0, 1],
                             cropbox[1, 0]:cropbox[1, 1],
                             cropbox[2, 0]:cropbox[2, 1]]
    
    height, width = ct_crop_images.shape[1:3]
    new_height, new_width = output_size
    ct_zoom_images = zoom(ct_crop_images, 
                          (1, new_height/height, new_width/width))
    ct_zoom_masks = zoom(ct_crop_masks,
                          (1, new_height/height, new_width/width))
    supplement_info = zoom(supplement_info,
                          (1, new_height/height, new_width/width))

    global crop_box 
    # (minx, miny, maxx, maxy)
    crop_box = [cropbox[2, 0], cropbox[1, 0], cropbox[2, 1], cropbox[1, 1]]
    # if True:
    #     logger.debug("Dumping zoomed images and masks to {}...\
    #                  ".format("im_mask.avi"))
    #     from zqlib import imgs2vid
    #     img_mask = np.concatenate([ct_zoom_images, ct_zoom_masks*255], axis=2)
    #     imgs2vid(img_mask, "im_mask.avi")

    return ct_zoom_images, ct_zoom_masks

def cam_mask(ct_images, ct_masks, model, filepath, uniq_id):
    global origin_images
    global crop_box
    global zz,yy,xx
    global left
    global supplement_info,ct_zoom_images
    from zqlib import imgs2vid
    ct_immasks = combine_ct_images_and_masks(ct_images, ct_masks)
    preds,features = model(ct_immasks.cuda())
    # caculate cam
    pool = torch.nn.AdaptiveAvgPool3d(output_size=(1,1,1))
    weight1 = pool(model.module.head[1].weight).squeeze()
    weight2 = pool(model.module.head[4].weight).squeeze()
    weight3 = pool(model.module.head[8].weight).squeeze()
    weight4 = model.module.classifier[0].weight
    weight5 = model.module.classifier[1].weight
    weight = torch.mm(weight1.t(), weight2.t())
    weight = torch.mm(weight,weight3.t())
    weight = torch.mm(weight,weight4.t())
    weight = torch.mm(weight,weight5.t())
    final_weight = weight[:,1].unsqueeze(0)
    features = features.squeeze()
    channel,t,h,w = features.shape
    features = features.reshape(channel,t*h*w)
    cam = torch.mm(final_weight,features).reshape(t,h,w)
    cam = cam.detach().cpu().numpy()
    ori_h, ori_w = crop_box[3]-crop_box[1], crop_box[2]-crop_box[0]

    from scipy.ndimage import zoom
    import cv2
    slice,H,W = ct_images.shape
    cam = zoom(cam,(slice/t,H/h,W/w), order=1)
    cam = (cam>0)*cam
    cam = cam - np.min(cam)
    cam = cam / np.max(cam)

    num,thresholds = np.histogram(cam,bins=50)
    temp = 0
    for x in range(len(thresholds)):
        temp = temp + num[len(thresholds)-2-x]
        if temp > cam.shape[0]*cam.shape[1]*cam.shape[2]*0.05: 
            threshold = thresholds[len(thresholds)-1-x]
            break
    cam = np.round(cam*255).astype(np.uint8)
    fisrt_cam = cam.copy()
    cam = np.uint8(cam > threshold*255)
    cam  = np.where(ct_zoom_images>75,cam,0)

    global mask_left,mask_right
    import cc3d
    from skimage import measure
    labels_out = measure.label(cam[mask_left:-mask_right])
    vals,counts = np.unique(labels_out, return_counts=True)
    ids_ = counts.argsort()[-10:-1]
    extracted_images = []
    connected_comp_areas = []
    for segid in ids_:
        extracted_image = labels_out * (labels_out == vals[segid])*supplement_info[mask_left:-mask_right]
        connected_comp_areas.append((extracted_image).sum())
        extracted_images.append(extracted_image)
    temp = np.zeros(cam.shape)
    connected_comp_ids = np.argsort(connected_comp_areas)[-1]

    temp [mask_left:-mask_right]= temp[mask_left:-mask_right] + extracted_images[connected_comp_ids]
    cam = temp
    del extracted_images,labels_out


    cam = zoom(cam, (1, ori_h/H, ori_w/W),order=0)
    ori_cam = np.zeros((origin_images.shape[0], 512, 512), dtype=np.uint8)
    ori_cam[left+zz.min():left+zz.max(), crop_box[1]:crop_box[3], crop_box[0]:crop_box[2]] = cam
    cam = ori_cam.astype(np.uint8)

    origin_images = origin_images[:,:,:,np.newaxis]
    origin_images = np.concatenate((origin_images,origin_images,origin_images), axis=3)
    origin_images[:,:,:,2] = np.where(cam>0,255*np.ones(cam.shape),origin_images[:,:,:,2])
    origin_images = np.clip(origin_images, a_min=0, a_max=255)
    origin_images = np.uint8(origin_images)
    from zqlib import imgs2vid
    imgs2vid(origin_images, "visual-cam/"+uniq_id+".avi")
    np.save("cam-mask/"+uniq_id+".npy",cam)

############# Main Control #############
def cv_main(filepath, patient):
    base_filepath = filepath
    images_dir = filepath
    
    # Build models
    # model = import_module(f"model.{SEG_MODEL_UID}")
    # UNet = getattr(model, "UNet")
    # seg_model = UNet(n_channels=1, n_classes=2)
    # seg_model = torch.nn.DataParallel(seg_model)
    # seg_model.load_state_dict(
    #         torch.load(SEG_PRETRAINED_MODEL_PATH, map_location="cpu"))

    model = import_module(f"model.{CLS_MODEL_UID}")
    ENModel = getattr(model, "ENModel")
    cls_model = ENModel(n_channels=2, n_classes=2)
    cls_model = torch.nn.DataParallel(cls_model)
    cls_model.load_state_dict(
            torch.load(CLS_PRETRAINED_MODEL_PATH, map_location="cuda"))
    cls_model.cuda()

    # Get lung images and lung masks
    ct_images, ct_masks = get_ct_images_and_masks(
                            npy_filepath=filepath, 
                            patient=patient,
                            # model= seg_model,
                            seg_clip_range=SEG_CLIP_RANGE,
                            batch_num=SEG_BATCH_NUM)    
    # Crop and resize
    ct_images, ct_masks = crop_and_resize(ct_images, ct_masks,
                                          patient=patient,
                                          output_size=CROP_OUTPUT_SIZE,
                                          crop_clip_range=CROP_CLIP_RANGE)

    # generated CAM MASK
    cam_mask(ct_images, ct_masks, cls_model, filepath, patient)

if __name__ == "__main__":
    list_path = "/your/path/test.txt"
    img_ids = [i_id.strip() for i_id in open(list_path)]
    for img_id in img_ids:
        FILEPATH = "/your/path/JPEGImages/"
        cv_main(FILEPATH, img_id)


















