#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Tue Aug 28 10:23:47 2018

@author: tranthan
"""

import matplotlib.pyplot as plt
from scipy import misc
import numpy as np
import threading
import random
from keras.utils import to_categorical
from keras.callbacks import Callback as Callback
from keras import regularizers
from keras.applications.resnet50 import preprocess_input
from skimage import measure as Measure
import tensorflow as tf
from joblib import Parallel, delayed
from tqdm import tqdm
from scipy import ndimage as ndi
from skimage.morphology import watershed
from skimage.feature import peak_local_max

"""
This file contains most of the needed functions
"""

"""
Settings for sliding window and patch extraction
"""
FLOWER_INPUT_DIM = (64,64,3)
FLOWER_STRIDE = 5

"""
A function to save the best weights of a network
"""
class CustomCallback(Callback):
    """ callback to log best model parameter based on validation performance """

    def on_train_begin(self, logs={}): #set weights and monitor_quantity to zero at beginning
        self.model_weights = None
        self.monitor_quantity = 0

        return

    def on_epoch_end(self, epoch, logs={}):
        # get best model according to validation accuracy
        if logs['val_acc'] > self.monitor_quantity: #if the val_acc after THIS epoch is greater than previous monitoring_quantity:
            self.model_weights = self.model.get_weights() #save the weights after THIS epoch as optimal
            self.monitor_quantity = logs['val_acc']

        return

"""
Functions needed for generators
"""
class threadsafe_iter:
    """Takes an iterator/generator and makes it thread-safe by
    serializing call to the `next` method of given iterator/generator.
    """
    def __init__(self, it):
        self.it = it
        self.lock = threading.Lock()

    def __iter__(self):
        return self

    def __next__(self):
        with self.lock:
            return next(self.it)
        #   return self.it.next()

def threadsafe_generator(f):
    """A decorator that takes a generator function and makes it thread-safe.
    """
    def g(*a, **kw):
        return threadsafe_iter(f(*a, **kw))
    return g

"""
Function to read in the annotation.txt file
and creates arrays of them
"""
#5 class version
def process_annotation(annotation):
    """
    annotation: txt file with index, xmin, xmax, ymin, ymax, label
    return list of [index, xmin, xmax, ymin, ymax, label]
    """
    fid = open(annotation,'r')
    content = fid.read().split('\n')[:-1]
    fid.close()
    output = []
    labels = ['B','F','W','S','X']
    labels.sort()

    for row in content:
        index, xmin, xmax, ymin, ymax, label = row.split(',')
        output.append([int(index), int(xmin),int(xmax),int(ymin),int(ymax), labels.index(label)])

    return output

#2 class version
def process_annotation_2class(annotation):
    """
    annotation: txt file with index, xmin, xmax, ymin, ymax, label
    return list of [index, xmin, xmax, ymin, ymax, label]
    """
    fid = open(annotation,'r')
    content = fid.read().split('\n')[:-1]
    fid.close()
    output = []
    labels = ['F','X']
    labels.sort()

    for row in content:
        index, xmin, xmax, ymin, ymax, label = row.split(',')
        output.append([int(index), int(xmin),int(xmax),int(ymin),int(ymax), labels.index(label)])

    return output

def patch_flip(im):
    if random.choice([True,False]):
        im = np.flip(im, axis=1)

    return im

"""
Function to extract annotated bounding boxes from the frames
If augmentation=True, will perform random modifications of the box
"""
def extract_patch(images, annotation, target_shape, scale=0.1, augmentation=False):
    im = images[annotation[0]] #takes an image with the id in annotation
    H,W,C = im.shape #height, width, channels

    xmin, xmax, ymin, ymax = annotation[1:5]

    if augmentation:
        top_shift = int(np.random.uniform((xmin-xmax)*scale, (xmax-xmin)*scale))
        bottom_shift = int(np.random.uniform((xmin-xmax)*scale, (xmax-xmin)*scale))
        left_shift = int(np.random.uniform((ymin-ymax)*scale, (ymax-ymin)*scale))
        right_shift = int(np.random.uniform((ymin-ymax)*scale, (ymax-ymin)*scale))
        xmin = max(0, xmin+top_shift)
        xmax = min(H, xmax + bottom_shift)
        ymin = max(0, ymin + left_shift)
        ymax = min(W, ymax + right_shift)

    min_idx = min(xmin, ymin) #minimum of xmin, ymin
    max_idx = min(H-xmax, W-ymax) #maximum of xmax, ymax

    pad = max((min_idx<0)*(-min_idx), (max_idx < 0)*(-max_idx)) #checks how much outside the image the bounding box of the annotation is
    im = np.pad(im, ((pad,pad),(pad,pad),(0,0)),'constant') #zero padding the image
    patch = im[xmin+pad:xmax+pad,ymin+pad:ymax+pad,:] #takes the bb plus the padding
    patch = misc.imresize(patch, target_shape, interp='lanczos')

    if augmentation:
        patch = patch_flip(patch)
        scale=np.random.uniform(0.7,1.0)
        mini=0.5-0.5*scale
        maxi=0.5+0.5*scale
        xmin=ymin=int(mini*target_shape[0])
        xmax=ymax=int(maxi*target_shape[1])
        patchnew=patch[xmin:xmax,ymin:ymax,:]
        patch=misc.imresize(patchnew, target_shape, interp='lanczos')

    return patch

"""
Patch generator that does preprocessing, i.e. standardization based on mean and sd of
imagenet data base

Use with resnet50 model and all models pretrained on imagenet
"""
def get_patch_generator_imagenet(images, annotation, target_shape, batch_size=128, augmentation=False):
    """
    images: npy file NxHxWx3
    annotation: path to annotation text file with index, xmin, xmax, ymin, ymax, label
    target_shape: shape of the resized patch
    augmentation: flip, shifting
    """
    annotation = process_annotation(annotation) #outputs the annotations as a list
    steps = int(np.ceil(len(annotation)/float(batch_size))) #how many batches of annotated obs there are

    @threadsafe_generator
    def generator(annotation_list):
        while True:
            random.shuffle(annotation_list) #randomize the order
            for iteration in range(steps): #does this for each batch
                start_idx = iteration*batch_size
                stop_idx = min(len(annotation_list), (iteration+1)*batch_size)
                x = np.zeros((stop_idx-start_idx, target_shape[0], target_shape[1],3), dtype=np.float32) #makes a list for the patches in the batch
                y = np.zeros((stop_idx-start_idx, 5), dtype=np.float32) #makes a list for the labels/classes in the batch

                for sample_idx, idx in enumerate(range(start_idx, stop_idx)):
                    x[sample_idx] = extract_patch(images, annotation_list[idx], target_shape, augmentation) #takes the bounding boxes and turns them into patches
                    y[sample_idx] = to_categorical(annotation_list[idx][-1],5) #saves the labels of the patches

                yield preprocess_input(x), y #performs normalization etc. to the patches? yield is like return for generators.

    return generator(annotation), steps #returns the generator for the patches and the amount of batches (in one epoch)

def get_patch_generator_imagenet_2class(images, annotation, target_shape, batch_size=128, augmentation=False):
    """
    images: npy file NxHxWx3
    annotation: path to annotation text file with index, xmin, xmax, ymin, ymax, label
    target_shape: shape of the resized patch
    augmentation: flip, shifting
    """
    annotation = process_annotation_2class(annotation) #outputs the annotations as a list
    steps = int(np.ceil(len(annotation)/float(batch_size))) #how many batches of annotated obs there are

    @threadsafe_generator
    def generator(annotation_list):
        while True:
            random.shuffle(annotation_list) #randomize the order
            for iteration in range(steps): #does this for each batch
                start_idx = iteration*batch_size
                stop_idx = min(len(annotation_list), (iteration+1)*batch_size)
                x = np.zeros((stop_idx-start_idx, target_shape[0], target_shape[1],3), dtype=np.float32) #makes a list for the patches in the batch
                y = np.zeros((stop_idx-start_idx, 2), dtype=np.float32) #makes a list for the labels/classes in the batch

                for sample_idx, idx in enumerate(range(start_idx, stop_idx)):
                    x[sample_idx] = extract_patch(images, annotation_list[idx], target_shape, augmentation) #takes the bounding boxes and turns them into patches
                    y[sample_idx] = to_categorical(annotation_list[idx][-1],5) #saves the labels of the patches

                yield preprocess_input(x), y #performs normalization etc. to the patches? yield is like return for generators.

    return generator(annotation), steps #returns the generator for the patches and the amount of batches (in one epoch)

def get_patch_generator(images, annotation, target_shape, batch_size=128, scale=0.1, augmentation=False):
    """
    images: npy file NxHxWx3
    annotation: path to annotation text file with index, xmin, xmax, ymin, ymax, label
    target_shape: shape of the resized patch
    augmentation: flip, shifting
    """
    annotation = process_annotation(annotation) #outputs the annotations as a list
    steps = int(np.ceil(len(annotation)/float(batch_size))) #how many batches of annotated obs there are

    @threadsafe_generator
    def generator(annotation_list):
        while True:
            random.shuffle(annotation_list) #randomize the order
            for iteration in range(steps): #does this for each batch
                start_idx = iteration*batch_size
                stop_idx = min(len(annotation_list), (iteration+1)*batch_size)
                x = np.zeros((stop_idx-start_idx, target_shape[0], target_shape[1],3), dtype=np.float32) #makes a list for the patches in the batch
                y = np.zeros((stop_idx-start_idx, 5), dtype=np.float32) #makes a list for the labels/classes in the batch

                for sample_idx, idx in enumerate(range(start_idx, stop_idx)):
                    x[sample_idx] = extract_patch(images, annotation_list[idx], target_shape, scale, augmentation) #takes the bounding boxes and turns them into patches
                    y[sample_idx] = to_categorical(annotation_list[idx][-1],2) #saves the labels of the patches

                yield x/255.0, y #performs normalization etc. to the patches? yield is like return for generators.

    return generator(annotation), steps #returns the generator for the patches and the amount of batches (in one epoch)

def get_patch_generator_2class(images, annotation, target_shape, batch_size=128, scale=0.1, augmentation=False):
    """
    images: npy file NxHxWx3
    annotation: path to annotation text file with index, xmin, xmax, ymin, ymax, label
    target_shape: shape of the resized patch
    augmentation: flip, shifting
    """
    annotation = process_annotation_2class(annotation) #outputs the annotations as a list
    steps = int(np.ceil(len(annotation)/float(batch_size))) #how many batches of annotated obs there are

    @threadsafe_generator
    def generator(annotation_list):
        while True:
            random.shuffle(annotation_list) #randomize the order
            for iteration in range(steps): #does this for each batch
                start_idx = iteration*batch_size
                stop_idx = min(len(annotation_list), (iteration+1)*batch_size)
                x = np.zeros((stop_idx-start_idx, target_shape[0], target_shape[1],3), dtype=np.float32) #makes a list for the patches in the batch
                y = np.zeros((stop_idx-start_idx, 2), dtype=np.float32) #makes a list for the labels/classes in the batch

                for sample_idx, idx in enumerate(range(start_idx, stop_idx)):
                    x[sample_idx] = extract_patch(images, annotation_list[idx], target_shape, scale, augmentation) #takes the bounding boxes and turns them into patches
                    y[sample_idx] = to_categorical(annotation_list[idx][-1],2) #saves the labels of the patches

                yield x/255.0, y #performs normalization etc. to the patches? yield is like return for generators.

    return generator(annotation), steps #returns the generator for the patches and the amount of batches (in one epoch)

def merge_history(old, new):
    if len(old.keys()) > 0: #only do this if the old dictionary has something in it
        keys = new.keys()
        for k in keys:
            old[k] += new[k] #add the new to the old
        return old
    else:
        return new

"""
Create sliding window patches for test frames
"""
def get_test_generator_imagenet(image, stride, target_shape, window):
    """
    image: single image matrix of size H x W x 3
    stride: number of pixels to step
    window: [h,w] window size of hxw
    """

    H,W,C = image.shape
    h_pad = int(window[0] / 2)
    w_pad = int(window[1] / 2)

    H += 2*h_pad
    W += 2*w_pad

    h_steps = int((H-window[0])/stride)
    w_steps = int((W-window[1])/stride)

    image = np.pad(image, ((h_pad,h_pad), (w_pad, w_pad), (0,0)),'constant')
    @threadsafe_generator
    def generator():
        while True:
            for h_iter in range(h_steps):
                for w_iter in range(w_steps):
                    xmin = h_iter*stride
                    xmax = xmin + window[0]
                    ymin = w_iter*stride
                    ymax = ymin + window[1]
                    x = image[xmin:xmax, ymin:ymax,:]
                    x = misc.imresize(x, target_shape, interp='lanczos')
                    x = preprocess_input(x)
                    x.shape=[1,target_shape[0],target_shape[1],3]

                    #returns the preprocessed (=normalized) patch
                    yield x

    return generator(), h_steps, w_steps

def get_test_generator(image, stride, target_shape, window):
    """
    image: single image matrix of size H x W x 3
    stride: number of pixels to step
    window: [h,w] window size of hxw
    """

    H,W,C = image.shape
    h_pad = int(window[0] / 2)
    w_pad = int(window[1] / 2)

    H += 2*h_pad
    W += 2*w_pad

    h_steps = int((H-window[0])/stride)
    w_steps = int((W-window[1])/stride)

    image = np.pad(image, ((h_pad,h_pad), (w_pad, w_pad), (0,0)),'constant')
    @threadsafe_generator
    def generator():
        while True:
            for h_iter in range(h_steps):
                for w_iter in range(w_steps):
                    xmin = h_iter*stride
                    xmax = xmin + window[0]
                    ymin = w_iter*stride
                    ymax = ymin + window[1]
                    x = image[xmin:xmax, ymin:ymax,:]
                    x = misc.imresize(x, target_shape, interp='lanczos')
                    x = x/255.0
                    x.shape=[1,target_shape[0],target_shape[1],3]

                    #returns the patch
                    yield x

    return generator(), h_steps, w_steps

"""
Function to draw bounding boxes based on the probability map
predicted by a model

