import os
import sys
import random
import math
import re
import time
import numpy as np
import cv2
import matplotlib
import matplotlib.pyplot as plt
import glob
import skimage

# Root directory of the project
ROOT_DIR = os.path.abspath("../../")
print(ROOT_DIR)

from os import listdir
from os.path import isfile, join
from mrcnn.config import Config
from mrcnn import utils
import mrcnn.model as modellib
from mrcnn import visualize
from mrcnn.model import log

# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")

# Local path to trained weights file
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
# Download COCO trained weights from Releases if needed
if not os.path.exists(COCO_MODEL_PATH):
    utils.download_trained_weights(COCO_MODEL_PATH)

# Results Directory
# Save Submission files here
RESULTS_DIR = os.path.join(ROOT_DIR, "Results/Buildings/")
# Currently a temp set up, just identify valdiation set, contains 171 images from the new formed zoom 20 sample training set
VAL_IMAGE_IDS = ['nia_825a50_396_530115_484370_20.png','nia_825a50_397_530116_484370_20.png','nia_825a50_398_530117_484370_20.png','nia_825a50_399_530118_484370_20.png',
                 'nia_825a50_400_530119_484370_20.png','nia_825a50_401_530120_484370_20.png','nia_825a50_402_530121_484370_20.png','nia_825a50_403_530122_484370_20.png',
                 'nia_825a50_404_530123_484370_20.png','nia_825a50_405_530124_484370_20.png','nia_825a50_406_530125_484370_20.png','nia_825a50_407_530126_484370_20.png',
                 'nia_825a50_408_530127_484370_20.png','nia_825a50_409_530128_484370_20.png','nia_825a50_410_530129_484370_20.png','nia_825a50_411_530130_484370_20.png',
                 'nia_825a50_412_530131_484370_20.png','nia_825a50_413_530132_484370_20.png','nia_825a50_414_530133_484370_20.png','nia_825a50_415_530115_484371_20.png',
                 'nia_825a50_416_530116_484371_20.png','nia_825a50_417_530117_484371_20.png','nia_825a50_418_530118_484371_20.png','nia_825a50_419_530119_484371_20.png',
                 'nia_825a50_420_530120_484371_20.png','nia_825a50_421_530121_484371_20.png','nia_825a50_422_530122_484371_20.png','nia_825a50_423_530123_484371_20.png',
                 'nia_825a50_424_530124_484371_20.png','nia_825a50_425_530125_484371_20.png','nia_825a50_426_530126_484371_20.png','nia_825a50_427_530127_484371_20.png',
                 'nia_825a50_428_530128_484371_20.png','nia_825a50_429_530129_484371_20.png','nia_825a50_430_530130_484371_20.png','nia_825a50_431_530131_484371_20.png',
                 'nia_825a50_432_530115_484372_20.png','nia_825a50_433_530116_484372_20.png','nia_825a50_434_530117_484372_20.png','nia_825a50_435_530118_484372_20.png',
                 'nia_825a50_436_530119_484372_20.png','nia_825a50_437_530120_484372_20.png','nia_825a50_438_530121_484372_20.png','nia_825a50_439_530122_484372_20.png',
                 'nia_825a50_440_530123_484372_20.png','nia_825a50_441_530124_484372_20.png','nia_825a50_442_530125_484372_20.png','nia_825a50_443_530126_484372_20.png',
                 'nia_825a50_444_530127_484372_20.png','nia_825a50_445_530128_484372_20.png','nia_825a50_446_530129_484372_20.png','nia_825a50_447_530116_484373_20.png',
                 'nia_825a50_448_530117_484373_20.png','nia_825a50_449_530118_484373_20.png','nia_825a50_450_530119_484373_20.png','nia_825a50_451_530120_484373_20.png',
                 'nia_825a50_452_530121_484373_20.png','nia_825a50_453_530122_484373_20.png','nia_825a50_454_530123_484373_20.png','nia_825a50_455_530124_484373_20.png',
                 'nia_825a50_456_530125_484373_20.png','nia_825a50_457_530126_484373_20.png','nia_825a50_458_530127_484373_20.png','nia_825a50_459_530116_484374_20.png',
                 'nia_825a50_460_530117_484374_20.png','nia_825a50_461_530118_484374_20.png','nia_825a50_462_530119_484374_20.png','nia_825a50_463_530120_484374_20.png',
                 'nia_825a50_464_530121_484374_20.png','nia_825a50_465_530122_484374_20.png','nia_825a50_466_530123_484374_20.png','nia_825a50_467_530124_484374_20.png',
                 'nia_825a50_468_530125_484374_20.png','nia_825a50_469_530126_484374_20.png','nia_825a50_470_530116_484375_20.png','nia_825a50_471_530117_484375_20.png',
                 'nia_825a50_472_530118_484375_20.png','nia_825a50_473_530119_484375_20.png','nia_825a50_474_530120_484375_20.png','nia_825a50_475_530121_484375_20.png',
                 'nia_825a50_476_530122_484375_20.png','nia_825a50_477_530123_484375_20.png','nia_825a50_478_530124_484375_20.png','nia_825a50_479_530125_484375_20.png',
                 'nia_825a50_480_530126_484375_20.png','nia_825a50_481_530117_484376_20.png','nia_825a50_482_530118_484376_20.png','nia_825a50_483_530119_484376_20.png',
                 'nia_825a50_484_530120_484376_20.png','nia_825a50_485_530121_484376_20.png','nia_825a50_486_530122_484376_20.png','nia_825a50_487_530123_484376_20.png',
                 'nia_825a50_488_530124_484376_20.png','nia_825a50_489_530125_484376_20.png','nia_825a50_490_530126_484376_20.png','nia_825a50_491_530127_484376_20.png',
                 'nia_825a50_492_530118_484377_20.png','nia_825a50_493_530119_484377_20.png','nia_825a50_494_530120_484377_20.png','nia_825a50_495_530121_484377_20.png',
                 'nia_825a50_496_530122_484377_20.png','nia_825a50_497_530123_484377_20.png','nia_825a50_498_530124_484377_20.png','nia_825a50_499_530125_484377_20.png',
                 'nia_825a50_500_530126_484377_20.png','nia_825a50_501_530127_484377_20.png','nia_825a50_502_530128_484377_20.png','nia_825a50_503_530118_484378_20.png',
                 'nia_825a50_504_530119_484378_20.png','nia_825a50_505_530120_484378_20.png','nia_825a50_506_530121_484378_20.png','nia_825a50_507_530122_484378_20.png',
                 'nia_825a50_508_530123_484378_20.png','nia_825a50_509_530124_484378_20.png','nia_825a50_510_530125_484378_20.png','nia_825a50_511_530126_484378_20.png',
                 'nia_825a50_512_530127_484378_20.png','nia_825a50_513_530128_484378_20.png','nia_825a50_514_530129_484378_20.png','nia_825a50_515_530119_484379_20.png',
                 'nia_825a50_516_530120_484379_20.png','nia_825a50_517_530121_484379_20.png','nia_825a50_518_530122_484379_20.png','nia_825a50_519_530123_484379_20.png',
                 'nia_825a50_520_530124_484379_20.png','nia_825a50_521_530125_484379_20.png','nia_825a50_522_530126_484379_20.png','nia_825a50_523_530127_484379_20.png',
                 'nia_825a50_524_530128_484379_20.png','nia_825a50_525_530129_484379_20.png','nia_825a50_526_530120_484380_20.png','nia_825a50_527_530121_484380_20.png',
                 'nia_825a50_528_530122_484380_20.png','nia_825a50_529_530123_484380_20.png','nia_825a50_530_530124_484380_20.png','nia_825a50_531_530125_484380_20.png',
                 'nia_825a50_532_530126_484380_20.png','nia_825a50_533_530127_484380_20.png','nia_825a50_534_530128_484380_20.png','nia_825a50_535_530129_484380_20.png',
                 'nia_825a50_536_530121_484381_20.png','nia_825a50_537_530122_484381_20.png','nia_825a50_538_530123_484381_20.png','nia_825a50_539_530124_484381_20.png',
                 'nia_825a50_540_530125_484381_20.png','nia_825a50_541_530126_484381_20.png','nia_825a50_542_530127_484381_20.png','nia_825a50_543_530128_484381_20.png',
                 'nia_825a50_544_530129_484381_20.png','nia_825a50_545_530122_484382_20.png','nia_825a50_546_530123_484382_20.png','nia_825a50_547_530124_484382_20.png',
                 'nia_825a50_548_530125_484382_20.png','nia_825a50_549_530126_484382_20.png','nia_825a50_550_530127_484382_20.png','nia_825a50_551_530128_484382_20.png',
                 'nia_825a50_552_530123_484383_20.png','nia_825a50_553_530124_484383_20.png','nia_825a50_554_530125_484383_20.png','nia_825a50_555_530126_484383_20.png',
                 'nia_825a50_556_530127_484383_20.png','nia_825a50_557_530128_484383_20.png','nia_825a50_558_530124_484384_20.png','nia_825a50_559_530125_484384_20.png',
                 'nia_825a50_560_530126_484384_20.png','nia_825a50_561_530127_484384_20.png','nia_825a50_562_530125_484385_20.png','nia_825a50_563_530126_484385_20.png',
                 'nia_825a50_564_530127_484385_20.png','nia_825a50_565_530126_484386_20.png','nia_825a50_566_530127_484386_20.png',]
#################################################
# Configurations
#################################################

class BuildingConfig(Config):
    """Configuration for training on the sample Building  dataset.
    Derives from the base Config class and overrides values specific
    to the building shapes dataset. See the CONFIG FILE for more attributes
    """

    #Name config
    NAME = "Building"

    #Train on 1 GPU and 8 images per PGU. We can put multiple images on each
    #GPU because the images are small. Batch size is 8 (GPUs *images/GPU)
    # GPU_COUNT = 1 since training on CPU
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

    #Number of classes (including background)
    NUM_CLASSES = 1 + 1 #1 background + 1 Building class

    # Use small images for faster training. Set the limits of the small side
    # the large side, and that determines the image shape.
    IMAGE_MIN_DIM = 256
    IMAGE_MAX_DIM = 256

    # Use a small epoch since the data is simple
    # Total images / batch size = Steps
    # Another way to think about it train_length / batch_size
    #
    STEPS_PER_EPOCH = 200

    # use small validation steps since the epoch is small
    VALIDATION_STEPS = 10

    # Image mean (RGB)
    # Need to Update
    MEAN_PIXEL = np.array([123.7, 116.8, 103.9])

    #IMAGE_RESIZE_MODE = "none"

config = BuildingConfig()
config.display()

#################################################
# Dataset
#################################################

