import os
import warnings
import pdb
import sys
import shutil
import pandas as pd
import numpy as np
import time
import multiprocessing as mp
import astropy.coordinates as astroCoords
import astropy.units as u
from kbmodpy import kbmod as kb
from astropy.io import fits
from astropy.wcs import WCS
from sklearn.cluster import DBSCAN
from skimage import measure
from analysis_utils import Interface, PostProcess

class region_search:
    """
    CLASS CURRENTLY DOES NOT WORK
    """
    def __init__(self,v_guess,radius,num_obs):
        """
        INPUT-
            v_guess : float array
                Initial object velocity guess. Given as an array or tuple.
                Algorithm will search velocities within 'radius' of 'v_guess'
            radius : float
                radius in velocity space to search, centered around 'v_guess'
            num_obs : int
                The minimum number of observations required to keep the object
        """
        self.v_guess = v_guess
        self.radius = radius
        self.num_obs = num_obs
        return

    def run_search(self, im_filepath, res_filepath, out_suffix, time_file,
                   likelihood_level=10., mjd_lims=None):
        # Initialize some values
        start = time.time()

        memory_error = False
        # Load images to search
        search,image_params = self.load_images(
            im_filepath, time_file, mjd_lims=mjd_lims)

        # Run the region search
        # Save values in image_params for use in filter_results

        print("Starting Search")
        print('---------------------------------------')
        param_headers = ("X Velocity Guess","Y Velocity Guess",
                         "Radius in velocity space")
        param_values = (*self.v_guess,self.radius)
        for header, val in zip(param_headers, param_values):
            print('%s = %.4f' % (header, val))
        results = search.region_search(
            *self.v_guess, self.radius, likelihood_level, int(self.num_obs))
        duration = image_params['times'][-1]-image_params['times'][0]
        # Convert the results to the grid formatting
        grid_results = kb.region_to_grid(results,duration)
        # Process the search results
        keep = self.process_region_results(
            search, image_params, res_filepath, likelihood_level, grid_results)
        del(search)

        # Cluster the results
        #keep = self.filter_results(keep,image_params)

        # Save the results
        self.save_results(res_filepath, out_suffix, keep)

        end = time.time()

        del(keep)
        return

    def process_region_results(
        self,search,image_params,res_filepath,likelihood_level,results):
        """
        Processes results that are output by the gpu search.
        """

        keep = {'stamps': [], 'new_lh': [], 'results': [], 'times': [],
                'lc': [], 'final_results': []}

        print('---------------------------------------')
        print("Processing Results")
        print('---------------------------------------')
        print('Starting pooling...')
        pool = mp.Pool(processes=16)
        print('Getting results...')

        psi_curves = []
        phi_curves = []
        # print(results)
        for line in results:
            psi_curve, phi_curve = search.lightcurve(line)
            psi_curves.append(np.array(psi_curve).flatten())
            phi_curve = np.array(phi_curve).flatten()
            phi_curve[phi_curve == 0.] = 99999999.
            phi_curves.append(phi_curve)

        keep_idx_results = pool.starmap_async(
            return_indices,
            zip(psi_curves, phi_curves, [j for j in range(len(psi_curves))]))
        pool.close()
        pool.join()
        keep_idx_results = keep_idx_results.get()
        if (len(keep_idx_results) < 1):
            keep_idx_results = [(0,[-1],0.)]

        if (len(keep_idx_results[0]) < 3):
            keep_idx_results = [(0, [-1], 0.)]

        for result_on in range(len(psi_curves)):

            if keep_idx_results[result_on][1][0] == -1:
                continue
            elif len(keep_idx_results[result_on][1]) < 3:
                continue
            elif keep_idx_results[result_on][2] < likelihood_level:
                continue
            else:
                keep_idx = keep_idx_results[result_on][1]
                new_likelihood = keep_idx_results[result_on][2]
                keep['results'].append(results[result_on])
                keep['new_lh'].append(new_likelihood)
                stamps = search.sci_stamps(results[result_on], 10)
                stamp_arr = np.array(
                    [np.array(stamps[s_idx]) for s_idx in keep_idx])
                keep['stamps'].append(np.sum(stamp_arr, axis=0))
                keep['lc'].append(
                    (psi_curves[result_on]/phi_curves[result_on])[keep_idx])
                keep['times'].append(image_params['mjd'][keep_idx])
        print(len(keep['results']))
        # Needed for compatibility with grid_search save functions
        keep['final_results'] = range(len(keep['results']))

        return(keep)

class run_search:
    """
    This class runs the grid search for kbmod.
    """
    def __init__(self, input_parameters):

        """
        INPUT-
            input_parameters : dictionary
                Dictionary containing input parameters. Merged with the
                defaults dictionary. MUST include 'im_filepath',
                'res_filepath', and 'time_file'. These are the filepaths to the
                image directory, results directory, and time file,
                respectively. Should contain 'v_arr', and 'ang_arr', which are
                lists containing the lower and upper velocity and angle limits.
        """

        defaults = { # Mandatory values
            'im_filepath':None, 'res_filepath':None, 'time_file':None,
            # Suggested values
            'v_arr':[92.,526.,256], 'ang_arr':[np.pi/15,np.pi/15,128],
            # Optional values
            'output_suffix':'search', 'mjd_lims':None, 'average_angle':None,
            'do_mask':True, 'mask_num_images':2, 'mask_threshold':120.,
            'lh_level':10., 'psf_val':1.4, 'num_obs':10, 'num_cores':30,
            'visit_in_filename':[0,6], 'file_format':'{0:06d}.fits',
            'sigmaG_lims':[25,75], 'chunk_size':500000, 'max_lh':1000.,
            'filter_type':'clipped_sigmaG', 'center_thresh':0.00,
            'peak_offset':[2.,2.], 'mom_lims':[35.5,35.5,2.0,0.3,0.3],
            'stamp_type':'sum', 'eps':0.03, 'gpu_filter':False,
            'do_clustering':True, 'do_stamp_filter':True,
            'clip_negative':False, 'sigmaG_filter_type':'lh'
        }
        # Make sure input_parameters contains valid input options
        for key, val in input_parameters.items():
            if key in defaults:
                defaults[key] = val
            else:
                warnings.warn('Key "{}" is not a valid option. It is being ignored.'.format(key))
        self.config = defaults
        #self.config = {**defaults, **input_parameters}
        if (self.config['im_filepath'] is None):
            raise ValueError('Image filepath not set')
        if (self.config['res_filepath'] is None):
            raise ValueError('Results filepath not set')
        if (self.config['time_file'] is None):
            raise ValueError('Time filepath not set')
        return

    def do_gpu_search(self, search, image_params, post_process):

        # Run the grid search
        # Set min and max values for angle and velocity
        if self.config['average_angle'] == None:
            average_angle = image_params['ec_angle']
        else:
            average_angle = self.config['average_angle']
        ang_min = average_angle - self.config['ang_arr'][0]
        ang_max = average_angle + self.config['ang_arr'][1]
        vel_min = self.config['v_arr'][0]
        vel_max = self.config['v_arr'][1]
        image_params['ang_lims'] = [ang_min, ang_max]
        image_params['vel_lims'] = [vel_min, vel_max]

        search_start = time.time()
        print("Starting Search")
        print('---------------------------------------')
        param_headers = ("Ecliptic Angle", "Min. Search Angle",
                         "Max Search Angle", "Min Velocity", "Max Velocity")
        param_values = (image_params['ec_angle'], *image_params['ang_lims'],
                        *image_params['vel_lims'])
        for header, val in zip(param_headers, param_values):
            print('%s = %.4f' % (header, val))
        if self.config['gpu_filter']:
            print('Using in-line GPU filtering methods', flush=True)
            self.config['sigmaG_coeff'] = post_process._find_sigmaG_coeff(
                self.config['sigmaG_lims'])
            search.gpuFilter(
                int(self.config['ang_arr'][2]), int(self.config['v_arr'][2]),
                *image_params['ang_lims'], *image_params['vel_lims'],
                int(self.config['num_obs']),
                np.array(self.config['sigmaG_lims'])/100.0,
                self.config['sigmaG_coeff'], self.config['mom_lims'],
                self.config['lh_level'])
        else:
            search.gpu(
                int(self.config['ang_arr'][2]), int(self.config['v_arr'][2]),
                *image_params['ang_lims'], *image_params['vel_lims'],
                int(self.config['num_obs']))
        print(
            'Search finished in {0:.3f}s'.format(time.time()-search_start),
            flush=True)
        return(search, image_params)

    def run_search(self):
        """
        This function serves as the highest-level python interface for starting
        a KBMOD search.
        INPUT-
            im_filepath : string
                Path to the folder containing the images to be ingested into
                KBMOD and searched over.
            res_filepath : string
                Path to the folder that will contain the results from the
                search.
            out_suffix : string
                Suffix to append to the output files. Used to differentiate
                between different searches over the same stack of images.
            time_file : string
                Path to the file containing the image times.
            lh_level : float
                Minimum acceptable likelihood level for a trajectory.
                Trajectories with likelihoods below this value will be
                discarded.
            psf_val : float
                Determines the size of the psf generated by the kbmod stack.
            mjd_lims : numpy array
                Limits the search to images taken within the limits input by
                mjd_lims.
            average_angle : float
                Overrides the ecliptic angle calculation and instead centers
                the average search around average_angle.
        """

        start = time.time()
        kb_interface = Interface()
        kb_post_process = PostProcess(self.config)

        # Load images to search
        stack,image_params = kb_interface.load_images(
            self.config['im_filepath'], self.config['time_file'],
            self.config['mjd_lims'], self.config['visit_in_filename'],
            self.config['file_format'])
        # Save values in image_params for later use
        if self.config['do_mask']:
            stack = kb_post_process.apply_mask(
                stack, mask_num_images=self.config['mask_num_images'],
                mask_threshold=self.config['mask_threshold'])
        psf = kb.psf(self.config['psf_val'])
        search = kb.stack_search(stack, psf)

        search, image_params = self.do_gpu_search(
            search, image_params, kb_post_process)
        # Load the KBMOD results into Python and apply a filter based on
        # 'filter_type'
        image_params['sigmaG_filter_type'] = self.config['sigmaG_filter_type']
        keep = kb_post_process.load_results(
            search, image_params, self.config['lh_level'],
            chunk_size=self.config['chunk_size'],
            filter_type=self.config['filter_type'],
            max_lh=self.config['max_lh'])
        if self.config['do_stamp_filter']:
            keep = kb_post_process.apply_stamp_filter(
                keep, search, center_thresh=self.config['center_thresh'],
                peak_offset=self.config['peak_offset'],
                mom_lims=self.config['mom_lims'],
                stamp_type=self.config['stamp_type'])
        if self.config['do_clustering']:
            keep = kb_post_process.apply_clustering(keep, image_params)
        keep = kb_post_process.get_all_stamps(keep, search)
        del(search)
        # Save the results
        kb_interface.save_results(
        self.config['res_filepath'], self.config['output_suffix'], keep)
        end = time.time()
        del(keep)
        print("Time taken for patch: ", end-start)

import numpy,import astropy,from astropy


import numpy as np
import tensorflow as tf
from keras import backend as K
from keras.backend import epsilon,clip,mean
from nets.ious import box_ciou

import numpy as np
from keras.backend import exp
from keras.layers import Softmax

import math