Uses non-maximum suppression to leave out overlapping boxes
"""
def generate_box_prediction_nms(class_prediction, lb, threshold = 0.5, extension_percentage=0.05, minbox=0, buffer=[0,0,0,0]):

    bbox = []

    # get the probability map of this label
    lb_prob_map = class_prediction[lb]
    # threshold the probability map
    mask = (lb_prob_map > threshold)*1 #Keep only the probabilities over 50 % #changing all the remaining probabilities to 1 so that ANY connected regions will have the same label

    H,W = mask.shape

    img_labels = Measure.label(mask, background=0) #labels connected regions
    unique_labels = np.unique(img_labels)[1:] #counts the number of these connected regions = objects

    #This loop creates the coordinates for bounding boxes around the connected regions
    for idx, lbs in enumerate(unique_labels):
        x,y = np.where(img_labels==lbs)
        x_max = np.max(x)
        x_min = np.min(x)
        y_max = np.max(y)
        y_min = np.min(y)
        x_extension = int((x_max-x_min)*extension_percentage)
        y_extension = int((y_max-y_min)*extension_percentage)

        x_max = min(x_max + x_extension,H)
        x_min = max(x_min - x_extension,0)
        y_max = min(y_max + y_extension,W)
        y_min = max(y_min - y_extension,0)

        #This checks that the predicted bounding box is over some minimum
        #Checks that bounding box center is not outside buffer
        area = (x_max-x_min+1)*(y_max-y_min+1)
        x_center=(x_min+x_max)/2
        y_center=(y_min+y_max)/2

        if x_max > x_min and y_max > y_min and area>=(0.9*minbox) and buffer[0]<x_center<buffer[1] and buffer[2]<y_center<buffer[3]:
            bbox.append([[int(x_min), int(x_max), int(y_min), int(y_max)], lb,lbs])

    bboxes = non_max_suppression_fast(bbox,0.5)

    return bboxes

def non_max_suppression_fast(boxes, overlapThresh):
    # if there are no boxes, return an empty list
    if len(boxes) == 0:
        return []

	# if the bounding boxes integers, convert them to floats --
	# this is important since we'll be doing a bunch of divisions
    bbs = np.matrix([x[0] for x in boxes])
    if bbs.dtype.kind == "i":
        bbs = bbs.astype("float")

    # initialize the list of picked indexes
    pick = []

    # grab the coordinates of the bounding boxes
    x1 = np.ravel(np.array(bbs[:,0]))
    y1 = np.ravel(np.array(bbs[:,2]))
    x2 = np.ravel(np.array(bbs[:,1]))
    y2 = np.ravel(np.array(bbs[:,3]))

    # compute the area of the bounding boxes and sort the bounding
    # boxes by the bottom-right y-coordinate of the bounding box
    area = (x2 - x1 + 1) * (y2 - y1 + 1)
    idxs = np.argsort(y2)

    # keep looping while some indexes still remain in the indexes
    # list
    while len(idxs) > 0:
        #grab the last index in the indexes list and add the
        #index value to the list of picked indexes
        last = len(idxs) - 1
        i = idxs[last]
        pick.append(i)

        # find the largest (x, y) coordinates for the start of
        # the bounding box and the smallest (x, y) coordinates
        # for the end of the bounding box
        xx1 = np.maximum(x1[i], x1[idxs[:last]])
        yy1 = np.maximum(y1[i], y1[idxs[:last]])
        xx2 = np.minimum(x2[i], x2[idxs[:last]])
        yy2 = np.minimum(y2[i], y2[idxs[:last]])

        # compute the width and height of the bounding box
        w = np.maximum(0, xx2 - xx1 + 1)
        h = np.maximum(0, yy2 - yy1 + 1)

        # compute the ratio of overlap
        overlap = (w * h) / area[idxs[:last]]

        # delete all indexes from the index list that have
        idxs = np.delete(idxs, np.concatenate(([last], np.where(overlap > overlapThresh)[0])))

    # return only the bounding boxes that were picked using the
    # integer data type
    return [boxes[i] for i in pick]

"""
Calculates average precision
"""
def evaluate_ap_score(annotation, class_prediction, prob_thresholds, iou_thresholds,min_box=0, buffer=[0,0,0,0]):
    labels = [f for f in class_prediction.keys() if f!= 'X']
    """
    AP is a dictionary with key is the class name
    e.g., AP['F'] is the AP vector the class 'F', with AP[0] corresponds to the 1st iou threshold, AP[1] corresponds to the 2nd iou threshold...
    """
    AP = {}

    for lb in labels:
        AP[lb] = []
        P = np.zeros((len(iou_thresholds),len(prob_thresholds)))
        R = np.zeros((len(iou_thresholds),len(prob_thresholds)))
        for idj, prob in enumerate(prob_thresholds):
            print('generating predicted boxes')
            #Uses buffer to leave out edge predictions, nms to leave out overlapping predictions
            #and minbox to leave out anything smaller than that
            predicted_boxes = generate_box_prediction_nms(class_prediction, lb, threshold=prob, minbox=min_box, buffer=buffer)
            print('calculating iou scores')
            iou_scores = generate_iou(predicted_boxes, annotation, lb)
            for idi, iou in enumerate(iou_thresholds):
                precision, recall = evaluate_precision_recall(predicted_boxes, annotation, iou_scores, iou)
                P[idi][idj] = precision
                R[idi][idj] = recall
        for idi, iou in enumerate(iou_thresholds):
            AP[lb].append(evaluate_average_precision(P[idi],R[idi]))

    return AP

def generate_iou(predicted_boxes, annotation, lb):
    """
    predicted_boxes: [[xmin,xmax, ymin, ymax, lb, id],.....]
    annotation: [[xmin,xmax,ymin,ymax,lb,id],....]
    lb: label of the class
    iou: IOU score
    """
    annotation = [box for box in annotation if box[-1] == lb or box[-1] == 'NA'] # <- take into account NA class

    iou_mat = np.zeros((len(predicted_boxes),len(annotation)))
    for idi, pred_box in enumerate(predicted_boxes):
        for idj, gt_box in enumerate(annotation):
            iou_mat[idi][idj]=IOU(pred_box[0],gt_box[0])

    return iou_mat

def generate_iom(predicted_boxes, annotation, lb):
    """
    predicted_boxes: [[xmin,xmax, ymin, ymax, lb, id],.....]
    annotation: [[xmin,xmax,ymin,ymax,lb,id],....]
    lb: label of the class
    iou: IOU score
    """
    annotation = [box for box in annotation if box[-1] == lb or box[-1] == 'NA'] # <- take into account NA class

    iom_mat = np.zeros((len(predicted_boxes),len(annotation)))
    for idi, pred_box in enumerate(predicted_boxes):
        for idj, gt_box in enumerate(annotation):
            iom_mat[idi][idj]=IOM(pred_box[0],gt_box[0])

    return iom_mat

def IOU(bb1, bb2):

    """ intersection over union """
    x1_min, x1_max, y1_min, y1_max = bb1
    x2_min, x2_max, y2_min, y2_max = bb2

    vertical_overlap = 0
    horizontal_overlap = 0
    if x1_min<=x2_max<=x1_max or x2_min<=x1_max<=x2_max:
        sort = [x1_min,x2_min,x1_max,x2_max]
        sort.sort()
        vertical_overlap = np.abs(sort[2] - sort[1])
    if y1_min<=y2_max<=y1_max or y2_min<=y1_max<=y2_max:
        sort = [y1_min,y2_min,y1_max,y2_max]
        sort.sort()
        horizontal_overlap = np.abs(sort[2]-sort[1])

    overlap_area = horizontal_overlap * vertical_overlap
    union_area = (x1_max - x1_min)*(y1_max-y1_min) + (x2_max - x2_min)*(y2_max-y2_min) - overlap_area

    return overlap_area/float(union_area)

def IOM(bb1, bb2):

    """ intersection over minimum """
    x1_min, x1_max, y1_min, y1_max = bb1
    x2_min, x2_max, y2_min, y2_max = bb2

    vertical_overlap = 0
    horizontal_overlap = 0
    if x1_min<=x2_max<=x1_max or x2_min<=x1_max<=x2_max:
        sort = [x1_min,x2_min,x1_max,x2_max]
        sort.sort()
        vertical_overlap = np.abs(sort[2] - sort[1])
    if y1_min<=y2_max<=y1_max or y2_min<=y1_max<=y2_max:
        sort = [y1_min,y2_min,y1_max,y2_max]
        sort.sort()
        horizontal_overlap = np.abs(sort[2]-sort[1])

    overlap_area = horizontal_overlap * vertical_overlap
    min_area = min((x1_max - x1_min)*(y1_max-y1_min),(x2_max - x2_min)*(y2_max-y2_min))

    return overlap_area/float(min_area)

def evaluate_precision_recall(predicted_boxes, annotation, iou_matrix, iou):
    TP = 0
    no_predicted = len(predicted_boxes)
    no_objects = len([f for f in annotation if f[-1] != 'NA'])

    for idi, pred_box in enumerate(predicted_boxes):
        for idj, gt_box in enumerate(annotation):
            if iou_matrix[idi][idj] >= iou:
                if pred_box[-2] == gt_box[-1] or gt_box[-1] == 'NA':
                    TP += 1
                    if gt_box[-1] == 'NA':
                        no_objects +=1
    if no_predicted == 0:
        precision = 0
    else:
        precision = TP / float(no_predicted)
    if no_objects == 0:
        recall = 0
    else:
        recall = TP / float(no_objects)

    return precision, recall

def evaluate_average_precision(precisions, recalls):
    AP = 0
    for index in range(1,len(precisions)):
        if precisions[index]>0:
            AP += (recalls[index]  + recalls[index-1]) / precisions[index]

    return AP

"""
Function to evaluate detection results
"""
def evaluate(bbox, bbox_gt):
    # bbox : Nx4 numpy array
    # bbox_gt list of 4-tupple
    correct_prediction = np.zeros((len(bbox_gt),))
    false_positive = np.zeros((len(bbox),))
    for bb_index, bb in enumerate(bbox):
        prediction_matched = False
        for bb_gt_index, bb_gt in enumerate(bbox_gt):
            if bb_gt[0][1] <= bb_gt[0][0] or bb_gt[0][3] <= bb_gt[0][2]:
                print('ground truth wrong')
            if IOM(bb[0], bb_gt[0]) > 0.5:
                correct_prediction[bb_gt_index] = 1
                prediction_matched = True
        if not prediction_matched:
            false_positive[bb_index] = 1
    no_detection = np.sum(correct_prediction)
    no_misdetection = correct_prediction.size - no_detection
    no_false_positive = np.sum(false_positive)
    no_objects = len(bbox_gt)

    return no_detection, no_misdetection, no_false_positive, no_objects

"""
Function for plotting AND saving the images with
annotated boxes and predicted boxes
and the buffer
"""
def im_with_boxes(boxes, predic, image, path, thickness=5, buffer=[0,0,0,0]):
    im = np.copy(image) #takes a copy so doesn't change the original
    H, W, C = im.shape #height, width, channels
    max_intensity = 1.0
    if np.max(im.flatten()) > 1.0: #checks all the pixels to be more than 1.0
        max_intensity = 255

    for box in boxes:
        xmin,xmax,ymin,ymax = box[0]
        im[xmin:xmax, ymin:ymin+thickness,:] = [max_intensity,0,0] #draws a box around the coordinates
        im[xmin:xmax, ymax-thickness:ymax,:] = [max_intensity,0,0]
        im[xmin:xmin+thickness, ymin:ymax,:] = [max_intensity,0,0]
        im[xmax-thickness:xmax, ymin:ymax,:] = [max_intensity,0,0]

    for pred in predic:
        #xmin,xmax,ymin,ymax=pred[0] #This one is for using my generate_box_prediction
        xmin,xmax,ymin,ymax=pred[0] #This one is for using Dat's generate_box_prediction
        im[xmin:xmax, ymin:ymin+thickness,:] = [0,max_intensity,0] #draws a box around the coordinates
        im[xmin:xmax, ymax-thickness:ymax,:] = [0,max_intensity,0]
        im[xmin:xmin+thickness, ymin:ymax,:] = [0,max_intensity,0]
        im[xmax-thickness:xmax, ymin:ymax,:] = [0,max_intensity,0]

    xmin,xmax,ymin,ymax = buffer
    im[xmin:xmax, ymin:ymin+thickness,:] = [max_intensity,0,0] #draws a box around the coordinates
    im[xmin:xmax, ymax-thickness:ymax,:] = [max_intensity,0,0]
    im[xmin:xmin+thickness, ymin:ymax,:] = [max_intensity,0,0]
    im[xmax-thickness:xmax, ymin:ymax,:] = [max_intensity,0,0]

    plt.figure()
    plt.imshow(im)
    plt.show()
    misc.imsave(path, im)
    return


#!/usr/bin/python
# -*- coding: utf-8 -*-

"""
pycopanbehave -- An adaptive network mode of behaviour selection in Python

