from data import COCODetection, get_label_map, MEANS, COLORS
from yolact.yolact import Yolact
from utils.augmentations import BaseTransform, FastBaseTransform, Resize
from utils.functions import MovingAverage, ProgressBar
from layers.box_utils import jaccard, center_size, mask_iou
from utils import timer
from utils.functions import SavePath
from layers.output_utils import postprocess, undo_image_transformation
import pycocotools
import random
from torch import nn
from data import cfg, set_cfg, set_dataset

import numpy as np
import torch
import torch.backends.cudnn as cudnn
from torch.autograd import Variable
import argparse
import time
import random
import cProfile
import pickle
import json
import os
from collections import defaultdict
from pathlib import Path
from collections import OrderedDict
from PIL import Image
import operator, functools
from dataclasses import dataclass
import re
import matplotlib.pyplot as plt
import cv2

import faiss

@dataclass
class ImageData:
    x: float
    y: float
    rotation_z: float
    rotation_w: float
    side: str

def get_image_data(image_file):
    regular_expression = "IMAGE_(-?\d+)_(-?\d+)_(-?\d+)_(-?\d+)_-serena\-camera\-(right|left)\-color\-image_raw_frame_\d+.png"
    result = re.search(regular_expression, image_file)
    if result:
        data = result.group(1).split()
        return ImageData(float(result.group(1).split()[0]) / 1000.0,
                        float(result.group(2).split()[0]) / 1000.0,
                        float(result.group(3).split()[0]) / 1000.0,
                        float(result.group(4).split()[0]) / 1000.0,
                        result.group(5).split())
    return None