class BuildingDataset(utils.Dataset):

    def load_building(self, dataset_dir, subset):
        """Load a subset of the building dataset.
               dataset_dir: Root directory of the dataset.
               subset: Subset to load: train or val
               """
        # This will need to be changed based on the images loaded
        #dataset_dir = "/Users/Sebastian/Documents/GitHub/DSC672/"
        # Add classes. We have only one class to add.
        self.add_class("building", 1, "building")

        # Which Dataset are you loading in?
         # Val:Use validation test est
         # train : use data from temp_data/sample/tile -256
         # TO DO: ADD other training/testing sources
        assert subset in ["sample/", "val"]
        subset_dir = "sample/tile-256/"
        dataset_dir = os.path.join(dataset_dir, subset_dir)
        if subset == "val":
            image_ids = VAL_IMAGE_IDS
        else:
            # Get image Ids from directory names
            # print(dataset_dir) ## Remove latter DEBUG
            image_ids = next(os.walk(dataset_dir))[2]
            # image_ids = [f for f in listdir(dataset_dir) if isfile(join(dataset_dir, f))]
            # print(image_ids) ## Remove Later DEBUG
            if subset == "sample/":
                image_ids = list(set(image_ids) - set(VAL_IMAGE_IDS))

        # dir_name = join(dataset_dir + "temp_data/sample/tile-256") # for loading in regualr tiles
        # annotations_b = [f for f in listdir(dir_name) if isfile(join(dir_name, f))]

        # dir_name = join(dataset_dir + "temp_data/sample/mask-256") # for loading in mask tiles
        # annotations_m = [f for f in listdir(dir_name) if isfile(join(dir_name, f))]

        #Add images
        for image_id in image_ids:
            self.add_image(
                "building",
                image_id=image_id,
                path = os.path.join(dataset_dir, image_id.format(image_id)))

    def load_mask(self, image_id):
        info = self.image_info[image_id]
        # print("This is info ",str(info))
        # Get mask directory from image path
        mask_dir = os.path.join(os.path.dirname(os.path.dirname(info['path'])), "mask-256")
        # print("this is ", mask_dir) # REMOVE LATER FOR DEBUG
        # Read mask files from .png image
        mask = []
        # image_ids = [f for f in listdir(dataset_dir) if isfile(join(dataset_dir, f))]
        for f in next(os.walk(mask_dir))[2]:
            if f == info['id']: # match image id to mask id
                m = skimage.io.imread(os.path.join(mask_dir, f)).astype(np.bool)
                mask.append(m)

        # print("masks has {} items".format(len(mask)))
        mask = np.stack(mask, axis=-1)
        # print("mask's shape is {}".format(mask.shape))
        # Return mask, and array of class IDs of each instance. Since we have
        # one class ID, we return an array of ones
        return mask, np.ones([mask.shape[-1]], dtype=np.uint8)

    def image_reference(self, image_id):
        """Return the path of the image."""
        info = self.image_info[image_id]
        if info["source"] == "building":
            return info["id"]
        else:
            super(self.__class__, self).image_reference(image_id)

# Function used for display purposes
def get_ax(rows=1, cols=1, size=8):
    """Return a Matplotlib Axes array to be used in
    all visualizations in the notebook. Provide a
    central point to control graph sizes.

    Change the default size attribute to control the size
    of rendered images
    """
    _, ax = plt.subplots(rows, cols, figsize=(size * cols, size * rows))
    return ax

############################################
# DATA SET LOAD IN

# Training Dataset
dataset_train = BuildingDataset()

dataset_dir = "/Users/Sebastian/Documents/GitHub/DSC672/temp_data/"
subset = "sample/"
dataset_train.load_building(dataset_dir, subset)
dataset_train.prepare()

# Print out dataset information regarding the training set
print("Image Count: {}".format(len(dataset_train.image_ids)))
print("Class Count: {}".format(dataset_train.num_classes))
for i, info in enumerate(dataset_train.class_info):
    print("{:3}. {:50}".format(i, info['name']))

# Load validation dataset
dataset_val = BuildingDataset()
dataset_val.load_building(dataset_dir, "val")
dataset_val.prepare()

# Print out Validation dataset info
print("Image Count: {}".format(len(dataset_val.image_ids)))
print("Class Count: {}".format(dataset_val.num_classes))
for i, info in enumerate(dataset_val.class_info):
    print("{:3}. {:50}".format(i, info['name']))

# Load and display random samples
#image_ids = np.random.choice(dataset_train.image_ids, 4)
#for image_id in image_ids:
  #  image = dataset_train.load_image(image_id)
   # mask, class_ids = dataset_train.load_mask(image_id)
   # visualize.display_top_masks(image, mask, class_ids, dataset_train.class_names, limit= 1) #Coded provided to visualize outputs

############################################

###########################################
# TRAINING

# Create model in training mode
model = modellib.MaskRCNN(mode="training", config=config, model_dir=MODEL_DIR)

# Which weights to start with?
init_with = "coco"  # imagenet, coco, or last

if init_with == "imagenet":
    model.load_weights(model.get_imagenet_weights(), by_name=True)
elif init_with == "coco":
    # Load weights trained on MS COCO, but skip layers that
    # are different due to the different number of classes
    # See README for instructions to download the COCO weights
    model.load_weights(COCO_MODEL_PATH, by_name=True,
                       exclude=["mrcnn_class_logits", "mrcnn_bbox_fc",
                                "mrcnn_bbox", "mrcnn_mask"])
elif init_with == "last":
    # Load the last model you trained and continue training
    model.load_weights(model.find_last(), by_name=True)

#### Turn this on if you want to train the model
# Train the head branches
# Passing layers="heads" freezes all layers except the head
# layers. You can also pass a regular expression to select
# which layers to train by name pattern.
model.train(dataset_train, dataset_val,
            learning_rate=config.LEARNING_RATE,
            epochs=3,
            layers='heads')

# Training - Stage 2
# Finetune layers from ResNet stage 4 and up
print("Fine tune Resnet stage 4 and up")
model.train(dataset_train, dataset_val,
            learning_rate=config.LEARNING_RATE,
            epochs=10,
            layers='4+')

#### Turn this on if you want to train the model
# Fine tune all layers
# Passing layers="all" trains all layers. You can also
# pass a regular expression to select which layers to
# train by name pattern.
model.train(dataset_train, dataset_val,
            learning_rate=config.LEARNING_RATE / 10,
            epochs=10,
            layers="all")

###########################################
# INFERENCE MODE
inference_config = config

# Recreate the model in inference mode
model = modellib.MaskRCNN(mode="inference",
                          config=inference_config,
                          model_dir=MODEL_DIR)

# /Users/Sebastian/Documents/GitHub/logs/building20200209T1400
# Get path to saved weights
# Either set a specific path or find last trained weights
log_dir = "/Users/Sebastian/Documents/GitHub/logs/building20200213T2133/"
model_path = os.path.join(log_dir, "mask_rcnn_building_0005.h5")
# model_path = model.find_last()

# Load trained weights
print("Loading weights from ", model_path)
model.load_weights(model_path, by_name=True)

# Test on a random image
# Image_ID = 59
# Image_ID = 117 street sample
# Image Id = 42 Building Sample
# Image ID = 44
image_id = random.choice(dataset_val.image_ids)
image_id = 44

print("This is the image_ID", str(image_id)) #Find out which image is showing
original_image, image_meta, gt_class_id, gt_bbox, gt_mask =\
    modellib.load_image_gt(dataset_val, inference_config,
                           image_id, use_mini_mask=False)

log("original_image", original_image)
log("image_meta", image_meta)
log("gt_class_id", gt_class_id)
log("gt_bbox", gt_bbox)
log("gt_mask", gt_mask)

#SHOW RANDOM IMAGE with MASK ON
visualize.display_instances(original_image, gt_bbox, gt_mask, gt_class_id,
                            dataset_train.class_names, figsize=(8, 8))

# DETECT MASK FOR RANDOM IMAGE
results = model.detect([original_image], verbose=1)

r = results[0]
visualize.display_instances(original_image, r['rois'], r['masks'], r['class_ids'],
                            dataset_val.class_names, r['scores'], figsize =(8,8))

## TODO: CLEAN UP PATHS WITHIN FILE TO BE PROJECT FREIENDLY
## TODO: TUNE PARAMETERS WITHIN MASK RCNN TO FIT SAMPLE SET


"""ASE interface for MST module implemented in MuST package available at
 https://github.com/mstsuite/MuST"""

import numpy as np
from ase.calculators.calculator import FileIOCalculator, SCFError
import os
import subprocess
import glob
from ase.units import Bohr, Rydberg
from default_params import defaults
from ase.data import atomic_numbers

magmoms = {'Fe': 2.1,
           'Co': 1.4,
           'Ni': 0.6}

def write_positions_input(atoms, method):
    """ Function that writes the positions data input file based on
    atoms object and selected calculation method"""
    with open('position.dat', 'w') as filehandle:
        filehandle.write(str(1.0) + '\n\n')

        for i in range(3):
            filehandle.write('%s\n' % str(atoms.get_cell()[i] / Bohr)[1:-1])
        filehandle.write('\n')

        if method == 3:

            for site in atoms.info['CPA']:
                sitestring = 'CPA  %s' % str(atoms[site['index']].position
                                             / Bohr)[1:-1]

                for key in site.keys():
                    if key == 'index':
                        pass
                    else:
                        sitestring += '  %s %s' % (key, str(site[key]))
                sitestring += '\n'
                filehandle.write(sitestring)

        else:
            for index in range(len(atoms)):
                filehandle.write('%s %s\n'
                                 % (atoms[index].symbol,
                                    str(atoms[index].position / Bohr)[1:-1]))

def write_atomic_pot_input(symbol, nspins, moment, xc, niter, mp):
    """
    Function to write input file for generating
    atomic potential using 'newa' command
    Parameters
    ----------
    symbol: str
        Chemical symbol of the element.
    nspins: int
        Number of spins.
    moment: float
        Magnetic moment.
    xc: int
        ex-cor type (1=vb-hedin,2=vosko).
    niter: int
        Maximum number of SCF iterations.
    mp: float
        SCF mixing parameter
    """
    title = symbol + ' Atomic Potential'
    output_file = symbol + '_a_out'
    pot_file = symbol + '_a_pot'
    z = atomic_numbers[symbol]

    if moment == 0. and nspins == 2 and symbol in ['Fe', 'Co', 'Ni']:
        moment = magmoms[symbol]

    space = 20 * ' '
    contents = [('Title', title),
                (output_file,
                 'Output file name. If blank, data will show on screen'),
                (z, 'Atomic number'),
                (moment, 'Magnetic moment'),
                (nspins, 'Number of spins'),
                (xc, 'Exchange-correlation type (1=vb-hedin,2=vosko)'),
                (niter, 'Number of Iterations'),
                (mp, 'Mixing parameter'),
                (pot_file, 'Output potential file')]

    with open(symbol + '_a_in', 'w') as filehandle:
        for entry in contents:
            filehandle.write('%s %s %s\n' % (entry[0], space, entry[1]))