Copyright (C) 2011--2016 Potsdam Institute for Climate Impact Research
Authors: Jonathan F. Donges <donges@pik-potsdam.de>,
         Carl-Friedrich Schleussner <schleussner@pik-potsdam.de>,
         Denis Engemann <denis.engemann@gmail.com>
URL:     <http://www.pik-potsdam.de/copan/software>

PLOT SCRIPTS
"""

#
#  Imports
#

import os.path as op
import os
import pickle

import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

#
#  Initializations
#

input_file = 'full_trans_smok.pkl'
# input_file='trans_smok_output_full.pkl'
with open(input_file) as w:
    data = pickle.load(w)

if not os.path.isdir('figures'):
    os.mkdir('figures')

to_plot = [1, 2, 3]
set_titles = False
smokers = data['no_of_smokers']
decim = 1
ci = 66
n_smokers = 1000.
start_time = 0.
n_xticks = [0, 200, 400, 600, 790]
ma_window_size = 30
keys = ['full', 'dyn', 'infl', 'mean_field']
mapping = dict(zip(keys, ['coupled', 'network', 'interaction', 'mean-field']))
colors_context = "#55247A", "#3368A5", "#D33C3E", '#37442A'
color_smoker = '#A44791'
color_nonsmoker = '#BECD00'
sns.set(style='ticks', context='poster', font_scale=1.4)
my_gray = '#6e6a6a'

plt.rcParams.update({k: my_gray for k in (
    'axes.labelcolor',
    'axes.edgecolor',
    'text.color',
    'xtick.color',
    'ytick.color'
)})

plt.rcParams.update({k: 2.2 for k in (
    'axes.linewidth',
    'ytick.major.width',
    'xtick.major.width'
)})

#
#  DEFINE PLOTTING FUNCTIONS
#

def moving_average(values, window, axis):
    """
    Add docstring!
    """
    weights = np.repeat(1.0, window) / window
    sma = np.apply_along_axis(
        lambda m: np.convolve(m, weights, mode='valid'), arr=values, axis=axis)
    return sma

def preprocess_array(X):
    """
    Add docstring!
    """
    return moving_average(np.transpose(X, (1, 2, 0)), ma_window_size, 1)

def rescale(x, axis, method='percent'):
    """
    Add docstring!
    """
    if method == 'divide':
        x = x / x[:, 0:2, :]
    elif method == 'percent':
        x = (x - x[:, 0:1, :]) / x[:, 0:1, :]
        # x = (x - x[:, 0:20, :].mean(axis=1)) / x[:, 0:20, :].mean(axis=1)
        x *= 100
    return x

#
#  MAIN PLOTTING SCRIPT
#

#  FIGURE 1
if 1 in to_plot:
    X = [np.asarray(smokers[k])[:1000, start_time:] for k in keys]
    X = preprocess_array(X)
    X /= n_smokers

    times = np.arange(X.shape[1])

    fig = plt.figure()
    for i, (key, color) in enumerate(zip(keys, colors_context)):
        sns.tsplot(X[::decim, :, i:i + 1], color=color,
                   condition=mapping[key])

    plt.xlabel('Time [AU]')
    plt.xticks(n_xticks)
    plt.xlim(0, max(n_xticks))
    plt.ylabel('Proportion of smokers')
    if set_titles:
        plt.title('Evolution of Smoking prevalence')
    sns.despine(offset=10, trim=True)
    plt.gca().set_xticklabels(np.linspace(0, 1, len(n_xticks)))
    plt.subplots_adjust(bottom=0.15)
    plt.show()
    fig.savefig(op.join('figures', 'n_smokers-%i.png' % ci), dpi=300)

####################
#  FIGURE 2
if 2 in to_plot:

    centrality = data['centrality']
    centrality_smoker = np.array([
        np.asarray(centrality['smoker'][k])[:1000, start_time:] for k in keys])
    centrality_nosmoker = np.array([
        np.asarray(centrality['non_smoker'][k])[:1000, start_time:]
        for k in keys])
    centrality_smoker = rescale(preprocess_array(centrality_smoker), axis=1)
    centrality_nosmoker = rescale(preprocess_array(centrality_nosmoker),
                                  axis=1)
    centrality_diff = centrality_smoker - centrality_nosmoker

    fig, axes = plt.subplots(1, 2, figsize=(12, 8), sharex=True, sharey=True)
    ax1, ax2 = axes.ravel()
    ax1.set_xlim(0, max(n_xticks))
    ax2.set_xlim(0, max(n_xticks))
    ax1.set_ylim(-40, 30)
    ax2.set_ylim(-40, 30)
    ax1.text(20, 27, 'a', weight='bold', fontsize=24)
    ax2.text(20, 27, 'b', weight='bold', fontsize=24)

    X = centrality_smoker
    for i, (key, color) in enumerate(zip(keys, colors_context)):
        sns.tsplot(X[::decim, :, i:i + 1],
                   color=color, condition=mapping[key], ax=ax1, ci=[95, ci],
                   legend=False)
    if set_titles:
        fig.suptitle('Evolution of eigenvector centrality', fontsize=32)

    plt.xticks(n_xticks)
    # if set_titles:
    ax1.set_title('smokers')
    ax1.set_xlabel('Time [AU]')
    ax1.set_ylabel('Relative change in centrality [percent]')
    ax1.set_xlim(0, max(n_xticks))

    X = centrality_nosmoker
    for i, (key, color) in enumerate(zip(keys, colors_context)):
        sns.tsplot(X[::decim, :, i:i + 1],
                   color=color, condition=mapping[key], ax=ax2, ci=[95, ci],
                   legend=True)

    plt.xticks(n_xticks)
    ax2.set_xticklabels(np.linspace(0, 1, len(n_xticks)))
    # if set_titles:
    ax2.set_title('non-smokers')
    ax2.set_xlabel('Time [AU]')
    ax2.set_xlim(0, max(n_xticks))

    sns.despine(offset=10, trim=True)
    plt.subplots_adjust(bottom=0.15, top=.85, right=.96, left=0.14, wspace=.27)
    plt.show()
    fig.savefig(op.join('figures', 'centrality_smokers--%i.png' % ci), dpi=300)

####################
#  FIGURE 3

X = np.array([np.asarray(data['conditional_prob'][k])[:1000, start_time:, :]
              for k in keys])
X = np.transpose(X, (-1, 1, 2, 0))
X = moving_average(X, ma_window_size, axis=2)

if 3 in to_plot:
    fig, axes = plt.subplots(2, 2, figsize=(16, 12), sharex=False,
                             sharey=False)
    axes = axes.ravel()
    n_yticks = [0, 200, 400, 700]
    for i, (key, color) in enumerate(zip(keys, colors_context)):
        ax = axes[i]
        ax.set_ylim(-100, 700)
        ax.text(50, 710, 'abcd'[i], weight='bold', fontsize=24)
        for degree, (x, subcolor) in enumerate(
                zip(X,
                    sns.light_palette(color, len(X))[::-1]), 1):
            sns.tsplot(rescale(x[::decim, :, i:i + 1] + 1, axis=1),
                       color=subcolor, condition='degree %i' % degree,
                       ax=ax, ci=ci)
        # plt.xticks(times, times[::1000])
        if i == 0:
            ax.set_ylabel('Relative change in CP [percent]', labelpad=10)
        ax.set_title(mapping[key])
        ax.set_xlim(0, max(n_xticks))
        ax.set_xticks(n_xticks)
        if i in (1, 3):
            ax.set_ylabel('')
            ax.set_yticks(n_yticks, ['', '', '', ''])
        if i in (0, 1):
            ax.set_xlabel('')
            ax.set_yticks(n_yticks, ['', '', '', ''])
        if i in (2, 3):
            ax.set_xlabel('Time [AU]')
        ax.set_xticklabels(np.linspace(0, 1, len(n_xticks)))

    if set_titles:
        fig.suptitle('Conditional probability of smoking', fontsize=32)
    sns.despine(offset=10, trim=True)
    plt.subplots_adjust(left=0.1, bottom=0.15, right=0.97, top=0.86,
                        wspace=.41, hspace=.41)
    fig.show()
    fig.savefig(op.join('figures', 'conditional_prob_all-%i.png' % ci),
                dpi=300)
# fig.savefig(op.join('figures', 'conditional_prob_all.png'))


import numpy as np
from astropy.io import fits
from scipy import integrate
import matplotlib.pyplot as plt

def blockshaped(arr, nrow, ncol):  # CONFIRMED
    h, w = arr.shape
    return arr.reshape(h // nrow, nrow, -1, ncol).swapaxes(1, 2).reshape(-1, nrow, ncol)

def rebin(data, n):
    rebinned = []
    for z in range(len(data)):
        subarrays = blockshaped(data[z, :, :], n, n)  # bin the data in groups of nxn (4x4) pixels
        # each pixel in the new, rebinned data cube is the mean of each 4x4 set of original pixels
        # reshaped = np.mean(np.mean(subarrays, axis=-1), axis=-1).reshape((int(len(data[0]) / 4.),
        #                                                                   int(len(data[0][0]) / 4.)))
        reshaped = n**2 * np.mean(np.mean(subarrays, axis=-1), axis=-1).reshape((int(len(data[0]) / n),
                                                                                int(len(data[0][0]) / n)))
        rebinned.append(reshaped)
    print('rebinned')
    return np.asarray(rebinned)

def compare(data, model, z_ax, inds_to_try2, v_sys, n):
    data_4 = rebin(data, n)
    ap_4 = rebin(model, n)

    for i in range(len(inds_to_try2)):
        print(inds_to_try2[i][0], inds_to_try2[i][1])
        plt.plot(z_ax, ap_4[:, inds_to_try2[i][1], inds_to_try2[i][0]], 'r+', label=r'Model')  # r-
        plt.plot(z_ax, data_4[:, inds_to_try2[i][1], inds_to_try2[i][0]], 'b+', label=r'Data')  # b:
        plt.axvline(x=v_sys, color='k', label=r'v$_{\text{sys}}$')
        # plt.title(str(inds_to_try2[i][0]) + ', ' + str(inds_to_try2[i][1]))  # ('no x,y offset')
        plt.legend()
        plt.xlabel(r'Frequency [GHz]')
        plt.ylabel(r'Flux Density [Jy/beam]')
        plt.show()
        plt.close()

if __name__ == "__main__":

    base = '/Users/jonathancohn/Documents/dyn_mod/'

    import pickle

    import dyn_model as dm
    params, priors = dm.par_dicts(base + 'param_files/ngc_3258bl_params.txt', q=False)  # get dicts of params and file names from parameter file

    ax_lab = [r'$\log_{10}$(M$_{\odot}$)', 'deg', 'deg', 'pixels', 'pixels', 'km/s', 'km/s', 'km/s', 'pc', 'pc', 'unitless', r'M$_{\odot}$/L$_{\odot}$']
    pars = ['mbh', 'inc', 'PAdisk', 'xloc', 'yloc', 'vsys', 'sig1', 'sig0', 'mu', 'r0', 'f', 'ml_ratio']
    for par in range(len(pars)):
        chains = 'emcee_out/flatchain_' + pars[par] + '_100_1_50.pkl'

        with open(base + chains, 'rb') as pk:
            u = pickle._Unpickler(pk)
            u.encoding = 'latin1'
            chain = u.load()

            if pars[par] == 'mbh':
                plt.hist(np.log10(chain), 100, color="k", histtype="step")  # axes[i]
                percs = np.percentile(np.log10(chain), [16., 50., 84.])
                plt.axvline(np.log10(params[pars[par]]), ls='-', color='k')
            else:
                plt.hist(chain, 100, color="k", histtype="step")  # axes[i]
                percs = np.percentile(chain, [16., 50., 84.])
                print(params[pars[par]])
                plt.axvline(params[pars[par]], ls='-', color='k')

            plt.axvline(percs[1], ls='-', color='b')  # axes[i]
            plt.axvline(percs[0], ls='--', color='b')  # axes[i]
            plt.axvline(percs[2], ls='--', color='b')  #
            plt.tick_params('both', labelsize=16)
            plt.xlabel(ax_lab[par])
            # plt.title("Dimension {0:d}".format(i))
            plt.title(pars[par] + ': ' + str(round(percs[1],4)) + ' (+'
                                     + str(round(percs[2] - percs[1], 2)) + ', -'
                                     + str(round(percs[1] - percs[0], 2)) + ')', fontsize=16)
            plt.show()
    print(oop)

    # hdu = fits.open(base + 'NGC_1332_newfiles/NGC1332_CO21_C3_MS_bri_20kms.pbcor.fits')
    hdu = fits.open(base + 'ngc_3258/ngc3258_CO21_bri.pbcor.fits')
    data_in = hdu[0].data[0]
    z_len = len(data_in)  # store the number of velocity slices in the data cube
    freq1 = float(hdu[0].header['CRVAL3'])
    f_step = float(hdu[0].header['CDELT3'])
    # f_0 = 2.305380000000e11  # 2.29369132e11  # intrinsic frequency of CO(2-1) line  # 2.30538e11
    f_0 = float(hdu[0].header['RESTFRQ'])
    hdu.close()
    freq_axis = np.arange(freq1, freq1 + (z_len * f_step), f_step)  # [bluest, less blue, ..., reddest]
    # freq_axis = freq_axis[:-1]
    # print(freq_axis[0], freq_axis[-1])  # 2.29909, 2.28777 [e11]
    v_sys = 2760.76  # 1562.2  # km/s
    z_ax = np.asarray([((f_0 - freq) / freq) * (3.*10**5) for freq in freq_axis])  # v_opt, km/s
    # print(z_ax)
    # print(oop)

    # hdu = fits.open(base + 'ngc1332_things/NGC_1332_fullsize_idl_n5_beam31_s1.fits')
    # idl_out = hdu[0].data
    # hdu.close()

    # hdu = fits.open(base + 'ngc1332_things/NGC_1332_fullsize_filtconv_n5_beam31_s1.fits')
    # filt_out = hdu[0].data
    # hdu.close()

    # hdu = fits.open(base + 'NGC_1332_fullsize_apconv_n5_beam31fwhm_s1_weightdiv_corrsig.fits')  # _zred
    # hdu = fits.open(base + 'NGC_1332_fullsize_apconv_n5_beam31fwhm_s1.fits')
    # hdu = fits.open(base+ 'NGC_1332_freqcube_summed_apconv_n5_beam31fwhm_s1.fits')
    # hdu = fits.open(base + 'NGC_1332_freqcube_summed_apconv_n5_beam31fwhm_spline_s1.fits')
    # hdu = fits.open(base + 'outputs/NGC_3258_general_7.42_2386000000.0_1.02_166.8_354.89_362.04_-0.46_4.75_2760.76_0.14_3.16_46.0__benlucycorr.fits')
    # hdu = fits.open(base + 'outputs/NGC_3258_general_7.42_3.16_2386000000.0_1.02_166.8_2760.76_-0.46_4.75_0.14_362.04_354.89_46.0__pcscale.fits')
    # hdu = fits.open(base + 'outputs/NGC_3258_general_7.42_2386000000.0_1.02_166.8_354.89_362.04_-0.46_4.75_2760.76_0.14_3.16_46.0__subcube.fits')
    # hdu = fits.open(base + 'outputs/NGC_3258_general_7.42_2386000000.0_1.02_166.8_354.89_362.04_-0.46_4.75_2760.76_0.14_3.16_46.0__subcube_ellmask_bl2.fits')
    hdu = fits.open(base + 'outputs/NGC_3258_general_7.42_2386000000.0_1.02_166.8_354.89_362.04_-0.46_4.75_2760.76_0.14_3.16_46.0__subcube_ellmask_bl2_noell.fits')
    # hdu = fits.open(base + 'outputs/ngc_3258_benlucy_outcube.fits')

    ap_out = hdu[0].data
    hdu.close()

    hdu = fits.open('/Users/jonathancohn/Documents/dyn_mod/outputs/NGC_3258_fitting_ellipse.fits')
    ell_fits = hdu[0].data
    hdu.close()

    plt.imshow(ell_fits, origin='lower')
    plt.colorbar()
    plt.show()

    print(ap_out.shape)
    print(data_in.shape)
    data_in = data_in[39:85, 310:410, 310:410]
    print(data_in.shape)

    data_in *= ell_fits
    ap_out *= ell_fits

    idx = 59  # 26
    idx = 20
    print(np.amin((ap_out[idx] - data_in[idx])/data_in[idx]), np.amax((ap_out[idx] - data_in[idx])/data_in[idx]),
          np.median((ap_out[idx] - data_in[idx])/data_in[idx]))
    plt.imshow((ap_out[idx] - data_in[idx])/data_in[idx], origin='lower')
    # plt.imshow((idl_out[26] - data_in[26])/data_in[26], origin='lower')
    plt.colorbar()
    # plt.xlim(700,750)
    # plt.ylim(720,750)
    # plt.plot(848, 720, 'w*')
    plt.show()
    # print(oop)
    plt.imshow(data_in[idx], origin='lower')
    plt.colorbar()
    # plt.plot(848, 720, 'w*')
    plt.show()

    data_4 = rebin(data_in, 4)
    # idl_4 = rebin(idl_out, 4)
    # filt_4 = rebin(filt_out)
    ap_4 = rebin(ap_out, 4)

    # CREATED ONCE
    # hdu = fits.PrimaryHDU(data_4)
    # hdul = fits.HDUList([hdu])
    # hdul.writeto(base + 'NGC_1332_newfiles/NGC1332_CO21_C3_MS_bri_20kms_bin4x4.pbcor.fits')

    '''
    # UNCOMMENT THIS LATER
    hdu = fits.PrimaryHDU(ap_4)
    hdul = fits.HDUList([hdu])
    hdul.writeto(base + 'outputs/NGC_3258_general_4x4binned.fits')
    # hdul.writeto(base + 'NGC_1332_freqcube_summed_apconv_n5_beam31fwhm_spline_s1_4x4binned.fits')
    # hdul.writeto(base + 'NGC_1332_freqcube_summed_apconv_n5_beam31fwhm_s1_4x4binned.fits')

    hdu = fits.PrimaryHDU(data_4 - ap_4)
    hdul = fits.HDUList([hdu])
    hdul.writeto(base + 'NGC_3258_data_minus_freqcube_4x4binned.fits')
    '''
    # hdul.writeto(base + 'NGC_1332_data_minus_freqcube_spline_4x4binned.fits')
    # print(oops)
    #hdu = fits.PrimaryHDU(data_4 - ap_4)
    #hdul = fits.HDUList([hdu])
    #hdul.writeto('data_minus_apconv_4x4_fixed.fits')
    # print(oop)
    # 212, 180 --> 212*4
    # x=2 --> x_0 = 2, 3
    # x=2, n=4 --> x_0 = 4,5,6,7 = (x-1)*n
    # x=4, n=4 --> x_0 = 12,13,14,15 = (x-1)*n:(x*n) = 12:16 YAY!
    # y
    # 3, 4 (n=2): [(3/2):(3+1)/2] 22 23 32 33
    # 0, 0 (n=2): [(0/2):(0/2)+1] 00, 01, 10, 11

    #plt.imshow(idl_4[23], origin='lower')
    #plt.colorbar()
    #plt.plot(212, 180, 'w*')
    #plt.show()
    print(ap_4.shape, data_4.shape)
    chi_sq_num = 0.
    for z in range(len(data_4)):
        chi_sq_num += np.sum((ap_4[z] - data_4[z]) ** 2)
    print(chi_sq_num)  # ~ -1160 for take2, ~7.2 for pcscale. Good?

    # inds_to_try2 = np.asarray([[212, 180], [159, 159], [165, 155], [155, 165], [160, 160], [170, 170], [155, 155], [165, 165]])
    # inds_to_try2 = np.asarray([[178, 178], [200, 172], [168, 155], [155, 168], [135, 140], [125, 150]])
    # inds_to_try2 = np.asarray([[89, 92], [95, 89], [92, 92], [91, 86]])  # red side, blue side, center top, center bot
    inds_to_try2 = np.asarray([[10,10], [10,15], [15,10]])

    for i in range(len(inds_to_try2)):
        print(inds_to_try2[i][0], inds_to_try2[i][1])

        '''
        data_full = []
        for ii in range((inds_to_try2[i][1])*4,((inds_to_try2[i][1]+1)*4)):
            for jj in range((inds_to_try2[i][0])*4, ((inds_to_try2[i][0]+1)*4)):
                data_full.append(data_in[:, ii, jj])
        data_full = np.asarray(data_full)
        # print(data_full.shape)  # 4x4 x 75 --> 16, 75

        data_stack = np.zeros(shape=(data_full[0].shape))
        for l in range(len(data_full)):
            # plt.plot(z_ax, data_full[l], 'k--')  # confirmed
            data_stack += data_full[l]
        # plt.plot(z_ax, data_stack, 'g--')  # confirmed identical to how Data is calculated
        '''
        plt.plot(z_ax[39:85], ap_4[:, inds_to_try2[i][1], inds_to_try2[i][0]], 'r-', label=r'Astropy conv')  # ap_out  # 0.01217/.00031 *
                 # / np.amax(ap_out[:, inds_to_try2[i][1], inds_to_try2[i][0]]), 'r-', label=r'Astropy conv')
        # plt.plot(z_ax, idl_out[:, inds_to_try2[i][1], inds_to_try2[i][0]], 'k--', label=r'IDL conv')  # idl_out
                 # / np.amax(idl_out[:, inds_to_try2[i][1], inds_to_try2[i][0]]), 'k--', label=r'IDL conv')
        # plt.plot(z_ax, filt_out[:, inds_to_try2[i][1], inds_to_try2[i][0]], 'b:', label=r'Scipy conv')
                 # / np.amax(filt_out[:, inds_to_try2[i][1], inds_to_try2[i][0]]), 'b:', label=r'Scipy conv')
        plt.plot(z_ax[39:85], data_4[:, inds_to_try2[i][1], inds_to_try2[i][0]], 'b:', label=r'Data')  # data_in
        plt.axvline(x=v_sys, color='k')
        plt.title(str(inds_to_try2[i][0]) + ', ' + str(inds_to_try2[i][1]))  # ('no x,y offset')
        plt.legend()
        plt.show()
        # base = '/Users/jonathancohn/Documents/dyn_mod/lps/sum_div_filtconv_'  # div_'
        # plt.savefig(base + 'newfiles_lp_s' + str(s) + '_' + str(inds_to_try2[i][0]) + '_' + str(inds_to_try2[i][1]))
        plt.close()
    print(oop)

    '''  #
    idl_filt = []
    idl_ap = []
    for z in range(len(idl_out)):
        print(z)
        idl_filt.append((idl_out[z] - filt_out[z]))
        idl_ap.append((idl_out[z] - ap_out[z]))
        if 35 < z < 39:
            plt.imshow(idl_filt[z], origin='lower')
            plt.colorbar()
            plt.show()

            plt.imshow(idl_ap[z], origin='lower')
            plt.colorbar()
            plt.show()

    idl_filt = np.asarray(idl_filt)
    idl_ap = np.asarray(idl_ap)

    idl_filt_coll = integrate.simps(idl_filt, axis=0)
    idl_ap_coll = integrate.simps(idl_ap, axis=0)

    plt.imshow(idl_filt_coll, origin='lower')
    plt.colorbar()
    plt.show()

    plt.imshow(idl_ap_coll, origin='lower')
    plt.colorbar()
    plt.show()
    print(oops)
    # '''  #

    obs3d = np.asarray([[1., 2., 3., 4.], [1., 2., 3., 4.], [5., 6., 7., 8.], [9., 10., 11., 12.]])
    s = 2
    print(obs3d.shape)

    '''
    subarrays = blockshaped(obs3d[:, :], 4, 4)  # bin the data in groups of 4x4 pixels
    # data[z, ::-1, :] flips the y-axis, which is stored in python in the reverse order (problem.png)

    # Take the mean along the first axis of each subarray in subarrays, then take the mean along the
    # other (s-length) axis, such that what remains is a 1d array of the means of the sxs subarrays. Then reshape
    # into the correct lengths
    reshaped = np.mean(np.mean(subarrays, axis=-1), axis=-1).reshape((int(len(data[0])/4.),
                                                                      int(len(data[0][0])/4.)))
    reshaped_m = np.mean(np.mean(subarrays_m, axis=-1), axis=-1).reshape((int(len(mask[0])/4.),
                                                                          int(len(mask[0][0])/4.)))
    rebinned.append(reshaped)
    rebinned_m.append(reshaped_m)
    print("Rebinning the cube done in {0} s".format(time.time() - t_rebin))  # 0.5 s
    '''

    intrinsic_cube = []  # np.zeros(shape=(len(z_ax), len(fluxes), len(fluxes[0])))
    # break each (s*len(realx), s*len(realy))-sized velocity slice of obs3d into an array comprised of sxs subarrays
    # i.e. break the 300s x 300s array into blocks of sxs arrays, so that I can take the means of each block
    subarrays = blockshaped(obs3d[:, :], s, s)

    print('sub:')
    print(subarrays)

    # Take the mean along the first (s-length) axis of each subarray in subarrays, then take the mean along the
    # other (s-length) axis, such that what remains is a 1d array of the means of the sxs subarrays. Then reshape
    # into the correct real x_pixel by real y_pixel lengths
    reshaped = s**2 * np.mean(np.mean(subarrays, axis=-1), axis=-1).reshape((2, 2))
    reshaped2 = np.sum(np.sum(subarrays, axis=-1), axis=-1).reshape((2, 2))
    print('reshape:')
    print(reshaped)

    print('reshaped2:')
    print(reshaped2)


import sqlite3
import os
import datetime
import time
from shutil import copyfile
from . import models
import calendar
from .models import Entities

class DatabaseHelper:

    DATA_DIR = os.path.expanduser('~/.liquidity/')
    DB_FILE = os.path.expanduser('~/.liquidity/liquidity.db')
    DB_FLAG_FILE = os.path.expanduser('~/.liquidity/.liquidity_db_exists')
    DB_SCHEMA_FILE = os.path.expanduser('~/.liquidity/liquidity_schema.sql')
    OFFSET_FILE = os.path.expanduser('~/.liquidity/.offset')

    def __init__(self):
        if not os.path.exists(self.DATA_DIR):
            os.mkdir(self.DATA_DIR)

        assert os.path.exists(self.DATA_DIR)
        self.conn = sqlite3.connect(self.DB_FILE)
        if not os.path.exists(self.DB_FLAG_FILE):
            self.create_db()
            print('\n\nCREATED LIQUIDITY TABLES\n\n')

        # self.open_database()
        self.transactions = self.get_transactions()
        self.load_offset()

    def open_database(self):
        self.conn = sqlite3.connect(self.DB_FILE)

    def create_db(self):
        copyfile('./app/db/schemas/liquidity_schema.sql', self.DATA_DIR + 'liquidity_schema.sql')
        assert os.path.exists(self.DB_SCHEMA_FILE)
        with open(self.DB_SCHEMA_FILE, 'r') as schema_file:
            queries = schema_file.read()
            self.conn.executescript(queries)
            self.conn.close()
            with open(self.DB_FLAG_FILE, 'w') as f:
                pass

    def get_current_month(self):
        return datetime.datetime.now().strftime('%B')

    def get_last_month(self):
        today = datetime.datetime.now()
        first = today.replace(day = 1)
        last_month = (first - datetime.timedelta(days = 1)).strftime('%B')

        return last_month

    def get_entities(self, type, month):
        self.open_database()

        query = '''SELECT
        id,
        e_name,
        e_value,
        e_type
        FROM  entities
        WHERE e_type = {type}
        '''.format(type = type)

        entities = []

        cursor = self.conn.execute(query)
        for row in cursor:
            entity = models.Entities()

            entity.id = int(row[0])
            entity.name = row[1]
            entity.value = float(row[2])
            entity.type = int(row[3])

            if entity.type == Entities.TYPE_EXPENSE or entity.type == Entities.TYPE_INCOME:
                entity.value =self.get_entity_monthly_value(entity, month)

            entities.append(entity)

        self.conn.close()
        return entities

    def get_entity_monthly_value(self, entity, month):
        transactions = self.transactions
        value = 0

        for transaction in transactions:
            t_month = datetime.datetime.fromtimestamp(transaction.timestamp).strftime('%B')
            if t_month == month:
                if entity.type == Entities.TYPE_INCOME:
                    if transaction.from_entity.type == Entities.TYPE_INCOME and entity.id == transaction.from_entity.id:
                        value += transaction.value
                elif entity.type == Entities.TYPE_EXPENSE:
                    if transaction.to_entity.type == Entities.TYPE_EXPENSE  and entity.id == transaction.to_entity.id:
                        value += transaction.value

        return value

    def add_entity(self, entity):
        self.open_database()

        query = '''INSERT INTO entities(
            e_name,
            e_value,
            e_type
        ) VALUES (
            "{name}",
            {value},
            {type}
        )
        '''.format(
            name = entity.name,
            value = entity.value,
            type = entity.type
        )

        self.conn.execute(query)
        self.conn.commit()

        self.conn.close()

    def update_entity(self, entity):
        self.open_database()

        query = '''UPDATE entities SET
            e_name = "{name}",
            e_value = {value},
            e_type = {type}
            WHERE id = {id}'''.format(
            name = entity.name,
            value = entity.value,
            type = entity.type,
            id = entity.id
        )

        self.conn.execute(query)
        self.conn.commit()
        self.conn.close()

    def delete_entity(self, entity):
        self.open_database()

        query = '''DELETE FROM entities WHERE id = {id}'''.format(
            id = entity.id
        )

        self.conn.execute(query)
        self.conn.commit()

        self.conn.close()
        pass

    def get_entity(self, id):
        self.open_database()

        query = '''SELECT
        id,
        e_name,
        e_value,
        e_type
        FROM  entities
        WHERE id = {id}
        '''.format(id = id)

        entities = []

        cursor = self.conn.execute(query)
        for row in cursor:
            entity = models.Entities()

            entity.id = int(row[0])
            entity.name = row[1]
            entity.value = float(row[2])
            entity.type = int(row[3])

            entities.append(entity)

        self.conn.close()
        if len(entities) > 0:
            return entities[0]
        else:
            return None

    def reset_monthly_entities(self):

        print('Backing up data...\n')

        today = datetime.datetime.now()
        first = today.replace(day = 1)
        last_month = (first - datetime.timedelta(days = 1)).strftime('%B')

        copyfile('db/databases/liquidity.db', 'db/databases/{}.db'.format(last_month))

        print('Resetting income and expense values...\n')

        query = '''UPDATE entities
        SET e_value = 0
        WHERE e_type = {} OR e_type = {}'''.format(
            models.Entities.TYPE_EXPENSE,
            models.Entities.TYPE_INCOME
        )

        self.open_database()
        self.conn.execute(query)
        self.conn.commit()
        self.conn.close()

    def refresh_transactions(self):
        self.transactions = self.get_transactions()

    def get_transactions(self):
        self.open_database()

        query = '''SELECT
        id,
        t_from_entity_id,
        t_to_entity_id,
        t_value,
        t_comment,
        t_timestamp
        FROM transactions
        ORDER BY t_timestamp
        '''

        transactions = []

        for row in self.conn.execute(query):
            transaction = models.Transactions()

            transaction.id = int(row[0])
            transaction.from_entity = self.get_entity(int(row[1]))
            transaction.to_entity = self.get_entity(int(row[2]))
            transaction.value = float(row[3])
            transaction.comment = row[4]
            transaction.timestamp = int(row[5])

            transactions.append(transaction)

        self.transactions = transactions
        return transactions

    def add_transaction(self, transaction):

        from_entity = transaction.from_entity
        to_entity = transaction.to_entity

        from_entity.value -= transaction.value
        to_entity.value += transaction.value

        self.update_entity(from_entity)
        self.update_entity(to_entity)

        query = '''INSERT INTO transactions(
            t_from_entity_id,
            t_to_entity_id,
            t_value,
            t_timestamp,
            t_comment
        ) VALUES (
            {from_e_id},
            {to_e_id},
            {value},
            {timestamp},
            "{comment}"
        )
        '''.format(
            from_e_id = transaction.from_entity.id,
            to_e_id = transaction.to_entity.id,
            value = transaction.value,
            timestamp = transaction.timestamp,
            comment = transaction.comment
        )

        try:
            self.open_database()
            self.conn.execute(query)
            self.conn.commit()
        except Exception as e:
            print(e)

        self.conn.close()

    def update_transaction(self, transaction):
        self.open_database()

        query = '''UPDATE transactions SET
            t_from_entity_id = {from_e_id},
            t_to_entity_id = {to_e_id},
            t_value = {value},
            t_timestamp = {timestamp},
            t_comment = "{comment}"
            WHERE id = {id}
        '''.format(
            from_e_id = transaction.from_entity.id,
            to_e_id = transaction.to_entity.id,
            value = transaction.value,
            timestamp = transaction.timestamp,
            comment = transaction.comment,
            id = transaction.id
        )

        self.conn.execute(query)
        self.conn.commit()

        self.conn.close()
        pass

    def get_transaction(self, id):
        t_id = int(id)

        for t in self.transactions:
            if int(t.id) == t_id:
                return t

    def delete_transaction(self, transaction):
        self.open_database()

        query = '''DELETE FROM transactions WHERE id = {id}'''.format(
            id = transaction.id
        )

        self.conn.execute(query)
        self.conn.commit()

        self.conn.close()
        pass

    #HOME FUNCTIONS
    def get_net_worth(self, month):
        assets = self.get_entities(models.Entities.TYPE_ASSET, month)
        debts = self.get_entities(models.Entities.TYPE_DEBT, month)

        assets_total = 0
        debts_total = 0

        for asset in assets:
            assets_total += asset.value

        for debt in debts:
            debts_total += debt.value

        return assets_total - debts_total

    def get_month_income(self, month):
        transactions = self.transactions
        current_month = month

        total_income = 0

        for transaction in transactions:
            t_month = datetime.datetime.fromtimestamp(transaction.timestamp).strftime('%B')
            if current_month == t_month:
                if transaction.from_entity.type == Entities.TYPE_INCOME:
                    total_income += transaction.value

        return total_income

    def get_month_expense(self, month):
        transactions = self.transactions
        current_month = month

        total_expense = 0

        for transaction in transactions:
            t_month = datetime.datetime.fromtimestamp(transaction.timestamp).strftime('%B')
            if current_month == t_month:
                if transaction.to_entity.type == Entities.TYPE_EXPENSE:
                    total_expense += transaction.value

        return total_expense

    def get_month_net(self, month):
        total_income = self.get_month_income(month)
        total_expense = self.get_month_expense(month)

        return total_income - total_expense

    def get_budget_dict(self, amount):
        dictionary = {
            'needs': float(amount) * 0.5,
            'wants': float(amount) * 0.2,
            'save': float(amount) * 0.3
        }

        return dictionary

    def get_spending_capacity(self, amount, month):
        spending_capacity = 0
        liquidity = amount
        expense = self.get_month_expense(month)

        spending_capacity = liquidity - expense
        return spending_capacity

    def get_daily_expense_limit(self, spending_capacity):
        return round(spending_capacity / 31, 2)

    def get_monthly_avg_income(self):
        incomes = []
        for i in range(1, 13):
            month = calendar.month_name[i]
            income = self.get_month_income(month)
            if income != 0:
                incomes.append(income)
        incomes.sort()
        try:
            return incomes[int(len(incomes)/2)] if len(incomes) % 2 != 0 else (incomes[int(len(incomes)/2)] + incomes[int(len(incomes)/2) -1 ])/2
        except IndexError:
            return 0

    def get_monthly_avg_expense(self):
        expenses = []
        for i in range(1, 13):
            month = calendar.month_name[i]
            expense = self.get_month_expense(month)
            if expense != 0:
                expenses.append(expense)
        expenses.sort()
        try:
            return expenses[int(len(expenses)/2)] if len(expenses) % 2 != 0 else (expenses[int(len(expenses)/2)] + expenses[int(len(expenses)/2) -1 ])/2
        except IndexError:
            return 0

    def get_monthly_avg_net(self):
        return self.get_monthly_avg_income() - self.get_monthly_avg_expense()

    def set_offset(self, days = 0):
        with open(self.OFFSET_FILE, 'w') as o_file:
            o_file.write(str(days))
        self.load_offset()

    def load_offset(self):
        try:
            with open(self.OFFSET_FILE, 'r') as o_file:
                self.offset = int(o_file.read())
        except:
            self.set_offset()

    def get_offset(self):
        return self.offset

    def get_daily_average_expense(self):
        expenses = []
        transactions = self.transactions

        for d in range(self.offset, datetime.datetime.now().day + 1):
            expense = 0
            for transaction in transactions:
                if datetime.datetime.fromtimestamp(transaction.timestamp).day == d:
                    if transaction.to_entity.type == Entities.TYPE_EXPENSE:
                        expense += transaction.value

            expenses.append(expense)

        expenses.sort()
        try:
            return expenses[int(len(expenses)/2)] if len(expenses) % 2 != 0 else (expenses[int(len(expenses)/2)] + expenses[int(len(expenses)/2) -1 ])/2
        except IndexError:
            return 0


import os
import sys
import math
import pygame
import pygame.mixer
from pygame.locals import *

class Balloon(pygame.sprite.Sprite):
    def __init__(self):
        pygame.sprite.Sprite.__init__(self) #call Sprite initializer
        self.image = pygame.image.load("balloon.jpg")
        self.rect = self.image.get_rect()
        self.mask = pygame.mask.from_surface(self.image)

width = 500
height = 300
pygame.init()
pygame.display.set_mode((width, height))

b1 = Balloon()
b2 = Balloon()

while True:
    event = pygame.event.poll()
    if event.type == pygame.QUIT:
        running = 0
    elif  pygame.sprite.spritecollide(b1, b2, False, pygame.sprite.collide_mask):
        print("sprites have collided!")


__author__ = 'RetroPy'
###Dedicated to my little Matey Potatey, AKA Mr Po. Feb 2011 - 23/09/2014

import pygame
import sys
from pygame.locals import *

# set up pygame
pygame.init()

# set up the window
windowSurface = pygame.display.set_mode((500, 400), 0, 32)
pygame.display.set_caption('RPN Calc')

# set up the colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)

# set up fonts
basicFont = pygame.font.SysFont(None, 48)

# set up the text
text = basicFont.render('Hello world!', True, WHITE, BLUE)
textRect = text.get_rect()
textRect.centerx = windowSurface.get_rect().centerx
textRect.centery = windowSurface.get_rect().centery

# draw the text's background rectangle onto the surface
pygame.draw.rect(windowSurface, RED, (textRect.left - 20, textRect.top - 20, textRect.width + 40, textRect.height + 40))

# get a pixel array of the surface
pixArray = pygame.PixelArray(windowSurface)
pixArray[480][380] = BLACK
del pixArray

# draw the text onto the surface
windowSurface.blit(text, textRect)

# draw the window onto the screen
pygame.display.update()

# run the game loop
while True:
    for event in pygame.event.get():
        if event.type == QUIT:
            pygame.quit()
            sys.exit()


import numpy as np

import logging

from matplotlib import pyplot as plt

from time import time

from sklearn.svm import SVC
from sklearn.model_selection import train_test_split
from sklearn.datasets import fetch_olivetti_faces
from sklearn.metrics import confusion_matrix
from sklearn.metrics import classification_report

# loading dataset
logging.basicConfig(level = logging.INFO, format="%(asctime)s %(message)s")
data_home = 'code/datasets/'
logging.info("start loading data")
faces = fetch_olivetti_faces(data_home=data_home)
logging.info("Complete")

# label
X = faces.data
y = faces.target
targets = np.unique(faces.target)
target_names = np.array(["p%d" % t for t in targets])
n_targets = target_names.shape[0]
n_samples, h, w = faces.images.shape

# show image
def plot_gallery(images, titles, h, w, n_row=2, n_col=5):

    plt.figure(figsize=(2*n_col, 2.2*n_row),dpi=140)
    plt.subplots_adjust(bottom=0, left=.01, right=.99, top=.90, hspace=.01)
    for i in range(n_row * n_col):
        plt.subplot(n_row, n_col, i+1)
        plt.imshow(images[i].reshape((h,w)), cmap=plt.cm.gray)
        plt.title(titles[i])
        plt.axis('off')

sample_images = None
sample_titles = []
for i in range(n_targets):
    people_images = X[y == i]
    people_sample_index = np.random.randint(0, people_images.shape[0], 1)
    people_sample_image = people_images[people_sample_index, :]
    if sample_images is not None:
        sample_images = np.concatenate((sample_images, people_sample_image), axis=0)
    else:
        sample_images = people_sample_image
    sample_titles.append(target_names[i])

# loading data
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=4)

# svm
t = time()
clf = SVC(C=1.0, class_weight='balanced', coef0=0.0,
    decision_function_shape=None, degree=3, gamma='auto', kernel='rbf',
    max_iter=-1, probability=False, random_state=None, shrinking=True,
    tol=0.001, verbose=False)
clf.fit(X_train, y_train)
print("time：{}s".format(time() - t))

# classification report
y_pred = clf.predict(X_test)
print(classification_report(y_test, y_pred))

# confusion matrix
cm = confusion_matrix(y_test, y_pred, labels=range(n_targets))
print("confusion_matrix:\n")
np.set_printoptions(threshold=np.inf)
print(cm[:])


#!/usr/bin/python3

import collections
import json

import locale

from lxml import html
from lxml.cssselect import CSSSelector

import requests
import re

locale.setlocale(locale.LC_ALL, 'en_US.UTF-8')

class FrozenDict(collections.Mapping):
    """An immutable dict, see http://stackoverflow.com/a/2704866"""

    def __init__(self, *args, **kwargs):
        self._d = dict(*args, **kwargs)
        self._hash = None

    def __iter__(self):
        return iter(self._d)

    def __len__(self):
        return len(self._d)

    def __getitem__(self, key):
        return self._d[key]

    def __hash__(self):
        # It would have been simpler and maybe more obvious to
        # use hash(tuple(sorted(self._d.iteritems()))) from this discussion
        # so far, but this solution is O(n). I don't know what kind of
        # n we are going to run into, but sometimes it's hard to resist the
        # urge to optimize when it will gain improved algorithmic performance.
        if self._hash is None:
            self._hash = 0
            for pair in self .iteritems():
                self._hash ^= hash(pair)
        return self._hash

