# coding=utf-8  python3.6
# ================================================================
#   Copyright (C) 2019 * Ltd. All rights reserved.
#   license     : MIT License
#   Author      : haibingshuai
#   Created date: 2019/10/29 15:19
#   Description :
# ================================================================

import tensorflow as tf
import numpy as np
import cv2
import random
import colorsys
from core.model_config import config

def bbox_iou(boxes1, boxes2):
    boxes1 = np.array(boxes1)
    boxes2 = np.array(boxes2)

    boxes1_area = boxes1[..., 2] * boxes1[..., 3]
    boxes2_area = boxes2[..., 2] * boxes2[..., 3]

    boxes1 = np.concatenate([boxes1[..., :2] - boxes1[..., 2:] * 0.5,
                             boxes1[..., :2] + boxes1[..., 2:] * 0.5], axis=-1)
    boxes2 = np.concatenate([boxes2[..., :2] - boxes2[..., 2:] * 0.5,
                             boxes2[..., :2] + boxes2[..., 2:] * 0.5], axis=-1)

    left_up = np.maximum(boxes1[..., :2], boxes2[..., :2])
    right_down = np.minimum(boxes1[..., 2:], boxes2[..., 2:])

    inter_section = np.maximum(right_down - left_up, 0.0)
    inter_area = inter_section[..., 0] * inter_section[..., 1]
    union_area = boxes1_area + boxes2_area - inter_area

    return inter_area / union_area

def bbox_giou(boxes1, boxes2):
    boxes1 = tf.concat([boxes1[..., :2] - boxes1[..., 2:] * 0.5,
                        boxes1[..., :2] + boxes1[..., 2:] * 0.5], axis=-1)
    boxes2 = tf.concat([boxes2[..., :2] - boxes2[..., 2:] * 0.5,
                        boxes2[..., :2] + boxes2[..., 2:] * 0.5], axis=-1)

    boxes1 = tf.concat([tf.minimum(boxes1[..., :2], boxes1[..., 2:]),
                        tf.maximum(boxes1[..., :2], boxes1[..., 2:])], axis=-1)
    boxes2 = tf.concat([tf.minimum(boxes2[..., :2], boxes2[..., 2:]),
                        tf.maximum(boxes2[..., :2], boxes2[..., 2:])], axis=-1)

    boxes1_area = (boxes1[..., 2] - boxes1[..., 0]) * (boxes1[..., 3] - boxes1[..., 1])
    boxes2_area = (boxes2[..., 2] - boxes2[..., 0]) * (boxes2[..., 3] - boxes2[..., 1])

    left_up = tf.maximum(boxes1[..., :2], boxes2[..., :2])
    right_down = tf.minimum(boxes1[..., 2:], boxes2[..., 2:])

    inter_section = tf.maximum(right_down - left_up, 0.0)
    inter_area = inter_section[..., 0] * inter_section[..., 1]
    union_area = boxes1_area + boxes2_area - inter_area
    iou = inter_area / union_area

    enclose_left_up = tf.minimum(boxes1[..., :2], boxes2[..., :2])
    enclose_right_down = tf.maximum(boxes1[..., 2:], boxes2[..., 2:])
    enclose = tf.maximum(enclose_right_down - enclose_left_up, 0.0)
    enclose_area = enclose[..., 0] * enclose[..., 1]
    giou = iou - 1.0 * (enclose_area - union_area) / enclose_area

    return giou

def focal(target, actual, alpha=1, gamma=2):
    focal_loss = alpha * tf.pow(tf.abs(target - actual), gamma)
    return focal_loss

# iou_threshold相交门阀控制（超过了就剔除）
def nms(pro_need_cls_bboxes, iou_threshold, sigma=0.3, method='nms'):
    need_bboxs = []

    # 流程1: 判断边界框的数目是否大于0
    while len(pro_need_cls_bboxes) > 0:
        # 流程2: 按照 socre 排序选出评分最大的边界框 A,并将边界框 A 取出并剔除
        max_ind = np.argmax(pro_need_cls_bboxes[:, 4])
        need_bbox = pro_need_cls_bboxes[max_ind]
        need_bboxs.append(need_bbox)
        pro_need_cls_bboxes = np.concatenate([pro_need_cls_bboxes[: max_ind], pro_need_cls_bboxes[max_ind + 1:]])

        # 流程3: 计算这个边界框 A 与剩下所有边界框的 iou 并剔除那些 iou 值高于阈值的边界框
        iou = bboxes_iou(need_bbox[np.newaxis, :4], pro_need_cls_bboxes[:, :4])
        weight = np.ones((len(iou),), dtype=np.float32)

        assert method in ['nms', 'soft-nms']

        if method == 'nms':
            iou_mask = iou > iou_threshold
            weight[iou_mask] = 0.0

        if method == 'soft-nms':
            weight = np.exp(-(1.0 * iou ** 2 / sigma))

        pro_need_cls_bboxes[:, 4] = pro_need_cls_bboxes[:, 4] * weight
        score_mask = pro_need_cls_bboxes[:, 4] > 0.
        pro_need_cls_bboxes = pro_need_cls_bboxes[score_mask]
    return need_bboxs

def read_class_names(class_file_name):
    names = {}
    with open(class_file_name, 'r') as data:
        for ID, name in enumerate(data):
            names[ID] = name.strip('\n')
    return names

def get_anchors(anchors_path):
    with open(anchors_path) as f:
        anchors = f.readline()
    anchors = np.array(anchors.split(','), dtype=np.float32)
    return anchors.reshape(3, 3, 2)

def image_pretreat_process(org_image, target_size, gt_boxes=None):
    h, w = org_image.shape[:2]
    tar_h, tar_w = target_size, target_size
    h_resize, w_resize, scale = [tar_h, tar_h * w // h, target_size / h] if h > w else [tar_w * h // w, tar_w,
                                                                                target_size / w]

    org_image_temp = cv2.resize(org_image, (w_resize, h_resize))
    h_f, w_f = (tar_h - h_resize) // 2, (tar_w - w_resize) // 2
    target_image = np.full(shape=[tar_h, tar_w, 3], fill_value=127, dtype=np.uint8)
    target_image[h_f:h_f + h_resize, w_f:w_f + w_resize] = org_image_temp

    if gt_boxes is None:
        return target_image

    else:
        gt_boxes[:, [0, 2]] = gt_boxes[:, [0, 2]] * scale + w_f
        gt_boxes[:, [1, 3]] = gt_boxes[:, [1, 3]] * scale + h_f
        return target_image, gt_boxes

def random_horizontal_flip(image, bboxes):
    if random.random() < 0.5:
        _, w, _ = image.shape
        image = image[:, ::-1, :]
        bboxes[:, [0, 2]] = w - bboxes[:, [2, 0]]

    return image, bboxes

def random_crop(image, bboxes):
    if random.random() < 0.5:
        h, w, _ = image.shape
        max_bbox = np.concatenate([np.min(bboxes[:, 0:2], axis=0), np.max(bboxes[:, 2:4], axis=0)], axis=-1)

        max_l_trans = max_bbox[0]
        max_u_trans = max_bbox[1]
        max_r_trans = w - max_bbox[2]
        max_d_trans = h - max_bbox[3]

        crop_xmin = max(0, int(max_bbox[0] - random.uniform(0, max_l_trans)))
        crop_ymin = max(0, int(max_bbox[1] - random.uniform(0, max_u_trans)))
        crop_xmax = max(w, int(max_bbox[2] + random.uniform(0, max_r_trans)))
        crop_ymax = max(h, int(max_bbox[3] + random.uniform(0, max_d_trans)))

        image = image[crop_ymin: crop_ymax, crop_xmin: crop_xmax]

        bboxes[:, [0, 2]] = bboxes[:, [0, 2]] - crop_xmin
        bboxes[:, [1, 3]] = bboxes[:, [1, 3]] - crop_ymin

    return image, bboxes

def random_translate(image, bboxes):
    if random.random() < 0.5:
        h, w, _ = image.shape
        max_bbox = np.concatenate([np.min(bboxes[:, 0:2], axis=0), np.max(bboxes[:, 2:4], axis=0)], axis=-1)

        max_l_trans = max_bbox[0]
        max_u_trans = max_bbox[1]
        max_r_trans = w - max_bbox[2]
        max_d_trans = h - max_bbox[3]

        tx = random.uniform(-(max_l_trans - 1), (max_r_trans - 1))
        ty = random.uniform(-(max_u_trans - 1), (max_d_trans - 1))

        M = np.array([[1, 0, tx], [0, 1, ty]])
        image = cv2.warpAffine(image, M, (w, h))

        bboxes[:, [0, 2]] = bboxes[:, [0, 2]] + tx
        bboxes[:, [1, 3]] = bboxes[:, [1, 3]] + ty

    return image, bboxes

def read_pb_return_tensors(graph, pb_file, return_elements):
    with tf.gfile.FastGFile(pb_file, 'rb') as f:
        frozen_graph_def = tf.GraphDef()
        frozen_graph_def.ParseFromString(f.read())

    with graph.as_default():
        return_elements = tf.import_graph_def(frozen_graph_def,
                                              return_elements=return_elements)
    return return_elements

def postprocess_boxes(pred_bbox, org_img_shape, input_size, score_threshold):
    valid_scale = [0, np.inf]
    pred_bbox = np.array(pred_bbox)

    pred_xywh = pred_bbox[:, 0:4]
    pred_conf = pred_bbox[:, 4]
    pred_prob = pred_bbox[:, 5:]

    # # (1) (x, y, w, h) --> (xmin, ymin, xmax, ymax)
    pred_coor = np.concatenate([pred_xywh[:, :2] - pred_xywh[:, 2:] * 0.5,
                                pred_xywh[:, :2] + pred_xywh[:, 2:] * 0.5], axis=-1)
    # # (2) (xmin, ymin, xmax, ymax) -> (xmin_org, ymin_org, xmax_org, ymax_org)
    org_h, org_w = org_img_shape
    resize_ratio = min(input_size / org_w, input_size / org_h)

    dw = (input_size - resize_ratio * org_w) / 2
    dh = (input_size - resize_ratio * org_h) / 2

    pred_coor[:, 0::2] = 1.0 * (pred_coor[:, 0::2] - dw) / resize_ratio
    pred_coor[:, 1::2] = 1.0 * (pred_coor[:, 1::2] - dh) / resize_ratio

    # # (3) clip some boxes those are out of range
    pred_coor = np.concatenate([np.maximum(pred_coor[:, :2], [0, 0]),
                                np.minimum(pred_coor[:, 2:], [org_w - 1, org_h - 1])], axis=-1)
    invalid_mask = np.logical_or((pred_coor[:, 0] > pred_coor[:, 2]), (pred_coor[:, 1] > pred_coor[:, 3]))
    pred_coor[invalid_mask] = 0

    # # (4) discard some invalid boxes
    bboxes_scale = np.sqrt(np.multiply.reduce(pred_coor[:, 2:4] - pred_coor[:, 0:2], axis=-1))
    scale_mask = np.logical_and((valid_scale[0] < bboxes_scale), (bboxes_scale < valid_scale[1]))

    # # (5) discard some boxes with low scores
    classes = np.argmax(pred_prob, axis=-1)
    scores = pred_conf * pred_prob[np.arange(len(pred_coor)), classes]
    score_mask = scores > score_threshold
    mask = np.logical_and(scale_mask, score_mask)
    coors, scores, classes = pred_coor[mask], scores[mask], classes[mask]

    return np.concatenate([coors, scores[:, np.newaxis], classes[:, np.newaxis]], axis=-1)