def DRloss(targets, logits):
    pos_lambda = 1
    margin = 0.5
    neg_lambda = 0.1 / math.log(3.5)
    L = 6.
    tau = 4.
    # 标注相关信息
    num_classes = logits.shape[1]
    dtype = targets.dtype
    device = targets.device
    class_range = np.arange(1, num_classes + 1, dtype=dtype, device=device).unsqueeze(0)
    t = targets.unsqueeze(1)
    # 获得正负样本id
    pos_ind = (t == class_range)
    neg_ind = (t != class_range) * (t >= 0)
    # 概率p使用sigmoid求得
    pos_prob = logits[pos_ind].sigmoid()
    neg_prob = logits[neg_ind].sigmoid()
    # 对应于式(3.9)
    neg_q = Softmax(neg_prob/neg_lambda, dim=0)
    neg_dist = np.sum(neg_q * neg_prob)
    # 论文中提到，如果图像中没有正样本，则正样本的P使用1代替
    if pos_prob.numel() > 0:
        # 对应于式(3.9)
        pos_q = Softmax(-pos_prob/pos_lambda, dim=0)
        pos_dist = np.sum(pos_q * pos_prob)
        # 对应于式(3.12)
        loss = tau*np.log(1.+exp(L*(neg_dist - pos_dist+margin)))/L
    else:
        # 对应于式(3.12)
        loss = tau*np.log(1.+exp(L*(neg_dist - 1. + margin)))/L
    return loss

#---------------------------------------------------#
#   平滑标签
#---------------------------------------------------#
def _smooth_labels(y_true, label_smoothing):
    num_classes = tf.cast(K.shape(y_true)[-1], dtype=K.floatx())

    label_smoothing = K.constant(label_smoothing, dtype=K.floatx())

    return y_true * (1.0 - label_smoothing) + label_smoothing / num_classes

def focal(alpha=0.25, gamma=2.0):
    def _focal(y_true, y_pred):
        # y_true [batch_size, num_anchor, num_classes+1]
        # y_pred [batch_size, num_anchor, num_classes]
        labels         = y_true[:, :, :-1]
        anchor_state   = y_true[:, :, -1]  # -1 是需要忽略的, 0 是背景, 1 是存在目标
        classification = y_pred

        # 找出存在目标的先验框
        indices_for_object        = K.where(K.equal(anchor_state, 1))
        labels_for_object         = K.gather_nd(labels, indices_for_object)
        classification_for_object = K.gather_nd(classification, indices_for_object)

        # 计算每一个先验框应该有的权重
        alpha_factor_for_object = K.ones_like(labels_for_object) * alpha
        alpha_factor_for_object = K.where(K.equal(labels_for_object, 1), alpha_factor_for_object, 1 - alpha_factor_for_object)
        focal_weight_for_object = K.where(K.equal(labels_for_object, 1), 1 - classification_for_object, classification_for_object)
        focal_weight_for_object = alpha_factor_for_object * focal_weight_for_object ** gamma

        # 将权重乘上所求得的交叉熵
        cls_loss_for_object = focal_weight_for_object * K.binary_crossentropy(labels_for_object, classification_for_object)

        # 找出实际上为背景的先验框
        indices_for_back        = K.where(K.equal(anchor_state, 0))
        labels_for_back         = K.gather_nd(labels, indices_for_back)
        classification_for_back = K.gather_nd(classification, indices_for_back)

        # 计算每一个先验框应该有的权重
        alpha_factor_for_back = K.ones_like(labels_for_back) * (1 - alpha)
        focal_weight_for_back = classification_for_back
        focal_weight_for_back = alpha_factor_for_back * focal_weight_for_back ** gamma

        # 将权重乘上所求得的交叉熵
        cls_loss_for_back = focal_weight_for_back * K.binary_crossentropy(labels_for_back, classification_for_back)

        # 标准化，实际上是正样本的数量
        normalizer = tf.where(K.equal(anchor_state, 1))
        normalizer = K.cast(K.shape(normalizer)[0], K.floatx())
        normalizer = K.maximum(K.cast_to_floatx(1.0), normalizer)

        # 将所获得的loss除上正样本的数量
        cls_loss_for_object = K.sum(cls_loss_for_object)
        cls_loss_for_back = K.sum(cls_loss_for_back)

        # 总的loss
        loss = (cls_loss_for_object + cls_loss_for_back)/normalizer

        return loss
    return _focal

def binary_focal_loss(true_label,probs ):
    gamma = 2
    alpha = 0.25
    alpha = tf.constant(alpha, dtype=tf.float32)
    gamma = tf.constant(gamma, dtype=tf.float32)

    epsilon = 1.e-8
    # 得到y_true和y_pred
    y_true = tf.one_hot(true_label, 2)
    # probs = tf.nn.sigmoid(logits)
    y_pred = tf.clip_by_value(probs, epsilon, 1. - epsilon)
    # 得到调节因子weight和alpha
    ## 先得到y_true和1-y_true的概率【这里是正负样本的概率都要计算哦！】
    p_t = y_true * y_pred \
          + (tf.ones_like(y_true) - y_true) * (tf.ones_like(y_true) - y_pred)
    ## 然后通过p_t和gamma得到weight
    weight = tf.pow((tf.ones_like(y_true) - p_t), gamma)
    ## 再得到alpha，y_true的是alpha，那么1-y_true的是1-alpha
    alpha_t = y_true * alpha + (tf.ones_like(y_true) - y_true) * (1 - alpha)
    # 最后就是论文中的公式，相当于：- alpha * (1-p_t)^gamma * log(p_t)
    focal_loss = - alpha_t * weight * tf.log(p_t)
    return tf.reduce_mean(focal_loss)

#---------------------------------------------------#
#   将预测值的每个特征层调成真实值
#---------------------------------------------------#
def yolo_head(feats, anchors, num_classes, input_shape, calc_loss=False):
    num_anchors = len(anchors)
    #---------------------------------------------------#
    #   [1, 1, 1, num_anchors, 2]
    #---------------------------------------------------#
    anchors_tensor = K.reshape(K.constant(anchors), [1, 1, 1, num_anchors, 2])

    #---------------------------------------------------#
    #   获得x，y的网格
    #   (13, 13, 1, 2)
    #---------------------------------------------------#
    grid_shape = K.shape(feats)[1:3]
    grid_y = K.tile(K.reshape(K.arange(0, stop=grid_shape[0]), [-1, 1, 1, 1]),
        [1, grid_shape[1], 1, 1])
    grid_x = K.tile(K.reshape(K.arange(0, stop=grid_shape[1]), [1, -1, 1, 1]),
        [grid_shape[0], 1, 1, 1])
    grid = K.concatenate([grid_x, grid_y])
    grid = K.cast(grid, K.dtype(feats))

    #---------------------------------------------------#
    #   将预测结果调整成(batch_size,13,13,3,85)
    #   85可拆分成4 + 1 + 80
    #   4代表的是中心宽高的调整参数
    #   1代表的是框的置信度
    #   80代表的是种类的置信度
    #---------------------------------------------------#
    feats = K.reshape(feats, [-1, grid_shape[0], grid_shape[1], num_anchors, num_classes + 5])

    #---------------------------------------------------#
    #   将预测值调成真实值
    #   box_xy对应框的中心点
    #   box_wh对应框的宽和高
    #---------------------------------------------------#
    box_xy = (K.sigmoid(feats[..., :2]) + grid) / K.cast(grid_shape[::-1], K.dtype(feats))
    box_wh = K.exp(feats[..., 2:4]) * anchors_tensor / K.cast(input_shape[::-1], K.dtype(feats))
    box_confidence = K.sigmoid(feats[..., 4:5])
    box_class_probs = K.sigmoid(feats[..., 5:])

    #---------------------------------------------------------------------#
    #   在计算loss的时候返回grid, feats, box_xy, box_wh
    #   在预测的时候返回box_xy, box_wh, box_confidence, box_class_probs
    #---------------------------------------------------------------------#
    if calc_loss == True:
        return grid, feats, box_xy, box_wh
    return box_xy, box_wh, box_confidence, box_class_probs

#---------------------------------------------------#
#   用于计算每个预测框与真实框的iou
#---------------------------------------------------#
def box_iou(b1, b2):
    # 13,13,3,1,4
    # 计算左上角的坐标和右下角的坐标
    b1 = K.expand_dims(b1, -2)
    b1_xy = b1[..., :2]
    b1_wh = b1[..., 2:4]
    b1_wh_half = b1_wh/2.
    b1_mins = b1_xy - b1_wh_half
    b1_maxes = b1_xy + b1_wh_half

    # 1,n,4
    # 计算左上角和右下角的坐标
    b2 = K.expand_dims(b2, 0)
    b2_xy = b2[..., :2]
    b2_wh = b2[..., 2:4]
    b2_wh_half = b2_wh/2.
    b2_mins = b2_xy - b2_wh_half
    b2_maxes = b2_xy + b2_wh_half

    # 计算重合面积
    intersect_mins = K.maximum(b1_mins, b2_mins)
    intersect_maxes = K.minimum(b1_maxes, b2_maxes)
    intersect_wh = K.maximum(intersect_maxes - intersect_mins, 0.)
    intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]
    b1_area = b1_wh[..., 0] * b1_wh[..., 1]
    b2_area = b2_wh[..., 0] * b2_wh[..., 1]
    iou = intersect_area / (b1_area + b2_area - intersect_area)

    return iou