def parse_product(element):
    select_name = CSSSelector('[itemprop=name]')
    select_url = CSSSelector('[itemprop=url]')
    select_price = CSSSelector('[itemprop=price]')
    select_desc = CSSSelector('.product_descr p')

    name = select_name(element)
    url = select_url(element)
    price = select_price(element)
    desc = select_desc(element)

    assert len(name) == 1
    assert len(desc) == 1
    assert len(price) == 1
    assert len(url) == 1

    if (len(name) != 1 or len(desc) != 1 or len(price) != 1 or len(url) != 1):
        return None

    return (name[0].text.strip(), {
        'url': url[0].get('href'),
        'desc': re.sub(re.compile(r"\r\s*"), "\n", desc[0].text.strip()),
        'price': locale.atof(price[0].get('content')),
    })

def get_products():
    page = requests.get("https://routerboard.com")
    root = html.fromstring(page.content)

    select_products = CSSSelector(".product_entry:not(.hideBox):not(.hist)")

    products = select_products(root)

    return {name: details for name, details in (
        parse_product(x) for x in products)}

def format_product(name, product):
    return "%s: %s\n  %s\n  %s" % (
        name,
        locale.currency(product['price']),
        product['desc'].replace("\n", "\n  "),
        product['url']
    )

if __name__ == '__main__':
    settings = json.loads(open("settings.json").read())

    db_old = None
    try:
        db_file = open(settings['db_file'], 'r')
        db_old = json.load(db_file)
    except (json.JSONDecodeError, FileNotFoundError):
        pass

    db_new = get_products()
    db_file = open(settings['db_file'], 'w')

    if db_old is not None:
        for name in (db_new.keys() - db_old.keys()):
            print("+ %s\n" % format_product(name, db_new[name]))
        for name in (db_old.keys() - db_new.keys()):
            print("- %s\n" % format_product(name, db_old[name]))

    json.dump(db_new, db_file, sort_keys=True, indent=4)

    db_file.close()