def write_single_site_pot_input(symbol, crystal_type, a, nspins, moment, xc,
                                lmax, print_level, ncomp, conc, mt_radius,
                                ws_radius, egrid, ef, niter, mp):
    """
    Function to write input file for generating single site
    potential using 'newss' command
    Parameters
    ----------

    symbol: str
        Chemical symbol of the element
    crystal_type: int
                1 for FCC, 2 for BCC.
    a: float
        The lattice constant.
    nspins: int
        number of spins.
    moment: float
        Magnetic moment. If nspins = 2 and moment = 0 during input,
        moment will be changed to values from
        this dictionary: {Fe': 2.1, 'Co': 1.4, 'Ni': 0.6}
    xc: int
        ex-cor type (1=vb-hedin,2=vosko).
    lmax: int
        angular momentum quantum number cutoff value.
    print_level: int
        Print level.
    ncomp: int
        Number of components.
    conc: float
        Concentrations.
    mt_radius: float
            mt_radius.
    ws_radius: float
        ws_radius.
    egrid: vector
        e-grid vector of form (ndiv(=#div/0.1Ryd), bott, eimag).
    ef: float
        Estomate of fermi energy.
    niter: int
        Maximum number of SCF iterations.
    mp: float
        Mixing parameter for SCF iterations.
    """

    title = symbol + ' Single Site Potential'
    output_file = symbol + '_ss_out'
    input_file = symbol + '_a_pot'
    pot_file = symbol + '_ss_pot'
    keep_file = symbol + '_ss_k'

    z = atomic_numbers[symbol]
    a = a / Bohr

    if moment == 0.:
        if nspins == 2:
            if symbol in ['Fe', 'Co', 'Ni']:
                moment = magmoms[symbol]

    space = '                    '
    contents = [('Title', title),
                (output_file,
                 'Output file name. If blank, data will show on screen'),
                (print_level, 'Print level'),
                (crystal_type, 'Crystal type (1=FCC,2=BCC)'),
                (lmax, 'lmax'),
                (round(a, 3), 'Lattice constant'),
                (nspins, 'Number of spins'),
                (xc, 'Exchange Correlation type (1=vb-hedin,2=vosko)'),
                (ncomp, 'Number of components'),
                (str(z) + '  ' + str(moment),
                 'Atomic number, Magnetic moment'),
                (conc, 'Concentrations'),
                (str(mt_radius / Bohr) + '  ' + str(ws_radius / Bohr),
                 'mt radius, ws radius'),
                (input_file, 'Input potential file'),
                (pot_file, 'Output potential file'),
                (keep_file, 'Keep file'),
                (str(egrid[0]) + ' ' + str(egrid[1]) + ' ' + str(egrid[2]),
                 'e-grid: ndiv(=#div/0.1Ryd), bott, eimag'),
                (str(round(ef / Rydberg, 3)) + ' '
                 + str(round(ef / Rydberg, 3)),
                 'Fermi energy (estimate)'),
                (str(niter) + ' ' + str(mp),
                 'Number of scf iterations, Mixing parameter')]

    with open(str(symbol) + '_ss_in', 'w') as filehandle:
        for entry in contents:
            filehandle.write('%s %s %s\n' % (entry[0], space, entry[1]))

def write_input_parameters_file(atoms, parameters):
    """Write the main input file for 'mst2' command. This file contains all
    essential input parameters required for calculation"""
    energy_params = ['etol', 'ptol', 'ftol',
                     'offset_energy_pt',
                     'em_switch']  # Parameters with units of energy
    spatial_params = ['liz_cutoff', 'max_core_radius',
                      'max_mt_radius', 'core_radius',
                      'mt_radius']  # Parameters with units of length
    vector_params = ['uniform_grid', 'grid_origin', 'grid_1',
                     'grid_2', 'grid_3', 'grid_pts', 'kpts',
                     'moment_direction', 'constrain_field',
                     'liz_shell_lmax', 'em_mix_param']  # vector parameters
    # Header
    hline = 80 * '='
    separator = 18 * ' ' + 3 * ('* * *' + 14 * ' ')
    header = [hline, '{:^80s}'.format('Input Parameter Data File'),
              hline, separator, hline,
              '{:^80}'.format('System Related Parameters'), hline]

    natoms = ['No. Atoms in System (> 0)  ::  '
              + str(len(atoms)), hline, separator, hline]

    # Get number of atoms from CPA sites if self.parameters['method'] == 3
    if 'method' in parameters.keys():
        if parameters['method'] == 3:
            natoms = ['No. Atoms in System (> 0)  ::  '
                      + str(len(atoms.info['CPA'])), hline, separator, hline]

    header += natoms

    with open('i_new', 'w') as filehandle:
        for entry in header:
            filehandle.write('%s\n' % entry)

    # Rest of the parameters:
    contents = []

    for key in parameters.keys():
        if key in energy_params:
            parameters[key] = parameters[key] / Rydberg

        if key in spatial_params:
            parameters[key] = parameters[key] / Bohr

        if key in vector_params:
            parameters[key] = str(parameters[key])[1:-1]

        if key == 'in_pot':
            for index in parameters['in_pot'].keys():
                contents.append(defaults[key] + '  ::  '
                                + index + ' ' + parameters['in_pot'][index])
        else:
            contents.append(defaults[key] + '  ::  ' + str(parameters[key]))

    with open('i_new', 'a') as filehandle:
        for entry in contents:
            filehandle.write('%s\n' % entry)

def generate_starting_potentials(atoms, crystal_type, a, cpa=False, nspins=1,
                                 moment=0., xc=1, lmax=3, print_level=1,
                                 ncomp=1, conc=1., mt_radius=0., ws_radius=0,
                                 egrid=(10, -0.4, 0.3), ef=9.5, niter=50,
                                 mp=0.1):
    """
    Function to generate single site starting potentials for
    all elements in an atoms object

    Parameters
    ----------
    atoms: The Atoms object to generate starting potentials.
    crystal_type: int
        1 for FCC, 2 for BCC.
    a: float
        The lattice constant.
    cpa: bool
        If True, use atoms.info['CPA'] to generate starting potentials
        for all elements in CPA formalism
    nspins: int
        number of spins.
    moment: float
        Magnetic moment. If nspins = 2 and moment = 0 during input,
        moment will be changed to values from this
        dictionary: {Fe': 2.1, 'Co': 1.4, 'Ni': 0.6}
    xc: int
        ex-cor type (1=vb-hedin,2=vosko).
    lmax: int
        angular momentum quantum number cutoff value.
    print_level: int
        Print level.
    ncomp: int
        Number of components.
    conc: float
        Concentrations.
    mt_radius: float
        mt_radius.
    ws_radius: float
        ws_radius.
    egrid: vector
        e-grid vector of form (ndiv(=#div/0.1Ryd), bott, eimag).
    ef: float
        Estimate of fermi energy.
    niter: int
        Maximum number of SCF iterations.
    mp: float
        Mixing parameter for SCF iterations.
    """
    if cpa:
        species = []
        for site in atoms.info['CPA']:
            for element in site:
                if element == 'index':
                    pass
                else:
                    species.append(element)
        species = np.unique(species)
    else:
        species = np.unique(atoms.get_chemical_symbols())

    for symbol in species:
        # Generate atomic potential
        write_atomic_pot_input(symbol, nspins=nspins, moment=moment,
                               xc=xc, niter=niter, mp=mp)

        subprocess.run('newa', stdin=open(symbol + '_a_in'), check=True)

        # Generate single site potential
        write_single_site_pot_input(symbol=symbol, crystal_type=crystal_type,
                                    a=a, nspins=nspins, moment=moment, xc=xc,
                                    lmax=lmax, print_level=print_level,
                                    ncomp=ncomp, conc=conc,
                                    mt_radius=mt_radius,
                                    ws_radius=ws_radius,
                                    egrid=egrid, ef=ef, niter=niter, mp=mp)

        subprocess.run('newss', stdin=open(symbol + '_ss_in'), check=True)

class MuST(FileIOCalculator):
    """
    Multiple Scattering Theory based ab-initio calculator.
    Capable of performing LSMS, KKR and KKR-CPA calculations.
    """

    implemented_properties = ['energy']
    command = 'mst2 < i_new'

    def __init__(self, restart=None, ignore_bad_restart_file=False,
                 label='mst', atoms=None, **kwargs):
        FileIOCalculator.__init__(self, restart, ignore_bad_restart_file,
                                  label, atoms, **kwargs)

        if 'ntasks' in self.parameters:
            self.command = 'mpirun -np ' + str(self.parameters['ntasks']) \
                           + ' ' + self.command

    def clean(self):
        """
        Delete the input parameters file, positions data file
        and default output files generated by MuST
        """

        files = ['o_n00000_D', 'k_n00000_D', 'position.dat', 'i_new']
        for f in files:
            try:
                os.remove(f)
            except OSError:
                pass

    def write_input(self, atoms, properties=None, system_changes=None):
        """ Write input files"""

        FileIOCalculator.write_input(self, atoms,
                                     properties=None, system_changes=None)
        # Write positions using CPA sites if self.parameters['method'] == 3
        if 'method' in self.parameters.keys():
            if self.parameters['method'] == 3:
                method = self.parameters['method']
            else:
                method = None

        else:
            method = None

        write_positions_input(atoms, method=method)
        write_input_parameters_file(atoms=atoms, parameters=self.parameters)

    def read_results(self):
        """ Read results from output files"""
        outfile = glob.glob('k_n00000_D')[0]
        with open(outfile, 'r') as file:
            lines = file.readlines()

        e_offset = float(lines[7].split()[-1])

        results = {tag: value for tag, value in zip(lines[9].split(),
                                                    lines[-1].split())}
        read_energy = (float(results['Energy']) + e_offset)

        convergence = False

        outfile = glob.glob('o_n00000_D')[0]
        with open(outfile, 'r') as file:
            lines = file.readlines()

        for line in lines:
            if 'SCF Convergence is reached' in line:
                convergence = True
                break

        if not convergence:
            raise SCFError()
        else:
            self.results['energy'] = read_energy * Rydberg


import numpy as np
import argparse
from time import time, strftime
import pygame
import go
import playmodel
import cProfile

parser = argparse.ArgumentParser()
parser.add_argument("--size", "-s", dest="size", default=19, type=int)
parser.add_argument("--playouts", "-p", dest="playouts", default=1024, type=int)
parser.add_argument("--selfplay", "-S", dest="self_play", action="store_true")
parser.add_argument("--use-model", "-m", dest="use_model", type=str, default="", action="store")
parser.add_argument("--playas", type=str, dest="playas", default="random", action="store")

args = parser.parse_args()

WHITE = go.WHITE
BLACK = go.BLACK
COLOR = go.COLOR
BOARD_SIZE = args.size

if args.playas == "b":
    PLAYAS = BLACK
elif args.playas == "w":
    PLAYAS = WHITE
else:
    PLAYAS = np.random.randint(2)

BACKGROUND = 'images/ramin.jpg'
GRID_SIZE = 20

class Stone(go.Stone):
    def __init__(self, board, point):
        """Create, initialize and draw a stone."""
        super(Stone, self).__init__(board, point)
        if self.islegal:
            board.draw_stones()