#---------------------------------------------------#
#   loss值计算
#---------------------------------------------------#
def yolo_loss(args, anchors, num_classes, ignore_thresh=.5, label_smoothing=0.1, print_loss=False, normalize=True):
    # 一共有两层
    num_layers = len(anchors)//3

    #---------------------------------------------------------------------------------------------------#
    #   将预测结果和实际ground truth分开，args是[*model_body.output, *y_true]
    #   y_true是一个列表，包含两个特征层，shape分别为(m,13,13,3,85),(m,26,26,3,85)
    #   yolo_outputs是一个列表，包含两个特征层，shape分别为(m,13,13,3,85),(m,26,26,3,85)
    #---------------------------------------------------------------------------------------------------#
    y_true = args[num_layers:]
    yolo_outputs = args[:num_layers]

    #-----------------------------------------------------------#
    #   13x13的特征层对应的anchor是[81,82], [135,169], [344,319]
    #   26x26的特征层对应的anchor是[23,27], [37,58], [81,82]
    #-----------------------------------------------------------#
    anchor_mask = [[6,7,8], [3,4,5], [0,1,2]] if num_layers==3 else [[3,4,5], [1,2,3]]

    # 得到input_shpae为416,416
    input_shape = K.cast(K.shape(yolo_outputs[0])[1:3] * 32, K.dtype(y_true[0]))

    loss = 0
    num_pos = 0
    #-----------------------------------------------------------#
    #   取出每一张图片
    #   m的值就是batch_size
    #-----------------------------------------------------------#
    m = K.shape(yolo_outputs[0])[0]
    mf = K.cast(m, K.dtype(yolo_outputs[0]))

    #---------------------------------------------------------------------------------------------------#
    #   y_true是一个列表，包含两个特征层，shape分别为(m,13,13,3,85),(m,26,26,3,85)
    #   yolo_outputs是一个列表，包含两个特征层，shape分别为(m,13,13,3,85),(m,26,26,3,85)
    #---------------------------------------------------------------------------------------------------#
    for l in range(num_layers):
        #-----------------------------------------------------------#
        #   以第一个特征层(m,13,13,3,85)为例子
        #   取出该特征层中存在目标的点的位置。(m,13,13,3,1)
        #-----------------------------------------------------------#
        object_mask = y_true[l][..., 4:5]#[…, 4: 5]指的是该框的置信度,用1或0表示https://www.cnblogs.com/learningcaiji/p/14077315.html

        #-----------------------------------------------------------#
        #   取出其对应的种类(m,13,13,3,80)
        #-----------------------------------------------------------#
        true_class_probs = y_true[l][..., 5:]
        if label_smoothing:
            true_class_probs = _smooth_labels(true_class_probs, label_smoothing)

        #-----------------------------------------------------------#
        #   将yolo_outputs的特征层输出进行处理、获得四个返回值
        #   其中：
        #   grid        (13,13,1,2) 网格坐标
        #   raw_pred    (m,13,13,3,85) 尚未处理的预测结果
        #   pred_xy     (m,13,13,3,2) 解码后的中心坐标
        #   pred_wh     (m,13,13,3,2) 解码后的宽高坐标
        #-----------------------------------------------------------#
        grid, raw_pred, pred_xy, pred_wh = yolo_head(yolo_outputs[l],
             anchors[anchor_mask[l]], num_classes, input_shape, calc_loss=True)

        #-----------------------------------------------------------#
        #   pred_box是解码后的预测的box的位置
        #   (m,13,13,3,4)
        #-----------------------------------------------------------#
        pred_box = K.concatenate([pred_xy, pred_wh])

        #-----------------------------------------------------------#
        #   找到负样本群组，第一步是创建一个数组，[]
        #-----------------------------------------------------------#
        ignore_mask = tf.TensorArray(K.dtype(y_true[0]), size=1, dynamic_size=True)
        object_mask_bool = K.cast(object_mask, 'bool')

        #-----------------------------------------------------------#
        #   对每一张图片计算ignore_mask
        #-----------------------------------------------------------#
        def loop_body(b, ignore_mask):
            #-----------------------------------------------------------#
            #   取出n个真实框：n,4
            #-----------------------------------------------------------#
            true_box = tf.boolean_mask(y_true[l][b,...,0:4], object_mask_bool[b,...,0])
            #-----------------------------------------------------------#
            #   计算预测框与真实框的iou
            #   pred_box    13,13,3,4 预测框的坐标
            #   true_box    n,4 真实框的坐标
            #   iou         13,13,3,n 预测框和真实框的iou
            #-----------------------------------------------------------#
            iou = box_iou(pred_box[b], true_box)

            #-----------------------------------------------------------#
            #   best_iou    13,13,3 每个特征点与真实框的最大重合程度
            #-----------------------------------------------------------#
            best_iou = K.max(iou, axis=-1)

            #-----------------------------------------------------------#
            #   判断预测框和真实框的最大iou小于ignore_thresh
            #   则认为该预测框没有与之对应的真实框
            #   该操作的目的是：
            #   忽略预测结果与真实框非常对应特征点，因为这些框已经比较准了
            #   不适合当作负样本，所以忽略掉。
            #-----------------------------------------------------------#
            ignore_mask = ignore_mask.write(b, K.cast(best_iou<ignore_thresh, K.dtype(true_box)))
            return b+1, ignore_mask

        #-----------------------------------------------------------#
        #   在这个地方进行一个循环、循环是对每一张图片进行的
        #-----------------------------------------------------------#
        _, ignore_mask = K.control_flow_ops.while_loop(lambda b,*args: b<m, loop_body, [0, ignore_mask])

        #-----------------------------------------------------------#
        #   ignore_mask用于提取出作为负样本的特征点
        #   (m,13,13,3)
        #-----------------------------------------------------------#
        ignore_mask = ignore_mask.stack()
        #   (m,13,13,3,1)
        ignore_mask = K.expand_dims(ignore_mask, -1)

        #-----------------------------------------------------------#
        #   真实框越大，比重越小，小框的比重更大。
        #-----------------------------------------------------------#
        box_loss_scale = 2 - y_true[l][...,2:3]*y_true[l][...,3:4]  #预测值

        #-----------------------------------------------------------#
        #   计算Ciou loss
        #-----------------------------------------------------------#
        raw_true_box = y_true[l][...,0:4]
        ciou = box_ciou(pred_box, raw_true_box)
        ciou_loss = object_mask * box_loss_scale * (1 - ciou)

        #------------------------------------------------------------------------------#
        #   如果该位置本来有框，那么计算1与置信度的交叉熵
        #   如果该位置本来没有框，那么计算0与置信度的交叉熵
        #   在这其中会忽略一部分样本，这些被忽略的样本满足条件best_iou<ignore_thresh
        #   该操作的目的是：
        #   忽略预测结果与真实框非常对应特征点，因为这些框已经比较准了
        #   不适合当作负样本，所以忽略掉。
        #------------------------------------------------------------------------------#

        #object_mask代表的是True
        #列表中的三个点代表前面所有的维数
        # print('object_ mask',object_mask.shape)
        # print('raw',raw_pred[...,4:5].shape)

        #confidence_loss，实际存在的框，预测结果中置信度的值与1对比；实际不存在的框，预测结果中置信度的值与0对比，该部分要去除被忽略的不包含目标的框。
        confidence_loss = object_mask * K.binary_crossentropy(object_mask, raw_pred[...,4:5],from_logits=True)+ \
            (1-object_mask) * K.binary_crossentropy(object_mask, raw_pred[...,4:5],from_logits=True) * ignore_mask
        # confidence_loss = binary_focal_loss(object_mask,raw_pred[...,4:5])
        # print('confidence_loss',confidence_loss)
        #置信度loss
        '''正样本有坐标，置信度和类别损失函数，而负样本只有置信度损失函数'''

        '''
                alpha = 0.25
        gamma = 2
        alpha_factor = K.ones_like(object_mask) * alpha
        alpha_factor = tf.where(K.equal(object_mask, 1), alpha_factor, 1 - alpha_factor)
        focal_weight = tf.where(K.equal(object_mask, 1), 1 - raw_pred[..., 4:5], raw_pred[..., 4:5])
        focal_weight = alpha_factor * focal_weight ** gamma
        confidence_loss = focal_weight * K.binary_crossentropy(object_mask, raw_pred[..., 4:5],from_logits=True)
        '''

        class_loss = object_mask * K.binary_crossentropy(true_class_probs, raw_pred[...,5:],from_logits=True) #分类loss

        location_loss = K.sum(ciou_loss) #回归框
        confidence_loss = K.sum(confidence_loss)
        class_loss = K.sum(class_loss)
        #-----------------------------------------------------------#
        #   计算正样本数量
        #-----------------------------------------------------------#
        num_pos += tf.maximum(K.sum(K.cast(object_mask, tf.float32)), 1)
        loss += location_loss + confidence_loss + class_loss
        # if print_loss:
        #   loss = tf.Print(loss, [loss, location_loss, confidence_loss, class_loss, K.sum(ignore_mask)], message='loss: ')

    if normalize:
        loss = loss / num_pos
    else:
        loss = loss / mf
    return loss


#!/usr/bin/env python

"""Variational Autoencoder class

@author: Matt Whiteway, June 2017
VAE class implements a variational autoencoder

"""

from __future__ import print_function
from __future__ import division

import os
import numpy as np
import tensorflow as tf

import sys
sys.path.append('..')
import utils.exceptions as exc