# -*- coding: utf-8 -*-
"""
Created on Mon Apr 10 14:02:33 2017

@author: Dell
"""
import numpy as np
from matplotlib.font_manager import FontProperties
from matplotlib import pyplot as plt

class code_spectrum(object):
    def __init__(self,alpha_max,Tg,xi):
        gamma=0.9+(0.05-xi)/(0.3+6*xi)
        eta1=0.02+(0.05-xi)/(4+32*xi)
        eta1=eta1 if eta1>0 else 0
        eta2=1+(0.05-xi)/(0.08+1.6*xi)
        eta2=eta2 if eta2>0.55 else 0.55
        T=np.linspace(0,6,601)
        alpha=[]
        for t in T:
            if t<0.1:
                alpha.append(np.interp(t,[0,0.1],[0.45*alpha_max,eta2*alpha_max]))
            elif t<Tg:
                alpha.append(eta2*alpha_max)
            elif t<5*Tg:
                alpha.append((Tg/t)**gamma*eta2*alpha_max)
            else:
                alpha.append((eta2*0.2**gamma-eta1*(t-5*Tg))*alpha_max)
        self.__spectrum={'T':T,'alpha':alpha}

    @property
    def spectrum(self):
        return self.__spectrum

if __name__=='__main__':
    cs1=code_spectrum(0.12,0.45,0.02)
    cs2=code_spectrum(0.34,0.45,0.02)
    cs3=code_spectrum(0.72,0.45,0.02)

    font = FontProperties(fname=r"C:\\WINDOWS\\Fonts\\simsun.ttc", size=14)#C:\WINDOWS\Fonts
    plt.xlabel('T(s)',fontproperties=font)
    plt.ylabel('α',fontproperties=font)
    plt.plot(cs1.spectrum['T'],cs1.spectrum['alpha'],label=u'小震')
    plt.plot(cs2.spectrum['T'],cs2.spectrum['alpha'],label=u'中震')
    plt.plot(cs3.spectrum['T'],cs3.spectrum['alpha'],label=u'大震')
    plt.legend(prop=font)