class Board(go.Board):
    def __init__(self, size, komi):
        """Create, initialize and draw an empty board.
        """
        super(Board, self).__init__(size, komi)
        self.outline = pygame.Rect(GRID_SIZE+5, GRID_SIZE+5, DRAW_BOARD_SIZE[0]-GRID_SIZE*2, DRAW_BOARD_SIZE[1]-GRID_SIZE*2)
        self.draw_board()

    def draw_board(self):
        """Draw the board to the background and blit it to the screen.

        The board is drawn by first drawing the outline, then the 19x19
        grid and finally by adding hoshi to the board. All these
        operations are done with pygame's draw functions.

        This method should only be called once, when initializing the
        board.
        """
        pygame.draw.rect(background, BLACK, self.outline, 3)
        # Outline is inflated here for future use as a collidebox for the mouse
        self.outline.inflate_ip(20, 20)
        for i in range(self.size-1):
            for j in range(self.size-1):
                rect = pygame.Rect(5+GRID_SIZE+(GRID_SIZE*i), 5+GRID_SIZE+(GRID_SIZE*j), GRID_SIZE, GRID_SIZE)
                pygame.draw.rect(background, COLOR[BLACK], rect, 1)
        if self.size >= 13:
            for i in range(3):
                for j in range(3):
                    coords = (5+4*GRID_SIZE+(GRID_SIZE*6*i), 5+4*GRID_SIZE+(GRID_SIZE*6*j))
                    pygame.draw.circle(background, COLOR[BLACK], coords, 5, 0)
        screen.blit(background, (0, 0))
        pygame.display.update()

    def draw_stones(self):
        """
        This method is called every time the board is updated by add_stone() and it is succesful
        """
        screen.blit(background, (0, 0))
        for g in self.groups:
            for p in g.stones:
                coords = (5+GRID_SIZE+p[0]*GRID_SIZE, 5+GRID_SIZE+p[1]*GRID_SIZE)
                pygame.draw.circle(screen, COLOR[g.color], coords, GRID_SIZE//2, 0)
        pygame.display.update()

    def write_game_log(self, logfile):
        COLOR_CHAR = ("B", "W")
        logstr = ""
        for entry in self.log:
            logstr += COLOR_CHAR[entry[0]] + " " + str(entry[1]) + "\n"
        logfile.write(logstr)
        win, score_diff, outstr = self.score(output=True)
        logfile.write(outstr+"\n")

def test(ai_play_as):
    print("begin test")
    pass_count = 0
    steps = 0
    while True:
        pygame.time.wait(250)
        if ai_play_as == board.next:
            x, y = model.decide_monte_carlo(board, args.playout)
            if x >= BOARD_SIZE or y >= BOARD_SIZE:
                print("model passes")
                pass_count += 1
                board.pass_move()
            else:
                print("model choose (%d, %d)" % (x, y))
                pass_count = 0
                added_stone = Stone(board, (x, y))
                if not added_stone.islegal:
                    print("model tried illegal move")
                    break
            steps += 1
        else:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    exit()
                if event.type == pygame.MOUSEBUTTONDOWN:
                    if event.button == 1 and board.outline.collidepoint(event.pos):
                        x = int(round(((event.pos[0]-GRID_SIZE-5) / GRID_SIZE), 0))
                        y = int(round(((event.pos[1]-GRID_SIZE-5) / GRID_SIZE), 0))
                        added_stone = Stone(board, (x, y))
                        pass_count = 0
                        steps += 1
                        print("player choose (%d, %d)"%(x, y))
                if event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_p:
                        pass_count += 1
                        break
        if pass_count >= 2:
            break
    print("game over")
    winner, score_diff, out_str = board.score(output=True)
    print(out_str)
    # end while True

def self_play():
    print("begin self play")
    pass_count = 0
    measure_times = []
    game_over = False
    while not game_over and pass_count < 2:
        pygame.time.wait(250)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                exit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_p:
                    game_over = True
                    continue
                t = time()
                # cProfile.run("x, y = model.decide(board, args.playouts)", "mcts_search.profile")
                x, y = model.decide_monte_carlo(board, args.playouts, output = True)
                policy, value = model.get_policy_value(board.grid, board.next)
                print("time", time()-t)
                measure_times.append(time()-t)
                if y == BOARD_SIZE:
                    board.pass_move()
                    play_as = "B" if board.next == WHITE else "W"
                    print(play_as, "pass (p,v) =", policy, value)
                    pass_count += 1
                    if pass_count >= 2: break
                elif y < 0:
                    board.pass_move()
                    play_as = "B" if board.next == WHITE else "W"
                    print(play_as, "resigned (p,v) =", policy, value)
                    game_over = True
                else:
                    added_stone = Stone(board, (x, y))
                    if not added_stone.islegal:
                        print("B" if play_as == "W" else "W", "tried illegal move (%d, %d)"%(x, y))
                        # game_over = True
                        # break
                    else:
                        pass_count = 0
                        play_as = "B" if board.next == WHITE else "W"
                        print(play_as, "choose (%d, %d)"%(x, y), "(p,v) =", policy, value)

    print("game over")
    winner, score_diff, out_str = board.score(output=True)
    print(out_str)
    print("decision time", np.sum(measure_times), np.mean(measure_times), np.std(measure_times))

if __name__ == '__main__':

    model = playmodel.ActorCritic(args.size, args.use_model)
    BOARD_SIZE = model.board_size
    DRAW_BOARD_SIZE = (GRID_SIZE * BOARD_SIZE + 20, GRID_SIZE * BOARD_SIZE + 20)
    if BOARD_SIZE <= 9:
        KOMI = 5.5
    elif BOARD_SIZE <= 13:
        KOMI = 6.5
    else:
        KOMI = 7.5

    print("Board size:", BOARD_SIZE)
    print("Komi:", KOMI)

    if args.self_play:
        pygame.init()
        pygame.display.set_caption('Go-Ai')
        screen = pygame.display.set_mode(DRAW_BOARD_SIZE, 0, 32)
        background = pygame.image.load(BACKGROUND).convert()
        board = Board(size=BOARD_SIZE, komi=KOMI)
        self_play()
    else:
        pygame.init()
        pygame.display.set_caption('Go-Ai')
        screen = pygame.display.set_mode(DRAW_BOARD_SIZE, 0, 32)
        background = pygame.image.load(BACKGROUND).convert()
        board = Board(size=BOARD_SIZE, komi=KOMI)
        test(PLAYAS)


"""Contains the merge generator class."""

from flow.core.generator import Generator

from numpy import pi, sin, cos

INFLOW_EDGE_LEN = 100  # length of the inflow edges (needed for resets)
VEHICLE_LENGTH = 5

class MergeGenerator(Generator):
    """Generator for merge networks."""

    def __init__(self, net_params, base):
        merge = net_params.additional_params["merge_length"]
        premerge = net_params.additional_params["pre_merge_length"]
        postmerge = net_params.additional_params["post_merge_length"]
        length = merge + premerge + postmerge + 2 * INFLOW_EDGE_LEN + 8.1
        h_lanes = net_params.additional_params["highway_lanes"]
        m_lanes = net_params.additional_params["merge_lanes"]
        self.name = "{}-{}m{}l-{}l".format(base, length, h_lanes, m_lanes)
        super().__init__(net_params, base)

    def specify_nodes(self, net_params):
        """See parent class."""
        angle = pi / 4
        merge = net_params.additional_params["merge_length"]
        premerge = net_params.additional_params["pre_merge_length"]
        postmerge = net_params.additional_params["post_merge_length"]

        nodes = [
            {
                "id": "inflow_highway",
                "x": repr(-INFLOW_EDGE_LEN),
                "y": repr(0)
            },
            {
                "id": "left",
                "y": repr(0),
                "x": repr(0)
            },
            {
                "id": "center",
                "y": repr(0),
                "x": repr(premerge)
            },
            {
                "id": "right",
                "y": repr(0),
                "x": repr(premerge + postmerge)
            },
            {
                "id": "inflow_merge",
                "x": repr(premerge - (merge + INFLOW_EDGE_LEN) * cos(angle)),
                "y": repr(-(merge + INFLOW_EDGE_LEN) * sin(angle))
            },
            {
                "id": "bottom",
                "x": repr(premerge - merge * cos(angle)),
                "y": repr(-merge * sin(angle))
            },
        ]

        return nodes

    def specify_edges(self, net_params):
        """See parent class."""
        merge = net_params.additional_params["merge_length"]
        premerge = net_params.additional_params["pre_merge_length"]
        postmerge = net_params.additional_params["post_merge_length"]

        edges = [{
            "id": "inflow_highway",
            "type": "highwayType",
            "from": "inflow_highway",
            "to": "left",
            "length": repr(INFLOW_EDGE_LEN)
        }, {
            "id": "left",
            "type": "highwayType",
            "from": "left",
            "to": "center",
            "length": repr(premerge)
        }, {
            "id": "inflow_merge",
            "type": "mergeType",
            "from": "inflow_merge",
            "to": "bottom",
            "length": repr(INFLOW_EDGE_LEN)
        }, {
            "id": "bottom",
            "type": "mergeType",
            "from": "bottom",
            "to": "center",
            "length": repr(merge)
        }, {
            "id": "center",
            "type": "highwayType",
            "from": "center",
            "to": "right",
            "length": repr(postmerge)
        }]

        return edges

    def specify_types(self, net_params):
        """See parent class."""
        h_lanes = net_params.additional_params["highway_lanes"]
        m_lanes = net_params.additional_params["merge_lanes"]
        speed = net_params.additional_params["speed_limit"]

        types = [{
            "id": "highwayType",
            "numLanes": repr(h_lanes),
            "speed": repr(speed)
        }, {
            "id": "mergeType",
            "numLanes": repr(m_lanes),
            "speed": repr(speed)
        }]

        return types

    def specify_routes(self, net_params):
        """See parent class."""
        rts = {
            "inflow_highway": ["inflow_highway", "left", "center"],
            "left": ["left", "center"],
            "center": ["center"],
            "inflow_merge": ["inflow_merge", "bottom", "center"],
            "bottom": ["bottom", "center"]
        }

        return rts


import logging
import dynet as dy
from baseline.model import (
    TaggerModel,
    register_model
)
import numpy as np
from baseline.dy.dynety import CRF, Linear, DynetModel, rnn_forward, ConvEncoderStack, viterbi
from baseline.utils import Offsets

logger = logging.getLogger('baseline')

class TaggerModelBase(DynetModel, TaggerModel):
    def __init__(self, embeddings_set, labels, **kwargs):
        super(TaggerModelBase, self).__init__(kwargs['pc'])
        self.labels = labels
        self.hsz = int(kwargs['hsz'])
        self.pdrop = kwargs.get('dropout', 0.5)
        self.rnntype = kwargs.get('rnntype', 'blstm')
        self.do_crf = bool(kwargs.get('crf', False))
        self.constraint = kwargs.get('constraint')
        self.lengths_key = kwargs.get('lengths_key')
        dsz = self.init_embed(embeddings_set)
        nc = len(self.labels)

        if self.do_crf:
            self.crf = CRF(
                nc, self.pc, idxs=(Offsets.GO, Offsets.EOS), mask=self.constraint
            )

        self.activation_type = kwargs.get('activation', 'tanh')
        self.init_encoder(dsz, **kwargs)
        self.output = self.init_output(self.hsz, nc)
        self.train = True

    def init_encoder(self, input_sz, **kwargs):
        pass

    def dropout(self, input_):
        if self.train:
            return dy.dropout(input_, self.pdrop)
        return input_

    def __str__(self):
        str_ = super(TaggerModelBase, self).__str__()
        return "Auto-batching: \n{}".format(str_)

    def init_embed(self, embeddings):
        dsz = 0
        self.embeddings = embeddings
        for embedding in self.embeddings.values():
            dsz += embedding.get_dsz()
        return dsz

    def embed(self, batch_dict):
        all_embeddings_lists = []
        for k, embedding in self.embeddings.items():
            all_embeddings_lists.append(embedding.encode(batch_dict[k], self.train))

        embedded = dy.concatenate(all_embeddings_lists, d=1)
        embed_list = [self.dropout(e) for e in embedded]
        return embed_list

    def make_input(self, batch_dict):
        example_dict = dict({})

        for k, embedding in self.embeddings.items():
            example_dict[k] = batch_dict[k].T

        lengths = batch_dict[self.lengths_key]
        example_dict['lengths'] = lengths

        y = batch_dict.get('y')
        if y is not None:
            example_dict['y'] = y

        ids = batch_dict.get('ids')
        if ids is not None:
            example_dict['ids'] = ids

        return example_dict

    def compute_unaries(self, batch_dict):
        embed_list = self.embed(batch_dict)
        exps = self.encode(embed_list)
        return exps

    def encode(self, embed_list):
        pass

    def predict(self, batch_dict):
        dy.renew_cg()
        inputs = self.make_input(batch_dict)
        lengths = inputs['lengths']
        unaries = self.compute_unaries(inputs)
        if self.do_crf is True:
            best_path, path_score = self.crf.decode(unaries)
        elif self.constraint is not None:
            best_path, path_score = viterbi(
                unaries,
                dy.log_softmax(dy.inputTensor(self.constraint[1] * -1e4)),
                Offsets.GO, Offsets.EOS,
                norm=True
            )
        else:
            best_path = [np.argmax(x.npvalue(), axis=0) for x in unaries]
        # TODO: RN using autobatching, so none of this is really useful
        # If we want to support batching in this function we have to either loop over the batch
        # or we can just simplify all this code here
        best_path = np.stack(best_path).reshape(-1, 1)  # (T, B)

        best_path = best_path.transpose(1, 0)
        results = []

        for b in range(best_path.shape[0]):
            sentence = best_path[b, :lengths[b]]
            results.append(sentence)
        return results

    def loss(self, preds, y):
        if self.do_crf is True:
            return self.crf.neg_log_loss(preds, y.squeeze(0))
        else:
            element_loss = dy.pickneglogsoftmax
            errs = []

            for pred, y_i in zip(preds, y.T):
                err = element_loss(pred, y_i)
                errs.append(err)
            return dy.esum(errs)

    def dropout(self, input_):
        if self.train:
            return dy.dropout(input_, self.pdrop)
        return input_

    def init_output(self, input_dim, n_classes):
        return Linear(n_classes, input_dim, self.pc, name="output")

    @classmethod
    def create(cls, embeddings_set, labels, **kwargs):
        model = cls(embeddings_set, labels, **kwargs)
        logger.info(model)
        return model

    def save(self, file_name):
        self.pc.save(file_name)
        return self

    def load(self, file_name):
        self.pc.populate(file_name)
        return self

@register_model(task='tagger', name='default')
class RNNTaggerModel(TaggerModelBase):

    def __init__(self, embeddings_set, labels, **kwargs):
        super(RNNTaggerModel, self).__init__(embeddings_set, labels, **kwargs)

    def init_encoder(self, input_sz, **kwargs):
        layers = kwargs.get('layers', 1)
        self.rnn = dy.BiRNNBuilder(layers, input_sz, self.hsz, self.pc, dy.VanillaLSTMBuilder)

    def encode(self, embed_list):
        return [self.output(out) for out in rnn_forward(self.rnn, embed_list)]

@register_model(task='tagger', name='cnn')
class CNNTaggerModel(TaggerModelBase):

    def __init__(self, *args, **kwargs):
        super(CNNTaggerModel, self).__init__(*args, **kwargs)

    def init_encoder(self, input_sz, **kwargs):
        layers = int(kwargs.get('layers', 1))
        pdrop = float(kwargs.get('dropout', 0.5))
        filtsz = kwargs.get('wfiltsz', 5)
        activation_type = kwargs.get('activation_type', 'relu')
        hsz = int(kwargs['hsz'])
        self.encoder = ConvEncoderStack(filtsz, hsz, input_sz, pdrop, self.pc, layers, activation_type)
        return hsz

    def encode(self, embed_list):
        embed_list = dy.transpose(dy.concatenate_cols(embed_list))
        return [self.output(out) for out in self.encoder(embed_list, self.train)]


import numpy as np

class Dynamics:
	"""docstring for Dynamics"""
	def __init__(self, dt):
		self.dt = dt

	def create_motion_model_params(motion_model_cov):
		raise NotImplementedError

	def create_meas_model_params(meas_cov):
		raise NotImplementedError

if __name__ == '__main__':
	main()


#HeaderInfo
#type=ROITool
#name=Import from CSV
#version=2021.08.05
#author=aheyd@berkeley.edu
#EndHeaderInfo

import osirix
import dicom
import numpy as np
import shutil
import os
import subprocess
import csv

csv_file_path = "./Users/aarash/Documents/annotations.csv"

vc = osirix.frontmostViewer()
pixs = vc.pixList()
instanceToPixPosition = {pix.imageObj().instanceNumber:i for i, pix in enumerate(pixs)}

def import_from_csvs():
    vc = osirix.frontmostViewer()
    for scans, dcmseriesuid in [load_scan(x.paths()) for x in osirix.currentBrowser().databaseSelection()]:
        spacing, origin = scans.values()[0].PixelSpacing, scans.values()[0].ImagePositionPatient
        with open(csv_file_path) as csvfile:
            csvreader = csv.reader(csvfile, delimiter=' ', quotechar='|')
            columns = next(csvreader) # Read the first row, which is names of columns rather than data.

            # This dict records number of overlays already in a slice, where key = slice Z position
            NextOverlayIndex = {}
            for i, row in enumerate(csvreader):
                try:
                     s = row[0].split(",")
                     csvseriesuid, coordX, coordY, coordZ, diameter_mm = s[0], s[1], s[2], s[3].strip("\'"), float(s[4])
                except:
                     print("Failed with len row: ", len(row[0].split(",")), "at i = %d" % i)
                     print(row[0])
                     continue
                if dcmseriesuid == csvseriesuid:
                     print("equality at csvrow = %d" % i)
                     vX, vY = world_to_voxel(float(coordX), float(coordY), spacing, origin)
                     print("ROI at voxel coordinates", coordX, coordY, vX, vY, coordZ)
                     Zlist = [z for z in scans.keys()]
                     scanZ = sorted(Zlist, key=lambda x: np.abs(float(x) - np.round(float(coordZ.rstrip('0')))))[0]
                     print(scanZ)
                     ds = scans[scanZ]
                     overlayIndex = NextOverlayIndex.get(scanZ, 0)
                     mask = create_circular_mask(512, 512, vX, vY, radius=diameter_mm/2*(1/ds.PixelSpacing[0]))
                     embed(mask, ds, overlayIndex, OVERLAY_CSV_NAME_SUFFIX)
                     roiNew = osirix.ROI(itype='tPlain',buffer=mask.T.astype(np.bool),name="ROI read from CSV row %d" % i,ipixelSpacing=spacing, iimageOrigin=origin[:2])
                     vc.setROI(roiNew,position=instanceToPixPosition[ds.InstanceNumber])
                     NextOverlayIndex[scanZ] = overlayIndex + 1
                     ds.save_as(pixs[instanceToPixPosition[ds.InstanceNumber]].sourceFile)
                     print("saved file in %s for instance %f with overlayIndex = %d, radius=%f\n\n\n" % (pixs[instanceToPixPosition[ds.InstanceNumber]].sourceFile, ds.InstanceNumber, overlayIndex, diameter_mm))

#############
### Utils ###
#############
OVERLAY_CSV_NAME_SUFFIX = "Annotation read from CSV"
# Embeds a mask into a DICOM file at the given overlay index
def embed(mask, dcm_file, overlay_index, name):
    if overlay_index > 16:
        return # cannot support more than 16 overlays.
    # packbits converts array of integer 1s and 0s to array of numbers.
    # It interprets a run of 8 bits as a number.
    # i.e. [1, 0, 0, 0, 0, 0, 0, 0] -> [128]
    # The reshaping and flattening is to accomodate DICOM Header's expected format.
    w = np.packbits(mask.reshape(-1,8)[:,::-1].flatten('C'))
    dcm_file.add_new(0x60000010 + overlay_index*0x20000 , 'US', 512)
    dcm_file.add_new(0x60000011 + overlay_index*0x20000 , 'US', 512)
    dcm_file.add_new(0x60000022 + overlay_index*0x20000 , 'LO', name)
    dcm_file.add_new(0x60000040 + overlay_index*0x20000 , 'CS', 'R')
    dcm_file.add_new(0x60000050 + overlay_index*0x20000 , 'SS', [1,1])
    dcm_file.add_new(0x60000100 + overlay_index*0x20000 , 'US', 1)
    dcm_file.add_new(0x60000102 + overlay_index*0x20000 , 'US', 0)
    dcm_file.add_new(0x60003000 + overlay_index*0x20000 , 'OW', w)

# Source: https://stackoverflow.com/questions/44865023/circular-masking-an-image-in-python-using-numpy-arrays
def create_circular_mask(h, w, centerX, centerY, radius):
    # Grid of indices
    Y, X = np.ogrid[:h, :w]
    # Grid of distances
    dist_from_center = np.sqrt((X - centerX)**2 + (Y-centerY)**2)
    # Grid of 1s and 0s for distances <= radius
    mask = np.array([[1 if dist <= radius else 0 for dist in dist_row] for dist_row in dist_from_center])
    return mask

# Returns the list of DICOM files for the patient sorted by Z position
def load_scan(paths):
   paths2 = [s for s in paths if '.dcm' in s and hasattr(dicom.read_file(s, force=True), "InstanceNumber") and hasattr(dicom.read_file(s, force=True), "ImagePositionPatient")]
   slices = [dicom.read_file(s, force=True) for s in paths2]
   slices2 = [x for x in slices if hasattr(x, "InstanceNumber") and hasattr(x, "ImagePositionPatient")]
   slices=slices2
   indices = sorted(list(range(len(paths2))), key=lambda x: slices[x].InstanceNumber)
   sortedPaths = np.array(paths)[indices]
   slices.sort(key=lambda x: int(x.InstanceNumber))
   slices2 = {float(str(x.ImagePositionPatient[2]).rstrip('0')): x for x in slices}
   pathDict = {float(str(x.ImagePositionPatient[2]).rstrip('0')): sortedPaths[count] for count, x in enumerate(slices)}
   return slices2, slices[0].SeriesInstanceUID

def add_row_to_csv_file(csvpath, seriesuid, x, y, z, diam_mm):
    columns = ["seriesuid", "coordX", "coordY", "coordZ", "diameter_mm"]
    with open(csvpath, 'ab') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow([seriesuid, x, y, z, diam_mm])

def area_to_diameter(area):
    return 2 * np.sqrt(area / np.pi)

# Note: The relation between world coordinates and voxel
# coordinates is as follows:
# World coordinate (0, 0) ~= voxel coordinate (256, 256)
# in the middle of the photo
# Positive X coordinates are the patient's right, which is
# is the viewer's left side of an image, which trends toward
# voxel-X-coordinate 0
# Contrarywise, as world-X -> negative values :: voxel-X -> 512
# Similarly for Y coordinates,
# world-Y -> positive (anterior) :: voxel-Y -> 0
# world-Y -> negative (posterior) :: voxel-Y -> 512
#
# Because of this inverse relationship, both voxel_to_world
# and world_to_voxel negate the coordinate input.
@staticmethod
def world_to_voxel(wX, wY, spacing, origin):
    vX = -(wX - origin[0]) / spacing[0]
    vY = -(wY - origin[1]) / spacing[1]
    return round(vX), round(vY)

if __name__ == "__main__":
    import_from_csvs()
    vc.needsDisplayUpdate()


#!/usr/bin/python()
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import numpy as np
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--folder', '-f', help='Path to data folder')
parser.add_argument("--dest")
args = parser.parse_args()
BASE_DIR = args.folder

object_dict = {
    'clutter':   0,
    'ceiling':   1,
    'floor':     2,
    'wall':      3,
    'beam':      4,
    'column':    5,
    'door':      6,
    'window':    7,
    'table':     8,
    'chair':     9,
    'sofa':     10,
    'bookcase': 11,
    'board':    12
}

path_Dir_Areas = os.listdir(BASE_DIR)

for Area in path_Dir_Areas:
    if not os.path.isdir(os.path.join(BASE_DIR, Area)):
        continue
    path_Dir_Rooms = os.listdir(os.path.join(BASE_DIR, Area))
    for Room in path_Dir_Rooms:

        xyz_Room = np.zeros((1, 6))
        label_Room = np.zeros((1, 1))
        path_Annotations = os.path.join(BASE_DIR, Area, Room, "Annotations")
        if not os.path.exists(path_Annotations):
            print("Error {} does not exists".format(path_Annotations))
            continue

        print(path_Annotations)
        # make store directories
        path_prepare_label = os.path.join(args.dest, Area, Room)
        if not os.path.exists(path_prepare_label):
            os.makedirs(path_prepare_label)
        path_objects = os.listdir(path_Annotations)
        for Object in path_objects:
            if Object.split("_", 1)[0] in object_dict:
                print(Object.split("_", 1)[0] + " value:", object_dict[Object.split("_", 1)[0]])
                xyz_object = np.loadtxt(os.path.join(path_Annotations, Object))[:, :]#(N,6)
                label_object = np.tile([object_dict[Object.split("_", 1)[0]]], (xyz_object.shape[0], 1))#(N,1)
            else:
                continue

            xyz_Room = np.vstack((xyz_Room, xyz_object))
            label_Room = np.vstack((label_Room, label_object))

        xyz_Room = np.delete(xyz_Room, [0], 0)
        label_Room = np.delete(label_Room, [0], 0)

        np.save(path_prepare_label+"/xyzrgb.npy", xyz_Room.astype(np.float16))
        np.save(path_prepare_label+"/label.npy", label_Room)


from picamera.array import PiRGBArray
from picamera import PiCamera
import numpy as np
import os
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile
import cv2
import threading
import time
import RPi.GPIO as GPIO
import serial

from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image

GPIO.setmode(GPIO.BCM)
GPIO.setup(18, GPIO.OUT)
pwm = GPIO.PWM(18, 100)
pwm.start(50)

# sudo modprobe uvcvideo nodrop=1 timeout=6000
p_det = False
image_det = None

f = open("mscoco_label_map.pbtxt","r")
data = f.read()
r1=data.split('item')
r1.pop(0)

map_label = {}
for i in range(len(r1)):
    id = r1[i].split('id:')[1].split('\n')[0].strip()
    name = r1[i].split('display_name:')[1].split('\n')[0].strip().replace('"','')
    map_label[int(id)] = name

print map_label

detection_graph = tf.Graph()
with detection_graph.as_default():
    od_graph_def = tf.GraphDef()
    with tf.gfile.GFile('frozen_inference_graph.pb', 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')

def load_image_into_numpy_array(image):
    im_height, im_width = image.shape[:2]
    aux = np.array(image).reshape((im_height, im_width, 3)).astype(np.uint8)
    return aux

def th_camera():
    global image_det, p_det
    camera = cv2.VideoCapture(0) # usar webcam
    camera.set(cv2.cv.CV_CAP_PROP_FPS, 5)
    camera.set(cv2.cv.CV_CAP_PROP_FRAME_WIDTH, 320);
    camera.set(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT, 240);
    #camera = PiCamera()
    #camera.resolution = (320, 240)
    #camera.framerate = 4
    #capture = PiRGBArray(camera, size=(320, 240))
    #for frame in camera.capture_continuous(capture, format="bgr", use_video_port=True):
    while (True):
        ret, image = camera.read()
        cv2.imshow('odet', image)
        #image = frame.array
        #frame.truncate(0)

        if(p_det == False):
            image_det = image
            p_det = True
        if cv2.waitKey(1) == 27:
            break

def th_detect():
    global image_det, p_det
    err = 0
    err0 = 0
    px = 0
    py = 0
    pm = 0
    with detection_graph.as_default():
        with tf.Session(graph=detection_graph) as sess:

            image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
            detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
            detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
            detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
            num_detections = detection_graph.get_tensor_by_name('num_detections:0')

            while True:

                if(p_det == True):
                    print '...'
                    image_np = load_image_into_numpy_array(image_det)
                    image_np_expanded = np.expand_dims(image_np, axis=0)
                    (boxes, scores, classes, num) = sess.run([detection_boxes, detection_scores, detection_classes, num_detections],
                      feed_dict={image_tensor: image_np_expanded})

                    result = [[map_label[int(item)],box] for item, per, box in zip(classes[0],scores[0],boxes[0]) if per >= 0.65]
                    if(len(result) > 0):
                        height, width = image_det.shape[:2]
                        for obj in result:
                            xmin = int(obj[1][1]*width) # posicao do objeto
                            ymin = int(obj[1][0]*height)
                            xmax = int(obj[1][3]*width)
                            ymax = int(obj[1][2]*height)
                            px, py = xmin+(xmax-xmin)/2, ymin+(ymax-ymin)/2 # centro do objeto
                            cv2.rectangle(image_det, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
                            cv2.circle(image_det,(px,py), 5, (0,0,255), -1)
                            #print xmin,ymin,xmax,ymax
                            err = width/2 - px
                            print obj[0]
                            print px, py
                            break
                        cv2.imshow('odet', image_det)
                    p_det = False

threading.Thread(target=th_camera).start()
threading.Thread(target=th_detect).start()


from __future__ import absolute_import, division, print_function, unicode_literals
import logging
import copy
import numpy as np
import astropy.units as u
from ..utils.scripts import make_path
from ..utils.fitting import fit_minuit
from .. import stats
from .utils import CountsPredictor
from . import SpectrumObservationList, SpectrumObservation

__all__ = [
    'SpectrumFit',
]

log = logging.getLogger(__name__)

class SpectrumFit(object):
    """Orchestrate a 1D counts spectrum fit.

    After running the :func:`~gammapy.spectrum.SpectrumFit.fit` and
    :func:`~gammapy.spectrum.SpectrumFit.est_errors` methods, the fit results
    are available in :func:`~gammapy.spectrum.SpectrumFit.result`. For usage
    examples see :ref:`spectral_fitting`

    Parameters
    ----------
    obs_list : `~gammapy.spectrum.SpectrumObservationList`, `~gammapy.spectrum.SpectrumObservation`
        Observation(s) to fit
    model : `~gammapy.spectrum.models.SpectralModel`
        Source model. Should return counts if ``forward_folded`` is False and a flux otherwise
    stat : {'wstat', 'cash'}
        Fit statistic
    forward_folded : bool, default: True
        Fold ``model`` with the IRFs given in ``obs_list``
    fit_range : tuple of `~astropy.units.Quantity`
        Fit range, will be convolved with observation thresholds. If you want to
        control which bins are taken into account in the fit for each
        observation, use :func:`~gammapy.spectrum.PHACountsSpectrum.quality`
    background_model : `~gammapy.spectrum.models.SpectralModel`, optional
        Background model to be used in cash fits
    method : {'sherpa', 'iminuit'}
        Optimization backend for the fit
    err_method : {'sherpa'}
        Optimization backend for error estimation
    """

    def __init__(self, obs_list, model, stat='wstat', forward_folded=True,
                 fit_range=None, background_model=None,
                 method='sherpa', err_method='sherpa'):
        self.obs_list = obs_list
        self._model = model
        self.stat = stat
        self.forward_folded = forward_folded
        self.fit_range = fit_range
        self._background_model = background_model
        self.method = method
        self.err_method = err_method

        self._predicted_counts = None
        self._statval = None

        self.covar_axis = None
        self.covariance = None
        self._result = None

        self._check_valid_fit()
        self._apply_fit_range()

    def __str__(self):
        ss = self.__class__.__name__
        ss += '\nSource model {}'.format(self.model)
        ss += '\nStat {}'.format(self.stat)
        ss += '\nForward Folded {}'.format(self.forward_folded)
        ss += '\nFit range {}'.format(self.fit_range)
        if self.background_model is not None:
            ss += '\nBackground model {}'.format(self.background_model)
        ss += '\nBackend {}'.format(self.method)
        ss += '\nError Backend {}'.format(self.err_method)

        return ss

    @property
    def result(self):
        """Fit result

        The result is a list of length ``n``, where ``n`` ist the number of
        observations that participated in the fit. The best fit model is
        usually the same for all observations but the results differ in the
        fitted energy range, predicted counts, final fit statistic value
        etc.
        """
        return self._result

    @property
    def model(self):
        """Source model

        The model parameters change every time the likelihood is evaluated. In
        order to access the best-fit model parameters, use
        :func:`gammapy.spectrum.SpectrumFit.result`
        """
        return self._model

    @model.setter
    def model(self, model):
        self._model = model

    @property
    def background_model(self):
        """Background model

        The model parameters change every time the likelihood is evaluated. In
        order to access the best-fit model parameters, use
        :func:`gammapy.spectrum.SpectrumFit.result`
        """
        return self._background_model

    @background_model.setter
    def background_model(self, model):
        self._background_model = model

    @property
    def obs_list(self):
        """Observations participating in the fit"""
        return self._obs_list

    @obs_list.setter
    def obs_list(self, obs_list):
        if isinstance(obs_list, SpectrumObservation):
            obs_list = SpectrumObservationList([obs_list])

        self._obs_list = SpectrumObservationList(obs_list)

    @property
    def bins_in_fit_range(self):
        """Bins participating in the fit for each observation."""
        return self._bins_in_fit_range

    @property
    def predicted_counts(self):
        """Current value of predicted counts.

        For each observation a tuple to counts for the on and off region is
        returned.
        """
        return self._predicted_counts

    @property
    def statval(self):
        """Current value of statval.

        For each observation the statval per bin is returned.
        """
        return self._statval

    @property
    def fit_range(self):
        """Fit range."""
        return self._fit_range

    @fit_range.setter
    def fit_range(self, fit_range):
        self._fit_range = fit_range
        self._apply_fit_range()

    @property
    def true_fit_range(self):
        """True fit range for each observation.

        True fit range is the fit range set in the
        `~gammapy.spectrum.SpectrumFit` with observation threshold taken into
        account.
        """
        true_range = []
        for binrange, obs in zip(self.bins_in_fit_range, self.obs_list):
            idx = np.where(binrange)[0]
            if len(idx) == 0:
                true_range.append(None)
                continue
            e_min = obs.e_reco[idx[0]]
            e_max = obs.e_reco[idx[-1] + 1]
            fit_range = u.Quantity((e_min, e_max))
            true_range.append(fit_range)
        return true_range

    def _apply_fit_range(self):
        """Mark bins within desired fit range for each observation."""
        self._bins_in_fit_range = []
        for obs in self.obs_list:
            # Take into account fit range
            energy = obs.e_reco
            valid_range = np.zeros(energy.nbins)

            if self.fit_range is not None:

                precision = 1e-3  # to avoid floating round precision
                idx_lo = np.where(energy * (1 + precision) < self.fit_range[0])[0]
                valid_range[idx_lo] = 1

                idx_hi = np.where(energy[:-1] * (1 - precision) > self.fit_range[1])[0]
                if len(idx_hi) != 0:
                    idx_hi = np.insert(idx_hi, 0, idx_hi[0] - 1)
                valid_range[idx_hi] = 1

            # Take into account thresholds
            try:
                quality = obs.on_vector.quality
            except AttributeError:
                quality = np.zeros(obs.e_reco.nbins)

            convolved = np.logical_and(1 - quality, 1 - valid_range)

            self._bins_in_fit_range.append(convolved)

    def predict_counts(self):
        """Predict counts for all observations.

        The result is stored as ``predicted_counts`` attribute.
        """
        predicted_counts = []
        for obs in self.obs_list:
            mu_sig = self._predict_counts_helper(obs,
                                                 self.model,
                                                 self.forward_folded)
            mu_bkg = None
            if self.background_model is not None:
                # For now, never fold background model with IRFs
                mu_bkg = self._predict_counts_helper(obs,
                                                     self.background_model,
                                                     False)
            counts = [mu_sig, mu_bkg]
            predicted_counts.append(counts)
        self._predicted_counts = predicted_counts

    def _predict_counts_helper(self, obs, model, forward_folded=True):
        """Predict counts for one observation.

        Parameters
        ----------
        obs : `~gammapy.spectrum.SpectrumObservation`
            Response functions
        model : `~gammapy.spectrum.SpectralModel`
            Source or background model
        forward_folded : bool, default: True
            Fold model with IRFs

        Returns
        ------
        predicted_counts: `np.array`
            Predicted counts for one observation
        """
        predictor = CountsPredictor(model=model)
        if forward_folded:
            predictor.aeff = obs.aeff
            predictor.edisp = obs.edisp
        else:
            predictor.e_true = obs.e_reco

        predictor.livetime = obs.livetime

        predictor.run()
        counts = predictor.npred.data.data

        # Check count unit (~unit of model amplitude)
        if counts.unit.is_equivalent(''):
            counts = counts.value
        else:
            raise ValueError('Predicted counts {}'.format(counts))

        # Apply AREASCAL column
        counts *= obs.on_vector.areascal

        return counts

    def calc_statval(self):
        """Calc statistic for all observations.

        The result is stored as attribute ``statval``, bin outside the fit
        range are set to 0.
        """
        statval = []
        for obs, npred in zip(self.obs_list, self.predicted_counts):
            on_stat, off_stat = self._calc_statval_helper(obs, npred)
            statvals = (on_stat, off_stat)
            statval.append(statvals)
        self._statval = statval
        self._restrict_statval()

    def _calc_statval_helper(self, obs, prediction):
        """Calculate ``statval`` for one observation.

        Parameters
        ----------
        obs : `~gammapy.spectrum.SpectrumObservation`
            Measured counts
        prediction : tuple of `~numpy.ndarray`
            Predicted (on counts, off counts)

        Returns
        ------
        statsval : tuple of `~numpy.ndarray`
            Statval for (on, off)
        """
        stats_func = getattr(stats, self.stat)

        # Off stat = 0 by default
        off_stat = np.zeros(obs.e_reco.nbins)
        if self.stat == 'cash' or self.stat == 'cstat':
            if self.background_model is not None:
                mu_on = prediction[0] + prediction[1]
                on_stat = stats_func(n_on=obs.on_vector.data.data.value,
                                     mu_on=mu_on)
                mu_off = prediction[1] / obs.alpha
                off_stat = stats_func(n_on=obs.off_vector.data.data.value,
                                      mu_on=mu_off)
            else:
                mu_on = prediction[0]
                on_stat = stats_func(n_on=obs.on_vector.data.data.value,
                                     mu_on=mu_on)
                off_stat = np.zeros_like(on_stat)

        elif self.stat == 'wstat':
            kwargs = dict(n_on=obs.on_vector.data.data.value,
                          n_off=obs.off_vector.data.data.value,
                          alpha=obs.alpha,
                          mu_sig=prediction[0])
            # Store the result of the profile likelihood as bkg prediction
            mu_bkg = stats.get_wstat_mu_bkg(**kwargs)
            prediction[1] = mu_bkg * obs.alpha
            on_stat_ = stats_func(**kwargs)
            # The on_stat sometime contains nan values
            # TODO: Handle properly
            on_stat = np.nan_to_num(on_stat_)
            off_stat = np.zeros_like(on_stat)
        else:
            raise NotImplementedError('{}'.format(self.stat))

        return on_stat, off_stat

    def total_stat(self, parameters):
        """Statistic summed over all bins and all observations.

        This is the likelihood function that is passed to the optimizers

        Parameters
        ----------
        parameters : `~gammapy.utils.fitting.ParameterList`
            Model parameters
        """
        self.model.parameters = parameters
        self.predict_counts()
        self.calc_statval()
        total_stat = np.sum([np.sum(v) for v in self.statval], dtype=np.float64)
        return total_stat

    def _restrict_statval(self):
        """Apply valid fit range to statval.
        """
        for statval, valid_range in zip(self.statval, self.bins_in_fit_range):
            # Find bins outside safe range
            idx = np.where(np.invert(valid_range))[0]
            statval[0][idx] = 0
            statval[1][idx] = 0

    def _check_valid_fit(self):
        """Helper function to give useful error messages."""
        # Assume that settings are the same for all observations
        test_obs = self.obs_list[0]
        irfs_exist = test_obs.aeff is not None or test_obs.edisp is not None
        if self.forward_folded and not irfs_exist:
            raise ValueError('IRFs required for forward folded fit')
        if self.stat == 'wstat' and self.obs_list[0].off_vector is None:
            raise ValueError('Off vector required for WStat fit')
        try:
            test_obs.livetime
        except KeyError:
            raise ValueError('No observation livetime given')

    def likelihood_1d(self, model, parname, parvals):
        """Compute likelihood profile.

        TODO: Replace by something more generic

        Parameters
        ----------
        model : `~gammapy.spectrum.models.SpectralModel`
            Model to draw likelihood profile for
        parname : str
            Parameter to calculate profile for
        parvals : `~astropy.units.Quantity`
            Parameter values
        """
        likelihood = []
        self._model = model
        for val in parvals:
            self.model.parameters[parname].value = val
            stat = self.total_stat(self.model.parameters)
            likelihood.append(stat)
        return np.array(likelihood)

    def plot_likelihood_1d(self, ax=None, **kwargs):
        """Plot 1-dim likelihood profile.

        See :func:`~gammapy.spectrum.SpectrumFit.likelihood_1d`
        """
        import matplotlib.pyplot as plt
        ax = plt.gca() if ax is None else ax

        yy = self.likelihood_1d(**kwargs)
        ax.plot(kwargs['parvals'], yy)
        ax.set_xlabel(kwargs['parname'])

        return ax

    def fit(self):
        """Run the fit."""
        if self.method == 'sherpa':
            self._fit_sherpa()
        elif self.method == 'iminuit':
            self._fit_iminuit()
        else:
            raise NotImplementedError('method: {}'.format(self.method))

    def _fit_sherpa(self):
        """Wrapper around sherpa minimizer."""
        from sherpa.fit import Fit
        from sherpa.data import Data1DInt
        from sherpa.optmethods import NelderMead
        from .sherpa_utils import SherpaModel, SherpaStat

        binning = self.obs_list[0].e_reco
        # The sherpa data object is not usued in the fit. It is set to the
        # first observation for debugging purposes, see below
        data = self.obs_list[0].on_vector.data.data.value
        data = Data1DInt('Dummy data', binning[:-1].value,
                         binning[1:].value, data)
        # DEBUG
        # from sherpa.models import PowLaw1D
        # from sherpa.stats import Cash
        # model = PowLaw1D('sherpa')
        # model.ref = 0.1
        # fit = Fit(data, model, Cash(), NelderMead())

        # NOTE: We cannot use the Levenbergr-Marquart optimizer in Sherpa
        # because it relies on the fvec return value of the fit statistic (we
        # return None). The computation of fvec is not straightforwad, not just
        # stats per bin. E.g. for a cash fit the sherpa stat computes it
        # according to cstat
        # see https://github.com/sherpa/sherpa/blob/master/sherpa/include/sherpa/stats.hh#L122

        self._sherpa_fit = Fit(data,
                               SherpaModel(self),
                               SherpaStat(self),
                               NelderMead())
        fitresult = self._sherpa_fit.fit()
        log.debug(fitresult)
        self._make_fit_result(self.model.parameters)

    def _fit_iminuit(self):
        """Iminuit minimization"""
        parameters, minuit = fit_minuit(parameters=self.model.parameters,
                                        function=self.total_stat)
        log.debug(minuit)
        self._make_fit_result(parameters)

    def _make_fit_result(self, parameters):
        """Bundle fit results into `~gammapy.spectrum.SpectrumFitResult`.

        Parameters
        ----------
        parameters : `~gammapy.utils.modeling.ParameterList`
            Best fit parameters
        """
        from . import SpectrumFitResult

        # run again with best fit parameters
        self.total_stat(parameters)
        model = self.model.copy()

        if self.background_model is not None:
            bkg_model = self.background_model.copy()
        else:
            bkg_model = None

        statname = self.stat

        results = []
        for idx, obs in enumerate(self.obs_list):
            fit_range = self.true_fit_range[idx]
            statval = np.sum(self.statval[idx])
            stat_per_bin = self.statval[idx]
            npred_src = copy.deepcopy(self.predicted_counts[idx][0])
            npred_bkg = copy.deepcopy(self.predicted_counts[idx][1])

            results.append(SpectrumFitResult(
                model=model,
                fit_range=fit_range,
                statname=statname,
                statval=statval,
                stat_per_bin=stat_per_bin,
                npred_src=npred_src,
                npred_bkg=npred_bkg,
                background_model=bkg_model,
                obs=obs
            ))

        self._result = results

    def est_errors(self):
        """Estimate parameter errors."""
        if self.err_method == 'sherpa':
            self._est_errors_sherpa()
        else:
            raise NotImplementedError('{}'.format(self.err_method))

        for res in self.result:
            res.covar_axis = self.covar_axis
            res.covariance = self.covariance
            res.model.parameters.set_parameter_covariance(self.covariance, self.covar_axis)

    def _est_errors_sherpa(self):
        """Wrapper around Sherpa error estimator."""
        covar = self._sherpa_fit.est_errors()
        self.covar_axis = [par.split('.')[-1] for par in covar.parnames]
        self.covariance = copy.deepcopy(covar.extra_output)

    def run(self, outdir=None):
        """Run all steps and write result to disk.

        Parameters
        ----------
        outdir : Path, str
            directory to write results files to (if given)
        """
        log.info('Running {}'.format(self))

        self.fit()
        self.est_errors()

        if outdir is not None:
            self._write_result(outdir)

    def _write_result(self, outdir):
        outdir = make_path(outdir)
        outdir.mkdir(exist_ok=True, parents=True)

        # Assume only one model is fit to all data
        modelname = self.result[0].model.__class__.__name__
        filename = outdir / 'fit_result_{}.yaml'.format(modelname)
        log.info('Writing {}'.format(filename))
        self.result[0].to_yaml(filename)


from txros import util
import tf
import rospy
import numpy as np
from mil_ros_tools import rosmsg_to_numpy
import visualization_msgs.msg as visualization_msgs
from visualization_msgs.msg import Marker, MarkerArray
from geometry_msgs.msg import Point, Vector3
from mil_misc_tools import FprintFactory

MISSION = 'Torpedo Challenge'

class Target(object):

    def __init__(self):
        self.position = None
        self.destroyed = False

    def set_destroyed(self):
        self.destroyed = True

    def update_position(self, pos):
        self.position = pos

class FireTorpedos(object):
    '''
    Mission to solve the torpedo RoboSub challenge.

    This code was based off of the Buoy mission code written by Kevin Allen.
    Its goal is to search for a target on the torpedo board and fire at it.
    '''
    TIMEOUT_SECONDS = 30
    Z_PATTERN_RADIUS = 0.5
    Y_PATTERN_RADIUS = 1.0
    X_OFFSET = .2
    Z_OFFSET = .2
    BACKUP_METERS = 3.0
    BLIND = True

    def __init__(self, sub):
        self.sub = sub
        self.print_info = FprintFactory(title=MISSION).fprint
        self.print_bad = FprintFactory(title=MISSION, msg_color="red").fprint
        self.print_good = FprintFactory(
            title=MISSION, msg_color="green").fprint
        self.tf_listener = tf.TransformListener()

        # B = bottom; T = Top; L = left; R = right; C = center; O = unblocked;
        # X = blocked;
        self.targets = {
            'TCX': Target(),
            'TRX': Target(),
            'TLX': Target(),
            'BCO': Target()
        }
        self.pattern_done = False
        self.done = False
        self.ltime = None
        self.generate_pattern()

    def generate_pattern(self):
        z = self.Z_PATTERN_RADIUS
        y = self.Y_PATTERN_RADIUS
        self.moves = [[0, 0, -z], [0, y, 0], [0, 0, -2 * z], [0, -2 * y, 0]]
        self.move_index = 0

    @util.cancellableInlineCallbacks
    def search(self):
        global markers
        markers = MarkerArray()
        pub_markers = yield self.sub.nh.advertise('/torpedo/rays', MarkerArray)
        while True:
            info = 'CURRENT TARGETS: '

            target = 'BCO'
            pub_markers.publish(markers)
            '''
            In the event we want to attempt other targets beyond bare minimum
            for target in self.targets. Eventually we will want to take the
            best target, which is the TCX target. Priority targetting will
            come once we confirm we can actually unblock currently blocked
            targets.
            '''
            res = yield self.sub.vision_proxies.arm_torpedos.get_pose(
                target='board')
            if res.found:
                self.ltime = res.pose.header.stamp
                self.targets[target].update_position(
                    rosmsg_to_numpy(res.pose.pose.position))
                marker = Marker(
                    ns='torp_board',
                    action=visualization_msgs.Marker.ADD,
                    type=Marker.ARROW,
                    scale=Vector3(0.2, 0.5, 0),
                    points=np.array([Point(0, 0, 0),
                                     res.pose.pose.position]))
                marker.id = 3
                marker.header.frame_id = '/base_link'
                marker.color.r = 1
                marker.color.g = 0
                marker.color.a = 1
                markers.markers.append(marker)
            if self.targets[target].position is not None:
                info += target + ' '
            yield self.sub.nh.sleep(0.5)  # Throttle service calls
            self.print_info(info)

    @util.cancellableInlineCallbacks
    def pattern(self):
        self.print_info('Descending to Depth...')
        yield self.sub.move.depth(1.5).go(blind=self.BLIND, speed=0.1)
        def err():
            self.print_info('Search pattern canceled')

        self.pattern_done = False
        for i, move in enumerate(self.moves[self.move_index:]):
            move = self.sub.move.relative(np.array(move)).go(blind=self.BLIND, speed=0.1)
            move.addErrback(err)
            yield move
            self.move_index = i + 1
        self.print_bad('Pattern finished. Firing at any locked targets.')
        self.pattern_done = True

    @util.cancellableInlineCallbacks
    def fire(self, target):
        self.print_info("FIRING {}".format(target))
        target_pose = self.targets[target].position
        yield self.sub.move.go(blind=self.BLIND, speed=0.1)  # Station hold
        transform = yield self.sub._tf_listener.get_transform('/map', '/base_link')
        target_position = transform._q_mat.dot(
                target_pose - transform._p)

        sub_pos = yield self.sub.tx_pose()
        print('Current Sub Position: ', sub_pos)

        # sub_pos = transform._q_mat.dot(
                # (sub_pos[0]) - transform._p)
        # target_position = target_position - sub_pos[0]
        # yield self.sub.move.look_at_without_pitching(target_position).go(
            # blind=self.BLIND, speed=.1)
        print('Map Position: ', target_position)
        yield self.sub.move.relative(np.array([0, target_position[1], 0])).go(blind=True, speed=.1)
        yield self.sub.move.relative(np.array([target_position[0] + self.X_OFFSET, 0, target_position[2] + self.Z_OFFSET])).go(
            blind=self.BLIND, speed=.1)

        self.print_good(
            "{} locked. Firing torpedos. Hit confirmed, good job Commander.".
            format(target))
        sub_pos = yield self.sub.tx_pose()
        print('Current Sub Position: ', sub_pos)
        yield self.sub.actuators.shoot_torpedo1()
        yield self.sub.actuators.shoot_torpedo2()
        self.done = True

    def get_target(self):
        target = 'BCO'
        '''
        Returns the target we are going to focus on. Loop through priorities.
        Highest priority is the TCX target, followed by the TRX and TLX.
        Lowest priority is the BCO target. Targets are ordered in the list
        already, so this will take the first target it can actually find.
        Currently this will always by the BCO target. This is because we don't
        differentiate between targets currently as we cannot unblock the
        blocked ones. This means there is only one target for us to lock onto.
        '''
        if self.targets[target].destroyed:
            pass
            # temp = target
        elif self.targets[target].position is not None:
            return target
        elif self.pattern_done and self.targets[target].position is not None:
            return target
        else:
            return None

    @util.cancellableInlineCallbacks
    def run(self):
        # start_time = yield self.sub.nh.get_time()  # Store time mission
        # starts for timeout

        self.print_info("Enabling Perception")
        self.sub.vision_proxies.arm_torpedos.start()

        pattern = self.pattern()
        self.do_search = True
        search = self.search()
        while not self.done:
            t = self.get_target()
            if t is not None:
                pattern.cancel()
                yield self.fire(t)
                self.targets[t].set_destroyed()
                if not self.done:
                    pattern = self.pattern()
            elif self.pattern_done:
                break
            yield self.sub.nh.sleep(0.1)
        search.cancel()
        pattern.cancel()
        self.sub.vision_proxies.arm_torpedos.stop()
        self.print_good('Done!')

@util.cancellableInlineCallbacks
def run(sub):
    # print('running')
    rospy.init_node('arm_torpedos', anonymous=False)
    mission = FireTorpedos(sub)
    yield mission.run()


import numpy as np
import soundfile as sf
import yaml

import tensorflow as tf

from tensorflow_tts.inference import AutoConfig
from tensorflow_tts.inference import TFAutoModel
from tensorflow_tts.inference import AutoProcessor

import zipfile

def main():

  # initialize fastspeech model.
  fs_config = AutoConfig.from_pretrained('./examples/fastspeech/conf/fastspeech.v1.yaml')
  fastspeech = TFAutoModel.from_pretrained(
      config=fs_config,
      pretrained_path="./examples/fastspeech/pretrained/model-195000.h5"
  )

  # initialize melgan model
  melgan_config = AutoConfig.from_pretrained('./examples/melgan/conf/melgan.v1.yaml')
  melgan = TFAutoModel.from_pretrained(
      config=melgan_config,
      pretrained_path="./examples/melgan/checkpoints/generator-1500000.h5"
  )

  # inference
  processor = AutoProcessor.from_pretrained(pretrained_path="./test/files/ljspeech_mapper.json")

  # text to be converted
  #ids = processor.text_to_sequence("Hi boss, the current version of our TTS has been able to convert text into speech. I am still working on APIs for easier usage.")
  with open('./Text.txt', 'r') as file:
    data = file.read().replace('\n', '')

  ids = processor.text_to_sequence(data)
  ids = tf.expand_dims(ids, 0)

  # fastspeech inference
  masked_mel_before, masked_mel_after, duration_outputs = fastspeech.inference(
    ids,
    speaker_ids=tf.zeros(shape=[tf.shape(ids)[0]], dtype=tf.int32),
    speed_ratios=tf.constant([1.0], dtype=tf.float32)
  )

  # melgan inference
  #audio_before = melgan.inference(masked_mel_before)[0, :, 0]
  audio_after = melgan.inference(masked_mel_after)[0, :, 0]

  # save to file
  #sf.write('./audio_before.wav', audio_before, 22050, "PCM_16")
  sf.write('./speech.wav', audio_after, 22050, "PCM_16")

if __name__=='__main__':
    main()


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import os
import sys

import tensorflow as tf
import importlib
import itertools
import argparse
from tensorflow.python.framework import graph_util
import numpy as np
import tensorflow.contrib.slim as slim

FLAGS = None

def main(intput_model_name,output_model_name):

    network = importlib.import_module('inception_resnet_v1')
    with tf.Graph().as_default():
        with tf.Session() as sess:
            input_tensor = tf.placeholder(dtype=tf.float32, shape=(1, 112, 112, 3), name='input')
            # Load the model metagraph and checkpoint
            #meta_file ="./ToBeConvertModels/%s.cpkt.meta"%intput_model_name
            cpkt_file="./ToBeConvertModels/%s.ckpt-90"%intput_model_name

			# Build the inference graph
            prelogits, _ = network.inference(input_tensor, 1.0, phase_train=False, bottleneck_layer_size=512,weight_decay=0.0)

            embeddings = tf.nn.l2_normalize(prelogits, 1, 1e-10, name='embeddings')
            saver = tf.train.Saver()
            tf.get_default_session().run(tf.global_variables_initializer())
            tf.get_default_session().run(tf.local_variables_initializer())
            saver.restore(tf.get_default_session(), cpkt_file)

            # Retrieve the protobuf graph definition and fix the batch norm nodes
            input_graph_def = sess.graph.as_graph_def()
            # Freeze the graph def
            output_graph_def = freeze_graph_def(sess, input_graph_def, 'embeddings')
        # Serialize and dump the output graph to the filesystem
        with tf.gfile.GFile("%s.pb"%output_model_name, 'wb') as f:
            f.write(output_graph_def.SerializeToString())

def freeze_graph_def(sess, input_graph_def, output_node_names):
    for node in input_graph_def.node:
        if node.op == 'RefSwitch':
            node.op = 'Switch'
            for index in xrange(len(node.input)):
                if 'moving_' in node.input[index]:
                    node.input[index] = node.input[index] + '/read'
        elif node.op == 'AssignSub':
            node.op = 'Sub'
            if 'use_locking' in node.attr: del node.attr['use_locking']
        elif node.op == 'AssignAdd':
            node.op = 'Add'
            if 'use_locking' in node.attr: del node.attr['use_locking']
        elif node.op == 'ArgMax':
            print (" network conain ArgMax ops, which is not support by snpe" )

    # Get the list of important nodes
    whitelist_names = []
    for node in input_graph_def.node:
        if (node.name.startswith('InceptionResnetV1') or node.name.startswith('embeddings') or node.name.startswith('phase_train') or node.name.startswith('Bottleneck') or node.name.startswith('Logits')):
            whitelist_names.append(node.name)

    # Replace all the variables in the graph with constants of the same values
    output_graph_def = graph_util.convert_variables_to_constants(
        sess, input_graph_def, output_node_names.split(","),
        variable_names_whitelist=whitelist_names)
    return output_graph_def

if __name__ == '__main__':
    parser=argparse.ArgumentParser()
    parser.add_argument("-i","--intput_model_name",type=str,help="intput model name",default="")
    parser.add_argument("-o","--output_model_name",type=str,help="output model name",default="result")
    args=parser.parse_args(sys.argv[1:])
    main( args.intput_model_name,args.output_model_name )