class NearestNeighbourIndex():

    EPSILLON = 0.00001

    def __init__(self, output_folder,
            embedding_sizes=[
                [256, 69, 69],
                [256, 35, 35],
                [256, 18, 18],
                [256, 9, 9],
                [256, 5, 5]
            ],
            embedding_enabled=[ False, False, True, True, True ]):
        self.output_folder = output_folder
        self.embedding_sizes = embedding_sizes
        self.embedding_enabled = embedding_enabled
        self.embedding_flat_sizes = self.get_flat_sizes()
        self.number_layers = len([e for e in self.embedding_enabled if e])

    def create_indices(self, number_of_trees=20):
        indices = []
        for layer_index, embedding_size in enumerate(self.embedding_flat_sizes):
            if self.embedding_enabled[layer_index]:
                quantizer = faiss.IndexFlatL2(embedding_size)
                bytes_per_vector = 8
                bits_per_byte = 8
                nlist = 100
                index = faiss.IndexIVFPQ(quantizer, embedding_size, nlist, bytes_per_vector, bits_per_byte)
                index.cp.min_points_per_centroid = 5   # quiet warning
                index.quantizer_trains_alone = 2
                indices.append(index)
        return indices

    def create_pcas(self):
        pcas = []
        for layer_index, embedding_size in enumerate(self.embedding_flat_sizes):
            if self.embedding_enabled[layer_index]:
                pca = faiss.PCAMatrix (embedding_size, embedding_size // 4)
                pcas.append(index)
        return pcas

    def get_indices_file_paths(self):
        file_paths = []
        for layer_index, _ in enumerate(range(self.number_layers)):
            file_paths.append(f'{self.output_folder}/embeddings_{layer_index}.faiss')
        return file_paths

    def save_indices(self, indices):
        file_paths = self.get_indices_file_paths()
        for layer_index, index in enumerate(indices):
            faiss.write_index(index, file_paths[layer_index])

    def indices_created(self):
        file_paths =self.get_indices_file_paths()
        for path in file_paths:
            if not os.path.exists(path):
                return False
        return True

    def get_embeddings(self, indices, net:Yolact, path:str, sample_index_offset):
        frame = torch.from_numpy(cv2.imread(path)).cuda().float()
        batch = FastBaseTransform()(frame.unsqueeze(0))
        preds, embeddings = net(batch)
        embeddings = [embedding for layer_index, embedding in enumerate(embeddings)
            if self.embedding_enabled[layer_index]]
        return embeddings

    def normalize_embeddings(self, embeddings_batch_in_layer, means, stds):
        embeddings_batch_in_layer = embeddings_batch_in_layer.view(embeddings_batch_in_layer.size()[0],
            embeddings_batch_in_layer.size()[1],
            -1)
        embeddings_batch_in_layer = (embeddings_batch_in_layer - means) / (stds + NearestNeighbourIndex.EPSILLON)
        embeddings_batch_in_layer = embeddings_batch_in_layer.view(embeddings_batch_in_layer.size()[0], -1)
        return embeddings_batch_in_layer

    def add_to_index(self, embeddings_batch, indices, means, stds):
        embeddings_count = len(embeddings_batch)
        for layer_index, embeddings_batch_in_layer in enumerate(embeddings_batch):
            index = indices[layer_index]
            if len(means) < embeddings_count or len(stds) < embeddings_count:
                mean = StyleExtractor.mean(embeddings_batch_in_layer)
                std  = StyleExtractor.standard_deviation(embeddings_batch_in_layer, mean)
                embeddings_batch_in_layer = self.normalize_embeddings(embeddings_batch_in_layer, mean, std)
                embeddings_batch_in_layer = embeddings_batch_in_layer.cpu().numpy()
                index.train(embeddings_batch_in_layer)
                index.add(embeddings_batch_in_layer)
                means.append(mean)
                stds.append(std)
            else:
                embeddings_batch_in_layer = self.normalize_embeddings(embeddings_batch_in_layer,
                    means[layer_index], stds[layer_index])
                embeddings_batch_in_layer = embeddings_batch_in_layer.cpu().numpy()
                index.add(embeddings_batch_in_layer)
        return means, stds

    def create_content_embedding_indices(self, input_folder, indices, batch_size=1000):
        print('**** START EMBEDDING INDICES GENERATION ****')
        sample_index_offset = 0
        embeddings_batch = None
        means, stds = [], []
        with open(f'{self.output_folder}/images.txt', 'w') as images_file:
            for file_path in Path(input_folder).glob('*'):
                path = str(file_path)
                name = os.path.basename(path)
                image_data = get_image_data(name)
                print(f"{sample_index_offset} - Processing: {image_data}")
                images_file.write(f"{name}\n")
                name = '.'.join(name.split('.')[:-1]) + '.png'
                out_path = os.path.join(self.output_folder, name)
                embeddings = self.get_embeddings(indices, net, path, sample_index_offset)
                if embeddings_batch is not None:
                    embeddings_batch = [torch.cat([embedding_batch, embeddings[layer_index]], 0) \
                        for layer_index, embedding_batch in enumerate(embeddings_batch)]
                else:
                    embeddings_batch = embeddings
                if sample_index_offset > 0 and sample_index_offset % ( batch_size - 1 ) == 0:
                    means, stds = self.add_to_index(embeddings_batch, indices, means, stds)
                    embeddings_batch = None
                sample_index_offset = sample_index_offset + 1
            if embeddings_batch is not None:
                means, stds = self.add_to_index(embeddings_batch, indices, means, stds)
        self.save_indices(indices)
        print('**** END EMBEDDING INDICES GENERATION ****')
        return means, stds

    def get_flat_sizes(self):
        embedding_flat_sizes = []
        for layer_index, embedding_size in enumerate(self.embedding_sizes):
            embedding_flat_sizes.append(functools.reduce(operator.mul, embedding_size, 1))
        return embedding_flat_sizes

    def load_indices(self):
        indices = []
        file_paths = self.get_indices_file_paths()
        for layer_index, file_path in enumerate(file_paths):
            indices.append(faiss.read_index(file_path))
        return indices

    def load_statistics(self):
        indices = []
        means = torch.load(f"{self.output_folder}/means.torch")
        stds = torch.load(f"{self.output_folder}/stds.torch")
        return (means, stds)

    def load_image_files(self):
        with open(f'{self.output_folder}/images.txt', 'r') as images_file:
	        images = images_file.readlines()
	        images = [get_image_data(image) for image in images]
	        return images
        return []

    def index_images(self, net:Yolact, input_folder:str):
        if not os.path.exists(self.output_folder):
            os.mkdir(self.output_folder)
        if not self.indices_created():
            indices = self.create_indices(self.embedding_flat_sizes)
            means, stds = self.create_content_embedding_indices(input_folder, indices)
            torch.save(means, f"{self.output_folder}/means.torch")
            torch.save(stds,  f"{self.output_folder}/stds.torch")
        indices = self.load_indices()

class StyleExtractor(nn.Module):

    def __init__(self):
        super(StyleExtractor, self).__init__()

    def mean(embeddings):
        embeddings = embeddings.view(embeddings.size()[0], embeddings.size()[1], -1)
        means = embeddings.std(0)
        means = means.unsqueeze(0)
        return means

    def standard_deviation(embeddings, mean):
        embeddings = embeddings.view(embeddings.size()[0], embeddings.size()[1], -1)
        embeddings = embeddings - mean
        embeddings = embeddings.permute(0, 2, 1).contiguous()
        embeddings = embeddings.view(-1, embeddings.size()[2])
        stds = embeddings.std(0)
        stds = stds.unsqueeze(0).unsqueeze(2)
        return stds

    def style(embeddings, means, stds):
        embeddings = embeddings.view(embeddings.size()[0], embeddings.size()[1], -1)
        embeddings = embeddings - means
        embeddings = embeddings / ( stds + NearestNeighbourIndex.EPSILLON )
        style = torch.bmm(embeddings, embeddings.transpose(1, 2))
        style = style / embeddings.size(2)
        return style

from multiprocessing.pool import ThreadPool
from queue import Queue

def str2bool(v):
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

def parse_args(argv=None):
    parser = argparse.ArgumentParser(
        description='YOLACT COCO Evaluation')
    parser.add_argument('--trained_model',
                        default='weights/ssd300_mAP_77.43_v2.pth', type=str,
                        help='Trained state_dict file path to open. If "interrupt", this will open the interrupt file.')
    parser.add_argument('--cuda', default=True, type=str2bool,
                        help='Use cuda to evaulate model')
    parser.add_argument('--fast_nms', default=True, type=str2bool,
                        help='Whether to use a faster, but not entirely correct version of NMS.')
    parser.add_argument('--cross_class_nms', default=False, type=str2bool,
                        help='Whether compute NMS cross-class or per-class.')
    parser.add_argument('--shuffle', dest='shuffle', action='store_true',
                        help='Shuffles the images when displaying them. Doesn\'t have much of an effect when display is off though.')
    parser.add_argument('--ap_data_file', default='results/ap_data.pkl', type=str,
                        help='In quantitative mode, the file to save detections before calculating mAP.')
    parser.add_argument('--resume', dest='resume', action='store_true',
                        help='If display not set, this resumes mAP calculations from the ap_data_file.')
    parser.add_argument('--max_images', default=-1, type=int,
                        help='The maximum number of images from the dataset to consider. Use -1 for all.')
    parser.add_argument('--output_coco_json', dest='output_coco_json', action='store_true',
                        help='If display is not set, instead of processing IoU values, this just dumps detections into the coco json file.')
    parser.add_argument('--bbox_det_file', default='results/bbox_detections.json', type=str,
                        help='The output file for coco bbox results if --coco_results is set.')
    parser.add_argument('--mask_det_file', default='results/mask_detections.json', type=str,
                        help='The output file for coco mask results if --coco_results is set.')
    parser.add_argument('--config', default=None,
                        help='The config object to use.')
    parser.add_argument('--output_web_json', dest='output_web_json', action='store_true',
                        help='If display is not set, instead of processing IoU values, this dumps detections for usage with the detections viewer web thingy.')
    parser.add_argument('--web_det_path', default='web/dets/', type=str,
                        help='If output_web_json is set, this is the path to dump detections into.')
    parser.add_argument('--no_bar', dest='no_bar', action='store_true',
                        help='Do not output the status bar. This is useful for when piping to a file.')
    parser.add_argument('--display_lincomb', default=False, type=str2bool,
                        help='If the config uses lincomb masks, output a visualization of how those masks are created.')
    parser.add_argument('--benchmark', default=False, dest='benchmark', action='store_true',
                        help='Equivalent to running display mode but without displaying an image.')
    parser.add_argument('--no_sort', default=False, dest='no_sort', action='store_true',
                        help='Do not sort images by hashed image ID.')
    parser.add_argument('--seed', default=None, type=int,
                        help='The seed to pass into random.seed. Note: this is only really for the shuffle and does not (I think) affect cuda stuff.')
    parser.add_argument('--mask_proto_debug', default=False, dest='mask_proto_debug', action='store_true',
                        help='Outputs stuff for scripts/compute_mask.py.')
    parser.add_argument('--no_crop', default=False, dest='crop', action='store_false',
                        help='Do not crop output masks with the predicted bounding box.')
    parser.add_argument('--image', default=None, type=str,
                        help='A path to an image to use for display.')
    parser.add_argument('--images', default=None, type=str,
                        help='An input folder of images and output folder to save detected images. Should be in the format input->output.')
    parser.add_argument('--video', default=None, type=str,
                        help='A path to a video to evaluate on. Passing in a number will use that index webcam.')
    parser.add_argument('--video_multiframe', default=1, type=int,
                        help='The number of frames to evaluate in parallel to make videos play at higher fps.')
    parser.add_argument('--score_threshold', default=0, type=float,
                        help='Detections with a score under this threshold will not be considered. This currently only works in display mode.')
    parser.add_argument('--dataset', default=None, type=str,
                        help='If specified, override the dataset specified in the config with this one (example: coco2017_dataset).')
    parser.add_argument('--detect', default=False, dest='detect', action='store_true',
                        help='Don\'t evauluate the mask branch at all and only do object detection. This only works for --display and --benchmark.')
    parser.add_argument('--display_fps', default=False, dest='display_fps', action='store_true',
                        help='When displaying / saving video, draw the FPS on the frame')
    parser.add_argument('--emulate_playback', default=False, dest='emulate_playback', action='store_true',
                        help='When saving a video, emulate the framerate that you\'d get running in real-time mode.')

    parser.set_defaults(no_bar=False, display=False, resume=False, output_coco_json=False, output_web_json=False, shuffle=False,
                        benchmark=False, no_sort=False, no_hash=False, mask_proto_debug=False, crop=True, detect=False, display_fps=False,
                        emulate_playback=False)

    global args
    args = parser.parse_args(argv)

    if args.output_web_json:
        args.output_coco_json = True

    if args.seed is not None:
        random.seed(args.seed)

def evaluate(net:Yolact, dataset, train_mode=False):
    net.detect.use_fast_nms = args.fast_nms
    net.detect.use_cross_class_nms = args.cross_class_nms
    cfg.mask_proto_debug = args.mask_proto_debug
    nn_index = NearestNeighbourIndex('indices')
    if args.images is not None:
        nn_index.index_images(net, args.images)
        return
if __name__ == '__main__':
    parse_args()

    if args.config is not None:
        set_cfg(args.config)

    if args.trained_model == 'interrupt':
        args.trained_model = SavePath.get_interrupt('weights/')
    elif args.trained_model == 'latest':
        args.trained_model = SavePath.get_latest('weights/', cfg.name)

    if args.config is None:
        model_path = SavePath.from_str(args.trained_model)
        # TODO: Bad practice? Probably want to do a name lookup instead.
        args.config = model_path.model_name + '_config'
        print('Config not specified. Parsed %s from the file name.\n' % args.config)
        set_cfg(args.config)

    if args.detect:
        cfg.eval_mask_branch = False

    with torch.no_grad():
        if not os.path.exists('results'):
            os.makedirs('results')

        if args.cuda:
            cudnn.fastest = True
            torch.set_default_tensor_type('torch.cuda.FloatTensor')
        else:
            torch.set_default_tensor_type('torch.FloatTensor')

        dataset = None

        print('Loading model...', end='')
        net = Yolact()
        net.load_weights(args.trained_model)
        net.eval()
        print(' Done.')

        if args.cuda:
            net = net.cuda()

        evaluate(net, dataset)


import torch
from torch.autograd import Variable
import torch.nn as nn
from torchvision.transforms import ToPILImage
from torchvision.utils import save_image
import cv2
from PIL import Image
from PIL import ImageDraw
import numpy as np
from matplotlib import pyplot as plt

import os

# from net import resnet50
import torchvision.transforms as transforms

VOC_CLASSES = (    # always index 0
    'ok', 'ng', '', '',
    '', '', '', '', '',
    '', '', '', '',
    '', '', '',
'', '', '', '')

Color = [[0, 0, 0],
        [128, 0, 0],
        [0, 128, 0],
        [128, 128, 0],
        [0, 0, 128],
        [128, 0, 128],
        [0, 128, 128],
        [128, 128, 128],
        [64, 0, 0],
        [192, 0, 0],
        [64, 128, 0],
        [192, 128, 0],
        [64, 0, 128],
        [192, 0, 128],
        [64, 128, 128],
        [192, 128, 128],
        [0, 64, 0],
        [128, 64, 0],
        [0, 192, 0],
        [128, 192, 0],
        [0, 64, 128]]

def decoder(pred, output_size):
    '''
    pred (tensor) 1x7x7x30
    return (tensor) box[[x1,y1,x2,y2]] label[...]
    '''
    # grid_num =
    boxes=[]
    cls_indexs=[]
    probs = []
    # cell_size = 1./grid_num
    pred = pred.data
    pred = pred.squeeze(0) #7x7x30
    contain = pred[:,:,4].unsqueeze(2)
    # contain2 = pred[:,:,9].unsqueeze(2)
    # contain = torch.cat((contain1,contain2),2)
    mask1 = contain > 0.1 #大于阈值
    mask2 = (contain==contain.max()) #we always select the best contain_prob what ever it>0.9
    mask = (mask1+mask2).gt(0)
    # min_score,min_index = torch.min(contain,2) #每个cell只选最大概率的那个预测框
    for i in range(output_size[1]):
        for j in range(output_size[0]):
            for b in range(1):
                # index = min_index[i,j]
                # mask[i,j,index] = 0
                if mask[i,j,b] == 1:
                    #print(i,j,b)
                    box = pred[i,j,b*5:b*5+4]
                    contain_prob = torch.FloatTensor([pred[i,j,b*5+4]])
                    xy = torch.FloatTensor([j,i])/torch.Tensor(output_size) #cell左上角  up left of cell
                    box[:2] = box[:2]/torch.Tensor(output_size) + xy # return cxcy relative to image
                    box_xy = torch.FloatTensor(box.size())#转换成xy形式    convert[cx,cy,w,h] to [x1,xy1,x2,y2]
                    box_xy[:2] = box[:2] - 0.5*box[2:]
                    box_xy[2:] = box[:2] + 0.5*box[2:]
                    max_prob,cls_index = torch.max(pred[i,j,4],0)
                    if float((contain_prob*max_prob)[0]) > 0.1:
                        boxes.append(box_xy.view(1,4))
                        cls_indexs.append(cls_index)
                        probs.append(contain_prob*max_prob)
    _cls_indexs = torch.Tensor(len(cls_indexs))
    for i in range(len(cls_indexs)):
        _cls_indexs[i] = cls_indexs[i]
    cls_indexs = _cls_indexs
    if len(boxes) ==0:
        boxes = torch.zeros((1,4))
        probs = torch.zeros(1)
        cls_indexs = torch.zeros(1)
    else:
        boxes = torch.cat(boxes,0) #(n,4)
        probs = torch.cat(probs,0) #(n,)
        # cls_indexs = torch.cat(cls_indexs,torch.tensor([0])) #(n,)
    keep = nms(boxes,probs)
    # return boxes[keep],cls_indexs[keep],probs[keep]
    return boxes,cls_indexs,probs

def nms(bboxes,scores,threshold=0.9):
    '''
    bboxes(tensor) [N,4]
    scores(tensor) [N,]
    '''
    x1 = bboxes[:,0]
    y1 = bboxes[:,1]
    x2 = bboxes[:,2]
    y2 = bboxes[:,3]
    areas = (x2-x1) * (y2-y1)

    _,order = scores.sort(0,descending=True)
    keep = []
    while order.numel() > 0:
        if len(order.size()) == 0:
            break
        i = order[0].item()
        keep.append(i)

        if order.numel() == 1:
            break

        xx1 = x1[order[1:]].clamp(min=x1[i])
        yy1 = y1[order[1:]].clamp(min=y1[i])
        xx2 = x2[order[1:]].clamp(max=x2[i])
        yy2 = y2[order[1:]].clamp(max=y2[i])

        w = (xx2-xx1).clamp(min=0)
        h = (yy2-yy1).clamp(min=0)
        inter = w*h

        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        ids = (ovr<=threshold).nonzero().squeeze()
        if ids.numel() == 0:
            break
        order = order[ids+1]
    return torch.LongTensor(keep)

def show_bb(img, x, y, w, h, text, textcolor, bbcolor):
    draw = ImageDraw.Draw(img)
    text_w, text_h = draw.textsize(text)
    label_y = y if y <= text_h else y - text_h
    draw.rectangle((x, label_y, x+w, label_y+h), outline=bbcolor)
    draw.rectangle((x, label_y, x+text_w, label_y+text_h), outline=bbcolor, fill=bbcolor)
    draw.text((x, label_y), text, fill=textcolor)

def draw_bb(img, pred, text, bbox, label, output_size):
    """get pil image and return bbox pil image

    args
    img: pil image
    pred: bbox, label tensor
    text: text color
    bbox: bbox color

    return
    img: pil image bboxed image
    """
    pred = pred.cpu()
    pred_boxes, pred_cls_indexs, probs = decoder(pred, output_size)

    h,w = img.size
    # print("h:{},w:{}".format(h,w))
    result = []
    box_len = len(pred_boxes)
    for i in range(box_len):
        box = pred_boxes[i]
        x1 = int(box[0]*h)
        x2 = int(box[2]*h)
        y1 = int(box[1]*w)
        y2 = int(box[3]*w)
        cls_index = pred_cls_indexs[i]
        cls_index = int(cls_index) # convert LongTensor to int
        prob = probs[i]
        if prob <= 0.7:
            continue
        if (x2-x1) * (y2-y1) > h*w*0.2:
            continue
        result.append([(x1,y1),(x2,y2),prob, cls_index])

    for (x_min, y_min), (x_max, y_max), prob, cls_index in result:
        w = x_max - x_min; h = y_max - y_min
        textcolor = (255,255,255)
        # label = class_name+str(round(prob.item(),2))
        label = VOC_CLASSES[cls_index]+str(round(prob.item(), 2))
        bbcolor = Color[cls_index]
        bbcolor = tuple(bbcolor)
        show_bb(img, x_min, y_min, w, h, label, text, bbcolor)
    return img
    # img.save(os.path.join(root, name))

def save_res_im(img, pred, ans, name, root, output_size, draw_ans=True):
    """draw pred bboxes and ans bboxes and save to root/name
    args
    img: torch tensor
    pred: torch tensor model output
    ans: torch tensor target
    name: file name you want to save im
    root: save dir
    """
    save_image(img,"/tmp/tmp.jpg")
    img = Image.open("/tmp/tmp.jpg")

    # pred bboxes
    img = draw_bb(img, pred, (255,255,255),(255,0,0), "pred", output_size)
    # ans bboxes
    if draw_ans:
        img = draw_bb(img, ans, (255,255,255), (0,255,0), "ans", output_size)

    img.save(os.path.join(root, name))


import numpy as np
import socket, struct
from ismore import settings, udp_feedback_client
import time
from ismore import common_state_lists, ismore_bmi_lib
import pandas as pd
import pickle
import os

class Patient(object):

    def __init__(self, targets_matrix_file):
        self.addrs = [settings.ARMASSIST_UDP_SERVER_ADDR, settings.REHAND_UDP_SERVER_ADDR]

        self.socks = [socket.socket(socket.AF_INET, socket.SOCK_DGRAM), socket.socket(socket.AF_INET, socket.SOCK_DGRAM)]
        self.n_dofs = [range(3), range(3, 7)]
        self.plant_types = ['ArmAssist', 'ReHand']

        self.aa_p = range(3) #common_state_lists.aa_pos_states
        self.rh_p = range(4) #common_state_lists.rh_pos_states
        self.rh_v = range(4, 8) #common_state_lists.rh_vel_states
        self.aa_v = range(3, 6)

        #self.aa = udp_feedback_client.ArmAssistData()
        #self.rh = udp_feedback_client.ReHandData()
        #self.aa.start()
        #self.last_aa_pos = pd.Series(np.zeros((3, )), dtype=self.aa_p) #self.aa.get()['data'][self.aa_p]
        #self.last_aa_pos_t = time.time()

        #self.rh.start()

        assister_kwargs = {
            'call_rate': 20,
            'xy_cutoff': 5,
        }
        self.assister = ismore_bmi_lib.ASSISTER_CLS_DICT['IsMore'](**assister_kwargs)
        self.targets_matrix = pickle.load(open(targets_matrix_file))

    def send_vel(self, vel):
        for i, (ia, sock, ndof, plant) in enumerate(zip(self.addrs, self.socks, self.n_dofs, self.plant_types)):
            self._send_command('SetSpeed %s %s\r' % (plant, self.pack_vel(vel[ndof], ndof)), ia, sock)

    def pack_vel(self, vel, n_dof):
        format_str = "%f " * len(n_dof)
        return format_str % tuple(vel)

    def _send_command(self, command, addr, sock):
        sock.sendto(command, addr)

    def _get_current_state(self):
        #aa_data = self.aa.get()['data']

        with open(os.path.expandvars('$HOME/code/bmi3d/log/armassist.txt'), 'r') as f:
            lines = f.read().splitlines()
            last_line = lines[-2]
        aa_data = np.array([float(i) for i in last_line.split(',')])

        with open(os.path.expandvars('$HOME/code/bmi3d/log/rehand.txt'), 'r') as f:
            lines = f.read().splitlines()
            last_line = lines[-2]
        rh_data = np.array([float(i) for i in last_line.split(',')])

        #daa = np.array([aa_data[0][i] - self.last_aa_pos[0][i] for i in range(3)])
        #aa_vel = daa/(time.time() - self.last_aa_pos_t)
        #self.last_aa_pos = aa_data[self.aa_p]
        #rh_data = self.rh.get()['data']

        pos = np.hstack(( aa_data[self.aa_p], rh_data[self.rh_p] ))
        vel = np.hstack(( aa_data[self.aa_v], rh_data[self.rh_v] ))
        return np.hstack((pos, vel))

    def get_to_target(self, target_pos):
        current_state = np.mat(self._get_current_state()).T
        target_state = np.mat(np.hstack((target_pos, np.zeros((7, ))))).T
        assist_kwargs = self.assister(current_state, target_state, 1., mode=None)
        self.send_vel(10*np.squeeze(np.array(assist_kwargs['Bu'][7:14])))
        return np.sum((np.array(current_state)-np.array(target_state))**2)

    def go_to_target(self, target_name, tix=0):
        if len(self.targets_matrix[target_name].shape) > 1:
            targ = self.targets_matrix[target_name][tix]
        else:
            targ = self.targets_matrix[target_name]
        d = 100
        while d > 20:
            d = self.get_to_target(targ)
            print d


import igraph
import itertools
import functools
import numpy as np
import os
import pyinter
import sys
from bisect import bisect_left
from operator import attrgetter

from arcsv.breakpoint_merge import Breakpoint
from arcsv.helper import normcdf, not_primary, load_genome_gaps
from arcsv.unif_ci import uniform_ci

class DiscordantPair:
    def __init__(self, chrom, pos1, pos2, insert, qname):
        self.chrom = chrom
        self.pos1 = pos1
        self.pos2 = pos2
        self.insert = insert
        self.qname = qname

    def __repr__(self):
        return str((self.chrom, self.pos1, self.pos2, self.insert))

    def __hash__(self):
        return hash((self.chrom, self.pos1, self.pos2, self.insert, self.qname))

    def __lt__(self, other):
        return (self.pos1, self.pos2) < (other.pos1, other.pos2)

    def __gt__(self, other):
        return (self.pos1, self.pos2) > (other.pos1, other.pos2)

def process_discordant_pair(aln1, aln2, chrom, discordant_pairs, min_mapq, ilen,
                            min_insert, max_insert, is_rf=False):
    if (aln1.is_reverse != aln2.is_reverse) and (ilen is not None) and \
       (ilen >= min_insert) and (ilen <= max_insert):
        return None
    if aln1.mapq < min_mapq or aln2.mapq < min_mapq or aln1.is_unmapped or \
       aln2.is_unmapped or not_primary(aln1) or not_primary(aln2):
        return None
    # "First" is -> if FR (->  <-) and <- if RF (<-  ->)
    # i.e. the read we expect on the "left" in ref. coords
    if aln1.is_reverse != aln2.is_reverse:
        second = aln1 if (aln1.is_reverse ^ is_rf) else aln2
        first = aln1 if second is aln2 else aln2
        if ilen > max_insert and (first.reference_end < second.reference_start):
            # second check is needed for libraries with huge number of readthroughs,
            # e.g. with 2 x 150 reads and insert size cutoff < 300
            # TODO add conditions on other discordant read types?
            dtype = 'Del'
            disc = DiscordantPair(chrom, first.reference_end, second.reference_start,
                                  ilen, first.qname)
        elif (first.reference_start > second.reference_start) or \
             (first.reference_end > second.reference_end):
            dtype = 'Dup'
            disc = DiscordantPair(chrom, second.reference_start, first.reference_end,
                                  ilen, second.qname)
        elif ilen < min_insert:
            dtype = 'Ins'
            disc = DiscordantPair(chrom, first.reference_end, second.reference_start,
                                  ilen, first.qname)
    else:
        dtype = 'InvR' if (aln1.is_reverse ^ is_rf) else 'InvL'
        if dtype == 'InvL':
            pos1, pos2 = sorted([aln1.reference_end, aln2.reference_end])
        else:
            pos1, pos2 = sorted([aln1.reference_start, aln2.reference_start])
        disc = DiscordantPair(chrom, pos1, pos2, ilen, aln1.qname)
    discordant_pairs[dtype] = discordant_pairs.get(dtype, []) + [disc]
    if disc.pos1 > disc.pos2 and dtype != 'Ins':
        raise Warning('[process_disc_pair] discordant type {0} pos1 > pos2'.format(dtype))
    return dtype

# for each library in discordant_pairs_list, determine SV-specific
# cutoffs and cluster the discordant PE reads
# LATER allow to combine libraries of the same "type"
def apply_discordant_clustering(opts, discordant_pairs_list,
                                insert_mu, insert_sigma,
                                insert_min, insert_max, gap_file,
                                lr_cond=False, bp_confidence_level=0.95):
    nlib = opts['nlib']
    breakpoints = []
    for i in range(nlib):
        lib_name = opts['library_names'][i]
        for (dtype, pairs) in discordant_pairs_list[i].items():
            if opts['verbosity'] > 0:
                print('[pecluster] clustering {0}'.format(dtype))
            clusters, pairs_clustered = cluster_pairs(opts, pairs, dtype, i,
                                                      insert_mu[i], insert_sigma[i])

            if opts['verbosity'] > 0:
                print('[pecluster] computing null distribution for {0} clusters'.format(dtype))
            lr_null_clusters = compute_null_dist(opts, pairs_clustered, dtype,
                                                 insert_mu[i], insert_sigma[i],
                                                 gap_file, lib_idx=i, lr_cond=lr_cond)

            insert_cutoff = insert_max[i] if dtype == 'Del' else insert_min[i]
            fdc_out = fdr_discordant_clusters(opts, clusters, lr_null_clusters, dtype,
                                              insert_mu[i], insert_sigma[i],
                                              insert_cutoff, lr_cond)
            clusters_pass, clusters_fail, lr_pairs, first_reject = fdc_out
            for cl in clusters_pass:
                # print('passing cluster:')
                # print(cl)
                breakpoints.extend(cluster_to_bp(cl, bp_confidence_level, dtype, lib_name))
                # print(breakpoints[-2])
                # print(breakpoints[-1])
                # print('')
            if opts['verbosity'] > 0:
                print('[pecluster] {0}: {1} discordant {2} reads'
                      .format(opts['library_names'][i], len(pairs), dtype))
                print('[pecluster] {0}: {1} clusters, {2} passing {3} failing'
                      .format(opts['library_names'][i], dtype,
                              len(clusters_pass), len(clusters_fail)))
            outname = '{0}_{1}_cluster.txt'.format(lib_name, dtype)
            fname = os.path.join(opts['outdir'], 'logging', outname)
            write_clustering_results(fname, lr_pairs, first_reject)

    return breakpoints

def is_deldupinv_compatible(opts, pairs, max_distance,
                            insert_mu=None, insert_sigma=None, adjust=None):
    # close and intersecting?
    min_pos1 = min(p.pos1 for p in pairs)
    max_pos1 = max(p.pos1 for p in pairs)
    min_pos2 = min(p.pos2 for p in pairs)
    max_pos2 = max(p.pos2 for p in pairs)
    return max(max_pos1 - min_pos1, max_pos2 - min_pos2) <= max_distance \
        and max_pos1 < min_pos2

def is_ins_compatible(opts, pairs, max_distance, insert_mu, insert_sigma):
    # if adjust:
    #     est_insertion_size = insert_mu - (pair1.insert + pair2.insert)/2
    #     overlap = max(0, max(pair1.pos1, pair2.pos1) - min(pair1.pos2, pair2.pos2))
    #     adjustment = max(0, est_insertion_size - overlap - 3/sqrt(2)*insert_sigma)
    #     # TODO use this
    #     adjusted_max_distance = max_distance - adjustment
    #     # if adjustment > 0:
    #     #     print('ins_compatible adjust mu={0} est={1} adjusted={2}'
    #     #           .format(insert_mu, est_insertion_size, adjusted_max_distance))
    # else:
    #     adjusted_max_distance = max_distance  # INSERTIONS THESIS
    min_pos1 = min(p.pos1 for p in pairs)
    max_pos1 = max(p.pos1 for p in pairs)
    min_pos2 = min(p.pos2 for p in pairs)
    max_pos2 = max(p.pos2 for p in pairs)

    is_close = max(max_pos1 - min_pos1, max_pos2 - min_pos2) <= max_distance
    overlap = max(0, max_pos1 - min_pos2)

    # no intersection requirement because possible homologous flanking sequences
    return is_close and overlap <= opts['max_ins_pair_slop']

def is_ins_cluster_compatible(opts, cluster):
    overlap = max(0, max([p.pos1 for p in cluster]) - min([p.pos2 for p in cluster]))
    # LATER do we need this or is it guaranteed?
    return overlap <= opts['max_ins_cluster_slop']

compatibility_fun = {'Del': is_deldupinv_compatible,
                     'Ins': is_ins_compatible,
                     'Dup': is_deldupinv_compatible,
                     'InvR': is_deldupinv_compatible,
                     'InvL': is_deldupinv_compatible}

def cluster_pairs(opts, pairs, dtype, lib_idx, insert_mu, insert_sigma):
    pairs.sort(key=attrgetter('pos1'))
    max_compatible_distance = insert_mu + opts['cluster_max_distance_sd'] * insert_sigma
    max_cluster_size = opts['max_pecluster_size'][lib_idx]
    is_compatible = functools.partial(compatibility_fun[dtype],
                                      opts=opts,
                                      max_distance=max_compatible_distance,
                                      insert_mu=insert_mu, insert_sigma=insert_sigma)
    if opts['verbosity'] > 1:
        print('clustering {0} pairs'.format(dtype))
        print('max cluster size: {0}'.format(max_cluster_size))

    cur_comps = []              # pairs in the current connected components
    cur_maxpos = []             # max(pair.pos1) over pairs in cur_comps
    clusters = []               # list of components
    pairs_clustered = []
    for pair in pairs:
        # check for components we've moved past
        passed_comps = [i for i in range(len(cur_comps)) if
                        abs(cur_maxpos[i] - pair.pos1) > max_compatible_distance]
        if passed_comps != sorted(passed_comps):
            raise Warning('passed_comps not sorted? {0}'.format(passed_comps))
        offset = 0
        for i in passed_comps:
            idx = i - offset    # adjust for deleting other stuff
            # print('passed component with maxpos %d' % cur_maxpos[idx])

            if len(cur_comps[idx]) > max_cluster_size and \
               is_compatible(pairs=cur_comps[idx]):
                # for huge clusters, we just skip unless all discordant pairs
                #   are mutually compatible
                new_clusters = [cur_comps[idx]]
            elif 1 < len(cur_comps[idx]) <= max_cluster_size:
                new_clusters = cluster_handle_component(cur_comps[idx], is_compatible)
            else:
                new_clusters = []
            clusters.extend(new_clusters)
            # add to list of pairs to use for null distribution simulation,
            #   unless this is a huge cluster we skipped
            if len(new_clusters) > 0 or len(cur_comps[idx]) == 1:
                pairs_clustered.extend(cur_comps[idx])

            del cur_comps[idx]
            del cur_maxpos[idx]
            offset += 1
        # check whether pair is connected to existing components
        adjacent_comps = [i for i in range(len(cur_comps))
                          if any(is_compatible(pairs=(pair, p))
                                 for p in reversed(cur_comps[i]))]
        # add pair to existing component, else make new component
        if len(adjacent_comps) > 0:
            cur_comps[adjacent_comps[0]].append(pair)
            cur_maxpos[adjacent_comps[0]] = max(cur_maxpos[adjacent_comps[0]], pair.pos1)
        else:
            cur_comps.append([pair])
            cur_maxpos.append(pair.pos1)
        # and merge adjacent comps if necessary
        if len(adjacent_comps) > 1:
            merged = list(itertools.chain(*(cur_comps[i] for i in adjacent_comps)))
            merged_maxpos = max(cur_maxpos[i] for i in adjacent_comps)
            offset = 0
            for i in adjacent_comps:
                del cur_comps[i - offset]
                del cur_maxpos[i - offset]
                offset += 1
            cur_comps.append(merged)
            cur_maxpos.append(merged_maxpos)
    # handle remaining components
    for comp in cur_comps:
        if len(comp) > max_cluster_size and \
           is_compatible(pairs=comp):
            new_clusters = [comp]
        elif 1 < len(comp) <= max_cluster_size:
            new_clusters = cluster_handle_component(comp, is_compatible)
        else:
            new_clusters = []
        clusters.extend(new_clusters)
        if len(new_clusters) > 0 or len(comp) == 1:
            pairs_clustered.extend(comp)

    # insertions need some additional filtering
    if dtype == 'Ins':
        clusters = [c for c in clusters if is_ins_cluster_compatible(opts, c)]
    return clusters, pairs_clustered

# component: guaranteed length >= 2
def cluster_handle_component(component, is_compatible):
    # print('handling component:\n\t%s' % component)
    # create a graph and fill in all the edges
    g = igraph.Graph(len(component))
    g.vs['pairs'] = component
    iter_pairs = range(len(component))
    compatible_pairs = [(i, j) for (i, j) in itertools.product(iter_pairs, repeat=2) if
                        i != j and is_compatible(pairs=(component[i], component[j]))]
    for cp in compatible_pairs:
        g.add_edge(*cp)
    # get max cliques and add to result
    result = []
    for clique in g.largest_cliques():
        # print('\tclique of size %d' % len(clique))
        result.append(g.vs[clique]['pairs'])
    return result

# discordant_pairs: list of tuples corresponding to discordant pairs
# non_gaps: list of intervals where we can place the discordant pairs
def shuffle_discordant_pairs(discordant_pairs, chrom_len_no_gaps, max_insert_size=np.Inf):
    shuffled = []
    for pair in discordant_pairs:
        if max_insert_size < np.Inf and pair.insert > max_insert_size:
            continue
        pair_len = pair.pos2 - pair.pos1
        # ignoring read length, but doesn't matter for chrom_len >> read_len
        if pair_len < chrom_len_no_gaps and pair_len > -chrom_len_no_gaps:
            new_pos1 = np.random.random_integers(max(0, -pair_len),
                                                 chrom_len_no_gaps - max(0, pair_len))
            new_pair = DiscordantPair(pair.chrom, new_pos1, new_pos1 + pair_len,
                                      pair.insert, pair.qname)
            shuffled.append(new_pair)
        else:
            continue
    return shuffled

# LATER proper mle with empirical dist
# LATER LR including proportion of discordant reads in this area
# cluster: list of mutually compatible del-type discordants
def lr_del(cluster, insert_mu, insert_sigma, cutoff, conditioning=False):
    max_size = min([p.pos2 for p in cluster]) - max([p.pos1 for p in cluster])
    del_size_mle = min(max_size, np.mean([pair.insert for pair in cluster]) - insert_mu)

    lr = sum([(p.insert - insert_mu)**2 for p in cluster]) \
        - sum([(p.insert - insert_mu - del_size_mle)**2 for p in cluster])
    lr = lr / (2 * insert_sigma**2)

    if lr == -np.inf:
        sys.stderr.write('[lr_del] DEL likelihood ratio = -inf\n')
        sys.stderr.write('{0}\n'.format(cluster))

    if conditioning:
        n = len(cluster)
        lr += n * (np.log(1 - normcdf(cutoff, insert_mu, insert_sigma))
                   - np.log(1 - normcdf(cutoff - del_size_mle, insert_mu, insert_sigma)))

    return lr

def lr_ins(cluster, insert_mu, insert_sigma, cutoff, conditioning=False):
    # -->   <--
    #      -->   <--
    # overlap like this is possible with homologous sequences flanking insertion
    overlap = max(0, max([p.pos1 for p in cluster]) - min([p.pos2 for p in cluster]))
    ins_size_mle = max(overlap, insert_mu - np.mean([pair.insert for pair in cluster]))
    # if overlap > 0:
    #     print('insertion cluster w/overlap: {0}'.format(cluster))
    #     print('overlap {0} mle {1}\n'.format(overlap, ins_size_mle))

    lr = sum([(p.insert - insert_mu)**2 for p in cluster]) \
        - sum([(p.insert - insert_mu + ins_size_mle)**2 for p in cluster])
    lr = lr / (2 * insert_sigma**2)

    if conditioning:
        n = len(cluster)
        lr += n * (np.log(normcdf(cutoff, insert_mu, insert_sigma))
                   - np.log(normcdf(cutoff + ins_size_mle, insert_mu, insert_sigma)))

    return lr

# no cutoff or conditioning
def lr_inv(cluster, insert_mu, insert_sigma, cutoff=None, conditioning=None):
    return len(cluster)

# no cutoff or conditioning
def lr_dup(cluster, insert_mu, insert_sigma, cutoff=None, conditioning=None):
    return len(cluster)

lr_fun = {'Del': lr_del, 'Ins': lr_ins, 'InvL': lr_inv, 'InvR': lr_inv, 'Dup': lr_dup}

# does opts['pecluster_null_reps'] replicates of null cluster simulation and returns
# a sorted list of the null cluster likelihood ratios
# returns sorted list of null likelihood ratios under a permutation simulation
def compute_null_dist(opts, discordant_pairs, dtype,
                      insert_mu, insert_sigma,
                      gap_file, lib_idx, lr_cond):
    nreps = opts['pecluster_null_reps']
    chrom_name, start, end = opts['chromosome'], opts['region_start'], opts['region_end']
    gaps_inter = load_genome_gaps(gap_file, chrom_name)
    chrom_inter = pyinter.IntervalSet()
    chrom_inter.add(pyinter.closedopen(start, end))
    non_gaps_inter = chrom_inter.difference(gaps_inter)
    non_gaps = [(i.lower_value, i.upper_value) for i in non_gaps_inter]
    total_len = sum([i[1] - i[0] for i in non_gaps])

    # For deletion null clusters, don't use pairs that are obviously too large.
    # (for normal data the discordant read cutoff for deletion supports
    #  is like mu + 3 sigma ~ mu + .3mu, and we're excluding stuff bigger than 3mu)
    if dtype == 'Del':
        max_null_insert = insert_mu * opts['insert_max_mu_multiple']
    else:
        max_null_insert = np.Inf

    null_clusters = []
    lr_null_clusters = np.array([], float)
    for _ in range(nreps):
        shuffled = shuffle_discordant_pairs(discordant_pairs, total_len,
                                            max_insert_size=max_null_insert)
        clusters_tmp, _ = cluster_pairs(opts, shuffled, dtype, lib_idx,
                                        insert_mu, insert_sigma)
        null_clusters.extend(clusters_tmp)
        lr_tmp = np.fromiter((lr_fun[dtype](c, insert_mu, insert_sigma,
                                            opts['insert_cutoff'], lr_cond)
                              for c in clusters_tmp),
                             float)
        lr_null_clusters = np.append(lr_null_clusters, lr_tmp)
    if opts['verbosity'] > 1:
        print('[compute_null_dist] {0}'.format(dtype))
        print('shuffled lr:')
        print(lr_null_clusters)
        print('')

    outname = ('{0}_{1}_null_cluster_{2}reps.txt'
               .format(opts['library_names'][lib_idx], dtype, nreps))
    fname = os.path.join(opts['outdir'], 'logging', outname)
    write_clustering_results(fname, list(zip(lr_null_clusters, null_clusters)), first_reject=0)

    # print('there were {0} {1} clusters after shuffling'.format(len(clusters),
    #                                                            dtype))

    lr_null_clusters.sort()
    return lr_null_clusters

def fdr_discordant_clusters(opts, clusters, lr_null_clusters, dtype,
                            insert_mu, insert_sigma, insert_cutoff, lr_cond):
    # SPEEDUP this should be computed already
    lr_clusters = [lr_fun[dtype](c, insert_mu, insert_sigma, insert_cutoff, lr_cond)
                   for c in clusters]
    lr_pairs = list(zip(lr_clusters, clusters))
    lr_pairs.sort()

    # find FDR cutoff
    target_fdr = opts['pecluster_target_fdr']
    null_sim_nreps = opts['pecluster_null_reps']
    M = len(lr_clusters)
    M_0 = len(lr_null_clusters) / null_sim_nreps

    # P(null score >= x)
    def p_null_exceeds(x):
        return 1 - bisect_left(lr_null_clusters, x) / len(lr_null_clusters)

    if M > 0:
        pi_0 = min(1, M_0 / M)
    else:
        pi_0 = None

    if M_0 == 0:                # no nulls, reject everything
        j = 0
    else:
        for j in range(M + 1):
            if j == M:              # failed, can't reject anything
                break
            # estimated fdr from rejecting lr_pairs[j:]
            num_rej = M - j
            s_j = lr_pairs[j][0]
            est_fdr = M * pi_0 * p_null_exceeds(s_j) / num_rej
            if est_fdr <= target_fdr:
                break

    clusters_pass = [lr_pairs[i][1] for i in range(j, M)]
    clusters_fail = [lr_pairs[i][1] for i in range(0, j)]

    return clusters_pass, clusters_fail, lr_pairs, j

def write_clustering_results(filename, lr_pairs, first_reject):
    fout = open(filename, 'w')
    if len(lr_pairs) > 0:
        lr_clusters, clusters = list(zip(*lr_pairs))
        for i in range(len(clusters)):
            npairs = len(clusters[i])
            pos1 = [p.pos1 for p in clusters[i]]
            pos2 = [p.pos2 for p in clusters[i]]
            lr = lr_clusters[i]
            passing = True if i >= first_reject else False
            qnames = ';'.join([p.qname for p in clusters[i]])
            fout.write('{0}\t{1}\t{2}\t{3}\t{4}\t{5}\t{6}\t{7}\n'
                       .format(min(pos1), max(pos1), min(pos2), max(pos2), npairs,
                               lr, passing, qnames))
    fout.close()

def cluster_to_bp(cluster, confidence, dtype, libname):
    left_ci, right_ci = cluster_to_ci(cluster, confidence, dtype)
    pe = [(p.qname, dtype) for p in cluster]
    libs = ['pe_' + libname] * len(cluster)
    left_bp = Breakpoint(left_ci, pe=pe, libs=libs)
    right_bp = Breakpoint(right_ci, pe=pe, libs=libs)
    return left_bp, right_bp

def cluster_to_ci(cluster, confidence, dtype):
    p1 = [p.pos1 for p in cluster]
    p2 = [p.pos2 for p in cluster]
    lower_1, upper_1 = uniform_ci(p1, confidence)
    lower_2, upper_2 = uniform_ci(p2, confidence)
    if dtype == 'Del':
        return upper_1, lower_2
    elif dtype == 'Ins':
        return upper_1, lower_2
    elif dtype == 'Dup':
        return lower_1, upper_2
    elif dtype == 'InvL':
        return upper_1, upper_2
    elif dtype == 'InvR':
        return lower_1, lower_2


"""
Milliscond X-ray shutter

The shutter is an aluminum block a slot for the X-ray beam
operated by a direct drive rotary servo motor driven by an Aerotech
Soloist server motor controller with an Etherner interface.

Friedrich Schotte, Feb 17, 2017 - Jun 30, 2017
"""
__version__ = "4.0.2" # top speed: closed_pos to closed_pos2

class Ms_Shutter(object):
    name = "ms_shutter"
    from persistent_property import persistent_property

    closed_pos = persistent_property("closed_pos",-0.3) # deg
    open_pos = persistent_property("open_pos",9.7) # deg
    # alternating closed position used only in pulsed mode
    closed_pos2 = persistent_property("closed_pos2",19.7) # deg
    dt = persistent_property("dt",0.008) # total operation time in s

    def PVT(self,Topen):
        """list of times and positions for a pulse sequence given by t
        T: opening times in seconds"""
        from numpy import array
        dt = self.dt
        closed_pos,open_pos,closed_pos2 = \
            self.closed_pos,self.open_pos,self.closed_pos2
        top_speed = 2*(self.closed_pos2-self.closed_pos)/self.dt
        P,V,T = [],[],[]
        for i,t in enumerate(Topen):
            if i%2 == 0: P += [closed_pos,open_pos,closed_pos2]
            else: P += [closed_pos2,open_pos,closed_pos]
            if i%2 == 0: V += [0,top_speed,0]
            else: V += [0,-top_speed,0]
            T += [t-dt/2,t,t+dt/2]
        P,V,T = array([P,V,T])
        return P,V,T

ms_shutter = Ms_Shutter()

if __name__ == "__main__": # for testing
    from numpy import arange
    self = ms_shutter
    dt = 0.0244388571428
    Topen = dt*arange(0,10)
    print('P,V,T = self.PVT(Topen)')
    print('p,v = self.pv(Topen)')


#!/usr/bin/env python
# -*- coding: utf8 -*-

# My imports
from __future__ import division, print_function
import logging
import os
import yaml
from .utils import fun_moog_synth as func
from .observations import read_obs_intervals, plot, snr
from .minimization import MinimizeSynth, getMic, getMac
from .synthetic import read_linelist, read_linelist_elem, save_synth_spec
import time

class synthMethod:
    def __init__(self, cfgfile='StarMe_synth.cfg', overwrite=None):
        '''The function that glues everything together. A log file is created
        with the list of processes (captain.log).

        Input
        -----
        cfgfile : str
          Configuration file (default: StarMe_synth.cfg)
        overwrite : bool
          Overwrite the synthresults.dat file (default: False)

        Output
        ------
        synthresults.dat : file
          Easy readable table with results
        '''

        self.cfgfile = cfgfile
        self.overwrite = overwrite
        # Setup of logger
        if os.path.isfile('captain.log'):  # Cleaning from previous runs
            os.remove('captain.log')
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(logging.DEBUG)
        handler = logging.FileHandler('captain.log')
        handler.setLevel(logging.DEBUG)
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        handler.setFormatter(formatter)
        self.logger.addHandler(handler)

    def _setup(self, line):
        '''Do the setup with initial parameters and options.

        Input
        -----
        line : list
          Each line from the configuration file after being split at spaces
          The format is spaced separated: linelist_file (params) (options)
        '''

        self.linelist = line[0]
        if len(line) == 1:
            self.initial = [5777, 4.44, 0.00, 1.00, 3.21, 1.90]
            self._options()
        elif len(line) == 2:
            self.initial = [5777, 4.44, 0.00, 1.00, 3.21, 1.90]
            self._options(line[-1])
        elif len(line) == 7:
            self.initial = list(map(float, line[1::]))
            self.initial[0] = int(self.initial[0])
            self._options()
        elif len(line) == 8:
            self.initial = list(map(float, line[1:-1]))
            self.initial[0] = int(self.initial[0])
            self._options(line[-1])

    def _genStar(self):
        '''A generator for the configuration file.
        '''

        lines = open(self.cfgfile, 'r')
        for line in lines:
            if line.startswith(('#', ' ')):
                self.logger.debug('Skipping header: %s' % line.strip())
                continue
            self.logger.info('Line processing: %s' % line.strip())
            line = line.strip()
            line = line.split(' ')
            # Check if configuration parameters are correct
            if len(line) not in [1, 2, 7, 8]:
                self.logger.error('Could not process this information: %s' % line)
                continue

            self._setup(line)
            yield self.initial, self.options, line

    def _prepare(self):
        '''Check if linelist exists and create the first synthetic spectrum with
        initial parameters.
        '''

        if not os.path.isfile(self.linelist):
            print('The line list does not exists!\n')
            self.logger.error('The line list does not exists!\n')
            return None

        if not os.path.isfile(self.options['inter_file']):
            print('The intervals list does not exists!\n')
            self.logger.error('The intervals list does not exists!\n')
            return None

        if self.options['element']:
            self.ranges, atomic = read_linelist_elem(
                self.linelist,
                element=self.options['element'],
                intname=self.options['inter_file'],
            )
            self.logger.info('Getting initial model grid')
            self.xspec, self.yspec = func(
                self.initial,
                atmtype=self.options['model'],
                abund=self.initial[2],
                elem=self.options['element'],
                ranges=self.ranges,
                driver='synth',
                version=self.options['MOOGv'],
                **self.options
            )
        else:
            self.ranges, atomic = read_linelist(
                self.linelist, intname=self.options['inter_file']
            )
            self.logger.info('Getting initial model grid')
            self.xspec, self.yspec = func(
                self.initial,
                atmtype=self.options['model'],
                ranges=self.ranges,
                driver='synth',
                version=self.options['MOOGv'],
                **self.options
            )

        if __name__ in ('__main__', 'synthDriver'):
            self.options['GUI'] = False  # Running batch mode
        else:
            self.options['GUI'] = True  # Running GUI mode

    def _options(self, options=None):
        '''Reads the options inside the config file otherwise set to defaults.
        '''

        defaults = {
            'model': 'apogee_kurucz',
            'MOOGv': 2014,
            'save': False,
            'element': False,
            'fix_teff': False,
            'fix_logg': False,
            'fix_feh': False,
            'fix_vt': False,
            'fix_vmac': False,
            'fix_vsini': False,
            'plot': False,
            'plot_res': False,
            'damping': 1,
            'step_wave': 0.01,
            'step_flux': 3.0,
            'minimize': False,
            'refine': False,
            'observations': False,
            'inter_file': 'intervals_hr10_15n.lst',
            'snr': None,
            'resolution': None,
            'limb': 0.6,
        }
        if not options:
            self.options = defaults
        else:
            for option in options.split(','):
                if ':' in option:
                    option = option.split(':')
                    defaults[option[0]] = option[1]
                else:
                    # Clever way to change the boolean
                    if option in ['teff', 'logg', 'mh', 'vt', 'vmac', 'vsini']:
                        option = 'fix_%s' % option
                    defaults[option] = False if defaults[option] else True
            defaults['model'] = defaults['model'].lower()
            defaults['step_wave'] = float(defaults['step_wave'])
            defaults['step_flux'] = float(defaults['step_flux'])
            defaults['limb'] = float(defaults['limb'])
            defaults['MOOGv'] = int(defaults['MOOGv'])
            if defaults['observations'] and (defaults['snr'] is None):
                if os.path.isfile(str(defaults['observations'])):
                    defaults['snr'] = snr(defaults['observations'])
                elif os.path.isfile('spectra/' + str(defaults['observations'])):
                    defaults['snr'] = snr('spectra/' + str(defaults['observations']))
                else:
                    print('The SNR was not measured.')
                    self.logger.error('Error: %s not found.' % defaults['observations'])
            if defaults['inter_file']:
                defaults['inter_file'] = str(defaults['inter_file'])
            if defaults['resolution'] is not None:
                defaults['resolution'] = int(float(defaults['resolution']))
            self.options = defaults

    def _output(
        self, overwrite=None, header=None, stellarparams=False, abundance=False
    ):
        '''Create the output file 'synthresults.dat'

        Input
        -----
        overwrite - Overwrite the file
        header    - Only use True if this is for the file to be created

        Output
        -----
        The 'synthresults.dat' file with the final parameters.
        '''
        if abundance:
            if header:
                hdr2 = [
                    'linelist',
                    'observations',
                    'Teff',
                    'erTeff',
                    'logg',
                    'erlogg',
                    '[M/H]',
                    'er[M/H]',
                    'vmic',
                    'ervmic',
                    'vmac',
                    'ervmac',
                    'vsini',
                    'ervsini',
                    'Na',
                    'erNa',
                    'Mg',
                    'erMg',
                    'Al',
                    'erAl',
                    'Si',
                    'erSi',
                    'Ca',
                    'erCa',
                    'Sc',
                    'erSc',
                    'Ti',
                    'erTi',
                    'V',
                    'erV',
                    'Cr',
                    'erCr',
                    'Mn',
                    'erMn',
                    'Ni',
                    'erNi',
                    'chi^2',
                    'time',
                    'model',
                    'resolution',
                    'snr',
                ]
                if overwrite:
                    with open('synthresults_elements.dat', 'w') as output:
                        output.write('\t'.join(hdr2) + '\n')
                else:
                    if not os.path.isfile('synthresults_elements.dat'):
                        with open('synthresults_elements.dat', 'w') as output:
                            output.write('\t'.join(hdr2) + '\n')
            else:
                with open('synthresults_elements.dat', 'a') as output:
                    output.write('\t'.join(map(str, self.parameters)) + '\n')

        if stellarparams:
            if header:
                hdr1 = [
                    'linelist',
                    'observations',
                    'Teff',
                    'erTeff',
                    'logg',
                    'erlogg',
                    '[M/H]',
                    'er[M/H]',
                    'vmic',
                    'ervmic',
                    'vmac',
                    'ervmac',
                    'vsini',
                    'ervsini',
                    'chi^2',
                    'time',
                    'model',
                    'resolution',
                    'snr',
                ]

                if overwrite:
                    with open('synthresults.dat', 'w') as output:
                        output.write('\t'.join(hdr1) + '\n')
                else:
                    if not os.path.isfile('synthresults.dat'):
                        with open('synthresults.dat', 'w') as output:
                            output.write('\t'.join(hdr1) + '\n')
            else:
                with open('synthresults.dat', 'a') as output:
                    output.write('\t'.join(map(str, self.parameters)) + '\n')

    def minizationRunner(self, p=None):
        '''A function to run the minimization routine

        Output
        ------
        params : output parameters
        '''

        print('Starting minimization...')
        self.logger.info('Starting the minimization procedure...')
        start_time = time.time()
        # Run the minimization routine first time
        if p is not None:
            function = MinimizeSynth(
                p, self.xobs, self.yobs, self.ranges, **self.options
            )
            self.params, self.xo, self.yo = function.minimize()
        else:
            function = MinimizeSynth(
                self.initial, self.xobs, self.yobs, self.ranges, **self.options
            )
            self.params, self.xo, self.yo = function.minimize()

        if self.options['refine']:
            print('Refining the parameters...')
            print('Patience is the key...')
            params2 = [
                self.params[0],
                self.params[2],
                self.params[4],
                self.params[6],
                self.params[8],
                self.params[10],
            ]
            params2[3] = getMic(self.params[0], self.params[2], self.params[4])
            params2[4] = getMac(self.params[0], self.params[2])
            if round(params2[1], 2) == 5.0:
                params2[1] = 4.9
                self.options['fix_logg'] = True
            self.options['fix_vt'] = True
            self.options['fix_vmac'] = True
            function = MinimizeSynth(
                params2, self.xo, self.yo, self.ranges, **self.options
            )
            self.params, xxo, yyo = function.minimize()

        self.end_time = int(time.time() - start_time)
        print('Minimization finished in %s sec' % int(self.end_time))
        self.logger.info('Minimization done.')
        status = 1
        return status

    def minizationElementRunner(self, p=None):
        '''A function to run the minimization routine for element abundances.

        Output
        ------
        params : output parameters
        '''
        import numpy as np

        print('Starting minimization for element abundance...')
        self.logger.info('Starting minimization for element abundance...')
        start_time = time.time()
        # Run the minimization routine first time
        if p is not None:
            function = MinimizeSynth(
                p, self.xobs, self.yobs, self.ranges, **self.options
            )
            self.elemabund, self.xo, self.yo = function.minimizeElement()
        else:
            function = MinimizeSynth(
                self.initial, self.xobs, self.yobs, self.ranges, **self.options
            )
            self.elemabund, self.xo, self.yo = function.minimizeElement()

        self.end_time = int(time.time() - start_time)
        print('Minimization finished in %s sec' % int(self.end_time))
        self.logger.info('Minimization done.')

        # Make the output array to save results.
        el = ['Na', 'Mg', 'Al', 'Si', 'Ca', 'Sc', 'Ti', 'V', 'Cr', 'Mn', 'Ni']
        a = np.empty(2 * len(el) + 1)
        a[:] = np.nan
        ind1 = el.index(self.options['element'])
        ind2 = ind1 + 1
        a[ind1] = self.elemabund[0]
        a[ind2] = self.elemabund[1]
        a[-1] = self.elemabund[2]
        self.abund = a
        status = 1
        return status

    def plotRunner(self, x=None, y=None, xs=None, ys=None, xf=None, yf=None, res=False):
        '''A function to plot spectra with or without residuals.

        Output
        ------
        plots
        '''

        if self.options['minimize']:
            # Create final spectrum with final parameters.
            p1 = [
                self.params[0],
                self.params[2],
                self.params[4],
                self.params[6],
                self.params[8],
                self.params[10],
            ]
            xf, yf = func(
                p1,
                atmtype=self.options['model'],
                ranges=self.ranges,
                driver='synth',
                version=self.options['MOOGv'],
                **self.options
            )
        elif self.options['element']:
            # Create final spectrum with final parameters.
            xf, yf = func(
                self.initial,
                atmtype=self.options['model'],
                abund=self.elemabund[0],
                elem=self.options['element'],
                ranges=self.ranges,
                driver='synth',
                version=self.options['MOOGv'],
                **self.options
            )
        else:
            xf, yf = (None, None)

        if self.options['plot_res']:
            plot(x, y, xs, ys, xf, yf, res=True)
        else:
            plot(x, y, xs, ys, xf, yf)

    def saveRunner(self):
        '''A function to save spectra in a fits like format in the results/ folder.
        If initial spectrum exists, save it.

        Output
        ------
        saved spectrum
        '''

        if not os.path.isdir('results'):
            os.mkdir('results')
            self.logger.info('results directory was created')

        if self.xspec is not None:
            save_synth_spec(
                self.xspec, self.yspec, initial=self.initial, **self.options
            )
        else:
            print('No spectrum was created.')

    def synthdriver(self):
        '''A function to connect all. This function applies the options set by
        the user in the StarMe_synth.cfg file.
        '''

        # Creating the output file
        self._output(header=True, stellarparams=True, abundance=True)

        for (self.initial, self.options, line) in self._genStar():
            self.logger.info(
                'Initial parameters: {:.0f}, {:.2f}, {:.2f}, {:.2f}'.format(
                    *self.initial
                )
            )

            # Check here if element exists
            if self.options['element']:
                el = ['Na', 'Mg', 'Al', 'Si', 'Ca', 'Sc', 'Ti', 'V', 'Cr', 'Mn', 'Ni']
                if self.options['element'] not in el:
                    print(
                        'This element is not in the line list:', self.options['element']
                    )
                    self.logger.error('This element is not in the line list.')
                    continue

            self._prepare()

            if self.options is None:
                self.logger.error('The line list does not exists!\n')
                continue
            if self.options['minimize'] and self.options['element']:
                print('I am confused! First derive parameters and then abundances')
                self.logger.error(
                    'Minimization of parameters and abundances at the same time!\n'
                )
                continue

            if self.options['observations']:
                if os.path.isfile(self.options['observations']):
                    self.xobs, self.yobs, delta_l = read_obs_intervals(self.options['observations'],
                    self.ranges, snr=self.options['snr'])
                    print('Observed spectrum contains %s points' % len(self.xobs))
                    self.logger.info('Observed spectrum read.')
                elif os.path.isfile('spectra/' + self.options['observations']):
                    print(
                        'This is your observed spectrum: %s'
                        % self.options['observations']
                    )
                    self.xobs, self.yobs, delta_l = read_obs_intervals(
                        'spectra/' + self.options['observations'],
                        self.ranges,
                        snr=self.options['snr']
                    )
                    print('Observed spectrum contains %s points' % len(self.xobs))
                    self.logger.info('Observed spectrum read.')
                else:
                    print(
                        'The observed spectrum %s was not found.'
                        % self.options['observations']
                    )
                    self.logger.error(
                        'Error: %s not found.' % self.options['observations']
                    )
                    continue
            else:
                self.xobs, self.yobs = (None, None)

            if self.options['minimize']:
                status = self.minizationRunner()
                if status is None:
                    self.logger.error(
                        'The minimization routine did not finish succesfully.'
                    )
                    continue  # Problem with the minimization routine
                else:
                    self.logger.info('The minimization routine finished succesfully.')
                self.parameters = (
                    [self.linelist]
                    + [self.options['observations']]
                    + self.params
                    + [self.end_time]
                    + [
                        self.options['model'],
                        self.options['resolution'],
                        self.options['snr']
                    ]
                )
                # Save the results in the output file.
                self._output(stellarparams=True)
                self.xobs, self.yobs = self.xo, self.yo
                return self.params

            if self.options['element']:
                status = self.minizationElementRunner()
                if status is None:
                    self.logger.error(
                        'The minimization routine did not finish succesfully.'
                    )
                    continue  # Problem with the minimization routine
                else:
                    self.logger.info('The minimization routine finished succesfully.')
                self.parameters = (
                    [self.linelist]
                    + [self.options['observations']]
                    + self.initial
                    + self.abund.tolist()
                    + [self.end_time]
                    + [
                        self.options['model'],
                        self.options['resolution'],
                        self.options['snr'],
                    ]
                )
                self._output(abundance=True)
                return self.abund.tolist()

            if self.options['save']:
                self.logger.info('Save synthetic spectrum.')
                self.saveRunner()

            if self.options['plot']:
                self.logger.info('Plotting results.')
                self.plotRunner(x=self.xobs, y=self.yobs, xs=self.xspec, ys=self.yspec)

if __name__ == '__main__':
    import sys

    if len(sys.argv) > 1:
        cfgfile = sys.argv[1]
    else:
        cfgfile = 'StarMe_synth.cfg'

    driver = synthMethod(cfgfile=cfgfile, overwrite=None)
    driver.synthdriver()