# MGZF 蘑菇租房
import requests
from lxml import etree
import re
import csv
import pandas as pd
import time
import os

class Get_infor():
    def __init__(self):
        self.headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0',' WOW64) AppleWebKit/537.cipg (KHTML, like Gecko) Chrome/45.0.2454.101 Safari/537.36"
        }
        self.URL = 'http://www.mgzf.com/'
        self.start_url = self.URL + 'list/pg1/'
        self.folder = './SH/'

        self.title_list=[]#标题列表
        self.address_list=[]#地址列表
        self.shape_list=[]#户型面积列表
        self.money_list=[]#价格列表
        self.href_list=[]#链接列表
        self.detail_list =[]#详细

    def run(self, useLocal=False):
        '''
        useLoal: use local html file
        '''
        isExists = os.path.exists(self.folder)
        if not isExists:  # 目录不存在，则创建
            os.makedirs(self.folder)
        if useLocal == True:
            # use local html file
            try:
                with open(self.folder+'urlList.html', 'r', encoding='utf-8') as f:
                    html = ''
                    for line in f.readlines():
                        html += line
            except:
                pass
        else:
            html = self.getHTMLText(self.start_url)
            with open(self.folder + 'urlList.html','w',encoding='utf-8') as f:
                f.write(html)
        pattern = re.compile('<span data-v-68579cc8>共([\s\S]*?)页</span>')
        allPage = re.findall(pattern, html)
        print("allPage: " + allPage[0])
        for page in range(1, int(allPage[0])+1):
            url = re.sub(r'pg\d', 'pg'+ str(page), self.start_url)
            self.parseUrl(page, url, useLocal)

        #字典中的key值即为csv中列名
        dataframe=pd.DataFrame({
            '标题':self.title_list,   '详细':self.detail_list,
            '地址':self.address_list, '房型':self.shape_list,
            '价格':self.money_list,   '链接':self.href_list})
        # 将DataFrame存储为csv,index表示是否显示行名，default=True
        dataframe.to_csv(self.folder+"./蘑菇网.csv",index=False,sep=',', encoding='utf_8_sig')

    def parseUrl(self, page, url, useLocal):
        print('parse url: ' + url)
        folder = self.folder +'Detail/'
        isExists = os.path.exists(folder)
        if not isExists:  # 目录不存在，则创建
            os.makedirs(folder)
        if useLocal == True:
            # use local html file
            with open(folder + str(page)+'.html', 'r', encoding='utf-8') as f:
                html = ''
                for line in f.readlines():
                    html += line
            html = etree.HTML(html)
        else:
            # save to local
            with open(folder + str(page)+'.html', 'w', encoding='utf-8') as f:
                f.write(self.getHTMLText(url))
            html = self.getHTMLText(url)
            html = etree.HTML(html)

        # use xpath
        ele = html.xpath('//*[@id="__layout"]/div/div/div[1]/div[5]/div[1]/div[2]/div')
        for e in ele:
            la = e.xpath('.//a')
            print(len(la))
            for i in la:
                title = i.xpath('.//div/div[2]/h1/span//text()')
                shape = i.xpath('.//div/div[2]/h2[1]//text()')
                detail = i.xpath('.//div/div[2]/h2[2]//text()')
                address = i.xpath('.//div/div[2]/p//text()')
                money = i.xpath('.//div/div[3]/p[1]/span//text()')
                self.title_list.append(title[0])
                self.shape_list.append(shape[0])
                if len(detail) > 0:
                    self.detail_list.append(detail[0])
                else:
                    self.detail_list.append('')
                a = i.xpath('.//@href')
                self.href_list.append(a[0])
                self.address_list.append(address[0])
                self.money_list.append(money[0])
        if useLocal == False:
            time.sleep(0.5)   # 延时