class VAE(object):
    """Variational Autoencoder class

    Attributes:
        layers_encoder (list of ints): size of each layer in encoder, excluding
            input layer
        input_size (int): size of input
        layer_latent (int): size of latent layer
        layers_decoder (list of ints): size of each layer in decoder, including
            output layer
        num_lvs (int): size of latent layer
        num_layers_enc (int): number of layers in encoder, including input
            layer
        num_layers_dec (int): number of layers in decoder, including latent
            layer and output layer
        act_func (str): activation function for network layers

        weights_enc (list of tf.Variable): weights of the encoding network
        biases_enc (list of tf.Variable): biases of the encoding network
        weights_mean (tf.Variable): weights from encoding network to latent
            variable distribution mean
        biases_mean (tf.Variable): biases from encoding network to latent
            variable distribution mean
        weights_log_var (tf.Variable): weights from encoding network to log of
            latent variable distribution variance
        biases_log_var (tf.Variable): biases from encoding network to log of
            latent variable distribution variance
        weights_dec (list of tf.Variable): weights of the decoding network
        biases_dec (list of tf.Variable): biases of the decoding network

        x (tf placeholder): ph for input to model
        z_mean (tf op): mean value for each latent variable
        z_log_var (tf op): log of the variance for each latent variable
        z (tf op): sample value of latent variable
        eps (tf placeholder): ph for N(0,1) input to stochastic layer
        x_recon (tf op): reconstructed input values
        cost (tf op): evaluates the cost function of the network

        learning_rate (float): global learning rate used by gradient descent
            optimizers
        train_step (tf op): evaluates one training step using the specified
            cost function and learning algorithm

        graph (tf.Graph): dataflow graph for the network
        saver (tf.train.Saver): for saving and restoring variables
        merge_summaries (tf op): op that merges all summary ops
        init (tf op): op that initializes global variables in graph

    """

    def __init__(
            self,
            layers_encoder=None,
            layer_latent=None,
            layers_decoder=None,
            act_func='relu',
            learning_rate=1e-3):
        """Constructor for VAE class

        Args:
            layers_encoder (list of ints): size of each layer in encoder,
                including input layer
            layer_latent (int): size of latent layer
            layers_decoder (list of ints): size of each layer in decoder,
                including output layer
            act_func (str): activation function for network layers
                ['relu'] | 'sigmoid' | 'tanh' | 'linear' | 'softplus' | 'elu'
            learning_rate (scalar): global learning rate for gradient descent
                methods

        Raises:
            InputError if layers_encoder is not specified
            InputError if layers_latent is not specified
            InputError if layers_decoder is not specified
            InputError if act_func is not a valid string

        """

        # input checking
        if layers_encoder is None:
            raise exc.InputError('Must specify layer sizes for encoder')
        if layer_latent is None:
            raise exc.InputError('Must specify number of latent dimensions')
        if layers_decoder is None:
            raise exc.InputError('Must specify layer sizes for decoder')

        self.input_size = layers_encoder[0]
        self.layers_encoder = layers_encoder[1:]
        self.layer_latent = layer_latent
        self.layers_decoder = layers_decoder

        if act_func == 'relu':
            self.act_func = tf.nn.relu
        elif act_func == 'sigmoid':
            self.act_func = tf.sigmoid
        elif act_func == 'tanh':
            self.act_func = tf.tanh
        elif act_func == 'linear':
            self.act_func = tf.identity
        elif act_func == 'softplus':
            self.act_func = tf.nn.softplus
        elif act_func == 'elu':
            self.act_func = tf.nn.elu
        else:
            raise exc.InputError('Invalid activation function')

        self.learning_rate = learning_rate

        # define useful constants
        self.num_lvs = self.layer_latent
        self.num_layers_enc = len(self.layers_encoder)
        self.num_layers_dec = len(self.layers_decoder)

        # for saving and restoring models
        self.graph = tf.Graph()  # must be initialized before graph creation

        # build model graph
        with self.graph.as_default():

            # define pipeline for feeding data into model
            with tf.variable_scope('data'):
                self._initialize_data_pipeline()

            # initialize weights and create encoder model
            with tf.variable_scope('encoder'):
                self._define_recognition_network()

            # initialize weights and create decoder model
            with tf.variable_scope('decoder'):
                self._define_generator_network()

            # define loss function
            with tf.variable_scope('loss'):
                self._define_loss()

            # define optimizer
            with tf.variable_scope('optimizer'):
                self._define_optimizer()

            # add additional ops
            # for saving and restoring models
            self.saver = tf.train.Saver()  # must be init after var creation
            # collect all summaries into a single op
            self.merge_summaries = tf.summary.merge_all()
            # add variable initialization op to graph
            self.init = tf.global_variables_initializer()

    def _initialize_data_pipeline(self):
        """Create placeholders for input and random values"""

        self.x = tf.placeholder(
            dtype=tf.float32,
            shape=[None, self.input_size],
            name='input_ph')
        self.eps = tf.placeholder(
            dtype=tf.float32,
            shape=[None, self.num_lvs],
            name='rand_ph')

    def _define_recognition_network(self):
        """
        Create a recognition network to transform inputs into its latent
        representation
        """

        # push data through the encoding function to determine mean and std
        # of latent vars
        self.weights_enc = []
        self.biases_enc = []
        z_enc = [self.x]
        for layer in range(self.num_layers_enc):
            with tf.variable_scope(str('layer_%01i' % layer)):

                # initialize weights
                if layer == 0:
                    in_size = self.input_size
                else:
                    in_size = self.layers_encoder[layer - 1]
                out_size = self.layers_encoder[layer]
                self.weights_enc.append(tf.get_variable(
                    shape=[in_size, out_size],
                    name='weights',
                    initializer=tf.truncated_normal_initializer(stddev=0.1)))

                # initialize biases
                self.biases_enc.append(tf.get_variable(
                    initializer=tf.zeros(shape=[1, out_size]),
                    name='biases'))

                # calculate layer activations
                pre = tf.add(
                    tf.matmul(z_enc[layer], self.weights_enc[layer]),
                    self.biases_enc[layer])
                post = self.act_func(pre)
                z_enc.append(post)

                # save summaries of layer activations
                tf.summary.histogram('pre_act', pre)
                tf.summary.histogram('post_act', post)

        with tf.variable_scope('latent_layer'):

            with tf.variable_scope('means'):
                # initialize weights/biases for means of stochastic layer
                self.weights_mean = tf.get_variable(
                    shape=[self.layers_encoder[-1], self.num_lvs],
                    name='weights',
                    initializer=tf.truncated_normal_initializer(stddev=0.1))
                self.biases_mean = tf.get_variable(
                    initializer=tf.zeros(shape=[1, self.num_lvs]),
                    name='biases')
                # weights to estimate mean of normally distributed latent vars
                self.z_mean = tf.add(
                    tf.matmul(z_enc[-1], self.weights_mean), self.biases_mean,
                    name='z_means')

            with tf.variable_scope('log_vars'):
                # initialize weights/biases for log vars of stochastic layer
                self.weights_log_var = tf.get_variable(
                    shape=[self.layers_encoder[-1], self.num_lvs],
                    name='weights',
                    initializer=tf.truncated_normal_initializer(stddev=0.1))
                self.biases_log_var = tf.get_variable(
                    initializer=tf.zeros(shape=[1, self.num_lvs]),
                    name='biases')
                # estimating log of the variance is easier since the latent
                # loss has a log determinant term
                self.z_log_var = tf.add(
                    tf.matmul(z_enc[-1], self.weights_log_var),
                    self.biases_log_var,
                    name='z_log_vars')

            # transform estimated mean and log variance into a sampled value
            # of the latent state using z = mu + sigma*epsilon
            self.z = tf.add(
                self.z_mean,
                tf.multiply(tf.sqrt(tf.exp(self.z_log_var)), self.eps))

            # save summaries of means and log_vars
            tf.summary.histogram('means', self.z_mean)
            tf.summary.histogram('log_vars', self.z_log_var)

    def _define_generator_network(self):
        """
        Create a generator network to transform a random sample
        in the latent space into an image
        """

        self.weights_dec = []
        self.biases_dec = []
        z_dec = [self.z]
        for layer in range(self.num_layers_dec):
            with tf.variable_scope(str('layer_%01i' % layer)):

                # initialize weights
                if layer == 0:
                    in_size = self.num_lvs
                else:
                    in_size = self.layers_decoder[layer - 1]
                out_size = self.layers_decoder[layer]
                self.weights_dec.append(tf.get_variable(
                    shape=[in_size, out_size],
                    name='weights',
                    initializer=tf.truncated_normal_initializer(stddev=0.1)))

                # initialize biases
                self.biases_dec.append(tf.get_variable(
                    initializer=tf.zeros(shape=[1, out_size]),
                    name='biases'))

                # calculate layer activations
                pre = tf.add(
                    tf.matmul(z_dec[layer], self.weights_dec[layer]),
                    self.biases_dec[layer])
                post = self.act_func(pre)
                z_dec.append(post)

                # save summaries of layer activations
                tf.summary.histogram('pre_act', pre)
                tf.summary.histogram('post_act', post)

        # define this for easier access later
        self.x_recon = z_dec[-1]

    def _define_loss(self):
        """Define loss function that will be used to optimize model params"""

        # define reconstruction loss
        loss_recon = 0.5 * tf.reduce_sum(tf.square(self.x_recon - self.x), 1)

        # define latent loss
        loss_latent = 0.5 * tf.reduce_sum(tf.exp(self.z_log_var)
                                          + tf.square(self.z_mean)
                                          - 1 - self.z_log_var, 1)

        # define cost
        self.cost = tf.reduce_mean(loss_recon + loss_latent)
        # save summaries of cost
        tf.summary.scalar('cost', self.cost)

    def _define_optimizer(self):
        """Define one step of the optimization routine"""
        self.train_step = tf.train.AdamOptimizer(self.learning_rate). \
            minimize(self.cost)

    def train(
            self,
            sess,
            data=None,
            batch_size=128,
            epochs_training=10,
            epochs_disp=None,
            epochs_ckpt=None,
            epochs_summary=None,
            output_dir=None):
        """Network training

        Args:
            sess (tf.Session object): current session object to run graph
            data (DataReader object): input to network
            batch_size (int, optional): batch size used by the gradient
                descent-based optimizers
            epochs_training (int, optional): number of epochs for gradient
                descent-based optimizers
            epochs_disp (int, optional): number of epochs between updates to
                the console
            epochs_ckpt (int, optional): number of epochs between saving
                checkpoint files
            epochs_summary (int, optional): number of epochs between saving
                network summary information
            output_dir (string, optional): absolute path for saving checkpoint
                files and summary files; must be present if either epochs_ckpt
                or epochs_summary is not 'None'.

        Returns:
            None

        Raises:
            InputError: If epochs_ckpt is not None and output_dir is None
            InputError: If epochs_summary is not None and output_dir is None

        """

        # check input
        if data is None:
            raise exc.InputError('data reader must be specified')
        if epochs_ckpt is not None and output_dir is None:
            raise exc.InputError('output_dir must be specified to save model')
        if epochs_summary is not None and output_dir is None:
            raise exc.InputError('output_dir must be specified to save ' +
                                 'summaries')

        # initialize file writers
        if epochs_summary is not None:
            test_writer = tf.summary.FileWriter(
                os.path.join(output_dir, 'summaries', 'test'),
                sess.graph)

        # begin training
        with self.graph.as_default():

            num_batches = int(data.train.num_examples / batch_size)

            # start training loop
            for epoch in range(epochs_training):

                for batch in range(num_batches):

                    # get batch of data for this training step
                    x = data.train.next_batch(batch_size)

                    # draw random samples for latent layer
                    eps = np.random.normal(size=(batch_size, self.num_lvs))

                    # one step of optimization routine
                    sess.run(
                        self.train_step,
                        feed_dict={self.x: x[0], self.eps: eps})

                # print training updates
                if epochs_disp is not None and epoch % epochs_disp == 0:

                    # print updates using test set
                    x = data.test.next_batch(data.test.num_examples)
                    eps = np.random.normal(
                        size=(data.test.num_examples, self.num_lvs))
                    cost = sess.run(
                        self.cost,
                        feed_dict={self.x: x[0], self.eps: eps})
                    print('Epoch %03d:' % epoch)
                    print('   test cost = %2.5f' % cost)

                # save model checkpoints
                if epochs_ckpt is not None and epoch % epochs_ckpt == 0:
                    save_file = os.path.join(
                        output_dir, 'ckpts',
                        str('epoch_%05g.ckpt' % epoch))
                    self.save_model(sess, save_file)

                # save model summaries
                if epochs_summary is not None and \
                        epoch % epochs_summary == 0:
                    # output summaries using test set
                    x = data.test.next_batch(data.test.num_examples)
                    eps = np.random.normal(
                        size=(data.test.num_examples, self.num_lvs))
                    summary = sess.run(
                        self.merge_summaries,
                        feed_dict={self.x: x[0], self.eps: eps})
                    test_writer.add_summary(summary, epoch)

    def train_iters(
            self,
            sess,
            data=None,
            batch_size=128,
            iters_training=1000,
            iters_disp=None,
            iters_ckpt=None,
            iters_summary=None,
            output_dir=None):
        """
        Network training by specifying number of iterations rather than
        epochs. Used for easily generating sample outputs during training

        Args:
            sess (tf.Session object): current session object to run graph
            data (DataReader object): input to network
            batch_size (int, optional): batch size used by the gradient
                descent-based optimizers
            iters_training (int, optional): number of iters for gradient
                descent-based optimizers
            iters_disp (int, optional): number of iters between updates to
                the console
            iters_ckpt (int, optional): number of iters between saving
                checkpoint files
            iters_summary (int, optional): number of iters between saving
                network summary information
            output_dir (string, optional): absolute path for saving checkpoint
                files and summary files; must be present if either iters_ckpt
                or iters_summary is not 'None'.

        Returns:
            None

        Raises:
            InputError: If iters_ckpt is not None and output_dir is None
            InputError: If iters_summary is not None and output_dir is None

        """

        # check input
        if data is None:
            raise exc.InputError('data reader must be specified')
        if iters_ckpt is not None and output_dir is None:
            raise exc.InputError('output_dir must be specified to save model')
        if iters_summary is not None and output_dir is None:
            raise exc.InputError('output_dir must be specified to save ' +
                                 'summaries')

        # initialize file writers
        if iters_summary is not None:
            test_writer = tf.summary.FileWriter(
                os.path.join(output_dir, 'summaries', 'test'),
                sess.graph)

        # begin training
        with self.graph.as_default():

            # start training loop
            for iter_ in range(iters_training):

                # get batch of data for this training step
                x = data.train.next_batch(batch_size)

                # draw random samples for latent layer
                eps = np.random.normal(size=(batch_size, self.num_lvs))

                # one step of optimization routine
                sess.run(
                    self.train_step,
                    feed_dict={self.x: x[0], self.eps: eps})

                # print training updates
                if iters_disp is not None and iter_ % iters_disp == 0:
                    # print updates using test set
                    x = data.test.next_batch(data.test.num_examples)
                    eps = np.random.normal(
                        size=(data.test.num_examples, self.num_lvs))
                    cost = sess.run(
                        self.cost,
                        feed_dict={self.x: x[0], self.eps: eps})
                    print('Iter %03d:' % iter_)
                    print('   test cost = %2.5f' % cost)

                # save model checkpoints
                if iters_ckpt is not None and iter_ % iters_ckpt == 0:
                    save_file = os.path.join(
                        output_dir, 'ckpts',
                        str('epoch_%05g.ckpt' % iter_))
                    self.save_model(sess, save_file)

                # save model summaries
                if iters_summary is not None and iter_ % iters_summary == 0:
                    # output summaries using test set
                    x = data.test.next_batch(data.test.num_examples)
                    eps = np.random.normal(
                        size=(data.test.num_examples, self.num_lvs))
                    summary = sess.run(
                        self.merge_summaries,
                        feed_dict={self.x: x[0], self.eps: eps})
                    test_writer.add_summary(summary, iter_)

    def generate(self, sess, z_mean=None):
        """Sample the network and generate an image

        If z_mean is None, a random point is generated using the prior in
        the latent space, else z_mean is used as the point in latent space
        """

        if z_mean is None:
            z_mean = np.random.normal(size=(1, self.num_lvs))

        return sess.run(self.x_recon, feed_dict={self.z: z_mean})

    def recognize(self, sess, x):
        """Transform a given input into its latent representation"""
        return sess.run(self.z_mean, feed_dict={self.x: x})

    def reconstruct(self, sess, x, eps):
        """Transform a given input into its reconstruction"""
        return sess.run(self.x_recon, feed_dict={self.x: x, self.eps: eps})

    def save_model(self, sess, save_file):
        """Save model parameters

        Args:
            sess (tf.Session object): current session object to run graph
            save_file (str): full path to output file

        """

        if not os.path.isdir(os.path.dirname(save_file)):
            os.makedirs(os.path.dirname(save_file))

        self.saver.save(sess, save_file)
        print('Model saved to %s' % save_file)

    def load_model(self, sess, save_file):
        """Load previously saved model parameters

        Args:
            sess (tf.Session object): current session object to run graph
            save_file (str): full path to saved model

        """

        if not os.path.isfile(save_file + '.meta'):
            raise exc.InputError(str('%s is not a valid filename' % save_file))

        self.saver.restore(sess, save_file)
        print('Model loaded from %s' % save_file)