def draw_bbox(image, bboxes, classes=read_class_names(config.YOLO.CLASSES), show_label=True):
    """
    bboxes: [x_min, y_min, x_max, y_max, probability, cls_id] format coordinates.
    """

    num_classes = len(classes)
    image_h, image_w, _ = image.shape
    hsv_tuples = [(1.0 * x / num_classes, 1., 1.) for x in range(num_classes)]
    colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
    colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors))

    random.seed(0)
    random.shuffle(colors)
    random.seed(None)

    for i, bbox in enumerate(bboxes):
        coor = np.array(bbox[:4], dtype=np.int32)
        fontScale = 0.5
        score = bbox[4]
        class_ind = int(bbox[5])
        bbox_color = colors[class_ind]
        bbox_thick = int(0.6 * (image_h + image_w) / 600)
        c1, c2 = (coor[0], coor[1]), (coor[2], coor[3])
        cv2.rectangle(image, c1, c2, bbox_color, bbox_thick)

        if show_label:
            bbox_mess = '%s: %.2f' % (classes[class_ind], score)
            t_size = cv2.getTextSize(bbox_mess, 0, fontScale, thickness=bbox_thick // 2)[0]
            cv2.rectangle(image, c1, (c1[0] + t_size[0], c1[1] - t_size[1] - 3), bbox_color, -1)  # filled

            cv2.putText(image, bbox_mess, (c1[0], c1[1] - 2), cv2.FONT_HERSHEY_SIMPLEX,
                        fontScale, (0, 0, 0), bbox_thick // 2, lineType=cv2.LINE_AA)

    return image

def bboxes_iou(boxes1, boxes2):
    boxes1 = np.array(boxes1)
    boxes2 = np.array(boxes2)

    boxes1_area = (boxes1[..., 2] - boxes1[..., 0]) * (boxes1[..., 3] - boxes1[..., 1])
    boxes2_area = (boxes2[..., 2] - boxes2[..., 0]) * (boxes2[..., 3] - boxes2[..., 1])

    left_up = np.maximum(boxes1[..., :2], boxes2[..., :2])
    right_down = np.minimum(boxes1[..., 2:], boxes2[..., 2:])

    inter_section = np.maximum(right_down - left_up, 0.0)
    inter_area = inter_section[..., 0] * inter_section[..., 1]
    union_area = boxes1_area + boxes2_area - inter_area
    ious = np.maximum(1.0 * inter_area / union_area, np.finfo(np.float32).eps)

    return ious


"""Look-up table for GGF computation"""
import pathlib
from pkg_resources import resource_filename
import warnings

import h5py
import numpy as np
from scipy.interpolate import interpn

class NotInLUTError(BaseException):
    pass

def _match_lut(kwargs, lut_path=None):
    """Find matching LUT for a set of keyword arguments

    Parameters
    ----------
    kwargs: dict
        keyword arguments for ggf.get_ggf
    lut_path: str, pathlib.Path or None
        path to an hdf5 file containing a LUT; if None,
        the default LUTs will be searched

    Returns
    -------
    lut_path: pathlib.Path
        path to matching LUT hdf5 file
    """
    if lut_path:
        paths = [pathlib.Path(lut_path)]
    else:
        paths = get_lut_paths()
    for path in paths:
        with h5py.File(path, mode="r") as h5:
            for key in kwargs:
                if _param_in_lut(key, kwargs[key], h5["lut"].attrs):
                    pass  # everything OK so far
                else:
                    break  # try next LUT
            else:
                break  # this is the right LUT
    else:
        msg = "No matching LUT found for: {} in {}".format(kwargs, paths)
        raise NotInLUTError(msg)
    return path

def _param_in_lut(key, value, h5_attrs):
    """Test whether the given parameter is accessible via a LUT

    Parameters
    ----------
    key: str
        keyword name (see ggf.get_ggf)
    value: str or float
        keyword value to check
    h5_attrs: dict
        attributes from a LUT hdf5 file

    Returns
    -------
    inlut: bool
        True, if the key-value pair is covered by the LUT
    """
    inlut = False
    if key == "n_poly" and value is None:  # default value is given in LUT
        inlut = True
    elif key in h5_attrs:  # fixed parameter
        if h5_attrs[key] == value or value is None:
            inlut = True
    elif isinstance(value, str):  # model
        inlut = False
    else:  # range
        vmin = h5_attrs["{} min".format(key)]
        vmax = h5_attrs["{} max".format(key)]
        if value <= vmax and value >= vmin:
            inlut = True
        else:
            inlut = False
    return inlut

def get_lut_paths():
    """Return a list of look-up table hdf5 files in ggf"""
    lutpath = pathlib.Path(resource_filename("ggf", "lut"))
    paths = lutpath.glob("*.h5")
    return sorted(paths)

def get_ggf_lut(model, semi_major, semi_minor, object_index, medium_index,
                effective_fiber_distance, mode_field_diameter,
                power_per_fiber, wavelength, poisson_ratio,
                n_poly=None, lut_path=None, verbose=False):
    """Linear interpolation of the GGF from a look-up table

    Parameters
    ----------
    model: str
        Model to use, one of: `boyde2009`
    semi_major: float
        Semi-major axis of an ellipse fit to the object perimeter [m]
    semi_minor: float
        Semi-minor axis of an ellipse fit to the object perimeter [m]
    object_index: float
        Refractive index of the object
    medium_index: float
        Refractive index of the surrounding medium
    effective_fiber_distance: float
        Effective distance between the two trapping fibers relative
        to the medium refractive index [m]. For an open setup, this is
        the physical distance between the fibers. For a closed setup
        (capillary), this distance takes into account the refractive
        indices and thicknesses of the glass capillary and index
        matching gel. For the closed setup, the convenience function
        :func:`ggf.fiber_distance_capillary` can be used.
    mode_field_diameter: float
        The mode field diameter MFD of the fiber used [m]. Note that
        the MFD is dependent on the wavelength used. If the
        manufacturer did not provide a value for the MFD, the MFD
        can be approximated as ``3*wavelenth`` for a single-mode
        fiber.
    power_per_fiber: float
        The laser power coupled into each of the fibers [W]
    wavelength: float
        The laser wavelength used for the trap [m]
    poisson_ratio: float
        The Poisson's ratio of the stretched material. Set this
        to 0.5 for volume conservation.
    n_poly: int
        Number of Legendre polynomials to use for computing the GGF.
        Note that only even Legendre polynomials are used and thus,
        this number is effectively halved.
    lut_path: str or pathlib.Path
        Path to a LUT hdf5 file. If `None`, the internal LUTs are
        used.
    verbose: int
        Increases verbosity

    Returns
    -------
    ggf: float
        Linearly interpolated global geometric factor

    Notes
    -----
    - To avoid invalid values in the look-up table (LUT), such as
      `semi_major < semi_minor` or `object_index < medium_index`,
      the LUT is not built using the exact same keyword arguments
      as this method:

      - `object_index` is stored as
        ``relative_object_index = object_index / medium_index``
      - `semi_major` is stored as
        ``stretch_ratio = (semi_major - semi_minor) / semi_minor``
    - The following keywords are not interpolated in the LUT:

      - `model`
      - `wavelength`: the OS uses a fixed wavelength
      - `mode_field_diameter`: the fiber geometry is fixed
      - `power_per_fiber`: usually fixed for reproducibility
      - `n_poly`: set to a high number (e.g. 120)
    - The following are approximate guiding values for when a keyword
      can be considered linear:

      - stretch_ratio: linear only within interval of 0.004
      - semi_minor: linear only within interval of 0.08µm
      - relative_object_index: linear only within interval of 0.003
      - medium_index: linear only within interval of 0.005
      - poisson_ratio: good linearity
      - power_per_fiber: good linearity
      - effective_fiber_distance: linear only within interval of 15µm
    """
    # convert major_axis to stretch ratio
    stretch_ratio = (semi_major - semi_minor) / semi_minor
    # normalize object index with medium_index
    relative_object_index = object_index / medium_index
    # determine the correct LUT
    kwargs = {"model": model,
              "stretch_ratio": stretch_ratio,
              "semi_minor": semi_minor,
              "relative_object_index": relative_object_index,
              "medium_index": medium_index,
              "effective_fiber_distance": effective_fiber_distance,
              "mode_field_diameter": mode_field_diameter,
              "power_per_fiber": power_per_fiber,
              "wavelength": wavelength,
              "poisson_ratio": poisson_ratio,
              "n_poly": n_poly}
    # get the right LUT
    lut_path = _match_lut(kwargs, lut_path)  # raises NotInLutError
    if verbose:
        print("Using LUT path: {}".format(lut_path))
    # reproduce warning in boyde2009.core
    if model == "boyde2009" and stretch_ratio > 0.15:
        warnings.warn('Stretching ratio is high: {}'.format(stretch_ratio))
    # get LUT data
    with h5py.File(lut_path, mode="r") as h5:
        values = h5["lut"][:]
        meta = dict(h5["lut"].attrs)
    # order of interpolation dimensions
    order = meta["dimension_order"].split(",")
    # grid points
    points = []
    for label in order:
        points.append(np.linspace(meta["{} min".format(label)],
                                  meta["{} max".format(label)],
                                  meta["{} num".format(label)]))
    # interpolation coordinates
    xi = [kwargs[kk] for kk in order]
    # perform interpolation
    ggfval = interpn(points=points, values=values, xi=xi, method="linear",
                     bounds_error=True)
    ggfval = ggfval.item()
    if np.isnan(ggfval):
        raise ValueError("The value to be estimated in the LUT is `nan`. "
                         + "Either the LUT is incomplete or the specified "
                         + "model '{}' cannot handle ".format(model)
                         + "the given input parameters.")
    return ggfval


#!python
# coding=utf-8
"""
----------------------------------------------------------------
| Swarm-of-Trajectory Dynamic Simulation Package
|
| Version: 0.1
| Program Language: Python 3.6 (also support py2.7)
| Developer: Xinyan Wang
| Homepage: https://github.com/WangXinyan940/swarm-of-trajectory
----------------------------------------------------------------
"""
from __future__ import unicode_literals
from __future__ import division
from __future__ import print_function
import sys
import os
import json
import numpy as np
NAME = ""
BOHR = 5.291772108e-11  # Bohr -> m
ANGSTROM = 1e-10  # angstrom -> m
AMU = 1.660539040e-27  # amu -> kg
FS = 1e-15  # fs -> s
EH = 4.35974417e-18  # Hartrees -> J
H = 6.626069934e-34
KB = 1.38064852e-23
kconst = 500000.0e3 * (1. / 6.02e23 / ANGSTROM ** 2)

def printTitle():
    """
    Print the title of the output result.
    """
    print("""
----------------------------------------------------------------
| Swarm-of-Trajectory Dynamic Simulation Package
|
| Version: 0.1
| Program Language: Python 3.6
| Developer: Xinyan Wang
| Homepage: https://github.com/WangXinyan940/swarm-of-trajectory
----------------------------------------------------------------
    """)

def printHelp():
    """
    Print help.
    """
    print("""
    Usage:
        python sot.py -j json -t

        -j        configuration .json file
        -t        test the accuracy of template .gjf file
    """)

def readFile(fname, func):
    """
    Read text file and deal with different functions.
    """
    with open(fname, "r") as f:
        text = [i for i in f if i.strip()]
    return func(text)

def readXYZ(text):
    """
    Read xyz format text.
    """
    natoms = int(text[0].strip())
    body = text[2:natoms + 2]
    body = [i.strip().split() for i in body]
    atom = [i[0] for i in body]
    crd = [[float(j) for j in i[1:]] for i in body]
    return atom, np.array(crd)

def readMultiXYZ(text):
    """
    Read XYZ file with multi conformations.
    """
    xyzs = []
    ip = 0
    while True:
        natom = int(text[ip].strip())
        xyzs.append(text[ip:ip + natom + 2])
        ip = ip + natom + 2
        if ip >= len(text):
            break
    return [readXYZ(i) for i in xyzs]

def readGauGrad(text, natoms):
    """
    Read Gaussian output and find energy gradient.
    """
    ener = [i for i in text if "SCF Done:" in i]
    if len(ener) != 0:
        ener = ener[-1]
        ener = np.float64(ener.split()[4])
    else:
        ener = np.float64([i for i in text if "Energy=" in i][-1].split()[1])
    for ni, li in enumerate(text):
        if "Forces (Hartrees/Bohr)" in li:
            break
    forces = text[ni + 3:ni + 3 + natoms]
    forces = [i.strip().split()[-3:] for i in forces]
    forces = [[np.float64(i[0]), np.float64(i[1]), np.float64(i[2])]
              for i in forces]
    return ener * EH, - np.array(forces) * EH / BOHR

def genQMInput(atom, crd, temp, pre=False, nstep=-1):
    """
    Generate QM Input file for force calculation.
    """
    with open(temp, "r") as f:
        temp = f.readlines()
    wrt = []
    if pre == True:
        wrt.append("%oldchk=old.chk\n")
    for line in temp:
        if "[title]" in line:
            wrt.append("Temparary input file for step %i\n" % nstep)
        elif "[coord]" in line:
            for ni in range(len(atom)):
                wrt.append("%s  %16.8f %16.8f %16.8f\n" %
                           (atom[ni], crd[ni][0], crd[ni][1], crd[ni][2]))
        elif line[0] == "#":
            wrt.append(line)
            if pre == True:
                wrt.append(" guess=read\n")
        else:
            wrt.append(line)
    return "".join(wrt)

def writeXYZ(fname, atom, xyz, title="Title", append=False):
    """
    Write file with XYZ format.
    """
    if append:
        f = open(fname, "a")
    else:
        f = open(fname, "w")
    f.write("%i\n" % len(atom))
    f.write("%s\n" % (title.rstrip()))
    for i in range(len(atom)):
        x, y, z = xyz[i, :]
        f.write("%s  %12.8f %12.8f %12.8f\n" % (atom[i], x, y, z))
    f.close()

def distance(crd, i, j):
    """
    Calc distance of two points.
    """
    return np.sqrt(((crd[i, :] - crd[j, :]) ** 2).sum())

def bondforce(vi, vj, b, k):
    """
    Calculate force on bond
    """
    r = np.sqrt(((vi - vj) ** 2).sum())
    fr = 2 * k * abs(r - b)
    if r < b:
        gi = (vi - vj) / r * fr
    else:
        gi = (vj - vi) / r * fr
    gj = - gi
    return gi, gj

def angle(crd, i, j, k):
    """
    Calculate i-j-k angle in rad unit.
    """
    pass

def angleforce(vi, vj, vk, b, k):
    """
    Harmonic force to fix i-j-k angle at value b (rad).
    """
    pass

def genMassMat(atom):
    """
    Generate matrix of mass.
    """
    massd = {"H": 1.008,
             "C": 12.011,
             "N": 14.007,
             "O": 15.999,
             "S": 32.066,
             "CL": 35.453,
             "BR": 79.904, }
    massv = np.array([massd[i.upper()] for i in atom])
    massm = np.zeros((len(atom), 3))
    massm[:, 0] = massv
    massm[:, 1] = massv
    massm[:, 2] = massv
    return massm

def testTemplate(conf):
    """
    Test whether the template file is correct. (Use water)
    """
    t_atom = ["H", "O", "H"]
    t_xyz = [[1.0, 0.0, 0.0],
             [0.0, 0.0, 0.0],
             [0.0, 1.0, 0.0]]
    template = conf["force"]["template"]
    print(">>> Generate template below:\n++++++++++++++++++++++++++++")
    print(genQMInput(t_atom, t_xyz, template, pre=True))

def argparse():
    """
    Parse the args.
    """
    # Test qm template
    if "-j" not in sys.argv and "-t" not in sys.argv:
        printHelp()
        exit()

    with open(sys.argv[2], "r") as f:
        text = "".join(f)
        conf = json.loads(text)

    if "-t" in sys.argv:
        testTemplate(conf)
        exit()
    return conf

def calcGauGrad(atom, crd, template, nstep, path="g09"):
    """
    Calculate gradient using Gaussian.
    """
    with open("tmp.gjf", "w") as f:
        f.write(genQMInput(atom, crd, template,
                           pre=True if nstep > 0 else False, nstep=nstep))
    os.system("{} tmp.gjf".format(path))
    grad = readFile("tmp.log", lambda x: readGauGrad(x, len(atom)))
    os.system("cp tmp.chk old.chk")
    return grad

def genGrad(conf, template):
    """
    Generate function used to calculate energy gradient.
    """
    if conf["engine"].upper() == "GAUSSIAN":
        return lambda atom, crd, nstep: calcGauGrad(atom, crd / ANGSTROM, template, nstep, conf["path"])

def setInitMotion(conf):
    """
    Set init coord and velocities.
    """
    xyzs = readFile(conf["coordinate"], readMultiXYZ)
    if "velocity" in conf and conf["velocity"] is not None:
        vels = readFile(conf["velocity"], readMultiXYZ)
    else:
        vels = None
    if "start" not in conf or conf["start"] is None or conf["start"] > len(xyzs):
        atom, crd = xyzs[-1]
        crd = crd * ANGSTROM
        _, vel = None, vels[-1][1] if vels is not None else None
    else:
        rdm = np.random.randint(conf["start"], len(xyzs))
        atom, crd = xyzs[rdm]
        crd = crd * ANGSTROM
        _, vel = None, vels[rdm][1] if vels is not None else None
    if vel is None:
        T = conf["temperature"]
        massm = genMassMat(atom) * AMU
        vel = np.random.normal(0.0, np.sqrt(KB * T / massm))
    else:
        vel = vel * ANGSTROM / FS
    return atom, crd, vel

def dynamics(atom, initx, initv, grad=None, conf=None):
    """
    Run dynamics. Using Velocity verlet algorithm.
    """
    md, prt, cons, chk, stop = conf["md"], conf["print"], conf[
        "constraint"], conf["check"], conf["stop"]
    dt = md["deltat"] * FS
    massm = genMassMat(atom) * AMU
    if md["type"].upper() == "NVT":
        T = md["temperature"]
        friction = md["friction"] / (1000.0 * FS)
        kT = KB * T
        vscale = np.exp(- dt * friction)
        fscale = dt if friction == 0.0 else (1 - vscale) / friction
        noisescale = np.sqrt(kT * (1.0 - vscale ** 2))
        invmass = 1. / massm
        sqrtinvmass = np.sqrt(invmass)

    crd = initx
    vel = initv

    e, f = grad(atom, crd, nstep=-1)
    f = -f

    for nstep in range(md["nsteps"]):
        # print
        if "freq" in prt and nstep % prt["freq"] == 0:
            KE = (0.5 * massm * vel * vel).sum() / vel.shape[0] / vel.shape[1]
            Tseq = KE * 2.0 / KB
            if prt["coordinate"]:
                writeXYZ("%s-traj.xyz" % NAME, atom, crd / ANGSTROM,
                         title="NSTEP:%i E:%10.6f T:%8.4f" % (nstep, e / EH, Tseq), append=True if nstep > 0 else False)
            if prt["coordinate"]:
                writeXYZ("%s-vel.xyz" % NAME, atom, vel / (ANGSTROM / FS),
                         title="NSTEP:%i E:%10.6f T:%8.4f" % (nstep, e / EH, Tseq), append=True if nstep > 0 else False)
        if md["type"].upper() == "NVE":
            # velocity verlet
            crd = crd + vel * dt + 0.5 * (f / massm) * dt ** 2
            f_old = f
            e, f = grad(atom, crd, nstep=nstep)
            f = -f
            print(">>> step: %i    e:%10.4f" % (nstep, e / EH))
            vel = vel + 0.5 * (f_old + f) / massm * dt

        elif md["type"].upper() == "NVT":
            # langevin dynamics
            KE = (0.5 * massm * vel * vel).sum() / vel.shape[0] / vel.shape[1]
            Tseq = KE * 2.0 / KB
            print(">>> step: %i    e:%10.4f    T:%8.4f" %
                  (nstep, e / EH, Tseq))
            # step1
            p1 = vscale * vel
            p2 = fscale * invmass * f
            p3 = noisescale * sqrtinvmass * \
                np.random.normal(0., np.ones(vel.shape))
            vel = p1 + p2 + p3
            if "fixcom" in md and md["fixcom"]:
                vcom = vel.mean(axis=0)
                ke_pre = (0.5 * massm * vel * vel).sum() / \
                    vel.shape[0] / vel.shape[1]
                v_remove = vel - vcom
                ke_after = (0.5 * massm * v_remove * v_remove).sum() / \
                    v_remove.shape[0] / v_remove.shape[1]
                vel = v_remove * np.sqrt(ke_pre / ke_after)
                print(">>> remove COM motion", vcom / (ANGSTROM / FS))
            #print(p1, p2, p3)
            # step2
            pre_crd = crd

            # change cartesian coordinate
            crd = crd + vel * dt

            # LINCS algorithm
            # Build B, d, S, A
            B = np.zeros((len(cons), crd.ravel().shape[0]))
            d = np.zeros((len(cons), 1))
            S = np.zeros((len(cons), len(cons)))
            A = np.zeros(S.shape)
            for n, cv in enumerate(cons):
                if cv["type"].upper() == "B":
                    ia, ib = cv["index"]
                    value = cv["value"] * ANGSTROM
                    r = np.sqrt(np.power(crd[ia,:] - crd[ib,:], 2).sum())
                    B[n, ia * 3] = (crd[ia, 0] - crd[ib, 0]) / r
                    B[n, ia * 3 + 1] = (crd[ia, 1] - crd[ib, 1]) / r
                    B[n, ia * 3 + 2] = (crd[ia, 2] - crd[ib, 2]) / r
                    B[n, ib * 3] = (crd[ib, 0] - crd[ia, 0]) / r
                    B[n, ib * 3 + 1] = (crd[ib, 1] - crd[ia, 1]) / r
                    B[n, ib * 3 + 2] = (crd[ib, 2] - crd[ia, 2]) / r
                    d[n] = value
                    S[n, n] = 1. / np.sqrt(1. / massm[ia][0] + 1. / massm[ib][0])
            # Calculate A
            #A = np.diag(np.ones(len(cons),)) - np.dot(S, np.dot(B,np.dot(np.diag(invmass.ravel()), np.dot(B.T, S))))
            invMmat = np.diag(invmass.ravel())
            Tm = np.linalg.inv(np.dot(np.dot(B, invMmat), B.T))

            #inverse = np.diag(np.ones(len(cons),)) + A + A ** 2 + A ** 3 + A ** 4
            # Update new coord
            crd = crd - np.dot(np.dot(invMmat, B.T), np.dot(Tm, np.dot(B, crd.reshape((-1,1))) - d)).reshape((-1,3))

            # step3
            vel = (crd - pre_crd) / dt
            e, f = grad(atom, crd, nstep=nstep)
            f = -f
        # check_traj
        if "time" in chk and nstep == chk["time"]:
            for cv in chk["cv"]:
                if cv["type"].upper() == "B":
                    r = distance(crd / ANGSTROM,
                                 cv["index"][0], cv["index"][1])
                    if r < cv["range"][0] or r > cv["range"][1]:
                        print(">>> Bond %i-%i out of range. Stop." %
                              (cv["index"][0], cv["index"][1]))
                        return
        # check_stop
        for state in stop:
            ifquit = True
            for cv in state["cv"]:
                if cv["type"].upper() == "B":
                    r = distance(crd / ANGSTROM,
                                 cv["index"][0], cv["index"][1])
                    if r < cv["range"][0] or r > cv["range"][1]:
                        ifquit = False
                        break
            if ifquit:
                print(">>> Get state %s. Stop." % state["name"])
                exit()

def main():
    """
    The main function.
    """
    printTitle()
    conf = argparse()

    global NAME
    NAME = conf["name"]

    if "restraint" not in conf:
        conf["restraint"] = []
    if "constraint" not in conf:
        conf["constraint"] = []

    # build template for qm engine
    template = conf["force"]["template"]

    grad = genGrad(conf["force"], template)
    # select init crd and vel
    atom, crd, vel = setInitMotion(conf["init"])
    # run dynamics
    dynamics(atom, crd, vel, grad=grad, conf=conf)
    print(">>> STDSP is finished.")

if __name__ == '__main__':
    main()


'''
This code is to implement embedding loss.
'''
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.autograd import Function, Variable

class EmbLoss_v1(nn.Module):
    def __init__(self, feature_dim=4, loss_weight=1.0):
        super(EmbLoss_v1, self).__init__()
        self.feature_dim = feature_dim
        self.loss_weight = loss_weight
        self.delta_v = 0.5 # delta_agg
        self.delta_d = 1.5 # delta_dis
        self.weights = (1.0, 1.0)

    def forward_single(self, emb, instance, kernel, training_mask, bboxes):
        training_mask = (training_mask > 0.5).long()
        kernel = (kernel > 0.5).long()
        instance = instance * training_mask
        instance_kernel = (instance * kernel).view(-1)
        instance = instance.view(-1)
        emb = emb.view(self.feature_dim, -1)

        unique_labels, unique_ids = torch.unique(instance_kernel, sorted=True, return_inverse=True)
        num_instance = unique_labels.size(0)
        if num_instance <= 1:
            return 0

        emb_mean = emb.new_zeros((self.feature_dim, num_instance), dtype=torch.float32)
        for i, lb in enumerate(unique_labels):
            if lb == 0:
                continue
            ind_k = instance_kernel == lb
            emb_mean[:, i] = torch.mean(emb[:, ind_k], dim=1)

        l_agg = emb.new_zeros(num_instance, dtype=torch.float32)  # bug
        for i, lb in enumerate(unique_labels):
            if lb == 0:
                continue
            ind = instance == lb
            emb_ = emb[:, ind]
            dist = (emb_ - emb_mean[:, i:i + 1]).norm(p=2, dim=0)
            dist = F.relu(dist - self.delta_v) ** 2
            l_agg[i] = torch.mean(torch.log(dist + 1.0))
        l_agg = torch.mean(l_agg[1:])

        if num_instance > 2:
            emb_interleave = emb_mean.permute(1, 0).repeat(num_instance, 1)
            emb_band = emb_mean.permute(1, 0).repeat(1, num_instance).view(-1, self.feature_dim)
            # print(seg_band)

            mask = (1 - torch.eye(num_instance, dtype=torch.int8)).view(-1, 1).repeat(1, self.feature_dim)
            mask = mask.view(num_instance, num_instance, -1)
            mask[0, :, :] = 0
            mask[:, 0, :] = 0
            mask = mask.view(num_instance * num_instance, -1)
            # print(mask)

            dist = emb_interleave - emb_band
            dist = dist[mask > 0].view(-1, self.feature_dim).norm(p=2, dim=1)
            dist = F.relu(2 * self.delta_d - dist) ** 2
            l_dis = torch.mean(torch.log(dist + 1.0))
        else:
            l_dis = 0

        l_agg = self.weights[0] * l_agg
        l_dis = self.weights[1] * l_dis
        l_reg = torch.mean(torch.log(torch.norm(emb_mean, 2, 0) + 1.0)) * 0.001
        loss = l_agg + l_dis + l_reg
        return loss

    def forward(self, emb, instance, kernel, training_mask, bboxes, reduce=True):
        # TO CHECK: bboxes needs to be removed?
        loss_batch = emb.new_zeros((emb.size(0)), dtype=torch.float32)

        for i in range(loss_batch.size(0)):
            loss_batch[i] = self.forward_single(emb[i], instance[i], kernel[i], training_mask[i], bboxes[i])

        loss_batch = self.loss_weight * loss_batch

        if reduce:
            loss_batch = torch.mean(loss_batch)

        return loss_batch


__all__ = [
    "check_dp",
    "t_eff",
    "paramget",
    "berv",
    "v_orb",
    "astropyberv",
    "calculateberv",
    "phase",
    "RV_star",
    "transit",
    "RV",
    "dRV"
]

def check_dp(dp):
    """
    This is a helper function that checks for the presence of a config file
    and an obs_times file at path dp. It also checks that dp is a Path object and returns
    it as such.
    """
    from tayph.util import check_path
    from pathlib import Path
    check_path(dp)
    if isinstance(dp,str):
        dp = Path(dp)
    p1=dp/'obs_times'
    p2=dp/'config'
    check_path(p1,exists=True)
    check_path(p2,exists=True)
    return(dp)

def paramget(keyword,dp,full_path=False):
    """This code queries a planet system parameter from a config file located in the folder
    specified by the path dp; or run configuration parameters from a file speciefied by the full
    path dp, if full_path is set to True.

    Parameters
    ----------
    keyword : str
        A keyword present in the cofig file.

    dp : str, Path
        Output filename/path.

    full_path: bool
        If set, dp refers to the actual file, not the location of a folder with a config.dat;
        but the actual file itself.

    Returns
    -------
    value : int, float, bool, str
        The value corresponding to the requested keyword.

    """
    from tayph.vartests import typetest
    from tayph.util import check_path
    import pathlib
    import distutils.util
    import pdb

    #This is probably the only case where I don't need obs_times and config to exist together...
    dp=check_path(dp)
    typetest(keyword,str,'keyword in paramget()')

    if isinstance(dp,str) == True:
        dp=pathlib.Path(dp)
    try:
        if full_path == False:
            dp = dp/'config'
        f = open(dp, 'r')
    except FileNotFoundError:
        raise FileNotFoundError('parameter file does not exist at %s' % str(dp)) from None
    x = f.read().splitlines()
    f.close()
    n_lines=len(x)
    keywords={}
    for i in range(0,n_lines):
        line=x[i].split()
        if len(line) > 1:
            try:
                value=float(line[1])
            except ValueError:
                try:
                    value=bool(distutils.util.strtobool(line[1]))
                except ValueError:
                    value=(line[1])
            keywords[line[0]] = value
    try:
        return(keywords[keyword])
    except KeyError:
        # print(keywords)
        raise Exception('Keyword %s is not present in parameter file at %s' % (keyword,dp)) from None

def t_eff(M,R):
    """This function computes the mass and radius of a star given its mass and radius relative to solar."""
    from tayph.vartests import typetest
    import numpy as np
    import astropy.constants as const

    typetest(M,[int,float],'M in t_eff()')
    typetest(R,[int,float],'R in t_eff()')
    M=float(M)
    R=float(R)

    Ms = const.M_sun
    Rs = const.R_sun
    Ls = const.L_sun
    sb = const.sigma_sb

    if M < 0.43:
        a = 0.23
        b = 2.3
    elif M < 2:
        a = 1.0
        b = 4.0
    elif M < 55:
        a = 1.4
        b = 3.5
    else:
        a = 32000.0
        b = 1.0

    T4 = a*M**b * Ls / (4*np.pi*R**2*Rs**2*sb)
    return(T4**0.25)

def berv(dp):
    """This retrieves the BERV corrcetion tabulated in the obs_times table.
    Example: brv=berv('data/Kelt-9/night1/')
    The output is an array with length N, corresponding to N exposures. These values
    are / should be taken from the FITS header.
    """
    from astropy.io import ascii
    from pathlib import Path
    dp=check_dp(dp)#Path object

    d=ascii.read(dp/'obs_times',comment="#")
    try:
        berv = d['col5']#Needs to be in col 5.
    except:
        raise Exception(f'Runtime error in sp.berv(): col5 could not be indexed. Check the integrity of your obst_times file located at {dp}.')
    return berv.data

def v_orb(dp):
    """
    This program calculates the orbital velocity in km/s for the planet in the
    data sequence provided in dp, the data-path. dp starts in the root folder,
    i.e. it starts with data/projectname/. This assumes a circular orbit.

    The output is a number in km/s.

    Parameters
    ----------
    dp : str, path like
        The path to the dataset containing the config file.

    Returns
    -------
    v_orb : float
        The planet's orbital velocity.

    list_of_sigmas_corrected : list
        List of 2D error matrices, telluric corrected.

    """
    import numpy as np
    import pdb
    import astropy.units as u
    from tayph.vartests import typetest,postest

    dp=check_dp(dp)#Path object
    P=paramget('P',dp)
    r=paramget('a',dp)
    typetest(P,float,'P in sp.v_orb()')
    typetest(r,float,'r in sp.v_orb()')
    postest(P,'P in sp.v_orb()')
    postest(r,'r in sp.v_orb()')

    return (2.0*np.pi*r*u.AU/(P*u.d)).to('km/s').value

def astropyberv(dp):
    """
    This does the same as berv(dp), but uses astropy to compute the BERV for the
    dates of observation given a data parameter file.
    Useful if the BERV keyword was somehow wrong or missing, or if you wish to
    cross-validate. Requires latitude, longitude, ra, dec and elevation to be provided in
    the config file as lat, long, RA, DEC and elev in units of degrees and meters.
    Date should be provided as mjd.
    """
    from tayph.vartests import typetest
    from pathlib import Path
    import numpy as np
    from astropy.io import ascii
    from astropy.time import Time
    from astropy import units as u
    from astropy.coordinates import SkyCoord, EarthLocation
    dp=check_dp(dp)#Path object
    d=ascii.read(dp/'obs_times',comment="#")#,names=['mjd','time','exptime','airmass'])
    #Not using named columns because I may not know for sure how many columns
    #there are, and read-ascii breaks if only some columns are named.
    #The second column has to be an MJD date array though.
    dates = d['col1']
    RA=paramget('RA',dp)
    DEC=paramget('DEC',dp)
    typetest(RA,str,'RA in sp.astropyberv()')
    typetest(DEC,str,'DEC in sp.astropyberv()')
    berv = []
    observatory = EarthLocation.from_geodetic(lat=paramget('lat',dp)*u.deg, lon=paramget('long',dp)*u.deg, height=paramget('elev',dp)*u.m)
    sc = SkyCoord(RA+' '+DEC, unit=(u.hourangle, u.deg))
    for date in dates:
        barycorr = sc.radial_velocity_correction(obstime=Time(date,format='mjd'), location=observatory).to(u.km/u.s)
        berv.append(barycorr.value)
    return berv

def calculateberv(date,lat,long,elev,ra,dec):
    """This is a copy of the astropyberv above, but as a function for a single
    date in mjd. lat, long, RA, DEC and elev are in units of degrees and meters."""
    from astropy.time import Time
    from astropy import units as u
    from astropy.coordinates import SkyCoord, EarthLocation
    observatory = EarthLocation.from_geodetic(lat=lat*u.deg, lon=long*u.deg, height=elev*u.m)
    sc = SkyCoord(f'{ra} {dec}', unit=(u.hourangle, u.deg))
    barycorr = sc.radial_velocity_correction(obstime=Time(date,format='mjd'), location=observatory).to(u.km/u.s)
    return(barycorr.value)

def phase(dp):
    """
    Calculates the orbital phase of the planet in the data
    sequence provided using the parameters in dp/config and the timings in
    dp/obstimes.

    The output is an array with length N, corresponding to N exposures.

    Be CAREFUL: This program provides a time difference of ~1 minute compared
    to IDL/calctimes. This likely has to do with the difference between HJD
    and BJD, and the TDB timescale. In the future you should have a thorough
    look at the time-issue, because you should be able to get this right to the
    second. At the very least, make sure that the time conventions are ok.

    More importantly: The transit center time needs to be provided in config
    in BJD.
    """
    from tayph.vartests import typetest
    import numpy as np
    from astropy.io import ascii
    from astropy.time import Time
    from astropy import units as u, coordinates as coord
    import tayph.util as ut
    dp=check_dp(dp)#Path object
    d=ascii.read(dp/'obs_times',comment="#")#,names=['mjd','time','exptime','airmass'])
    #Not using the named columns because I may not know for sure how many columns
    #there are, and read-ascii breaks if only some columns are named.
    #The second column has to be a date array though.

    # t = Time(d['col2'],scale='utc', location=coord.EarthLocation.of_site('paranal'))# I determined that the difference between this and geodetic 0,0,0 is zero.
    t = Time(d['col2'],scale='utc', location=coord.EarthLocation.from_geodetic(0,0,0))

    jd = t.jd
    P=paramget('P',dp)
    RA=paramget('RA',dp)
    DEC=paramget('DEC',dp)
    Tc=paramget('Tc',dp)#Needs to be given in BJD!

    typetest(P,float,'P in sp.phase()')
    typetest(Tc,float,'Tc in sp.phase()')
    typetest(RA,str,'RA in sp.phase()')
    typetest(DEC,str,'DEC in sp.phase()')

    ip_peg = coord.SkyCoord(RA,DEC,unit=(u.hourangle, u.deg), frame='icrs')
    ltt_bary = t.light_travel_time(ip_peg)

    n=0.0
    Tc_n=Time(Tc,format='jd',scale='tdb')
    while Tc_n.jd >= min(jd):
        Tc_n=Time(Tc-100.0*n*P,format='jd',scale='tdb')#This is to make sure that the Transit central time PRECEDES the observations (by tens or hundreds or thousands of years). Otherwise, the phase could pick up a minus sign somewhere and be flipped. I wish to avoid that.
        n+=1
    BJD = t.tdb + ltt_bary
    diff = BJD-Tc_n
    phase=((diff.jd) % P)/P
    return phase

def transit(dp):
    """This code uses Ians astro python routines for the approximate Mandel &
    Agol transit lightcurve to produce the predicted transit lightcurve for the
    planet described by the configfile located at dp/config.
    This all assumes a circular orbit.
    ===========
    Derivation:
    ===========
    occultnonlin_small(z,p, cn) is the algorithm of the Mandel&Agol derivation.
    z = d/R_star, where d is the distance of the planet center to the LOS to the
    center of the star.
    sin(alpha) = d/a, with a the orbital distance (semi-major axis).
    so sin(alpha)*a/Rstar = d/a*a/Rstar = d/Rstar = z.
    a/Rstar happens to be a quantity that is well known from the transit light-
    curve. So z = sin(2pi phase)*a/Rstar. But this is in the limit of i = 90.

    From Cegla 2016 it follows that z = sqrt(xp^2 + yp^2). These are given
    as xp = a/Rstar sin(2pi phase) and yp = -a/Rstar * cos(2pi phase) * cos(i).

    The second quantity, p, is Rp/Rstar, also well known from the transit light-
    curve.

    cn is a four-element vector with the nonlinear limb darkening coefficients.
    If a shorter sequence is entered, the later values will be set to zero.
    By default I made it zero; i.e. the injected model does not take into
    account limb-darkening.
    """
    from tayph.vartests import typetest
    import tayph.util as ut
    import tayph.iansastropy as iap
    import numpy as np
    import pdb
    dp=ut.check_path(dp)
    p=phase(dp)
    a_Rstar=paramget('aRstar',dp)
    Rp_Rstar=paramget('RpRstar',dp)
    i=paramget('inclination',dp)
    typetest(a_Rstar,float,'Rp_Rstar')
    typetest(a_Rstar,float,'a_Rstar')
    typetest(i,float,'i')

    xp=np.sin(p*2.0*np.pi)*a_Rstar
    yp=np.cos(p*2.0*np.pi)*np.cos(np.radians(i))*a_Rstar
    z=np.sqrt(xp**2.0 + yp**2.0)
    transit=iap.occultnonlin_small(z,Rp_Rstar,[0.0,0.0])
    return transit

def RV_star(dp):
    """
    This calculates the radial velocity in km/s for the star in the
    data sequence provided in dp. The output is an array with length N,
    corresponding to N exposures. The radial velocity is provided in km/s.
    This is meant to be used to correct (align) the stellar spectra to the same
    reference frame. It requires K (the RV-semi amplitude to be provided in the
    config file, in km/s as well. Often this value is given in discovery papers.
    Like all my routines, this assumes a circular orbit.
    """
    from tayph.vartests import typetest
    import numpy as np
    dp=check_dp(dp)
    p=phase(dp)
    K=paramget('K',dp)
    typetest(K,float,'K in sp.RV_star()')
    rv=K*np.sin(2.0*np.pi*p) * (-1.0)
    return(rv)

def RV(dp,vorb=None,vsys=False):
    """This program calculates the radial velocity in km/s for the planet in the
    data sequence provided in dp, the data-path. dp starts in the root folder,
    i.e. it starts with data/projectname/, and it ends with a slash.

    Example: v=RV('data/Kelt-9/night1/')
    The output is an array with length N, corresponding to N exposures.
    The radial velocity is provided in km/s."""
    import tayph.util as ut
    import numpy as np
    from tayph.vartests import typetest
    dp=ut.check_path(dp)
    p=phase(dp)
    i=paramget('inclination',dp)
    typetest(i,float,'i')
    if vorb == None:
        vorb=v_orb(dp)
    typetest(vorb,float,'vorb in sp.RV')
    rv=vorb*np.sin(2.0*np.pi*p)*np.sin(np.radians(i))

    if vsys == True:
        vs=paramget('vsys',dp)
        rv+=vs
    return rv#In km/s.

def dRV(dp):
    """This program calculates the change in radial velocity in km/s for the
    planet in the data sequence provided in dp, the data-path. dp starts in the
    root folder,i.e. it starts with data/projectname/, and it ends with a slash.

    Example: dv=dRV('data/Kelt-9/night1/')
    The output is an array with length N, corresponding to N exposures.
    The change in radial velocity is calculated using the first derivative of the
    formula for RV, multiplied by the exposure time provided in obs_times.
    The answer is provided in units of km/s change within each exposure."""
    from tayph.vartests import typetest
    import numpy as np
    import astropy.units as u
    from astropy.io import ascii
    import pdb
    import tayph.util as ut
    dp=ut.check_path(dp,exists=True)
    obsp=ut.check_path(dp/'obs_times',exists=True)

    d=ascii.read(obsp,comment="#")
    #Texp=d['exptime'].astype('float')
    Texp=d['col3'].data#astype('float')
    vorb=v_orb(dp)
    p=phase(dp)
    P=paramget('P',dp)
    i=paramget('inclination',dp)
    typetest(P,float,'P in dRV()')
    typetest(i,float,'i in dRV()')

    dRV=vorb*np.cos(2.0*np.pi*p)*2.0*np.pi/((P*u.d).to('s').value)*np.sin(np.radians(i))
    return abs(dRV*Texp)

import numpy,import astropy,from astropy


# adjust python path so we may import things from peer packages
import sys
import os
import numpy as np
import pandas as pd
import datetime
from optparse import OptionParser
from version import VERSION_DESC

import terrain
import connection
import zone
import scenario
import curve
import curve_log
import dbimport
import logger
import database
import debris
import wateringress
import engine
from house import zoneByLocationMap, connByZoneTypeMap, ctgMap, \
    connByTypeGroupMap, inflZonesByConn, connByTypeMap, connByIDMap, zoneByIDMap

class CurvePlot(object):
    def __init__(self, coeff, method, xmin, xmax, label='_nolegend_', col='b'):
        self.coeff = coeff
        self.method = method
        self.x_arr = np.linspace(xmin, xmax, 500)
        self.label = label
        self.col = col

    def plot_vuln(self, faint=False):
        alpha = 0.3 if faint else 1.0
        col = 'r' if faint else 'b'
        if self.method == 'lognormal':
            obs = curve_log.generate_observations(self.coeff, self.x_arr)
            output.plot_fitted_curve(self.x_arr, obs, self.label, alpha, col)
        else:
            obs = curve.generate_observations(self.coeff, self.x_arr)
            output.plot_fitted_curve(self.x_arr, obs, self.label, alpha, col)

    def plot_frag(self, faint=False):
        alpha = 0.3 if faint else 1.0
        obs = curve_log.generate_observations(self.coeff, self.x_arr)
        output.plot_fragility_curve(self.x_arr, obs, self.label, alpha,
                                    self.col)

class WindDamageSimulator(object):
    """
    WindDamageSimulator: Stores sampled state (PDF) for house and wind for
    current simulation loop

    # results (for v) stored in dictionary bucket keyed with wind_speed
    # each entry has list form: (FLD_MEAN, [FLD_DIARRAY], [FLD_FRAGILITIES],
    # FLD_PRESSURIZED_COUNT, [FLD_DEBRIS_AT])
    # subarrays are indexed by house(iteration)

    """
    FLD_MEAN = 0
    FLD_DIARRAY = 1
    FLD_FRAGILITIES = 2
    FLD_PRESSURIZED_COUNT = 3
    FLD_DEBRIS_AT = 4
    FLD_WI_AT = 5
    FLD_DEBRIS_NV_AT = 6
    FLD_DEBRIS_NUM_AT = 7
    result_buckets = {}

    # record: fragility level name, level, plot color, CurvePlot
    fragility_thresholds = [['slight', 0.0, 'b', None],
                            ['medium', 0.0, 'g', None],
                            ['severe', 0.0, 'y', None],
                            ['complete', 0.0, 'r', None]]

    def __init__(self, options, diCallback=None, mplDict=None):
        self.s = None
        self.debrisManager = None
        self.A_final = None
        self.diCallback = diCallback
        self.wind_orientation = 0
        self.mzcat = 0
        self.di = 0
        self.mplDict = mplDict
        self.prevCurvePlot = None
        self.options = options

        # Undefined and later added
        self.result_wall_collapse = None
        self.cols = None
        self.rows = None
        self.house_number = None
        self.qz = 0.0
        self.Ms = 1.0
        self.di = 0.0
        self.fragilities = []
        self.profile = None
        self.mzcat = None
        self.cpiAt = None
        self.cpi = None
        self.internally_pressurized = None
        self.construction_level = None
        self.cpiAt = None
        self.prev_di = None
        self.water_ingress_cost = None
        self.file_cpis = None
        self.file_debris = None
        self.file_damage = None
        self.file_dmg = None
        self.file_frag = None
        self.file_water = None
        self.speeds = None
        self.dmg_map = None
        self.frag_levels = None
        self.wind_speeds = None
        self.di_means = None
        self.ss = None

        global output
        if mplDict is not None:
            output = __import__('gui.output').output
            output.hookupWidget(mplDict)
            self.mplDict = mplDict
        else:
            output = __import__('core.output').output

        terrain.populate_wind_profile_by_terrain()
        self.clear_loop_results()

    @staticmethod
    def set_fragility_thresholds(thresholds):
        # thresholds is a dict in form {'slight': 0.0}
        for thresh_key in thresholds:
            for tup in WindDamageSimulator.fragility_thresholds:
                if tup[0] == thresh_key:
                    tup[1] = thresholds[thresh_key]
                    break

    def set_scenario(self, s):
        """
        simulate an intact house
        Args:
            s: an instance of Scenario class

        Returns:

        """
        self.s = s
        self.cols = [chr(x) for x in range(ord('A'), ord('A') +
                                           self.s.house.roof_columns)]
        self.rows = range(1, self.s.house.roof_rows + 1)
        self.set_fragility_thresholds(s.fragility_thresholds)
        self.s.house.clear_sim_results()

    def clear_loop_results(self):
        self.qz = 0.0
        self.Ms = 1.0
        self.di = 0.0
        self.fragilities = []

    def set_wind_direction(self):
        self.wind_orientation = self.s.getWindDirIndex()
        if self.debrisManager:
            self.debrisManager.set_wind_direction_index(self.wind_orientation)

    def set_wind_profile(self):
        self.profile = np.random.random_integers(1, 10)
        self.mzcat = terrain.calculateMZCAT(self.s.terrain_category,
                                            self.profile,
                                            self.s.getHouseHeight())

    def calculate_qz(self, wind_speed):
        if self.s.regional_shielding_factor <= 0.85:
            thresholds = np.array([63, 63+15])
            ms_dic = {0: 1.0, 1: 0.85, 2: 0.95}
            idx = sum(thresholds <= np.random.random_integers(0, 100))
            self.Ms = ms_dic[idx]
            Vmod = (wind_speed * self.Ms) / self.s.regional_shielding_factor
            self.qz = 0.6 * 1.0e-3 * (Vmod * self.mzcat)**2
        else:
            self.qz = 0.6 * 1.0e-3 * (wind_speed * self.mzcat)**2

    def check_pressurized_failure(self, v):
        if self.s.getOpt_Debris():
            self.debrisManager.run(v)
            if self.cpi == 0 and self.debrisManager.get_breached():
                self.cpi = 0.7
                self.cpiAt = v
                self.file_cpis.write('%d,%.3f\n' % (self.house_number + 1, v))

    def sample_house_and_wind_params(self):
        self.cpi = 0
        self.cpiAt = 0
        self.internally_pressurized = False
        self.set_wind_profile()
        self.s.house.reset_results()
        self.prev_di = 0

        self.construction_level = 'na'
        mean_factor = 1.0
        cov_factor = 1.0
        if self.s.getOpt_ConstructionLevels():
            self.construction_level, mean_factor, cov_factor = \
                self.s.sampleConstructionLevel()

        connection.assign_connection_strengths(self.s.house.connections,
                                               mean_factor,
                                               cov_factor)

        connection.assign_connection_deadloads(self.s.house.connections)

        zone.sample_zone_pressures(self.s.house.zones,
                                   self.wind_orientation,
                                   self.s.house.cpe_V,
                                   self.s.house.cpe_k,
                                   self.s.house.cpe_struct_V)
        # we don't want to collapse multiple times (no need)
        self.result_wall_collapse = False

    def check_house_collapse(self, wind_speed):
        if not self.result_wall_collapse:
            for ctg in self.s.house.conn_type_groups:
                if ctg.trigger_collapse_at > 0:
                    perc_damaged = 0
                    for ct in ctg.conn_types:
                        perc_damaged += ct.perc_damaged()
                    if perc_damaged >= ctg.trigger_collapse_at:
                        for c in self.s.house.connections:
                            c.damage(wind_speed, 99.9, inflZonesByConn[c])
                        for z in self.s.house.zones:
                            z.result_effective_area = 0
                        self.result_wall_collapse = True

    def calculate_connection_group_areas(self):
        for ctg in self.s.house.conn_type_groups:
            ctg.result_area = 0.0
        for c in self.s.house.connections:
            c.ctype.group.result_area += c.ctype.costing_area

    def calculate_damage_ratio(self, wind_speed):

        # calculate damage percentages
        for ctg in self.s.house.conn_type_groups:
            ctg.result_percent_damaged = 0.0
            if ctg.group_name == 'debris':
                if not self.debrisManager:
                    ctg.result_percent_damaged = 0
                else:
                    ctg.result_percent_damaged = \
                        self.debrisManager.result_dmgperc
            else:
                for ct in ctg.conn_types:
                    for c in ct.connections_of_type:
                        if c.result_damaged:
                            ctg.result_percent_damaged += \
                                c.ctype.costing_area / float(ctg.result_area)

        # calculate repair cost
        repair_cost = 0
        for ctg in self.s.house.conn_type_groups:
            ctg_perc = ctg.result_percent_damaged
            if ctg_perc > 0:
                fact_arr = [0]
                for factor in self.s.house.factorings:
                    if factor.parent_id == ctg.id:
                        factor_perc = factor.factor.result_percent_damaged
                        if factor_perc:
                            fact_arr.append(factor_perc)
                max_factor_perc = max(fact_arr)
                if ctg_perc > max_factor_perc:
                    ctg_perc = ctg_perc - max_factor_perc
                    repair_cost += ctg.costing.calculate_damage(ctg_perc)

        # calculate initial envelope repair cost before water ingress is added
        self.di = repair_cost / self.s.house.replace_cost
        if self.di > 1.0:
            self.di = 1.0
        else:
            self.water_ingress_cost = 0
            if self.s.getOpt_WaterIngress():
                self.water_ingress_cost = \
                    wateringress.get_costing_for_envelope_damage_at_v(
                        self.di,
                        wind_speed,
                        self.s.house.water_groups,
                        self.file_water)
                repair_cost += self.water_ingress_cost

        # combined internal + envelope damage costing can now be calculated
        self.di = repair_cost / self.s.house.replace_cost
        if self.di > 1.0:
            self.di = 1.0
        self.prev_di = self.di

    def redistribute_damage(self, ctg):
        # setup for distribution
        if ctg.distribution_order <= 0:
            return
        distByCol = ctg.distribution_direction == 'col'
        primaryDir = self.cols
        secondaryDir = self.rows
        if not distByCol:
            primaryDir = self.rows
            secondaryDir = self.cols

        # walk the zone grid for current group
        # (only one conn of each group per zone)
        for i in primaryDir:
            for j in secondaryDir:
                # determine zoneLocation and then zone
                zoneLoc = i if distByCol else j
                zoneLoc += str(j) if distByCol else str(i)

                # not all grid locations have a zone
                if zoneLoc not in zoneByLocationMap:
                    continue

                # not all zones have area or connections remaining
                z = zoneByLocationMap[zoneLoc]
                if z.result_effective_area == 0.0 or len(z.located_conns) == 0:
                    continue

                # not all zones have connections of all types
                if ctg.group_name not in connByZoneTypeMap[z.zone_name]:
                    continue

                # grab appropriate connection from zone
                c = connByZoneTypeMap[z.zone_name][ctg.group_name]

                # if that connection is (newly) damaged then redistribute
                # load/infl/area
                if c.result_damaged and not c.result_damage_distributed:
                    # print 'Connection: %s newly damaged' % c

                    if ctg.patch_distribution == 1:
                        patchList = \
                            database.db.qryConnectionPatchesFromDamagedConn(c.id)

                        for patch in patchList:
                            patch_zone = zoneByIDMap[patch[1]]
                            patch_conn = connByIDMap[patch[0]]
                            curr_infl = inflZonesByConn[patch_conn].get(
                                patch_zone, None)
                            if curr_infl is None:
                                # print 'discarding patch as no existing patch present: %s-->%s = %f' % (patch_conn, patch_zone, patch[2])
                                continue
                            # print 'patching: %s-->%s = %f' % (patch_conn, patch_zone, patch[2])
                            inflZonesByConn[patch_conn][patch_zone] = patch[2]
                    else:
                        gridCol, gridRow = zone.getGridFromZoneLoc(z.zone_name)
                        if c.edge != 3:
                            if distByCol:
                                if c.edge == 0:
                                    k = 0.5
                                    if not self.redistribute_to_nearest_zone(z, range(gridRow+1, self.s.house.roof_rows), k, ctg, gridCol, gridRow, distByCol):
                                        k = 1.0
                                    if not self.redistribute_to_nearest_zone(z, reversed(range(0, gridRow)), k, ctg, gridCol, gridRow, distByCol):
                                        self.redistribute_to_nearest_zone(z, range(gridRow+1, self.s.house.roof_rows), k, ctg, gridCol, gridRow, distByCol)
                                elif c.edge == 2:
                                    k = 1.0
                                    self.redistribute_to_nearest_zone(z, range(gridRow+1, self.s.house.roof_rows), k, ctg, gridCol, gridRow, distByCol)
                                elif c.edge == 1:
                                    k = 1.0
                                    self.redistribute_to_nearest_zone(z, reversed(range(0, gridRow)), k, ctg, gridCol, gridRow, distByCol)
                            else:
                                if c.edge == 0:
                                    k = 0.5
                                    if not self.redistribute_to_nearest_zone(z, range(gridCol+1, self.s.house.roof_columns), k, ctg, gridCol, gridRow, distByCol):
                                        k = 1.0
                                    if not self.redistribute_to_nearest_zone(z, reversed(range(0, gridCol)), k, ctg, gridCol, gridRow, distByCol):
                                        self.redistribute_to_nearest_zone(z, range(gridCol+1, self.s.house.roof_columns), k, ctg, gridCol, gridRow, distByCol)
                                elif c.edge == 2:
                                    k = 1.0
                                    self.redistribute_to_nearest_zone(z, range(gridCol+1, self.s.house.roof_columns), k, ctg, gridCol, gridRow, distByCol)
                                elif c.edge == 1:
                                    k = 1.0
                                    self.redistribute_to_nearest_zone(z, reversed(range(0, gridCol)), k, ctg, gridCol, gridRow, distByCol)

                    if ctg.set_zone_to_zero > 0:
                        z.result_effective_area = 0.0
                    c.result_damage_distributed = True

    @staticmethod
    def redistribute_to_nearest_zone(zoneSrc, connRange, k, ctgroup, gridCol,
                                     gridRow, distByCol):
        for line in connRange:
            r = line
            c = gridCol
            if not distByCol:
                r = gridRow
                c = line
            zoneDest = zoneByLocationMap[zone.getZoneLocFromGrid(c, r)]
            conn = connByZoneTypeMap[zoneDest.zone_name].get(ctgroup.group_name)
            if conn:
                if not conn.result_damaged and zoneDest.result_effective_area > 0:
                    zoneDest.sampled_cpe = ((zoneDest.result_effective_area * zoneDest.sampled_cpe) +
                                            (k * zoneSrc.result_effective_area * zoneSrc.sampled_cpe)) / (zoneDest.result_effective_area +
                                            (k * zoneSrc.result_effective_area))
                    zoneDest.result_effective_area = zoneDest.result_effective_area + (k * zoneSrc.result_effective_area)
                    return True
                if conn.edge > 0:
                    return False
        return False

    def run_simulation(self, wind_speed):
        self.check_pressurized_failure(wind_speed)

        self.calculate_qz(wind_speed)
        zone.calc_zone_pressures(self.s.house.zones,
                                 self.wind_orientation,
                                 self.cpi,
                                 self.qz,
                                 self.Ms,
                                 self.s.building_spacing,
                                 self.s.getOpt_DiffShielding())

        self.file_damage.write('%d,%.3f,%s' % (self.house_number+1,
                                               wind_speed,
                                               scenario.Scenario.dirs[self.wind_orientation]))
        for ctg in self.s.house.conn_type_groups:
            connection.calc_connection_loads(wind_speed,
                                             ctg,
                                             self.s.house,
                                             self.file_damage,
                                             self.dmg_map,
                                             inflZonesByConn,
                                             connByTypeMap)
        if self.s.getOpt_DmgDistribute():
            for ctg in self.s.house.conn_type_groups:
                self.redistribute_damage(ctg)
        self.file_damage.write('\n')

        self.check_house_collapse(wind_speed)
        self.calculate_damage_ratio(wind_speed)

    def simulator_mainloop(self, verbose=False):
        date_run = datetime.datetime.now()

        # setup file based reporting (files must exist and be runnable)
        if not os.path.exists(self.options.output_folder):
            os.makedirs(self.options.output_folder)

        self.file_cpis = open(os.path.join(self.options.output_folder,
                                           'house_cpi.csv'), 'w')
        self.file_cpis.write('Simulated House #, Cpi Changed At\n')
        self.file_debris = open(os.path.join(self.options.output_folder,
                                             'wind_debris.csv'), 'w')
        header_ = ('Wind Speed(m/s),% Houses Internally Pressurized,'
                   '% Debris Damage Mean\n')
        self.file_debris.write(header_)
        self.file_damage = open(os.path.join(self.options.output_folder,
                                             'house_damage.csv'), 'w')
        self.file_dmg_idx = open(os.path.join(self.options.output_folder,
                                             'house_dmg_idx.csv'), 'w')
        self.file_dmg = open(os.path.join(self.options.output_folder,
                                          'houses_damaged_at_v.csv'), 'w')
        self.file_frag = open(os.path.join(self.options.output_folder,
                                           'fragilities.csv'), 'w')
        header_ = ('Slight Median,Slight Beta,Medium Median,Median Beta,'
                   'Severe Median,Severe Beta,Complete Median,Complete Beta\n')
        self.file_frag.write(header_)
        self.file_water = open(os.path.join(self.options.output_folder,
                                            'wateringress.csv'), 'w')
        header_ = ('V,Envelope DI,Water Damage,Damage Scenario,'
                   'Water Damage Cost,WaterCosting\n')
        self.file_water.write(header_)

        header = 'Simulated House #,Wind Speed(m/s),Wind Direction,'

        list_ = [ct.connection_type for ctg in self.s.house.conn_type_groups
                 if ctg.enabled for ct in ctg.conn_types]
        header += ','.join(list_)
        header += '\n'
        self.file_damage.write(header)

        # optionally seed random numbers
        if self.s.getOpt_SampleSeed():
            np.random.seed(42)
            zone.seed_scipy(42)
            engine.seed(42)

        # setup speeds and buckets
        self.speeds = np.linspace(self.s.wind_speed_min,
                                  self.s.wind_speed_max,
                                  self.s.wind_speed_num_steps)

        for wind_speed in self.speeds:
            type(self).result_buckets[wind_speed] = \
                [0., [], [], 0, [], [], [], []]

        # setup connections and groups
        self.s.house.clear_sim_results()
        self.calculate_connection_group_areas()

        # optionally create the debris manager and
        # make sure a wind orientation is set
        bDebris = self.s.getOpt_Debris()
        if bDebris:
            self.debrisManager = debris.DebrisManager(
                self.s.house,
                self.s.region,
                self.s.wind_speed_min,
                self.s.wind_speed_max,
                self.s.wind_speed_num_steps,
                self.s.getOpt_DebrisStaggeredSources(),
                self.s.debris_radius,
                self.s.debris_angle,
                self.s.debris_extension,
                self.s.building_spacing,
                self.s.source_items,
                self.s.flighttime_mean,
                self.s.flighttime_stddev)
        self.set_wind_direction()

        # gui bookkeeping
        if self.diCallback:
            totalLoops = self.s.num_iters * len(self.speeds)
            currentLoop = 1

        # SIMULATE HOUSES
        house_results = []
        keep_looping = True

        # iteration over samples
        for house_number in range(self.s.num_iters):
            self.house_number = house_number
            if not keep_looping:
                break

            # sample new house and wind direction (if random)
            if self.s.wind_dir_index == 8:
                self.set_wind_direction()

            self.sample_house_and_wind_params()

            # prime damage map where we track min() V that damage occurs
            # across types for this house (reporting)
            self.dmg_map = {}
            for conn_type in self.s.house.conn_types:
                self.dmg_map[conn_type.connection_type] = 99999

            # iteration over wind speed list
            for wind_speed in self.speeds:

                # simulate sampled house
                self.clear_loop_results()
                self.run_simulation(wind_speed)

                # collect results
                type(self).result_buckets[wind_speed][type(self).FLD_WI_AT].append(self.water_ingress_cost)
                type(self).result_buckets[wind_speed][type(self).FLD_DIARRAY].append(self.di)
                if bDebris:
                    type(self).result_buckets[wind_speed][type(self).FLD_DEBRIS_AT].append(self.debrisManager.result_dmgperc)
                    type(self).result_buckets[wind_speed][type(self).FLD_DEBRIS_NV_AT].append(self.debrisManager.result_nv)
                    type(self).result_buckets[wind_speed][type(self).FLD_DEBRIS_NUM_AT].append(self.debrisManager.result_num_items)
                else:
                    type(self).result_buckets[wind_speed][type(self).FLD_DEBRIS_AT].append(0.0)
                    type(self).result_buckets[wind_speed][type(self).FLD_DEBRIS_NV_AT].append(0)
                    type(self).result_buckets[wind_speed][type(self).FLD_DEBRIS_NUM_AT].append(0)

                # for all houses, count the number that were pressurized at this wind_speed
                if self.cpi != 0:
                    type(self).result_buckets[wind_speed][type(self).FLD_PRESSURIZED_COUNT] += 1

                # interact with GUI listener
                if self.diCallback:
                    currentLoop += 1
                    percLoops = (float(currentLoop) / float(totalLoops)) * 100.0
                    keep_looping = self.diCallback(wind_speed, self.di, percLoops)
                    if not keep_looping:
                        break

            # collect results to be used by the GUI client
            zone_results = {}
            for z in self.s.house.zones:
                zone_results[z.zone_name] = [z.zone_name,
                                             z.sampled_cpe,
                                             z.sampled_cpe_struct,
                                             z.sampled_cpe_eaves]

            conn_results = []
            for c in self.s.house.connections:
                conn_results.append([c.ctype.connection_type,
                                     c.location_zone.zone_name,
                                     c.result_failure_v_raw,
                                     c.result_strength,
                                     c.result_deadload,
                                     c.result_damaged_report,
                                     c.ctype.group.group_name,
                                     c.id])

            house_results.append([zone_results,
                                  self.dmg_map,
                                  self.wind_orientation,
                                  self.cpiAt,
                                  conn_results,
                                  self.construction_level])

        if keep_looping:
            # post processing of results (aggregations)
            for wind_speed in self.speeds:

                # write debris output file
                debrisarray = np.array(type(self).result_buckets[wind_speed][type(self).FLD_DEBRIS_AT])
                perc = type(self).result_buckets[wind_speed][type(self).FLD_PRESSURIZED_COUNT]/float(self.s.num_iters) * 100.0
                self.file_debris.write('%.3f,%.3f,%.3f\n' % (wind_speed,
                                                             perc,
                                                             np.mean(debrisarray) * 100.0))

                # calculate and store DI mean
                diarray = np.array(type(self).result_buckets[wind_speed][type(self).FLD_DIARRAY])
                type(self).result_buckets[wind_speed][type(self).FLD_MEAN] = np.mean(diarray)

                # calculate fragilities
                for thr in type(self).fragility_thresholds:
                    filter_ = diarray > thr[1]
                    p = (float(diarray[filter_].size)) / (float(diarray.size))
                    type(self).result_buckets[wind_speed][type(self).FLD_FRAGILITIES].append(p)

        # produce damage map report
        self.file_dmg.write('Number of Damaged Houses\n')
        self.file_dmg.write('Num Houses,%d\n' % self.s.num_iters)
        self.file_dmg.write('Wind Direction,%s\n' % scenario.Scenario.dirs[self.wind_orientation])
        self.file_dmg.write('Wind Speed(m/s)')

        # setup headers and counts
        str_ = [conn_type.connection_type for conn_type in
                self.s.house.conn_types]
        self.file_dmg.write(','.join(str_))
        self.file_dmg.write('\n')

        # we need to count houses damaged by type for each v
        counts = {}
        for wind_speed in self.speeds:
            self.file_dmg.write(str(wind_speed))

            # initialise damage counts for each conn_type to zero
            for conn_type in self.s.house.conn_types:
                counts[conn_type.connection_type] = 0

            # for all houses, increment type counts
            # if wind_speed exceeds minimum observed damages.
            for hr in house_results:
                dmg_map = hr[1]
                for conn_type in self.s.house.conn_types:
                    dmg_min = dmg_map[conn_type.connection_type]
                    if wind_speed >= dmg_min:
                        counts[conn_type.connection_type] += 1

            # write accumulated counts for this wind speed
            str_ = [str(counts[conn_type.connection_type]) for conn_type
                    in self.s.house.conn_types]
            self.file_dmg.write(','.join(str_))
            self.file_dmg.write('\n')

        # cleanup: close output files
        self.file_cpis.close()
        self.file_debris.close()
        self.file_damage.close()
        self.file_dmg.close()
        self.file_water.close()
        self.debrisManager = None

        # temporary
        di_summary = []
        for speed, value in type(self).result_buckets.iteritems():
            tmp = np.array(value[1])
            tmp_added = np.append(tmp, [speed, tmp.mean()])
            di_summary.append(tmp_added)

        columns_str = [str(x) for x in range(self.s.num_iters)]
        columns_str.append('speed')
        columns_str.append('mean')
        df_di_summary = pd.DataFrame(di_summary, columns=columns_str)
        df_di_summary.to_csv(self.file_dmg_idx, index=False)
        self.file_dmg_idx.close()

        if keep_looping:
            self.fit_fragility_curves()
            self.file_frag.close()
            runTime = (datetime.datetime.now() - date_run)
            return runTime, house_results
        else:
            self.file_frag.close()
            return None, None

    def get_windresults_perc_houses_breached(self):
        breaches = []
        for wind_speed in self.speeds:
            perc = type(self).result_buckets[wind_speed][type(self).FLD_PRESSURIZED_COUNT]/float(self.s.num_iters) * 100.0
            breaches.append(perc)
        return self.speeds, breaches

    def get_windresults_samples_perc_debris_damage(self):
        samples = {}
        for wind_speed in self.speeds:
            samples[wind_speed] = type(self).result_buckets[wind_speed][type(self).FLD_DEBRIS_AT]
        return self.speeds, samples

    def get_windresults_samples_nv(self):
        samples = {}
        for wind_speed in self.speeds:
            samples[wind_speed] = type(self).result_buckets[wind_speed][type(self).FLD_DEBRIS_NV_AT]
        return self.speeds, samples

    def get_windresults_samples_num_items(self):
        samples = {}
        for wind_speed in self.speeds:
            samples[wind_speed] = type(self).result_buckets[wind_speed][type(self).FLD_DEBRIS_NUM_AT]
        return self.speeds, samples

    def get_windresults_samples_perc_water_ingress(self):
        samples = {}
        for wind_speed in self.speeds:
            samples[wind_speed] = type(self).result_buckets[wind_speed][type(self).FLD_WI_AT]
        return self.speeds, samples

    def plot_connection_damage(self, vRed, vBlue):
        for ctg_name in ['sheeting', 'batten', 'rafter', 'piersgroup',
                         'wallracking', 'Truss']:
            if ctgMap.get(ctg_name, None) is None:
                continue
            vgrid = np.ones((self.s.house.roof_rows, self.s.house.roof_columns),
                            dtype=np.float32) * vBlue + 10.0
            for conn in connByTypeGroupMap[ctg_name]:
                gridCol, gridRow = \
                    zone.getGridFromZoneLoc(conn.location_zone.zone_name)
                if conn.result_failure_v > 0:
                    vgrid[gridRow][gridCol] = conn.result_failure_v
            output.plot_damage_show(ctg_name, vgrid, self.s.house.roof_columns,
                                    self.s.house.roof_rows, vRed, vBlue)

        if 'plot_wall_damage_show' in output.__dict__:
            wall_major_rows = 2
            wall_major_cols = self.s.house.roof_columns
            wall_minor_rows = 2
            wall_minor_cols = 8

            for ctg_name in ('wallcladding', 'wallcollapse'):
                if ctgMap.get(ctg_name, None) is None:
                    continue

                v_south_grid = np.ones((wall_major_rows, wall_major_cols),
                                       dtype=np.float32) * vBlue + 10.0
                v_north_grid = np.ones((wall_major_rows, wall_major_cols),
                                       dtype=np.float32) * vBlue + 10.0
                v_west_grid = np.ones((wall_minor_rows, wall_minor_cols),
                                      dtype=np.float32) * vBlue + 10.0
                v_east_grid = np.ones((wall_minor_rows, wall_minor_cols),
                                      dtype=np.float32) * vBlue + 10.0

                # construct south grid
                for gridCol in range(0, wall_major_cols):
                    for gridRow in range(0, wall_major_rows):
                        colChar = chr(ord('A')+gridCol)
                        loc = 'WS%s%d' % (colChar, gridRow+1)
                        conn = connByZoneTypeMap[loc].get(ctg_name)
                        if conn and conn.result_failure_v > 0:
                            v_south_grid[gridRow][gridCol] = \
                                conn.result_failure_v

                # construct north grid
                for gridCol in range(0, wall_major_cols):
                    for gridRow in range(0, wall_major_rows):
                        colChar = chr(ord('A')+gridCol)
                        loc = 'WN%s%d' % (colChar, gridRow+1)
                        conn = connByZoneTypeMap[loc].get(ctg_name)
                        if conn and conn.result_failure_v > 0:
                            v_north_grid[gridRow][gridCol] = \
                                conn.result_failure_v

                # construct west grid
                for gridCol in range(0, wall_minor_cols):
                    for gridRow in range(0, wall_minor_rows):
                        loc = 'WW%d%d' % (gridCol+2, gridRow+1)
                        conn = connByZoneTypeMap[loc].get(ctg_name)
                        if conn and conn.result_failure_v > 0:
                            v_west_grid[gridRow][gridCol] = \
                                conn.result_failure_v

                # construct east grid
                for gridCol in range(0, wall_minor_cols):
                    for gridRow in range(0, wall_minor_rows):
                        loc = 'WE%d%d' % (gridCol+2, gridRow+1)
                        conn = connByZoneTypeMap[loc].get(ctg_name)
                        if conn and conn.result_failure_v > 0:
                            v_east_grid[gridRow][gridCol] = \
                                conn.result_failure_v

                output.plot_wall_damage_show(
                    ctg_name,
                    v_south_grid, v_north_grid, v_west_grid, v_east_grid,
                    wall_major_cols, wall_major_rows,
                    wall_minor_cols, wall_minor_rows,
                    vRed, vBlue)

    # This needs to be done outside of the plotting function
    # as these coefficients are
    # the final output of this program in batch... they are all that matters.
    #
    def fit_fragility_curves(self):
        # unpack the fragility means into seperate arrays for fit/plot.
        self.frag_levels = []
        for thr in type(self).fragility_thresholds:
            self.frag_levels.append(np.zeros(len(self.speeds)))
        for i, wind_speed in enumerate(self.speeds):
            for frag_ind in range(len(self.frag_levels)):
                self.frag_levels[frag_ind][i] = \
                    type(self).result_buckets[wind_speed][type(self).FLD_FRAGILITIES][frag_ind]

        # fit curves and store results
        coeff_arr = []
        ss = 0
        for frag_ind in range(len(self.frag_levels)):
            try:
                coeff_arr, ss = curve_log.fit_curve(self.speeds,
                                                    self.frag_levels[frag_ind],
                                                    False)
                if frag_ind > 0:
                    self.file_frag.write(',')

                label = '%s(%.2f)' % (type(self).fragility_thresholds[frag_ind][0],
                                    type(self).fragility_thresholds[frag_ind][1])
                self.file_frag.write('%f,%f' % (coeff_arr[0], coeff_arr[1]))
                type(self).fragility_thresholds[frag_ind][3] = \
                    CurvePlot(coeff_arr,
                              'lognormal',
                              self.s.wind_speed_min,
                              self.s.wind_speed_max,
                              label,
                              type(self).fragility_thresholds[frag_ind][2])
            except Exception, e:
                print 'fit_fragility_curves failed to fit: coeff_arr: %s' % coeff_arr
                print e

        self.file_frag.write('\n')

    def plot_fragility(self, output_folder):
        for frag_ind in range(len(self.frag_levels)):
            output.plot_fragility_curve(self.speeds,
                                        self.frag_levels[frag_ind],
                                        '_nolegend_',
                                        0.3,
                                        type(self).fragility_thresholds[frag_ind][2])
            plot_obj = type(self).fragility_thresholds[frag_ind][3]
            if plot_obj:
                plot_obj.plot_frag()

        output.plot_fragility_show(self.s.num_iters,
                                   self.s.wind_speed_min,
                                   self.s.wind_speed_max, output_folder)

    def fit_vuln_curve(self):
        self.wind_speeds = np.zeros(len(self.speeds))
        self.di_means = np.zeros(len(self.speeds))

        ss = 0
        for i, wind_speed in enumerate(self.speeds):
            self.wind_speeds[i] = wind_speed
            self.di_means[i] = type(self).result_buckets[wind_speed][type(self).FLD_MEAN]

        if self.s.getOpt_VulnFitLog():
            self.A_final, self.ss = curve_log.fit_curve(self.wind_speeds,
                                                        self.di_means)
        else:
            self.A_final, self.ss = curve.fit_curve(self.wind_speeds,
                                                    self.di_means)

    def plot_vulnerability(self, output_folder, label="Fitted Curve"):
        # fit current observations
        self.fit_vuln_curve()

        # plot means
        if self.s.num_iters <= 100:
            for wind_speed in self.speeds:
                damage_indexes = type(self).result_buckets[wind_speed][type(self).FLD_DIARRAY]
                output.plot_wind_event_damage([wind_speed]*len(damage_indexes),
                                              damage_indexes)
        output.plot_wind_event_mean(self.wind_speeds, self.di_means)

        # plot fitted curve (with previous dimmed red)
        if self.s.getOpt_VulnFitLog():
            fn_form = 'lognormal'
        else:
            fn_form = 'original'

        cp = CurvePlot(self.A_final,
                       fn_form,
                       self.s.wind_speed_min,
                       self.s.wind_speed_max,
                       "Fitted Curve")

        cp.plot_vuln()
        if self.prevCurvePlot:
            self.prevCurvePlot.plot_vuln(True)
        self.prevCurvePlot = cp

    def show_results(self, output_folder=None, vRed=40, vBlue=80):
        if self.mplDict:
            self.mplDict['fragility'].axes.cla()
            self.mplDict['fragility'].axes.figure.canvas.draw()
            self.mplDict['vulnerability'].axes.cla()
            self.mplDict['vulnerability'].axes.figure.canvas.draw()
        if self.s.getOpt_DmgPlotFragility():
            self.plot_fragility(output_folder)
        if self.s.getOpt_DmgPlotVuln():
            self.plot_vulnerability(output_folder)
            output.plot_wind_event_show(self.s.num_iters,
                                        self.s.wind_speed_min,
                                        self.s.wind_speed_max,
                                        output_folder)
        self.plot_connection_damage(vRed, vBlue)

    def clear_connection_damage(self):
        v = np.ones((self.s.house.roof_rows, self.s.house.roof_columns),
                    dtype=np.float32) * self.s.wind_speed_max
        for ctname in ['sheeting', 'batten', 'rafter', 'piersgroup',
                       'wallracking']:
            output.plot_damage_show(ctname, v, self.s.house.roof_columns,
                                    self.s.house.roof_rows,
                                    self.s.wind_speed_min,
                                    self.s.wind_speed_max)
        for ctname in ['wallcladding', 'wallcollapse']:
            wall_major_rows = 2
            wall_major_cols = self.s.house.roof_columns
            wall_minor_rows = 2
            wall_minor_cols = 8
            v_major_grid = np.ones((wall_major_rows, wall_major_cols),
                                   dtype=np.float32) * self.s.wind_speed_max
            v_minor_grid = np.ones((wall_minor_rows, wall_minor_cols),
                                   dtype=np.float32) * self.s.wind_speed_max
            output.plot_wall_damage_show(
                ctname,
                v_major_grid, v_major_grid, v_minor_grid, v_minor_grid,
                wall_major_cols, wall_major_rows,
                wall_minor_cols, wall_minor_rows,
                self.s.wind_speed_min, self.s.wind_speed_max)

lastiPerc = -1

def simProgressCallback(V, di, percLoops):
    global lastiPerc
    iPerc = int(percLoops)
    if iPerc != lastiPerc:
        lastiPerc = iPerc
        sys.stdout.write('.')
    return True

# @profile
def simulate(s, options):
    if options.verbose:
        arg = simProgressCallback
    else:
        arg = None
    mySim = WindDamageSimulator(options, arg, None)
    mySim.set_scenario(s)
    runTime, hr = mySim.simulator_mainloop(options.verbose)
    if runTime:
        if options.plot_frag:
            mySim.plot_fragility(options.output_folder)
        if options.plot_vuln:
            mySim.plot_vulnerability(options.output_folder, None)
            output.plot_wind_event_show(mySim.s.num_iters,
                                        mySim.s.wind_speed_min,
                                        mySim.s.wind_speed_max,
                                        options.output_folder)
    return runTime

def main():
    USAGE = ('%prog -s <scenario_file> [-m <model database file>] '
             '[-o <output_folder>]')
    parser = OptionParser(usage=USAGE, version=VERSION_DESC)
    parser.add_option("-s", "--scenario",
                      dest="scenario_filename",
                      help="read scenario description from FILE",
                      metavar="FILE")
    parser.add_option("-m", "--model",
                      dest="model_database",
                      help="Use Model Database from FILE",
                      metavar="FILE")
    parser.add_option("-o", "--output",
                      dest="output_folder",
                      help="folder name to store simulation results",
                      metavar="FOLDER")
    parser.add_option("-v", "--verbose",
                      action="store_true",
                      dest="verbose",
                      default=False,
                      help="show verbose simulator output")
    parser.add_option("-i", "--import",
                      dest="data_folder",
                      help="data folder to import into model.db",
                      metavar="FOLDER")
    parser.add_option("--plot_vuln",
                      action="store_true",
                      dest="plot_vuln",
                      default=False,
                      help="show vulnerability plot")
    parser.add_option("--plot_frag",
                      action="store_true",
                      dest="plot_frag",
                      default=False,
                      help="show fragility plot")

    (options, args) = parser.parse_args()

    path_, _ = os.path.split(sys.argv[0])

    if options.model_database is None:
        model_db = None
    else:
        model_db = os.path.abspath(os.path.join(os.getcwd(),
                                                options.model_database))
    if options.output_folder is None:
        options.output_folder = os.path.abspath(os.path.join(path_,
                                                             './outputs'))
    else:
        options.output_folder = os.path.abspath(os.path.join(
            os.getcwd(), options.output_folder))
    print 'output directory: %s' % options.output_folder

    if options.verbose:
        logger.configure(logger.LOGGING_CONSOLE)
    else:
        logger.configure(logger.LOGGING_NONE)

    if options.data_folder:
        print ('Importing database from folder: {} '
               'to: {}').format(options.data_folder, options.model_database)

        database.configure(model_db, flag_make=True)
        dbimport.import_model(options.data_folder, options.model_database)
        database.db.close()
        return

    if options.scenario_filename:
        database.configure(model_db)
        s = scenario.loadFromCSV(options.scenario_filename)
        simulate(s, options)
        database.db.close()
    else:
        print '\nERROR: Must provide as scenario file to run simulator...\n'
        parser.print_help()

if __name__ == '__main__':
    main()