def test():
    # TODO
    # use Selenium (don't need add Cookie)
    work = Get_infor()
    work.run(useLocal=True)
    print('OK')

if __name__=='__main__':
    start = time.time()
    test()
    end = time.time()
    print("耗时： {}".format(end-start))


import cv2
import numpy as np
import os
import pandas as pd
from configparser import ConfigParser
from generator import AugmentedImageSequence
from models.keras import ModelFactory
from tensorflow.keras import backend as kb

def get_output_layer(model, layer_name):
    # get the symbolic outputs of each "key" layer (we gave them unique names).
    layer_dict = dict([(layer.name, layer) for layer in model.layers])
    layer = layer_dict[layer_name]
    return layer

def create_cam(df_g, output_dir, image_source_dir, model, generator, class_names, targets=None):
    """
    Create a CAM overlay image for the input image

    :param targets: producs CAMs only for these labels. if left empty, CAMs are generated for all labels
    :param df_g: pandas.DataFrame with file paths
    :param output_dir: str
    :param image_source_dir: str
    :param model: keras model
    :param generator: generator.AugmentedImageSequence
    :param class_names: list of str
    """
    file_name = df_g["Path"]
    print(f"**process image: {file_name}**")

    # draw cam with labels
    img_ori = cv2.imread(filename=os.path.join(image_source_dir, file_name))
    """
    label = df_g["label"]
    if label == "Infiltrate":
        label = "Infiltration"
    """
    # label = "Lung Opacity"
    if targets is None:
        targets = class_names

    img_transformed = generator.load_image(file_name)

    # CAM overlay
    # Get the 1024 input weights to the softmax.
    class_weights = model.layers[-1].get_weights()[0]
    # print(class_weights.shape)
    # print(index)
    # print(class_weights[..., index].shape)
    final_conv_layer = get_output_layer(model, "bn")
    get_output = kb.function([model.layers[0].input], [final_conv_layer.output, model.layers[-1].output])
    [conv_outputs, predictions] = get_output([np.array([img_transformed])])
    conv_outputs = conv_outputs[0, :, :, :]

    for label in targets:
        # print(f"process label: {label}")
        index = class_names.index(label)
        output_file = f"{file_name.split('/')[-1].split('.')[-2]}_{label}.{file_name.split('.')[-1]}"
        output_path = os.path.join(output_dir, output_file)

        # Create the class activation map.
        cam = np.zeros(dtype=np.float32, shape=(conv_outputs.shape[:2]))
        for i, w in enumerate(class_weights[..., index]):
            cam += w * conv_outputs[:, :, i]
        # print(f"predictions: {predictions}")
        cam /= np.max(cam)
        # cam[np.where(cam < 0.2)] = 0
        cam = cv2.resize(cam, img_ori.shape[:2][::-1])  # flip image dimensions, see https://stackoverflow.com/questions/21248245/opencv-image-resize-flips-dimensions

        heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET)
        heatmap[np.where(cam < 0.2)] = 0
        img = heatmap * 0.5 + img_ori

        # add label & predicted probability

        cv2.putText(img, text=label + f": {str(predictions[...,index])}", org=(5, 20), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                    fontScale=0.8, color=(0, 0, 255), thickness=1)
        cv2.imwrite(output_path, img)