from os import stat
from typing import Dict, List, Tuple
from treehmm import initHMM, baumWelch
from treehmm import fwd_seq_gen, forward
from TreeHMM4Glycan.Glycan import Glycan
import csv
from scipy.linalg import block_diag
import numpy as np
import re
import argparse
from scipy.sparse.csr import csr_matrix
from scipy.special import logsumexp
import logging
import sys

#
# Method read a file and then retunr a Dict of iupac names
#
def get_iupcas(iupac_name_file:str) -> Dict[int, str]:
    iupacs = {}
    with open(iupac_name_file) as file_in:
        csv_reader = csv.reader(file_in)
        for idx,row in enumerate(csv_reader):
            if idx == 0:
                continue
            id = int(row[0])
            # remove right most part'(a1-sp14'
            iupac = re.split(r"\([^\)]*$", row[1], 1)[0]
            iupacs[id] = iupac
    return iupacs

# method to get a dict of glycans form input iupac snfg
def get_glycans(iupacs:Dict[int, str], start = None, end = None) -> Tuple[Dict[int, Glycan], List[str], List[str]]:
    gylcans_dict = {}
    monos = []
    links = []

    for id in iupacs:
        if start is not None and id < start:
            continue
        if end is not None and id > end:
            continue
        inpuac_text = iupacs[id]
        gylcan = Glycan(inpuac_text)

        #if gylcan.get_num_nosaccharides() > 2:
        gylcans_dict[id] = gylcan
        monos += gylcan.get_filtered_monosaccharide_emssions()
        links += gylcan.get_filtered_linkage_emssions()

    mono_emissions = list(set(monos))
    link_emissions = list(set(links))

    return gylcans_dict, mono_emissions, link_emissions