def main():
    # parser config
    config_file = "./config.ini"
    cp = ConfigParser()
    cp.read(config_file)

    # default config
    output_dir = cp["DEFAULT"].get("output_dir")
    base_model_name = cp["DEFAULT"].get("base_model_name")
    class_names = cp["DEFAULT"].get("class_names").split(",")
    image_source_dir = cp["DEFAULT"].get("image_source_dir")
    image_dimension = cp["TRAIN"].getint("image_dimension")

    # parse weights file path
    output_weights_name = cp["TRAIN"].get("output_weights_name")
    weights_path = os.path.join(output_dir, output_weights_name)
    best_weights_path = os.path.join(output_dir, f"best_{output_weights_name}")

    # CAM config
    # cam_list_file = cp["CAM"].get("cam_list_file")
    cam_folder = cp["CAM"].get("cam_folder")
    use_best_weights = cp["CAM"].getboolean("use_best_weights")

    print("** load model **")
    if use_best_weights:
        print("** use best weights **")
        model_weights_path = best_weights_path
    else:
        print("** use last weights **")
        model_weights_path = weights_path

    model_factory = ModelFactory()
    model = model_factory.get_model(
        class_names,
        model_name=base_model_name,
        use_base_weights=False,
        weights_path=model_weights_path)

    # print(model.summary())

    print("read contents of cam folder")
    cam_files = [f for f in os.listdir(cam_folder) if os.path.isfile(os.path.join(cam_folder, f))]
    df_images = pd.DataFrame(cam_files)
    df_images.columns = ["Path"]

    print("create a generator for loading transformed images")
    cam_sequence = AugmentedImageSequence(
        dataset_csv_file=os.path.join(output_dir, "valid.csv"),  # variable must be passed, but is not used for CAMs
        class_names=class_names,
        source_image_dir=cam_folder,
        batch_size=1,
        target_size=(image_dimension, image_dimension),
        augmenter=None,
        steps=1,
        shuffle_on_epoch_end=False,
    )

    image_output_dir = os.path.join(output_dir, "cam_output")
    if not os.path.isdir(image_output_dir):
        os.makedirs(image_output_dir)

    print("create CAM")
    df_images.apply(
        lambda g: create_cam(
            df_g=g,
            output_dir=image_output_dir,
            image_source_dir=cam_folder,
            model=model,
            generator=cam_sequence,
            class_names=class_names,
            # targets=["Lung Lesion"]
        ),
        axis=1,
    )

if __name__ == "__main__":
    main()