# Method create foreset iputs from a dict collection of glycans
# Input:
#   glycans_dict: dict of glycans
# Return:
#   joint_adj_matrix - joint adjcent matrix
#   joint_monosaccharide_emission_observations - joint emissions for monosaccharides types
#   joint_linkage_emission_observations - joint emissions for linkage types
def create_forest_inputs(glycans_dict:Dict[int, Glycan]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    adj_matrices = []
    joint_monosaccharide_emission_observations = []
    joint_linkage_emission_observations = []

    for id in glycans_dict:
        glyan = glycans_dict[id]
        # update list of adj_matrices we are going to join along the diagonal
        adj_matrices.append(glyan.get_adj_matrix())
        # update joint_monosaccharide_emission_observations and joint_linkage_emission_observations
        joint_monosaccharide_emission_observations = joint_monosaccharide_emission_observations + glyan.get_filtered_monosaccharide_emssions()
        joint_linkage_emission_observations = joint_linkage_emission_observations + glyan.get_filtered_linkage_emssions()

    # join adj_martrix along diagonal
    joint_adj_matrix = block_diag(*adj_matrices)
    return joint_adj_matrix, joint_monosaccharide_emission_observations, joint_linkage_emission_observations

# Method learn a treehmm from a forset
# Input:
#   joint_adj_matrix - joint adjcent matrix
#   joint_emissions_observations - emissions observations for this foreset N * M, N is the number of emissions group eg,monosaccharide, linkage,
#       M is number of emssions groups eg.  monosaccharide_emissions for group 1 and linkage_emissions for group 2
#   number_state - number of states
#   possible_emissions -  possible_ emissions for this foreset N * M, N is the number of emissions group eg,monosaccharide, linkage,
#       M is number of emssions groups   eg.  monosaccharide_emissions for group 1 and linkage_emissions for group 2
def create_and_train_treehmm(number_state:int, iupacs, include_linkage = False, max_iterations=50, delta=1e-5):

    gylcans, possible_monosaccharide_emissions, possible_linkage_emissions = get_glycans(iupacs, 50, 75)

    joint_adj_matrix, joint_monosaccharide_emission_observations, joint_linkage_emission_observations = create_forest_inputs(gylcans)

    #print(possible_monosaccharide_emissions)
    # monosaccharide_emission_observations_counts = {}
    # for item in possible_monosaccharide_emissions:
    #     monosaccharide_emission_observations_counts[item] = 0
    # for item in joint_monosaccharide_emission_observations:
    #     monosaccharide_emission_observations_counts[item] += 1
    # print(monosaccharide_emission_observations_counts)

    # possible_linkage_emissions_counts = {}
    # for item in possible_linkage_emissions:
    #     possible_linkage_emissions_counts[item] = 0
    # for item in joint_linkage_emission_observations:
    #     possible_linkage_emissions_counts[item] += 1
    # print(possible_linkage_emissions_counts)

    # we only use monosaccharide
    if not include_linkage:
        joint_emissions_observations = [joint_monosaccharide_emission_observations]
        possible_emissions = [possible_monosaccharide_emissions]
    else:
    # if we want to use both of them
        joint_emissions_observations = [joint_monosaccharide_emission_observations, joint_linkage_emission_observations]
        possible_emissions = [possible_monosaccharide_emissions, possible_linkage_emissions]

    forest = csr_matrix(joint_adj_matrix)
    # create states
    states = [ str(i) for i in range(number_state)]

    #state_transition_probabilities = np.array([0.1,0.9,0.1,0.9]).reshape(2,2)
    hmm = initHMM.initHMM(states, possible_emissions)
    newparam = None
    for i in range(5):
        # The baumWelch part: To find the new parameters and result statistics
        newparam = baumWelch.hmm_train_and_test(hmm, forest, joint_emissions_observations, maxIterations = max_iterations, delta = delta)
        #newparam = baumWelch.baumWelchRecursion(hmm, emission_observation)

        hmm = initHMM.initHMM(states, possible_emissions, state_transition_probabilities=newparam['hmm']['state_transition_probabilities'],
                                                emission_probabilities=newparam['hmm']['emission_probabilities'])

        #print(newparam['hmm']['state_transition_probabilities'])
        #print(newparam['hmm']['emission_probabilities'])
        #fwd_tree_sequence = fwd_seq_gen.forward_sequence_generator(forest)
        #bind_fwd_probs = forward.forward(hmm_trained, forest, joint_emissions_observations, fwd_tree_sequence)
        #print(joint_emissions_observations)

        ll = 0
        for glycan_idx in gylcans:
            glycan = gylcans[glycan_idx]
            #print(glycan.get_filtered_monosaccharide_emssions())
            glycan_tree = csr_matrix(glycan.get_adj_matrix())
            fwd_tree_sequence = fwd_seq_gen.forward_sequence_generator(glycan_tree)
            emssions = glycan.get_filtered_monosaccharide_emssions()
            if include_linkage:
                emssions += glycan.get_filtered_linkage_emssions()
            if fwd_tree_sequence[-1] != glycan.get_num_nosaccharides() - 1:
                print('oops')
            bind_fwd_probs = forward.forward(hmm, glycan_tree, [emssions], fwd_tree_sequence)
            case_ll = logsumexp(bind_fwd_probs.iloc[:, -1])
            ll += case_ll

        print(max_iterations * (i + 1), ll)
    return newparam

if __name__ == "__main__":

    #log_format = '%(asctime)s %(message)s'
    #$logging.basicConfig(stream=sys.stdout, level=logging.INFO,
    #                    format=log_format, datefmt='%m/%d %I:%M:%S %p')

    # get arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('--num_states', help='num of hidden states', type = int, default = 2)
    parser.add_argument('--include_linkage', help='0 for no, 1 for yes', type = int, default = 0)
    args = parser.parse_args()

    num_states = args.num_states
    include_linkage = True if args.include_linkage > 0 else False

    # read iupac and get gylcans
    iupac_name_file = './Data/IUPAC.csv'
    iupacs = get_iupcas(iupac_name_file)
    #create_and_train_treehmm(num_states, iupacs, max_iterations = 5)
    #create_and_train_treehmm(num_states, iupacs, max_iterations = 10)
    #create_and_train_treehmm(num_states, iupacs, max_iterations = 20)
    #create_and_train_treehmm(num_states, iupacs, max_iterations = 30)
    #create_and_train_treehmm(num_states, iupacs, max_iterations = 40)

    gylcans, possible_monosaccharide_emissions, possible_linkage_emissions = get_glycans(iupacs)
    joint_adj_matrix, joint_monosaccharide_emission_observations, joint_linkage_emission_observations = create_forest_inputs(gylcans)

    #print(possible_monosaccharide_emissions)
    monosaccharide_emission_observations_counts = {}
    for item in possible_monosaccharide_emissions:
        monosaccharide_emission_observations_counts[item] = 0
    for item in joint_monosaccharide_emission_observations:
        monosaccharide_emission_observations_counts[item] += 1
    print(monosaccharide_emission_observations_counts)

    possible_linkage_emissions_counts = {}
    for item in possible_linkage_emissions:
        possible_linkage_emissions_counts[item] = 0
    for item in joint_linkage_emission_observations:
        possible_linkage_emissions_counts[item] += 1
    print(possible_linkage_emissions_counts)


#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Thu May 25 10:10:28 2017

@author: Katherine
"""

#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Tue May 23 16:04:21 2017

@author: Katherine
"""

# -*- coding: utf-8 -*-
"""
Photosynthesis and Stomatal Conductance Model
Created 9/27/2016
Katherine Wentz

This is a program that runs photosynthesis and
stomatal conductance models given changes in leaf-
level traits.

The end product is graphs of NUE vs. WUE.

Update: I am going to run the model for plants with
traits that are distinctive of the meadow moisture
gradient in the alpine tundra.

Fix: correct for atmospheric pressure differences in co2, o2, and vapor pressure

Fix: vcmax temp dependence (pg 63 in plant physiological ecology book)

Fix: NEW VARIBALE TRAIT-->make the fraction of leaf N in rubisco go down with increasing SLA,
chlorophyll content, and decreasing light (wet meadow)--more N is allocated
to thylakoids. The only way for chl/m2 to increase even when g N/m2 goes down
or is constant is for the leaf to allocate more of leaf N to chl...also, note
that there is more organic N designated to photo in leaf when SLA goes up
because less N is used in structure. see "Photosynthesis or persistence: N allocation
in leaves of evergreen and deciduous... by Takashima et al. 2004. Also see Photosynthetic
nitrogen-use efficiency of species...by Poorter and Evans 1998

Note to self: NUE and WUE relationship flipflops with change in air temperature;
NUE makes sense because C:N decreases from dry to wet meadows; WUE increasing
in snowbed does not necessarilly make sense--look in the literature for this

herbs have a higher NUE

"""

#---------------Import Modules---------------#

import itertools as it
import numpy as np
from matplotlib import pyplot as plt

#Import combinations of variable parameters
from uncertain_params import monte_carlo

#Import photosynthesis model
from Photosynthesis_Model import photo_bound_meso_eqstom as photo

#Import functions to switch between Pa and umol/mol at sea level
from photo_functions import pa_con_atmfrac

#import timeseries of vwc and temp
from time_dep_params import surtemp_dm, surtemp_wm, vwc_dm, vwc_wm, na_dm_min_inter,na_wm_min_inter,na_dm_max_inter,na_wm_max_inter

#---------------Determine if I Want to Keep Any of the Variable Parameters Constant---------------#

const_params=[]
for xxx in it.combinations(['ht'],0): #keep ht and t constant for constant vpd
    const_params+=[xxx]

#do this when I do not put any of the variable parameters as constant. instead I
#vary each parameter one at a time while keeping the other parameters constant.
if const_params==[()]:
    const_params=[[-999999]]

#---------------Begin Looping Through Photosynthesis Model---------------#

#each loop is for a constant value, or combinatin of constant values, of variable parameter as determined above
for ii in range(len(const_params)):

    #---------------Run through time series---------------#

    days=np.linspace(1,365,365)

    #dry meadow
    tot_nue_dm_avg=[]
    tot_wue_dm_avg=[]
    tot_nue_dm_min=[]
    tot_wue_dm_min=[]
    tot_nue_dm_max=[]
    tot_wue_dm_max=[]
    tot_A_dm_avg=[]
    tot_gs_dm_avg=[]

    #moist meadow
    tot_nue_mm_avg=[]
    tot_wue_mm_avg=[]
    tot_nue_mm_min=[]
    tot_wue_mm_min=[]
    tot_nue_mm_max=[]
    tot_wue_mm_max=[]
    tot_A_mm_avg=[]
    tot_gs_mm_avg=[]

    #wet meadow
    tot_nue_wm_avg=[]
    tot_wue_wm_avg=[]
    tot_nue_wm_min=[]
    tot_wue_wm_min=[]
    tot_nue_wm_max=[]
    tot_wue_wm_max=[]
    tot_A_wm_avg=[]
    tot_gs_wm_avg=[]

    #---------------Photosynthesis + Stomatal Conductance Model---------------#

    ##---Constant Parameter Arrays for Model---##

    #----Params Used in Model Currently----#

    tk_25=298.16; #absolute temperature at 25 C
    ekc=80500.0 #Activation energy for K of CO2 (J mol-1)
    eko=14500.0 #Activation energy for K of O2 (J mol-1)
    etau=-29000.0  #Activation energy for tau (???) (J mol-1)
    ev=55000.0 #Activation energy for carboxylation (J mol-1)
    ej=55000.0 #Activation energy for electron transport (J mol-1)
    toptv=303.0 #Optimum temperature for maximum carboxylation (K)
    toptj=303.0 #Optimum temperature for maximum electron transport (K)
    ra=np.zeros(shape=1)+20.7 #specific rubisco activity (umol CO2/g Rub s)
    flnr=np.zeros(shape=1)+0.1 #fraction of leaf nitrogen in rubisco (g N Rub/g N leaf)
    frnr=np.zeros(shape=1)+6.25 #weight fraction of nitrogen in rubisco molecule (g Rub/g N Rub)
    rh=np.zeros(shape=1)+0.5 #relative humidity (kPa/kPa)
    ca=np.zeros(shape=1)+405 #ambient carbon dioxide (umol CO2/mol air)
    ko25=np.zeros(shape=1)+30000 #Michaelis-Menten kinetic coefficient for oxygen at 25 C(Pa)
    kc25=np.zeros(shape=1)+30 #Michaelis-Menten kinetic coefficient for carbon dioxide at 25 C (Pa)
    o=np.zeros(shape=1)+210000 #concentration of ambient oxygen (umol/mol)
    g0=np.zeros(shape=1)+0.002 #Ball-Berry stomatal conductance intercept parameter (mol H2O/m2s)
    a=np.zeros(shape=1)+1.6 #Conversion Coefficient between stomatal conductance to water and carbon dioxide (unitless)
    ij=np.zeros(shape=1)+1.0 #leaf angle index--downregulates jmax
    m=np.zeros(shape=1)+9.0 #ball-berry parameter (unitless)
    b=1.37 #Conversion Coefficient between boundary layer conductance to water and carbon dioxide
    u=5.0 #windspeed (m/s)
    qeff=0.32 #leaf quantum yield, electrons
    PAR=2000 #photosynthetic active radiation (umol/m2s)
    jm=2.68 #slope coefficient
    vwc_min=0.08 #minimum soil water content for photosynthesis to occur (permanent wilting point) (cm3/cm3)
    vwc_max=0.68 #maximum soil water content where increases in soil water do not affect photosynthesis (field capacity?) (cm3/cm3)
    q=0.2 #parameter for soil water affect on photosynthesis (unitless)

    #------constant variable params for sensitivty analysis-----#

    chl_c=np.zeros(shape=1)+(np.mean([396,465,476])) #Chlorophyll Content of the Leaf (umol chl/m2)
    ht_c=np.zeros(shape=1)+(np.mean([9.2,19.5,20.0])) #Temperature of the Leaf (K)
    dia_c=np.zeros(shape=1)+(np.mean([1.4,2.3,2.6])/100.) #Mean diameter or size of leaf (m)
    na_c=np.zeros(shape=1)+(np.mean([2.5,5.6,6.3])) #leaf nitrogen (g N/ m2)
    t_c=np.zeros(shape=1)+15.0 #temp (C)

    #-----which timeseries should I use--based on factorial meadow type---#

    na_min=[na_dm_min_inter, na_dm_min_inter, na_wm_min_inter, na_wm_min_inter]
    na_max=[na_dm_max_inter, na_dm_max_inter, na_wm_max_inter, na_wm_max_inter]

    vwc_type=[vwc_dm,vwc_dm,vwc_wm,vwc_wm]
    temp_type=[surtemp_dm,surtemp_dm,surtemp_wm,surtemp_wm]

    A_tot_all=[]

    chl_mean=[[395.7132],[475.8913],[395.7132],[475.8913]]
    chl_sd=[[24.410199999999975],[29.185099999999977],[24.410199999999975],[29.185099999999977]]
    dia_mean=[[1.6/100.],[3.0/100.],[1.6/100.],[3.0/100.]]
    dia_sd=[[0.9/100.0],[1.2/100.0],[0.9/100.0],[1.2/100.0]]
    ht_mean=[[9.183549],[19.98519],[9.183549],[19.98519]]
    ht_sd=[[1.5],[3.1],[1.5],[3.1]]

    depth=[0.2,0.2,0.2,0.4,0.4,0.4]

#---------------Import Variable Parameter Arrays from Leaf Parameter File---------------#
    for iii in range(len(chl_mean)):
        A_tot=0
        for time in range(129):
            params=monte_carlo(chl_mean[iii], chl_sd[iii], dia_mean[iii], dia_sd[iii], [na_min[iii][time]], [na_max[iii][time]], ht_mean[iii], ht_sd[iii])
            A_day=[]
            for xx in range(len(params)):
                for yy in range(len(params[xx])):
                    for key,val in params[xx][yy].items():
                        exec(key + '=val')

                    #set variable parameters constant if I specify this above
                    if 'na' in const_params[ii]:
                        na=na_c
                    if 'dia' in const_params[ii]:
                        dia=dia_c
                    if 'chl' in const_params[ii]:
                        chl=chl_c
                    if 'ht' in const_params[ii]:
                        ht=ht_c

                    #------calculate vapor pressure-----#
                    pa_v=611*np.exp((17.27*temp_type[iii][time])/(temp_type[iii][time]+237.3)) #saturation vapor pressure of air (Pa)
                    ea_str=pa_con_atmfrac(pa_v,3528) #saturation vapor pressure of air (Pa-->umol h20/mol air)
                    ea=rh*ea_str #vapor pressure (umol h2O/mol air)

                    #correct for leaf temperatures using leaf height

                    t_diff=18-0.4*ht

                    tl=temp_type[iii][time]+t_diff

                    z=depth[iii]

                    #---------------Photosynthesis Function---------------#

                    #alter this line of code for when implementing different photosynthesis functions
                    wue, nue, A, E, cs, ci, gsw, gs, gbw, gb, gm, cc,dd =photo(tk_25,ekc,eko,etau,ev,ej,toptv,toptj,na, qeff, PAR,tl,ea,chl,ij,kc25,ko25,o,ca,rh,m,a,frnr,flnr,ra,jm,g0,b,dia,u,q,vwc_min,vwc_max,vwc_type[iii][time],z)

                    #test to make sure wue and nue are positive at not 'nan'
                    if wue[0]==-999 and nue[0]==-999:

                        continue

                    if np.isnan(A[0]):
                        A[0]=0.0

                    A_day+=[(A[0]*3600*6)/1000000.*44.]

            A_tot+=np.mean(A_day)

        A_tot_all+=[A_tot]
        print A_tot_all

print A_tot_all


#Script to generate DSWE Version 2 Prototype 1 (V2_P1) Open Source Distribution Version 2 (distro_v3)
#
#
#
# *****************DEPENDENCIES*******************************************
# Python 2.7 64 bit, will not run on 32 bit Python
# GDAL 2.1+
# numpy 1.12
# lxml 3.7+
#
# *****************REVISION HISTORY***************************************
# Initiated on 08/13/16 by John W. Jones.
# Modified on 8/15/16 by Jacob Shermeyer to improve efficiency and activate
#  recoding scheme to produce the interpreted layer.
# Modified on 11/08/16 by John W. Jones to improve documentation as well as
#  revise the recode scheme to produce revised interpreted layer.
# Modified on 11/14/2016 to add slope masking.
# Modified on 1/18/2017 to enhance processing speed and remove a few unnecessary items
# Modified on 2/22/2017 to change to open source implementation (GDAL and numpy)
# Modified through March and April 2017 to add hillshade, modify slope masking, and update to the
#   new SR naming convention
# Modified April 27-28 2017 to remove NLCD masking, remove redundant opening and closing of files,
#   generally improve performance, and change output bands to only those that will be distributed.
#   Note also, that the calculation of percent slope and hillshade layers are assumed to occur
#   outside of this code (in a pre-processing program).
#
# *****************INPUT FILE STRUCTURE***********************************
#  Required inputs for DSWE are LSRP data, the CFMask file generated by EROS, percent slope and hillshade
#   The input file directory should contain subdirectories in which all the geotiff formatted bands of a single
#  input scene have been extracted from their zip files and are stored in seperate folders.
#  Naming conventions of these files must not be edited in ANY way!  Use the default filenames as they were distributed by EROS.
#   04/28/2017
#  With this version, a percent slope file for the US is stored at the top level of the input directory. And hillshade data
#   calculated using the scene-specific solar geometry included in the mtl or xml file is included in
#   the same folder as each set of unzipped LSRP file associated with each scene.
# This data derived from the DEM should be at 30m spatial resolution, snapped to your Landsat image and MUST be projected to match
#   the Landsat UTM zone.  It should be slighlty larger than your path/row.
#
#
# *****************VARIABLE DECLARATIONS***********************************
#
# Input (data) variables
# BLUE = Landsat TM/ETM Band 1 or OLI Band 2
# GREEN = Landsat TM/ETM Band 2 or OLI Band 3
# NIR = Landsat TM/ETM Band 4 or OLI Band 5
# SWIR1 = Landsat TM/ETM Band 5 or OLI Band 6
# SWIR2 = Landsat TM/ETM Band 7 or OLI Band 7
# CF= Landsat Cfmask- Cloud, Cloud Shadow, and Snow Mask distributed with LSRP data
#
# *****************CALCULATED VARIABLES ON WHICH DECISIONS ARE MADE*********
#
# NDVI = Normalized Difference Vegetation Index
# MNDWI = Modified Normalized Difference Wetness Index
# MBSRV = Multiband Spectral Relationship Visible
# MBSRN = Multiband Spectral Relationship Near-infrared
# AWEsh = Automated raster Extent Shadow
# PSW1 = Partial Surface Water 1
# PSW2 = Partial Surface Water 2
#
# *****************OUTPUT INFORMATION*****************************
# Diagnostic layer- Layer indicating the combination of outcomes of the 5 water tests
# Interpreted layer- Interpretation of the diagnostic layer:
#     0= Non-Water, 1= High Confidence water, 2= Medium Confidence Water, 3= Potential Wetland, 4= Low Confidence Water or Wetland
# Masked Interpreted layer- Cloud, Snow, Slope, and Hillshade
# Mask Layer- A combination of all masks, combinations of these values indicate a pixel was masked multiple times:
# 0- No masks applied
# 2- CFMask Cloud Shadow Mask
# 3- CFMask Snow Mask
# 4- CFMask Cloud Mask
# 10- Slope Mask
# 20- Hillshade Mask
# 30- Slope and Hillshade Combination Mask

#Imports
import numpy as np
import gdal
import os
import datetime
import glob
import sys
from lxml import etree

# ******************REQUIRED INPUT PARAMETERS********************************
WIGT = 0.124
AWGT = 0
AWE_param1 = 2.5
AWE_param2 = -1.5
AWE_param3 = -0.25
PSW1_MNDWI = -0.44
PSW1_NIR = 1500
PSW1_SWIR1 = 900
PSW1_NDVI = 0.7
PSW2_MNDWI = -0.5
PSW2_BLUE = 1000
PSW2_NIR = 2500
PSW2_SWIR1 = 3000
PSW2_SWIR2 = 1000
# Cutoff value for hillshade, values for hillshade range from 1-255, the lower the more shaded.
HS=110

# ***************************REQUIRED PATH INFORMATION**************************
# Specify the path for input and output data,  remember to use / insted of \ if there are any numbers in the path
input_dir = "H:\DSWE\DSWE_V2\Cal_Val_Inputs_subset"
output_dir = "H:\DSWE\DSWE_V2\Cal_Val_Outputs_subset"

# ***************************REQUIRED SLOPE INFORMATION**************************
# Slope file you have created must be stored in the input_dir along with the LSRP data
# The file must have the word "perslp" anywhere in the filename and be a .tif image.

#Set the overwrite & output environment
driver = gdal.GetDriverByName("GTiff")
os.chdir(input_dir)
paths = glob.glob('*/')
print paths
for folder in paths:
    output_dir2=output_dir + "/" + folder
    if not os.path.exists(output_dir2):
        os.makedirs(output_dir2)
    folder=os.path.join(input_dir,folder)
    os.chdir(folder)
    # Identify our percent slope file
    PerSlpL=[]
    PerSlpL=glob.glob('*perslp*.tif')
    if len(PerSlpL) > 1 or len(PerSlpL) < 1:
        print "multiple or no PerSlpL files, please place only one PerSlpL file in the LSRP input_dir"
        exit()
    else:
        print "Using PerSlpL file:"
        print PerSlpL[0]
        Perslp= os.path.abspath(PerSlpL[0])
        print Perslp
    count = 1
    #Loop through the input directory and identitify our LSRP scenes and save them to the list
    print datetime.datetime.now()
    subdirs= glob.glob('*/')
    print('Found directories: %s' % subdirs)
    subdirnumber=0
    subdirnumber+= len(subdirs)
    print subdirnumber
    cwd = os.getcwd()
    print cwd
    if count <= subdirnumber:
        T1=datetime.datetime.now()
        for directory in subdirs:
            print ""
            print ""
            print ""
            working_dir=folder + "//" + directory
            print working_dir
            os.chdir(working_dir)
            inputRaster = []
            inputRaster= glob.glob('L*')
            print inputRaster

            #Identify bands  based upon satellite
            for raster in inputRaster:
                filetrim=raster[:5]
                numtrim=filetrim[4:]
                if numtrim == "_":
                    filetrim=raster[:4]
                    numtrim=filetrim[3:]
                    LSnum=int(numtrim)
                    length=40
                    length2=21
                else:
                    filetrim=raster[:3]
                    numtrim=filetrim[2:]
                    LSnum=int(numtrim)
                    length=21
                    length2=13

            #Landsat 4,5, and 7
                if LSnum<=7 :
                    Blue= working_dir +  '\%s_sr_band1.tif'%(raster)[0:length]
                    Green= working_dir +  '\%s_sr_band2.tif'%(raster)[0:length]
                    Red= working_dir +  '\%s_sr_band3.tif'%(raster)[0:length]
                    NIR= working_dir +  '\%s_sr_band4.tif'%(raster)[0:length]
                    SWIR1= working_dir +  '\%s_sr_band5.tif'%(raster)[0:length]
                    SWIR2= working_dir +  '\%s_sr_band7.tif'%(raster)[0:length]
                    CF = working_dir + '\%s_cfmask.tif'%(raster)[0:length]
                    Hillshade= working_dir + '\%s_hillshade.tif'%(raster)[0:length]
                    Metadata = working_dir + '\%s.xml'%(raster)[0:length]
                    yeartrim=raster[:length2]
                    year=yeartrim[-4:]
                    year=int(year)
            #Landsat 8
                elif LSnum==8:
                    Blue= working_dir +  '\%s_sr_band2.tif'%(raster)[0:length]
                    Green= working_dir +  '\%s_sr_band3.tif'%(raster)[0:length]
                    Red= working_dir +  '\%s_sr_band4.tif'%(raster)[0:length]
                    NIR= working_dir +  '\%s_sr_band5.tif'%(raster)[0:length]
                    SWIR1= working_dir +  '\%s_sr_band6.tif'%(raster)[0:length]
                    SWIR2= working_dir +  '\%s_sr_band7.tif'%(raster)[0:length]
                    CF = working_dir + '\%s_cfmask.tif'%(raster)[0:length]
                    Hillshade= working_dir + '\%s_hillshade.tif'%(raster)[0:length]
                    Metadata = working_dir + '\%s.xml'%(raster)[0:length]
                    yeartrim=raster[:length2]
                    year=yeartrim[-4:]
                    year=int(year)
                else:
                    print "Wrong satellite or wonky input, check data."
                    quit()
                out_name = 'DSWE_V2_P1'

            # Perform Water Index (MNDWI only based) Test- Working
            GreenB= gdal.Open(Green)
            SWIR1B= gdal.Open(SWIR1)
            geo = GreenB.GetGeoTransform()
            proj = GreenB.GetProjection()
            GreenB= GreenB.GetRasterBand(1).ReadAsArray()
            SWIR1B= SWIR1B.GetRasterBand(1).ReadAsArray()
            shape = GreenB.shape
            MNDWI = ((GreenB - SWIR1B) /np.float32(GreenB + SWIR1B))
            con_MNDWI=MNDWI.copy()
            con_MNDWI[con_MNDWI <= WIGT] = 0
            con_MNDWI[con_MNDWI > WIGT] = 1
            con_MNDWI=np.int8(con_MNDWI)
            #Optional to output MNDWI test
            #con_MNDWI1_out = driver.Create( "con_MNDWI1.tif", shape[1], shape[0], 1, gdal.GDT_UInt16)
            #con_MNDWI1_out.SetGeoTransform( geo )
            #con_MNDWI1_out.SetProjection( proj )
            #con_MNDWI1_out.GetRasterBand(1).WriteArray(con_MNDWI)
            #con_MNDWI1_out=None

            # Perform MBSR Test- Working
            RedB= gdal.Open(Red)
            RedB= RedB.GetRasterBand(1).ReadAsArray()
            MBSV = (GreenB + RedB)
            NIRB= gdal.Open(NIR)
            NIRB= NIRB.GetRasterBand(1).ReadAsArray()
            MBSRN = (NIRB + SWIR1B)
            con_MBSR=np.zeros(shape)
            con_MBSR[MBSV > MBSRN] = 10
            con_MBSR=np.int8(con_MBSR)
            del MBSV
            #Optional to output MBSR test
            #con_MBSR_out = driver.Create( "con_MBSR.tif", shape[1], shape[0], 1, gdal.GDT_UInt16)
            #con_MBSR_out.SetGeoTransform( geo )
            #con_MBSR_out.SetProjection( proj )
            #con_MBSR_out.GetRasterBand(1).WriteArray(con_MBSR)
            #con_MBSR_out=None

            #Calculate AWEsh values- working
            con_AWEsh = (np.float32(GreenB) * AWE_param1)
            con_AWEsh += (np.float32(MBSRN) * AWE_param2)
            MBSRN_out = driver.Create( "MBSRN.tif", shape[1], shape[0], 1, gdal.GDT_UInt32)
            MBSRN_out.SetGeoTransform( geo )
            MBSRN_out.SetProjection( proj )
            MBSRN_out.GetRasterBand(1).WriteArray(MBSRN)
            MBSRN_out=None
            SWIR2B= gdal.Open(SWIR2)
            SWIR2B= SWIR2B.GetRasterBand(1).ReadAsArray()
            con_AWEsh += (np.float32(SWIR2B) * AWE_param3)
            BlueB= gdal.Open(Blue)
            BlueB= BlueB.GetRasterBand(1).ReadAsArray()
            con_AWEsh += np.float32(BlueB)
            con_AWEsh[con_AWEsh > 0] = 100
            con_AWEsh[con_AWEsh <= 0] = 0
            con_AWEsh=np.int8(con_AWEsh)
            #Optional to output AWEsh test
            #con_AWEsh_out = driver.Create( "con_AWEsh.tif", shape[1], shape[0], 1, gdal.GDT_UInt16)
            #con_AWEsh_out.SetGeoTransform( geo )
            #con_AWEsh_out.SetProjection( proj )
            #con_AWEsh_out.GetRasterBand(1).WriteArray(con_AWEsh)
            #con_AWEsh_out=None

            # Perform Partial Surface Water test 1 (PSW1)-working
            con_PSW1_MNDWI=((GreenB - SWIR1B) /np.float32(GreenB + SWIR1B))

            con_PSW1_MNDWI[con_PSW1_MNDWI > PSW1_MNDWI] = 1
            con_PSW1_MNDWI[con_PSW1_MNDWI <= PSW1_MNDWI] = 0
            con_PSW1_MNDWI=np.int8(con_PSW1_MNDWI)
            NIRB[NIRB < PSW1_NIR] = 1
            NIRB[NIRB >= PSW1_NIR] = 0
            PSW1= np.int8(con_PSW1_MNDWI) * np.int8(NIRB)
            PSW1_SWIR1B=SWIR1B.copy()
            PSW1_SWIR1B[PSW1_SWIR1B < PSW1_SWIR1] = 1
            PSW1_SWIR1B[PSW1_SWIR1B >= PSW1_SWIR1] = 0
            PSW1= np.int8(PSW1) * np.int8(PSW1_SWIR1B)
            del PSW1_SWIR1B

            # Calculate NDVI
            NIRB= gdal.Open(NIR)
            NIRB= NIRB.GetRasterBand(1).ReadAsArray()
            NDVI = ((NIRB - RedB) /np.float32(NIRB + RedB))
            con_NDVI=np.zeros(shape)
            con_NDVI[NDVI < PSW1_NDVI] = 1
            PSW1= np.int16(PSW1) * np.int16(con_NDVI) * 1000
            #Optional to output PSW1 test
            #PSW1_out = driver.Create( "PSW1.tif", shape[1], shape[0], 1, gdal.GDT_UInt16)
            #PSW1_out.SetGeoTransform( geo )
            #PSW1_out.SetProjection( proj )
            #PSW1_out.GetRasterBand(1).WriteArray(PSW1)
            #PSW1_out=None

            #Perform PSW2 test- working
            con_PSW2_MNDWI=((GreenB - SWIR1B) /np.float32(GreenB + SWIR1B))
            con_PSW2_MNDWI[con_PSW2_MNDWI > PSW2_MNDWI] = 1
            con_PSW2_MNDWI[con_PSW2_MNDWI <= PSW2_MNDWI] = 0
            con_PSW2_MNDWI=np.int8(con_PSW2_MNDWI)
            BlueB[BlueB < PSW2_BLUE] = 1
            BlueB[BlueB >= PSW2_BLUE] = 0
            BlueB=np.int8(BlueB)
            NIRB[NIRB < PSW2_NIR] = 1
            NIRB[NIRB >= PSW2_NIR] = 0
            NIRB=np.int8(NIRB)
            SWIR1B[SWIR1B < PSW2_SWIR1] = 1
            SWIR1B[SWIR1B >= PSW2_SWIR1] = 0
            SWIR1B=np.int8(SWIR1B)
            SWIR2B[SWIR2B < PSW2_SWIR2] = 1
            SWIR2B[SWIR2B >= PSW2_SWIR2] = 0
            SWIR2B=np.int8(SWIR2B)
            PSW2= con_PSW2_MNDWI * BlueB * NIRB * SWIR1B *  SWIR2B * 10000
            #Optional to output PSW2 test
            #PSW2_out = driver.Create( "PSW2.tif", shape[1], shape[0], 1, gdal.GDT_UInt16)
            #PSW2_out.SetGeoTransform( geo )
            #PSW2_out.SetProjection( proj )
            #PSW2_out.GetRasterBand(1).WriteArray(PSW2)
            #PSW2_out=None
            del GreenB, SWIR1B, MNDWI, RedB, NIRB, MBSRN, SWIR2B, BlueB, con_PSW1_MNDWI, con_NDVI, con_PSW2_MNDWI
            diagmap= con_MNDWI + con_MBSR + con_AWEsh + PSW1 + PSW2
            del con_MNDWI, con_MBSR, con_AWEsh, PSW1, PSW2

            os.chdir(input_dir)

            # Terrain Correction, here we are simply recoding slopes >=x degrees to PS (specified above), which will be reclassified as non-water.

            perslp = gdal.Open(Perslp)
            #Get extent and projection
            perslpgeo = perslp.GetGeoTransform()
            perslpproj = perslp.GetProjection()
            # Calculate percent slope
            #perslp= perslp.GetRasterBand(1).ReadAsArray()
            os.chdir(working_dir)
            CF = gdal.Open(CF)
            #Get extent of CF mask (and our LS image by proxy)
            geoTransform = CF.GetGeoTransform()
            minx = geoTransform[0]
            maxy = geoTransform[3]
            maxx = minx + geoTransform[1] * CF.RasterXSize
            miny = maxy + geoTransform[5] * CF.RasterYSize
            #Clip our slope and hillshade mask to the image (have to do this for each image due to the variable extents of each image)
            GeoClip= [minx, maxy, maxx, miny]
            print GeoClip
            hillshade= gdal.Open(Hillshade)
            perslp_clip=gdal.Translate('perslp_clip.tif', perslp, projWin = GeoClip)
            hillshade_clip=gdal.Translate('hillshade_clip.tif', hillshade, projWin = GeoClip)
            hillshade= hillshade_clip.GetRasterBand(1).ReadAsArray()
            hillshade[hillshade <= HS] = 0
            hillshade[hillshade > HS] = 1

            #Read in the slope and hillshade masks and convert to an array
            workingoutput_dir=output_dir2 + "\\" + directory
            if not os.path.exists(workingoutput_dir):
                os.makedirs(workingoutput_dir)

            #calculate our unmasked diagnostic map
            CF= CF.GetRasterBand(1).ReadAsArray()
            diagoutput=output_dir2 +  "\\" + directory + "\\" + out_name + "_diag" + "_%s.tif"%"_".join(raster.split('_'))[0:length]
            diagmap[np.where((CF == 255))] = -9999
            diagmap_out = driver.Create( diagoutput, shape[1], shape[0], 1, gdal.GDT_Int16)
            diagmap_out.SetGeoTransform( geo )
            diagmap_out.SetProjection( proj )
            diagmap_out.GetRasterBand(1).WriteArray(diagmap)
            diagmap_out.GetRasterBand(1).SetNoDataValue(-9999)

            # Output unmasked interpreted map
            interpoutput=output_dir2 +  "\\" + directory + "\\" + out_name + "_interp" + "_%s.tif"%"_".join(raster.split('_'))[0:length]
            interpmap=diagmap.copy()
            interpmap[np.where((diagmap == 0) | (diagmap == 1) | (diagmap == 10) | (diagmap == 100) | (diagmap == 1000))] = 0
            interpmap[np.where((diagmap == 1111) | (diagmap == 10111) | (diagmap == 11011) | (diagmap == 11101) | (diagmap == 11110) | (diagmap == 11111)) ] = 1
            interpmap[np.where((diagmap == 111) | (diagmap == 1011) | (diagmap == 1101) | (diagmap == 1110) | (diagmap == 10011) | (diagmap == 10101) | (diagmap == 10110) | (diagmap == 11001) | (diagmap == 11010) | (diagmap == 11100)) ] = 2
            interpmap[np.where((diagmap == 11000)) ] = 3
            interpmap[np.where((diagmap == 11) | (diagmap == 101) | (diagmap == 110) | (diagmap == 1001) | (diagmap == 1010) | (diagmap == 1100) | (diagmap == 10000) | (diagmap == 10001) | (diagmap == 10010) | (diagmap == 10100)) ] = 4
            interpmap[np.where((diagmap == -9999)) ] = 255
            interpmap_out = driver.Create( interpoutput, shape[1], shape[0], 1, gdal.GDT_Byte)
            interpmap_out.SetGeoTransform( geo )
            interpmap_out.SetProjection( proj )
            interpmap_out.GetRasterBand(1).WriteArray(interpmap)
            interpmap_out.GetRasterBand(1).SetNoDataValue(255)

            # Mask interpreted map using the slope file
            perslp= perslp_clip.GetRasterBand(1).ReadAsArray()
            interpmap_copy=interpmap.copy()
            interpmap[np.where((perslp >= 30) & (interpmap == 2))] = 0
            interpmap[np.where((perslp >= 20) & (interpmap == 3))] = 0
            interpmap[np.where((perslp >= 10) & (interpmap == 4))] = 0
            interpmap[np.where((perslp >= 30) & (interpmap == 1))] = 0
            interpmap_masked= interpmap * hillshade

            #Cloud and snow masking and outputing masked interpreted layer
            interpmaskoutput=output_dir2 +  "\\" + directory + "\\" + out_name + "_interp_masked" + "_%s.tif"%"_".join(raster.split('_'))[0:length]
            interpmap_masked=interpmap_masked.copy()
            interpmap_masked[np.where((CF == 255))] = 255
            interpmap_masked[np.where((CF >= 2) & (CF < 255))] = 9
            interpmap_masked_out = driver.Create( interpmaskoutput, shape[1], shape[0], 1, gdal.GDT_Byte)
            interpmap_masked_out.SetGeoTransform( geo )
            interpmap_masked_out.SetProjection( proj )
            interpmap_masked_out.GetRasterBand(1).WriteArray(interpmap_masked)
            interpmap_masked_out.GetRasterBand(1).SetNoDataValue(255)
            del interpmap_masked_out
            del diagmap
            del interpmap

            #Create a full mask layer
            CF[np.where((CF == 1)) ] = 0
            perslp[np.where((perslp < 10) & (interpmap_copy == 4)) ] = 0
            perslp[np.where((perslp < 20) & (interpmap_copy == 3)) ] = 0
            perslp[np.where((perslp < 30) & (interpmap_copy == 2)) ] = 0
            perslp[np.where((perslp < 30) & (interpmap_copy == 1)) ] = 0
            perslp[np.where((perslp >= 10) & (interpmap_copy == 4)) ] = 10
            perslp[np.where((perslp >= 20) & (interpmap_copy == 3)) ] = 10
            perslp[np.where((perslp >= 30) & (interpmap_copy == 2)) ] = 10
            perslp[np.where((perslp >= 30) & (interpmap_copy == 1)) ] = 10
            perslp[np.where((interpmap_copy == 0)) ] = 0
            hillshade[np.where((hillshade == 0)) ] = 20
            hillshade[np.where((hillshade == 1)) ] = 0
            MaskLayer=CF + perslp + hillshade
            MaskLayerOut=output_dir2 +  "\\" + directory + "\\" + out_name + "_MaskLayer" + "_%s.tif"%"_".join(raster.split('_'))[0:length]
            MaskLayer[np.where((CF == 255))] = 255
            MaskLayer_Output = driver.Create(MaskLayerOut, shape[1], shape[0], 1, gdal.GDT_Byte)
            MaskLayer_Output.SetGeoTransform( geo )
            MaskLayer_Output.SetProjection( proj )
            MaskLayer_Output.GetRasterBand(1).WriteArray(MaskLayer)
            MaskLayer_Output.GetRasterBand(1).SetNoDataValue(255)
            del hillshade_clip, hillshade, MaskLayer, CF, MaskLayerOut, MaskLayer_Output, diagoutput
            del Metadata, diagmap_out, interpmap_masked, interpmap_out, interpmaskoutput, interpoutput, perslp, perslp_clip

            itemlist=["hillshade*.tif", "perslp*.tif"]
            for item in itemlist:
                i=glob.glob(item)
                for r in i:
                    os.remove(r)

           #Timer
            print datetime.datetime.now()
            T2=datetime.datetime.now()
            T3=T2-T1
            print T3

            if count == subdirnumber:
                break
            else:
                count=count+1
                print count